From 21d22b88d87a171bde95701011f549cea0957e5f Mon Sep 17 00:00:00 2001 From: bowencui123 Date: Sat, 29 Aug 2026 03:01:06 +0000 Subject: [PATCH 1/7] nki(interleave): 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 #212 (older per-operator 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 --- benchmarks/operators/interleave/impl_nki.py | 76 +++++++++++++++++++++ 1 file changed, 76 insertions(+) create mode 100644 benchmarks/operators/interleave/impl_nki.py diff --git a/benchmarks/operators/interleave/impl_nki.py b/benchmarks/operators/interleave/impl_nki.py new file mode 100644 index 00000000..02c7f2ae --- /dev/null +++ b/benchmarks/operators/interleave/impl_nki.py @@ -0,0 +1,76 @@ +import torch + +try: + import nki + import nki.language as nl + import nki.isa as nisa + PMAX = nl.tile_size.pmax +except ImportError: + nki = None + +@nki.jit +def interleave_kernel(a_input, b_input): + """Interleave two [P, F] tensors into a [P, 2*F] tensor. + + Output layout: out[p, 2*c] = a[p, c], out[p, 2*c + 1] = b[p, c]. + Flattened row-major this yields a[0], b[0], a[1], b[1], ... which matches + ``out[0::2] = A; out[1::2] = B`` on the flattened inputs. + + The interleaving is done inside SBUF with strided destination slices + (``out_tile[:, 0::2]`` / ``out_tile[:, 1::2]``) so that both the HBM loads + and the HBM store stay fully contiguous DMAs. + """ + P, F = a_input.shape + + num_blocks = (P + PMAX - 1) // PMAX + + # Per free block SBUF footprint is (a_tile + b_tile + out_tile) = + # 4 * free_tile_size elements per partition, so keep this well under the + # per-partition SBUF budget for fp32. + free_tile_size = 4096 + num_free_blocks = (F + free_tile_size - 1) // free_tile_size + + hbm_result_tile = nl.ndarray((P, 2 * F), dtype=a_input.dtype, buffer=nl.shared_hbm) + + for i in range(num_blocks): + p_start = i * PMAX + p_end = min(p_start + PMAX, P) + p_sz = p_end - p_start + + for j in range(num_free_blocks): + f_start = j * free_tile_size + f_end = min(f_start + free_tile_size, F) + f_sz = f_end - f_start + + a_tile = nl.ndarray((p_sz, f_sz), dtype=a_input.dtype, buffer=nl.sbuf) + nisa.dma_copy(dst=a_tile, src=a_input[p_start:p_end, f_start:f_end]) + + b_tile = nl.ndarray((p_sz, f_sz), dtype=b_input.dtype, buffer=nl.sbuf) + nisa.dma_copy(dst=b_tile, src=b_input[p_start:p_end, f_start:f_end]) + + out_tile = nl.ndarray((p_sz, 2 * f_sz), dtype=a_input.dtype, buffer=nl.sbuf) + nisa.tensor_copy(dst=out_tile[0:p_sz, 0:2 * f_sz:2], src=a_tile) + nisa.tensor_copy(dst=out_tile[0:p_sz, 1:2 * f_sz:2], src=b_tile) + + nisa.dma_copy( + dst=hbm_result_tile[p_start:p_end, 2 * f_start:2 * f_end], + src=out_tile, + ) + + return hbm_result_tile + +def run(a: torch.Tensor, b: torch.Tensor, n: int, block_size: int = 1024, autotune=False, **kwargs) -> torch.Tensor: + free_dim = (n + (PMAX - 1)) // PMAX + padded_size = PMAX * free_dim + + if padded_size > n: + a = torch.nn.functional.pad(a, (0, padded_size - n)) + b = torch.nn.functional.pad(b, (0, padded_size - n)) + + a_2d = a.reshape(PMAX, free_dim) + b_2d = b.reshape(PMAX, free_dim) + result = interleave_kernel(a_2d, b_2d) + return result.reshape(-1)[:2 * n] + +def get_last_config() -> dict | None: + return None From d8f464e80d7d9e14be081264dcf290d1e3d012dc Mon Sep 17 00:00:00 2001 From: bowencui123 Date: Sat, 29 Aug 2026 08:36:57 +0000 Subject: [PATCH 2/7] nki(interleave): NkiAutotuner wiring (`block_size`) Tunables mirror the Triton search space (`BLOCK_SIZE`); defaults are the previous constants, so autotune=False is unchanged. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_012Q38kGmXvyoeM1qtCbheSL --- benchmarks/operators/interleave/impl_nki.py | 28 ++++++++++++++++++--- 1 file changed, 24 insertions(+), 4 deletions(-) diff --git a/benchmarks/operators/interleave/impl_nki.py b/benchmarks/operators/interleave/impl_nki.py index 02c7f2ae..4814db25 100644 --- a/benchmarks/operators/interleave/impl_nki.py +++ b/benchmarks/operators/interleave/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 @@ -9,7 +13,7 @@ nki = None @nki.jit -def interleave_kernel(a_input, b_input): +def interleave_kernel(a_input, b_input, block_size): """Interleave two [P, F] tensors into a [P, 2*F] tensor. Output layout: out[p, 2*c] = a[p, c], out[p, 2*c + 1] = b[p, c]. @@ -27,7 +31,7 @@ def interleave_kernel(a_input, b_input): # Per free block SBUF footprint is (a_tile + b_tile + out_tile) = # 4 * free_tile_size elements per partition, so keep this well under the # per-partition SBUF budget for fp32. - free_tile_size = 4096 + free_tile_size = block_size num_free_blocks = (F + free_tile_size - 1) // free_tile_size hbm_result_tile = nl.ndarray((P, 2 * F), dtype=a_input.dtype, buffer=nl.shared_hbm) @@ -59,6 +63,12 @@ def interleave_kernel(a_input, b_input): return hbm_result_tile +_DEFAULT_CONFIG = SimpleNamespace(block_size=4096) +_SEARCH_SPACE = [SimpleNamespace(block_size=b) for b in (1024, 2048, 4096, 8192)] +_tuner = NkiAutotuner(interleave_kernel) if nki is not None else None +_last_autotune_config: dict = {} + + def run(a: torch.Tensor, b: torch.Tensor, n: int, block_size: int = 1024, autotune=False, **kwargs) -> torch.Tensor: free_dim = (n + (PMAX - 1)) // PMAX padded_size = PMAX * free_dim @@ -69,8 +79,18 @@ def run(a: torch.Tensor, b: torch.Tensor, n: int, block_size: int = 1024, autotu a_2d = a.reshape(PMAX, free_dim) b_2d = b.reshape(PMAX, free_dim) - result = interleave_kernel(a_2d, b_2d) + if autotune: + cfg = _tuner.tune_or_cached( + shape_key=(tuple(a_2d.shape), str(a_2d.dtype)), + search_space=_SEARCH_SPACE, + args_fn=lambda cfg: (a_2d, b_2d, cfg.block_size), + ) + _last_autotune_config.clear() + _last_autotune_config.update(vars(cfg)) + else: + cfg = _DEFAULT_CONFIG + result = interleave_kernel(a_2d, b_2d, cfg.block_size) return result.reshape(-1)[:2 * n] def get_last_config() -> dict | None: - return None + return dict(_last_autotune_config) or None From 665a5a9203016fb9422a8f623f3c0017dbdfcd61 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Mon, 31 Aug 2026 17:28:34 +0000 Subject: [PATCH 3/7] nki autotune results for interleave --- results/csv/interleave_autotune.csv | 162 ++++++++++++++-------------- 1 file changed, 81 insertions(+), 81 deletions(-) diff --git a/results/csv/interleave_autotune.csv b/results/csv/interleave_autotune.csv index ddec9b70..e744f09d 100644 --- a/results/csv/interleave_autotune.csv +++ b/results/csv/interleave_autotune.csv @@ -1,81 +1,81 @@ -params,dtype,torch_ms,triton_ms,cutile_ms,speedup_triton,speedup_cutile,triton_vs_cutile,tilelang_ms,speedup_tilelang -n=1000000,fp16,0.0088,0.0026,0.0027,3.38,3.28,1.0385,0.0028,3.11 -n=1000000,bf16,0.0087,0.0025,0.0026,3.42,3.32,1.0400,0.0026,3.34 -n=1000000,fp32,0.0090,0.0035,0.0038,2.57,2.34,1.0857,0.0036,2.49 -n=1000000,int8,0.0083,0.0021,0.0022,3.94,3.77,1.0476,0.0024,3.44 -n=2000000,fp16,0.0127,0.0037,0.0038,3.46,3.31,1.0270,0.0037,3.44 -n=2000000,bf16,0.0124,0.0034,0.0037,3.64,3.36,1.0882,0.0040,3.06 -n=2000000,fp32,0.0146,0.0064,0.0070,2.29,2.07,1.0938,0.0067,2.19 -n=2000000,int8,0.0117,0.0026,0.0027,4.56,4.26,1.0385,0.0033,3.60 -n=3000000,fp16,0.0171,0.0053,0.0057,3.24,3.00,1.0755,0.0054,3.16 -n=3000000,bf16,0.0167,0.0051,0.0054,3.26,3.11,1.0588,0.0056,2.99 -n=3000000,fp32,0.0219,0.0098,0.0098,2.24,2.25,1.0000,0.0096,2.28 -n=3000000,int8,0.0143,0.0029,0.0035,4.99,4.12,1.2069,0.0035,4.08 -n=4000000,fp16,0.0217,0.0069,0.0075,3.15,2.87,1.0870,0.0072,3.03 -n=4000000,bf16,0.0216,0.0066,0.0067,3.29,3.25,1.0152,0.0070,3.09 -n=4000000,fp32,0.0299,0.0117,0.0124,2.55,2.42,1.0598,0.0122,2.46 -n=4000000,int8,0.0181,0.0035,0.0037,5.18,4.91,1.0571,0.0042,4.35 -n=5000000,fp16,0.0266,0.0081,0.0085,3.28,3.15,1.0494,0.0087,3.04 -n=5000000,bf16,0.0266,0.0080,0.0081,3.32,3.27,1.0125,0.0086,3.09 -n=5000000,fp32,0.0373,0.0145,0.0150,2.57,2.49,1.0345,0.0151,2.46 -n=5000000,int8,0.0211,0.0056,0.0041,3.74,5.11,0.7321,0.0046,4.54 -n=6000000,fp16,0.0318,0.0103,0.0096,3.09,3.31,0.9320,0.0105,3.04 -n=6000000,bf16,0.0320,0.0095,0.0106,3.39,3.04,1.1158,0.0105,3.05 -n=6000000,fp32,0.0429,0.0170,0.0172,2.52,2.49,1.0118,0.0171,2.50 -n=6000000,int8,0.0256,0.0049,0.0052,5.24,4.95,1.0612,0.0053,4.79 -n=7000000,fp16,0.0369,0.0108,0.0108,3.41,3.42,1.0000,0.0110,3.35 -n=7000000,bf16,0.0374,0.0109,0.0110,3.44,3.40,1.0092,0.0111,3.38 -n=7000000,fp32,0.0493,0.0194,0.0194,2.54,2.54,1.0000,0.0193,2.56 -n=7000000,int8,0.0303,0.0056,0.0064,5.42,4.77,1.1429,0.0061,4.95 -n=8000000,fp16,0.0421,0.0123,0.0124,3.43,3.40,1.0081,0.0132,3.20 -n=8000000,bf16,0.0425,0.0132,0.0124,3.23,3.43,0.9394,0.0127,3.35 -n=8000000,fp32,0.0557,0.0216,0.0216,2.57,2.58,1.0000,0.0219,2.55 -n=8000000,int8,0.0348,0.0064,0.0067,5.44,5.18,1.0469,0.0074,4.72 -n=9000000,fp16,0.0475,0.0134,0.0137,3.53,3.47,1.0224,0.0146,3.24 -n=9000000,bf16,0.0470,0.0133,0.0137,3.54,3.44,1.0301,0.0137,3.44 -n=9000000,fp32,0.0621,0.0241,0.0242,2.58,2.57,1.0041,0.0246,2.53 -n=9000000,int8,0.0393,0.0071,0.0075,5.50,5.22,1.0563,0.0084,4.68 -n=10000000,fp16,0.0518,0.0145,0.0146,3.56,3.55,1.0069,0.0149,3.47 -n=10000000,bf16,0.0515,0.0145,0.0145,3.55,3.55,1.0000,0.0150,3.44 -n=10000000,fp32,0.0694,0.0263,0.0265,2.65,2.62,1.0076,0.0266,2.61 -n=10000000,int8,0.0442,0.0077,0.0081,5.72,5.43,1.0519,0.0086,5.13 -n=11000000,fp16,0.0560,0.0170,0.0160,3.30,3.51,0.9412,0.0173,3.24 -n=11000000,bf16,0.0561,0.0170,0.0159,3.29,3.53,0.9353,0.0164,3.42 -n=11000000,fp32,0.0796,0.0287,0.0287,2.77,2.78,1.0000,0.0287,2.77 -n=11000000,int8,0.0493,0.0088,0.0090,5.61,5.48,1.0227,0.0095,5.19 -n=12000000,fp16,0.0607,0.0170,0.0173,3.57,3.52,1.0176,0.0170,3.56 -n=12000000,bf16,0.0608,0.0169,0.0171,3.59,3.56,1.0118,0.0171,3.55 -n=12000000,fp32,0.0882,0.0310,0.0310,2.84,2.85,1.0000,0.0314,2.81 -n=12000000,int8,0.0540,0.0091,0.0105,5.96,5.16,1.1538,0.0105,5.16 -n=13000000,fp16,0.0654,0.0182,0.0182,3.59,3.59,1.0000,0.0184,3.56 -n=13000000,bf16,0.0652,0.0183,0.0184,3.57,3.54,1.0055,0.0184,3.54 -n=13000000,fp32,0.0952,0.0330,0.0333,2.88,2.86,1.0091,0.0335,2.84 -n=13000000,int8,0.0582,0.0104,0.0104,5.62,5.58,1.0000,0.0112,5.20 -n=14000000,fp16,0.0697,0.0193,0.0194,3.62,3.60,1.0052,0.0194,3.60 -n=14000000,bf16,0.0699,0.0193,0.0195,3.62,3.59,1.0104,0.0196,3.57 -n=14000000,fp32,0.1015,0.0354,0.0357,2.87,2.84,1.0085,0.0358,2.83 -n=14000000,int8,0.0632,0.0109,0.0113,5.80,5.59,1.0367,0.0113,5.59 -n=15000000,fp16,0.0739,0.0205,0.0207,3.60,3.58,1.0098,0.0204,3.62 -n=15000000,bf16,0.0743,0.0205,0.0206,3.63,3.61,1.0049,0.0206,3.61 -n=15000000,fp32,0.1080,0.0376,0.0380,2.87,2.84,1.0106,0.0378,2.86 -n=15000000,int8,0.0675,0.0125,0.0128,5.41,5.27,1.0240,0.0125,5.42 -n=16000000,fp16,0.0787,0.0218,0.0219,3.61,3.59,1.0046,0.0218,3.61 -n=16000000,bf16,0.0789,0.0218,0.0218,3.62,3.62,1.0000,0.0236,3.34 -n=16000000,fp32,0.1146,0.0399,0.0403,2.87,2.84,1.0100,0.0399,2.87 -n=16000000,int8,0.0716,0.0123,0.0124,5.82,5.76,1.0081,0.0131,5.46 -n=17000000,fp16,0.0835,0.0228,0.0229,3.66,3.65,1.0044,0.0229,3.64 -n=17000000,bf16,0.0835,0.0230,0.0231,3.64,3.62,1.0043,0.0234,3.56 -n=17000000,fp32,0.1217,0.0421,0.0426,2.89,2.86,1.0119,0.0432,2.82 -n=17000000,int8,0.0758,0.0129,0.0130,5.88,5.83,1.0078,0.0131,5.77 -n=18000000,fp16,0.0881,0.0241,0.0242,3.66,3.64,1.0041,0.0241,3.66 -n=18000000,bf16,0.0877,0.0239,0.0240,3.68,3.65,1.0042,0.0240,3.65 -n=18000000,fp32,0.1277,0.0443,0.0448,2.88,2.85,1.0113,0.0456,2.80 -n=18000000,int8,0.0799,0.0143,0.0146,5.58,5.45,1.0210,0.0141,5.68 -n=19000000,fp16,0.0931,0.0251,0.0252,3.70,3.69,1.0040,0.0251,3.71 -n=19000000,bf16,0.0930,0.0250,0.0251,3.73,3.70,1.0040,0.0251,3.71 -n=19000000,fp32,0.1344,0.0469,0.0473,2.87,2.84,1.0085,0.0475,2.83 -n=19000000,int8,0.0842,0.0149,0.0142,5.65,5.92,0.9530,0.0141,5.96 -n=20000000,fp16,0.0985,0.0263,0.0264,3.75,3.73,1.0038,0.0263,3.74 -n=20000000,bf16,0.0983,0.0262,0.0263,3.75,3.74,1.0038,0.0265,3.71 -n=20000000,fp32,0.1404,0.0489,0.0495,2.87,2.84,1.0123,0.0499,2.81 -n=20000000,int8,0.0882,0.0146,0.0148,6.05,5.94,1.0137,0.0147,5.99 +params,dtype,torch_ms,triton_ms,cutile_ms,speedup_triton,speedup_cutile,triton_vs_cutile,torch_nki_ms,nki_ms,speedup_nki +n=1000000,fp16,0.0088,0.0026,0.0027,3.38,3.28,1.0385,12.4575,0.5737,21.72 +n=1000000,bf16,0.0087,0.0025,0.0026,3.42,3.32,1.0400,12.4024,0.5841,21.23 +n=1000000,fp32,0.0090,0.0035,0.0038,2.57,2.34,1.0857,12.4399,1.1782,10.56 +n=1000000,int8,0.0083,0.0021,0.0022,3.94,3.77,1.0476,13.2777,0.3616,36.72 +n=2000000,fp16,0.0127,0.0037,0.0038,3.46,3.31,1.0270,0.0846,0.1343,0.63 +n=2000000,bf16,0.0124,0.0034,0.0037,3.64,3.36,1.0882,0.0851,0.1442,0.59 +n=2000000,fp32,0.0146,0.0064,0.0070,2.29,2.07,1.0938,0.0964,0.2485,0.39 +n=2000000,int8,0.0117,0.0026,0.0027,4.56,4.26,1.0385,0.0787,0.0840,0.94 +n=3000000,fp16,0.0171,0.0053,0.0057,3.24,3.00,1.0755,230.0416,1.5152,151.82 +n=3000000,bf16,0.0167,0.0051,0.0054,3.26,3.11,1.0588,230.1485,1.5168,151.73 +n=3000000,fp32,0.0219,0.0098,0.0098,2.24,2.25,1.0000,230.4213,2.3901,96.41 +n=3000000,int8,0.0143,0.0029,0.0035,4.99,4.12,1.2069,230.9414,0.9354,246.88 +n=4000000,fp16,0.0217,0.0069,0.0075,3.15,2.87,1.0870,0.1381,0.2468,0.56 +n=4000000,bf16,0.0216,0.0066,0.0067,3.29,3.25,1.0152,0.1383,0.2488,0.56 +n=4000000,fp32,0.0299,0.0117,0.0124,2.55,2.42,1.0598,0.1735,0.4900,0.35 +n=4000000,int8,0.0181,0.0035,0.0037,5.18,4.91,1.0571,0.1378,0.1359,1.01 +n=5000000,fp16,0.0266,0.0081,0.0085,3.28,3.15,1.0494,383.0173,2.7416,139.70 +n=5000000,bf16,0.0266,0.0080,0.0081,3.32,3.27,1.0125,383.2611,2.7481,139.46 +n=5000000,fp32,0.0373,0.0145,0.0150,2.57,2.49,1.0345,384.0977,5.0566,75.96 +n=5000000,int8,0.0211,0.0056,0.0041,3.74,5.11,0.7321,384.0517,1.9465,197.30 +n=6000000,fp16,0.0318,0.0103,0.0096,3.09,3.31,0.9320,0.1923,0.3595,0.53 +n=6000000,bf16,0.0320,0.0095,0.0106,3.39,3.04,1.1158,0.1951,0.3699,0.53 +n=6000000,fp32,0.0429,0.0170,0.0172,2.52,2.49,1.0118,0.2238,0.7081,0.32 +n=6000000,int8,0.0256,0.0049,0.0052,5.24,4.95,1.0612,0.1735,0.1936,0.90 +n=7000000,fp16,0.0369,0.0108,0.0108,3.41,3.42,1.0000,535.6650,2.6642,201.06 +n=7000000,bf16,0.0374,0.0109,0.0110,3.44,3.40,1.0092,536.0594,2.6430,202.82 +n=7000000,fp32,0.0493,0.0194,0.0194,2.54,2.54,1.0000,537.4404,5.2243,102.87 +n=7000000,int8,0.0303,0.0056,0.0064,5.42,4.77,1.1429,537.2765,2.0574,261.14 +n=8000000,fp16,0.0421,0.0123,0.0124,3.43,3.40,1.0081,0.2498,0.4936,0.51 +n=8000000,bf16,0.0425,0.0132,0.0124,3.23,3.43,0.9394,0.2499,0.4945,0.51 +n=8000000,fp32,0.0557,0.0216,0.0216,2.57,2.58,1.0000,0.2950,0.9369,0.31 +n=8000000,int8,0.0348,0.0064,0.0067,5.44,5.18,1.0469,0.2400,0.2622,0.92 +n=9000000,fp16,0.0475,0.0134,0.0137,3.53,3.47,1.0224,689.6865,4.4340,155.54 +n=9000000,bf16,0.0470,0.0133,0.0137,3.54,3.44,1.0301,688.3414,4.4502,154.68 +n=9000000,fp32,0.0621,0.0241,0.0242,2.58,2.57,1.0041,685.5587,6.6911,102.46 +n=9000000,int8,0.0393,0.0071,0.0075,5.50,5.22,1.0563,690.5807,2.2912,301.40 +n=10000000,fp16,0.0518,0.0145,0.0146,3.56,3.55,1.0069,0.3023,0.6215,0.49 +n=10000000,bf16,0.0515,0.0145,0.0145,3.55,3.55,1.0000,0.3025,0.6243,0.48 +n=10000000,fp32,0.0694,0.0263,0.0265,2.65,2.62,1.0076,0.4012,1.2067,0.33 +n=10000000,int8,0.0442,0.0077,0.0081,5.72,5.43,1.0519,0.2866,0.3119,0.92 +n=11000000,fp16,0.0560,0.0170,0.0160,3.30,3.51,0.9412,839.4678,5.2473,159.98 +n=11000000,bf16,0.0561,0.0170,0.0159,3.29,3.53,0.9353,838.5246,5.2138,160.83 +n=11000000,fp32,0.0796,0.0287,0.0287,2.77,2.78,1.0000,873.8916,16.7334,52.22 +n=11000000,int8,0.0493,0.0088,0.0090,5.61,5.48,1.0227,838.7255,4.7972,174.84 +n=12000000,fp16,0.0607,0.0170,0.0173,3.57,3.52,1.0176,0.3574,0.7314,0.49 +n=12000000,bf16,0.0608,0.0169,0.0171,3.59,3.56,1.0118,0.3573,0.7299,0.49 +n=12000000,fp32,0.0882,0.0310,0.0310,2.84,2.85,1.0000,0.4464,1.4446,0.31 +n=12000000,int8,0.0540,0.0091,0.0105,5.96,5.16,1.1538,0.3316,0.3702,0.90 +n=13000000,fp16,0.0654,0.0182,0.0182,3.59,3.59,1.0000,993.0337,4.9868,199.13 +n=13000000,bf16,0.0652,0.0183,0.0184,3.57,3.54,1.0055,995.5750,5.0278,198.01 +n=13000000,fp32,0.0952,0.0330,0.0333,2.88,2.86,1.0091,1038.6870,9.7063,107.01 +n=13000000,int8,0.0582,0.0104,0.0104,5.62,5.58,1.0000,999.4420,2.6469,377.59 +n=14000000,fp16,0.0697,0.0193,0.0194,3.62,3.60,1.0052,0.4098,0.8425,0.49 +n=14000000,bf16,0.0699,0.0193,0.0195,3.62,3.59,1.0104,0.4090,0.8467,0.48 +n=14000000,fp32,0.1015,0.0354,0.0357,2.87,2.84,1.0085,0.5690,1.6889,0.34 +n=14000000,int8,0.0632,0.0109,0.0113,5.80,5.59,1.0367,0.3915,0.4281,0.91 +n=15000000,fp16,0.0739,0.0205,0.0207,3.60,3.58,1.0098,1143.8018,7.2508,157.75 +n=15000000,bf16,0.0743,0.0205,0.0206,3.63,3.61,1.0049,1145.1144,7.2602,157.72 +n=15000000,fp32,0.1080,0.0376,0.0380,2.87,2.84,1.0106,1191.9748,23.1621,51.46 +n=15000000,int8,0.0675,0.0125,0.0128,5.41,5.27,1.0240,1142.0914,6.6936,170.63 +n=16000000,fp16,0.0787,0.0218,0.0219,3.61,3.59,1.0046,0.4728,0.9629,0.49 +n=16000000,bf16,0.0789,0.0218,0.0218,3.62,3.62,1.0000,0.4732,0.9786,0.48 +n=16000000,fp32,0.1146,0.0399,0.0403,2.87,2.84,1.0100,0.6408,1.9248,0.33 +n=16000000,int8,0.0716,0.0123,0.0124,5.82,5.76,1.0081,0.4363,0.4891,0.89 +n=17000000,fp16,0.0835,0.0228,0.0229,3.66,3.65,1.0044,1300.5112,8.1159,160.24 +n=17000000,bf16,0.0835,0.0230,0.0231,3.64,3.62,1.0043,1300.3009,8.1421,159.70 +n=17000000,fp32,0.1217,0.0421,0.0426,2.89,2.86,1.0119,1358.8154,12.6698,107.25 +n=17000000,int8,0.0758,0.0129,0.0130,5.88,5.83,1.0078,1303.5753,4.1178,316.57 +n=18000000,fp16,0.0881,0.0241,0.0242,3.66,3.64,1.0041,0.5341,1.0891,0.49 +n=18000000,bf16,0.0877,0.0239,0.0240,3.68,3.65,1.0042,0.5387,1.0911,0.49 +n=18000000,fp32,0.1277,0.0443,0.0448,2.88,2.85,1.0113,0.6454,2.1714,0.30 +n=18000000,int8,0.0799,0.0143,0.0146,5.58,5.45,1.0210,0.5107,0.5467,0.93 +n=19000000,fp16,0.0931,0.0251,0.0252,3.70,3.69,1.0040,1455.6185,9.4625,153.83 +n=19000000,bf16,0.0930,0.0250,0.0251,3.73,3.70,1.0040,1453.9965,9.5088,152.91 +n=19000000,fp32,0.1344,0.0469,0.0473,2.87,2.84,1.0085,1520.2285,36.6460,41.48 +n=19000000,int8,0.0842,0.0149,0.0142,5.65,5.92,0.9530,1457.7301,5.4803,266.00 +n=20000000,fp16,0.0985,0.0263,0.0264,3.75,3.73,1.0038,0.5906,1.2130,0.49 +n=20000000,bf16,0.0983,0.0262,0.0263,3.75,3.74,1.0038,0.5911,1.2459,0.47 +n=20000000,fp32,0.1404,0.0489,0.0495,2.87,2.84,1.0123,0.8019,2.4118,0.33 +n=20000000,int8,0.0882,0.0146,0.0148,6.05,5.94,1.0137,0.5486,0.6083,0.90 From af2e68a8dc2181a521d7ba3e7692a35ca560b507 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Wed, 9 Sep 2026 11:55:42 +0000 Subject: [PATCH 4/7] nki(interleave): add NKI benchmark results (trn2.3xlarge, LNC2) Merges NKI backend timing into results/csv/interleave_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"). All cases pass correctness verification. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01AQseF7nyesBh8KZAp8g7Cm --- results/csv/interleave_default.csv | 162 ++++++++++++++--------------- 1 file changed, 81 insertions(+), 81 deletions(-) diff --git a/results/csv/interleave_default.csv b/results/csv/interleave_default.csv index 1f9ce2df..09f4e009 100644 --- a/results/csv/interleave_default.csv +++ b/results/csv/interleave_default.csv @@ -1,81 +1,81 @@ -params,dtype,torch_ms,triton_ms,cutile_ms,speedup_triton,speedup_cutile,triton_vs_cutile,tilelang_ms,speedup_tilelang -n=1000000,fp16,0.0088,0.0026,0.0027,3.36,3.28,1.0385,0.0026,3.34 -n=1000000,bf16,0.0087,0.0025,0.0026,3.52,3.35,1.0400,0.0026,3.38 -n=1000000,fp32,0.0089,0.0037,0.0039,2.44,2.31,1.0541,0.0039,2.28 -n=1000000,int8,0.0081,0.0021,0.0024,3.78,3.37,1.1429,0.0023,3.53 -n=2000000,fp16,0.0129,0.0038,0.0038,3.42,3.37,1.0000,0.0038,3.43 -n=2000000,bf16,0.0126,0.0036,0.0037,3.49,3.41,1.0278,0.0038,3.27 -n=2000000,fp32,0.0147,0.0064,0.0071,2.29,2.08,1.1094,0.0068,2.15 -n=2000000,int8,0.0117,0.0026,0.0029,4.42,3.99,1.1154,0.0028,4.11 -n=3000000,fp16,0.0169,0.0052,0.0053,3.26,3.20,1.0192,0.0053,3.18 -n=3000000,bf16,0.0165,0.0053,0.0054,3.10,3.07,1.0189,0.0052,3.17 -n=3000000,fp32,0.0224,0.0096,0.0107,2.33,2.10,1.1146,0.0104,2.15 -n=3000000,int8,0.0143,0.0034,0.0037,4.19,3.86,1.0882,0.0037,3.89 -n=4000000,fp16,0.0218,0.0068,0.0069,3.22,3.16,1.0147,0.0068,3.22 -n=4000000,bf16,0.0213,0.0067,0.0068,3.17,3.15,1.0149,0.0068,3.14 -n=4000000,fp32,0.0304,0.0122,0.0135,2.48,2.25,1.1066,0.0128,2.37 -n=4000000,int8,0.0175,0.0039,0.0042,4.50,4.14,1.0769,0.0039,4.51 -n=5000000,fp16,0.0267,0.0084,0.0085,3.18,3.15,1.0119,0.0086,3.11 -n=5000000,bf16,0.0268,0.0084,0.0085,3.20,3.16,1.0119,0.0086,3.12 -n=5000000,fp32,0.0373,0.0146,0.0164,2.56,2.28,1.1233,0.0149,2.50 -n=5000000,int8,0.0211,0.0046,0.0050,4.54,4.24,1.0870,0.0046,4.56 -n=6000000,fp16,0.0321,0.0098,0.0101,3.26,3.18,1.0306,0.0098,3.27 -n=6000000,bf16,0.0320,0.0098,0.0099,3.25,3.23,1.0102,0.0099,3.22 -n=6000000,fp32,0.0438,0.0170,0.0191,2.58,2.29,1.1235,0.0174,2.52 -n=6000000,int8,0.0256,0.0056,0.0058,4.59,4.39,1.0357,0.0057,4.50 -n=7000000,fp16,0.0373,0.0111,0.0111,3.37,3.35,1.0000,0.0114,3.28 -n=7000000,bf16,0.0368,0.0110,0.0110,3.34,3.33,1.0000,0.0111,3.33 -n=7000000,fp32,0.0498,0.0192,0.0219,2.60,2.27,1.1406,0.0199,2.51 -n=7000000,int8,0.0297,0.0064,0.0067,4.68,4.47,1.0469,0.0065,4.54 -n=8000000,fp16,0.0421,0.0123,0.0124,3.43,3.40,1.0081,0.0122,3.47 -n=8000000,bf16,0.0421,0.0123,0.0124,3.43,3.41,1.0081,0.0121,3.47 -n=8000000,fp32,0.0560,0.0217,0.0248,2.59,2.26,1.1429,0.0221,2.54 -n=8000000,int8,0.0353,0.0074,0.0076,4.79,4.61,1.0270,0.0072,4.88 -n=9000000,fp16,0.0471,0.0136,0.0136,3.47,3.45,1.0000,0.0136,3.47 -n=9000000,bf16,0.0473,0.0138,0.0139,3.43,3.41,1.0072,0.0137,3.45 -n=9000000,fp32,0.0618,0.0240,0.0275,2.57,2.25,1.1458,0.0245,2.52 -n=9000000,int8,0.0395,0.0079,0.0082,4.99,4.82,1.0380,0.0082,4.79 -n=10000000,fp16,0.0518,0.0145,0.0145,3.57,3.57,1.0000,0.0149,3.48 -n=10000000,bf16,0.0515,0.0148,0.0148,3.49,3.48,1.0000,0.0147,3.50 -n=10000000,fp32,0.0692,0.0264,0.0302,2.63,2.30,1.1439,0.0267,2.59 -n=10000000,int8,0.0446,0.0088,0.0091,5.06,4.91,1.0341,0.0092,4.87 -n=11000000,fp16,0.0561,0.0158,0.0159,3.54,3.54,1.0063,0.0159,3.52 -n=11000000,bf16,0.0558,0.0159,0.0160,3.51,3.48,1.0063,0.0156,3.57 -n=11000000,fp32,0.0785,0.0288,0.0329,2.73,2.39,1.1424,0.0289,2.72 -n=11000000,int8,0.0492,0.0098,0.0100,5.03,4.90,1.0204,0.0100,4.92 -n=12000000,fp16,0.0607,0.0170,0.0170,3.58,3.56,1.0000,0.0173,3.50 -n=12000000,bf16,0.0607,0.0170,0.0171,3.57,3.55,1.0059,0.0171,3.55 -n=12000000,fp32,0.0873,0.0310,0.0356,2.81,2.46,1.1484,0.0312,2.80 -n=12000000,int8,0.0537,0.0104,0.0107,5.19,5.03,1.0288,0.0104,5.16 -n=13000000,fp16,0.0656,0.0181,0.0182,3.62,3.60,1.0055,0.0184,3.56 -n=13000000,bf16,0.0652,0.0181,0.0182,3.60,3.59,1.0055,0.0182,3.57 -n=13000000,fp32,0.0948,0.0334,0.0382,2.84,2.48,1.1437,0.0338,2.81 -n=13000000,int8,0.0584,0.0109,0.0112,5.33,5.20,1.0275,0.0112,5.19 -n=14000000,fp16,0.0698,0.0194,0.0195,3.60,3.59,1.0052,0.0193,3.61 -n=14000000,bf16,0.0697,0.0193,0.0194,3.60,3.59,1.0052,0.0194,3.60 -n=14000000,fp32,0.1019,0.0357,0.0406,2.85,2.51,1.1373,0.0361,2.82 -n=14000000,int8,0.0631,0.0114,0.0118,5.51,5.34,1.0351,0.0119,5.28 -n=15000000,fp16,0.0742,0.0206,0.0207,3.61,3.59,1.0049,0.0205,3.61 -n=15000000,bf16,0.0741,0.0207,0.0207,3.59,3.57,1.0000,0.0206,3.59 -n=15000000,fp32,0.1077,0.0380,0.0434,2.83,2.48,1.1421,0.0383,2.81 -n=15000000,int8,0.0671,0.0126,0.0130,5.33,5.18,1.0317,0.0126,5.32 -n=16000000,fp16,0.0786,0.0217,0.0217,3.62,3.62,1.0000,0.0216,3.63 -n=16000000,bf16,0.0789,0.0216,0.0216,3.66,3.65,1.0000,0.0218,3.62 -n=16000000,fp32,0.1148,0.0404,0.0460,2.85,2.49,1.1386,0.0408,2.81 -n=16000000,int8,0.0717,0.0132,0.0136,5.43,5.29,1.0303,0.0132,5.42 -n=17000000,fp16,0.0833,0.0227,0.0227,3.67,3.67,1.0000,0.0228,3.65 -n=17000000,bf16,0.0833,0.0229,0.0229,3.64,3.63,1.0000,0.0229,3.63 -n=17000000,fp32,0.1211,0.0426,0.0486,2.85,2.49,1.1408,0.0430,2.82 -n=17000000,int8,0.0759,0.0134,0.0139,5.65,5.47,1.0373,0.0140,5.41 -n=18000000,fp16,0.0882,0.0241,0.0241,3.67,3.66,1.0000,0.0242,3.64 -n=18000000,bf16,0.0884,0.0241,0.0241,3.67,3.67,1.0000,0.0243,3.63 -n=18000000,fp32,0.1282,0.0450,0.0513,2.85,2.50,1.1400,0.0452,2.83 -n=18000000,int8,0.0800,0.0145,0.0148,5.51,5.39,1.0207,0.0145,5.52 -n=19000000,fp16,0.0933,0.0252,0.0253,3.70,3.69,1.0040,0.0254,3.68 -n=19000000,bf16,0.0932,0.0251,0.0252,3.70,3.70,1.0040,0.0256,3.64 -n=19000000,fp32,0.1347,0.0474,0.0538,2.84,2.50,1.1350,0.0478,2.82 -n=19000000,int8,0.0842,0.0151,0.0156,5.57,5.41,1.0331,0.0152,5.56 -n=20000000,fp16,0.0984,0.0264,0.0264,3.73,3.73,1.0000,0.0268,3.67 -n=20000000,bf16,0.0983,0.0265,0.0265,3.72,3.71,1.0000,0.0263,3.74 -n=20000000,fp32,0.1404,0.0496,0.0564,2.83,2.49,1.1371,0.0502,2.80 -n=20000000,int8,0.0883,0.0155,0.0160,5.71,5.53,1.0323,0.0156,5.66 +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 +n=1000000,fp16,0.0088,0.0026,0.0027,3.36,3.28,1.0385,0.0026,3.34,10.3637,0.3514,29.49 +n=1000000,bf16,0.0087,0.0025,0.0026,3.52,3.35,1.0400,0.0026,3.38,10.3799,0.3507,29.60 +n=1000000,fp32,0.0089,0.0037,0.0039,2.44,2.31,1.0541,0.0039,2.28,10.6759,0.6075,17.57 +n=1000000,int8,0.0081,0.0021,0.0024,3.78,3.37,1.1429,0.0023,3.53,9.8533,0.2229,44.21 +n=2000000,fp16,0.0129,0.0038,0.0038,3.42,3.37,1.0000,0.0038,3.43,0.0779,0.1316,0.59 +n=2000000,bf16,0.0126,0.0036,0.0037,3.49,3.41,1.0278,0.0038,3.27,0.0779,0.1317,0.59 +n=2000000,fp32,0.0147,0.0064,0.0071,2.29,2.08,1.1094,0.0068,2.15,0.0930,0.2478,0.38 +n=2000000,int8,0.0117,0.0026,0.0029,4.42,3.99,1.1154,0.0028,4.11,0.0725,0.0793,0.91 +n=3000000,fp16,0.0169,0.0052,0.0053,3.26,3.20,1.0192,0.0053,3.18,226.0327,0.9852,229.44 +n=3000000,bf16,0.0165,0.0053,0.0054,3.10,3.07,1.0189,0.0052,3.17,224.6891,0.9842,228.29 +n=3000000,fp32,0.0224,0.0096,0.0107,2.33,2.10,1.1146,0.0104,2.15,225.5857,1.8522,121.80 +n=3000000,int8,0.0143,0.0034,0.0037,4.19,3.86,1.0882,0.0037,3.89,225.5373,0.7045,320.15 +n=4000000,fp16,0.0218,0.0068,0.0069,3.22,3.16,1.0147,0.0068,3.22,0.1330,0.2477,0.54 +n=4000000,bf16,0.0213,0.0067,0.0068,3.17,3.15,1.0149,0.0068,3.14,0.1329,0.2478,0.54 +n=4000000,fp32,0.0304,0.0122,0.0135,2.48,2.25,1.1066,0.0128,2.37,0.1757,0.4726,0.37 +n=4000000,int8,0.0175,0.0039,0.0042,4.50,4.14,1.0769,0.0039,4.51,0.1304,0.1320,0.99 +n=5000000,fp16,0.0267,0.0084,0.0085,3.18,3.15,1.0119,0.0086,3.11,374.5145,1.8147,206.38 +n=5000000,bf16,0.0268,0.0084,0.0085,3.20,3.16,1.0119,0.0086,3.12,375.0674,1.8149,206.66 +n=5000000,fp32,0.0373,0.0146,0.0164,2.56,2.28,1.1233,0.0149,2.50,374.8177,3.3103,113.23 +n=5000000,int8,0.0211,0.0046,0.0050,4.54,4.24,1.0870,0.0046,4.56,376.0774,1.1847,317.46 +n=6000000,fp16,0.0321,0.0098,0.0101,3.26,3.18,1.0306,0.0098,3.27,0.1854,0.3628,0.51 +n=6000000,bf16,0.0320,0.0098,0.0099,3.25,3.23,1.0102,0.0099,3.22,0.1854,0.3629,0.51 +n=6000000,fp32,0.0438,0.0170,0.0191,2.58,2.29,1.1235,0.0174,2.52,0.2374,0.6939,0.34 +n=6000000,int8,0.0256,0.0056,0.0058,4.59,4.39,1.0357,0.0057,4.50,0.1674,0.1900,0.88 +n=7000000,fp16,0.0373,0.0111,0.0111,3.37,3.35,1.0000,0.0114,3.28,526.4570,2.1328,246.83 +n=7000000,bf16,0.0368,0.0110,0.0110,3.34,3.33,1.0000,0.0111,3.33,525.5231,2.1325,246.43 +n=7000000,fp32,0.0498,0.0192,0.0219,2.60,2.27,1.1406,0.0199,2.51,524.8790,4.2200,124.38 +n=7000000,int8,0.0297,0.0064,0.0067,4.68,4.47,1.0469,0.0065,4.54,528.1355,1.5258,346.13 +n=8000000,fp16,0.0421,0.0123,0.0124,3.43,3.40,1.0081,0.0122,3.47,0.2435,0.4745,0.51 +n=8000000,bf16,0.0421,0.0123,0.0124,3.43,3.41,1.0081,0.0121,3.47,0.2434,0.4746,0.51 +n=8000000,fp32,0.0560,0.0217,0.0248,2.59,2.26,1.1429,0.0221,2.54,0.3156,0.9214,0.34 +n=8000000,int8,0.0353,0.0074,0.0076,4.79,4.61,1.0270,0.0072,4.88,0.2302,0.2441,0.94 +n=9000000,fp16,0.0471,0.0136,0.0136,3.47,3.45,1.0000,0.0136,3.47,676.6584,3.4669,195.18 +n=9000000,bf16,0.0473,0.0138,0.0139,3.43,3.41,1.0072,0.0137,3.45,674.9231,3.4668,194.68 +n=9000000,fp32,0.0618,0.0240,0.0275,2.57,2.25,1.1458,0.0245,2.52,677.3427,5.4605,124.04 +n=9000000,int8,0.0395,0.0079,0.0082,4.99,4.82,1.0380,0.0082,4.79,676.2451,1.7709,381.87 +n=10000000,fp16,0.0518,0.0145,0.0145,3.57,3.57,1.0000,0.0149,3.48,0.2985,0.5901,0.51 +n=10000000,bf16,0.0515,0.0148,0.0148,3.49,3.48,1.0000,0.0147,3.50,0.2986,0.5904,0.51 +n=10000000,fp32,0.0692,0.0264,0.0302,2.63,2.30,1.1439,0.0267,2.59,0.4045,1.1497,0.35 +n=10000000,int8,0.0446,0.0088,0.0091,5.06,4.91,1.0341,0.0092,4.87,0.2809,0.3010,0.93 +n=11000000,fp16,0.0561,0.0158,0.0159,3.54,3.54,1.0063,0.0159,3.52,824.6380,4.2182,195.50 +n=11000000,bf16,0.0558,0.0159,0.0160,3.51,3.48,1.0063,0.0156,3.57,829.2845,4.2163,196.69 +n=11000000,fp32,0.0785,0.0288,0.0329,2.73,2.39,1.1424,0.0289,2.72,856.4410,8.0895,105.87 +n=11000000,int8,0.0492,0.0098,0.0100,5.03,4.90,1.0204,0.0100,4.92,831.5321,2.9111,285.64 +n=12000000,fp16,0.0607,0.0170,0.0170,3.58,3.56,1.0000,0.0173,3.50,0.3544,0.6961,0.51 +n=12000000,bf16,0.0607,0.0170,0.0171,3.57,3.55,1.0059,0.0171,3.55,0.3544,0.6964,0.51 +n=12000000,fp32,0.0873,0.0310,0.0356,2.81,2.46,1.1484,0.0312,2.80,0.4842,1.3741,0.35 +n=12000000,int8,0.0537,0.0104,0.0107,5.19,5.03,1.0288,0.0104,5.16,0.3242,0.3624,0.89 +n=13000000,fp16,0.0656,0.0181,0.0182,3.62,3.60,1.0055,0.0184,3.56,983.7040,4.1385,237.69 +n=13000000,bf16,0.0652,0.0181,0.0182,3.60,3.59,1.0055,0.0182,3.57,980.5048,4.1386,236.92 +n=13000000,fp32,0.0948,0.0334,0.0382,2.84,2.48,1.1437,0.0338,2.81,1017.2364,7.9952,127.23 +n=13000000,int8,0.0584,0.0109,0.0112,5.33,5.20,1.0275,0.0112,5.19,984.2373,2.1860,450.24 +n=14000000,fp16,0.0698,0.0194,0.0195,3.60,3.59,1.0052,0.0193,3.61,0.4036,0.8092,0.50 +n=14000000,bf16,0.0697,0.0193,0.0194,3.60,3.59,1.0052,0.0194,3.60,0.4036,0.8092,0.50 +n=14000000,fp32,0.1019,0.0357,0.0406,2.85,2.51,1.1373,0.0361,2.82,0.5683,1.6018,0.35 +n=14000000,int8,0.0631,0.0114,0.0118,5.51,5.34,1.0351,0.0119,5.28,0.3812,0.4212,0.90 +n=15000000,fp16,0.0742,0.0206,0.0207,3.61,3.59,1.0049,0.0205,3.61,1125.8076,5.8701,191.79 +n=15000000,bf16,0.0741,0.0207,0.0207,3.59,3.57,1.0000,0.0206,3.59,1126.5679,5.8702,191.91 +n=15000000,fp32,0.1077,0.0380,0.0434,2.83,2.48,1.1421,0.0383,2.81,1170.8232,11.0465,105.99 +n=15000000,int8,0.0671,0.0126,0.0130,5.33,5.18,1.0317,0.0126,5.32,1126.4175,3.8837,290.03 +n=16000000,fp16,0.0786,0.0217,0.0217,3.62,3.62,1.0000,0.0216,3.63,0.4691,0.9212,0.51 +n=16000000,bf16,0.0789,0.0216,0.0216,3.66,3.65,1.0000,0.0218,3.62,0.4692,0.9216,0.51 +n=16000000,fp32,0.1148,0.0404,0.0460,2.85,2.49,1.1386,0.0408,2.81,0.6559,1.8170,0.36 +n=16000000,int8,0.0717,0.0132,0.0136,5.43,5.29,1.0303,0.0132,5.42,0.4268,0.4751,0.90 +n=17000000,fp16,0.0833,0.0227,0.0227,3.67,3.67,1.0000,0.0228,3.65,1281.0027,6.4801,197.68 +n=17000000,bf16,0.0833,0.0229,0.0229,3.64,3.63,1.0000,0.0229,3.63,1279.1394,6.4795,197.41 +n=17000000,fp32,0.1211,0.0426,0.0486,2.85,2.49,1.1408,0.0430,2.82,1326.3975,10.2278,129.69 +n=17000000,int8,0.0759,0.0134,0.0139,5.65,5.47,1.0373,0.0140,5.41,1279.3219,3.3126,386.19 +n=18000000,fp16,0.0882,0.0241,0.0241,3.67,3.66,1.0000,0.0242,3.64,0.5254,1.0355,0.51 +n=18000000,bf16,0.0884,0.0241,0.0241,3.67,3.67,1.0000,0.0243,3.63,0.5254,1.0355,0.51 +n=18000000,fp32,0.1282,0.0450,0.0513,2.85,2.50,1.1400,0.0452,2.83,0.6864,2.0515,0.33 +n=18000000,int8,0.0800,0.0145,0.0148,5.51,5.39,1.0207,0.0145,5.52,0.4957,0.5362,0.92 +n=19000000,fp16,0.0933,0.0252,0.0253,3.70,3.69,1.0040,0.0254,3.68,1425.8746,7.3739,193.37 +n=19000000,bf16,0.0932,0.0251,0.0252,3.70,3.70,1.0040,0.0256,3.64,1425.6330,7.3765,193.27 +n=19000000,fp32,0.1347,0.0474,0.0538,2.84,2.50,1.1350,0.0478,2.82,1489.0733,11.6871,127.41 +n=19000000,int8,0.0842,0.0151,0.0156,5.57,5.41,1.0331,0.0152,5.56,1432.0092,4.2696,335.40 +n=20000000,fp16,0.0984,0.0264,0.0264,3.73,3.73,1.0000,0.0268,3.67,0.5852,1.1501,0.51 +n=20000000,bf16,0.0983,0.0265,0.0265,3.72,3.71,1.0000,0.0263,3.74,0.5852,1.1495,0.51 +n=20000000,fp32,0.1404,0.0496,0.0564,2.83,2.49,1.1371,0.0502,2.80,0.8085,2.2774,0.36 +n=20000000,int8,0.0883,0.0155,0.0160,5.71,5.53,1.0323,0.0156,5.66,0.5365,0.5946,0.90 From 157f21493ab335c76b226ba46b9d2a9d96a51f58 Mon Sep 17 00:00:00 2001 From: Bowen Cui Date: Thu, 17 Sep 2026 00:05:13 +0000 Subject: [PATCH 5/7] nki(interleave): LNC2 kernel with on-chip strided interleave and in-kernel tails; XLA torch baseline = strided assignment on a 128-padded 2-D view (odd lengths lower to scatters); rerun benchmarks Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01ScXYNjrrKGgDUVNHxv7HJt --- benchmarks/operators/interleave/impl_nki.py | 172 ++++++++----- benchmarks/operators/interleave/impl_torch.py | 35 ++- results/csv/interleave_autotune.csv | 160 ++++++------ results/csv/interleave_default.csv | 160 ++++++------ .../autotune_logs/interleave_autotune.json | 242 +++++++++++++++++- 5 files changed, 539 insertions(+), 230 deletions(-) diff --git a/benchmarks/operators/interleave/impl_nki.py b/benchmarks/operators/interleave/impl_nki.py index 4814db25..c396f40b 100644 --- a/benchmarks/operators/interleave/impl_nki.py +++ b/benchmarks/operators/interleave/impl_nki.py @@ -1,3 +1,30 @@ +"""NKI (AWS Trainium) implementation of interleave: ``out[0::2] = A; out[1::2] = B``. + +Design (mirrors the Triton / cuTile kernels): + +* The flat inputs are processed in fixed-size blocks. A Triton program handles + ``BLOCK_SIZE`` contiguous elements of ``A`` and ``B`` and writes ``2 * BLOCK_SIZE`` + outputs; the NKI analogue is one ``[128, BLOCK_SIZE]`` SBUF tile per input -- + 128 partitions (the 128 parallel lanes of the Vector engine) times + ``BLOCK_SIZE`` contiguous elements per partition. The default and the autotune + search space use the same ``BLOCK_SIZE`` values as ``impl_triton.py`` / + ``impl_cutile.py`` (default 1024; search 1024/2048/4096/8192). +* The interleaving itself is done on-chip with strided SBUF destination slices + (``out_tile[:, 0::2] = a``, ``out_tile[:, 1::2] = b``), so every HBM load and + the HBM store stay fully contiguous DMAs. +* Work is split across the NeuronCores of the logical core (LNC2 on trn2): + each program instance (``nl.program_id(0)``) owns a contiguous half of the + element range, so both physical cores' DMA engines are used. +* Tails are handled inside the kernel: the last block of a core may cover fewer + than ``128 * BLOCK_SIZE`` elements, in which case it is processed as one + ``[128, q]`` tile plus one ``[1, r]`` tile (``r < 128``). No host-side + padding / reshape / slicing is needed -- those torch ops would be fused into + the same NEFF as the kernel and dominate its runtime. +""" +import functools +import os +import re +import subprocess from types import SimpleNamespace import torch @@ -6,91 +33,106 @@ 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 + PMAX = 128 -@nki.jit -def interleave_kernel(a_input, b_input, block_size): - """Interleave two [P, F] tensors into a [P, 2*F] tensor. - Output layout: out[p, 2*c] = a[p, c], out[p, 2*c + 1] = b[p, c]. - Flattened row-major this yields a[0], b[0], a[1], b[1], ... which matches - ``out[0::2] = A; out[1::2] = B`` on the flattened inputs. +@functools.lru_cache(maxsize=1) +def _lnc_degree() -> int: + """Logical-NeuronCore degree the kernel is launched with (``kernel[lnc]``). - The interleaving is done inside SBUF with strided destination slices - (``out_tile[:, 0::2]`` / ``out_tile[:, 1::2]``) so that both the HBM loads - and the HBM store stay fully contiguous DMAs. + 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. """ - P, F = a_input.shape - - num_blocks = (P + PMAX - 1) // PMAX - - # Per free block SBUF footprint is (a_tile + b_tile + out_tile) = - # 4 * free_tile_size elements per partition, so keep this well under the - # per-partition SBUF budget for fp32. - free_tile_size = block_size - num_free_blocks = (F + free_tile_size - 1) // free_tile_size - - hbm_result_tile = nl.ndarray((P, 2 * F), dtype=a_input.dtype, buffer=nl.shared_hbm) - - for i in range(num_blocks): - p_start = i * PMAX - p_end = min(p_start + PMAX, P) - p_sz = p_end - p_start - - for j in range(num_free_blocks): - f_start = j * free_tile_size - f_end = min(f_start + free_tile_size, F) - f_sz = f_end - f_start - - a_tile = nl.ndarray((p_sz, f_sz), dtype=a_input.dtype, buffer=nl.sbuf) - nisa.dma_copy(dst=a_tile, src=a_input[p_start:p_end, f_start:f_end]) - - b_tile = nl.ndarray((p_sz, f_sz), dtype=b_input.dtype, buffer=nl.sbuf) - nisa.dma_copy(dst=b_tile, src=b_input[p_start:p_end, f_start:f_end]) - - out_tile = nl.ndarray((p_sz, 2 * f_sz), dtype=a_input.dtype, buffer=nl.sbuf) - nisa.tensor_copy(dst=out_tile[0:p_sz, 0:2 * f_sz:2], src=a_tile) - nisa.tensor_copy(dst=out_tile[0:p_sz, 1:2 * f_sz:2], src=b_tile) - - nisa.dma_copy( - dst=hbm_result_tile[p_start:p_end, 2 * f_start:2 * f_end], - src=out_tile, - ) - - return hbm_result_tile - -_DEFAULT_CONFIG = SimpleNamespace(block_size=4096) + 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 interleave_kernel(a_input, b_input, block_size): + """``out[2i] = a[i]; out[2i+1] = b[i]`` over flat ``(n,)`` HBM tensors. + + Args: + a_input, b_input: flat ``(n,)`` tensors in HBM (same dtype). + block_size: elements per partition per tile (compile-time constant). + """ + n = a_input.shape[0] + out = nl.ndarray((2 * n,), dtype=a_input.dtype, buffer=nl.shared_hbm) + + # Contiguous element range owned by this program instance (NeuronCore). + num_programs = nl.num_programs() + per_core = (n + num_programs - 1) // num_programs + lo = nl.program_id(0) * per_core + hi = min(n, lo + per_core) + + chunk = PMAX * block_size + for j in range(max(0, (hi - lo + chunk - 1) // chunk)): + start = lo + j * chunk + count = min(chunk, hi - start) + q = count // PMAX # full partitions: [PMAX, q] tiles + r = count - q * PMAX # remaining < PMAX elements: [1, r] tiles + if q > 0: + _tile_body(PMAX, q, start, a_input, b_input, out) + if r > 0: + _tile_body(1, r, start + q * PMAX, a_input, b_input, out) + return out + + def _tile_body(p, f, start, a_input, b_input, out): + """Interleave the ``[p, f]`` tiles of a and b at flat input offset ``start``.""" + a_tile = nl.ndarray((p, f), dtype=a_input.dtype, buffer=nl.sbuf) + nisa.dma_copy(dst=a_tile, src=a_input.ap(pattern=[[f, p], [1, f]], offset=start)) + b_tile = nl.ndarray((p, f), dtype=b_input.dtype, buffer=nl.sbuf) + nisa.dma_copy(dst=b_tile, src=b_input.ap(pattern=[[f, p], [1, f]], offset=start)) + # On-chip interleave into the even / odd free-dim slots of one [p, 2f] tile. + out_tile = nl.ndarray((p, 2 * f), dtype=a_input.dtype, buffer=nl.sbuf) + nisa.tensor_copy(dst=out_tile[0:p, 0:2 * f:2], src=a_tile) + nisa.tensor_copy(dst=out_tile[0:p, 1:2 * f:2], src=b_tile) + # Partition p of the tile holds inputs [start + p*f, start + (p+1)*f), whose + # outputs are the contiguous range [2*(start + p*f), 2*(start + (p+1)*f)). + nisa.dma_copy(dst=out.ap(pattern=[[2 * f, p], [1, 2 * f]], offset=2 * start), src=out_tile) + + +# Same block sizes as impl_triton.py / impl_cutile.py (default 1024; search 1024/2048/4096/8192). +_DEFAULT_CONFIG = SimpleNamespace(block_size=1024) _SEARCH_SPACE = [SimpleNamespace(block_size=b) for b in (1024, 2048, 4096, 8192)] -_tuner = NkiAutotuner(interleave_kernel) if nki is not None else None +_kernel = interleave_kernel[_lnc_degree()] if nki is not None else None +_tuner = NkiAutotuner(_kernel) if nki is not None else None _last_autotune_config: dict = {} -def run(a: torch.Tensor, b: torch.Tensor, n: int, block_size: int = 1024, autotune=False, **kwargs) -> torch.Tensor: - free_dim = (n + (PMAX - 1)) // PMAX - padded_size = PMAX * free_dim - - if padded_size > n: - a = torch.nn.functional.pad(a, (0, padded_size - n)) - b = torch.nn.functional.pad(b, (0, padded_size - n)) - - a_2d = a.reshape(PMAX, free_dim) - b_2d = b.reshape(PMAX, free_dim) +def run(A: torch.Tensor, B: torch.Tensor, N: int, block_size: int = 1024, + autotune: bool = False, **kwargs) -> torch.Tensor: + a_flat, b_flat = A.reshape(-1), B.reshape(-1) + n = a_flat.numel() if autotune: cfg = _tuner.tune_or_cached( - shape_key=(tuple(a_2d.shape), str(a_2d.dtype)), + shape_key=(n, str(a_flat.dtype)), search_space=_SEARCH_SPACE, - args_fn=lambda cfg: (a_2d, b_2d, cfg.block_size), + args_fn=lambda cfg: (a_flat, b_flat, cfg.block_size), ) _last_autotune_config.clear() _last_autotune_config.update(vars(cfg)) else: cfg = _DEFAULT_CONFIG - result = interleave_kernel(a_2d, b_2d, cfg.block_size) - return result.reshape(-1)[:2 * n] + return _kernel(a_flat, b_flat, cfg.block_size) + def get_last_config() -> dict | None: return dict(_last_autotune_config) or None diff --git a/benchmarks/operators/interleave/impl_torch.py b/benchmarks/operators/interleave/impl_torch.py index 572bd07a..afe7f46b 100644 --- a/benchmarks/operators/interleave/impl_torch.py +++ b/benchmarks/operators/interleave/impl_torch.py @@ -1,8 +1,35 @@ import torch +import torch.nn.functional as F -def run(A: torch.Tensor, B: torch.Tensor, N: int, **kwargs): - output = torch.empty(2 * N, dtype=A.dtype, device=A.device) - output[0::2] = A - output[1::2] = B +def _view_cols(n: int, lo: int = 128, hi: int = 8192) -> int | None: + """Largest divisor of ``n`` in ``[lo, hi]`` (row length of the 2-D view), or None.""" + for c in range(hi, lo - 1, -1): + if n % c == 0: + return c + return None + + +def _interleave(a: torch.Tensor, b: torch.Tensor, n: int) -> torch.Tensor: + output = torch.empty(2 * n, dtype=a.dtype, device=a.device) + output[0::2] = a + output[1::2] = b return output + + +def run(A: torch.Tensor, B: torch.Tensor, N: int, **kwargs): + if A.device.type == "xla" and N % 128 != 0: + # Neuron/XLA: the strided assignments compile to one streaming pass when the + # length is a multiple of 128 (2M elements: 0.08 ms) but to element-wise + # scatters otherwise (1M: 10 ms, 3M: 226 ms). Pad the row length of a 2-D view + # to a multiple of 128 (strided DMAs), interleave the padded flat tensors and + # drop the padding again. + cols = _view_cols(N) + if cols is not None: + rows = N // cols + pad = (-cols) % 128 + a2 = F.pad(A.reshape(rows, cols), (0, pad)).reshape(-1) + b2 = F.pad(B.reshape(rows, cols), (0, pad)).reshape(-1) + out = _interleave(a2, b2, a2.numel()).reshape(rows, 2 * (cols + pad)) + return out[:, :2 * cols].reshape(-1) + return _interleave(A, B, N) diff --git a/results/csv/interleave_autotune.csv b/results/csv/interleave_autotune.csv index e744f09d..2c4d8e2b 100644 --- a/results/csv/interleave_autotune.csv +++ b/results/csv/interleave_autotune.csv @@ -1,81 +1,81 @@ params,dtype,torch_ms,triton_ms,cutile_ms,speedup_triton,speedup_cutile,triton_vs_cutile,torch_nki_ms,nki_ms,speedup_nki -n=1000000,fp16,0.0088,0.0026,0.0027,3.38,3.28,1.0385,12.4575,0.5737,21.72 -n=1000000,bf16,0.0087,0.0025,0.0026,3.42,3.32,1.0400,12.4024,0.5841,21.23 -n=1000000,fp32,0.0090,0.0035,0.0038,2.57,2.34,1.0857,12.4399,1.1782,10.56 -n=1000000,int8,0.0083,0.0021,0.0022,3.94,3.77,1.0476,13.2777,0.3616,36.72 -n=2000000,fp16,0.0127,0.0037,0.0038,3.46,3.31,1.0270,0.0846,0.1343,0.63 -n=2000000,bf16,0.0124,0.0034,0.0037,3.64,3.36,1.0882,0.0851,0.1442,0.59 -n=2000000,fp32,0.0146,0.0064,0.0070,2.29,2.07,1.0938,0.0964,0.2485,0.39 -n=2000000,int8,0.0117,0.0026,0.0027,4.56,4.26,1.0385,0.0787,0.0840,0.94 -n=3000000,fp16,0.0171,0.0053,0.0057,3.24,3.00,1.0755,230.0416,1.5152,151.82 -n=3000000,bf16,0.0167,0.0051,0.0054,3.26,3.11,1.0588,230.1485,1.5168,151.73 -n=3000000,fp32,0.0219,0.0098,0.0098,2.24,2.25,1.0000,230.4213,2.3901,96.41 -n=3000000,int8,0.0143,0.0029,0.0035,4.99,4.12,1.2069,230.9414,0.9354,246.88 -n=4000000,fp16,0.0217,0.0069,0.0075,3.15,2.87,1.0870,0.1381,0.2468,0.56 -n=4000000,bf16,0.0216,0.0066,0.0067,3.29,3.25,1.0152,0.1383,0.2488,0.56 -n=4000000,fp32,0.0299,0.0117,0.0124,2.55,2.42,1.0598,0.1735,0.4900,0.35 -n=4000000,int8,0.0181,0.0035,0.0037,5.18,4.91,1.0571,0.1378,0.1359,1.01 -n=5000000,fp16,0.0266,0.0081,0.0085,3.28,3.15,1.0494,383.0173,2.7416,139.70 -n=5000000,bf16,0.0266,0.0080,0.0081,3.32,3.27,1.0125,383.2611,2.7481,139.46 -n=5000000,fp32,0.0373,0.0145,0.0150,2.57,2.49,1.0345,384.0977,5.0566,75.96 -n=5000000,int8,0.0211,0.0056,0.0041,3.74,5.11,0.7321,384.0517,1.9465,197.30 -n=6000000,fp16,0.0318,0.0103,0.0096,3.09,3.31,0.9320,0.1923,0.3595,0.53 -n=6000000,bf16,0.0320,0.0095,0.0106,3.39,3.04,1.1158,0.1951,0.3699,0.53 -n=6000000,fp32,0.0429,0.0170,0.0172,2.52,2.49,1.0118,0.2238,0.7081,0.32 -n=6000000,int8,0.0256,0.0049,0.0052,5.24,4.95,1.0612,0.1735,0.1936,0.90 -n=7000000,fp16,0.0369,0.0108,0.0108,3.41,3.42,1.0000,535.6650,2.6642,201.06 -n=7000000,bf16,0.0374,0.0109,0.0110,3.44,3.40,1.0092,536.0594,2.6430,202.82 -n=7000000,fp32,0.0493,0.0194,0.0194,2.54,2.54,1.0000,537.4404,5.2243,102.87 -n=7000000,int8,0.0303,0.0056,0.0064,5.42,4.77,1.1429,537.2765,2.0574,261.14 -n=8000000,fp16,0.0421,0.0123,0.0124,3.43,3.40,1.0081,0.2498,0.4936,0.51 -n=8000000,bf16,0.0425,0.0132,0.0124,3.23,3.43,0.9394,0.2499,0.4945,0.51 -n=8000000,fp32,0.0557,0.0216,0.0216,2.57,2.58,1.0000,0.2950,0.9369,0.31 -n=8000000,int8,0.0348,0.0064,0.0067,5.44,5.18,1.0469,0.2400,0.2622,0.92 -n=9000000,fp16,0.0475,0.0134,0.0137,3.53,3.47,1.0224,689.6865,4.4340,155.54 -n=9000000,bf16,0.0470,0.0133,0.0137,3.54,3.44,1.0301,688.3414,4.4502,154.68 -n=9000000,fp32,0.0621,0.0241,0.0242,2.58,2.57,1.0041,685.5587,6.6911,102.46 -n=9000000,int8,0.0393,0.0071,0.0075,5.50,5.22,1.0563,690.5807,2.2912,301.40 -n=10000000,fp16,0.0518,0.0145,0.0146,3.56,3.55,1.0069,0.3023,0.6215,0.49 -n=10000000,bf16,0.0515,0.0145,0.0145,3.55,3.55,1.0000,0.3025,0.6243,0.48 -n=10000000,fp32,0.0694,0.0263,0.0265,2.65,2.62,1.0076,0.4012,1.2067,0.33 -n=10000000,int8,0.0442,0.0077,0.0081,5.72,5.43,1.0519,0.2866,0.3119,0.92 -n=11000000,fp16,0.0560,0.0170,0.0160,3.30,3.51,0.9412,839.4678,5.2473,159.98 -n=11000000,bf16,0.0561,0.0170,0.0159,3.29,3.53,0.9353,838.5246,5.2138,160.83 -n=11000000,fp32,0.0796,0.0287,0.0287,2.77,2.78,1.0000,873.8916,16.7334,52.22 -n=11000000,int8,0.0493,0.0088,0.0090,5.61,5.48,1.0227,838.7255,4.7972,174.84 -n=12000000,fp16,0.0607,0.0170,0.0173,3.57,3.52,1.0176,0.3574,0.7314,0.49 -n=12000000,bf16,0.0608,0.0169,0.0171,3.59,3.56,1.0118,0.3573,0.7299,0.49 -n=12000000,fp32,0.0882,0.0310,0.0310,2.84,2.85,1.0000,0.4464,1.4446,0.31 -n=12000000,int8,0.0540,0.0091,0.0105,5.96,5.16,1.1538,0.3316,0.3702,0.90 -n=13000000,fp16,0.0654,0.0182,0.0182,3.59,3.59,1.0000,993.0337,4.9868,199.13 -n=13000000,bf16,0.0652,0.0183,0.0184,3.57,3.54,1.0055,995.5750,5.0278,198.01 -n=13000000,fp32,0.0952,0.0330,0.0333,2.88,2.86,1.0091,1038.6870,9.7063,107.01 -n=13000000,int8,0.0582,0.0104,0.0104,5.62,5.58,1.0000,999.4420,2.6469,377.59 -n=14000000,fp16,0.0697,0.0193,0.0194,3.62,3.60,1.0052,0.4098,0.8425,0.49 -n=14000000,bf16,0.0699,0.0193,0.0195,3.62,3.59,1.0104,0.4090,0.8467,0.48 -n=14000000,fp32,0.1015,0.0354,0.0357,2.87,2.84,1.0085,0.5690,1.6889,0.34 -n=14000000,int8,0.0632,0.0109,0.0113,5.80,5.59,1.0367,0.3915,0.4281,0.91 -n=15000000,fp16,0.0739,0.0205,0.0207,3.60,3.58,1.0098,1143.8018,7.2508,157.75 -n=15000000,bf16,0.0743,0.0205,0.0206,3.63,3.61,1.0049,1145.1144,7.2602,157.72 -n=15000000,fp32,0.1080,0.0376,0.0380,2.87,2.84,1.0106,1191.9748,23.1621,51.46 -n=15000000,int8,0.0675,0.0125,0.0128,5.41,5.27,1.0240,1142.0914,6.6936,170.63 -n=16000000,fp16,0.0787,0.0218,0.0219,3.61,3.59,1.0046,0.4728,0.9629,0.49 -n=16000000,bf16,0.0789,0.0218,0.0218,3.62,3.62,1.0000,0.4732,0.9786,0.48 -n=16000000,fp32,0.1146,0.0399,0.0403,2.87,2.84,1.0100,0.6408,1.9248,0.33 -n=16000000,int8,0.0716,0.0123,0.0124,5.82,5.76,1.0081,0.4363,0.4891,0.89 -n=17000000,fp16,0.0835,0.0228,0.0229,3.66,3.65,1.0044,1300.5112,8.1159,160.24 -n=17000000,bf16,0.0835,0.0230,0.0231,3.64,3.62,1.0043,1300.3009,8.1421,159.70 -n=17000000,fp32,0.1217,0.0421,0.0426,2.89,2.86,1.0119,1358.8154,12.6698,107.25 -n=17000000,int8,0.0758,0.0129,0.0130,5.88,5.83,1.0078,1303.5753,4.1178,316.57 -n=18000000,fp16,0.0881,0.0241,0.0242,3.66,3.64,1.0041,0.5341,1.0891,0.49 -n=18000000,bf16,0.0877,0.0239,0.0240,3.68,3.65,1.0042,0.5387,1.0911,0.49 -n=18000000,fp32,0.1277,0.0443,0.0448,2.88,2.85,1.0113,0.6454,2.1714,0.30 -n=18000000,int8,0.0799,0.0143,0.0146,5.58,5.45,1.0210,0.5107,0.5467,0.93 -n=19000000,fp16,0.0931,0.0251,0.0252,3.70,3.69,1.0040,1455.6185,9.4625,153.83 -n=19000000,bf16,0.0930,0.0250,0.0251,3.73,3.70,1.0040,1453.9965,9.5088,152.91 -n=19000000,fp32,0.1344,0.0469,0.0473,2.87,2.84,1.0085,1520.2285,36.6460,41.48 -n=19000000,int8,0.0842,0.0149,0.0142,5.65,5.92,0.9530,1457.7301,5.4803,266.00 -n=20000000,fp16,0.0985,0.0263,0.0264,3.75,3.73,1.0038,0.5906,1.2130,0.49 -n=20000000,bf16,0.0983,0.0262,0.0263,3.75,3.74,1.0038,0.5911,1.2459,0.47 -n=20000000,fp32,0.1404,0.0489,0.0495,2.87,2.84,1.0123,0.8019,2.4118,0.33 -n=20000000,int8,0.0882,0.0146,0.0148,6.05,5.94,1.0137,0.5486,0.6083,0.90 +n=1000000,fp16,0.0088,0.0026,0.0027,3.38,3.28,1.0385,0.0755,0.0277,2.73 +n=1000000,bf16,0.0087,0.0025,0.0026,3.42,3.32,1.0400,0.0754,0.0286,2.64 +n=1000000,fp32,0.0090,0.0035,0.0038,2.57,2.34,1.0857,0.1205,0.0382,3.16 +n=1000000,int8,0.0083,0.0021,0.0022,3.94,3.77,1.0476,0.0525,0.0236,2.22 +n=2000000,fp16,0.0127,0.0037,0.0038,3.46,3.31,1.0270,0.0777,0.0388,2.00 +n=2000000,bf16,0.0124,0.0034,0.0037,3.64,3.36,1.0882,0.0776,0.0389,1.99 +n=2000000,fp32,0.0146,0.0064,0.0070,2.29,2.07,1.0938,0.0922,0.0606,1.52 +n=2000000,int8,0.0117,0.0026,0.0027,4.56,4.26,1.0385,0.0723,0.0290,2.49 +n=3000000,fp16,0.0171,0.0053,0.0057,3.24,3.00,1.0755,0.1936,0.0503,3.85 +n=3000000,bf16,0.0167,0.0051,0.0054,3.26,3.11,1.0588,0.1940,0.0506,3.83 +n=3000000,fp32,0.0219,0.0098,0.0098,2.24,2.25,1.0000,0.3324,0.0844,3.94 +n=3000000,int8,0.0143,0.0029,0.0035,4.99,4.12,1.2069,0.1588,0.0343,4.64 +n=4000000,fp16,0.0217,0.0069,0.0075,3.15,2.87,1.0870,0.1326,0.0605,2.19 +n=4000000,bf16,0.0216,0.0066,0.0067,3.29,3.25,1.0152,0.1327,0.0610,2.17 +n=4000000,fp32,0.0299,0.0117,0.0124,2.55,2.42,1.0598,0.1752,0.1063,1.65 +n=4000000,int8,0.0181,0.0035,0.0037,5.18,4.91,1.0571,0.1303,0.0373,3.49 +n=5000000,fp16,0.0266,0.0081,0.0085,3.28,3.15,1.0494,0.3018,0.0736,4.10 +n=5000000,bf16,0.0266,0.0080,0.0081,3.32,3.27,1.0125,0.3020,0.0766,3.94 +n=5000000,fp32,0.0373,0.0145,0.0150,2.57,2.49,1.0345,0.5474,0.1340,4.08 +n=5000000,int8,0.0211,0.0056,0.0041,3.74,5.11,0.7321,0.2655,0.0461,5.76 +n=6000000,fp16,0.0318,0.0103,0.0096,3.09,3.31,0.9320,0.1850,0.0847,2.18 +n=6000000,bf16,0.0320,0.0095,0.0106,3.39,3.04,1.1158,0.1849,0.0846,2.18 +n=6000000,fp32,0.0429,0.0170,0.0172,2.52,2.49,1.0118,0.2369,0.1581,1.50 +n=6000000,int8,0.0256,0.0049,0.0052,5.24,4.95,1.0612,0.1675,0.0506,3.31 +n=7000000,fp16,0.0369,0.0108,0.0108,3.41,3.42,1.0000,0.4157,0.0965,4.31 +n=7000000,bf16,0.0374,0.0109,0.0110,3.44,3.40,1.0092,0.4161,0.1019,4.08 +n=7000000,fp32,0.0493,0.0194,0.0194,2.54,2.54,1.0000,0.7504,0.1852,4.05 +n=7000000,int8,0.0303,0.0056,0.0064,5.42,4.77,1.1429,0.3461,0.0569,6.08 +n=8000000,fp16,0.0421,0.0123,0.0124,3.43,3.40,1.0081,0.2431,0.1099,2.21 +n=8000000,bf16,0.0425,0.0132,0.0124,3.23,3.43,0.9394,0.2430,0.1127,2.16 +n=8000000,fp32,0.0557,0.0216,0.0216,2.57,2.58,1.0000,0.3154,0.2075,1.52 +n=8000000,int8,0.0348,0.0064,0.0067,5.44,5.18,1.0469,0.2300,0.0605,3.80 +n=9000000,fp16,0.0475,0.0134,0.0137,3.53,3.47,1.0224,0.5377,0.1221,4.40 +n=9000000,bf16,0.0470,0.0133,0.0137,3.54,3.44,1.0301,0.5376,0.1310,4.10 +n=9000000,fp32,0.0621,0.0241,0.0242,2.58,2.57,1.0041,0.9769,0.2307,4.24 +n=9000000,int8,0.0393,0.0071,0.0075,5.50,5.22,1.0563,0.4223,0.0677,6.24 +n=10000000,fp16,0.0518,0.0145,0.0146,3.56,3.55,1.0069,0.2982,0.1343,2.22 +n=10000000,bf16,0.0515,0.0145,0.0145,3.55,3.55,1.0000,0.2983,0.1346,2.22 +n=10000000,fp32,0.0694,0.0263,0.0265,2.65,2.62,1.0076,0.4032,0.2610,1.54 +n=10000000,int8,0.0442,0.0077,0.0081,5.72,5.43,1.0519,0.2807,0.0735,3.82 +n=11000000,fp16,0.0560,0.0170,0.0160,3.30,3.51,0.9412,0.6431,0.1464,4.39 +n=11000000,bf16,0.0561,0.0170,0.0159,3.29,3.53,0.9353,0.6434,0.1464,4.39 +n=11000000,fp32,0.0796,0.0287,0.0287,2.77,2.78,1.0000,1.1971,0.2781,4.30 +n=11000000,int8,0.0493,0.0088,0.0090,5.61,5.48,1.0227,0.5630,0.0791,7.11 +n=12000000,fp16,0.0607,0.0170,0.0173,3.57,3.52,1.0176,0.3536,0.1580,2.24 +n=12000000,bf16,0.0608,0.0169,0.0171,3.59,3.56,1.0118,0.3536,0.1619,2.18 +n=12000000,fp32,0.0882,0.0310,0.0310,2.84,2.85,1.0000,0.4807,0.3038,1.58 +n=12000000,int8,0.0540,0.0091,0.0105,5.96,5.16,1.1538,0.3242,0.0834,3.89 +n=13000000,fp16,0.0654,0.0182,0.0182,3.59,3.59,1.0000,0.7870,0.1713,4.60 +n=13000000,bf16,0.0652,0.0183,0.0184,3.57,3.54,1.0055,0.7880,0.1793,4.39 +n=13000000,fp32,0.0952,0.0330,0.0333,2.88,2.86,1.0091,1.4637,0.3262,4.49 +n=13000000,int8,0.0582,0.0104,0.0104,5.62,5.58,1.0000,0.6391,0.0903,7.08 +n=14000000,fp16,0.0697,0.0193,0.0194,3.62,3.60,1.0052,0.4033,0.1851,2.18 +n=14000000,bf16,0.0699,0.0193,0.0195,3.62,3.59,1.0104,0.4034,0.1873,2.15 +n=14000000,fp32,0.1015,0.0354,0.0357,2.87,2.84,1.0085,0.5679,0.3532,1.61 +n=14000000,int8,0.0632,0.0109,0.0113,5.80,5.59,1.0367,0.3812,0.0966,3.95 +n=15000000,fp16,0.0739,0.0205,0.0207,3.60,3.58,1.0098,0.8741,0.1984,4.41 +n=15000000,bf16,0.0743,0.0205,0.0206,3.63,3.61,1.0049,0.8740,0.2005,4.36 +n=15000000,fp32,0.1080,0.0376,0.0380,2.87,2.84,1.0106,1.6289,0.3748,4.35 +n=15000000,int8,0.0675,0.0125,0.0128,5.41,5.27,1.0240,0.7064,0.1074,6.58 +n=16000000,fp16,0.0787,0.0218,0.0219,3.61,3.59,1.0046,0.4686,0.2074,2.26 +n=16000000,bf16,0.0789,0.0218,0.0218,3.62,3.62,1.0000,0.4686,0.2094,2.24 +n=16000000,fp32,0.1146,0.0399,0.0403,2.87,2.84,1.0100,0.6551,0.3915,1.67 +n=16000000,int8,0.0716,0.0123,0.0124,5.82,5.76,1.0081,0.4268,0.1126,3.79 +n=17000000,fp16,0.0835,0.0228,0.0229,3.66,3.65,1.0044,0.9881,0.2232,4.43 +n=17000000,bf16,0.0835,0.0230,0.0231,3.64,3.62,1.0043,0.9878,0.2261,4.37 +n=17000000,fp32,0.1217,0.0421,0.0426,2.89,2.86,1.0119,1.8419,0.4226,4.36 +n=17000000,int8,0.0758,0.0129,0.0130,5.88,5.83,1.0078,0.8645,0.1212,7.13 +n=18000000,fp16,0.0881,0.0241,0.0242,3.66,3.64,1.0041,0.5239,0.2348,2.23 +n=18000000,bf16,0.0877,0.0239,0.0240,3.68,3.65,1.0042,0.5239,0.2357,2.22 +n=18000000,fp32,0.1277,0.0443,0.0448,2.88,2.85,1.0113,0.6857,0.4461,1.54 +n=18000000,int8,0.0799,0.0143,0.0146,5.58,5.45,1.0210,0.4956,0.1267,3.91 +n=19000000,fp16,0.0931,0.0251,0.0252,3.70,3.69,1.0040,1.0995,0.2429,4.53 +n=19000000,bf16,0.0930,0.0250,0.0251,3.73,3.70,1.0040,1.0996,0.2462,4.47 +n=19000000,fp32,0.1344,0.0469,0.0473,2.87,2.84,1.0085,2.0164,0.4697,4.29 +n=19000000,int8,0.0842,0.0149,0.0142,5.65,5.92,0.9530,0.9226,0.1300,7.10 +n=20000000,fp16,0.0985,0.0263,0.0264,3.75,3.73,1.0038,0.5843,0.2567,2.28 +n=20000000,bf16,0.0983,0.0262,0.0263,3.75,3.74,1.0038,0.5843,0.2608,2.24 +n=20000000,fp32,0.1404,0.0489,0.0495,2.87,2.84,1.0123,0.8081,0.4942,1.64 +n=20000000,int8,0.0882,0.0146,0.0148,6.05,5.94,1.0137,0.5361,0.1315,4.08 diff --git a/results/csv/interleave_default.csv b/results/csv/interleave_default.csv index 09f4e009..dd01fcf9 100644 --- a/results/csv/interleave_default.csv +++ b/results/csv/interleave_default.csv @@ -1,81 +1,81 @@ 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 -n=1000000,fp16,0.0088,0.0026,0.0027,3.36,3.28,1.0385,0.0026,3.34,10.3637,0.3514,29.49 -n=1000000,bf16,0.0087,0.0025,0.0026,3.52,3.35,1.0400,0.0026,3.38,10.3799,0.3507,29.60 -n=1000000,fp32,0.0089,0.0037,0.0039,2.44,2.31,1.0541,0.0039,2.28,10.6759,0.6075,17.57 -n=1000000,int8,0.0081,0.0021,0.0024,3.78,3.37,1.1429,0.0023,3.53,9.8533,0.2229,44.21 -n=2000000,fp16,0.0129,0.0038,0.0038,3.42,3.37,1.0000,0.0038,3.43,0.0779,0.1316,0.59 -n=2000000,bf16,0.0126,0.0036,0.0037,3.49,3.41,1.0278,0.0038,3.27,0.0779,0.1317,0.59 -n=2000000,fp32,0.0147,0.0064,0.0071,2.29,2.08,1.1094,0.0068,2.15,0.0930,0.2478,0.38 -n=2000000,int8,0.0117,0.0026,0.0029,4.42,3.99,1.1154,0.0028,4.11,0.0725,0.0793,0.91 -n=3000000,fp16,0.0169,0.0052,0.0053,3.26,3.20,1.0192,0.0053,3.18,226.0327,0.9852,229.44 -n=3000000,bf16,0.0165,0.0053,0.0054,3.10,3.07,1.0189,0.0052,3.17,224.6891,0.9842,228.29 -n=3000000,fp32,0.0224,0.0096,0.0107,2.33,2.10,1.1146,0.0104,2.15,225.5857,1.8522,121.80 -n=3000000,int8,0.0143,0.0034,0.0037,4.19,3.86,1.0882,0.0037,3.89,225.5373,0.7045,320.15 -n=4000000,fp16,0.0218,0.0068,0.0069,3.22,3.16,1.0147,0.0068,3.22,0.1330,0.2477,0.54 -n=4000000,bf16,0.0213,0.0067,0.0068,3.17,3.15,1.0149,0.0068,3.14,0.1329,0.2478,0.54 -n=4000000,fp32,0.0304,0.0122,0.0135,2.48,2.25,1.1066,0.0128,2.37,0.1757,0.4726,0.37 -n=4000000,int8,0.0175,0.0039,0.0042,4.50,4.14,1.0769,0.0039,4.51,0.1304,0.1320,0.99 -n=5000000,fp16,0.0267,0.0084,0.0085,3.18,3.15,1.0119,0.0086,3.11,374.5145,1.8147,206.38 -n=5000000,bf16,0.0268,0.0084,0.0085,3.20,3.16,1.0119,0.0086,3.12,375.0674,1.8149,206.66 -n=5000000,fp32,0.0373,0.0146,0.0164,2.56,2.28,1.1233,0.0149,2.50,374.8177,3.3103,113.23 -n=5000000,int8,0.0211,0.0046,0.0050,4.54,4.24,1.0870,0.0046,4.56,376.0774,1.1847,317.46 -n=6000000,fp16,0.0321,0.0098,0.0101,3.26,3.18,1.0306,0.0098,3.27,0.1854,0.3628,0.51 -n=6000000,bf16,0.0320,0.0098,0.0099,3.25,3.23,1.0102,0.0099,3.22,0.1854,0.3629,0.51 -n=6000000,fp32,0.0438,0.0170,0.0191,2.58,2.29,1.1235,0.0174,2.52,0.2374,0.6939,0.34 -n=6000000,int8,0.0256,0.0056,0.0058,4.59,4.39,1.0357,0.0057,4.50,0.1674,0.1900,0.88 -n=7000000,fp16,0.0373,0.0111,0.0111,3.37,3.35,1.0000,0.0114,3.28,526.4570,2.1328,246.83 -n=7000000,bf16,0.0368,0.0110,0.0110,3.34,3.33,1.0000,0.0111,3.33,525.5231,2.1325,246.43 -n=7000000,fp32,0.0498,0.0192,0.0219,2.60,2.27,1.1406,0.0199,2.51,524.8790,4.2200,124.38 -n=7000000,int8,0.0297,0.0064,0.0067,4.68,4.47,1.0469,0.0065,4.54,528.1355,1.5258,346.13 -n=8000000,fp16,0.0421,0.0123,0.0124,3.43,3.40,1.0081,0.0122,3.47,0.2435,0.4745,0.51 -n=8000000,bf16,0.0421,0.0123,0.0124,3.43,3.41,1.0081,0.0121,3.47,0.2434,0.4746,0.51 -n=8000000,fp32,0.0560,0.0217,0.0248,2.59,2.26,1.1429,0.0221,2.54,0.3156,0.9214,0.34 -n=8000000,int8,0.0353,0.0074,0.0076,4.79,4.61,1.0270,0.0072,4.88,0.2302,0.2441,0.94 -n=9000000,fp16,0.0471,0.0136,0.0136,3.47,3.45,1.0000,0.0136,3.47,676.6584,3.4669,195.18 -n=9000000,bf16,0.0473,0.0138,0.0139,3.43,3.41,1.0072,0.0137,3.45,674.9231,3.4668,194.68 -n=9000000,fp32,0.0618,0.0240,0.0275,2.57,2.25,1.1458,0.0245,2.52,677.3427,5.4605,124.04 -n=9000000,int8,0.0395,0.0079,0.0082,4.99,4.82,1.0380,0.0082,4.79,676.2451,1.7709,381.87 -n=10000000,fp16,0.0518,0.0145,0.0145,3.57,3.57,1.0000,0.0149,3.48,0.2985,0.5901,0.51 -n=10000000,bf16,0.0515,0.0148,0.0148,3.49,3.48,1.0000,0.0147,3.50,0.2986,0.5904,0.51 -n=10000000,fp32,0.0692,0.0264,0.0302,2.63,2.30,1.1439,0.0267,2.59,0.4045,1.1497,0.35 -n=10000000,int8,0.0446,0.0088,0.0091,5.06,4.91,1.0341,0.0092,4.87,0.2809,0.3010,0.93 -n=11000000,fp16,0.0561,0.0158,0.0159,3.54,3.54,1.0063,0.0159,3.52,824.6380,4.2182,195.50 -n=11000000,bf16,0.0558,0.0159,0.0160,3.51,3.48,1.0063,0.0156,3.57,829.2845,4.2163,196.69 -n=11000000,fp32,0.0785,0.0288,0.0329,2.73,2.39,1.1424,0.0289,2.72,856.4410,8.0895,105.87 -n=11000000,int8,0.0492,0.0098,0.0100,5.03,4.90,1.0204,0.0100,4.92,831.5321,2.9111,285.64 -n=12000000,fp16,0.0607,0.0170,0.0170,3.58,3.56,1.0000,0.0173,3.50,0.3544,0.6961,0.51 -n=12000000,bf16,0.0607,0.0170,0.0171,3.57,3.55,1.0059,0.0171,3.55,0.3544,0.6964,0.51 -n=12000000,fp32,0.0873,0.0310,0.0356,2.81,2.46,1.1484,0.0312,2.80,0.4842,1.3741,0.35 -n=12000000,int8,0.0537,0.0104,0.0107,5.19,5.03,1.0288,0.0104,5.16,0.3242,0.3624,0.89 -n=13000000,fp16,0.0656,0.0181,0.0182,3.62,3.60,1.0055,0.0184,3.56,983.7040,4.1385,237.69 -n=13000000,bf16,0.0652,0.0181,0.0182,3.60,3.59,1.0055,0.0182,3.57,980.5048,4.1386,236.92 -n=13000000,fp32,0.0948,0.0334,0.0382,2.84,2.48,1.1437,0.0338,2.81,1017.2364,7.9952,127.23 -n=13000000,int8,0.0584,0.0109,0.0112,5.33,5.20,1.0275,0.0112,5.19,984.2373,2.1860,450.24 -n=14000000,fp16,0.0698,0.0194,0.0195,3.60,3.59,1.0052,0.0193,3.61,0.4036,0.8092,0.50 -n=14000000,bf16,0.0697,0.0193,0.0194,3.60,3.59,1.0052,0.0194,3.60,0.4036,0.8092,0.50 -n=14000000,fp32,0.1019,0.0357,0.0406,2.85,2.51,1.1373,0.0361,2.82,0.5683,1.6018,0.35 -n=14000000,int8,0.0631,0.0114,0.0118,5.51,5.34,1.0351,0.0119,5.28,0.3812,0.4212,0.90 -n=15000000,fp16,0.0742,0.0206,0.0207,3.61,3.59,1.0049,0.0205,3.61,1125.8076,5.8701,191.79 -n=15000000,bf16,0.0741,0.0207,0.0207,3.59,3.57,1.0000,0.0206,3.59,1126.5679,5.8702,191.91 -n=15000000,fp32,0.1077,0.0380,0.0434,2.83,2.48,1.1421,0.0383,2.81,1170.8232,11.0465,105.99 -n=15000000,int8,0.0671,0.0126,0.0130,5.33,5.18,1.0317,0.0126,5.32,1126.4175,3.8837,290.03 -n=16000000,fp16,0.0786,0.0217,0.0217,3.62,3.62,1.0000,0.0216,3.63,0.4691,0.9212,0.51 -n=16000000,bf16,0.0789,0.0216,0.0216,3.66,3.65,1.0000,0.0218,3.62,0.4692,0.9216,0.51 -n=16000000,fp32,0.1148,0.0404,0.0460,2.85,2.49,1.1386,0.0408,2.81,0.6559,1.8170,0.36 -n=16000000,int8,0.0717,0.0132,0.0136,5.43,5.29,1.0303,0.0132,5.42,0.4268,0.4751,0.90 -n=17000000,fp16,0.0833,0.0227,0.0227,3.67,3.67,1.0000,0.0228,3.65,1281.0027,6.4801,197.68 -n=17000000,bf16,0.0833,0.0229,0.0229,3.64,3.63,1.0000,0.0229,3.63,1279.1394,6.4795,197.41 -n=17000000,fp32,0.1211,0.0426,0.0486,2.85,2.49,1.1408,0.0430,2.82,1326.3975,10.2278,129.69 -n=17000000,int8,0.0759,0.0134,0.0139,5.65,5.47,1.0373,0.0140,5.41,1279.3219,3.3126,386.19 -n=18000000,fp16,0.0882,0.0241,0.0241,3.67,3.66,1.0000,0.0242,3.64,0.5254,1.0355,0.51 -n=18000000,bf16,0.0884,0.0241,0.0241,3.67,3.67,1.0000,0.0243,3.63,0.5254,1.0355,0.51 -n=18000000,fp32,0.1282,0.0450,0.0513,2.85,2.50,1.1400,0.0452,2.83,0.6864,2.0515,0.33 -n=18000000,int8,0.0800,0.0145,0.0148,5.51,5.39,1.0207,0.0145,5.52,0.4957,0.5362,0.92 -n=19000000,fp16,0.0933,0.0252,0.0253,3.70,3.69,1.0040,0.0254,3.68,1425.8746,7.3739,193.37 -n=19000000,bf16,0.0932,0.0251,0.0252,3.70,3.70,1.0040,0.0256,3.64,1425.6330,7.3765,193.27 -n=19000000,fp32,0.1347,0.0474,0.0538,2.84,2.50,1.1350,0.0478,2.82,1489.0733,11.6871,127.41 -n=19000000,int8,0.0842,0.0151,0.0156,5.57,5.41,1.0331,0.0152,5.56,1432.0092,4.2696,335.40 -n=20000000,fp16,0.0984,0.0264,0.0264,3.73,3.73,1.0000,0.0268,3.67,0.5852,1.1501,0.51 -n=20000000,bf16,0.0983,0.0265,0.0265,3.72,3.71,1.0000,0.0263,3.74,0.5852,1.1495,0.51 -n=20000000,fp32,0.1404,0.0496,0.0564,2.83,2.49,1.1371,0.0502,2.80,0.8085,2.2774,0.36 -n=20000000,int8,0.0883,0.0155,0.0160,5.71,5.53,1.0323,0.0156,5.66,0.5365,0.5946,0.90 +n=1000000,fp16,0.0088,0.0026,0.0027,3.36,3.28,1.0385,0.0026,3.34,0.0754,0.0277,2.72 +n=1000000,bf16,0.0087,0.0025,0.0026,3.52,3.35,1.0400,0.0026,3.38,0.0754,0.0277,2.72 +n=1000000,fp32,0.0089,0.0037,0.0039,2.44,2.31,1.0541,0.0039,2.28,0.1201,0.0390,3.08 +n=1000000,int8,0.0081,0.0021,0.0024,3.78,3.37,1.1429,0.0023,3.53,0.0524,0.0236,2.22 +n=2000000,fp16,0.0129,0.0038,0.0038,3.42,3.37,1.0000,0.0038,3.43,0.0776,0.0388,2.00 +n=2000000,bf16,0.0126,0.0036,0.0037,3.49,3.41,1.0278,0.0038,3.27,0.0776,0.0388,2.00 +n=2000000,fp32,0.0147,0.0064,0.0071,2.29,2.08,1.1094,0.0068,2.15,0.0923,0.0612,1.51 +n=2000000,int8,0.0117,0.0026,0.0029,4.42,3.99,1.1154,0.0028,4.11,0.0723,0.0322,2.25 +n=3000000,fp16,0.0169,0.0052,0.0053,3.26,3.20,1.0192,0.0053,3.18,0.1937,0.0503,3.85 +n=3000000,bf16,0.0165,0.0053,0.0054,3.10,3.07,1.0189,0.0052,3.17,0.1937,0.0503,3.85 +n=3000000,fp32,0.0224,0.0096,0.0107,2.33,2.10,1.1146,0.0104,2.15,0.3323,0.0846,3.93 +n=3000000,int8,0.0143,0.0034,0.0037,4.19,3.86,1.0882,0.0037,3.89,0.1587,0.0403,3.94 +n=4000000,fp16,0.0218,0.0068,0.0069,3.22,3.16,1.0147,0.0068,3.22,0.1327,0.0631,2.10 +n=4000000,bf16,0.0213,0.0067,0.0068,3.17,3.15,1.0149,0.0068,3.14,0.1326,0.0631,2.10 +n=4000000,fp32,0.0304,0.0122,0.0135,2.48,2.25,1.1066,0.0128,2.37,0.1753,0.1126,1.56 +n=4000000,int8,0.0175,0.0039,0.0042,4.50,4.14,1.0769,0.0039,4.51,0.1303,0.0475,2.74 +n=5000000,fp16,0.0267,0.0084,0.0085,3.18,3.15,1.0119,0.0086,3.11,0.3021,0.0765,3.95 +n=5000000,bf16,0.0268,0.0084,0.0085,3.20,3.16,1.0119,0.0086,3.12,0.3021,0.0765,3.95 +n=5000000,fp32,0.0373,0.0146,0.0164,2.56,2.28,1.1233,0.0149,2.50,0.5474,0.1348,4.06 +n=5000000,int8,0.0211,0.0046,0.0050,4.54,4.24,1.0870,0.0046,4.56,0.2654,0.0576,4.60 +n=6000000,fp16,0.0321,0.0098,0.0101,3.26,3.18,1.0306,0.0098,3.27,0.1850,0.0891,2.08 +n=6000000,bf16,0.0320,0.0098,0.0099,3.25,3.23,1.0102,0.0099,3.22,0.1852,0.0890,2.08 +n=6000000,fp32,0.0438,0.0170,0.0191,2.58,2.29,1.1235,0.0174,2.52,0.2369,0.1609,1.47 +n=6000000,int8,0.0256,0.0056,0.0058,4.59,4.39,1.0357,0.0057,4.50,0.1673,0.0652,2.56 +n=7000000,fp16,0.0373,0.0111,0.0111,3.37,3.35,1.0000,0.0114,3.28,0.4157,0.1018,4.08 +n=7000000,bf16,0.0368,0.0110,0.0110,3.34,3.33,1.0000,0.0111,3.33,0.4159,0.1020,4.08 +n=7000000,fp32,0.0498,0.0192,0.0219,2.60,2.27,1.1406,0.0199,2.51,0.7508,0.1866,4.02 +n=7000000,int8,0.0297,0.0064,0.0067,4.68,4.47,1.0469,0.0065,4.54,0.3461,0.0719,4.82 +n=8000000,fp16,0.0421,0.0123,0.0124,3.43,3.40,1.0081,0.0122,3.47,0.2430,0.1123,2.16 +n=8000000,bf16,0.0421,0.0123,0.0124,3.43,3.41,1.0081,0.0121,3.47,0.2430,0.1120,2.17 +n=8000000,fp32,0.0560,0.0217,0.0248,2.59,2.26,1.1429,0.0221,2.54,0.3154,0.2130,1.48 +n=8000000,int8,0.0353,0.0074,0.0076,4.79,4.61,1.0270,0.0072,4.88,0.2301,0.0792,2.91 +n=9000000,fp16,0.0471,0.0136,0.0136,3.47,3.45,1.0000,0.0136,3.47,0.5375,0.1293,4.16 +n=9000000,bf16,0.0473,0.0138,0.0139,3.43,3.41,1.0072,0.0137,3.45,0.5378,0.1293,4.16 +n=9000000,fp32,0.0618,0.0240,0.0275,2.57,2.25,1.1458,0.0245,2.52,0.9773,0.2361,4.14 +n=9000000,int8,0.0395,0.0079,0.0082,4.99,4.82,1.0380,0.0082,4.79,0.4224,0.0889,4.75 +n=10000000,fp16,0.0518,0.0145,0.0145,3.57,3.57,1.0000,0.0149,3.48,0.2982,0.1388,2.15 +n=10000000,bf16,0.0515,0.0148,0.0148,3.49,3.48,1.0000,0.0147,3.50,0.2982,0.1389,2.15 +n=10000000,fp32,0.0692,0.0264,0.0302,2.63,2.30,1.1439,0.0267,2.59,0.4032,0.2622,1.54 +n=10000000,int8,0.0446,0.0088,0.0091,5.06,4.91,1.0341,0.0092,4.87,0.2807,0.0972,2.89 +n=11000000,fp16,0.0561,0.0158,0.0159,3.54,3.54,1.0063,0.0159,3.52,0.6432,0.1545,4.16 +n=11000000,bf16,0.0558,0.0159,0.0160,3.51,3.48,1.0063,0.0156,3.57,0.6433,0.1546,4.16 +n=11000000,fp32,0.0785,0.0288,0.0329,2.73,2.39,1.1424,0.0289,2.72,1.1971,0.2863,4.18 +n=11000000,int8,0.0492,0.0098,0.0100,5.03,4.90,1.0204,0.0100,4.92,0.5630,0.1040,5.41 +n=12000000,fp16,0.0607,0.0170,0.0170,3.58,3.56,1.0000,0.0173,3.50,0.3536,0.1621,2.18 +n=12000000,bf16,0.0607,0.0170,0.0171,3.57,3.55,1.0059,0.0171,3.55,0.3536,0.1621,2.18 +n=12000000,fp32,0.0873,0.0310,0.0356,2.81,2.46,1.1484,0.0312,2.80,0.4812,0.3083,1.56 +n=12000000,int8,0.0537,0.0104,0.0107,5.19,5.03,1.0288,0.0104,5.16,0.3241,0.1097,2.95 +n=13000000,fp16,0.0656,0.0181,0.0182,3.62,3.60,1.0055,0.0184,3.56,0.7874,0.1789,4.40 +n=13000000,bf16,0.0652,0.0181,0.0182,3.60,3.59,1.0055,0.0182,3.57,0.7877,0.1789,4.40 +n=13000000,fp32,0.0948,0.0334,0.0382,2.84,2.48,1.1437,0.0338,2.81,1.4638,0.3389,4.32 +n=13000000,int8,0.0584,0.0109,0.0112,5.33,5.20,1.0275,0.0112,5.19,0.6391,0.1236,5.17 +n=14000000,fp16,0.0698,0.0194,0.0195,3.60,3.59,1.0052,0.0193,3.61,0.4033,0.1873,2.15 +n=14000000,bf16,0.0697,0.0193,0.0194,3.60,3.59,1.0052,0.0194,3.60,0.4033,0.1872,2.15 +n=14000000,fp32,0.1019,0.0357,0.0406,2.85,2.51,1.1373,0.0361,2.82,0.5679,0.3649,1.56 +n=14000000,int8,0.0631,0.0114,0.0118,5.51,5.34,1.0351,0.0119,5.28,0.3812,0.1290,2.96 +n=15000000,fp16,0.0742,0.0206,0.0207,3.61,3.59,1.0049,0.0205,3.61,0.8740,0.2048,4.27 +n=15000000,bf16,0.0741,0.0207,0.0207,3.59,3.57,1.0000,0.0206,3.59,0.8742,0.2048,4.27 +n=15000000,fp32,0.1077,0.0380,0.0434,2.83,2.48,1.1421,0.0383,2.81,1.6274,0.3898,4.17 +n=15000000,int8,0.0671,0.0126,0.0130,5.33,5.18,1.0317,0.0126,5.32,0.7062,0.1372,5.15 +n=16000000,fp16,0.0786,0.0217,0.0217,3.62,3.62,1.0000,0.0216,3.63,0.4686,0.2094,2.24 +n=16000000,bf16,0.0789,0.0216,0.0216,3.66,3.65,1.0000,0.0218,3.62,0.4686,0.2094,2.24 +n=16000000,fp32,0.1148,0.0404,0.0460,2.85,2.49,1.1386,0.0408,2.81,0.6556,0.4141,1.58 +n=16000000,int8,0.0717,0.0132,0.0136,5.43,5.29,1.0303,0.0132,5.42,0.4269,0.1439,2.97 +n=17000000,fp16,0.0833,0.0227,0.0227,3.67,3.67,1.0000,0.0228,3.65,0.9882,0.2270,4.35 +n=17000000,bf16,0.0833,0.0229,0.0229,3.64,3.63,1.0000,0.0229,3.63,0.9881,0.2268,4.36 +n=17000000,fp32,0.1211,0.0426,0.0486,2.85,2.49,1.1408,0.0430,2.82,1.8417,0.4333,4.25 +n=17000000,int8,0.0759,0.0134,0.0139,5.65,5.47,1.0373,0.0140,5.41,0.8646,0.1517,5.70 +n=18000000,fp16,0.0882,0.0241,0.0241,3.67,3.66,1.0000,0.0242,3.64,0.5239,0.2381,2.20 +n=18000000,bf16,0.0884,0.0241,0.0241,3.67,3.67,1.0000,0.0243,3.63,0.5240,0.2383,2.20 +n=18000000,fp32,0.1282,0.0450,0.0513,2.85,2.50,1.1400,0.0452,2.83,0.6857,0.4668,1.47 +n=18000000,int8,0.0800,0.0145,0.0148,5.51,5.39,1.0207,0.0145,5.52,0.4955,0.1612,3.07 +n=19000000,fp16,0.0933,0.0252,0.0253,3.70,3.69,1.0040,0.0254,3.68,1.0996,0.2532,4.34 +n=19000000,bf16,0.0932,0.0251,0.0252,3.70,3.70,1.0040,0.0256,3.64,1.0997,0.2532,4.34 +n=19000000,fp32,0.1347,0.0474,0.0538,2.84,2.50,1.1350,0.0478,2.82,2.0164,0.4939,4.08 +n=19000000,int8,0.0842,0.0151,0.0156,5.57,5.41,1.0331,0.0152,5.56,0.9225,0.1688,5.47 +n=20000000,fp16,0.0984,0.0264,0.0264,3.73,3.73,1.0000,0.0268,3.67,0.5842,0.2584,2.26 +n=20000000,bf16,0.0983,0.0265,0.0265,3.72,3.71,1.0000,0.0263,3.74,0.5842,0.2584,2.26 +n=20000000,fp32,0.1404,0.0496,0.0564,2.83,2.49,1.1371,0.0502,2.80,0.8081,0.5077,1.59 +n=20000000,int8,0.0883,0.0155,0.0160,5.71,5.53,1.0323,0.0156,5.66,0.5361,0.1749,3.07 diff --git a/results/logs/autotune_logs/interleave_autotune.json b/results/logs/autotune_logs/interleave_autotune.json index e4f22d57..60216b8d 100644 --- a/results/logs/autotune_logs/interleave_autotune.json +++ b/results/logs/autotune_logs/interleave_autotune.json @@ -13,6 +13,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 1024 } }, { @@ -29,6 +32,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 16 + }, + "nki_autotune_cfg": { + "block_size": 2048 } }, { @@ -45,6 +51,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 2048 } }, { @@ -61,6 +70,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 1024 } }, { @@ -77,6 +89,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 1024 } }, { @@ -93,6 +108,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 2048 } }, { @@ -109,6 +127,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 2048 } }, { @@ -125,6 +146,9 @@ "cutile_autotune_cfg": { "tile": 4096, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 2048 } }, { @@ -141,6 +165,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 1024 } }, { @@ -157,6 +184,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 16 + }, + "nki_autotune_cfg": { + "block_size": 2048 } }, { @@ -173,6 +203,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 1024 } }, { @@ -189,6 +222,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 16 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -205,6 +241,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -221,6 +260,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 8192 } }, { @@ -237,6 +279,9 @@ "cutile_autotune_cfg": { "tile": 4096, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 8192 } }, { @@ -253,6 +298,9 @@ "cutile_autotune_cfg": { "tile": 4096, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -269,6 +317,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 2048 } }, { @@ -285,6 +336,9 @@ "cutile_autotune_cfg": { "tile": 4096, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 1024 } }, { @@ -301,6 +355,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 2048 } }, { @@ -317,6 +374,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -333,6 +393,9 @@ "cutile_autotune_cfg": { "tile": 8192, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -349,6 +412,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 16 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -365,6 +431,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -381,6 +450,9 @@ "cutile_autotune_cfg": { "tile": 4096, "occupancy": 16 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -397,6 +469,9 @@ "cutile_autotune_cfg": { "tile": 8192, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -413,6 +488,9 @@ "cutile_autotune_cfg": { "tile": 8192, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 1024 } }, { @@ -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": 4096, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 8192 } }, { @@ -461,6 +545,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 8192 } }, { @@ -477,6 +564,9 @@ "cutile_autotune_cfg": { "tile": 4096, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 2048 } }, { @@ -493,6 +583,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 2048 } }, { @@ -509,6 +602,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 8192 } }, { @@ -525,6 +621,9 @@ "cutile_autotune_cfg": { "tile": 8192, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -541,6 +640,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 8192 } }, { @@ -557,6 +659,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -573,6 +678,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 8192 } }, { @@ -589,6 +697,9 @@ "cutile_autotune_cfg": { "tile": 4096, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -605,6 +716,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 16 + }, + "nki_autotune_cfg": { + "block_size": 2048 } }, { @@ -621,6 +735,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 2048 } }, { @@ -637,6 +754,9 @@ "cutile_autotune_cfg": { "tile": 4096, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 8192 } }, { @@ -653,6 +773,9 @@ "cutile_autotune_cfg": { "tile": 4096, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 2048 } }, { @@ -669,6 +792,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 2048 } }, { @@ -685,6 +811,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -701,6 +830,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 8192 } }, { @@ -717,6 +849,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -733,6 +868,9 @@ "cutile_autotune_cfg": { "tile": 4096, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 1024 } }, { @@ -749,6 +887,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 2048 } }, { @@ -765,6 +906,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 8192 } }, { @@ -781,6 +925,9 @@ "cutile_autotune_cfg": { "tile": 4096, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 8192 } }, { @@ -797,6 +944,9 @@ "cutile_autotune_cfg": { "tile": 4096, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 1024 } }, { @@ -813,6 +963,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 8192 } }, { @@ -829,6 +982,9 @@ "cutile_autotune_cfg": { "tile": 4096, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 8192 } }, { @@ -845,6 +1001,9 @@ "cutile_autotune_cfg": { "tile": 4096, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -861,6 +1020,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 1024 } }, { @@ -877,6 +1039,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 8192 } }, { @@ -893,6 +1058,9 @@ "cutile_autotune_cfg": { "tile": 4096, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 8192 } }, { @@ -909,6 +1077,9 @@ "cutile_autotune_cfg": { "tile": 4096, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -925,6 +1096,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 2048 } }, { @@ -941,6 +1115,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 8192 } }, { @@ -957,6 +1134,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -973,6 +1153,9 @@ "cutile_autotune_cfg": { "tile": 4096, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -989,6 +1172,9 @@ "cutile_autotune_cfg": { "tile": 4096, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 1024 } }, { @@ -1005,6 +1191,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 8192 } }, { @@ -1021,6 +1210,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -1037,6 +1229,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 2048 } }, { @@ -1053,6 +1248,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 8192 } }, { @@ -1069,6 +1267,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 8192 } }, { @@ -1085,6 +1286,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -1101,6 +1305,9 @@ "cutile_autotune_cfg": { "tile": 4096, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 8192 } }, { @@ -1117,6 +1324,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 16 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -1133,6 +1343,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 8192 } }, { @@ -1149,6 +1362,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 16 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -1165,6 +1381,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -1181,6 +1400,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 8192 } }, { @@ -1197,6 +1419,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 8192 } }, { @@ -1213,6 +1438,9 @@ "cutile_autotune_cfg": { "tile": 4096, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 8192 } }, { @@ -1229,6 +1457,9 @@ "cutile_autotune_cfg": { "tile": 4096, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 8192 } }, { @@ -1245,6 +1476,9 @@ "cutile_autotune_cfg": { "tile": 1024, "occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 2048 } }, { @@ -1261,6 +1495,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 4096 } }, { @@ -1277,6 +1514,9 @@ "cutile_autotune_cfg": { "tile": 2048, "occupancy": 16 + }, + "nki_autotune_cfg": { + "block_size": 8192 } } -] \ No newline at end of file +] From ba8a222a18861b1f7d2c613e399d32b56d628bcf Mon Sep 17 00:00:00 2001 From: Bowen Cui Date: Mon, 21 Sep 2026 20:37:06 +0000 Subject: [PATCH 6/7] nki(interleave): migrate to the tilebench package layout Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01ScXYNjrrKGgDUVNHxv7HJt --- tilebench/benchmarks/operators/interleave/impl_nki.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tilebench/benchmarks/operators/interleave/impl_nki.py b/tilebench/benchmarks/operators/interleave/impl_nki.py index c396f40b..5315ba3e 100644 --- a/tilebench/benchmarks/operators/interleave/impl_nki.py +++ b/tilebench/benchmarks/operators/interleave/impl_nki.py @@ -29,7 +29,7 @@ import torch -from core.nki_autotune import NkiAutotuner +from tilebench.core.nki_autotune import NkiAutotuner try: import nki From 2817f0012eb6c8fc700aa2e5ad9aa2f395bbd446 Mon Sep 17 00:00:00 2001 From: Bowen Cui Date: Tue, 22 Sep 2026 04:57:28 +0000 Subject: [PATCH 7/7] nki(interleave): drop docstrings and inline comments Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01ScXYNjrrKGgDUVNHxv7HJt --- .../operators/interleave/impl_nki.py | 44 +------------------ .../operators/interleave/impl_torch.py | 6 --- 2 files changed, 2 insertions(+), 48 deletions(-) diff --git a/tilebench/benchmarks/operators/interleave/impl_nki.py b/tilebench/benchmarks/operators/interleave/impl_nki.py index 5315ba3e..20c066be 100644 --- a/tilebench/benchmarks/operators/interleave/impl_nki.py +++ b/tilebench/benchmarks/operators/interleave/impl_nki.py @@ -1,26 +1,3 @@ -"""NKI (AWS Trainium) implementation of interleave: ``out[0::2] = A; out[1::2] = B``. - -Design (mirrors the Triton / cuTile kernels): - -* The flat inputs are processed in fixed-size blocks. A Triton program handles - ``BLOCK_SIZE`` contiguous elements of ``A`` and ``B`` and writes ``2 * BLOCK_SIZE`` - outputs; the NKI analogue is one ``[128, BLOCK_SIZE]`` SBUF tile per input -- - 128 partitions (the 128 parallel lanes of the Vector engine) times - ``BLOCK_SIZE`` contiguous elements per partition. The default and the autotune - search space use the same ``BLOCK_SIZE`` values as ``impl_triton.py`` / - ``impl_cutile.py`` (default 1024; search 1024/2048/4096/8192). -* The interleaving itself is done on-chip with strided SBUF destination slices - (``out_tile[:, 0::2] = a``, ``out_tile[:, 1::2] = b``), so every HBM load and - the HBM store stay fully contiguous DMAs. -* Work is split across the NeuronCores of the logical core (LNC2 on trn2): - each program instance (``nl.program_id(0)``) owns a contiguous half of the - element range, so both physical cores' DMA engines are used. -* Tails are handled inside the kernel: the last block of a core may cover fewer - than ``128 * BLOCK_SIZE`` elements, in which case it is processed as one - ``[128, q]`` tile plus one ``[1, r]`` tile (``r < 128``). No host-side - padding / reshape / slicing is needed -- those torch ops would be fused into - the same NEFF as the kernel and dominate its runtime. -""" import functools import os import re @@ -43,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()) @@ -67,16 +39,9 @@ def _lnc_degree() -> int: if nki is not None: @nki.jit def interleave_kernel(a_input, b_input, block_size): - """``out[2i] = a[i]; out[2i+1] = b[i]`` over flat ``(n,)`` HBM tensors. - - Args: - a_input, b_input: flat ``(n,)`` tensors in HBM (same dtype). - block_size: elements per partition per tile (compile-time constant). - """ n = a_input.shape[0] out = nl.ndarray((2 * n,), dtype=a_input.dtype, buffer=nl.shared_hbm) - # Contiguous element range owned by this program instance (NeuronCore). num_programs = nl.num_programs() per_core = (n + num_programs - 1) // num_programs lo = nl.program_id(0) * per_core @@ -86,8 +51,8 @@ def interleave_kernel(a_input, b_input, block_size): for j in range(max(0, (hi - lo + chunk - 1) // chunk)): start = lo + j * chunk count = min(chunk, hi - start) - q = count // PMAX # full partitions: [PMAX, q] tiles - r = count - q * PMAX # remaining < PMAX elements: [1, r] tiles + q = count // PMAX + r = count - q * PMAX if q > 0: _tile_body(PMAX, q, start, a_input, b_input, out) if r > 0: @@ -95,21 +60,16 @@ def interleave_kernel(a_input, b_input, block_size): return out def _tile_body(p, f, start, a_input, b_input, out): - """Interleave the ``[p, f]`` tiles of a and b at flat input offset ``start``.""" a_tile = nl.ndarray((p, f), dtype=a_input.dtype, buffer=nl.sbuf) nisa.dma_copy(dst=a_tile, src=a_input.ap(pattern=[[f, p], [1, f]], offset=start)) b_tile = nl.ndarray((p, f), dtype=b_input.dtype, buffer=nl.sbuf) nisa.dma_copy(dst=b_tile, src=b_input.ap(pattern=[[f, p], [1, f]], offset=start)) - # On-chip interleave into the even / odd free-dim slots of one [p, 2f] tile. out_tile = nl.ndarray((p, 2 * f), dtype=a_input.dtype, buffer=nl.sbuf) nisa.tensor_copy(dst=out_tile[0:p, 0:2 * f:2], src=a_tile) nisa.tensor_copy(dst=out_tile[0:p, 1:2 * f:2], src=b_tile) - # Partition p of the tile holds inputs [start + p*f, start + (p+1)*f), whose - # outputs are the contiguous range [2*(start + p*f), 2*(start + (p+1)*f)). nisa.dma_copy(dst=out.ap(pattern=[[2 * f, p], [1, 2 * f]], offset=2 * start), src=out_tile) -# Same block sizes as impl_triton.py / impl_cutile.py (default 1024; search 1024/2048/4096/8192). _DEFAULT_CONFIG = SimpleNamespace(block_size=1024) _SEARCH_SPACE = [SimpleNamespace(block_size=b) for b in (1024, 2048, 4096, 8192)] _kernel = interleave_kernel[_lnc_degree()] if nki is not None else None diff --git a/tilebench/benchmarks/operators/interleave/impl_torch.py b/tilebench/benchmarks/operators/interleave/impl_torch.py index afe7f46b..b8e405fa 100644 --- a/tilebench/benchmarks/operators/interleave/impl_torch.py +++ b/tilebench/benchmarks/operators/interleave/impl_torch.py @@ -3,7 +3,6 @@ def _view_cols(n: int, lo: int = 128, hi: int = 8192) -> int | None: - """Largest divisor of ``n`` in ``[lo, hi]`` (row length of the 2-D view), or None.""" for c in range(hi, lo - 1, -1): if n % c == 0: return c @@ -19,11 +18,6 @@ def _interleave(a: torch.Tensor, b: torch.Tensor, n: int) -> torch.Tensor: def run(A: torch.Tensor, B: torch.Tensor, N: int, **kwargs): if A.device.type == "xla" and N % 128 != 0: - # Neuron/XLA: the strided assignments compile to one streaming pass when the - # length is a multiple of 128 (2M elements: 0.08 ms) but to element-wise - # scatters otherwise (1M: 10 ms, 3M: 226 ms). Pad the row length of a 2-D view - # to a multiple of 128 (strided DMAs), interleave the padded flat tensors and - # drop the padding again. cols = _view_cols(N) if cols is not None: rows = N // cols