nki(streamk_matmul): NKI (Trainium) implementation

[bowencui123]

NKI (AWS Trainium) implementation of streamk_matmul, split out of the consolidated NKI branch cecilia/feature/nki-vector-add (nki-all-operators, #259) so each operator can be reviewed independently.

Files: A benchmarks/operators/streamk_matmul/impl_nki.py

Status: imports and exposes run()/get_last_config() on trn2 (nki 0.6.0); not individually re-benchmarked in this split — ⚠️ import check FAILED in this worktree:

Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/tmp/claude-1000/-home-ubuntu/10e45e99-5a58-44c3-a40e-687dadc1a343/scratchpad/wt/ops/streamk_matmul/benchmarks/operators/streamk_matmul/impl_nki.py", line 3, in <module>
    from benchmarks.operators.matmul_fp32_fp16_fp8.impl_nki import (
ModuleNotFoundError: No module named 'benchmarks.operators.matmul_fp32_fp16_fp8.impl_nki'

• imports its kernel from matmul_fp32_fp16_fp8/impl_nki.py -> merge after that PR

Implementation by @Cecilia123li. Timing/identity infrastructure: #261; Trainium peak/roofline infra: #262.

🤖 Generated with Claude Code

https://claude.ai/code/session_012Q38kGmXvyoeM1qtCbheSL

Autotune (3cd42a9)

NKI tunables	Triton counterpart	note
inherits matmul_fp32_fp16_fp8's block_size_m/n/k	BLOCK_M/BLOCK_N/BLOCK_K	delegates to the matmul kernel (no Stream-K on NKI)

autotune=False keeps the previous constants (default numbers unchanged). Validation on trn2, case 0 (default run + autotune code path with the candidate timer stubbed — no sweep; --autotune runs a real sweep):

# initial run
[streamk_matmul] default : verify=FAIL (3s) Comparing
[streamk_matmul] autotune: verify=FAIL (3s) last_config={'block_size_m': 256, 'block_size_n': 512, 'block_size_k': 512} trace_records=1 Comparing
STUB_EXIT=1
Params   |    Dtype |  Torch(ms) |      NKI(ms) |  Speedup(N)
n=4096   | fp16     |     0.2803 |       0.3202 |        0.88
# re-validation (tolerance parsing fixed in the checker; engine line above stands)
[streamk_matmul] default : verify=OK (1s) 
[streamk_matmul] autotune: verify=OK (2s) last_config={'block_size_m': 256, 'block_size_n': 512, 'block_size_k': 512} trace_records=1 
STUB_EXIT=0

[bowencui123]

Review pass (Bowen's assignment): Triton/cuTile alignment + scaffold removal + real autotune & default runs on trn2 (post-#309 runtime-trace timing framework, case 0, --warmup 2 --repeat 10; autotune numbers come from a real selector sweep replayed by the profile worker).

#303 streamk_matmul — NO CODE CHANGE

• Plain-matmul delegation kept; the header documents why real Stream-K is blocked (launch degree
  pinned to LNC, no cross-worker atomic accumulation — dma_compute requires unique indices).
  Same conclusion as before; nothing mislabeled.
• Tested (m1024 n4096 k4096 fp16): default 0.3442 ms (torch 0.2894), autotune 0.3442 ms —
  delegate's sweep timed 7 candidates (0.488–0.754 ms), winner 512/1024/1024.

🤖 Generated with Claude Code

https://claude.ai/code/session_011Wvs1tztZTGQFD78YdZaha