nki(matmul_fp32_fp16_fp8): NKI (Trainium) implementation

[bowencui123]

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

Files: M benchmarks/operators/matmul_fp32_fp16_fp8/config.yaml, A benchmarks/operators/matmul_fp32_fp16_fp8/impl_nki.py

Status: imports and exposes run()/get_last_config() on trn2 (nki 0.6.0); not individually re-benchmarked in this split

• also carries the operator's config.yaml change from the NKI branch

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

🤖 Generated with Claude Code

https://claude.ai/code/session_012Q38kGmXvyoeM1qtCbheSL

Autotune (4eb4f1a)

NKI tunables	Triton counterpart	note
block_size_m, block_size_n, block_size_k (= 128/512/128 PE tiles x tiles-in-block)	BLOCK_SIZE_M/BLOCK_SIZE_N/BLOCK_SIZE_K	candidates are filtered to divisors of M/N/K; candidate timing runs on one core

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
[matmul_fp32_fp16_fp8] default : verify=OK (3s) 
[matmul_fp32_fp16_fp8] autotune: verify=OK (3s) last_config={'block_size_m': 256, 'block_size_n': 512, 'block_size_k': 512} trace_records=1 
STUB_EXIT=0
Params          |    Dtype |  Torch(ms) |      NKI(ms) |  Speedup(N)
n=34359738368   | fp32     |     0.9270 |       0.9368 |        0.99

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

#288 matmul_fp32_fp16_fp8 — CHANGES PUSHED

• Approach ✓ (blocked PE tiling ≙ Triton/cuTile block GEMM; block_size_m/n/k ≙ BLOCK_SIZE_M/N/K).
• Removed scaffold (324→258 lines):
  • NKI_MATMUL_NUM_CORES env knob + double reconciliation dance → M blocks split across LNC when
    they divide evenly, else 1 core (same effective behaviour, no knob).
  • the entire fp8/double_row machinery (DOUBLE_ROW kernel branch, fp8 strided PSUM transpose,
    NKI_MATMUL_DOUBLE_ROW env): unreachable in the benchmark — the sweep's only fp8 dtype is
    e4m3fn, which neuronx-cc rejects before TRN3, so run() now rejects all fp8 with that message.
  • reverted the unrelated config.yaml plots addition (roofline etc.) — other matmul configs on
    main don't have it; single-purpose PR.
• Tested: fp32 default 1.0505 ms (torch 1.0424), fp16 default 0.2861 ms (torch 0.2589),
  autotune 1.0507 ms — real sweep, 7 candidates timed (1.827–1.928 ms baremetal), winner
  block_size_m/n/k = 512/1024/1024 (= the dynamic default).

🤖 Generated with Claude Code

https://claude.ai/code/session_011Wvs1tztZTGQFD78YdZaha