Add decode (flash-decoding) attention kernels to contributed/

[varuntej07]

Description of changes:

Adds a new community kernel file, contributed/decode_attention.py, implementing the decode (single-query) step of autoregressive attention. This is the memory-bound complement to the compute-bound prefill kernel in contributed/pipelined_attention.py.

Two kernels are included:

1. decode_attention_fwd - simplest correct version: one head, one KV tile, seqlen_kv <= 128. A single query position (seqlen_q = 1) attends over the full cached K/V, running QK -> scale -> softmax -> PV with the softmax scale applied.

2. decode_attention_gqa_fwd - lifts the length limit with KV tiling and a running online softmax (state carried across tiles), and adds grouped-query attention so query heads sharing a KV head also share its K/V loads (loaded once per group).

Both use the d-on-partition-axis layout from the existing attention tutorials.

IO layouts:

• decode_attention_fwd: q (d, 1), k (d, seqlen_kv), v (d, seqlen_kv) -> o (1, d)
• decode_attention_gqa_fwd: q (d, n_q_heads), k/v (n_kv_heads, d, seqlen_kv) -> o (n_q_heads, d)

Current limits: d <= 128; GQA requires n_q_heads % n_kv_heads == 0 and seqlen_kv % TILE_KV == 0 (no padding yet); single batch element. Split-KV flash-decoding for very long context is noted as planned future work.

This kernel targets contributed/, not src/reference/, so the baremetal /benchmark / integration requirements (which CONTRIBUTING.md scopes to src/reference/) do not apply. Correctness is validated on CPU via nki.simulate_kernel against NumPy reference implementations included in the file (numpy_decode_reference and numpy_decode_gqa_reference, the repeat_kv +per-head softmax oracle). Run python contributed/decode_attention.py to execute both checks; the GQA case exercises 4 KV tiles with group = 4 and asserts agreement within atol = rtol = 1e-2. The kernels are not yet run on Neuron hardware; this is stated in the module docstring WARNING.

[mrkcath-aws]

Thanks for this contribution! I ran both kernels on a Trainium2 instance (with the changes below) and the outputs match the NumPy reference, including the online-softmax GQA path.

Before this merges, can you please make the following updates, which works with the latest version of NKI?

These kernels import neuronxcc.nki, which is deprecated; the current package is the top-level import nki. New samples should use it so they keep working and model the right API for developers.

The change is small. Imports:

# before
import neuronxcc.nki as nki
import neuronxcc.nki.isa as nisa
import neuronxcc.nki.language as nl
  
# after
import nki
import nki.isa as nisa
import nki.language as nl

The only API difference is that nisa.tensor_copy and nisa.tensor_scalar now write into a destination tile instead of returning one:

# tensor_copy — before (return-style)
qk_sbuf = nisa.tensor_copy(qk_psum, dtype=nl.float32)

# tensor_copy — after (dest-style)
qk_sbuf = nl.ndarray(qk_psum.shape, dtype=nl.float32, buffer=nl.sbuf)
nisa.tensor_copy(qk_sbuf, qk_psum)

# tensor_scalar — before
qk_scaled = nisa.tensor_scalar(qk_sbuf, op0=nl.multiply, operand0=softmax_scale)

# tensor_scalar — after
qk_scaled = nl.ndarray(qk_sbuf.shape, dtype=qk_sbuf.dtype, buffer=nl.sbuf)
nisa.tensor_scalar(qk_scaled, qk_sbuf, op0=nl.multiply, operand0=softmax_scale)

The same transform applies to the remaining tensor_copy / tensor_scalar calls in both kernels (the dtype= argument moves onto the destination tile's allocation). Everything else carried over unchanged: nl.matmul(..., transpose_x=True), nl.transpose, nl.max/sum/exp/reciprocal/add, nl.load/store, and all the shapes and the online-softmax loop. Once these changes are made, I should be able to approve the merge.

[mrkcath-aws]

Changes summarized in previous comment

[varuntej07]

Thanks for the review and Trainium2 validation. I’ve updated both kernels to use the top-level nki imports and converted all tensor_copy and tensor_scalar calls to the destination-style API. Ready for another review.