From 3ff28aa25fb8dcdb9651d47266134c4ab0ed3771 Mon Sep 17 00:00:00 2001 From: bowencui123 Date: Sat, 29 Aug 2026 03:02:47 +0000 Subject: [PATCH 1/8] nki(weight_dequant): NKI (Trainium) implementation Split out of the consolidated NKI branch cecilia/feature/nki-vector-add (nki-all-operators, PR #259) so each operator can be reviewed on its own. Supersedes PR #192 (older per-operator branch). - also carries the operator's `impl_torch.py` change from the NKI branch Co-Authored-By: Cecilia123li <68335867+Cecilia123li@users.noreply.github.com> Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_012Q38kGmXvyoeM1qtCbheSL --- .../operators/weight_dequant/impl_nki.py | 120 ++++++++++++++++++ .../operators/weight_dequant/impl_torch.py | 9 +- 2 files changed, 126 insertions(+), 3 deletions(-) create mode 100644 benchmarks/operators/weight_dequant/impl_nki.py diff --git a/benchmarks/operators/weight_dequant/impl_nki.py b/benchmarks/operators/weight_dequant/impl_nki.py new file mode 100644 index 00000000..e52abc17 --- /dev/null +++ b/benchmarks/operators/weight_dequant/impl_nki.py @@ -0,0 +1,120 @@ +import torch + +try: + import nki + import nki.language as nl + import nki.isa as nisa + PMAX = nl.tile_size.pmax +except ImportError: + nki = None + + +def kernel_assert(condition: bool, error_text: str): + """Assert with NKI-formatted error message.""" + assert condition, f"[INTERNAL_ERROR] [NCC_INKI016] Kernel validation exception: {error_text}" + + +def div_ceil(n: int, d: int) -> int: + """Ceiling division: smallest integer >= n/d.""" + return (n + d - 1) // d + + +if nki is not None: + @nki.jit + def weight_dequant_kernel(X, S, TILE_SIZE): + """Block-wise weight dequantization: ``out[m, n] = X[m, n] * S[m // T, n // T]``. + + ``S`` holds one scale per ``TILE_SIZE x TILE_SIZE`` block of ``X``. + + Args: + X: [M, N] quantized weights in HBM. + S: [ceil(M/T), ceil(N/T)] per-block scales in HBM (same dtype as X). + TILE_SIZE: block edge length T, a multiple of PMAX (checked in ``run``). + + Returns: + [M, N] tensor in HBM with the same dtype as ``X``. + + Notes: + * Because ``TILE_SIZE`` is a multiple of ``PMAX``, a ``PMAX``-row tile never + straddles two scale rows, so one scale row of ``S`` serves every row tile + inside a ``TILE_SIZE``-row band. + * The scalar broadcast is done in two cheap steps rather than by + materializing a full [PMAX, TILE_SIZE] scale tile: + 1. one DMA per scale row with a partition stride of 0 + (``S.ap(pattern=[[0, PMAX], [1, s_cols]])``) replicates the whole + scale row to all 128 partitions in a single instruction; + 2. ``nisa.tensor_scalar`` with ``operand0`` = the [PMAX, 1] column for + that block broadcasts along the free axis in hardware. + ``operand0`` must be float32 (the MLIR verifier rejects a half-precision + ``operand0``), hence the fp32 copy of the broadcast row. + * Boundary tiles are clamped (tile sized to the surviving extent) instead of + masked, so no out-of-range element is ever loaded, multiplied or stored. + """ + kernel_assert(len(X.shape) == 2, "X must be 2D [M, N]") + kernel_assert(len(S.shape) == 2, "S must be 2D [M/T, N/T]") + + M, N = X.shape + s_rows = div_ceil(M, TILE_SIZE) + s_cols = div_ceil(N, TILE_SIZE) + kernel_assert(S.shape[0] >= s_rows and S.shape[1] >= s_cols, + "S is too small for the requested block grid") + + hbm_result = nl.ndarray((M, N), dtype=X.dtype, buffer=nl.shared_hbm) + + s_stride = S.shape[1] + + for sr in range(s_rows): + band_start = sr * TILE_SIZE + band_size = min(TILE_SIZE, M - band_start) + + # Replicate scale row ``sr`` to every partition: partition stride 0. + s_row_bcast = nl.ndarray((PMAX, s_cols), dtype=S.dtype, buffer=nl.sbuf) + nisa.dma_copy( + dst=s_row_bcast, + src=S.ap(pattern=[[0, PMAX], [1, s_cols]], offset=sr * s_stride), + ) + + s_row_f32 = nl.ndarray((PMAX, s_cols), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=s_row_f32, src=s_row_bcast) + + for rt in range(div_ceil(band_size, PMAX)): + row_start = band_start + rt * PMAX + row_size = min(PMAX, M - row_start) + + for sc in range(s_cols): + col_start = sc * TILE_SIZE + col_size = min(TILE_SIZE, N - col_start) + + x_tile = nl.ndarray((row_size, col_size), dtype=X.dtype, + buffer=nl.sbuf) + nisa.dma_copy( + dst=x_tile, + src=X[row_start:row_start + row_size, + col_start:col_start + col_size], + ) + + out_tile = nl.ndarray((row_size, col_size), dtype=X.dtype, + buffer=nl.sbuf) + nisa.tensor_scalar(dst=out_tile, data=x_tile, op0=nl.multiply, + operand0=s_row_f32[0:row_size, sc:sc + 1]) + + nisa.dma_copy( + dst=hbm_result[row_start:row_start + row_size, + col_start:col_start + col_size], + src=out_tile, + ) + + return hbm_result + + +def run(X: torch.Tensor, S: torch.Tensor, M: int, N: int, TILE_SIZE: int, + block_size: int = 1024, autotune: bool = False, **kwargs) -> torch.Tensor: + if TILE_SIZE % PMAX != 0: + raise NotImplementedError( + f"weight_dequant NKI: TILE_SIZE ({TILE_SIZE}) must be a multiple of {PMAX}" + ) + return weight_dequant_kernel(X, S, TILE_SIZE) + + +def get_last_config() -> dict | None: + return None diff --git a/benchmarks/operators/weight_dequant/impl_torch.py b/benchmarks/operators/weight_dequant/impl_torch.py index 97028010..ca6545e6 100644 --- a/benchmarks/operators/weight_dequant/impl_torch.py +++ b/benchmarks/operators/weight_dequant/impl_torch.py @@ -2,7 +2,10 @@ def run(X: torch.Tensor, S: torch.Tensor, M: int, N: int, TILE_SIZE: int, **kwargs): - row_idx = torch.arange(M, device=X.device) // TILE_SIZE - col_idx = torch.arange(N, device=X.device) // TILE_SIZE - scale = S[row_idx[:, None], col_idx[None, :]] + # repeat_interleave (broadcast+reshape) instead of S[row_idx[:, None], col_idx[None, :]]: + # that advanced-indexing gather compiles fine under torch_xla but faults at + # runtime on Neuron with "scatter/gather (indirect memory copy via vector DGE) + # out-of-bound access" -- see weight_dequant torch-baseline investigation. + scale = S.repeat_interleave(TILE_SIZE, dim=0).repeat_interleave(TILE_SIZE, dim=1) + scale = scale[:M, :N] return X * scale From 47dd19dc778ff30ecd2e08f0b0db28a2b33d89f9 Mon Sep 17 00:00:00 2001 From: bowencui123 Date: Sat, 29 Aug 2026 08:38:12 +0000 Subject: [PATCH 2/8] nki(weight_dequant): NkiAutotuner wiring (`block_size`) Tunables mirror the Triton search space (`BLOCK_SIZE`/`TILE_SIZE`); defaults are the previous constants, so autotune=False is unchanged. `TILE_SIZE` is the quantization block (semantic, not tuned) Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_012Q38kGmXvyoeM1qtCbheSL --- .../operators/weight_dequant/impl_nki.py | 32 +++++++++++++++---- 1 file changed, 26 insertions(+), 6 deletions(-) diff --git a/benchmarks/operators/weight_dequant/impl_nki.py b/benchmarks/operators/weight_dequant/impl_nki.py index e52abc17..e501b6ed 100644 --- a/benchmarks/operators/weight_dequant/impl_nki.py +++ b/benchmarks/operators/weight_dequant/impl_nki.py @@ -1,5 +1,9 @@ +from types import SimpleNamespace + import torch +from core.nki_autotune import NkiAutotuner + try: import nki import nki.language as nl @@ -21,7 +25,7 @@ def div_ceil(n: int, d: int) -> int: if nki is not None: @nki.jit - def weight_dequant_kernel(X, S, TILE_SIZE): + def weight_dequant_kernel(X, S, TILE_SIZE, block_size): """Block-wise weight dequantization: ``out[m, n] = X[m, n] * S[m // T, n // T]``. ``S`` holds one scale per ``TILE_SIZE x TILE_SIZE`` block of ``X``. @@ -77,9 +81,9 @@ def weight_dequant_kernel(X, S, TILE_SIZE): s_row_f32 = nl.ndarray((PMAX, s_cols), dtype=nl.float32, buffer=nl.sbuf) nisa.tensor_copy(dst=s_row_f32, src=s_row_bcast) - for rt in range(div_ceil(band_size, PMAX)): - row_start = band_start + rt * PMAX - row_size = min(PMAX, M - row_start) + for rt in range(div_ceil(band_size, block_size)): + row_start = band_start + rt * block_size + row_size = min(block_size, M - row_start) for sc in range(s_cols): col_start = sc * TILE_SIZE @@ -107,14 +111,30 @@ def weight_dequant_kernel(X, S, TILE_SIZE): return hbm_result +_DEFAULT_CONFIG = SimpleNamespace(block_size=128) +_SEARCH_SPACE = [SimpleNamespace(block_size=b) for b in (32, 64, 128)] +_tuner = NkiAutotuner(weight_dequant_kernel) if nki is not None else None +_last_autotune_config: dict = {} + + def run(X: torch.Tensor, S: torch.Tensor, M: int, N: int, TILE_SIZE: int, block_size: int = 1024, autotune: bool = False, **kwargs) -> torch.Tensor: if TILE_SIZE % PMAX != 0: raise NotImplementedError( f"weight_dequant NKI: TILE_SIZE ({TILE_SIZE}) must be a multiple of {PMAX}" ) - return weight_dequant_kernel(X, S, TILE_SIZE) + if autotune: + cfg = _tuner.tune_or_cached( + shape_key=(tuple(X.shape), str(X.dtype)), + search_space=_SEARCH_SPACE, + args_fn=lambda cfg: (X, S, TILE_SIZE, cfg.block_size), + ) + _last_autotune_config.clear() + _last_autotune_config.update(vars(cfg)) + else: + cfg = _DEFAULT_CONFIG + return weight_dequant_kernel(X, S, TILE_SIZE, cfg.block_size) def get_last_config() -> dict | None: - return None + return dict(_last_autotune_config) or None From cf9f7310b242b7903ee61d55b250ad3aeac795f3 Mon Sep 17 00:00:00 2001 From: bowencui123 Date: Mon, 31 Aug 2026 21:44:39 +0000 Subject: [PATCH 3/8] review(nki): device-dispatch the torch baseline (non-XLA path == main) The non-XLA path is byte-identical to main; only the XLA device takes the XLA-compatible variant. CPU equivalence of the two paths verified (including non-divisible shapes); on trn2 case 0 still verifies and times (torch 0.0232 ms, NKI 0.0349 ms). Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_011Wvs1tztZTGQFD78YdZaha --- benchmarks/operators/weight_dequant/impl_torch.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/benchmarks/operators/weight_dequant/impl_torch.py b/benchmarks/operators/weight_dequant/impl_torch.py index ca6545e6..a8c9652a 100644 --- a/benchmarks/operators/weight_dequant/impl_torch.py +++ b/benchmarks/operators/weight_dequant/impl_torch.py @@ -2,10 +2,10 @@ def run(X: torch.Tensor, S: torch.Tensor, M: int, N: int, TILE_SIZE: int, **kwargs): - # repeat_interleave (broadcast+reshape) instead of S[row_idx[:, None], col_idx[None, :]]: - # that advanced-indexing gather compiles fine under torch_xla but faults at - # runtime on Neuron with "scatter/gather (indirect memory copy via vector DGE) - # out-of-bound access" -- see weight_dequant torch-baseline investigation. - scale = S.repeat_interleave(TILE_SIZE, dim=0).repeat_interleave(TILE_SIZE, dim=1) - scale = scale[:M, :N] + if X.device.type == "xla": + scale = S.repeat_interleave(TILE_SIZE, dim=0).repeat_interleave(TILE_SIZE, dim=1) + return X * scale[:M, :N] + row_idx = torch.arange(M, device=X.device) // TILE_SIZE + col_idx = torch.arange(N, device=X.device) // TILE_SIZE + scale = S[row_idx[:, None], col_idx[None, :]] return X * scale From 798b77bdd36c830d8fee0ac936428a4dd3eeff34 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Tue, 8 Sep 2026 01:14:59 +0000 Subject: [PATCH 4/8] tilelang(weight_dequant): add NKI benchmark results (trn2.3xlarge, LNC2) Merges NKI backend timing into results/csv/weight_dequant_default.csv, run against this branch's impl_nki.py on trn2.3xlarge with the LNC2 execution contract (NEURON_LOGICAL_NC_CONFIG=2, NEURON_RT_NUM_CORES=1, NEURON_CC_FLAGS="--target trn2 --lnc 2"). Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01AQseF7nyesBh8KZAp8g7Cm --- results/csv/weight_dequant_default.csv | 122 ++++++++++++------------- 1 file changed, 61 insertions(+), 61 deletions(-) diff --git a/results/csv/weight_dequant_default.csv b/results/csv/weight_dequant_default.csv index 0c19f7ec..69ff522d 100644 --- a/results/csv/weight_dequant_default.csv +++ b/results/csv/weight_dequant_default.csv @@ -1,61 +1,61 @@ -params,dtype,torch_ms,triton_ms,cutile_ms,speedup_triton,speedup_cutile,triton_vs_cutile,tilelang_ms,speedup_tilelang -M=512,fp16,0.0185,0.0021,0.0022,8.78,8.22,1.0684,0.0019,9.84 -M=512,bf16,0.0173,0.0019,0.0021,8.97,8.27,1.0848,0.0017,10.43 -M=512,fp32,0.0185,0.0021,0.0025,8.66,7.53,1.1504,0.0017,10.64 -M=1024,fp16,0.0276,0.0025,0.0031,10.84,8.95,1.2111,0.0024,11.32 -M=1024,bf16,0.0266,0.0025,0.0031,10.80,8.53,1.2652,0.0021,12.75 -M=1024,fp32,0.0265,0.0028,0.0037,9.56,7.15,1.3378,0.0027,9.70 -M=1536,fp16,0.0390,0.0043,0.0064,9.16,6.04,1.5158,0.0036,10.75 -M=1536,bf16,0.0381,0.0043,0.0064,8.94,5.94,1.5039,0.0038,10.09 -M=1536,fp32,0.0388,0.0044,0.0068,8.73,5.69,1.5338,0.0048,8.02 -M=2048,fp16,0.0583,0.0046,0.0077,12.60,7.61,1.6569,0.0041,14.08 -M=2048,bf16,0.0570,0.0044,0.0077,13.09,7.38,1.7740,0.0042,13.50 -M=2048,fp32,0.0598,0.0072,0.0095,8.36,6.28,1.3315,0.0066,9.02 -M=2560,fp16,0.0920,0.0083,0.0140,11.04,6.57,1.6801,0.0078,11.80 -M=2560,bf16,0.0913,0.0085,0.0141,10.81,6.48,1.6664,0.0056,16.30 -M=2560,fp32,0.0969,0.0110,0.0161,8.81,6.02,1.4627,0.0114,8.51 -M=3072,fp16,0.1331,0.0089,0.0197,14.89,6.74,2.2069,0.0084,15.87 -M=3072,bf16,0.1322,0.0088,0.0197,15.01,6.73,2.2321,0.0084,15.66 -M=3072,fp32,0.1405,0.0145,0.0223,9.69,6.31,1.5347,0.0140,10.01 -M=3584,fp16,0.1792,0.0146,0.0263,12.26,6.82,1.7980,0.0150,11.92 -M=3584,bf16,0.1788,0.0146,0.0263,12.28,6.80,1.8042,0.0110,16.29 -M=3584,fp32,0.1886,0.0181,0.0289,10.41,6.52,1.5968,0.0199,9.48 -M=4096,fp16,0.2264,0.0146,0.0251,15.51,9.02,1.7200,0.0134,16.89 -M=4096,bf16,0.2264,0.0147,0.0251,15.44,9.03,1.7093,0.0137,16.57 -M=4096,fp32,0.2420,0.0229,0.0307,10.58,7.88,1.3436,0.0226,10.70 -M=4608,fp16,0.2830,0.0222,0.0405,12.76,6.99,1.8253,0.0228,12.42 -M=4608,bf16,0.2834,0.0222,0.0405,12.78,7.00,1.8254,0.0227,12.47 -M=4608,fp32,0.3021,0.0280,0.0457,10.78,6.60,1.6319,0.0306,9.89 -M=5120,fp16,0.3459,0.0209,0.0498,16.59,6.95,2.3879,0.0194,17.79 -M=5120,bf16,0.3458,0.0208,0.0498,16.59,6.94,2.3895,0.0195,17.70 -M=5120,fp32,0.3668,0.0335,0.0554,10.96,6.63,1.6542,0.0333,11.02 -M=5632,fp16,0.4140,0.0316,0.0592,13.08,7.00,1.8697,0.0328,12.61 -M=5632,bf16,0.4140,0.0316,0.0591,13.09,7.00,1.8702,0.0324,12.77 -M=5632,fp32,0.4416,0.0400,0.0669,11.05,6.60,1.6742,0.0436,10.12 -M=6144,fp16,0.4886,0.0287,0.0710,17.04,6.88,2.4769,0.0268,18.21 -M=6144,bf16,0.4884,0.0287,0.0710,17.02,6.88,2.4752,0.0269,18.18 -M=6144,fp32,0.5215,0.0465,0.0802,11.22,6.50,1.7256,0.0466,11.18 -M=6656,fp16,0.5704,0.0431,0.0815,13.25,7.00,1.8935,0.0446,12.78 -M=6656,bf16,0.5700,0.0430,0.0815,13.25,6.99,1.8943,0.0310,18.39 -M=6656,fp32,0.6108,0.0546,0.0922,11.20,6.62,1.6903,0.0594,10.28 -M=7168,fp16,0.6585,0.0375,0.0954,17.55,6.90,2.5431,0.0354,18.61 -M=7168,bf16,0.6582,0.0375,0.0954,17.54,6.90,2.5429,0.0353,18.64 -M=7168,fp32,0.7053,0.0621,0.1076,11.35,6.55,1.7325,0.0626,11.27 -M=7680,fp16,0.7533,0.0571,0.1112,13.20,6.77,1.9479,0.0401,18.80 -M=7680,bf16,0.7545,0.0570,0.1112,13.23,6.78,1.9494,0.0572,13.20 -M=7680,fp32,0.8091,0.0718,0.1220,11.27,6.63,1.6995,0.0772,10.47 -M=8192,fp16,0.8525,0.0490,0.0918,17.40,9.28,1.8742,0.0450,18.93 -M=8192,bf16,0.8515,0.0490,0.0918,17.38,9.27,1.8743,0.0456,18.65 -M=8192,fp32,0.9158,0.0799,0.1110,11.46,8.25,1.3891,0.0804,11.39 -M=8704,fp16,0.9572,0.0715,0.1374,13.39,6.97,1.9219,0.0730,13.12 -M=8704,bf16,0.9589,0.0716,0.1374,13.40,6.98,1.9205,0.0502,19.12 -M=8704,fp32,1.0309,0.0905,0.1555,11.39,6.63,1.7183,0.0978,10.54 -M=9216,fp16,1.0719,0.0593,0.1553,18.08,6.90,2.6192,0.0558,19.22 -M=9216,bf16,1.0724,0.0593,0.1553,18.09,6.91,2.6197,0.0557,19.27 -M=9216,fp32,1.1555,0.1002,0.1725,11.54,6.70,1.7223,0.1011,11.43 -M=9728,fp16,1.1920,0.0886,0.1710,13.45,6.97,1.9293,0.0901,13.23 -M=9728,bf16,1.1927,0.0885,0.1709,13.47,6.98,1.9306,0.0890,13.41 -M=9728,fp32,1.2885,0.1115,0.1934,11.56,6.66,1.7353,0.1122,11.49 -M=10240,fp16,1.3192,0.0721,0.1910,18.30,6.91,2.6506,0.0681,19.37 -M=10240,bf16,1.3175,0.0720,0.1910,18.31,6.90,2.6544,0.0679,19.39 -M=10240,fp32,1.4237,0.1229,0.2123,11.58,6.71,1.7272,0.1243,11.45 +params,dtype,torch_ms,triton_ms,cutile_ms,speedup_triton,speedup_cutile,triton_vs_cutile,tilelang_ms,speedup_tilelang,torch_nki_ms,nki_ms,speedup_nki +M=512,fp16,0.0185,0.0021,0.0022,8.78,8.22,1.0684,0.0019,9.84,2.5048,0.0351,71.41 +M=512,bf16,0.0173,0.0019,0.0021,8.97,8.27,1.0848,0.0017,10.43,2.5047,0.0350,71.65 +M=512,fp32,0.0185,0.0021,0.0025,8.66,7.53,1.1504,0.0017,10.64,2.5056,0.0356,70.34 +M=1024,fp16,0.0276,0.0025,0.0031,10.84,8.95,1.2111,0.0024,11.32,9.9260,0.1009,98.34 +M=1024,bf16,0.0266,0.0025,0.0031,10.80,8.53,1.2652,0.0021,12.75,9.9259,0.1009,98.33 +M=1024,fp32,0.0265,0.0028,0.0037,9.56,7.15,1.3378,0.0027,9.70,9.9266,0.1003,98.98 +M=1536,fp16,0.0390,0.0043,0.0064,9.16,6.04,1.5158,0.0036,10.75,20.2811,0.2074,97.77 +M=1536,bf16,0.0381,0.0043,0.0064,8.94,5.94,1.5039,0.0038,10.09,20.2811,0.2074,97.77 +M=1536,fp32,0.0388,0.0044,0.0068,8.73,5.69,1.5338,0.0048,8.02,13.9147,0.2079,66.92 +M=2048,fp16,0.0583,0.0046,0.0077,12.60,7.61,1.6569,0.0041,14.08,25.6457,0.3570,71.84 +M=2048,bf16,0.0570,0.0044,0.0077,13.09,7.38,1.7740,0.0042,13.50,25.6618,0.3570,71.88 +M=2048,fp32,0.0598,0.0072,0.0095,8.36,6.28,1.3315,0.0066,9.02,24.7112,0.3589,68.84 +M=2560,fp16,0.0920,0.0083,0.0140,11.04,6.57,1.6801,0.0078,11.80,38.5823,0.5483,70.37 +M=2560,bf16,0.0913,0.0085,0.0141,10.81,6.48,1.6664,0.0056,16.30,38.5821,0.5483,70.37 +M=2560,fp32,0.0969,0.0110,0.0161,8.81,6.02,1.4627,0.0114,8.51,38.6094,0.5491,70.31 +M=3072,fp16,0.1331,0.0089,0.0197,14.89,6.74,2.2069,0.0084,15.87,55.5372,0.7817,71.04 +M=3072,bf16,0.1322,0.0088,0.0197,15.01,6.73,2.2321,0.0084,15.66,55.5370,0.7818,71.04 +M=3072,fp32,0.1405,0.0145,0.0223,9.69,6.31,1.5347,0.0140,10.01,55.5738,0.7818,71.09 +M=3584,fp16,0.1792,0.0146,0.0263,12.26,6.82,1.7980,0.0150,11.92,75.5826,1.0570,71.51 +M=3584,bf16,0.1788,0.0146,0.0263,12.28,6.80,1.8042,0.0110,16.29,75.5828,1.0569,71.51 +M=3584,fp32,0.1886,0.0181,0.0289,10.41,6.52,1.5968,0.0199,9.48,75.6304,1.0578,71.49 +M=4096,fp16,0.2264,0.0146,0.0251,15.51,9.02,1.7200,0.0134,16.89,99.5658,1.3753,72.39 +M=4096,bf16,0.2264,0.0147,0.0251,15.44,9.03,1.7093,0.0137,16.57,99.5027,1.3755,72.34 +M=4096,fp32,0.2420,0.0229,0.0307,10.58,7.88,1.3436,0.0226,10.70,98.8365,1.3760,71.83 +M=4608,fp16,0.2830,0.0222,0.0405,12.76,6.99,1.8253,0.0228,12.42,124.9118,1.7339,72.04 +M=4608,bf16,0.2834,0.0222,0.0405,12.78,7.00,1.8254,0.0227,12.47,124.9119,1.7339,72.04 +M=4608,fp32,0.3021,0.0280,0.0457,10.78,6.60,1.6319,0.0306,9.89,125.0999,1.7367,72.03 +M=5120,fp16,0.3459,0.0209,0.0498,16.59,6.95,2.3879,0.0194,17.79,154.1932,2.1360,72.19 +M=5120,bf16,0.3458,0.0208,0.0498,16.59,6.94,2.3895,0.0195,17.70,154.1934,2.1361,72.19 +M=5120,fp32,0.3668,0.0335,0.0554,10.96,6.63,1.6542,0.0333,11.02,154.3176,2.1364,72.23 +M=5632,fp16,0.4140,0.0316,0.0592,13.08,7.00,1.8697,0.0328,12.61,186.6698,2.5780,72.41 +M=5632,bf16,0.4140,0.0316,0.0591,13.09,7.00,1.8702,0.0324,12.77,186.6701,2.5780,72.41 +M=5632,fp32,0.4416,0.0400,0.0669,11.05,6.60,1.6742,0.0436,10.12,186.8025,2.5789,72.44 +M=6144,fp16,0.4886,0.0287,0.0710,17.04,6.88,2.4769,0.0268,18.21,222.0560,3.0642,72.47 +M=6144,bf16,0.4884,0.0287,0.0710,17.02,6.88,2.4752,0.0269,18.18,222.0558,3.0642,72.47 +M=6144,fp32,0.5215,0.0465,0.0802,11.22,6.50,1.7256,0.0466,11.18,222.2288,3.0641,72.53 +M=6656,fp16,0.5704,0.0431,0.0815,13.25,7.00,1.8935,0.0446,12.78,260.5408,3.5904,72.57 +M=6656,bf16,0.5700,0.0430,0.0815,13.25,6.99,1.8943,0.0310,18.39,260.5411,3.5902,72.57 +M=6656,fp32,0.6108,0.0546,0.0922,11.20,6.62,1.6903,0.0594,10.28,260.8216,3.5915,72.62 +M=7168,fp16,0.6585,0.0375,0.0954,17.55,6.90,2.5431,0.0354,18.61,302.2148,4.1605,72.64 +M=7168,bf16,0.6582,0.0375,0.0954,17.54,6.90,2.5429,0.0353,18.64,302.2146,4.1606,72.64 +M=7168,fp32,0.7053,0.0621,0.1076,11.35,6.55,1.7325,0.0626,11.27,302.4038,4.1733,72.46 +M=7680,fp16,0.7533,0.0571,0.1112,13.20,6.77,1.9479,0.0401,18.80,346.9441,4.7720,72.70 +M=7680,bf16,0.7545,0.0570,0.1112,13.23,6.78,1.9494,0.0572,13.20,346.9451,4.7720,72.70 +M=7680,fp32,0.8091,0.0718,0.1220,11.27,6.63,1.6995,0.0772,10.47,347.1352,4.7817,72.60 +M=8192,fp16,0.8525,0.0490,0.0918,17.40,9.28,1.8742,0.0450,18.93,394.6603,5.4231,72.77 +M=8192,bf16,0.8515,0.0490,0.0918,17.38,9.27,1.8743,0.0456,18.65,394.6604,5.4231,72.77 +M=8192,fp32,0.9158,0.0799,0.1110,11.46,8.25,1.3891,0.0804,11.39,631.9603,5.4367,116.24 +M=8704,fp16,0.9572,0.0715,0.1374,13.39,6.97,1.9219,0.0730,13.12,445.8699,6.1203,72.85 +M=8704,bf16,0.9589,0.0716,0.1374,13.40,6.98,1.9205,0.0502,19.12,445.8691,6.1206,72.85 +M=8704,fp32,1.0309,0.0905,0.1555,11.39,6.63,1.7183,0.0978,10.54,446.3709,6.1300,72.82 +M=9216,fp16,1.0719,0.0593,0.1553,18.08,6.90,2.6192,0.0558,19.22,501.9925,6.8546,73.23 +M=9216,bf16,1.0724,0.0593,0.1553,18.09,6.91,2.6197,0.0557,19.27,501.9926,6.8546,73.23 +M=9216,fp32,1.1555,0.1002,0.1725,11.54,6.70,1.7223,0.1011,11.43,499.9652,6.8636,72.84 +M=9728,fp16,1.1920,0.0886,0.1710,13.45,6.97,1.9293,0.0901,13.23,556.8344,7.6341,72.94 +M=9728,bf16,1.1927,0.0885,0.1709,13.47,6.98,1.9306,0.0890,13.41,556.8358,7.6341,72.94 +M=9728,fp32,1.2885,0.1115,0.1934,11.56,6.66,1.7353,0.1122,11.49,557.5616,7.6355,73.02 +M=10240,fp16,1.3192,0.0721,0.1910,18.30,6.91,2.6506,0.0681,19.37,616.5930,8.4559,72.92 +M=10240,bf16,1.3175,0.0720,0.1910,18.31,6.90,2.6544,0.0679,19.39,616.5933,8.4559,72.92 +M=10240,fp32,1.4237,0.1229,0.2123,11.58,6.71,1.7272,0.1243,11.45,617.2312,8.4566,72.99 From e05630b0ec83034ecf23a5c1f80a4c3cabba9b7c Mon Sep 17 00:00:00 2001 From: Bowen Cui Date: Wed, 16 Sep 2026 22:33:54 +0000 Subject: [PATCH 5/8] nki(weight_dequant): LNC2 kernel, one Vector-engine instruction per block (broadcast per-tile scales), partition-stride-0 scale rows, no host pad; rerun default+autotune benchmarks Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01ScXYNjrrKGgDUVNHxv7HJt --- .../operators/weight_dequant/impl_nki.py | 201 +++++++++--------- results/csv/weight_dequant_autotune.csv | 122 +++++------ results/csv/weight_dequant_default.csv | 120 +++++------ .../weight_dequant_autotune.json | 182 +++++++++++++--- 4 files changed, 383 insertions(+), 242 deletions(-) diff --git a/benchmarks/operators/weight_dequant/impl_nki.py b/benchmarks/operators/weight_dequant/impl_nki.py index e501b6ed..bb048c69 100644 --- a/benchmarks/operators/weight_dequant/impl_nki.py +++ b/benchmarks/operators/weight_dequant/impl_nki.py @@ -1,3 +1,26 @@ +"""NKI (AWS Trainium) implementation of weight_dequant: ``out[m, n] = X[m, n] * S[m // T, n // T]``. + +Design (mirrors the Triton / cuTile kernels): + +* Triton walks the flat ``M * N`` elements in ``BLOCK_SIZE`` blocks, looks up the + block scale of every element and multiplies. NKI puts 128 rows on the 128 SBUF + partitions (``TILE_SIZE`` is a multiple of 128, so a row tile never straddles + two scale rows) and walks each row tile in ``BLOCK_SIZE``-column blocks: + - the scale row of the band is replicated to all partitions by a single + partition-stride-0 DMA and converted to fp32 once per row tile; + - per column block: one ``[128, BLOCK]`` DMA, then one Vector-engine + ``tensor_scalar`` per ``TILE_SIZE``-wide sub-block with the matching + ``[128, 1]`` scale column (broadcast along the free axis in hardware), and + one DMA store. + ``BLOCK_SIZE`` defaults / search space are those of ``impl_triton.py`` / + ``impl_cutile.py`` (default 1024; search 256/512/1024/2048/4096). +* Row tiles are split across the NeuronCores of the logical core (LNC2 on trn2) + by ``nl.program_id(0)``. Partial tiles / blocks are clamped with ``min()``. +""" +import functools +import os +import re +import subprocess from types import SimpleNamespace import torch @@ -6,114 +29,100 @@ try: import nki - import nki.language as nl import nki.isa as nisa + import nki.language as nl PMAX = nl.tile_size.pmax except ImportError: nki = None - - -def kernel_assert(condition: bool, error_text: str): - """Assert with NKI-formatted error message.""" - assert condition, f"[INTERNAL_ERROR] [NCC_INKI016] Kernel validation exception: {error_text}" - - -def div_ceil(n: int, d: int) -> int: - """Ceiling division: smallest integer >= n/d.""" - return (n + d - 1) // d + PMAX = 128 + + +@functools.lru_cache(maxsize=1) +def _lnc_degree() -> int: + """Logical-NeuronCore degree the kernel is launched with (``kernel[lnc]``). + + Must match the LNC the XLA module is compiled for: launching a ``kernel[2]`` + into an ``--lnc 1`` module silently computes only core 0's half. + """ + explicit = os.environ.get("NEURON_LOGICAL_NC_CONFIG", "") + if explicit.strip().isdigit(): + return int(explicit.strip()) + match = re.search(r"--lnc[=\s]+(\d+)", os.environ.get("NEURON_CC_FLAGS", "")) + if match: + return int(match.group(1)) + try: + out = subprocess.run(["neuron-ls"], capture_output=True, text=True, timeout=10).stdout + lnc = re.search(r"logical-neuroncore-config:\s*(\d+)", out) + if lnc: + return int(lnc.group(1)) + except (OSError, subprocess.SubprocessError): + pass + return 1 if nki is not None: @nki.jit def weight_dequant_kernel(X, S, TILE_SIZE, block_size): - """Block-wise weight dequantization: ``out[m, n] = X[m, n] * S[m // T, n // T]``. - - ``S`` holds one scale per ``TILE_SIZE x TILE_SIZE`` block of ``X``. + """Block-wise dequantization of ``[M, N]`` ``X`` with ``[ceil(M/T), ceil(N/T)]`` scales ``S``. Args: - X: [M, N] quantized weights in HBM. - S: [ceil(M/T), ceil(N/T)] per-block scales in HBM (same dtype as X). - TILE_SIZE: block edge length T, a multiple of PMAX (checked in ``run``). - - Returns: - [M, N] tensor in HBM with the same dtype as ``X``. - - Notes: - * Because ``TILE_SIZE`` is a multiple of ``PMAX``, a ``PMAX``-row tile never - straddles two scale rows, so one scale row of ``S`` serves every row tile - inside a ``TILE_SIZE``-row band. - * The scalar broadcast is done in two cheap steps rather than by - materializing a full [PMAX, TILE_SIZE] scale tile: - 1. one DMA per scale row with a partition stride of 0 - (``S.ap(pattern=[[0, PMAX], [1, s_cols]])``) replicates the whole - scale row to all 128 partitions in a single instruction; - 2. ``nisa.tensor_scalar`` with ``operand0`` = the [PMAX, 1] column for - that block broadcasts along the free axis in hardware. - ``operand0`` must be float32 (the MLIR verifier rejects a half-precision - ``operand0``), hence the fp32 copy of the broadcast row. - * Boundary tiles are clamped (tile sized to the surviving extent) instead of - masked, so no out-of-range element is ever loaded, multiplied or stored. + X: ``[M, N]`` quantized weights in HBM (fp16 / bf16 / fp32). + S: ``[ceil(M/T), ceil(N/T)]`` per-block scales in HBM (same dtype as X). + TILE_SIZE: block edge length T, a multiple of 128 (checked in ``run``). + block_size: columns per DMA block (compile-time constant, multiple of T or >= N). """ - kernel_assert(len(X.shape) == 2, "X must be 2D [M, N]") - kernel_assert(len(S.shape) == 2, "S must be 2D [M/T, N/T]") - M, N = X.shape - s_rows = div_ceil(M, TILE_SIZE) - s_cols = div_ceil(N, TILE_SIZE) - kernel_assert(S.shape[0] >= s_rows and S.shape[1] >= s_cols, - "S is too small for the requested block grid") - - hbm_result = nl.ndarray((M, N), dtype=X.dtype, buffer=nl.shared_hbm) - + out = nl.ndarray((M, N), dtype=X.dtype, buffer=nl.shared_hbm) + s_cols = (N + TILE_SIZE - 1) // TILE_SIZE s_stride = S.shape[1] - - for sr in range(s_rows): - band_start = sr * TILE_SIZE - band_size = min(TILE_SIZE, M - band_start) - - # Replicate scale row ``sr`` to every partition: partition stride 0. - s_row_bcast = nl.ndarray((PMAX, s_cols), dtype=S.dtype, buffer=nl.sbuf) - nisa.dma_copy( - dst=s_row_bcast, - src=S.ap(pattern=[[0, PMAX], [1, s_cols]], offset=sr * s_stride), - ) - - s_row_f32 = nl.ndarray((PMAX, s_cols), dtype=nl.float32, buffer=nl.sbuf) - nisa.tensor_copy(dst=s_row_f32, src=s_row_bcast) - - for rt in range(div_ceil(band_size, block_size)): - row_start = band_start + rt * block_size - row_size = min(block_size, M - row_start) - - for sc in range(s_cols): - col_start = sc * TILE_SIZE - col_size = min(TILE_SIZE, N - col_start) - - x_tile = nl.ndarray((row_size, col_size), dtype=X.dtype, - buffer=nl.sbuf) - nisa.dma_copy( - dst=x_tile, - src=X[row_start:row_start + row_size, - col_start:col_start + col_size], - ) - - out_tile = nl.ndarray((row_size, col_size), dtype=X.dtype, - buffer=nl.sbuf) - nisa.tensor_scalar(dst=out_tile, data=x_tile, op0=nl.multiply, - operand0=s_row_f32[0:row_size, sc:sc + 1]) - - nisa.dma_copy( - dst=hbm_result[row_start:row_start + row_size, - col_start:col_start + col_size], - src=out_tile, - ) - - return hbm_result - - -_DEFAULT_CONFIG = SimpleNamespace(block_size=128) -_SEARCH_SPACE = [SimpleNamespace(block_size=b) for b in (32, 64, 128)] -_tuner = NkiAutotuner(weight_dequant_kernel) if nki is not None else None + n_blocks = (N + block_size - 1) // block_size + + n_tiles = (M + PMAX - 1) // PMAX + num_programs = nl.num_programs() + per_core = (n_tiles + num_programs - 1) // num_programs + pid = nl.program_id(0) + for ti in range(pid * per_core, min(n_tiles, (pid + 1) * per_core)): + r0 = ti * PMAX + rs = min(PMAX, M - r0) + sr = r0 // TILE_SIZE # scale row of this band + # Replicate scale row ``sr`` to every partition (partition stride 0), then fp32 + # (tensor_scalar's [P, 1] operand must be fp32). + s_raw = nl.ndarray((rs, s_cols), dtype=S.dtype, buffer=nl.sbuf) + nisa.dma_copy(dst=s_raw, src=S.ap(pattern=[[0, rs], [1, s_cols]], offset=sr * s_stride)) + s_f32 = nl.ndarray((rs, s_cols), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=s_f32, src=s_raw) + for cb in range(n_blocks): + c0 = cb * block_size + cs = min(block_size, N - c0) + x_tile = nl.ndarray((rs, cs), dtype=X.dtype, buffer=nl.sbuf) + nisa.dma_copy(dst=x_tile, src=X[r0:r0 + rs, c0:c0 + cs]) + y_tile = nl.ndarray((rs, cs), dtype=X.dtype, buffer=nl.sbuf) + if c0 % TILE_SIZE == 0 and cs % TILE_SIZE == 0: + # Whole block in one Vector-engine instruction: view the block as + # [rs, n_sub, TILE_SIZE] and broadcast the [rs, n_sub] scales along the + # last axis (stride 0), instead of one instruction per 128-column sub-block. + n_sub = cs // TILE_SIZE + sc0 = c0 // TILE_SIZE + nisa.tensor_tensor(dst=y_tile.ap(pattern=[[cs, rs], [TILE_SIZE, n_sub], [1, TILE_SIZE]]), + data1=x_tile.ap(pattern=[[cs, rs], [TILE_SIZE, n_sub], [1, TILE_SIZE]]), + data2=s_f32.ap(pattern=[[s_cols, rs], [1, n_sub], [0, TILE_SIZE]], offset=sc0), + op=nl.multiply) + else: + # One scale per TILE_SIZE-wide column sub-block. + for sc in range(c0 // TILE_SIZE, (c0 + cs + TILE_SIZE - 1) // TILE_SIZE): + a = max(sc * TILE_SIZE, c0) - c0 + b = min((sc + 1) * TILE_SIZE, c0 + cs) - c0 + nisa.tensor_scalar(dst=y_tile[0:rs, a:b], data=x_tile[0:rs, a:b], + op0=nl.multiply, operand0=s_f32[0:rs, sc:sc + 1]) + nisa.dma_copy(dst=out[r0:r0 + rs, c0:c0 + cs], src=y_tile) + return out + + +# Same block sizes as impl_triton.py / impl_cutile.py (BLOCK_SIZE default 1024; search 256..4096). +_DEFAULT_CONFIG = SimpleNamespace(block_size=1024) +_SEARCH_SPACE = [SimpleNamespace(block_size=b) for b in (256, 512, 1024, 2048, 4096)] +_kernel = weight_dequant_kernel[_lnc_degree()] if nki is not None else None +_tuner = NkiAutotuner(_kernel) if nki is not None else None _last_autotune_config: dict = {} @@ -125,7 +134,7 @@ def run(X: torch.Tensor, S: torch.Tensor, M: int, N: int, TILE_SIZE: int, ) if autotune: cfg = _tuner.tune_or_cached( - shape_key=(tuple(X.shape), str(X.dtype)), + shape_key=(tuple(X.shape), str(X.dtype), TILE_SIZE), search_space=_SEARCH_SPACE, args_fn=lambda cfg: (X, S, TILE_SIZE, cfg.block_size), ) @@ -133,7 +142,7 @@ def run(X: torch.Tensor, S: torch.Tensor, M: int, N: int, TILE_SIZE: int, _last_autotune_config.update(vars(cfg)) else: cfg = _DEFAULT_CONFIG - return weight_dequant_kernel(X, S, TILE_SIZE, cfg.block_size) + return _kernel(X, S, TILE_SIZE, cfg.block_size) def get_last_config() -> dict | None: diff --git a/results/csv/weight_dequant_autotune.csv b/results/csv/weight_dequant_autotune.csv index a9b20d03..f5c8c9c8 100644 --- a/results/csv/weight_dequant_autotune.csv +++ b/results/csv/weight_dequant_autotune.csv @@ -1,61 +1,61 @@ -params,dtype,torch_ms,triton_ms,cutile_ms,speedup_triton,speedup_cutile,triton_vs_cutile,tilelang_ms,speedup_tilelang -M=512,fp16,0.0185,0.0021,0.0022,8.65,8.55,1.0118,0.0020,9.39 -M=512,bf16,0.0184,0.0019,0.0021,9.46,8.72,1.0849,0.0016,11.23 -M=512,fp32,0.0175,0.0021,0.0023,8.32,7.70,1.0808,0.0019,9.17 -M=1024,fp16,0.0265,0.0034,0.0029,7.73,9.15,0.8452,0.0023,11.59 -M=1024,bf16,0.0270,0.0022,0.0033,12.14,8.24,1.4740,0.0022,12.11 -M=1024,fp32,0.0280,0.0027,0.0035,10.48,8.09,1.2953,0.0028,10.10 -M=1536,fp16,0.0401,0.0039,0.0064,10.40,6.30,1.6492,0.0037,10.86 -M=1536,bf16,0.0401,0.0042,0.0064,9.62,6.23,1.5437,0.0038,10.63 -M=1536,fp32,0.0402,0.0055,0.0070,7.25,5.75,1.2606,0.0051,7.89 -M=2048,fp16,0.0570,0.0051,0.0078,11.23,7.29,1.5414,0.0037,15.21 -M=2048,bf16,0.0579,0.0058,0.0078,9.93,7.46,1.3312,0.0038,15.27 -M=2048,fp32,0.0599,0.0075,0.0085,7.99,7.07,1.1300,0.0072,8.36 -M=2560,fp16,0.0939,0.0083,0.0145,11.37,6.48,1.7527,0.0079,11.95 -M=2560,bf16,0.0927,0.0079,0.0145,11.68,6.38,1.8303,0.0062,15.00 -M=2560,fp32,0.0987,0.0112,0.0151,8.83,6.52,1.3547,0.0109,9.02 -M=3072,fp16,0.1339,0.0089,0.0198,14.98,6.77,2.2118,0.0078,17.24 -M=3072,bf16,0.1336,0.0090,0.0197,14.86,6.78,2.1925,0.0085,15.76 -M=3072,fp32,0.1409,0.0150,0.0209,9.42,6.75,1.3954,0.0145,9.70 -M=3584,fp16,0.1776,0.0146,0.0258,12.15,6.88,1.7656,0.0140,12.73 -M=3584,bf16,0.1785,0.0154,0.0258,11.56,6.92,1.6702,0.0109,16.35 -M=3584,fp32,0.1898,0.0193,0.0273,9.84,6.95,1.4142,0.0193,9.85 -M=4096,fp16,0.2270,0.0147,0.0249,15.44,9.11,1.6940,0.0125,18.10 -M=4096,bf16,0.2281,0.0148,0.0238,15.43,9.57,1.6125,0.0126,18.17 -M=4096,fp32,0.2428,0.0231,0.0274,10.52,8.85,1.1884,0.0233,10.43 -M=4608,fp16,0.2843,0.0207,0.0405,13.72,7.02,1.9538,0.0206,13.77 -M=4608,bf16,0.2838,0.0224,0.0405,12.70,7.01,1.8109,0.0216,13.15 -M=4608,fp32,0.3018,0.0278,0.0419,10.85,7.20,1.5081,0.0296,10.20 -M=5120,fp16,0.3446,0.0209,0.0487,16.52,7.07,2.3367,0.0193,17.88 -M=5120,bf16,0.3449,0.0209,0.0498,16.48,6.93,2.3793,0.0195,17.72 -M=5120,fp32,0.3677,0.0339,0.0510,10.86,7.21,1.5057,0.0334,11.01 -M=5632,fp16,0.4138,0.0294,0.0592,14.06,6.99,2.0096,0.0291,14.23 -M=5632,bf16,0.4145,0.0295,0.0591,14.07,7.01,2.0063,0.0314,13.20 -M=5632,fp32,0.4428,0.0402,0.0610,11.02,7.26,1.5179,0.0426,10.39 -M=6144,fp16,0.4882,0.0271,0.0710,18.03,6.87,2.6232,0.0251,19.44 -M=6144,bf16,0.4885,0.0271,0.0710,18.06,6.88,2.6238,0.0251,19.45 -M=6144,fp32,0.5225,0.0467,0.0737,11.19,7.09,1.5769,0.0475,11.00 -M=6656,fp16,0.5698,0.0399,0.0806,14.29,7.07,2.0206,0.0388,14.68 -M=6656,bf16,0.5696,0.0399,0.0819,14.26,6.95,2.0510,0.0311,18.34 -M=6656,fp32,0.6130,0.0540,0.0838,11.36,7.32,1.5531,0.0591,10.38 -M=7168,fp16,0.6586,0.0378,0.0935,17.40,7.05,2.4695,0.0363,18.13 -M=7168,bf16,0.6562,0.0375,0.0954,17.49,6.88,2.5431,0.0359,18.29 -M=7168,fp32,0.7073,0.0623,0.0978,11.35,7.23,1.5704,0.0625,11.31 -M=7680,fp16,0.7507,0.0537,0.1086,13.97,6.91,2.0213,0.0373,20.15 -M=7680,bf16,0.7519,0.0537,0.1086,14.01,6.92,2.0238,0.0375,20.03 -M=7680,fp32,0.8091,0.0713,0.1130,11.34,7.16,1.5847,0.0752,10.76 -M=8192,fp16,0.8524,0.0459,0.0919,18.58,9.28,2.0027,0.0419,20.33 -M=8192,bf16,0.8530,0.0458,0.0918,18.62,9.29,2.0052,0.0421,20.24 -M=8192,fp32,0.9187,0.0799,0.1006,11.50,9.14,1.2582,0.0807,11.39 -M=8704,fp16,0.9583,0.0659,0.1351,14.54,7.09,2.0496,0.0471,20.36 -M=8704,bf16,0.9572,0.0659,0.1375,14.52,6.96,2.0843,0.0501,19.12 -M=8704,fp32,1.0333,0.0902,0.1369,11.46,7.55,1.5179,0.0954,10.83 -M=9216,fp16,1.0729,0.0591,0.1517,18.16,7.07,2.5672,0.0558,19.24 -M=9216,bf16,1.0721,0.0593,0.1516,18.09,7.07,2.5583,0.0568,18.89 -M=9216,fp32,1.1569,0.0998,0.1510,11.59,7.66,1.5128,0.1013,11.42 -M=9728,fp16,1.1919,0.0816,0.1710,14.60,6.97,2.0951,0.0585,20.39 -M=9728,bf16,1.1929,0.0816,0.1682,14.62,7.09,2.0620,0.0583,20.46 -M=9728,fp32,1.2859,0.1111,0.1696,11.58,7.58,1.5268,0.1123,11.45 -M=10240,fp16,1.3179,0.0679,0.1864,19.42,7.07,2.7475,0.0636,20.73 -M=10240,bf16,1.3198,0.0677,0.1865,19.50,7.08,2.7561,0.0644,20.51 -M=10240,fp32,1.4256,0.1225,0.1909,11.64,7.47,1.5580,0.1241,11.49 +params,dtype,torch_ms,triton_ms,cutile_ms,speedup_triton,speedup_cutile,triton_vs_cutile,tilelang_ms,speedup_tilelang,torch_nki_ms,nki_ms,speedup_nki +M=512,fp16,0.0185,0.0021,0.0022,8.65,8.55,1.0118,0.0020,9.39,0.0231,0.0212,1.09 +M=512,bf16,0.0184,0.0019,0.0021,9.46,8.72,1.0849,0.0016,11.23,0.0231,0.0222,1.04 +M=512,fp32,0.0175,0.0021,0.0023,8.32,7.70,1.0808,0.0019,9.17,0.0260,0.0226,1.15 +M=1024,fp16,0.0265,0.0034,0.0029,7.73,9.15,0.8452,0.0023,11.59,0.0240,0.0306,0.78 +M=1024,bf16,0.0270,0.0022,0.0033,12.14,8.24,1.4740,0.0022,12.11,0.0240,0.0313,0.77 +M=1024,fp32,0.0280,0.0027,0.0035,10.48,8.09,1.2953,0.0028,10.10,0.0257,0.0334,0.77 +M=1536,fp16,0.0401,0.0039,0.0064,10.40,6.30,1.6492,0.0037,10.86,0.0292,0.0415,0.70 +M=1536,bf16,0.0401,0.0042,0.0064,9.62,6.23,1.5437,0.0038,10.63,0.0293,0.0417,0.70 +M=1536,fp32,0.0402,0.0055,0.0070,7.25,5.75,1.2606,0.0051,7.89,0.0438,0.0469,0.93 +M=2048,fp16,0.0570,0.0051,0.0078,11.23,7.29,1.5414,0.0037,15.21,0.0382,0.0461,0.83 +M=2048,bf16,0.0579,0.0058,0.0078,9.93,7.46,1.3312,0.0038,15.27,0.0382,0.0460,0.83 +M=2048,fp32,0.0599,0.0075,0.0085,7.99,7.07,1.1300,0.0072,8.36,0.0697,0.0682,1.02 +M=2560,fp16,0.0939,0.0083,0.0145,11.37,6.48,1.7527,0.0079,11.95,0.0535,0.0611,0.87 +M=2560,bf16,0.0927,0.0079,0.0145,11.68,6.38,1.8303,0.0062,15.00,0.0535,0.0613,0.87 +M=2560,fp32,0.0987,0.0112,0.0151,8.83,6.52,1.3547,0.0109,9.02,0.0926,0.0977,0.95 +M=3072,fp16,0.1339,0.0089,0.0198,14.98,6.77,2.2118,0.0078,17.24,0.0700,0.0786,0.89 +M=3072,bf16,0.1336,0.0090,0.0197,14.86,6.78,2.1925,0.0085,15.76,0.0699,0.1158,0.60 +M=3072,fp32,0.1409,0.0150,0.0209,9.42,6.75,1.3954,0.0145,9.70,0.1279,0.1365,0.94 +M=3584,fp16,0.1776,0.0146,0.0258,12.15,6.88,1.7656,0.0140,12.73,0.0911,0.1008,0.90 +M=3584,bf16,0.1785,0.0154,0.0258,11.56,6.92,1.6702,0.0109,16.35,0.0911,0.1045,0.87 +M=3584,fp32,0.1898,0.0193,0.0273,9.84,6.95,1.4142,0.0193,9.85,0.1729,0.1817,0.95 +M=4096,fp16,0.2270,0.0147,0.0249,15.44,9.11,1.6940,0.0125,18.10,0.1316,0.1292,1.02 +M=4096,bf16,0.2281,0.0148,0.0238,15.43,9.57,1.6125,0.0126,18.17,0.1316,0.1349,0.98 +M=4096,fp32,0.2428,0.0231,0.0274,10.52,8.85,1.1884,0.0233,10.43,0.2404,0.2307,1.04 +M=4608,fp16,0.2843,0.0207,0.0405,13.72,7.02,1.9538,0.0206,13.77,0.1436,0.1554,0.92 +M=4608,bf16,0.2838,0.0224,0.0405,12.70,7.01,1.8109,0.0216,13.15,0.1436,0.1573,0.91 +M=4608,fp32,0.3018,0.0278,0.0419,10.85,7.20,1.5081,0.0296,10.20,0.2761,0.2877,0.96 +M=5120,fp16,0.3446,0.0209,0.0487,16.52,7.07,2.3367,0.0193,17.88,0.1748,0.1872,0.93 +M=5120,bf16,0.3449,0.0209,0.0498,16.48,6.93,2.3793,0.0195,17.72,0.1748,0.1912,0.91 +M=5120,fp32,0.3677,0.0339,0.0510,10.86,7.21,1.5057,0.0334,11.01,0.3351,0.3525,0.95 +M=5632,fp16,0.4138,0.0294,0.0592,14.06,6.99,2.0096,0.0291,14.23,0.2132,0.2210,0.96 +M=5632,bf16,0.4145,0.0295,0.0591,14.07,7.01,2.0063,0.0314,13.20,0.2132,0.2278,0.94 +M=5632,fp32,0.4428,0.0402,0.0610,11.02,7.26,1.5179,0.0426,10.39,0.4145,0.4242,0.98 +M=6144,fp16,0.4882,0.0271,0.0710,18.03,6.87,2.6232,0.0251,19.44,0.2461,0.2603,0.95 +M=6144,bf16,0.4885,0.0271,0.0710,18.06,6.88,2.6238,0.0251,19.45,0.2461,0.2683,0.92 +M=6144,fp32,0.5225,0.0467,0.0737,11.19,7.09,1.5769,0.0475,11.00,0.4839,0.5043,0.96 +M=6656,fp16,0.5698,0.0399,0.0806,14.29,7.07,2.0206,0.0388,14.68,nan,nan,0.00 +M=6656,bf16,0.5696,0.0399,0.0819,14.26,6.95,2.0510,0.0311,18.34,nan,nan,0.00 +M=6656,fp32,0.6130,0.0540,0.0838,11.36,7.32,1.5531,0.0591,10.38,nan,nan,0.00 +M=7168,fp16,0.6586,0.0378,0.0935,17.40,7.05,2.4695,0.0363,18.13,nan,nan,0.00 +M=7168,bf16,0.6562,0.0375,0.0954,17.49,6.88,2.5431,0.0359,18.29,nan,nan,0.00 +M=7168,fp32,0.7073,0.0623,0.0978,11.35,7.23,1.5704,0.0625,11.31,nan,nan,0.00 +M=7680,fp16,0.7507,0.0537,0.1086,13.97,6.91,2.0213,0.0373,20.15,nan,nan,0.00 +M=7680,bf16,0.7519,0.0537,0.1086,14.01,6.92,2.0238,0.0375,20.03,nan,nan,0.00 +M=7680,fp32,0.8091,0.0713,0.1130,11.34,7.16,1.5847,0.0752,10.76,nan,nan,0.00 +M=8192,fp16,0.8524,0.0459,0.0919,18.58,9.28,2.0027,0.0419,20.33,nan,nan,0.00 +M=8192,bf16,0.8530,0.0458,0.0918,18.62,9.29,2.0052,0.0421,20.24,nan,nan,0.00 +M=8192,fp32,0.9187,0.0799,0.1006,11.50,9.14,1.2582,0.0807,11.39,nan,nan,0.00 +M=8704,fp16,0.9583,0.0659,0.1351,14.54,7.09,2.0496,0.0471,20.36,nan,nan,0.00 +M=8704,bf16,0.9572,0.0659,0.1375,14.52,6.96,2.0843,0.0501,19.12,nan,nan,0.00 +M=8704,fp32,1.0333,0.0902,0.1369,11.46,7.55,1.5179,0.0954,10.83,nan,nan,0.00 +M=9216,fp16,1.0729,0.0591,0.1517,18.16,7.07,2.5672,0.0558,19.24,nan,nan,0.00 +M=9216,bf16,1.0721,0.0593,0.1516,18.09,7.07,2.5583,0.0568,18.89,nan,nan,0.00 +M=9216,fp32,1.1569,0.0998,0.1510,11.59,7.66,1.5128,0.1013,11.42,nan,nan,0.00 +M=9728,fp16,1.1919,0.0816,0.1710,14.60,6.97,2.0951,0.0585,20.39,nan,nan,0.00 +M=9728,bf16,1.1929,0.0816,0.1682,14.62,7.09,2.0620,0.0583,20.46,nan,nan,0.00 +M=9728,fp32,1.2859,0.1111,0.1696,11.58,7.58,1.5268,0.1123,11.45,nan,nan,0.00 +M=10240,fp16,1.3179,0.0679,0.1864,19.42,7.07,2.7475,0.0636,20.73,nan,nan,0.00 +M=10240,bf16,1.3198,0.0677,0.1865,19.50,7.08,2.7561,0.0644,20.51,nan,nan,0.00 +M=10240,fp32,1.4256,0.1225,0.1909,11.64,7.47,1.5580,0.1241,11.49,nan,nan,0.00 diff --git a/results/csv/weight_dequant_default.csv b/results/csv/weight_dequant_default.csv index 69ff522d..aa7020f3 100644 --- a/results/csv/weight_dequant_default.csv +++ b/results/csv/weight_dequant_default.csv @@ -1,61 +1,61 @@ params,dtype,torch_ms,triton_ms,cutile_ms,speedup_triton,speedup_cutile,triton_vs_cutile,tilelang_ms,speedup_tilelang,torch_nki_ms,nki_ms,speedup_nki -M=512,fp16,0.0185,0.0021,0.0022,8.78,8.22,1.0684,0.0019,9.84,2.5048,0.0351,71.41 -M=512,bf16,0.0173,0.0019,0.0021,8.97,8.27,1.0848,0.0017,10.43,2.5047,0.0350,71.65 -M=512,fp32,0.0185,0.0021,0.0025,8.66,7.53,1.1504,0.0017,10.64,2.5056,0.0356,70.34 -M=1024,fp16,0.0276,0.0025,0.0031,10.84,8.95,1.2111,0.0024,11.32,9.9260,0.1009,98.34 -M=1024,bf16,0.0266,0.0025,0.0031,10.80,8.53,1.2652,0.0021,12.75,9.9259,0.1009,98.33 -M=1024,fp32,0.0265,0.0028,0.0037,9.56,7.15,1.3378,0.0027,9.70,9.9266,0.1003,98.98 -M=1536,fp16,0.0390,0.0043,0.0064,9.16,6.04,1.5158,0.0036,10.75,20.2811,0.2074,97.77 -M=1536,bf16,0.0381,0.0043,0.0064,8.94,5.94,1.5039,0.0038,10.09,20.2811,0.2074,97.77 -M=1536,fp32,0.0388,0.0044,0.0068,8.73,5.69,1.5338,0.0048,8.02,13.9147,0.2079,66.92 -M=2048,fp16,0.0583,0.0046,0.0077,12.60,7.61,1.6569,0.0041,14.08,25.6457,0.3570,71.84 -M=2048,bf16,0.0570,0.0044,0.0077,13.09,7.38,1.7740,0.0042,13.50,25.6618,0.3570,71.88 -M=2048,fp32,0.0598,0.0072,0.0095,8.36,6.28,1.3315,0.0066,9.02,24.7112,0.3589,68.84 -M=2560,fp16,0.0920,0.0083,0.0140,11.04,6.57,1.6801,0.0078,11.80,38.5823,0.5483,70.37 -M=2560,bf16,0.0913,0.0085,0.0141,10.81,6.48,1.6664,0.0056,16.30,38.5821,0.5483,70.37 -M=2560,fp32,0.0969,0.0110,0.0161,8.81,6.02,1.4627,0.0114,8.51,38.6094,0.5491,70.31 -M=3072,fp16,0.1331,0.0089,0.0197,14.89,6.74,2.2069,0.0084,15.87,55.5372,0.7817,71.04 -M=3072,bf16,0.1322,0.0088,0.0197,15.01,6.73,2.2321,0.0084,15.66,55.5370,0.7818,71.04 -M=3072,fp32,0.1405,0.0145,0.0223,9.69,6.31,1.5347,0.0140,10.01,55.5738,0.7818,71.09 -M=3584,fp16,0.1792,0.0146,0.0263,12.26,6.82,1.7980,0.0150,11.92,75.5826,1.0570,71.51 -M=3584,bf16,0.1788,0.0146,0.0263,12.28,6.80,1.8042,0.0110,16.29,75.5828,1.0569,71.51 -M=3584,fp32,0.1886,0.0181,0.0289,10.41,6.52,1.5968,0.0199,9.48,75.6304,1.0578,71.49 -M=4096,fp16,0.2264,0.0146,0.0251,15.51,9.02,1.7200,0.0134,16.89,99.5658,1.3753,72.39 -M=4096,bf16,0.2264,0.0147,0.0251,15.44,9.03,1.7093,0.0137,16.57,99.5027,1.3755,72.34 -M=4096,fp32,0.2420,0.0229,0.0307,10.58,7.88,1.3436,0.0226,10.70,98.8365,1.3760,71.83 -M=4608,fp16,0.2830,0.0222,0.0405,12.76,6.99,1.8253,0.0228,12.42,124.9118,1.7339,72.04 -M=4608,bf16,0.2834,0.0222,0.0405,12.78,7.00,1.8254,0.0227,12.47,124.9119,1.7339,72.04 -M=4608,fp32,0.3021,0.0280,0.0457,10.78,6.60,1.6319,0.0306,9.89,125.0999,1.7367,72.03 -M=5120,fp16,0.3459,0.0209,0.0498,16.59,6.95,2.3879,0.0194,17.79,154.1932,2.1360,72.19 -M=5120,bf16,0.3458,0.0208,0.0498,16.59,6.94,2.3895,0.0195,17.70,154.1934,2.1361,72.19 -M=5120,fp32,0.3668,0.0335,0.0554,10.96,6.63,1.6542,0.0333,11.02,154.3176,2.1364,72.23 -M=5632,fp16,0.4140,0.0316,0.0592,13.08,7.00,1.8697,0.0328,12.61,186.6698,2.5780,72.41 -M=5632,bf16,0.4140,0.0316,0.0591,13.09,7.00,1.8702,0.0324,12.77,186.6701,2.5780,72.41 -M=5632,fp32,0.4416,0.0400,0.0669,11.05,6.60,1.6742,0.0436,10.12,186.8025,2.5789,72.44 -M=6144,fp16,0.4886,0.0287,0.0710,17.04,6.88,2.4769,0.0268,18.21,222.0560,3.0642,72.47 -M=6144,bf16,0.4884,0.0287,0.0710,17.02,6.88,2.4752,0.0269,18.18,222.0558,3.0642,72.47 -M=6144,fp32,0.5215,0.0465,0.0802,11.22,6.50,1.7256,0.0466,11.18,222.2288,3.0641,72.53 -M=6656,fp16,0.5704,0.0431,0.0815,13.25,7.00,1.8935,0.0446,12.78,260.5408,3.5904,72.57 -M=6656,bf16,0.5700,0.0430,0.0815,13.25,6.99,1.8943,0.0310,18.39,260.5411,3.5902,72.57 -M=6656,fp32,0.6108,0.0546,0.0922,11.20,6.62,1.6903,0.0594,10.28,260.8216,3.5915,72.62 -M=7168,fp16,0.6585,0.0375,0.0954,17.55,6.90,2.5431,0.0354,18.61,302.2148,4.1605,72.64 -M=7168,bf16,0.6582,0.0375,0.0954,17.54,6.90,2.5429,0.0353,18.64,302.2146,4.1606,72.64 -M=7168,fp32,0.7053,0.0621,0.1076,11.35,6.55,1.7325,0.0626,11.27,302.4038,4.1733,72.46 -M=7680,fp16,0.7533,0.0571,0.1112,13.20,6.77,1.9479,0.0401,18.80,346.9441,4.7720,72.70 -M=7680,bf16,0.7545,0.0570,0.1112,13.23,6.78,1.9494,0.0572,13.20,346.9451,4.7720,72.70 -M=7680,fp32,0.8091,0.0718,0.1220,11.27,6.63,1.6995,0.0772,10.47,347.1352,4.7817,72.60 -M=8192,fp16,0.8525,0.0490,0.0918,17.40,9.28,1.8742,0.0450,18.93,394.6603,5.4231,72.77 -M=8192,bf16,0.8515,0.0490,0.0918,17.38,9.27,1.8743,0.0456,18.65,394.6604,5.4231,72.77 -M=8192,fp32,0.9158,0.0799,0.1110,11.46,8.25,1.3891,0.0804,11.39,631.9603,5.4367,116.24 -M=8704,fp16,0.9572,0.0715,0.1374,13.39,6.97,1.9219,0.0730,13.12,445.8699,6.1203,72.85 -M=8704,bf16,0.9589,0.0716,0.1374,13.40,6.98,1.9205,0.0502,19.12,445.8691,6.1206,72.85 -M=8704,fp32,1.0309,0.0905,0.1555,11.39,6.63,1.7183,0.0978,10.54,446.3709,6.1300,72.82 -M=9216,fp16,1.0719,0.0593,0.1553,18.08,6.90,2.6192,0.0558,19.22,501.9925,6.8546,73.23 -M=9216,bf16,1.0724,0.0593,0.1553,18.09,6.91,2.6197,0.0557,19.27,501.9926,6.8546,73.23 -M=9216,fp32,1.1555,0.1002,0.1725,11.54,6.70,1.7223,0.1011,11.43,499.9652,6.8636,72.84 -M=9728,fp16,1.1920,0.0886,0.1710,13.45,6.97,1.9293,0.0901,13.23,556.8344,7.6341,72.94 -M=9728,bf16,1.1927,0.0885,0.1709,13.47,6.98,1.9306,0.0890,13.41,556.8358,7.6341,72.94 -M=9728,fp32,1.2885,0.1115,0.1934,11.56,6.66,1.7353,0.1122,11.49,557.5616,7.6355,73.02 -M=10240,fp16,1.3192,0.0721,0.1910,18.30,6.91,2.6506,0.0681,19.37,616.5930,8.4559,72.92 -M=10240,bf16,1.3175,0.0720,0.1910,18.31,6.90,2.6544,0.0679,19.39,616.5933,8.4559,72.92 -M=10240,fp32,1.4237,0.1229,0.2123,11.58,6.71,1.7272,0.1243,11.45,617.2312,8.4566,72.99 +M=512,fp16,0.0185,0.0021,0.0022,8.78,8.22,1.0684,0.0019,9.84,0.0231,0.0212,1.09 +M=512,bf16,0.0173,0.0019,0.0021,8.97,8.27,1.0848,0.0017,10.43,0.0231,0.0212,1.09 +M=512,fp32,0.0185,0.0021,0.0025,8.66,7.53,1.1504,0.0017,10.64,0.0260,0.0224,1.16 +M=1024,fp16,0.0276,0.0025,0.0031,10.84,8.95,1.2111,0.0024,11.32,0.0240,0.0306,0.78 +M=1024,bf16,0.0266,0.0025,0.0031,10.80,8.53,1.2652,0.0021,12.75,0.0240,0.0306,0.78 +M=1024,fp32,0.0265,0.0028,0.0037,9.56,7.15,1.3378,0.0027,9.70,0.0257,0.0334,0.77 +M=1536,fp16,0.0390,0.0043,0.0064,9.16,6.04,1.5158,0.0036,10.75,0.0292,0.0418,0.70 +M=1536,bf16,0.0381,0.0043,0.0064,8.94,5.94,1.5039,0.0038,10.09,0.0292,0.0418,0.70 +M=1536,fp32,0.0388,0.0044,0.0068,8.73,5.69,1.5338,0.0048,8.02,0.0438,0.0475,0.92 +M=2048,fp16,0.0583,0.0046,0.0077,12.60,7.61,1.6569,0.0041,14.08,0.0382,0.0476,0.80 +M=2048,bf16,0.0570,0.0044,0.0077,13.09,7.38,1.7740,0.0042,13.50,0.0382,0.0476,0.80 +M=2048,fp32,0.0598,0.0072,0.0095,8.36,6.28,1.3315,0.0066,9.02,0.0698,0.0759,0.92 +M=2560,fp16,0.0920,0.0083,0.0140,11.04,6.57,1.6801,0.0078,11.80,0.0537,0.0638,0.84 +M=2560,bf16,0.0913,0.0085,0.0141,10.81,6.48,1.6664,0.0056,16.30,0.0535,0.0636,0.84 +M=2560,fp32,0.0969,0.0110,0.0161,8.81,6.02,1.4627,0.0114,8.51,0.0926,0.1029,0.90 +M=3072,fp16,0.1331,0.0089,0.0197,14.89,6.74,2.2069,0.0084,15.87,0.0699,0.0783,0.89 +M=3072,bf16,0.1322,0.0088,0.0197,15.01,6.73,2.2321,0.0084,15.66,0.0699,0.0782,0.89 +M=3072,fp32,0.1405,0.0145,0.0223,9.69,6.31,1.5347,0.0140,10.01,0.1279,0.1424,0.90 +M=3584,fp16,0.1792,0.0146,0.0263,12.26,6.82,1.7980,0.0150,11.92,0.0910,0.1036,0.88 +M=3584,bf16,0.1788,0.0146,0.0263,12.28,6.80,1.8042,0.0110,16.29,0.0910,0.1038,0.88 +M=3584,fp32,0.1886,0.0181,0.0289,10.41,6.52,1.5968,0.0199,9.48,0.1731,0.1811,0.96 +M=4096,fp16,0.2264,0.0146,0.0251,15.51,9.02,1.7200,0.0134,16.89,0.1318,0.1532,0.86 +M=4096,bf16,0.2264,0.0147,0.0251,15.44,9.03,1.7093,0.0137,16.57,0.1316,0.1530,0.86 +M=4096,fp32,0.2420,0.0229,0.0307,10.58,7.88,1.3436,0.0226,10.70,0.2404,0.2491,0.97 +M=4608,fp16,0.2830,0.0222,0.0405,12.76,6.99,1.8253,0.0228,12.42,0.1436,0.1577,0.91 +M=4608,bf16,0.2834,0.0222,0.0405,12.78,7.00,1.8254,0.0227,12.47,0.1436,0.1577,0.91 +M=4608,fp32,0.3021,0.0280,0.0457,10.78,6.60,1.6319,0.0306,9.89,0.2762,0.2892,0.96 +M=5120,fp16,0.3459,0.0209,0.0498,16.59,6.95,2.3879,0.0194,17.79,0.1748,0.2003,0.87 +M=5120,bf16,0.3458,0.0208,0.0498,16.59,6.94,2.3895,0.0195,17.70,0.1748,0.2004,0.87 +M=5120,fp32,0.3668,0.0335,0.0554,10.96,6.63,1.6542,0.0333,11.02,0.3350,0.3516,0.95 +M=5632,fp16,0.4140,0.0316,0.0592,13.08,7.00,1.8697,0.0328,12.61,0.2132,0.2231,0.96 +M=5632,bf16,0.4140,0.0316,0.0591,13.09,7.00,1.8702,0.0324,12.77,0.2133,0.2230,0.96 +M=5632,fp32,0.4416,0.0400,0.0669,11.05,6.60,1.6742,0.0436,10.12,0.4143,0.4248,0.98 +M=6144,fp16,0.4886,0.0287,0.0710,17.04,6.88,2.4769,0.0268,18.21,0.2458,0.2684,0.92 +M=6144,bf16,0.4884,0.0287,0.0710,17.02,6.88,2.4752,0.0269,18.18,0.2461,0.2686,0.92 +M=6144,fp32,0.5215,0.0465,0.0802,11.22,6.50,1.7256,0.0466,11.18,0.4844,0.5079,0.95 +M=6656,fp16,0.5704,0.0431,0.0815,13.25,7.00,1.8935,0.0446,12.78,0.2853,0.3046,0.94 +M=6656,bf16,0.5700,0.0430,0.0815,13.25,6.99,1.8943,0.0310,18.39,0.2855,0.3049,0.94 +M=6656,fp32,0.6108,0.0546,0.0922,11.20,6.62,1.6903,0.0594,10.28,0.5618,0.5814,0.97 +M=7168,fp16,0.6585,0.0375,0.0954,17.55,6.90,2.5431,0.0354,18.61,0.3300,0.3599,0.92 +M=7168,bf16,0.6582,0.0375,0.0954,17.54,6.90,2.5429,0.0353,18.64,0.3300,0.3598,0.92 +M=7168,fp32,0.7053,0.0621,0.1076,11.35,6.55,1.7325,0.0626,11.27,0.6460,0.6734,0.96 +M=7680,fp16,0.7533,0.0571,0.1112,13.20,6.77,1.9479,0.0401,18.80,0.3792,0.4118,0.92 +M=7680,bf16,0.7545,0.0570,0.1112,13.23,6.78,1.9494,0.0572,13.20,0.3789,0.4119,0.92 +M=7680,fp32,0.8091,0.0718,0.1220,11.27,6.63,1.6995,0.0772,10.47,0.7388,0.7658,0.96 +M=8192,fp16,0.8525,0.0490,0.0918,17.40,9.28,1.8742,0.0450,18.93,0.4583,0.5156,0.89 +M=8192,bf16,0.8515,0.0490,0.0918,17.38,9.27,1.8743,0.0456,18.65,0.4584,0.5145,0.89 +M=8192,fp32,0.9158,0.0799,0.1110,11.46,8.25,1.3891,0.0804,11.39,0.8366,0.8900,0.94 +M=8704,fp16,0.9572,0.0715,0.1374,13.39,6.97,1.9219,0.0730,13.12,0.4756,0.5299,0.90 +M=8704,bf16,0.9589,0.0716,0.1374,13.40,6.98,1.9205,0.0502,19.12,0.4759,0.5300,0.90 +M=8704,fp32,1.0309,0.0905,0.1555,11.39,6.63,1.7183,0.0978,10.54,0.9468,0.9730,0.97 +M=9216,fp16,1.0719,0.0593,0.1553,18.08,6.90,2.6192,0.0558,19.22,0.5416,0.5744,0.94 +M=9216,bf16,1.0724,0.0593,0.1553,18.09,6.91,2.6197,0.0557,19.27,0.5418,0.5743,0.94 +M=9216,fp32,1.1555,0.1002,0.1725,11.54,6.70,1.7223,0.1011,11.43,1.0612,1.0958,0.97 +M=9728,fp16,1.1920,0.0886,0.1710,13.45,6.97,1.9293,0.0901,13.23,0.5951,0.6384,0.93 +M=9728,bf16,1.1927,0.0885,0.1709,13.47,6.98,1.9306,0.0890,13.41,0.5950,0.6376,0.93 +M=9728,fp32,1.2885,0.1115,0.1934,11.56,6.66,1.7353,0.1122,11.49,1.1789,1.2131,0.97 +M=10240,fp16,1.3192,0.0721,0.1910,18.30,6.91,2.6506,0.0681,19.37,0.6570,0.6993,0.94 +M=10240,bf16,1.3175,0.0720,0.1910,18.31,6.90,2.6544,0.0679,19.39,0.6569,0.6993,0.94 +M=10240,fp32,1.4237,0.1229,0.2123,11.58,6.71,1.7272,0.1243,11.45,1.3004,1.3370,0.97 diff --git a/results/logs/autotune_logs/weight_dequant_autotune.json b/results/logs/autotune_logs/weight_dequant_autotune.json index 7cf25696..c90ae45f 100644 --- a/results/logs/autotune_logs/weight_dequant_autotune.json +++ b/results/logs/autotune_logs/weight_dequant_autotune.json @@ -13,6 +13,9 @@ "cutile_autotune_cfg": { "tile": 512, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 512 } }, { @@ -29,6 +32,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 256 } }, { @@ -45,6 +51,9 @@ "cutile_autotune_cfg": { "tile": 512, "occupancy": 16 + }, + "nki_autotune_cfg": { + "block_size": 2048 } }, { @@ -61,6 +70,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 2048 } }, { @@ -77,6 +89,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 16 + }, + "nki_autotune_cfg": { + "block_size": 512 } }, { @@ -93,6 +108,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 1024 } }, { @@ -109,6 +127,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -125,6 +146,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 2048 } }, { @@ -141,6 +165,9 @@ "cutile_autotune_cfg": { "tile": 512, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -157,6 +184,9 @@ "cutile_autotune_cfg": { "tile": 4096, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -173,6 +203,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 2048 } }, { @@ -189,6 +222,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 16 + }, + "nki_autotune_cfg": { + "block_size": 2048 } }, { @@ -205,6 +241,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 2048 } }, { @@ -221,6 +260,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 2048 } }, { @@ -237,6 +279,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -253,6 +298,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 2048 } }, { @@ -269,6 +317,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 512 } }, { @@ -285,6 +336,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -301,6 +355,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 16 + }, + "nki_autotune_cfg": { + "block_size": 2048 } }, { @@ -317,6 +374,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 16 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -333,6 +393,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 1024 } }, { @@ -349,6 +412,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 16 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -365,6 +431,9 @@ "cutile_autotune_cfg": { "tile": 4096, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 2048 } }, { @@ -381,6 +450,9 @@ "cutile_autotune_cfg": { "tile": 4096, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 2048 } }, { @@ -397,6 +469,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 2048 } }, { @@ -413,6 +488,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -429,6 +507,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 2048 } }, { @@ -445,6 +526,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 16 + }, + "nki_autotune_cfg": { + "block_size": 2048 } }, { @@ -461,6 +545,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -477,6 +564,9 @@ "cutile_autotune_cfg": { "tile": 4096, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -493,6 +583,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 2048 } }, { @@ -509,6 +602,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -525,6 +621,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 2048 } }, { @@ -541,6 +640,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -557,6 +659,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 1024 } }, { @@ -573,6 +678,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -589,7 +697,8 @@ "cutile_autotune_cfg": { "tile": 4096, "occupancy": 4 - } + }, + "nki_autotune_cfg": null }, { "params": { @@ -605,7 +714,8 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 16 - } + }, + "nki_autotune_cfg": null }, { "params": { @@ -621,7 +731,8 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 4 - } + }, + "nki_autotune_cfg": null }, { "params": { @@ -637,7 +748,8 @@ "cutile_autotune_cfg": { "tile": 4096, "occupancy": 8 - } + }, + "nki_autotune_cfg": null }, { "params": { @@ -653,7 +765,8 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 8 - } + }, + "nki_autotune_cfg": null }, { "params": { @@ -669,7 +782,8 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 4 - } + }, + "nki_autotune_cfg": null }, { "params": { @@ -685,7 +799,8 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 16 - } + }, + "nki_autotune_cfg": null }, { "params": { @@ -701,7 +816,8 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 16 - } + }, + "nki_autotune_cfg": null }, { "params": { @@ -717,7 +833,8 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 4 - } + }, + "nki_autotune_cfg": null }, { "params": { @@ -733,7 +850,8 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 8 - } + }, + "nki_autotune_cfg": null }, { "params": { @@ -749,7 +867,8 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 4 - } + }, + "nki_autotune_cfg": null }, { "params": { @@ -765,7 +884,8 @@ "cutile_autotune_cfg": { "tile": 4096, "occupancy": 16 - } + }, + "nki_autotune_cfg": null }, { "params": { @@ -781,7 +901,8 @@ "cutile_autotune_cfg": { "tile": 4096, "occupancy": 4 - } + }, + "nki_autotune_cfg": null }, { "params": { @@ -797,7 +918,8 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 4 - } + }, + "nki_autotune_cfg": null }, { "params": { @@ -813,7 +935,8 @@ "cutile_autotune_cfg": { "tile": 4096, "occupancy": 8 - } + }, + "nki_autotune_cfg": null }, { "params": { @@ -829,7 +952,8 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 16 - } + }, + "nki_autotune_cfg": null }, { "params": { @@ -845,7 +969,8 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 16 - } + }, + "nki_autotune_cfg": null }, { "params": { @@ -861,7 +986,8 @@ "cutile_autotune_cfg": { "tile": 4096, "occupancy": 16 - } + }, + "nki_autotune_cfg": null }, { "params": { @@ -877,7 +1003,8 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 4 - } + }, + "nki_autotune_cfg": null }, { "params": { @@ -893,7 +1020,8 @@ "cutile_autotune_cfg": { "tile": 4096, "occupancy": 4 - } + }, + "nki_autotune_cfg": null }, { "params": { @@ -909,7 +1037,8 @@ "cutile_autotune_cfg": { "tile": 4096, "occupancy": 8 - } + }, + "nki_autotune_cfg": null }, { "params": { @@ -925,7 +1054,8 @@ "cutile_autotune_cfg": { "tile": 4096, "occupancy": 4 - } + }, + "nki_autotune_cfg": null }, { "params": { @@ -941,7 +1071,8 @@ "cutile_autotune_cfg": { "tile": 4096, "occupancy": 4 - } + }, + "nki_autotune_cfg": null }, { "params": { @@ -957,6 +1088,7 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 4 - } + }, + "nki_autotune_cfg": null } -] \ No newline at end of file +] From f200a8414e9810c057b0d3eeced25acf7cac84c2 Mon Sep 17 00:00:00 2001 From: Bowen Cui Date: Thu, 17 Sep 2026 00:25:14 +0000 Subject: [PATCH 6/8] nki(weight_dequant): fill the 24 autotune cases lost to a full disk Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01ScXYNjrrKGgDUVNHxv7HJt --- results/csv/weight_dequant_autotune.csv | 48 +++++----- .../weight_dequant_autotune.json | 96 ++++++++++++++----- 2 files changed, 96 insertions(+), 48 deletions(-) diff --git a/results/csv/weight_dequant_autotune.csv b/results/csv/weight_dequant_autotune.csv index f5c8c9c8..2d0eb34b 100644 --- a/results/csv/weight_dequant_autotune.csv +++ b/results/csv/weight_dequant_autotune.csv @@ -35,27 +35,27 @@ M=5632,fp32,0.4428,0.0402,0.0610,11.02,7.26,1.5179,0.0426,10.39,0.4145,0.4242,0. M=6144,fp16,0.4882,0.0271,0.0710,18.03,6.87,2.6232,0.0251,19.44,0.2461,0.2603,0.95 M=6144,bf16,0.4885,0.0271,0.0710,18.06,6.88,2.6238,0.0251,19.45,0.2461,0.2683,0.92 M=6144,fp32,0.5225,0.0467,0.0737,11.19,7.09,1.5769,0.0475,11.00,0.4839,0.5043,0.96 -M=6656,fp16,0.5698,0.0399,0.0806,14.29,7.07,2.0206,0.0388,14.68,nan,nan,0.00 -M=6656,bf16,0.5696,0.0399,0.0819,14.26,6.95,2.0510,0.0311,18.34,nan,nan,0.00 -M=6656,fp32,0.6130,0.0540,0.0838,11.36,7.32,1.5531,0.0591,10.38,nan,nan,0.00 -M=7168,fp16,0.6586,0.0378,0.0935,17.40,7.05,2.4695,0.0363,18.13,nan,nan,0.00 -M=7168,bf16,0.6562,0.0375,0.0954,17.49,6.88,2.5431,0.0359,18.29,nan,nan,0.00 -M=7168,fp32,0.7073,0.0623,0.0978,11.35,7.23,1.5704,0.0625,11.31,nan,nan,0.00 -M=7680,fp16,0.7507,0.0537,0.1086,13.97,6.91,2.0213,0.0373,20.15,nan,nan,0.00 -M=7680,bf16,0.7519,0.0537,0.1086,14.01,6.92,2.0238,0.0375,20.03,nan,nan,0.00 -M=7680,fp32,0.8091,0.0713,0.1130,11.34,7.16,1.5847,0.0752,10.76,nan,nan,0.00 -M=8192,fp16,0.8524,0.0459,0.0919,18.58,9.28,2.0027,0.0419,20.33,nan,nan,0.00 -M=8192,bf16,0.8530,0.0458,0.0918,18.62,9.29,2.0052,0.0421,20.24,nan,nan,0.00 -M=8192,fp32,0.9187,0.0799,0.1006,11.50,9.14,1.2582,0.0807,11.39,nan,nan,0.00 -M=8704,fp16,0.9583,0.0659,0.1351,14.54,7.09,2.0496,0.0471,20.36,nan,nan,0.00 -M=8704,bf16,0.9572,0.0659,0.1375,14.52,6.96,2.0843,0.0501,19.12,nan,nan,0.00 -M=8704,fp32,1.0333,0.0902,0.1369,11.46,7.55,1.5179,0.0954,10.83,nan,nan,0.00 -M=9216,fp16,1.0729,0.0591,0.1517,18.16,7.07,2.5672,0.0558,19.24,nan,nan,0.00 -M=9216,bf16,1.0721,0.0593,0.1516,18.09,7.07,2.5583,0.0568,18.89,nan,nan,0.00 -M=9216,fp32,1.1569,0.0998,0.1510,11.59,7.66,1.5128,0.1013,11.42,nan,nan,0.00 -M=9728,fp16,1.1919,0.0816,0.1710,14.60,6.97,2.0951,0.0585,20.39,nan,nan,0.00 -M=9728,bf16,1.1929,0.0816,0.1682,14.62,7.09,2.0620,0.0583,20.46,nan,nan,0.00 -M=9728,fp32,1.2859,0.1111,0.1696,11.58,7.58,1.5268,0.1123,11.45,nan,nan,0.00 -M=10240,fp16,1.3179,0.0679,0.1864,19.42,7.07,2.7475,0.0636,20.73,nan,nan,0.00 -M=10240,bf16,1.3198,0.0677,0.1865,19.50,7.08,2.7561,0.0644,20.51,nan,nan,0.00 -M=10240,fp32,1.4256,0.1225,0.1909,11.64,7.47,1.5580,0.1241,11.49,nan,nan,0.00 +M=6656,fp16,0.5698,0.0399,0.0806,14.29,7.07,2.0206,0.0388,14.68,0.2854,0.3015,0.95 +M=6656,bf16,0.5696,0.0399,0.0819,14.26,6.95,2.0510,0.0311,18.34,0.2854,0.3016,0.95 +M=6656,fp32,0.6130,0.0540,0.0838,11.36,7.32,1.5531,0.0591,10.38,0.5621,0.5750,0.98 +M=7168,fp16,0.6586,0.0378,0.0935,17.40,7.05,2.4695,0.0363,18.13,0.3299,0.3459,0.95 +M=7168,bf16,0.6562,0.0375,0.0954,17.49,6.88,2.5431,0.0359,18.29,0.3300,0.3460,0.95 +M=7168,fp32,0.7073,0.0623,0.0978,11.35,7.23,1.5704,0.0625,11.31,0.6460,0.6736,0.96 +M=7680,fp16,0.7507,0.0537,0.1086,13.97,6.91,2.0213,0.0373,20.15,0.3791,0.4000,0.95 +M=7680,bf16,0.7519,0.0537,0.1086,14.01,6.92,2.0238,0.0375,20.03,0.3792,0.4121,0.92 +M=7680,fp32,0.8091,0.0713,0.1130,11.34,7.16,1.5847,0.0752,10.76,0.7386,0.7662,0.96 +M=8192,fp16,0.8524,0.0459,0.0919,18.58,9.28,2.0027,0.0419,20.33,0.4587,0.4540,1.01 +M=8192,bf16,0.8530,0.0458,0.0918,18.62,9.29,2.0052,0.0421,20.24,0.4586,0.4536,1.01 +M=8192,fp32,0.9187,0.0799,0.1006,11.50,9.14,1.2582,0.0807,11.39,0.8369,0.8901,0.94 +M=8704,fp16,0.9583,0.0659,0.1351,14.54,7.09,2.0496,0.0471,20.36,0.4758,0.5016,0.95 +M=8704,bf16,0.9572,0.0659,0.1375,14.52,6.96,2.0843,0.0501,19.12,0.4757,0.5013,0.95 +M=8704,fp32,1.0333,0.0902,0.1369,11.46,7.55,1.5179,0.0954,10.83,0.9467,0.9718,0.97 +M=9216,fp16,1.0729,0.0591,0.1517,18.16,7.07,2.5672,0.0558,19.24,0.5418,0.5585,0.97 +M=9216,bf16,1.0721,0.0593,0.1516,18.09,7.07,2.5583,0.0568,18.89,0.5418,0.5742,0.94 +M=9216,fp32,1.1569,0.0998,0.1510,11.59,7.66,1.5128,0.1013,11.42,1.0613,1.0950,0.97 +M=9728,fp16,1.1919,0.0816,0.1710,14.60,6.97,2.0951,0.0585,20.39,0.5949,0.6189,0.96 +M=9728,bf16,1.1929,0.0816,0.1682,14.62,7.09,2.0620,0.0583,20.46,0.5944,0.6255,0.95 +M=9728,fp32,1.2859,0.1111,0.1696,11.58,7.58,1.5268,0.1123,11.45,1.1791,1.2133,0.97 +M=10240,fp16,1.3179,0.0679,0.1864,19.42,7.07,2.7475,0.0636,20.73,0.6569,0.6869,0.96 +M=10240,bf16,1.3198,0.0677,0.1865,19.50,7.08,2.7561,0.0644,20.51,0.6569,0.6867,0.96 +M=10240,fp32,1.4256,0.1225,0.1909,11.64,7.47,1.5580,0.1241,11.49,1.3007,1.3244,0.98 diff --git a/results/logs/autotune_logs/weight_dequant_autotune.json b/results/logs/autotune_logs/weight_dequant_autotune.json index c90ae45f..ac34168e 100644 --- a/results/logs/autotune_logs/weight_dequant_autotune.json +++ b/results/logs/autotune_logs/weight_dequant_autotune.json @@ -698,7 +698,9 @@ "tile": 4096, "occupancy": 4 }, - "nki_autotune_cfg": null + "nki_autotune_cfg": { + "block_size": 4096 + } }, { "params": { @@ -715,7 +717,9 @@ "tile": 1024, "occupancy": 16 }, - "nki_autotune_cfg": null + "nki_autotune_cfg": { + "block_size": 4096 + } }, { "params": { @@ -732,7 +736,9 @@ "tile": 2048, "occupancy": 4 }, - "nki_autotune_cfg": null + "nki_autotune_cfg": { + "block_size": 2048 + } }, { "params": { @@ -749,7 +755,9 @@ "tile": 4096, "occupancy": 8 }, - "nki_autotune_cfg": null + "nki_autotune_cfg": { + "block_size": 4096 + } }, { "params": { @@ -766,7 +774,9 @@ "tile": 1024, "occupancy": 8 }, - "nki_autotune_cfg": null + "nki_autotune_cfg": { + "block_size": 4096 + } }, { "params": { @@ -783,7 +793,9 @@ "tile": 2048, "occupancy": 4 }, - "nki_autotune_cfg": null + "nki_autotune_cfg": { + "block_size": 1024 + } }, { "params": { @@ -800,7 +812,9 @@ "tile": 1024, "occupancy": 16 }, - "nki_autotune_cfg": null + "nki_autotune_cfg": { + "block_size": 2048 + } }, { "params": { @@ -817,7 +831,9 @@ "tile": 1024, "occupancy": 16 }, - "nki_autotune_cfg": null + "nki_autotune_cfg": { + "block_size": 1024 + } }, { "params": { @@ -834,7 +850,9 @@ "tile": 2048, "occupancy": 4 }, - "nki_autotune_cfg": null + "nki_autotune_cfg": { + "block_size": 1024 + } }, { "params": { @@ -851,7 +869,9 @@ "tile": 1024, "occupancy": 8 }, - "nki_autotune_cfg": null + "nki_autotune_cfg": { + "block_size": 4096 + } }, { "params": { @@ -868,7 +888,9 @@ "tile": 1024, "occupancy": 4 }, - "nki_autotune_cfg": null + "nki_autotune_cfg": { + "block_size": 4096 + } }, { "params": { @@ -885,7 +907,9 @@ "tile": 4096, "occupancy": 16 }, - "nki_autotune_cfg": null + "nki_autotune_cfg": { + "block_size": 1024 + } }, { "params": { @@ -902,7 +926,9 @@ "tile": 4096, "occupancy": 4 }, - "nki_autotune_cfg": null + "nki_autotune_cfg": { + "block_size": 2048 + } }, { "params": { @@ -919,7 +945,9 @@ "tile": 1024, "occupancy": 4 }, - "nki_autotune_cfg": null + "nki_autotune_cfg": { + "block_size": 2048 + } }, { "params": { @@ -936,7 +964,9 @@ "tile": 4096, "occupancy": 8 }, - "nki_autotune_cfg": null + "nki_autotune_cfg": { + "block_size": 1024 + } }, { "params": { @@ -953,7 +983,9 @@ "tile": 1024, "occupancy": 16 }, - "nki_autotune_cfg": null + "nki_autotune_cfg": { + "block_size": 2048 + } }, { "params": { @@ -970,7 +1002,9 @@ "tile": 1024, "occupancy": 16 }, - "nki_autotune_cfg": null + "nki_autotune_cfg": { + "block_size": 1024 + } }, { "params": { @@ -987,7 +1021,9 @@ "tile": 4096, "occupancy": 16 }, - "nki_autotune_cfg": null + "nki_autotune_cfg": { + "block_size": 1024 + } }, { "params": { @@ -1004,7 +1040,9 @@ "tile": 1024, "occupancy": 4 }, - "nki_autotune_cfg": null + "nki_autotune_cfg": { + "block_size": 4096 + } }, { "params": { @@ -1021,7 +1059,9 @@ "tile": 4096, "occupancy": 4 }, - "nki_autotune_cfg": null + "nki_autotune_cfg": { + "block_size": 2048 + } }, { "params": { @@ -1038,7 +1078,9 @@ "tile": 4096, "occupancy": 8 }, - "nki_autotune_cfg": null + "nki_autotune_cfg": { + "block_size": 1024 + } }, { "params": { @@ -1055,7 +1097,9 @@ "tile": 4096, "occupancy": 4 }, - "nki_autotune_cfg": null + "nki_autotune_cfg": { + "block_size": 2048 + } }, { "params": { @@ -1072,7 +1116,9 @@ "tile": 4096, "occupancy": 4 }, - "nki_autotune_cfg": null + "nki_autotune_cfg": { + "block_size": 2048 + } }, { "params": { @@ -1089,6 +1135,8 @@ "tile": 2048, "occupancy": 4 }, - "nki_autotune_cfg": null + "nki_autotune_cfg": { + "block_size": 2048 + } } ] From c8f3da38436211158448a5f2d311340877517b0a Mon Sep 17 00:00:00 2001 From: Bowen Cui Date: Mon, 21 Sep 2026 20:38:30 +0000 Subject: [PATCH 7/8] nki(weight_dequant): migrate to the tilebench package layout Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01ScXYNjrrKGgDUVNHxv7HJt --- tilebench/benchmarks/operators/weight_dequant/impl_nki.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tilebench/benchmarks/operators/weight_dequant/impl_nki.py b/tilebench/benchmarks/operators/weight_dequant/impl_nki.py index bb048c69..3294c1f6 100644 --- a/tilebench/benchmarks/operators/weight_dequant/impl_nki.py +++ b/tilebench/benchmarks/operators/weight_dequant/impl_nki.py @@ -25,7 +25,7 @@ import torch -from core.nki_autotune import NkiAutotuner +from tilebench.core.nki_autotune import NkiAutotuner try: import nki From ecc445dbe2dd3f94b50f55a26080e3c3129da992 Mon Sep 17 00:00:00 2001 From: Bowen Cui Date: Tue, 22 Sep 2026 04:57:48 +0000 Subject: [PATCH 8/8] nki(weight_dequant): drop docstrings and inline comments Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01ScXYNjrrKGgDUVNHxv7HJt --- .../operators/weight_dequant/impl_nki.py | 41 +------------------ 1 file changed, 1 insertion(+), 40 deletions(-) diff --git a/tilebench/benchmarks/operators/weight_dequant/impl_nki.py b/tilebench/benchmarks/operators/weight_dequant/impl_nki.py index 3294c1f6..1e0c5957 100644 --- a/tilebench/benchmarks/operators/weight_dequant/impl_nki.py +++ b/tilebench/benchmarks/operators/weight_dequant/impl_nki.py @@ -1,22 +1,3 @@ -"""NKI (AWS Trainium) implementation of weight_dequant: ``out[m, n] = X[m, n] * S[m // T, n // T]``. - -Design (mirrors the Triton / cuTile kernels): - -* Triton walks the flat ``M * N`` elements in ``BLOCK_SIZE`` blocks, looks up the - block scale of every element and multiplies. NKI puts 128 rows on the 128 SBUF - partitions (``TILE_SIZE`` is a multiple of 128, so a row tile never straddles - two scale rows) and walks each row tile in ``BLOCK_SIZE``-column blocks: - - the scale row of the band is replicated to all partitions by a single - partition-stride-0 DMA and converted to fp32 once per row tile; - - per column block: one ``[128, BLOCK]`` DMA, then one Vector-engine - ``tensor_scalar`` per ``TILE_SIZE``-wide sub-block with the matching - ``[128, 1]`` scale column (broadcast along the free axis in hardware), and - one DMA store. - ``BLOCK_SIZE`` defaults / search space are those of ``impl_triton.py`` / - ``impl_cutile.py`` (default 1024; search 256/512/1024/2048/4096). -* Row tiles are split across the NeuronCores of the logical core (LNC2 on trn2) - by ``nl.program_id(0)``. Partial tiles / blocks are clamped with ``min()``. -""" import functools import os import re @@ -39,11 +20,6 @@ @functools.lru_cache(maxsize=1) def _lnc_degree() -> int: - """Logical-NeuronCore degree the kernel is launched with (``kernel[lnc]``). - - Must match the LNC the XLA module is compiled for: launching a ``kernel[2]`` - into an ``--lnc 1`` module silently computes only core 0's half. - """ explicit = os.environ.get("NEURON_LOGICAL_NC_CONFIG", "") if explicit.strip().isdigit(): return int(explicit.strip()) @@ -63,14 +39,6 @@ def _lnc_degree() -> int: if nki is not None: @nki.jit def weight_dequant_kernel(X, S, TILE_SIZE, block_size): - """Block-wise dequantization of ``[M, N]`` ``X`` with ``[ceil(M/T), ceil(N/T)]`` scales ``S``. - - Args: - X: ``[M, N]`` quantized weights in HBM (fp16 / bf16 / fp32). - S: ``[ceil(M/T), ceil(N/T)]`` per-block scales in HBM (same dtype as X). - TILE_SIZE: block edge length T, a multiple of 128 (checked in ``run``). - block_size: columns per DMA block (compile-time constant, multiple of T or >= N). - """ M, N = X.shape out = nl.ndarray((M, N), dtype=X.dtype, buffer=nl.shared_hbm) s_cols = (N + TILE_SIZE - 1) // TILE_SIZE @@ -84,9 +52,7 @@ def weight_dequant_kernel(X, S, TILE_SIZE, block_size): for ti in range(pid * per_core, min(n_tiles, (pid + 1) * per_core)): r0 = ti * PMAX rs = min(PMAX, M - r0) - sr = r0 // TILE_SIZE # scale row of this band - # Replicate scale row ``sr`` to every partition (partition stride 0), then fp32 - # (tensor_scalar's [P, 1] operand must be fp32). + sr = r0 // TILE_SIZE s_raw = nl.ndarray((rs, s_cols), dtype=S.dtype, buffer=nl.sbuf) nisa.dma_copy(dst=s_raw, src=S.ap(pattern=[[0, rs], [1, s_cols]], offset=sr * s_stride)) s_f32 = nl.ndarray((rs, s_cols), dtype=nl.float32, buffer=nl.sbuf) @@ -98,9 +64,6 @@ def weight_dequant_kernel(X, S, TILE_SIZE, block_size): nisa.dma_copy(dst=x_tile, src=X[r0:r0 + rs, c0:c0 + cs]) y_tile = nl.ndarray((rs, cs), dtype=X.dtype, buffer=nl.sbuf) if c0 % TILE_SIZE == 0 and cs % TILE_SIZE == 0: - # Whole block in one Vector-engine instruction: view the block as - # [rs, n_sub, TILE_SIZE] and broadcast the [rs, n_sub] scales along the - # last axis (stride 0), instead of one instruction per 128-column sub-block. n_sub = cs // TILE_SIZE sc0 = c0 // TILE_SIZE nisa.tensor_tensor(dst=y_tile.ap(pattern=[[cs, rs], [TILE_SIZE, n_sub], [1, TILE_SIZE]]), @@ -108,7 +71,6 @@ def weight_dequant_kernel(X, S, TILE_SIZE, block_size): data2=s_f32.ap(pattern=[[s_cols, rs], [1, n_sub], [0, TILE_SIZE]], offset=sc0), op=nl.multiply) else: - # One scale per TILE_SIZE-wide column sub-block. for sc in range(c0 // TILE_SIZE, (c0 + cs + TILE_SIZE - 1) // TILE_SIZE): a = max(sc * TILE_SIZE, c0) - c0 b = min((sc + 1) * TILE_SIZE, c0 + cs) - c0 @@ -118,7 +80,6 @@ def weight_dequant_kernel(X, S, TILE_SIZE, block_size): return out -# Same block sizes as impl_triton.py / impl_cutile.py (BLOCK_SIZE default 1024; search 256..4096). _DEFAULT_CONFIG = SimpleNamespace(block_size=1024) _SEARCH_SPACE = [SimpleNamespace(block_size=b) for b in (256, 512, 1024, 2048, 4096)] _kernel = weight_dequant_kernel[_lnc_degree()] if nki is not None else None