nki(weight_dequant): NKI (Trainium) implementation

[bowencui123]

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

Files: A benchmarks/operators/weight_dequant/impl_nki.py, M benchmarks/operators/weight_dequant/impl_torch.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 impl_torch.py 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 (47dd19d)

NKI tunables	Triton counterpart	note
block_size (rows per tile inside a scale band, <=128)	BLOCK_SIZE/TILE_SIZE	TILE_SIZE is the quantization block (semantic, not tuned)

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
[weight_dequant] default : verify=OK (1s) 
[weight_dequant] autotune: verify=OK (0s) last_config={'block_size': 128} trace_records=1 
STUB_EXIT=0
Params     |    Dtype |  Torch(ms) |      NKI(ms) |  Speedup(N)
n=262144   | fp16     |     0.0304 |       0.0389 |        0.78

[bowencui123]

Baseline policy (same rule applied across the NKI PRs): impl_torch.py is now device-dispatched instead of replacing the advanced-indexing gather globally — the non-XLA path is byte-identical to main, only XLA (where the gather faults at runtime with the DGE out-of-bound access) takes the repeat_interleave variant. CPU equivalence of the two paths verified incl. non-divisible M/N and TILE_SIZE > M; re-validated on trn2 (torch 0.0232 ms, NKI 0.0349 ms).

🤖 Generated with Claude Code

https://claude.ai/code/session_011Wvs1tztZTGQFD78YdZaha