nki(top_k_selection): NKI (Trainium) implementation

[bowencui123]

NKI (AWS Trainium) implementation of top_k_selection, 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/top_k_selection/impl_nki.py, M benchmarks/operators/top_k_selection/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 (fef24d4)

NKI tunables	Triton counterpart	note
block_size (per-partition block cap of the hierarchy)	BLOCK_SIZE	timing proxy: first hierarchy level

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
[top_k_selection] default : verify=OK (2s) 
[top_k_selection] autotune: verify=OK (0s) last_config={'block': 256, 'K2': 16, 'levels': 2, 'block_size': 4096} trace_records=1 
STUB_EXIT=0
Params   |    Dtype |  Torch(ms) |      NKI(ms) |  Speedup(N)
n=4096   | fp32     |     0.0278 |       0.0706 |        0.39

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

#305 top_k_selection — NO CODE CHANGE

• Hierarchical block-top-k ≙ Triton's _run_hierarchy; .values→[0] torch fix is correct
  (torch_xla returns a list). Clean.
• Tested (case 0): default 0.0664 ms (torch 0.0261), autotune 0.0665 ms,
  winner namespace(block_size=2048) (first-level proxy).

🤖 Generated with Claude Code

https://claude.ai/code/session_011Wvs1tztZTGQFD78YdZaha

[bowencui123]

Baseline policy follow-up: impl_torch.py is now device-dispatched instead of replacing .values with [0] globally — the non-XLA path is byte-identical to main, only XLA (where the topk namedtuple comes back as a plain list) uses [0]. The two spellings return the same tensor; verified equal on CPU and re-validated on trn2 (torch 0.0260 ms, NKI 0.0662 ms).

🤖 Generated with Claude Code

https://claude.ai/code/session_011Wvs1tztZTGQFD78YdZaha