nki(batched_matmul): NKI (Trainium) implementation

[bowencui123]

NKI (AWS Trainium) implementation of batched_matmul, split out of the consolidated NKI branch cecilia/feature/nki-vector-add (nki-all-operators, #259) so each operator can be reviewed independently. Supersedes #174 (older per-operator branch: legacy neuronxcc.nki imports; this is the migrated import nki version).

Files: A benchmarks/operators/batched_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

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

🤖 Generated with Claude Code

https://claude.ai/code/session_012Q38kGmXvyoeM1qtCbheSL

Autotune (0ce4425)

NKI tunables	Triton counterpart	note
block_size_m, block_size_k, block_size_n (PE tile sizes <=128/128/512)	BLOCK_SIZE_M/BLOCK_SIZE_N/BLOCK_SIZE_K	—

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
[batched_matmul] default : verify=OK (1s) 
[batched_matmul] autotune: verify=OK (0s) last_config={'block_size_m': 32, 'block_size_k': 32, 'block_size_n': 32} trace_records=1 
STUB_EXIT=0
Params      |    Dtype |  Torch(ms) |      NKI(ms) |  Speedup(N)
n=2097152   | fp16     |     0.0252 |       0.0535 |        0.47

[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).

#269 batched_matmul — CHANGES PUSHED

• Approach ✓ (batch split across SPMD programs, per-element PE tiling; raw tile sizes as tunables,
  documented). Removed the NKI_BMM_NUM_CORES env knob → LNC-divisibility rule, same as nki(matmul_fp32_fp16_fp8): NKI (Trainium) implementation #288.
• Tested (BATCH=32, M=N=K=32 fp16): default 0.0493 ms (torch 0.0189 — tiny matrices, expected),
  autotune 0.0493 ms (the ≤M/K/N size filter leaves only the 32/32/32 default at this case).

🤖 Generated with Claude Code

https://claude.ai/code/session_011Wvs1tztZTGQFD78YdZaha