Fix decode attention for Neuron SDK 2.32 and validate on Inf2

[varuntej07]

Description of changes:

contributed/decode_attention.py, merged in #129, no longer runs on a current install. Its correctness checks call nki.simulate_kernel, which Neuron SDK 2.32 / NKI 0.6.0 removed along with nki.baremetal and nki.benchmark. #131 moved the repo to 2.32 and rewrote the attention_fwd_performance tutorial, but did not touch contributed/, so the break is latent in what is already on main:

$ python contributed/decode_attention.py
AttributeError: module 'nki' has no attribute 'simulate_kernel'

Updated per review to contain only the port and the range change. The device benchmark that was originally in this PR has been split out.

The kernels themselves are unchanged. decode_attention_fwd, decode_attention_gqa_fwd and both NumPy references differ only in the loop iterators and comments described below.

1. Harness ported to the 0.6.0 API

Pre-0.6.0	Now
nki.simulate_kernel(kernel, *args)	nki.simulate(kernel)(*args)
nki.baremetal()(kernel)(*args)	kernel(*args)

The second row is per @nki.jit's own docstring, "compiles and executes standalone kernel, without a framework": a plain call with NumPy arrays is the device path on 0.6.0, which is why baremetal was removed rather than renamed.

Adds --backend {simulate,baremetal}, defaulting to the device when /dev/neuron0 exists, so python contributed/decode_attention.py does the right thing on a laptop or on an instance.

Worth flagging for anyone porting another contributed/ kernel: simulate_kernel / baremetal / benchmark do still import from the deprecated neuronxcc.nki namespace, which makes them look like a working fallback. They are not. They cannot drive a kernel decorated with the current top-level @nki.jit and fail with AttributeError: 'Kernel' object has no attribute 'grid', because they expect the older TraceKernel.

2. nl.affine_range -> builtin range (closes the thread on #129)

@aws-mattmcm asked on #129 whether new content should use range instead of affine_range, and @mrkcath-aws merged as-is and left it to me. On 0.6.0 the answer is yes.

nl.affine_range, nl.sequential_range and nl.static_range are all deprecated, all three docstrings say to prefer range() directly, and their implementations are identical in both the tracer and the simulator backends:

def affine_range(start, stop, step):     return range(start, stop, step)
def sequential_range(start, stop, step): return range(start, stop, step)
def static_range(start, stop, step):     return range(start, stop, step)

nl.range is not the migration target either: it is the same function object as nl.affine_range and carries the same deprecation. So both loops become plain range, and the dependency reasoning moves into comments, since the iterator name no longer records it. The KV-head loop's iterations are genuinely independent (each gets its own softmax state and touches disjoint slices of q and out), and the KV-tile loop genuinely carries (m, l, acc) and relies on in-order unrolling.

Validation

Run on inf2.xlarge (NeuronCore-v2, gen2), NKI 0.6.0, fp32, before and after the range change, with identical results:

$ python contributed/decode_attention.py --backend baremetal
[check_correct]     baremetal float32  d=128 seqlen_kv=128            max|diff|=7.302e-07  PASS
[check_correct_gqa] baremetal float32  d=128 seqlen_kv=512 group=4    max|diff|=1.460e-06  PASS

$ python contributed/decode_attention.py --backend simulate
[check_correct]     simulate  float32  d=128 seqlen_kv=128            max|diff|=5.960e-08  PASS
[check_correct_gqa] simulate  float32  d=128 seqlen_kv=512 group=4    max|diff|=1.192e-07  PASS

Also clean under python -W error::DeprecationWarning.

This is the first validation of these kernels on Inferentia2; #129 was reviewed on Trainium2. That distinction caught one real thing, now gated rather than silently broken: bf16 inputs do not compile on gen2, because nl.matmul has no dtype parameter and infers its PSUM destination dtype from the operands, and the gen2 tensor engine rejects a non-fp32 matmul destination:

error: assertion failed: nc_matmul dst dtype must be float32 on gen2, got bfloat16

trn2 is gen3 and appears to permit it, which is why a Trn2 review would not surface it. The file now skips bf16 with an explanatory note on gen2 instead of failing to compile. Fixing it properly means explicit fp32 PSUM tiles plus nisa.nc_matmul at all four matmul sites, which changes the kernels, so it belongs in its own PR.

Per the README, contributed/ is experimental, and the baremetal / p99 / E2E requirements in CONTRIBUTING.md sit under "Requirements for Kernels Targeting src/reference/", so I have followed the in-file check_correct() / main() house style of the other contributed/ kernels rather than adding a test_*.py.

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

[mrkcath-aws]

Hi @varuntej07, can we split this into two PRs? This one with the decode_attention.py port and range change, as it is self-contained, correct, and uses public APIs. The solution for decode_attention_benchmark.py relies on APIs that are not part of the public contract and this will likely break again with the next release.

[varuntej07]

Thanks, that's fair and I've split it. This PR is now just the decode_attention.py port and the range change. I've dropped decode_attention_benchmark.py, the two .gitignore entries for its NEFF artifacts.

Before I open a second PR, I'd rather ask what you'd actually accept, because I couldn't find a supported route. There's no public timing API in 0.6.0, and src/ doesn't benchmark either: the attention_fwd_performance tutorial that #131 rewrote is correctness-only, and the one remaining reference to neuron-profile is stale now that it's been replaced by neuron-explorer. Timing a plain kernel(*args) call isn't an option because the standalone path re-runs the frontend on every invocation, so it measures the compiler (~1.5 s per call on Inf2 against tens of microseconds of actual kernel time).

Is there a supported way to get latency for a contributed/ kernel that I should be using instead? or would you rather this stay out of the repo entirely? Happy with either answer.

Separately, the measurement did turn up something worth recording regardless of whether the tool itself lands. Decode attention is usually assumed memory-bound, but on Inf2 this kernel is Vector-bound: it reaches 4.7% of the 410 GB/s per-core HBM peak at best, estimated_utilization puts Vector at 41.5% against DMA at 5.5%, and cost tracks KV tiles rather than bytes (~7.3 us per 128-wide tile, against the ~0.32 us that tile's 128 KiB would take at peak bandwidth). The likely cause is the online softmax running on [group, TILE_KV] tiles, so at Hkv=2 that's 4 rows on a 128-wide partition axis. That makes softmax partition occupancy, not the loads, the thing to target for the planned split-KV work. I can open an issue with the numbers if that's useful to have written down somewhere. Looking forward!