nki(histogramming): NKI (Trainium) implementation

[bowencui123]

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

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

NKI tunables	Triton counterpart	note
block_size (values chunk per pass)	BLOCK_SIZE/BLOCK_ROWS/BLOCK_BINS	Triton's bins/rows split has no NKI analog (histogram accumulates per partition)

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
[histogramming] default : verify=OK (1s) 
[histogramming] autotune: verify=OK (0s) last_config={'block_size': 8192} trace_records=1 
STUB_EXIT=0
core.nki_artifact.NkiArtifactIdentityError: 2 valid non-NKI (torch) NEFF/HLO pairs found — artifact identity is ambiguous
Params     |    Dtype |  Torch(ms) |      NKI(ms) |  Speedup(N)
n=262144   | int32    |        nan |          nan |        0.00
# re-validation (tolerance parsing fixed in the checker; engine line above stands)
Torch (Neuron) baseline FAILED: torch-on-Neuron verification failed: torch baseline phase failed: NkiArtifactIdentityError: 2 valid non-NKI (torch) NEFF/HLO pairs found — artifact identity is ambiguous
Params     |    Dtype |  Torch(ms) |      NKI(ms) |  Speedup(N)
n=262144   | int32    |        nan |       0.5025 |         nan
BATCH3_DONE

[bowencui123]

Lost result for N=67108864, num_bins=4096

[bowencui123]

I think the main performance issue here is algorithmic rather than an autotuning issue.

The current NKI implementation iterates over ceil(num_bins / 128) bin blocks. For each bin block, it scans the entire N-element input again, broadcasts/materializes each input chunk across up to 128 partitions, compares every value against the bin ID owned by each partition, and then reduces along the free axis. Therefore, the effective number of full input passes scales with ceil(num_bins / 128).

This is quite different from the Triton/cuTile implementation. Triton and cuTile shard the input across up to 256 partial workers, so each input element is read once and performs one atomic add into a per-worker partial histogram. A second kernel then reduces the partial histograms across workers. The main work is therefore approximately:

O(N) + O(num_partials * num_bins)

instead of repeatedly scanning N for every 128-bin block.

The measurements are consistent with this. For N=67,108,864, autotuned NKI goes from:

• 111.8 ms for 64 bins
• 217.1 ms for 256 bins
• 870.2 ms for 1024 bins

This almost exactly follows the number of 128-bin blocks:

1 -> 2 -> 8

while Triton stays around 0.8-1.0 ms.

Could we rework the NKI implementation to follow the Triton/cuTile algorithm more closely? In particular, it may be worth investigating input sharding plus nl.atomic_rmw(..., op=nl.add) into per-worker partial histograms, followed by a reduction kernel.

I don't think tuning block_size alone can fix the current scaling, because the repeated full-input scan is structural.

That's the reason why autotune take more than 3 hours.
In N=67108864, num_bins=1024, nki takes at least 870ms, while triton only takes 0.935ms