feat: Add DeepSeek-V4 CSA attention kernels

[aws-zifan-he]

Summary

This change adds the whole DeepSeek-V4 CSA attention block as NKI kernels: it takes the raw hidden state of a new token and returns the projected block output, running on one Trainium3 chip as 4 tensor-parallel ranks of 2 logical NeuronCores each (whole single device). One @nki.jit launch covers the lightning-indexer scoring, the top-k selection, the sparse attention and the output inverse RoPE, so neither the score row nor the selected-index array ever returns to the host. Over various sequence length, the kernel achieves 4x-170x speedup over the naive NKI baseline.

Highlights

• Full device solution with collectives, instead of single LNC kernel
• Mega-Kernel fusion between sparse attention and lightning indexer
• Sequence parallelism across LNCs to mitigate the insufficient parallelism on the head dim
• Parallel input preparation for top-k and D2D communication through nisa.sendrecv.
• DMA traffic shaping for TRN3

How the algorithm works

Standard attention compares each new token against every past token, so one decode step costs O(S) and the KV cache grows with the conversation. At a 128K context it is the cache, not the weights, that limits how many users one device can serve. CSA replaces the full comparison with a selection: the model keeps a compressed KV cache, a small network scores every compressed position, and the attention reads only the 1024 highest-scoring positions plus a local window of 128. Neither count changes with the context length, so the attention's cost is constant.

step	operation	engines
1. Compressor	fold 4 tokens into 1 with a gated weighted sum, once every 4 decode steps	Tensor, Vector
2. Lightning indexer	project indexer query and key into 128 channels; score all compressed positions with a multi-query product, a ReLU and a per-head weight	Tensor, Scalar
3. Top-k selection	select the 1024 highest-scoring compressed positions	GpSimd
4. Concatenate and attend	gather the selected KV rows, add the sliding window, run the multi-query attention and the output projection	Tensor, Vector, DMA

The 128 query heads are head-parallel across the 4 ranks, so each rank owns 32 query heads and 4 output-projection groups; one torch_neuronx.trace of one rank emits one NEFF holding the projections, the sparse attention and the cross-rank all_reduce. Decode issues three launches per step:

launch	LNC grid	work
nki_qkv_rms_rope_kernel	[1]	RMSNorm and RoPE for the query heads and the new KV token, on one packed tile
nki_indexer_qproj_gemv	[2]	the indexer query projection, as a hand-tiled GEMV
nki_indexer_score_topk_gather_2core	[2]	steps 2, 3 and 4 in one launch

The evaluated model is the production DeepSeek-V4-Pro-Max configuration.

Parallelism

An important question when deploying a kernel to multiple cores is the dimension of distribution: how should we split the work across cores? Since we know that communication between chips is more expensive than communication inside a chip, our strategy is splitting works that does not require synchronizations across chip, and splitting works that need synchronization between two NeuronCores inside a chip. Hence:

• We split the heads across the 4 chips, since the attention heads compute the partial results independently.
• If we further split the heads across the two NeuronCores in a chip, every core will only have 16 heads. Although the workload is half than processing 32 heads, the latency will not change as the parallelization along the head dimension is also halved and the tensor engine cannot be well utilized. As a result, we choose to split across sequence between two cores in a chip. This will require synchronization frequently (e.g., when computing softmax, we need to sync and compute maximum and sum of exponentials).

The top-k on the GPSIMD engine

The rotational TopK in NKI Library employs max8 and nc_find_index8 to perform repetitive Top-8 operations and generalize to arbitrary TopK. This is very inefficient and bottleneck the vector engine to proceed the attention computations. Here comes to the key component of our optimization: we move the selection itself onto GPSIMD and fuse the TopK with multi-query index score computation in the indexer for further execution overlapping.

• nisa.topk performs generic TopK operation in GPSIMD by consuming the 16 partition of data in SBUF, stored in the snake format: score 0-15 in column 0 of the free dimension, 16-31 in column 1, etc. The instruction restrict the sequence with a length of 64K. Empirically, we found that the instruction is unstable with sequence longer than 8K and can leads to garbage output.
• Fuse the indexer score computation with top-k inside the loop, for pipelining. Since we need to handle sequence of index scores longer than 8K, we segment the computation of indexer score in 8K-long chunks and perform TopK iteratively. In this case, the TopK of previous chunk can overlap with the indexer score computation in the next chunk. At the end, we will perform a final TopK to aggregate the intermediate TopK together.

Another optimization we did along the way is reduce the DMA of KV cache. In DeepSeek V4, K and V are the same embedding. Instead of fetching the KV entries where one is transposed and one is not, we read from HBM once for K and V and transpose in SBUF.

