feat: Add DeepSeek-V4 CSA attention kernels

README.md

@@ -61,6 +61,9 @@ More details can be found in the [NKI Library Documentation](https://awsdocs-neu
 | [MXFP8 Attention TKG Kernel](https://github.com/aws-neuron/nki-library/blob/main/src/nkilib_src/nkilib/experimental/attention_mxfp8/attention_mxfp8_tkg.py) | The kernel implements MXFP8 flash decode attention for token generation. |
 | [Sparse Attention Indexer Kernel](https://github.com/aws-neuron/nki-library/blob/main/src/nkilib_src/nkilib/experimental/sparse_attention_indexer/sparse_attention_indexer_mx_bf16score.py) | The kernel implements the DeepSeek sparse attention indexer: MX-quantized Q/K/W projections, a BF16 score matmul, and hardware top-K selection of the most relevant KV positions per query. |
 | [DeepSeek V3.2 MX MLP Kernel](https://github.com/aws-neuron/nki-library/blob/main/src/nkilib_src/nkilib/experimental/deepseekv32_mlp/mlp_deepseek_mx.py) | The kernel implements the DeepSeek V3.2 MLP for shared-experts and first dense layers with MX-prequantized packed block-scale input, auto-selecting hoisted or tiled weights with token or intermediate LNC sharding. |
+| [DeepSeek V4 CSA Decode Kernels](https://github.com/aws-neuron/nki-library/blob/main/src/nkilib_src/nkilib/experimental/deepseek_v4_csa/csa_decode_attention.py) | The kernels implement DeepSeek V4 Compressed Sparse Attention for decode, headlined by a megakernel that fuses the lightning-indexer scoring, GpSimd top-K selection and O(window + k) gathered sparse attention into one 2-core launch, so the attention cost is independent of the context length. |
+| [DeepSeek V4 CSA Prefill Kernels](https://github.com/aws-neuron/nki-library/blob/main/src/nkilib_src/nkilib/experimental/deepseek_v4_csa/csa_prefill_attention.py) | The kernels implement DeepSeek V4 Compressed Sparse Attention for prefill: a fused RMSNorm+RoPE projection tail, the gated-pooling KV compressor, the indexer's bisection top-K selection mask, and mask-predicated sparse attention with a compile-time per-tile causal bound. |
+| [DeepSeek V4 CSA Attention Block](https://github.com/aws-neuron/nki-library/blob/main/src/nkilib_src/nkilib/experimental/deepseek_v4_csa/csa_block.py) | The module composes the CSA kernels into complete prefill and decode attention blocks, head-parallel across tensor-parallel ranks with the cross-rank output all-reduce merged into the traced block. |
 
 ## Integration with the Neuron Compiler
 

pyproject.toml

@@ -126,3 +126,6 @@ dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
 # sparse_attention_indexer numpy torch references use `assert`. The kernel
 # files use kernel_assert instead.
 "src/nkilib_src/nkilib/experimental/sparse_attention_indexer/*_torch.py" = ["S101"]
+# deepseek_v4_csa torch references transcribe the model's own code, `assert`
+# included. The kernel files use kernel_assert instead.
+"src/nkilib_src/nkilib/experimental/deepseek_v4_csa/*_torch.py" = ["S101"]

src/nkilib_src/nkilib/experimental/deepseek_v4_csa/__init__.py

@@ -0,0 +1,103 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""DeepSeek-V4 Compressed Sparse Attention (CSA) kernels for Trainium3.
+
+CSA replaces attention's full comparison against every past token with a
+selection. The model keeps a compressed KV cache, a small "lightning indexer"
+scores every compressed position, and the attention reads only the
+``index_topk`` highest-scoring positions plus a local sliding window. The compute
+cost of the attention body is therefore O(``window_size`` + ``index_topk``) and
+does not grow with the context length -- at a 32K context one decode block takes
+0.337 ms on Trainium3 in BF16.
+
+Layout
+------
+``csa_common``
+    Config dataclasses and the host-side tables (RoPE, window bias) the kernels
+    take as inputs.
+``csa_decode_attention``
+    Decode kernels, headlined by ``nki_indexer_score_topk_gather_2core`` -- the
+    fused megakernel that runs the indexer score, the GpSimd top-k and the O(k)
+    sparse attention in a single ``[2]``-grid launch.
+``csa_prefill_attention``
+    Prefill kernels: the fused RMS+RoPE projection tail, the compressor, the
+    indexer's bisection top-k mask, and the two sparse-attention variants.
+``csa_tp_all_reduce``
+    The 2-LNC ``ncc.all_reduce`` that sums the head-parallel output partials
+    across tensor-parallel ranks.
+``csa_block``
+    The composition layer: complete prefill and decode attention blocks, plus a
+    runnable driver that grades them against ``csa_block_torch``.
+
+These kernels use ``priority=`` DMA class-of-service hints, which are
+NeuronCore-v4 only, so they target trn3.
+
+Each module's own docstring carries the design rationale for what it holds: why the
+sequence rather than the head axis is split below the rank boundary, why
+``name=`` on a ``shared_hbm`` allocation is load-bearing on a ``[2]``-grid kernel,
+and how the snake layout that ``nisa.topk`` requires is assembled.
+"""
+
+from .csa_common import (
+    CSAConfig,
+    CSAConfigFull,
+    precompute_freqs_cos_sin,
+    precompute_win_bias_parts,
+    shard_for_tp,
+)
+from .csa_decode_attention import (
+    nisa_topk_snake_kernel,
+    nki_decode_gather_ok_kernel,
+    nki_indexer_qproj_gemv,
+    nki_indexer_score_2core,
+    nki_indexer_score_kernel,
+    nki_indexer_score_topk_2core,
+    nki_indexer_score_topk_gather_2core,
+    nki_indexer_score_topk_kernel,
+    nki_qkv_rms_rope_kernel,
+)
+from .csa_prefill_attention import (
+    nki_compressor_core_kernel,
+    nki_fused_csa_attn_kernel,
+    nki_gather_csa_attn_kernel,
+    nki_indexer_score_mask_kernel,
+    nki_rms_rope_kernel,
+)
+from .csa_tp_all_reduce import TPAllReduceNKI, nki_tp_all_reduce_kernel, tp_all_reduce
+
+__all__ = [
+    "CSAConfig",
+    "CSAConfigFull",
+    "TPAllReduceNKI",
+    "nisa_topk_snake_kernel",
+    "nki_compressor_core_kernel",
+    "nki_decode_gather_ok_kernel",
+    "nki_fused_csa_attn_kernel",
+    "nki_gather_csa_attn_kernel",
+    "nki_indexer_qproj_gemv",
+    "nki_indexer_score_2core",
+    "nki_indexer_score_kernel",
+    "nki_indexer_score_mask_kernel",
+    "nki_indexer_score_topk_2core",
+    "nki_indexer_score_topk_gather_2core",
+    "nki_indexer_score_topk_kernel",
+    "nki_qkv_rms_rope_kernel",
+    "nki_rms_rope_kernel",
+    "nki_tp_all_reduce_kernel",
+    "precompute_freqs_cos_sin",
+    "precompute_win_bias_parts",
+    "shard_for_tp",
+    "tp_all_reduce",
+]