From 8c008020a171127259b300692aec54a7a82240e9 Mon Sep 17 00:00:00 2001 From: bowencui123 Date: Sat, 29 Aug 2026 03:01:37 +0000 Subject: [PATCH 1/2] nki(matmul_int8): 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 #181 (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/matmul_int8/impl_nki.py | 180 +++++++++++++++++++ 1 file changed, 180 insertions(+) create mode 100644 benchmarks/operators/matmul_int8/impl_nki.py diff --git a/benchmarks/operators/matmul_int8/impl_nki.py b/benchmarks/operators/matmul_int8/impl_nki.py new file mode 100644 index 00000000..a923086c --- /dev/null +++ b/benchmarks/operators/matmul_int8/impl_nki.py @@ -0,0 +1,180 @@ +import torch + +try: + import nki + import nki.isa as nisa + import nki.language as nl +except ImportError: + nki = None + +# --- Hardware constants (NeuronCore-v2/v3) --------------------------------- +# nc_matmul: dst[M, N] = stationary[K, M].T @ moving[K, N] +TILE_M = 128 # stationary free dim (output rows per PE tile) <= 128 +TILE_K = 128 # contraction dim per matmul (== partition dim) <= 128 +TILE_N = 512 # moving free dim == one fp32 PSUM bank (128 x 2KB) <= 512 + +# Tuning knobs (all must divide the corresponding problem dimension). +TILES_IN_BLOCK_M = 4 # 4 * 128 = 512 output rows share one K-loop / PSUM group +BLOCK_KB = 512 # packed-B rows (k_b) fetched per A/B slab + +if nki is not None: + @nki.jit + def matmul_int8_kernel(a_hbm, b_hbm, tiles_in_block_m, block_kb, tile_n): + M, K = a_hbm.shape + K_b, N = b_hbm.shape + + assert len(a_hbm.shape) == 2 and len(b_hbm.shape) == 2, "a and b must be 2D" + assert K == 4 * K_b, "A's K dim must be 4x B's packed K_b dim" + + block_m = TILE_M * tiles_in_block_m + assert M % block_m == 0, "M must be a multiple of TILE_M * tiles_in_block_m" + assert N % tile_n == 0, "N must be a multiple of tile_n" + assert K_b % block_kb == 0, "K_b must be a multiple of block_kb" + assert block_kb % TILE_K == 0, "block_kb must be a multiple of TILE_K" + assert tile_n <= 512, "tile_n exceeds the fp32 PSUM bank size" + + num_block_m = M // block_m + num_block_n = N // tile_n + num_kb_blocks = K_b // block_kb + num_j = block_kb // TILE_K # 128-row k_b tiles per slab + + out_hbm = nl.ndarray((M, N), dtype=nl.int32, buffer=nl.shared_hbm) + + # SBUF slab widths (flat 2D tiles: partition dim first, everything else + # packed into the free dim so all indexing stays 2D). + a_slab_w = tiles_in_block_m * 4 * block_kb # [TILE_M, .] + at_slab_w = tiles_in_block_m * 4 * num_j * TILE_M # [TILE_K, .] + + for m_blk in range(num_block_m): + m0 = m_blk * block_m + + for n_blk in range(num_block_n): + n0 = n_blk * tile_n + + # One fp32 PSUM bank per 128-row output tile; accumulates over + # the whole contraction (all kb-slabs x all j x all 4 packed + # fields). The per-tile stride is a full bank (TILE_N fp32 == + # 2KB) even when tile_n < TILE_N: an nc_matmul destination has + # to start on a PSUM bank boundary. + acc_psum = nl.ndarray((TILE_M, tiles_in_block_m * TILE_N), + dtype=nl.float32, buffer=nl.psum) + + for kbb in range(num_kb_blocks): + kb0 = kbb * block_kb + + # ---- A slab: the 4 packed fields read 4 disjoint column + # ranges of A (field i lives at columns i*K_b + .). + a_i8 = nl.ndarray((TILE_M, a_slab_w), dtype=nl.int8, buffer=nl.sbuf) + for bm in range(tiles_in_block_m): + row0 = m0 + bm * TILE_M + for i in range(4): + off = (bm * 4 + i) * block_kb + col0 = i * K_b + kb0 + nisa.dma_copy( + dst=a_i8[0:TILE_M, off:off + block_kb], + src=a_hbm[row0:row0 + TILE_M, col0:col0 + block_kb], + ) + + # int8 -> bf16 (nc_transpose / nc_matmul need a float dtype) + a_bf16 = nl.ndarray((TILE_M, a_slab_w), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.tensor_copy(dst=a_bf16, src=a_i8) + + # ---- transpose to lhsT form [K, M] for the stationary operand + a_t = nl.ndarray((TILE_K, at_slab_w), dtype=nl.bfloat16, buffer=nl.sbuf) + for bm in range(tiles_in_block_m): + for i in range(4): + src_off = (bm * 4 + i) * block_kb + for j in range(num_j): + t_psum = nl.ndarray((TILE_K, TILE_M), dtype=nl.bfloat16, + buffer=nl.psum) + nisa.nc_transpose( + dst=t_psum, + data=a_bf16[0:TILE_M, + src_off + j * TILE_K: + src_off + (j + 1) * TILE_K], + ) + dst_off = ((bm * 4 + i) * num_j + j) * TILE_M + nisa.tensor_copy( + dst=a_t[0:TILE_K, dst_off:dst_off + TILE_M], + src=t_psum, + ) + + # ---- K_b tile loop (Triton's `j`) ------------------------ + for j in range(num_j): + k0 = kb0 + j * TILE_K + + b_u8 = nl.ndarray((TILE_K, tile_n), dtype=nl.uint8, buffer=nl.sbuf) + nisa.dma_copy(dst=b_u8, + src=b_hbm[k0:k0 + TILE_K, n0:n0 + tile_n]) + + # ---- 4 packed 2-bit fields (Triton's `i`) ------------ + for i in range(4): + shift = 2 * i + + masked = nl.ndarray((TILE_K, tile_n), dtype=nl.uint8, buffer=nl.sbuf) + nisa.tensor_scalar(dst=masked, data=b_u8, op0=nl.bitwise_and, operand0=(3 << shift)) + + b_val = nl.ndarray((TILE_K, tile_n), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.tensor_scalar(dst=b_val, data=masked, op0=nl.multiply, operand0=1.0 / float(1 << shift), op1=nl.add, operand1=-1.0) + + first = (j == 0 and i == 0) + for bm in range(tiles_in_block_m): + stat_off = ((bm * 4 + i) * num_j + j) * TILE_M + nisa.nc_matmul( + dst=acc_psum[0:TILE_M, bm * TILE_N:bm * TILE_N + tile_n], + stationary=a_t[0:TILE_K, stat_off:stat_off + TILE_M], + moving=b_val, + accumulate=(kbb > 0) if first else True, + ) + + # ---- fp32 accumulator holds an exact integer -> int32 -------- + for bm in range(tiles_in_block_m): + out_sb = nl.ndarray((TILE_M, tile_n), dtype=nl.int32, buffer=nl.sbuf) + nisa.tensor_copy(dst=out_sb, src=acc_psum[0:TILE_M, bm * TILE_N:bm * TILE_N + tile_n]) + row0 = m0 + bm * TILE_M + nisa.dma_copy(dst=out_hbm[row0:row0 + TILE_M, n0:n0 + tile_n], src=out_sb) + + return out_hbm + + +def _largest_divisor(dim: int, candidates) -> int | None: + """Largest value in ``candidates`` that divides ``dim`` (None if none do).""" + for c in candidates: + if dim % c == 0: + return c + return None + + +def run(a: torch.Tensor, b: torch.Tensor, block_size: int = None, + autotune: bool = False, **kwargs) -> torch.Tensor: + """A (M, K) int8 @ unpack(B (K/4, N) uint8) -> (M, N) int32.""" + M, K = a.shape + K_b, N = b.shape + + if K != 4 * K_b: + raise ValueError("Incompatible dimensions: A's K must equal 4 * B's K_b") + + if a.dtype != torch.int8: + a = a.to(torch.int8) + if b.dtype != torch.uint8: + # Reinterpret the packed byte as unsigned (wraps for int8 inputs). + b = b.to(torch.uint8) + + tile_n = _largest_divisor(N, (TILE_N, 256, 128)) + tiles_in_block_m = _largest_divisor(M, tuple(TILE_M * t for t in (TILES_IN_BLOCK_M, 2, 1))) + block_kb = _largest_divisor(K_b, (BLOCK_KB, 256, TILE_K)) + + if tile_n is None or tiles_in_block_m is None or block_kb is None: + raise NotImplementedError( + f"matmul_int8 NKI: shape (M={M}, N={N}, K={K}) unsupported -- " + f"M must be a multiple of {TILE_M}, N a multiple of 128 and " + f"K a multiple of {4 * TILE_K}" + ) + + tiles_in_block_m //= TILE_M + + return matmul_int8_kernel(a, b, tiles_in_block_m, block_kb, tile_n) + + +def get_last_config() -> dict | None: + return None From 2083bdb00adbe4df5c84b879f1a9f520a2a26279 Mon Sep 17 00:00:00 2001 From: bowencui123 Date: Sat, 29 Aug 2026 08:37:22 +0000 Subject: [PATCH 2/2] nki(matmul_int8): NkiAutotuner wiring (`block_size_m`, `block_size_k`, `block_size_n`) Tunables mirror the Triton search space (`BLOCK_SIZE_M`/`BLOCK_SIZE_N`/`BLOCK_SIZE_K`); defaults are the previous constants, so autotune=False is unchanged. `block_size_k` counts unpacked K (4 x packed rows) Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_012Q38kGmXvyoeM1qtCbheSL --- benchmarks/operators/matmul_int8/impl_nki.py | 30 ++++++++++++++++++-- 1 file changed, 27 insertions(+), 3 deletions(-) diff --git a/benchmarks/operators/matmul_int8/impl_nki.py b/benchmarks/operators/matmul_int8/impl_nki.py index a923086c..0ec20eed 100644 --- a/benchmarks/operators/matmul_int8/impl_nki.py +++ b/benchmarks/operators/matmul_int8/impl_nki.py @@ -1,5 +1,9 @@ +from types import SimpleNamespace + import torch +from core.nki_autotune import NkiAutotuner + try: import nki import nki.isa as nisa @@ -145,6 +149,10 @@ def _largest_divisor(dim: int, candidates) -> int | None: return None +_tuner = NkiAutotuner(matmul_int8_kernel) if nki is not None else None +_last_autotune_config: dict = {} + + def run(a: torch.Tensor, b: torch.Tensor, block_size: int = None, autotune: bool = False, **kwargs) -> torch.Tensor: """A (M, K) int8 @ unpack(B (K/4, N) uint8) -> (M, N) int32.""" @@ -172,9 +180,25 @@ def run(a: torch.Tensor, b: torch.Tensor, block_size: int = None, ) tiles_in_block_m //= TILE_M - - return matmul_int8_kernel(a, b, tiles_in_block_m, block_kb, tile_n) + _default = SimpleNamespace(block_size_m=TILE_M * tiles_in_block_m, block_size_k=4 * block_kb, + block_size_n=tile_n) + if autotune: + _space = [SimpleNamespace(block_size_m=bm, block_size_k=bk, block_size_n=bn) + for bm in (128, 256, 512) for bk in (512, 1024, 2048) for bn in (128, 256, 512) + if M % bm == 0 and K_b % (bk // 4) == 0 and N % bn == 0] + if not any(vars(c) == vars(_default) for c in _space): + _space.append(_default) + cfg = _tuner.tune_or_cached( + shape_key=((M, N, K), str(a.dtype)), + search_space=_space, + args_fn=lambda cfg: (a, b, cfg.block_size_m // TILE_M, cfg.block_size_k // 4, cfg.block_size_n), + ) + _last_autotune_config.clear() + _last_autotune_config.update(vars(cfg)) + else: + cfg = _default + return matmul_int8_kernel(a, b, cfg.block_size_m // TILE_M, cfg.block_size_k // 4, cfg.block_size_n) def get_last_config() -> dict | None: - return None + return dict(_last_autotune_config) or None