After these updates, TopK selection is no longer the bottleneck. From 8K to 128K, the context grows 16x and the latency grows 1.24x, due to the scaling of computing indexer scores. The kernel achieves an overall 4.75x~16x speedup over the previous step.

DMA Traffic Shaping

The nisa.dma_copy API carries a priority. This priority defines how you want to allocate the DMA bandwidth when multiple DMA operations are executed. A lower priority number means a higher priority. In the CSA kernel, we tune the priority based on the criticality and efficient of each DMA access to balance the latency.

priority	tagged transfers
0	the top-k offset loads; the swdge gather; the indexer cache slice; the score store-back that the barrier waits on; the snake reformat; the q-projection weight bursts; the RMS tile assembly
1	the head-batched query; the attention sink; the RMSNorm gain; the indexer query; the output write-back
2	the window K^T and V, which are contiguous and efficient; the RoPE cos and sin of the projection kernel
3	the cos and sin of the output inverse RoPE, which are the last inputs consumed

Performance

Whole-block decode latency (ms), BF16. Each row is a development step, so the table also shows what each mechanism was worth.

development step	8K	16K	32K	64K	128K
Trn2, naive NKI baseline	1.303	1.767	6.106	60.564	out of memory
Trn2, optimized, no GpSimd	2.378	2.843	3.972	6.038	10.162
Trn2, GpSimd top-k + fusion	0.501	0.514	0.520	0.569	0.636
Trn3, the same code, no edit	0.338	0.339	0.339	0.378	0.416
Trn3, fusion + traffic shaping	0.326	0.325	0.337	0.355	0.404
Trn3, + collective	0.377	0.376	0.388	0.406	0.455
Trn3, D2D + megakernel	0.346	0.353	0.356	0.410	0.429

Launching a full block

The block is one PyTorch module. It owns its projection weights and the attention core, computes rank tp_rank's shard, and appends the cross-rank all-reduce as its final op, so the traced block returns the full output:

block = CSADecodeAttentionBlockNKI(config, tp_size=4, tp_rank=r, replica_ranks=[0, 1, 2, 3])

out = block(x,                 # [B, 1, dim]                 hidden state of the new token
            kv_window,         # [B, W, head_dim]            sliding-window KV cache
            kv_compress,       # [B, T_c, head_dim]          compressed KV cache, bf16
            indexer_kv_cache)  # [B, T_c, index_head_dim]
# -> [B, 1, dim]

config is one rank's view of the model, from shard_for_tp(CSAConfigFull(seq_len=...), tp_size): it divides both n_heads and o_groups by tp_size, which leaves group_in = n_heads * head_dim / o_groups at the full model's value — what wo_a expects, since the production ColumnParallelLinear is built from the global counts. A rank is then an ordinary block of its own size, constructed with tp_size=1. Passing replica_ranks=None instead returns the rank-local partial and leaves the caller to sum them.

csa_block.py drives this end to end against the CPU golden, in either of two modes:

# The real topology: one rank per worker, all-reduce merged into each NEFF.
torchrun --nproc_per_node=4 \
    -m nkilib.experimental.deepseek_v4_csa.csa_block --phase decode --seq-len 32768

Under torchrun each worker pins itself to its own two physical NeuronCores, so the 4 ranks occupy 8 cores and run concurrently. The driver reads the base index from NEURON_RT_VISIBLE_CORES or NEURON_RT_NUM_CORES.

---

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

[AakashShetty-aws]

/

[AakashShetty-aws]

[General comment] Can you add profiles with the PR. Decode and Prefill Profiles with few seqlen and some target sharding configs.

[AakashShetty-aws]

[General comment] Needs optimization, vectorize loops, reduce hbm round trips and fuse ops. Would be helpful to have profiles.

[AakashShetty-aws]

can intialize outside the loop

[aws-zifan-he]

There is no latency change when hoisting this buffer initialization

[aws-enidx]

Thank you for your contribution to the NKI Library! Please review our contribution guidelines if you haven’t already, and ensure your change follows our documented best practices. We are routing your PR to an appropriate reviewer and will follow up once it’s been assigned.

[AakashShetty-aws]

This looks like a torch code that run as part of the kernels run which should not live here. These should be written in nki if you want to add it here.

[aws-zifan-he]

What would be the right place to put this file? Having an example/ folder and put the launcher there?

[aws-zifan-he]

Updated. Remain torch reference due to performance regression when switch to NKI.

[AakashShetty-aws]

same comment as csa_block.