nki(radix_sort): NKI (Trainium) implementation

[bowencui123]

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

Files: A benchmarks/operators/radix_sort/impl_nki.py, M benchmarks/operators/radix_sort/impl_torch.py

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

• run() calls mark_step() per radix pass -> several compiled graphs; the exact-identity profiler (PR fix(nki): deterministic NEFF identity — exact winner replay in isolated workers, validated manifests #261) refuses to time a fragment (nki_ms=nan) until multi-graph aggregation lands
• 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 (11f19c5)

NKI tunables	Triton counterpart	note
block_size (elements per partition block S)	BLOCK_SIZE	timing proxy: one radix pass; end-to-end profiling of this multi-graph operator is a framework follow-up

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
[radix_sort] default : verify=OK (283s) 
[radix_sort] autotune: verify=OK (1020s) last_config={'block_size': 256} trace_records=1 
STUB_EXIT=0
  NKI FAILED: profile worker failed: NkiArtifactIdentityError: 2 valid NKI (marker-bearing) NEFF/HLO pairs found — artifact identity is ambiguous
Params      |    Dtype |  Torch(ms) |      NKI(ms) |  Speedup(N)
n=1000000   | int32    |        nan |          nan |        0.00

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

#296 radix_sort — CHANGES PUSHED

• Kernel ✓ (2-bit/pass ≙ Triton/cuTile; per-pass graphs are REQUIRED for correctness and the merged
  feat(nki): time cases from the Neuron runtime's execution trace (real inputs, warmup+repeat, multi-graph sums) #309 framework now times them as a sum — validated 327.7 ms before this review).
• Fixed: impl_torch.py had replaced torch.sort globally with a bitonic where-network — that
  would silently change the B200 GPU baseline. Now device-dispatched: non-XLA keeps
  torch.sort, only XLA (where XLA sort is NCC_EVRF029-unsupported) uses the network.
• Updated the stale "profiler times one pass only" comment (pre-feat(nki): time cases from the Neuron runtime's execution trace (real inputs, warmup+repeat, multi-graph sums) #309).
• Tested (n=1e6 int32): default NKI 327.8485 ms, torch-on-Neuron baseline now WORKS via the
  device dispatch (bitonic network, one graph, 2625.4415 ms) → speedup 8.01× where it used to be
  nan; autotune 327.7160 ms, winner namespace(block_size=1024) (17 graphs per iteration summed by the
  runtime-trace framework in both modes).

🤖 Generated with Claude Code

https://claude.ai/code/session_011Wvs1tztZTGQFD78YdZaha