diff --git a/README.md b/README.md index 3a065aa..b11705f 100644 --- a/README.md +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 9e5f793..8935835 100644 --- a/pyproject.toml +++ b/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"] diff --git a/src/nkilib_src/nkilib/experimental/deepseek_v4_csa/__init__.py b/src/nkilib_src/nkilib/experimental/deepseek_v4_csa/__init__.py new file mode 100644 index 0000000..925f530 --- /dev/null +++ b/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", +] diff --git a/src/nkilib_src/nkilib/experimental/deepseek_v4_csa/csa_block.py b/src/nkilib_src/nkilib/experimental/deepseek_v4_csa/csa_block.py new file mode 100644 index 0000000..0ce5c1a --- /dev/null +++ b/src/nkilib_src/nkilib/experimental/deepseek_v4_csa/csa_block.py @@ -0,0 +1,1891 @@ +# 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. + +"""Whole DeepSeek-V4 CSA attention blocks -- the composition layer over the kernels. + +This is where the kernels in ``csa_prefill_attention`` / ``csa_decode_attention`` +and the collective in ``csa_tp_all_reduce`` are wired into complete attention +blocks that take a raw hidden state and return the projected block output. It also +holds the host-side glue the kernels need (index reformatting, mask assembly, +tensor layout changes) and a ``main()`` that traces a block and checks it against +the CPU reference in ``csa_block_torch``. + +Topology +-------- +One chip holds the model as ``tp_size`` HEAD-PARALLEL tensor-parallel ranks, so +each rank owns ``n_heads / tp_size`` query heads and the matching output +projection groups. Each rank traces at ``--logical-nc-config=2`` and its +``ncc.all_reduce`` sums the RowParallelLinear partials across ranks, so one +``torch_neuronx.trace`` of one rank emits ONE NEFF holding the projections, the +sparse attention and the cross-rank reduction. + +Decode issues three ``@nki.jit`` launches per step: + +======================================== ======== ========================================= +launch grid work +======================================== ======== ========================================= +``nki_qkv_rms_rope_kernel`` ``[1]`` RMSNorm + 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]`` indexer score, top-k, and the O(k) + sparse attention, fused +======================================== ======== ========================================= + +Prefill instead calls the RMS+RoPE kernel three times (q, kv, output de-RoPE) +around the compressor, indexer and the two sparse-attention kernels. + +""" + +import os + +import torch +import torch.nn.functional as F +from torch import nn + +from .csa_common import ( + CSAConfig, + RMSNorm, + apply_rotary_emb_functional, + get_hadamard_matrix, + hadamard_transform, + precompute_freqs_cos_sin, + precompute_win_bias_parts, +) +from .csa_decode_attention import ( + NISA_TOPK_GROUP_SIZE, + NISA_TOPK_GROUPS_PER_CALL, + NISA_TOPK_PARTITIONS, + nisa_topk_snake_kernel, + nki_decode_gather_ok_kernel, + nki_indexer_qproj_rope_had_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_prefill_sparse_attn_kernel, + nki_prefill_topk_kernel, + nki_qb_rms_rope_kernel, + nki_rms_rope_kernel, +) +from .csa_tp_all_reduce import tp_all_gather_rows, tp_all_reduce + +_SPARSE_PREFILL_MODE = os.environ.get("CSA_SPARSE_PREFILL", "auto") +_SPARSE_MIN_T_C = 4096 +# Proven-safe nisa.topk width; see the note at the padding site below. +_SAFE_TOPK_N = 8192 +_SPARSE_TILE_Q = int(os.environ.get("CSA_SPARSE_TILE_Q", "1024")) + + +def sparse_prefill_q_range(seq_len: int, t_c: int, index_topk: int, ratio: int, tp_rank: int, tp_size: int): + """This rank's contiguous output-row range under SEQUENCE-parallel prefill. + + Queries below ``split_pos = index_topk * ratio`` have fewer causal compressed + positions than ``index_topk``, so they select ALL of them and there is no sparsity + to exploit; that region stays on the dense kernel. Rank 0 owns it, because keeping + every rank's rows CONTIGUOUS lets the driver concatenate the partials instead of + gathering them. + + Returns ``(lo, hi)`` half-open, in query positions. + """ + del t_c, index_topk, ratio + per = seq_len // tp_size + if tp_rank == 0: + return 0, per + share = (seq_len - per) // (tp_size - 1) + lo = per + (tp_rank - 1) * share + return lo, lo + share + + +def _pack_qb_weight(w, heads: int, head_dim: int, pmax: int = 128): + """A ``wq_b`` weight ``[heads * head_dim, R]`` in ``nki_qb_rms_rope_kernel``'s order. + + Returns ``[heads, pmax, R / pmax, head_dim]``: per head, the contraction axis split so + it lands on SBUF partitions, which is what ``nc_matmul`` contracts over. One head is + then one contiguous DMA. Runs on a parameter, so it constant-folds at trace time. + """ + _, r = w.shape + per_head = w.reshape(heads, head_dim, r) + return per_head.permute(0, 2, 1).reshape(heads, r // pmax, pmax, head_dim).permute(0, 2, 1, 3).contiguous() + + +def compress_sharded(compressor, x, start_pos, freqs_cos_sin, tp_shard): + """Compressed KV, computed on this rank's slice only and all-gathered to full. + + ``tp_shard`` is ``(tp_rank, replica_ranks)`` or None. This is the "each rank does a + gather of KV" half of the sequence-parallel design: a query's top-k may select ANY + compressed position, so every rank needs all ``T_c`` of them -- but only 1/tp of them + need to be COMPUTED locally. The compressor is pointwise up to a one-group halo, so + the shards are independent and the concatenation is bit-exact against computing the + whole thing (verified: max_abs 0.0 at seq_len 8192/16384/32768, tp=4). + """ + if tp_shard is None: + return compressor(x, start_pos, freqs_cos_sin) + tp_rank, replica_ranks = tp_shard + world = len(replica_ranks) + if world == 1: + return compressor(x, start_pos, freqs_cos_sin) + + t_c = x.shape[1] // compressor.compress_ratio + if t_c % world != 0: + # An uneven split needs all_gather_v; fall back rather than mis-gather. + return compressor(x, start_pos, freqs_cos_sin) + per = t_c // world + shard = compressor(x, start_pos, freqs_cos_sin, t_range=(tp_rank * per, (tp_rank + 1) * per)) + head_dim = shard.shape[-1] + gathered = tp_all_gather_rows(shard.reshape(per, head_dim), replica_ranks) + return gathered.reshape(1, t_c, head_dim) + + +def _seq_parallel_prefill(phase: str, full_config: CSAConfig) -> bool: + """Does this prefill run shard the SEQUENCE rather than the heads? + + Only when the sparse second half is selected, because that kernel needs all + ``n_heads`` on one core. Dense prefill stays head-parallel and byte-identical. + """ + if phase != "prefill": + return False + if os.environ.get("CSA_SEQ_PARALLEL", "") == "1": + return True + return _use_sparse_prefill(full_config.compressed_len) + + +def _use_sparse_prefill(t_c: int) -> bool: + """Trace-time: is the O(k) sparse second half the cheaper kernel at this T_c?""" + if _SPARSE_PREFILL_MODE == "1": + return True + if _SPARSE_PREFILL_MODE == "0": + return False + return t_c >= _SPARSE_MIN_T_C + + +def encode_snake(scores, n): + """Encode [rows, n] scores into snake layout [rows//8, 128, n//16] for nisa.topk.""" + rows = scores.shape[0] + src_x = n // NISA_TOPK_GROUP_SIZE + num_batches = rows // NISA_TOPK_GROUPS_PER_CALL + # scores[row, j] → snake[row_local*16 + j%16, j//16] + # Reshape [rows, n] → [rows, src_x, 16] → permute → [rows, 16, src_x] + snake = scores.reshape(rows, src_x, NISA_TOPK_GROUP_SIZE).permute(0, 2, 1).contiguous() + # Pack batches of 8 rows into 128 partitions: [num_batches, 128, src_x] + snake = snake.reshape(num_batches, NISA_TOPK_PARTITIONS, src_x) + # Flatten batch dim for kernel: [num_batches * 128, src_x] + return snake.reshape(num_batches * NISA_TOPK_PARTITIONS, src_x) + + +def decode_snake(vals_snake, idxs_snake, rows, k): + """Decode snake layout [rows//8 * 128, k] → [rows, k] values and indices.""" + num_batches = rows // NISA_TOPK_GROUPS_PER_CALL + k_cols = k // NISA_TOPK_GROUP_SIZE + # Reshape to [num_batches, 8, 16, k] and take meaningful columns + vals = vals_snake.reshape(num_batches, NISA_TOPK_GROUPS_PER_CALL, NISA_TOPK_GROUP_SIZE, k) + idxs = idxs_snake.reshape(num_batches, NISA_TOPK_GROUPS_PER_CALL, NISA_TOPK_GROUP_SIZE, k) + # Only first k_cols columns per partition are meaningful + vals = vals[:, :, :, :k_cols] # [num_batches, 8, 16, k_cols] + idxs = idxs[:, :, :, :k_cols] + # Snake decode: result j → partition j%16, column j//16 + # Permute [batch, group, part, col] → [batch, group, col, part] then reshape + vals = vals.permute(0, 1, 3, 2).reshape(rows, k) + idxs = idxs.permute(0, 1, 3, 2).reshape(rows, k) + return vals, idxs + + +def nisa_topk_batched(scores, k, n_cores=2): + """GpSimd top-k over [rows, n] scores -> (values [rows, k], indices [rows, k]). + + Uses nisa.topk (GPSIMD) with snake layout encode/decode. + rows must be divisible by 8. n must be divisible by 16. k must be divisible by 16. + """ + rows, n = scores.shape + scores_bf16 = scores.to(torch.bfloat16) + snake_input = encode_snake(scores_bf16, n) + vals_snake, idxs_snake = nisa_topk_snake_kernel[n_cores](snake_input, k, n) + vals, idxs = decode_snake(vals_snake, idxs_snake, rows, k) + return vals, idxs.int() + + +def prefill_second_half_attention( + selection: torch.Tensor, # [S_q, T_c] 0/-1e9 bias (dense) or [S_q, k] positions (sparse) + q_second: torch.Tensor, # [S_q, n_heads, head_dim] fp16, already softmax-scaled + win_K_T: torch.Tensor, # [head_dim, S_q + W] window K^T, front-padded by W + win_V: torch.Tensor, # [S_q + W, head_dim] window V, front-padded by W + compress_K_T: torch.Tensor, # [head_dim, T_c] compressed K^T + compress_V: torch.Tensor, # [T_c, head_dim] compressed V + win_bias_base: torch.Tensor, # [S_q, W] window causal bias + win_bias_sink_ind: torch.Tensor, # [S_q, W] sink-column indicator + attn_sink: torch.Tensor, # [1, n_heads] per-head sink scalar + s_lo: int, # global position of this block's first query row + ratio: int, # compress_ratio + sparse: bool, # which kernel; a trace-time choice on a compile-time shape + tile_q: int = _SPARSE_TILE_Q, +) -> torch.Tensor: + """Attention for the SCORED half of prefill. Returns [n_heads, S_q, head_dim]. + + One signature for both kernels, because they compute the same thing from the same + operands. What differs is internal and stays internal: + + * ``selection`` is a 0/-1e9 additive bias over all ``T_c`` for the dense kernel and a + list of ``k`` compressed POSITIONS for the sparse one -- the same operand slot, a + different encoding, which is the whole point of the sparse path. + * ``q_second`` arrives QUERY-major for both. The sparse kernel wants exactly that (a + query's heads are then contiguous, so it can ``dma_transpose`` them instead of + paying one DMA descriptor per element); the dense kernel wants head-major and gets + it from the permute below, which is the same permute it used to do at the call site. + * the sparse kernel is launched per query tile because it unrolls over its query + count, so one launch for the whole half would not compile. The dense kernel takes + the half in one launch. Either way the caller gets one ``[n_heads, S_q, head_dim]``. + """ + S_q, n_heads, head_dim = q_second.shape + win = win_K_T.shape[1] - S_q + + if not sparse: + # Head-major Q^T: [head_dim, n_heads * S_q] indexed h * S_q + s. + all_q_T = q_second.permute(2, 1, 0).reshape(head_dim, n_heads * S_q) + out_flat = nki_gather_csa_attn_kernel[2]( + selection.to(torch.bfloat16), + all_q_T, + win_K_T, + win_V, + compress_K_T, + compress_V, + win_bias_base, + win_bias_sink_ind, + attn_sink, + s_lo, + ratio, + ) + return out_flat.reshape(n_heads, S_q, head_dim) + + topk_idx = selection.to(torch.int32) + q_major = q_second.reshape(S_q * n_heads, head_dim) + compress_V_f16 = compress_V.to(torch.float16) + win_K_T_f16 = win_K_T.to(torch.float16) + win_V_f16 = win_V.to(torch.float16) + tiles = [] + for t0 in range(0, S_q, tile_q): + tq = min(tile_q, S_q - t0) + out_t = nki_prefill_sparse_attn_kernel[2]( + topk_idx[t0 : t0 + tq].transpose(0, 1).contiguous(), + q_major[t0 * n_heads : (t0 + tq) * n_heads], + win_K_T_f16[:, t0 : t0 + tq + win], + win_V_f16[t0 : t0 + tq + win], + compress_V_f16, + attn_sink, + ) + tiles.append(out_t.reshape(n_heads, tq, head_dim)) + return tiles[0] if len(tiles) == 1 else torch.cat(tiles, dim=1) + + +def nki_fused_csa_attn( + q, + kv_raw, + kv_compress, + attn_sink, + window_size, + compress_sel_mask, + softmax_scale, + win_bias_base, + win_bias_sink_ind, + split_pos=0, + T_c_first=0, + first_mask=None, +): + """NKI fused attention: window + compressed with online softmax. + + Uses multi-head kernel that processes all heads in a single call per sub-split, + sharing mask, K^T, and V loads across heads to reduce DMA bandwidth. + Window bias is computed inline in the kernel from base + attn_sink * sink_ind. + """ + B, S, n_heads, head_dim = q.shape + T_c = kv_compress.shape[1] + W = window_size + + q_scaled = (q * softmax_scale).to(torch.float16) + q_T = q_scaled.permute(0, 2, 3, 1).reshape(B * n_heads, head_dim, S) + + kv_raw_f16 = kv_raw.to(torch.float16) + kv_raw_bf16 = kv_raw.to(torch.bfloat16) + kv_compress_f16 = kv_compress.to(torch.float16) + kv_compress_bf16 = kv_compress.to(torch.bfloat16) + + kv_raw_K_T = kv_raw_f16.transpose(1, 2) + kv_compress_K_T = kv_compress_f16.transpose(1, 2) + + # attn_sink reshaped to [1, n_heads] for kernel consumption + attn_sink_2d = attn_sink.detach().view(1, n_heads).float().contiguous() + + if split_pos > 0 and split_pos < S and T_c_first > 0 and T_c_first < T_c: + raw_padded_K_T = F.pad(kv_raw_K_T, (W, 0)).reshape(head_dim, S + W) + raw_padded_V = F.pad(kv_raw_bf16, (0, 0, W, 0)).reshape(S + W, head_dim) + compress_K_T_2d = kv_compress_K_T.reshape(head_dim, T_c) + compress_V_2d = kv_compress_bf16.reshape(T_c, head_dim) + + # First half: T_c_first columns of compressed KV + first_mask_2d = first_mask.reshape(split_pos, T_c_first) + first_kt = torch.cat([raw_padded_K_T[:, : split_pos + W], compress_K_T_2d[:, :T_c_first]], dim=1) + first_v = torch.cat([raw_padded_V[: split_pos + W, :], compress_V_2d[:T_c_first, :]], dim=0) + + all_q_T_first = q_T[:, :, :split_pos].permute(1, 0, 2).reshape(head_dim, n_heads * split_pos) + out_first_flat = nki_fused_csa_attn_kernel[2]( + first_mask_2d, + all_q_T_first, + first_kt, + first_v, + win_bias_base[:split_pos], + win_bias_sink_ind[:split_pos], + attn_sink_2d, + ) + out_first_all = out_first_flat.reshape(n_heads, split_pos, head_dim) + + # Second half: full T_c compressed columns + second_mask_2d = compress_sel_mask.reshape(S - split_pos, T_c).to(torch.bfloat16) + second_kt = torch.cat([raw_padded_K_T[:, split_pos:], compress_K_T_2d], dim=1) + second_v = torch.cat([raw_padded_V[split_pos:, :], compress_V_2d], dim=0) + + S_second_len = S - split_pos + all_q_T_second = q_T[:, :, split_pos:].permute(1, 0, 2).reshape(head_dim, n_heads * S_second_len) + out_second_flat = nki_fused_csa_attn_kernel[2]( + second_mask_2d, + all_q_T_second, + second_kt, + second_v, + win_bias_base[split_pos:], + win_bias_sink_ind[split_pos:], + attn_sink_2d, + ) + out_second_all = out_second_flat.reshape(n_heads, S - split_pos, head_dim) + + out_all = torch.cat([out_first_all, out_second_all], dim=1) + out = out_all.permute(1, 0, 2) # [S, n_heads, head_dim] + else: + mask_2d = compress_sel_mask.reshape(S, T_c).to(torch.bfloat16) + all_K_T = F.pad(torch.cat([kv_raw_K_T, kv_compress_K_T], dim=2), (W, 0)) + all_V = F.pad(torch.cat([kv_raw_bf16, kv_compress_bf16], dim=1), (0, 0, W, 0)) + total_free = S + W + T_c + kt_2d = all_K_T.reshape(head_dim, total_free) + v_2d = all_V.reshape(total_free, head_dim) + + all_q_T_full = q_T.permute(1, 0, 2).reshape(head_dim, n_heads * S) + out_flat = nki_fused_csa_attn_kernel[2]( + mask_2d, all_q_T_full, kt_2d, v_2d, win_bias_base, win_bias_sink_ind, attn_sink_2d + ) + out = out_flat.reshape(n_heads, S, head_dim).permute(1, 0, 2) + + return out.unsqueeze(0) # [1, S, n_heads, head_dim] + + +# -------------------------------------------------------------------------- +# Compressor (stateless, functional) +# -------------------------------------------------------------------------- +class CompressorNKI(nn.Module): + def __init__(self, config, head_dim: int = 512, rotate: bool = False, use_nki: bool = True): + super().__init__() + self.dim = config.dim + self.head_dim = head_dim + self.rope_head_dim = config.rope_head_dim + self.compress_ratio = config.compress_ratio + self.overlap = config.compress_ratio == 4 + self.rotate = rotate + self.use_nki = use_nki + coff = 1 + self.overlap + + self.ape = nn.Parameter(torch.empty(config.compress_ratio, coff * self.head_dim, dtype=torch.float32)) + self.wkv = nn.Linear(self.dim, coff * self.head_dim, bias=False, dtype=torch.float32) + self.wgate = nn.Linear(self.dim, coff * self.head_dim, bias=False, dtype=torch.float32) + self.out_dim = coff * self.head_dim + self.norm = RMSNorm(self.head_dim, config.norm_eps) + + def overlap_transform_functional(self, tensor, value=0): + b, s, ratio, _ = tensor.size() + d = self.head_dim + first_half = tensor[..., :d] + second_half = tensor[..., d:] + top = F.pad(first_half[:, :-1], (0, 0, 0, 0, 1, 0), value=value) + return torch.cat([top, second_half], dim=2) + + def _compress_from_kv_score(self, kv_score, seqlen, freqs_cos_sin): + ratio = self.compress_ratio + rd = self.rope_head_dim + + kv = kv_score[..., : self.out_dim] + score = kv_score[..., self.out_dim :] + + remainder = seqlen % ratio + cutoff = seqlen - remainder + + if remainder > 0: + kv = kv[:, :cutoff] + score = score[:, :cutoff] + + kv = kv.unflatten(1, (-1, ratio)) + score_u = score.unflatten(1, (-1, ratio)) + + freqs_cos, freqs_sin = freqs_cos_sin + compress_cos = freqs_cos[:cutoff:ratio] + compress_sin = freqs_sin[:cutoff:ratio] + + if self.use_nki and self.overlap and kv.shape[0] == 1: + # ape is NOT added here: the kernel adds it in fp32, which is what lets the + # operands stay bf16 all the way in. Adding an fp32 parameter to the bf16 + # projection on the host would promote the whole tensor first. + return self._compress_core_nki( + self.overlap_transform_functional(kv, 0), + self.overlap_transform_functional(score_u, -1e9), + compress_cos, + compress_sin, + ) + + score = score_u + self.ape + if self.overlap: + kv = self.overlap_transform_functional(kv, 0) + score = self.overlap_transform_functional(score, -1e9) + + weights = score.softmax(dim=2) + kv = (kv * weights).sum(dim=2) + + kv = self.norm(kv.to(torch.bfloat16)) + + kv_nope = kv[..., :-rd] + kv_rope = apply_rotary_emb_functional(kv[..., -rd:], (compress_cos, compress_sin)) + kv = torch.cat([kv_nope, kv_rope], dim=-1) + + if self.rotate: + kv = hadamard_transform(kv) + + return kv + + def _compress_core_nki(self, kv, score, compress_cos, compress_sin): + """Run the gated-pooling + RMSNorm + RoPE core in a single NKI kernel. + + Args (post overlap_transform): + kv: [1, T_c, ratio2, head_dim] (bf16) + score: [1, T_c, ratio2, head_dim] (bf16), WITHOUT ape + compress_cos/sin: [T_c, rope_head_dim // 2] (fp32) + Returns: + [1, T_c, head_dim] bf16 + """ + T_c = kv.shape[1] + hd = self.head_dim + + # Drop the batch dim and make slot-major contiguous: [T_c, ratio2, head_dim]. + # BF16 at the handoff: halves what crosses HBM into the kernel, which widens on + # load. The fp32 projection accumulator is preserved upstream (see forward), so the + # only rounding here is of values the kernel is about to pool and normalize. + kv8 = kv[0].contiguous().to(torch.bfloat16) + score8 = score[0].contiguous().to(torch.bfloat16) + + # ape in post-overlap SLOT order: slots [0, ratio) took first_half channels, slots + # [ratio, 2*ratio) took second_half, so the two channel halves stack into rows. + ape_slots = torch.cat([self.ape[:, :hd], self.ape[:, hd:]], dim=0).float().contiguous() + + norm_weight = self.norm.weight.detach().view(1, hd).float().contiguous() + + # Repeat each per-pair cos/sin so adjacent channels share the same value: + # cos_rep[t, 2i] = cos_rep[t, 2i+1] = compress_cos[t, i]. + cos_rep = compress_cos.float().repeat_interleave(2, dim=-1).contiguous() + sin_rep = compress_sin.float().repeat_interleave(2, dim=-1).contiguous() + + TILE_P = 128 + num_tiles = (T_c + TILE_P - 1) // TILE_P + n_cores = 2 if (T_c % TILE_P == 0 and num_tiles % 2 == 0) else 1 + had = get_hadamard_matrix(hd, kv8.device, torch.bfloat16) if self.rotate else None + out = nki_compressor_core_kernel[n_cores]( + kv8, score8, norm_weight, cos_rep, sin_rep, float(self.norm.eps), ape_slots, had + ) + return out.unsqueeze(0) + + def forward(self, x, start_pos, freqs_cos_sin, t_range=None): + """Compress ``x`` into compressed cache positions. + + ``t_range=(t0, t1)`` computes only compressed positions ``[t0, t1)``, which is + what sequence-parallel prefill wants: each rank produces its own slice and the + ranks all-gather. It is NOT simply ``x[ratio*t0 : ratio*t1]``, because with + ``overlap`` (compress_ratio == 4) compressed position ``t`` pools the second half + of raw group ``t`` AND the first half of group ``t - 1`` -- see + ``overlap_transform_functional``, which shifts ``first_half`` down by one group. + So the slice carries a ONE-GROUP (``ratio`` raw tokens) halo at the front and the + halo's own compressed position is dropped afterwards. Its ``first_half`` would have + been zero-padded, which is only the right answer at the true sequence start, i.e. + for ``t0 == 0`` -- and there no halo is taken, so the padding is genuine. + """ + bsz, seqlen, _ = x.size() + + if seqlen < self.compress_ratio: + return None + + if t_range is not None: + ratio = self.compress_ratio + t0, t1 = t_range + halo = 1 if t0 > 0 else 0 + lo = ratio * (t0 - halo) + hi = ratio * t1 + freqs_cos, freqs_sin = freqs_cos_sin + shard = self.forward(x[:, lo:hi], start_pos + lo, (freqs_cos[lo:], freqs_sin[lo:])) + return shard if halo == 0 else shard[:, halo:] + + W = torch.cat([self.wkv.weight, self.wgate.weight], dim=0).to(torch.bfloat16) + kv_score = F.linear(x.to(torch.bfloat16), W).float() + return self._compress_from_kv_score(kv_score, seqlen, freqs_cos_sin) + + +# -------------------------------------------------------------------------- +# Indexer (stateless, functional) — with split scoring optimization +# -------------------------------------------------------------------------- +class IndexerNKI(nn.Module): + def __init__(self, config, use_nki: bool = True): + super().__init__() + # Sequence-parallel prefill: (lo, hi) rows this rank scores; None = all rows. + self.q_range = None + # (tp_rank, replica_ranks) when the indexer's compressed KV is sharded + gathered. + self.tp_shard = None + self.dim = config.dim + self.n_heads = config.index_n_heads + self.head_dim = config.index_head_dim + self.rope_head_dim = config.rope_head_dim + self.index_topk = config.index_topk + self.q_lora_rank = config.q_lora_rank + self.compress_ratio = config.compress_ratio + self.softmax_scale = self.head_dim**-0.5 + + self.wq_b = nn.Linear(self.q_lora_rank, self.n_heads * self.head_dim, bias=False, dtype=torch.bfloat16) + self.weights_proj = nn.Linear(self.dim, self.n_heads, bias=False, dtype=torch.bfloat16) + self.weight_scale = self.softmax_scale * (self.n_heads**-0.5) + self.compressor = CompressorNKI(config, self.head_dim, rotate=True, use_nki=use_nki) + + # Precompute causal bias masks for fixed seq_len. + # Collapsed (single-kernel) path: one full causal bias [S_q, T_c] covering + # the scored second half [split_pos:seqlen] against all T_c compressed kv. + seqlen = config.seq_len + ratio = config.compress_ratio + T_c_idx = seqlen // ratio + k = min(config.index_topk, seqlen // ratio) + if k < T_c_idx: + split_pos = k * ratio + + # first_mask: causal mask for positions 0..split_pos-1. Only need [split_pos, k] + # columns since positions beyond k are always -inf for those query rows. + kv_pos_first = torch.arange(k).view(1, -1) + vc_first = (torch.arange(1, split_pos + 1) // ratio).view(-1, 1) + first_mask_buf = torch.where(kv_pos_first < vc_first, torch.zeros(1), torch.tensor(-1e9)).to(torch.bfloat16) + self.register_buffer("first_mask_buf", first_mask_buf, persistent=False) + + # Full causal bias for the scored second half: rows = queries + # [split_pos:seqlen], cols = compressed kv [0:T_c_idx]. + S_q = seqlen - split_pos + kv_pos = torch.arange(T_c_idx).unsqueeze(0) + valid_counts_idx = (torch.arange(split_pos + 1, seqlen + 1) // ratio).unsqueeze(1) + causal_bias = torch.where(kv_pos < valid_counts_idx, torch.zeros(1), torch.tensor(-1e9)) + self.register_buffer("causal_bias_full", causal_bias.float().contiguous(), persistent=False) + # Zero bias used for start_pos != 0 (decode) — no causal masking on scores. + self.register_buffer("zero_bias_full", torch.zeros(S_q, T_c_idx, dtype=torch.float32), persistent=False) + + def forward(self, x, qr, start_pos, offset, freqs_cos_sin): + bsz, seqlen, _ = x.size() + freqs_cos, freqs_sin = freqs_cos_sin + ratio = self.compress_ratio + rd = self.rope_head_dim + end_pos = start_pos + seqlen + _NEG_INF = -1e9 + + T_c_idx = seqlen // ratio + k = min(self.index_topk, end_pos // ratio) + + if k >= T_c_idx: + if start_pos == 0: + kv_pos = torch.arange(T_c_idx, device=x.device).view(1, 1, -1) + valid_counts = (torch.arange(1, seqlen + 1, device=x.device) // ratio).view(1, -1, 1) + mask = torch.where(kv_pos < valid_counts, 0.0, _NEG_INF) + else: + mask = torch.zeros(bsz, seqlen, T_c_idx, device=x.device, dtype=torch.float32) + return None, mask + + split_pos = k * ratio + split_pos = min(split_pos, seqlen) + + # Sequence-parallel: score only THIS rank's slice of the scored region. + lo, hi = self.q_range if self.q_range is not None else (0, seqlen) + s_lo, s_hi = max(lo, split_pos), hi + S_q = s_hi - s_lo + + # --- First half: precomputed causal mask (queries select all valid kv) --- + first_mask = self.first_mask_buf + + if S_q == 0: + return first_mask, None + + # --- Second half: project + RoPE + Hadamard the queries [s_lo:s_hi] --- + seq_cos_second = freqs_cos[start_pos + s_lo : start_pos + s_hi] + seq_sin_second = freqs_sin[start_pos + s_lo : start_pos + s_hi] + + qr_second = qr[:, s_lo:s_hi, :] + q_second = self.wq_b(qr_second) + q_second = q_second.unflatten(-1, (self.n_heads, self.head_dim)) + q_rope = apply_rotary_emb_functional(q_second[..., -rd:], (seq_cos_second, seq_sin_second)) + q_second = torch.cat([q_second[..., :-rd], q_rope], dim=-1) + q_second = hadamard_transform(q_second) # [1, S_q, n_heads, head_dim] bf16 + + indexer_kv = compress_sharded(self.compressor, x, start_pos, freqs_cos_sin, self.tp_shard) + indexer_kv_t = indexer_kv.transpose(1, 2) # [1, head_dim, T_c_idx] + + # Match the ground-truth dtype: weights_proj (bf16) * scale -> bf16, F.linear in bf16. + weights_second = F.linear(x[:, s_lo:s_hi, :], (self.weights_proj.weight * self.weight_scale).to(torch.bfloat16)) + # [1, S_q, n_heads] bf16; widened to fp32 for the kernel's per-head accumulate. + + # --- Stack operands for the single fused NKI kernel --- + # q_T_all[d, h*S_q + s] = q_second[0, s, h, d]: need [head_dim, n_heads, S_q] + # then flatten the last two axes (h outer, s inner) to match the kernel's + # q_global = h * S_q + q_start indexing. permute(2, 1, 0) gives head_dim first. + q_T_all = q_second[0].permute(2, 1, 0).reshape(self.head_dim, self.n_heads * S_q).contiguous() + kv_t_2d = indexer_kv_t[0].contiguous() # [head_dim, T_c_idx] bf16 + weights_2d = weights_second[0].float().contiguous() # [S_q, n_heads] fp32 + + # The bias buffers cover rows [split_pos, seqlen); take this rank's rows. + b0, b1 = s_lo - split_pos, s_hi - split_pos + cbias = (self.causal_bias_full if start_pos == 0 else self.zero_bias_full)[b0:b1].contiguous() + + # Compute scores and bisection-based selection mask in one kernel + TILE_Q = 128 + num_q_tiles = S_q // TILE_Q + n_cores = 2 if (S_q % TILE_Q == 0 and num_q_tiles % 2 == 0) else 1 + if _use_sparse_prefill(kv_t_2d.shape[1]): + scores_2d = nki_indexer_score_kernel[n_cores](q_T_all, kv_t_2d, weights_2d, cbias) + topk_idx = nki_prefill_topk_kernel[2](scores_2d.to(torch.bfloat16), int(k), _SAFE_TOPK_N) + return first_mask, topk_idx.to(torch.int32).unsqueeze(0) + + second_mask_2d = nki_indexer_score_mask_kernel[n_cores](q_T_all, kv_t_2d, weights_2d, cbias, int(k)) + second_mask = second_mask_2d.unsqueeze(0) # [1, S_q, T_c_idx] + + return first_mask, second_mask + + +# -------------------------------------------------------------------------- +# Core Attention Module — NKI version +# -------------------------------------------------------------------------- +class CSAAttentionCoreNKI(nn.Module): + """Core attention with split window/compressed and NKI kernel integration.""" + + def __init__(self, config, use_dense_attn: bool = False, use_nki: bool = True): + super().__init__() + self.config = config + self.n_heads = config.n_heads + self.head_dim = config.head_dim + self.rope_head_dim = config.rope_head_dim + self.window_size = config.window_size + self.compress_ratio = config.compress_ratio + self.softmax_scale = config.head_dim**-0.5 + self.use_dense_attn = use_dense_attn + + self.attn_sink = nn.Parameter(torch.empty(config.n_heads, dtype=torch.float32)) + self.compressor = CompressorNKI(config, config.head_dim, rotate=False, use_nki=use_nki) + self.indexer = IndexerNKI(config, use_nki=use_nki) + # Sequence-parallel prefill: (lo, hi) output rows this rank owns; None = all rows. + self.q_range = None + # First global query row present in the `q` tensor handed to forward(). + self.q_row_base = 0 + # (tp_rank, replica_ranks) when the compressed KV is sharded + all-gathered. + self.tp_shard = None + self.out_head_major = False + + max_seq_len = config.seq_len + freqs_cos, freqs_sin = precompute_freqs_cos_sin( + self.rope_head_dim, + max_seq_len, + config.original_seq_len, + config.compress_rope_theta, + config.rope_factor, + config.beta_fast, + config.beta_slow, + ) + self.register_buffer("freqs_cos", freqs_cos, persistent=False) + self.register_buffer("freqs_sin", freqs_sin, persistent=False) + + win_bias_base, win_bias_sink_ind = precompute_win_bias_parts(max_seq_len, config.window_size) + self.register_buffer("win_bias_base", win_bias_base, persistent=False) + self.register_buffer("win_bias_sink_ind", win_bias_sink_ind, persistent=False) + + def forward(self, q, kv, x, qr, start_pos=0): + bsz, seqlen, _ = x.size() + win = self.window_size + ratio = self.compress_ratio + full_freqs_cs = (self.freqs_cos, self.freqs_sin) + + first_mask, second_mask = self.indexer(x, qr, start_pos, 0, full_freqs_cs) + + kv_compress = compress_sharded(self.compressor, x, start_pos, full_freqs_cs, self.tp_shard) + + T_c_idx = seqlen // ratio + k = min(self.config.index_topk, seqlen // ratio) + + if first_mask is not None: + split_pos = min(k * ratio, seqlen) + T_c_first = split_pos // ratio + # Sequence-parallel: this rank owns output rows [lo, hi). The dense first + # half and the scored second half are each intersected with that range; + # either intersection may be empty. + lo, hi = self.q_range if self.q_range is not None else (0, seqlen) + f_lo, f_hi = lo, min(hi, split_pos) + s_lo, s_hi = max(lo, split_pos), hi + S_q = s_hi - s_lo + + qb = self.q_row_base + n_q_local = q.shape[1] + q_scaled = (q * self.softmax_scale).to(torch.float16) + q_T = q_scaled.permute(0, 2, 3, 1).reshape(bsz * self.n_heads, self.head_dim, n_q_local) + + kv_f16 = kv.to(torch.float16) + kv_bf16 = kv.to(torch.bfloat16) + kv_compress_f16 = kv_compress.to(torch.float16) + kv_compress_bf16 = kv_compress.to(torch.bfloat16) + kv_raw_K_T = kv_f16.transpose(1, 2) + compress_V_2d = kv_compress_bf16.reshape(T_c_idx, self.head_dim) + compress_K_T_2d = kv_compress_f16.transpose(1, 2).reshape(self.head_dim, T_c_idx) + attn_sink_2d = self.attn_sink.detach().view(1, self.n_heads).float().contiguous() + + # The pad/cat operand assembly below is NOT serial wall time: in the profile its + # transfers stream continuously on the static DMA queue underneath the attention + # kernel's region rather than preceding it, so folding the padding into the + # kernels would remove overlapped bandwidth, not latency. + raw_padded_K_T = F.pad(kv_raw_K_T, (win, 0)).reshape(self.head_dim, seqlen + win) + raw_padded_V = F.pad(kv_bf16, (0, 0, win, 0)).reshape(seqlen + win, self.head_dim) + + # First half: mask-based kernel (unchanged) + parts = [] + if f_hi > f_lo: + n_first = f_hi - f_lo + first_mask_2d = first_mask.reshape(split_pos, T_c_first)[f_lo:f_hi].contiguous() + first_kt = torch.cat([raw_padded_K_T[:, f_lo : f_hi + win], compress_K_T_2d[:, :T_c_first]], dim=1) + first_v = torch.cat([raw_padded_V[f_lo : f_hi + win, :], compress_V_2d[:T_c_first, :]], dim=0) + all_q_T_first = ( + q_T[:, :, f_lo - qb : f_hi - qb].permute(1, 0, 2).reshape(self.head_dim, self.n_heads * n_first) + ) + num_q_tiles_first = n_first // 128 + n_cores_first = 2 if (num_q_tiles_first % 2 == 0 and num_q_tiles_first >= 2) else 1 + out_first_flat = nki_fused_csa_attn_kernel[n_cores_first]( + first_mask_2d, + all_q_T_first, + first_kt, + first_v, + self.win_bias_base[f_lo:f_hi], + self.win_bias_sink_ind[f_lo:f_hi], + attn_sink_2d, + ) + parts.append(out_first_flat.reshape(self.n_heads, n_first, self.head_dim)) + + # Second half: static causal-bound sparse attention (global-max softmax + # + sel_bias predication, per-tile compile-time causal chunk bound). + second_win_K_T = raw_padded_K_T[:, s_lo : s_lo + S_q + win] + second_win_V = raw_padded_V[s_lo : s_lo + S_q + win, :] + + if S_q > 0: + out_second_all = prefill_second_half_attention( + second_mask.reshape(S_q, -1), + q_scaled[0, s_lo - qb : s_hi - qb], + second_win_K_T, + second_win_V, + compress_K_T_2d, + compress_V_2d, + self.win_bias_base[s_lo:s_hi], + self.win_bias_sink_ind[s_lo:s_hi], + attn_sink_2d, + int(s_lo), + int(ratio), + sparse=_use_sparse_prefill(T_c_idx), + ) + parts.append(out_second_all) + + out_all = parts[0] if len(parts) == 1 else torch.cat(parts, dim=1) + if self.out_head_major: + return out_all + o = out_all.permute(1, 0, 2).unsqueeze(0) + else: + o = nki_fused_csa_attn( + q, + kv, + kv_compress, + self.attn_sink, + win, + second_mask, + self.softmax_scale, + self.win_bias_base, + self.win_bias_sink_ind, + ) + if self.out_head_major: + # Dense fallback (k >= T_c, i.e. short contexts only): reshape to the + # head-major contract, which costs a copy this path is small enough to pay. + o = o[0].permute(1, 0, 2) + + return o + + +class CSAAttentionXLA(nn.Module): + def __init__(self, config: CSAConfig, replica_ranks=None): + super().__init__() + self.config = config + self.replica_ranks = list(replica_ranks) if replica_ranks is not None else None + # Sequence-parallel prefill: (lo, hi) output rows this rank owns, or None for + # the head-parallel path where every rank produces all rows. + self._q_range = None + self.dim = config.dim + self.n_heads = config.n_heads + self.n_local_heads = config.n_heads + self.q_lora_rank = config.q_lora_rank + self.o_lora_rank = config.o_lora_rank + self.head_dim = config.head_dim + self.rope_head_dim = config.rope_head_dim + self.n_groups = config.o_groups + self.n_local_groups = config.o_groups + self.window_size = config.window_size + self.compress_ratio = config.compress_ratio + self.eps = config.norm_eps + + pdt = torch.bfloat16 + + # Query path + self.wq_a = nn.Linear(self.dim, self.q_lora_rank, bias=False, dtype=pdt) + self.q_norm = RMSNorm(self.q_lora_rank, self.eps) + self.wq_b = nn.Linear(self.q_lora_rank, self.n_heads * self.head_dim, bias=False, dtype=pdt) + + # KV path + self.wkv = nn.Linear(self.dim, self.head_dim, bias=False, dtype=pdt) + self.kv_norm = RMSNorm(self.head_dim, self.eps) + + # Output path (grouped low-rank) + self.group_in = self.n_heads * self.head_dim // self.n_groups + self.wo_a = nn.Linear(self.group_in, self.n_groups * self.o_lora_rank, bias=False, dtype=pdt) + self.wo_b = nn.Linear(self.n_groups * self.o_lora_rank, self.dim, bias=False, dtype=pdt) + + self.core = CSAAttentionCoreNKI(config, use_dense_attn=False, use_nki=True) + # Take the core's native head-major output; the de-RoPE kernel re-lays it out. + self.core.out_head_major = True + + # Precompute RoPE frequencies as real cos/sin (no complex on NeuronX) + freqs_cos, freqs_sin = precompute_freqs_cos_sin( + self.rope_head_dim, + config.seq_len, + config.original_seq_len, + config.compress_rope_theta, + config.rope_factor, + config.beta_fast, + config.beta_slow, + ) + self.register_buffer("freqs_cos", freqs_cos, persistent=False) + self.register_buffer("freqs_sin", freqs_sin, persistent=False) + + def set_q_range(self, q_range) -> None: + """Sequence-parallel prefill: restrict this rank to output rows ``[lo, hi)``. + + Pushed to the attention core and the indexer, which is where the query range + actually reduces work; ``None`` restores the head-parallel behaviour of every + rank producing every row. + """ + self._q_range = q_range + self.core.q_range = q_range + self.core.q_row_base = 0 if q_range is None else q_range[0] + self.core.indexer.q_range = q_range + + def set_tp_shard(self, tp_shard) -> None: + """Shard the COMPRESSED KV across ranks and all-gather it. + + ``tp_shard`` is ``(tp_rank, replica_ranks)``, or None to compute the whole + compressed cache locally on every rank. Pushed to the attention core and the + indexer, which own the two compressor call sites. This is orthogonal to + ``set_q_range``: that shards which QUERIES a rank produces, this shards which + compressed KV positions it COMPUTES -- the KV itself ends up complete on every + rank either way, because the top-k can select any position. + """ + self.core.tp_shard = tp_shard + self.core.indexer.tp_shard = tp_shard + + def forward(self, x: torch.Tensor, start_pos: int = 0): + """ + Args: + x: [B, S, dim] + start_pos: 0 for prefill + Returns: + output: [B, S, dim] + """ + bsz, seqlen, _ = x.size() + # Slice cos/sin for current positions + seq_cos = self.freqs_cos[start_pos : start_pos + seqlen] + seq_sin = self.freqs_sin[start_pos : start_pos + seqlen] + H, D = self.n_local_heads, self.head_dim + x_bf = x.to(torch.bfloat16) + + q_lo, q_hi = self._q_range if self._q_range is not None else (0, seqlen) + n_q = q_hi - q_lo + + cos_qs = seq_cos[q_lo:q_hi].float().contiguous() + sin_qs = seq_sin[q_lo:q_hi].float().contiguous() + cos_s = seq_cos.float().contiguous() + sin_s = seq_sin.float().contiguous() + + # ===== Query Path ===== + # RMSNorm of the q latent in NKI (do_rope=0). The torch module made ~6 passes over + # [S, q_lora_rank] fp32 -- float(), square(), mean(), rsqrt(), two multiplies, cast. + qr_lin = self.wq_a(x_bf) # [B, S, q_lora_rank] + qr = nki_rms_rope_kernel[2]( + qr_lin.reshape(seqlen, self.q_lora_rank).to(torch.bfloat16), + None, + None, + self.q_norm.weight.reshape(1, self.q_lora_rank).float().contiguous(), + self.eps, + do_rms=1, + inverse=0, + do_rope=0, + ).reshape(bsz, seqlen, self.q_lora_rank) + qr_q = qr if self._q_range is None else qr[:, q_lo:q_hi, :] + + q_out = nki_qb_rms_rope_kernel[2]( + qr_q.reshape(n_q, self.q_lora_rank).to(torch.bfloat16), + _pack_qb_weight(self.wq_b.weight, H, D), + cos_qs, + sin_qs, + self.eps, + ) + q = q_out.reshape(bsz, n_q, H, D) + + # ===== KV Path ===== + # Learnable-gain RMSNorm + RoPE, same kernel with gain_in = kv_norm.weight. + kv_lin = self.wkv(x_bf) # [B, S, D] + kv_out = nki_rms_rope_kernel[2]( + kv_lin.reshape(seqlen, D).to(torch.bfloat16), + cos_s, + sin_s, + self.kv_norm.weight.reshape(1, D).float().contiguous(), + self.eps, + do_rms=1, + inverse=0, + ) + kv = kv_out.reshape(bsz, seqlen, D) + + # ===== NKI Attention Core ===== + # Replaces window+compressed index computation, KV compression and + # sparse_attn_xla. Same operands and same [B, S, n_heads, head_dim] out. + o = self.core(q, kv, x_bf, qr, start_pos=start_pos) + + # Sequence-parallel: the core returned only THIS rank's rows, so everything + # downstream (de-RoPE positions, output projection) runs on that row count and + # at that position offset, not on the full sequence. + o_lo, o_hi = self._q_range if self._q_range is not None else (0, seqlen) + n_out = o_hi - o_lo + # Same rows as the q path (both ranges are self._q_range), so the same table. + cos_o, sin_o = cos_qs, sin_qs + + # ===== Output de-RoPE ===== + # Rotation only (do_rms=0) with inverse=1, same fused kernel. The attention core + # hands back HEAD-MAJOR [H, n_out, D] and the output projection wants QUERY-MAJOR + # [n_out, H*D], so the kernel reads one layout and writes the other: the transpose + # rides along in the DMA it was already issuing. The call site used to permute to + # query-major, back to head-major for this kernel, and to query-major again. + o_out = nki_rms_rope_kernel[2]( + o.reshape(H * n_out, D).to(torch.bfloat16), + cos_o, + sin_o, + None, + self.eps, + do_rms=0, + inverse=1, + heads=H, + in_head_major=1, + out_head_major=0, + ) + o = o_out.reshape(bsz, n_out, H * D) + + # ===== Output Projection (grouped low-rank) ===== + G, R, Din = self.n_local_groups, self.o_lora_rank, self.group_in + o = o.reshape(bsz, n_out, G, Din) + rows = bsz * n_out + fused_macs = self.dim * G * R * Din + rows * G * Din * self.dim + unfused_macs = rows * G * Din * R + rows * G * R * self.dim + + _FUSED_WEIGHT_BUDGET = 1 << 27 # 1.34e8 elements + if fused_macs <= unfused_macs and self.dim * G * Din <= _FUSED_WEIGHT_BUDGET: + wo_a = self.wo_a.weight.view(G, R, Din) + wo_b = self.wo_b.weight.view(self.dim, G, R) + wfused = torch.einsum("cgr,grd->cgd", wo_b, wo_a).reshape(self.dim, G * Din) + output = torch.matmul(o.reshape(bsz, n_out, G * Din), wfused.t()) + else: + wo_a = self.wo_a.weight.view(G, R, Din) + lat = torch.einsum("bsgd,grd->bsgr", o, wo_a) # [B, S, G, o_lora] + output = self.wo_b(lat.reshape(bsz, n_out, G * R)) + + # ===== Cross-rank all-reduce (merged): RowParallelLinear sum over ranks ===== + if self.replica_ranks is not None: + return tp_all_reduce(output, self.replica_ranks) + return output + + +# -------------------------------------------------------------------------- +# Decode Indexer — uses bisection mask for large T_c support +# -------------------------------------------------------------------------- +class DecodeIndexerGatheredNKI(nn.Module): + def __init__(self, config): + super().__init__() + self.dim = config.dim + self.n_heads = config.index_n_heads + self.head_dim = config.index_head_dim + self.rope_head_dim = config.rope_head_dim + self.index_topk = config.index_topk + self.q_lora_rank = config.q_lora_rank + self.compress_ratio = config.compress_ratio + self.softmax_scale = self.head_dim**-0.5 + + self.wq_b = nn.Linear(self.q_lora_rank, self.n_heads * self.head_dim, bias=False, dtype=torch.bfloat16) + self.weights_proj = nn.Linear(self.dim, self.n_heads, bias=False, dtype=torch.bfloat16) + self.weight_scale = self.softmax_scale * (self.n_heads**-0.5) + self.compressor = CompressorNKI(config, self.head_dim, rotate=True, use_nki=True) + + def _score_inputs(self, x, qr, start_pos, indexer_kv_cache, freqs_cos_sin): + """Host-side prep shared by `forward` and `fused_single_chunk_inputs`. + + Factored out (with its op sequence VERBATIM) so the fused + score+topk+attention path consumes bit-identical scoring inputs to the + two-launch path it replaces, rather than a re-derivation of them. + + Returns (q_T_all, weights_2d, indexer_kv_t, T_c, k, S_q). + """ + freqs_cos, freqs_sin = freqs_cos_sin + rd = self.rope_head_dim + T_c = indexer_kv_cache.shape[1] + k = min(self.index_topk, T_c) + TILE_Q = 128 + S_q = TILE_Q + + seq_cos = freqs_cos[start_pos : start_pos + 1] + seq_sin = freqs_sin[start_pos : start_pos + 1] + + wT = self.wq_b.weight.t().contiguous().reshape(self.q_lora_rank // 128, 128, self.n_heads * self.head_dim) + qr_2d = qr.reshape(1, self.q_lora_rank) + # The RoPE, the Hadamard rotation and the Q^T replication now happen INSIDE the + # launch that was already doing this projection, so five host ops and both + # intermediates disappear without adding a launch. The kernel returns q_T_all + # directly in the scorer's [head_dim, n_heads * S_q] layout. + q_T_all = nki_indexer_qproj_rope_had_gemv[2]( + wT, + qr_2d, + seq_cos.float().contiguous(), + seq_sin.float().contiguous(), + get_hadamard_matrix(self.head_dim, qr.device, torch.bfloat16), + int(S_q), + ) + + indexer_kv_t = indexer_kv_cache.transpose(1, 2) # [1, head_dim, T_c] + + weights = F.linear(x, (self.weights_proj.weight * self.weight_scale).to(torch.bfloat16)) + + weights_2d = weights[0, 0:1].float().expand(S_q, -1).contiguous() + + return q_T_all, weights_2d, indexer_kv_t, T_c, k, S_q + + IDX_CHUNK = 8192 + SCORE_CHUNK = 512 + SAFE_TOPK_N = 8192 + + def fused_single_chunk_inputs(self, x, qr, start_pos, indexer_kv_cache, freqs_cos_sin, gather_chunk): + """Scoring inputs for the FUSED score+topk+attention launch, or None. + + Returns (q_T_all, kv_t_seg, weights_2d, k, SAFE_TOPK_N) when this decode + step qualifies for `nki_indexer_score_topk_gather_2core[2]` — i.e. when it + already qualified for the merged `nki_indexer_score_topk_2core[2]` scoring + path AND the top-k width divides the attention gather's chunk. Returns None + otherwise, in which case the caller falls back to the unchanged + `forward()` -> `nki_decode_gather_ok_kernel[1]` two-launch pipeline. + """ + T_c = indexer_kv_cache.shape[1] + k = min(self.index_topk, T_c) + num_idx_chunks = (T_c + self.IDX_CHUNK - 1) // self.IDX_CHUNK + if not (num_idx_chunks == 1 and T_c % 128 == 0 and k % 16 == 0): + return None + if not (T_c % 2 == 0 and (T_c // 2) % self.SCORE_CHUNK == 0 and T_c <= self.SAFE_TOPK_N): + return None + if k % gather_chunk != 0: + return None + + q_T_all, weights_2d, indexer_kv_t, T_c, k, _ = self._score_inputs( + x, qr, start_pos, indexer_kv_cache, freqs_cos_sin + ) + kv_t_seg = indexer_kv_t[0, :, 0:T_c].contiguous() # [head_dim, T_c] + return q_T_all, kv_t_seg, weights_2d, k, self.SAFE_TOPK_N + + def forward(self, x, qr, start_pos, indexer_kv_cache, freqs_cos_sin): + """Score all T_c in chunks → bisection per chunk → merge top-k indices.""" + q_T_all, weights_2d, indexer_kv_t, T_c, k, S_q = self._score_inputs( + x, qr, start_pos, indexer_kv_cache, freqs_cos_sin + ) + + IDX_CHUNK = 8192 + num_idx_chunks = (T_c + IDX_CHUNK - 1) // IDX_CHUNK + + if num_idx_chunks == 1 and T_c % 128 == 0 and k % 16 == 0: + kv_t_seg = indexer_kv_t[0, :, 0:T_c].contiguous() # [head_dim, T_c] + SCORE_CHUNK = 512 + SAFE_TOPK_N = 8192 + if T_c % 2 == 0 and (T_c // 2) % SCORE_CHUNK == 0 and T_c <= SAFE_TOPK_N: + topk_idx_hbm = nki_indexer_score_topk_2core[2](q_T_all, kv_t_seg, weights_2d, int(k), SAFE_TOPK_N) + else: + topk_idx_hbm = nki_indexer_score_topk_kernel[1]( + q_T_all, kv_t_seg, weights_2d, int(k) + ) # [TOPK_ROWS, k] uint32 + topk_head = topk_idx_hbm.int() + return topk_head[0:1].contiguous() + + SCORE_CHUNK_2C = 512 + TOPK_ROWS = 8 + tail_len = T_c - (num_idx_chunks - 1) * IDX_CHUNK # IDX_CHUNK if T_c % IDX_CHUNK == 0 + use_2core_score = T_c % 128 == 0 and (T_c // 2) % SCORE_CHUNK_2C == 0 and k % 16 == 0 and tail_len >= k + if use_2core_score: + kv_t_full = indexer_kv_t[0, :, 0:T_c].contiguous() # [head_dim, T_c] + scores_full = nki_indexer_score_2core[2](q_T_all, kv_t_full, weights_2d) # [1, T_c] bf16 + + chunk_rows = [] + for c in range(num_idx_chunks): + seg_start = c * IDX_CHUNK + seg_end = min(seg_start + IDX_CHUNK, T_c) + row = scores_full[0:1, seg_start:seg_end] # [1, seg_len] + seg_len = seg_end - seg_start + if seg_len < IDX_CHUNK: + row = F.pad(row, (0, IDX_CHUNK - seg_len), value=-1e9) + chunk_rows.append(row) + scores_chunks = torch.cat(chunk_rows, dim=0) # [num_idx_chunks, IDX_CHUNK] + # Pad rows up to TOPK_ROWS=8 (nisa.topk needs rows % 8 == 0); only the + # first num_idx_chunks rows are consumed (pad with a copy of row 0). + if num_idx_chunks < TOPK_ROWS: + pad = scores_chunks[0:1].expand(TOPK_ROWS - num_idx_chunks, IDX_CHUNK) + scores_batched = torch.cat([scores_chunks, pad], dim=0).contiguous() + else: + scores_batched = scores_chunks.contiguous() + # Batched Pass-1: ONE launch, each row/group gets its own local top-k. + k_padded = ((k + 15) // 16) * 16 + pv, pl = nisa_topk_batched(scores_batched, k=int(k_padded)) # [8, k_padded] + # Merge candidates from the num_idx_chunks chunks (indices -> global). + merged_scores = torch.cat([pv[s : s + 1, :k] for s in range(num_idx_chunks)], dim=1) # [1, C*k] + merged_indices = torch.cat([pl[s : s + 1, :k] + s * IDX_CHUNK for s in range(num_idx_chunks)], dim=1) + # Pass-2: final top-k over the merged candidates (n = C*k, proven safe). + merge_n = merged_scores.shape[1] + pad_merge = (16 - merge_n % 16) % 16 + if pad_merge > 0: + merged_scores = F.pad(merged_scores, (0, pad_merge), value=-1e9) + merged_indices = F.pad(merged_indices, (0, pad_merge), value=0) + merged_scores_b = merged_scores.expand(TOPK_ROWS, -1).contiguous() + _, ml = nisa_topk_batched(merged_scores_b, k=int(k_padded)) + ml = ml[0:1, :k] # row 0 only (all rows identical in decode) + topk_head = torch.gather(merged_indices, dim=1, index=ml.long()).int() # [1, k] + # S_out=1 (see single-chunk path): kernel reads only column 0. + return topk_head[0:1].contiguous() + + # ---- Fallback (old proven path): per-chunk single-core fp32 scoring ------ + score_chunks = [] + for seg in range(num_idx_chunks): + seg_start = seg * IDX_CHUNK + seg_end = min(seg_start + IDX_CHUNK, T_c) + seg_len = seg_end - seg_start + + kv_t_seg = indexer_kv_t[0, :, seg_start:seg_end].contiguous() # [head_dim, seg_len] + zero_bias_seg = torch.zeros_like(kv_t_seg[0:1]).float().expand(S_q, -1).contiguous() + + scores_seg = nki_indexer_score_kernel[1](q_T_all, kv_t_seg, weights_2d, zero_bias_seg) + score_chunks.append(scores_seg) # [S_q, seg_len] + + TOPK_ROWS = 8 + + # Pass 1: top-k per chunk, collecting (scores, global_indices) + candidate_scores = [] + candidate_indices = [] + for seg_idx, scores_seg in enumerate(score_chunks): + seg_start = seg_idx * IDX_CHUNK + seg_len = scores_seg.shape[1] + seg_k = min(k, seg_len) + scores_head = scores_seg[0:TOPK_ROWS] # identical rows → only need 8 + if seg_len < IDX_CHUNK: + scores_padded = F.pad(scores_head, (0, IDX_CHUNK - seg_len), value=-1e9) + else: + pad_n = (16 - seg_len % 16) % 16 + scores_padded = F.pad(scores_head, (0, pad_n), value=-1e9) if pad_n > 0 else scores_head + # Pad seg_k to multiple of 16 if needed + seg_k_padded = ((seg_k + 15) // 16) * 16 + top_vals, top_local_idx = nisa_topk_batched(scores_padded, k=int(seg_k_padded)) + # Trim to actual seg_k + top_vals = top_vals[:, :seg_k] + top_local_idx = top_local_idx[:, :seg_k] + # Convert local indices to global: add segment offset + top_global_idx = top_local_idx + seg_start + candidate_scores.append(top_vals) # [TOPK_ROWS, seg_k] + candidate_indices.append(top_global_idx) # [TOPK_ROWS, seg_k] + + if num_idx_chunks == 1: + topk_head = candidate_indices[0][:, :k].int() + else: + # Pass 2: merge all candidates and take final top-k + merged_scores = torch.cat(candidate_scores, dim=1) # [TOPK_ROWS, num_chunks * k] + merged_indices = torch.cat(candidate_indices, dim=1) # [TOPK_ROWS, num_chunks * k] + + # nisa.topk on merged set: [TOPK_ROWS, num_chunks*k] → [TOPK_ROWS, k] + merge_n = merged_scores.shape[1] + pad_merge = (16 - merge_n % 16) % 16 + if pad_merge > 0: + merged_scores_padded = F.pad(merged_scores, (0, pad_merge), value=-1e9) + else: + merged_scores_padded = merged_scores + k_padded = ((k + 15) // 16) * 16 + _, merge_local_idx = nisa_topk_batched(merged_scores_padded, k=int(k_padded)) + merge_local_idx = merge_local_idx[:, :k] + # merge_local_idx[s, i] indexes into merged_indices[s, :] → use torch.gather + topk_head = torch.gather(merged_indices, dim=1, index=merge_local_idx.long()).int() + + return topk_head[0:1].contiguous() + + +# -------------------------------------------------------------------------- +# Decode Attention Module — O(k) with on-device gathered K+V +# -------------------------------------------------------------------------- +class CSADecodeAttentionGatheredNKI(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.n_heads = config.n_heads + self.head_dim = config.head_dim + self.rope_head_dim = config.rope_head_dim + self.window_size = config.window_size + self.compress_ratio = config.compress_ratio + self.softmax_scale = config.head_dim**-0.5 + + self.attn_sink = nn.Parameter(torch.empty(config.n_heads, dtype=torch.float32)) + self.compressor = CompressorNKI(config, config.head_dim, rotate=False, use_nki=True) + self.indexer = DecodeIndexerGatheredNKI(config) + + max_seq_len = config.seq_len + freqs_cos, freqs_sin = precompute_freqs_cos_sin( + self.rope_head_dim, + max_seq_len + 1, + config.original_seq_len, + config.compress_rope_theta, + config.rope_factor, + config.beta_fast, + config.beta_slow, + ) + self.register_buffer("freqs_cos", freqs_cos, persistent=False) + self.register_buffer("freqs_sin", freqs_sin, persistent=False) + + def forward(self, q, kv_window, kv_compress, x, qr, indexer_kv_cache): + """O(k) decode: gathered K scoring + gathered V matmul, all on device. + + Attention kernel cost: O(W + k) — independent of T_c. + Indexer cost: O(T_c) — unavoidable (must score all to find top-k). + """ + W = self.window_size + T_c = kv_compress.shape[1] + S = 1 + start_pos = self.config.seq_len + full_freqs_cs = (self.freqs_cos, self.freqs_sin) + + # ---- FUSED indexer-score + top-k + attention (single launch) ---------- + COMP_CHUNK = 128 # the attention kernel's gather chunk (k must divide it) + fused_inputs = self.indexer.fused_single_chunk_inputs( + x, qr, start_pos, indexer_kv_cache, full_freqs_cs, COMP_CHUNK + ) + + # --- Indexer (fallback only): score all T_c → top-k indices [S, k] --- + if fused_inputs is None: + topk_indices = self.indexer(x, qr, start_pos, indexer_kv_cache, full_freqs_cs) + + # --- Prepare Q: replicate to S --- + q_scaled = (q * self.softmax_scale).to(torch.float16) + q_single = q_scaled[0, 0] + all_q_T = ( + q_single.permute(1, 0) + .unsqueeze(2) + .expand(self.head_dim, self.n_heads, S) + .reshape(self.head_dim, self.n_heads * S) + .contiguous() + ) + + WIN_SIZE = W + kv_win_f16 = kv_window.to(torch.float16) + all_K_T_win = kv_win_f16.transpose(1, 2).reshape(self.head_dim, WIN_SIZE) + all_V_win = kv_win_f16.reshape(WIN_SIZE, self.head_dim) + + compress_kv = kv_compress.reshape(T_c, self.head_dim).contiguous() + + attn_sink_2d = self.attn_sink.detach().view(1, self.n_heads).float().contiguous() + + derope_cos = self.freqs_cos[start_pos : start_pos + 1].contiguous() # [1, half_rope] + derope_sin = self.freqs_sin[start_pos : start_pos + 1].contiguous() # [1, half_rope] + if fused_inputs is not None: + idx_q_T_all, idx_kv_t_seg, idx_weights_2d, fused_k, fused_n = fused_inputs + out_flat = nki_indexer_score_topk_gather_2core[2]( + idx_q_T_all, + idx_kv_t_seg, + idx_weights_2d, + int(fused_k), + int(fused_n), + all_q_T, + all_K_T_win, + all_V_win, + compress_kv, + attn_sink_2d, + derope_cos, + derope_sin, + ) + else: + topk_indices_T = topk_indices.t().contiguous() # [k, S] + out_flat = nki_decode_gather_ok_kernel[1]( + topk_indices_T, all_q_T, all_K_T_win, all_V_win, compress_kv, attn_sink_2d, derope_cos, derope_sin + ) + + # --- Extract --- + out_all = out_flat.reshape(self.n_heads, S, self.head_dim) + o = out_all[:, 0, :].unsqueeze(0).unsqueeze(1) + return o + + +class CSADecodeAttentionBlockNKI(nn.Module): + """Complete CSA attention block (compress_ratio=4), decode phase, on NKI. + + Owns the projection weights and a `CSADecodeAttentionGatheredNKI` core + (which itself owns attn_sink, the indexer, and the compressor). The output + projection is sharded for `tp_size`-way tensor parallelism; this module + holds and computes ONLY rank `tp_rank`'s shard. + """ + + def __init__(self, config, tp_size: int = 4, tp_rank: int = 0, replica_ranks=None): + super().__init__() + self.config = config + # None -> return the rank-local partial (host sums the partials). + # list -> append a 2-LNC ncc.all_reduce(op=add) over these ranks as the + # final forward op, so the traced block returns the full output. + self.replica_ranks = list(replica_ranks) if replica_ranks is not None else None + self.dim = config.dim + self.n_heads = config.n_heads + self.q_lora_rank = config.q_lora_rank + self.o_lora_rank = config.o_lora_rank + self.head_dim = config.head_dim + self.rope_head_dim = config.rope_head_dim + self.n_groups = config.o_groups + self.window_size = config.window_size + self.eps = config.norm_eps + self.softmax_scale = config.head_dim**-0.5 + + if self.n_groups % tp_size != 0: + raise ValueError(f"o_groups={self.n_groups} must be divisible by tp_size={tp_size}") + self.tp_size = tp_size + self.tp_rank = tp_rank + self.n_local_groups = self.n_groups // tp_size # groups this rank owns + self.group_in = self.n_heads * self.head_dim // self.n_groups # per-group wo_a input width + + pdt = torch.bfloat16 + + # ----- Query projection ----- + self.wq_a = nn.Linear(self.dim, self.q_lora_rank, bias=False, dtype=pdt) + self.q_norm = RMSNorm(self.q_lora_rank, self.eps) + self.wq_b = nn.Linear(self.q_lora_rank, self.n_heads * self.head_dim, bias=False, dtype=pdt) + + # ----- KV projection (for the new decode token) ----- + self.wkv = nn.Linear(self.dim, self.head_dim, bias=False, dtype=pdt) + self.kv_norm = RMSNorm(self.head_dim, self.eps) + + # ----- Output projection (grouped low-rank), rank-local shard ----- + self.wo_a = nn.Linear(self.group_in, self.n_local_groups * self.o_lora_rank, bias=False, dtype=pdt) + self.wo_b = nn.Linear(self.n_local_groups * self.o_lora_rank, self.dim, bias=False, dtype=pdt) + + # ----- Library core (gathered O(k) decode attention) ----- + self.core = CSADecodeAttentionGatheredNKI(config) + + # RoPE tables for the block's q/kv rotation (start_pos == seq_len needs + # index seq_len, so provision seq_len + 1 entries). + freqs_cos, freqs_sin = precompute_freqs_cos_sin( + self.rope_head_dim, + config.seq_len + 1, + config.original_seq_len, + config.compress_rope_theta, + config.rope_factor, + config.beta_fast, + config.beta_slow, + ) + self.register_buffer("freqs_cos", freqs_cos, persistent=False) + self.register_buffer("freqs_sin", freqs_sin, persistent=False) + + # ---- projection helpers ------------------------------------------------- + def _project_qkv(self, x, freqs_cs): + """Fused q-path AND kv-path RMS+RoPE via ONE merged NKI kernel. + + Both projection tails do the identical per-partition op (RMS over the free + axis + RoPE on the last rope_head_dim channels). Packing the n_heads q rows + and the single kv row onto one [n_heads+1, head_dim] partition tile lets a + single @nki.jit launch replace the two separate kernels — dropping a launch + boundary, a shared_hbm alloc, and an HBM round-trip. The q/kv path + difference (kv has a learnable gain, q does not) is a per-partition gain + tile: 1.0 on the q rows (x*1.0==x exact in fp32), kv_norm.weight on the kv + row. Bit-identical to the two-kernel path. + """ + # q-path pre-norm latent (shared with the indexer via qr). + qr = self.q_norm(self.wq_a(x)) + q_lin = self.wq_b(qr) # [1,1,n_heads*head_dim] bf16 + q_2d = q_lin.reshape(self.n_heads, self.head_dim).contiguous() # [n_heads, head_dim] bf16 + # kv-path pre-norm latent (the single decode-token KV). + kv_lin = self.wkv(x) # [1,1,head_dim] bf16 + kv_2d = kv_lin.reshape(1, self.head_dim).contiguous() # [1, head_dim] bf16 + weight = self.kv_norm.weight.reshape(1, self.head_dim) # [1, head_dim] fp32 gain + + cos, sin = freqs_cs + # ONE launch assembles the [n_heads+1, head_dim] tile on-chip and fuses + # RMS(+gain)+RoPE for both paths (q rows + kv row). + out = nki_qkv_rms_rope_kernel(q_2d, kv_2d, weight, cos, sin, self.eps) + + q = out[: self.n_heads].reshape(1, 1, self.n_heads, self.head_dim) + kv = out[self.n_heads : self.n_heads + 1].reshape(1, 1, self.head_dim) + return q, qr, kv + + def _output_projection(self, o, bsz, seqlen): + """Rank-local grouped low-rank output projection. + + o: [B, S, n_heads, head_dim] (post de-RoPE). Consumes only the + n_local_groups groups this rank owns. + + Decode is S=1, so both wo_a and wo_b are memory-bound GEMVs — latency is + dominated by the bf16 weight bytes streamed from HBM. There are two ways to + run the grouped low-rank projection, and which is cheaper depends on the + relation between group_in and o_lora_rank: + + two-step (wo_a then wo_b): reads n_local_groups*o_lora*group_in + + dim*n_local_groups*o_lora bytes + fused (compose wo_b@wo_a): reads dim*n_local_groups*group_in bytes + + """ + o = o.reshape(bsz, seqlen, self.n_groups, self.group_in) + g0 = self.tp_rank * self.n_local_groups + o_local = o[:, :, g0 : g0 + self.n_local_groups, :].contiguous() + G, R, D = self.n_local_groups, self.o_lora_rank, self.group_in + + fused_bytes = self.dim * G * D + twostep_bytes = G * R * D + self.dim * G * R + if fused_bytes <= twostep_bytes: + # Fused single matmul (constant-folded weight compose). Cheaper only + # when group_in <= o_lora_rank (reduced model). + wo_a = self.wo_a.weight.view(G, R, D) + wo_b = self.wo_b.weight.view(self.dim, G, R) + wfused = torch.einsum("cgr,grd->cgd", wo_b, wo_a).reshape(self.dim, G * D) + out_partial = torch.matmul(o_local.reshape(bsz, seqlen, G * D).to(torch.bfloat16), wfused.t()) + else: + wo_a = self.wo_a.weight.view(G, R, D) + lat = torch.einsum("bsgd,grd->bsgr", o_local.to(torch.bfloat16), wo_a) # [B,S,G,o_lora] + out_partial = self.wo_b(lat.reshape(bsz, seqlen, G * R)) # [B,S,dim] + + return out_partial + + # ---- decode forward ----------------------------------------------------- + def forward(self, x, kv_window, kv_compress, indexer_kv_cache): + """Single-token decode over the whole attention block. + + Args: + x: [B, 1, dim] raw hidden state for the new token + kv_window: [B, W, head_dim] window KV cache (post-prefill, pre-decode) + kv_compress: [B, T_c, head_dim] compressed KV cache + indexer_kv_cache: [B, T_c, index_head_dim] + Returns: + [B, 1, dim] — the full all-reduced output when `replica_ranks` was + given at construction, else rank `tp_rank`'s partial (host-sum path). + """ + bsz, seqlen, _ = x.size() + if seqlen != 1: + raise ValueError(f"Decode expects seqlen=1, got {seqlen}") + W = self.window_size + start_pos = self.config.seq_len + + seq_cos = self.freqs_cos[start_pos : start_pos + 1] + seq_sin = self.freqs_sin[start_pos : start_pos + 1] + freqs_cs = (seq_cos, seq_sin) + + # ----- Projections (q latent shared with the indexer via qr) ----- + # ONE merged NKI kernel fuses BOTH the q-path per-head RMS+RoPE and the + # kv-path learnable RMSNorm+RoPE (see _project_qkv). + q, qr, kv = self._project_qkv(x, freqs_cs) # q:[B,1,n_heads,head_dim] kv:[B,1,head_dim] + + # ----- Insert the new token into the window cache at (start_pos % W) ----- + p = start_pos % W + kv_slot = kv.reshape(bsz, 1, self.head_dim).to(kv_window.dtype) + kv_window = torch.cat([kv_window[:, :p], kv_slot, kv_window[:, p + 1 :]], dim=1) + + # ----- Core sparse attention (gathered O(k)) ----- + o = self.core(q, kv_window, kv_compress, x, qr, indexer_kv_cache) + + # ----- Rank-local output projection ----- + partial = self._output_projection(o, bsz, seqlen) # [B,1,dim] rank partial + + # ----- Cross-rank all-reduce (merged): RowParallelLinear sum over ranks ----- + if self.replica_ranks is not None: + return tp_all_reduce(partial, self.replica_ranks) + return partial + + +# ------------------------------------------------------------------------ +# Runnable driver: trace one rank's block and grade it against the CPU golden +# ------------------------------------------------------------------------ +# The blocks above cannot be graded by the kernel integration tests (they mix +# torch projections with NKI launches, and the multi-worker path spans ranks), so +# this is their end-to-end check. Two launch modes: +# +# sequential one process traces each rank in turn on one core-set and sums the +# partials on the host. `replica_ranks=None`, so no collective is +# traced -- this validates the block compute alone. +# distributed `torchrun --nproc_per_node=`; every rank traces its own +# block with the 2-LNC ncc.all_reduce MERGED in as the final op, so +# each rank's single NEFF returns the FULL all-reduced output. This +# is the real configuration, and the only one that exercises the +# collective. +# +# Both grade against the same 128-head CPU golden, so a head-parallel sharding +# mistake shows up as a correctness failure rather than as a plausible number. +_ATOL = 2e-3 + +_WEIGHT_GAIN = {"decode": 0.33, "prefill": 0.2} + + +def _rank_config(full_config: CSAConfig, tp_size: int) -> CSAConfig: + """One rank's self-contained config: ``n_heads`` and ``o_groups`` both divided by ``tp_size``. + + Dividing BOTH keeps ``group_in = n_heads * head_dim / o_groups`` at the full + model's value, which is what ``wo_a`` expects -- the production model's + ``ColumnParallelLinear`` is built from the global head and group counts. A rank + is then an ordinary block of its own size, constructed with ``tp_size=1``. + """ + from .csa_common import shard_for_tp + + return shard_for_tp(full_config, tp_size) + + +def _rank_config_for(phase: str, full_config: CSAConfig, tp_size: int) -> CSAConfig: + """Rank config for `phase`: unsharded under sequence parallelism, else head-sharded.""" + if _seq_parallel_prefill(phase, full_config): + return full_config + return _rank_config(full_config, tp_size) + + +def _load_rank_weights(model: nn.Module, rank_weights: dict) -> None: + """Copy one rank's reference shard into ``model``, reporting keys that found no home.""" + sd = model.state_dict() + missing = [k for k in rank_weights if k not in sd] + for k, v in rank_weights.items(): + if k in sd: + sd[k].copy_(v.to(sd[k].dtype)) + model.load_state_dict(sd, strict=False) + if missing: + print(f" [warn] {len(missing)} reference weights had no match in the NKI block:") + for k in missing[:12]: + print(f" {k}") + + +def _check(out: torch.Tensor, ref: torch.Tensor, label: str) -> bool: + """Report max/mean absolute and RMS-relative error against ``ref``; pass on ``_ATOL``.""" + diff = (out.float() - ref.float()).abs() + max_abs = diff.max().item() + rms_rel = (diff.pow(2).sum().sqrt() / (ref.float().pow(2).sum().sqrt() + 1e-12)).item() + passed = max_abs < _ATOL + print( + f" [{label}] ref std={ref.float().std().item():.4e} max_abs_diff={max_abs:.2e} " + f"mean_abs_diff={diff.mean().item():.2e} rms_rel={rms_rel:.2e}" + ) + print(f" [{label}] [{'PASS' if passed else 'FAIL'}] max_abs {max_abs:.2e} {'<' if passed else '>='} {_ATOL:.0e}") + return passed + + +def _build_reference(phase: str, full_config: CSAConfig, tp_size: int) -> dict: + """Run the CPU golden for ``phase`` and return its reference dict.""" + from .csa_block_torch import ( + generate_decode_block_reference_tp, + generate_prefill_block_reference_tp, + ) + + gen = generate_prefill_block_reference_tp if phase == "prefill" else generate_decode_block_reference_tp + # Sequence-parallel ranks each hold the FULL weights (all heads, all o_groups), so + # the reference is generated unsharded and the ranks differ only in which output + # rows they produce. + ref_tp = 1 if _seq_parallel_prefill(phase, full_config) else tp_size + return gen(full_config, tp_size=ref_tp, weight_gain=_WEIGHT_GAIN[phase]) + + +def _reference_inputs(phase: str, ref: dict) -> tuple: + """The trace inputs for ``phase``, taken from the reference so both see identical data.""" + if phase == "prefill": + return (ref["x"],) + return (ref["x_dec"], ref["kv_window"], ref["kv_compress"], ref["indexer_kv_cache"]) + + +def _trace_rank(phase, full_config, tp_size, tp_rank, ref, inputs, workdir, replica_ranks=None): + """Trace rank ``tp_rank``'s block into one NEFF, loading its reference weight shard. + + The rank is constructed as a self-contained ``tp_size=1`` block of its own + (already divided) size -- see ``_rank_config``. With ``replica_ranks`` set, the + cross-rank ``ncc.all_reduce`` is merged in and the traced block returns the + full output; otherwise it returns this rank's partial. + """ + import torch_neuronx + + seq_par = _seq_parallel_prefill(phase, full_config) + cfg = _rank_config_for(phase, full_config, tp_size) + if phase == "prefill": + model = CSAAttentionXLA(cfg, replica_ranks=replica_ranks) + if seq_par: + model.set_q_range( + sparse_prefill_q_range( + full_config.seq_len, + full_config.compressed_len, + full_config.index_topk, + full_config.compress_ratio, + tp_rank, + tp_size, + ) + ) + + if replica_ranks is not None: + model.set_tp_shard((tp_rank, list(replica_ranks))) + else: + model = CSADecodeAttentionBlockNKI(cfg, tp_size=1, tp_rank=0, replica_ranks=replica_ranks) + if ref is not None: + # Sequence-parallel ranks all load the SAME (full) weight set. + _load_rank_weights(model, ref["per_rank_weights"][0 if seq_par else tp_rank]) + model.eval() + return torch_neuronx.trace(model, inputs, compiler_workdir=workdir) + + +def warm_up(traced, inputs) -> None: + """Execute ``traced`` once and discard the result, before any graded execution.""" + traced(*inputs) + + +def run_sequential(phase: str, full_config: CSAConfig, tp_size: int) -> bool: + """Trace the ranks one at a time on one core-set and sum their partials on the host. + + No collective is traced, so this isolates the block compute: if it passes here + but fails under ``run_distributed``, the collective or the rank topology is at + fault rather than the kernels. + """ + print(f"=== CSA {phase} block, {tp_size}-rank head-parallel, sequential (host-summed) ===") + ref = _build_reference(phase, full_config, tp_size) + print(f" reference: ||sum_r partial - full||_inf = {ref['max_sum_err']:.3e}") + inputs = _reference_inputs(phase, ref) + + partials = [] + for r in range(tp_size): + traced = _trace_rank( + phase, full_config, tp_size, r, ref, inputs, workdir=f"./compiler_workdir_{phase}_block_rank{r}" + ) + warm_up(traced, inputs) + partials.append(traced(*inputs).float()) + print(f" rank {r} traced and run") + + if _seq_parallel_prefill(phase, full_config): + # Sequence-parallel: each rank produced a DISJOINT, contiguous block of output + # rows, in rank order, so the full output is their concatenation -- there is no + # cross-rank reduction to undo. + combined = torch.cat(partials, dim=1) + else: + combined = torch.stack(partials, 0).sum(0) + return _check(combined, ref["ref_output_full"], label=f"{phase}_full") + + +_LNC = 2 +"""Logical NeuronCores each rank runs on. The attention and indexer kernels use both.""" + + +def _pin_this_worker_to_its_cores() -> str | None: + """Give this ``torchrun`` worker its own ``_LNC`` physical cores, so the ranks run concurrently. + + Rank ``r`` takes cores ``[base + r * _LNC, base + r * _LNC + _LNC - 1]``, with + ``base`` read from the inherited ``NEURON_RT_VISIBLE_CORES`` (default 8, per the + launch convention) so 4 ranks fill cores 8 to 15. Without this every worker sees + the same cores and the ranks serialize instead of overlapping. + + If the inherited value ALREADY spans ``world * _LNC`` cores, it is left untouched + and this returns it as-is. That matters when the range is chosen to avoid cores + another job holds: the runtime allocates one logical core per process index out of + the visible set, so re-narrowing each worker to a single pair here would discard + the offset and every rank would fall back to logical cores ``0..world-1`` -- + observed as ``Requested:lnc0..lnc3 Available:0 (cores busy, ret=-16)`` in + ``nrt_allocate_neuron_cores`` while an unrelated job held the low cores. + + MUST run before ``torch_neuronx`` is imported, which is why the driver imports + it lazily inside ``_trace_rank`` rather than at module scope. Returns the pinned + range, or None outside a multi-worker launch. + """ + world = int(os.environ.get("WORLD_SIZE", "1")) + if world <= 1: + return None + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + visible = os.environ.get("NEURON_RT_VISIBLE_CORES", "8") + bounds = visible.split("-") + base = int(bounds[0]) + if len(bounds) == 2 and int(bounds[1]) - base + 1 >= world * _LNC: + return visible + first = base + local_rank * _LNC + pinned = f"{first}-{first + _LNC - 1}" + os.environ["NEURON_RT_VISIBLE_CORES"] = pinned + return pinned + + +def run_distributed(phase: str, full_config: CSAConfig, tp_size: int) -> bool: + """One rank per ``torchrun`` worker, each with the all-reduce merged into its NEFF. + + Every rank returns the FULL all-reduced output, so rank 0's result is graded + directly against the 128-head golden. + """ + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + if world != tp_size: + raise ValueError(f"launch with --nproc_per_node={tp_size} (got WORLD_SIZE={world})") + + # Pin FIRST, before anything that can touch the Neuron runtime: the process group + # and the tracer both read NEURON_RT_VISIBLE_CORES when they initialize, so a pin + # applied after them lands too late and every rank shares one core-set. + pinned = _pin_this_worker_to_its_cores() + + import torch.distributed as dist + + # Registers the "xla" process-group backend. Without it init_process_group raises + # `AssertionError: Unknown backend type xla` -- importing torch_xla alone is not + # enough, the backend module itself has to be imported for its side effect. + import torch_xla.distributed.xla_backend # noqa: F401 + + dist.init_process_group(backend="xla", init_method="env://", rank=rank, world_size=world) + + def barrier(): + try: + dist.barrier() + except Exception: + pass + + if rank == 0: + print(f"=== CSA {phase} block, {tp_size}-rank head-parallel x {_LNC} LNC, merged all-reduce ===") + print(f" rank {rank} pinned to physical cores [{pinned}]") + + # Every rank builds the reference itself: it is deterministic, so this is + # cheaper and simpler than broadcasting it, and it keeps the ranks independent. + ref = _build_reference(phase, full_config, tp_size) + if rank == 0: + print(f" reference: ||sum_r partial - full||_inf = {ref['max_sum_err']:.3e}") + inputs = _reference_inputs(phase, ref) + + traced = _trace_rank( + phase, + full_config, + tp_size, + rank, + ref, + inputs, + workdir=f"./compiler_workdir_{phase}_block_rank{rank}", + replica_ranks=list(range(tp_size)), + ) + barrier() + # Every rank warms before any rank grades: the collective is a rendezvous, so a + # warm-up on one rank has to be matched on all of them. + warm_up(traced, inputs) + barrier() + full_out = traced(*inputs) + barrier() + + passed = True + if rank == 0: + passed = _check(full_out, ref["ref_output_full"], label=f"{phase}_full") + barrier() + dist.destroy_process_group() + return passed + + +def emit_neff(phase: str, full_config: CSAConfig, tp_size: int, tp_rank: int, workdir: str, merged: bool) -> None: + """Trace ONE rank's block into ``workdir`` for profiling, then stop. + + ``merged`` picks WHAT gets profiled. With it on, the cross-rank + ``ncc.all_reduce`` is traced into the block, giving the real end-to-end artifact; + profile it with ``neuron-explorer capture --collectives-workers-per-node=``. + With it off the block is compute-only and captures as a single rank, which is the + lower-variance instrument for A/B-ing kernel changes -- a collective capture adds + hundreds of microseconds of cross-rank rendezvous to the reported total. + + Inputs are random rather than taken from the golden: nothing is graded here, and + building the CPU reference for a long context costs more than the trace does. + """ + import torch_neuronx # noqa: F401 (imported for its side effect on tracing) + + T_c = full_config.compressed_len + torch.manual_seed(200) + if phase == "prefill": + inputs = ((torch.randn(1, full_config.seq_len, full_config.dim) * 0.02).to(torch.bfloat16),) + else: + inputs = ( + (torch.randn(1, 1, full_config.dim) * 0.02).to(torch.bfloat16), + (torch.randn(1, full_config.window_size, full_config.head_dim) * 0.01).to(torch.bfloat16), + (torch.randn(1, T_c, full_config.head_dim) * 0.01).to(torch.bfloat16), + (torch.randn(1, T_c, full_config.index_head_dim) * 0.01).to(torch.bfloat16), + ) + + kind = "merged block + all-reduce" if merged else "compute-only block" + print(f"Emitting rank {tp_rank}/{tp_size} {phase} {kind} NEFF -> {workdir}") + _trace_rank( + phase, + full_config, + tp_size, + tp_rank, + ref=None, + inputs=inputs, + workdir=workdir, + replica_ranks=list(range(tp_size)) if merged else None, + ) + print(f" NEFF at {os.path.join(workdir, 'graph.neff')}") + + +def main(argv=None) -> int: + """Trace a CSA block and grade it against the CPU golden. Returns a process exit code.""" + import argparse + + from .csa_common import CSAConfigFull + + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("--phase", choices=("prefill", "decode"), default="decode") + parser.add_argument("--seq-len", type=int, default=8192) + parser.add_argument("--tp-size", type=int, default=4) + parser.add_argument( + "--sequential", + action="store_true", + help="trace the ranks one at a time and sum on the host, instead of one rank per torchrun worker", + ) + parser.add_argument( + "--emit-neff", metavar="WORKDIR", help="trace one rank's block to WORKDIR for profiling, then exit" + ) + parser.add_argument("--emit-rank", type=int, default=0) + parser.add_argument( + "--emit-merged", action="store_true", help="with --emit-neff, trace the all-reduce into the block" + ) + args = parser.parse_args(argv) + + full_config = CSAConfigFull(seq_len=args.seq_len) + if args.emit_neff: + emit_neff(args.phase, full_config, args.tp_size, args.emit_rank, args.emit_neff, args.emit_merged) + return 0 + + distributed = not args.sequential and "RANK" in os.environ + run = run_distributed if distributed else run_sequential + passed = run(args.phase, full_config, args.tp_size) + print("RESULT: PASSED" if passed else "RESULT: FAILED") + return 0 if passed else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/nkilib_src/nkilib/experimental/deepseek_v4_csa/csa_block_torch.py b/src/nkilib_src/nkilib/experimental/deepseek_v4_csa/csa_block_torch.py new file mode 100644 index 0000000..14cbc00 --- /dev/null +++ b/src/nkilib_src/nkilib/experimental/deepseek_v4_csa/csa_block_torch.py @@ -0,0 +1,964 @@ +# 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. + +"""CPU reference for the whole DeepSeek-V4 CSA attention block. + +A plain-torch transcription of the model's own attention block, following +DeepSeek-V4-Pro's inference code: it selects with ``torch.topk`` rather than +``nisa.topk``, attends densely over the gathered positions, and keeps the caches +as ordinary tensors. It is deliberately written for clarity over speed -- its job +is to be obviously correct so the fused kernels can be graded against it. + +Both phases share one ``CSAAttentionCore``: ``prefill()`` populates the window and +compressed caches over ``S`` positions, and ``forward()`` runs a single decode +step against them. That sharing is what makes the decode reference trustworthy -- +the caches the decode step reads are produced by the same code path that the +prefill reference validates. + +``generate_prefill_block_reference_tp`` / ``generate_decode_block_reference_tp`` +build a full 128-head model, shard its weights head-parallel over ``tp_size`` +ranks, and return each rank's weights alongside the FULL-model golden output. The +NKI block then loads one rank's shard and its partials are summed to compare +against that golden -- which is what makes the comparison a test of the tensor +parallelism as well as of the kernels. + +Used by ``csa_block.main()``. The integration tests grade individual kernels +against the per-kernel references in the ``*_torch`` modules instead. +""" + +import math + +import torch +import torch.nn.functional as F +from torch import nn + +from .csa_common import CSAConfig + + +# -------------------------------------------------------------------------- +# RoPE +# -------------------------------------------------------------------------- +def precompute_freqs_cis(dim, seqlen, original_seq_len, base, factor, beta_fast, beta_slow): + def find_correction_dim(num_rotations, dim, base, max_seq_len): + return dim * math.log(max_seq_len / (num_rotations * 2 * math.pi)) / (2 * math.log(base)) + + def find_correction_range(low_rot, high_rot, dim, base, max_seq_len): + low = math.floor(find_correction_dim(low_rot, dim, base, max_seq_len)) + high = math.ceil(find_correction_dim(high_rot, dim, base, max_seq_len)) + return max(low, 0), min(high, dim - 1) + + def linear_ramp_factor(min_val, max_val, dim): + if min_val == max_val: + max_val += 0.001 + linear_func = (torch.arange(dim, dtype=torch.float32) - min_val) / (max_val - min_val) + return torch.clamp(linear_func, 0, 1) + + freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + if original_seq_len > 0: + low, high = find_correction_range(beta_fast, beta_slow, dim, base, original_seq_len) + smooth = 1 - linear_ramp_factor(low, high, dim // 2) + freqs = freqs / factor * (1 - smooth) + freqs * smooth + + t = torch.arange(seqlen) + freqs = torch.outer(t, freqs) + freqs_cis = torch.polar(torch.ones_like(freqs), freqs) + return freqs_cis + + +def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False) -> torch.Tensor: + """Applies rotary positional embeddings in-place.""" + dtype = x.dtype + x_complex = torch.view_as_complex(x.float().unflatten(-1, (-1, 2))) + if inverse: + freqs_cis = freqs_cis.conj() + if x_complex.ndim == 3: + freqs_cis = freqs_cis.view(1, x_complex.size(1), x_complex.size(-1)) + else: + freqs_cis = freqs_cis.view(1, x_complex.size(1), 1, x_complex.size(-1)) + result = torch.view_as_real(x_complex * freqs_cis).flatten(-2) + x.copy_(result.to(dtype)) + return x + + +# -------------------------------------------------------------------------- +# RMSNorm +# -------------------------------------------------------------------------- +class RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim, dtype=torch.float32)) + + def forward(self, x: torch.Tensor): + dtype = x.dtype + x = x.float() + var = x.square().mean(-1, keepdim=True) + x = x * torch.rsqrt(var + self.eps) + return (self.weight * x).to(dtype) + + +# -------------------------------------------------------------------------- +# Hadamard Transform (for Indexer's rotate_activation) +# -------------------------------------------------------------------------- +def hadamard_transform_cpu(x: torch.Tensor) -> torch.Tensor: + """CPU implementation of Hadamard transform. Works for power-of-2 dims.""" + n = x.shape[-1] + assert n > 0 and (n & (n - 1)) == 0, f"Dim must be power of 2, got {n}" + scale = n**-0.5 + h = x.float() + step = 1 + while step < n: + idx_even = torch.arange(0, n, 2 * step) + for i in range(step): + a = h[..., idx_even + i].clone() + b = h[..., idx_even + step + i].clone() + h[..., idx_even + i] = a + b + h[..., idx_even + step + i] = a - b + step *= 2 + return (h * scale).to(x.dtype) + + +def rotate_activation(x: torch.Tensor) -> torch.Tensor: + return hadamard_transform_cpu(x) + + +# -------------------------------------------------------------------------- +# Sparse Attention (pure PyTorch) +# -------------------------------------------------------------------------- +def sparse_attn_cpu(q, kv, attn_sink, topk_idxs, softmax_scale): + """ + Pure PyTorch sparse attention. + + Args: + q: [B, S, n_heads, head_dim] (bf16) + kv: [B, T_kv, head_dim] (bf16) -- shared KV (K=V in CSA) + attn_sink: [n_heads] (fp32) -- per-head bias for position 0 + topk_idxs: [B, S, topk_count] -- indices into kv dim 1 (-1 = masked) + softmax_scale: float + + Returns: + o: [B, S, n_heads, head_dim] (bf16) + """ + B, S, n_heads, head_dim = q.shape + topk_count = topk_idxs.shape[-1] + + mask = topk_idxs == -1 + safe_idxs = topk_idxs.clamp(min=0) + + batch_idx = torch.arange(B, device=kv.device)[:, None, None].expand(B, S, topk_count) + gathered_kv = kv[batch_idx, safe_idxs] + + scores = torch.einsum("bshd,bstd->bsht", q.float(), gathered_kv.float()) * softmax_scale + + sink_mask = (safe_idxs == 0) & (~mask) + sink_bias = attn_sink[None, None, :, None].expand(B, S, n_heads, topk_count) + scores = scores + sink_bias * sink_mask[:, :, None, :].float() + + scores = scores.masked_fill(mask[:, :, None, :].expand_as(scores), float("-inf")) + + attn_weights = torch.softmax(scores, dim=-1) + attn_weights = attn_weights.masked_fill(mask[:, :, None, :].expand_as(attn_weights), 0.0) + + o = torch.einsum("bsht,bstd->bshd", attn_weights, gathered_kv.float()) + return o.to(torch.bfloat16) + + +# -------------------------------------------------------------------------- +# Index computation +# -------------------------------------------------------------------------- +def get_window_topk_idxs(window_size, bsz, seqlen, start_pos): + if start_pos >= window_size - 1: + start_pos_mod = start_pos % window_size + matrix = torch.cat([torch.arange(start_pos_mod + 1, window_size), torch.arange(0, start_pos_mod + 1)], dim=0) + elif start_pos > 0: + matrix = F.pad(torch.arange(start_pos + 1), (0, window_size - start_pos - 1), value=-1) + else: + base = torch.arange(seqlen).unsqueeze(1) + matrix = (base - window_size + 1).clamp(0) + torch.arange(min(seqlen, window_size)) + matrix = torch.where(matrix > base, -1, matrix) + return matrix.unsqueeze(0).expand(bsz, -1, -1) + + +# -------------------------------------------------------------------------- +# Compressor (compress_ratio=4, with overlap) +# -------------------------------------------------------------------------- +class Compressor(nn.Module): + def __init__(self, config: CSAConfig, head_dim: int = 512, rotate: bool = False): + super().__init__() + self.dim = config.dim + self.head_dim = head_dim + self.rope_head_dim = config.rope_head_dim + self.compress_ratio = config.compress_ratio + self.overlap = config.compress_ratio == 4 + self.rotate = rotate + coff = 1 + self.overlap + + self.ape = nn.Parameter(torch.empty(config.compress_ratio, coff * self.head_dim, dtype=torch.float32)) + self.wkv = nn.Linear(self.dim, coff * self.head_dim, bias=False, dtype=torch.float32) + self.wgate = nn.Linear(self.dim, coff * self.head_dim, bias=False, dtype=torch.float32) + self.norm = RMSNorm(self.head_dim, config.norm_eps) + + self.kv_cache = None + self.freqs_cis = None + + def overlap_transform(self, tensor: torch.Tensor, value=0): + b, s, _, _ = tensor.size() + ratio, d = self.compress_ratio, self.head_dim + new_tensor = tensor.new_full((b, s, 2 * ratio, d), value) + new_tensor[:, :, ratio:] = tensor[:, :, :, d:] + new_tensor[:, 1:, :ratio] = tensor[:, :-1, :, :d] + return new_tensor + + def forward(self, x: torch.Tensor, start_pos: int): + assert self.kv_cache is not None + bsz, seqlen, _ = x.size() + ratio = self.compress_ratio + rd = self.rope_head_dim + + x_float = x.float() + kv = self.wkv(x_float) + score = self.wgate(x_float) + + if start_pos == 0: + should_compress = seqlen >= ratio + remainder = seqlen % ratio + cutoff = seqlen - remainder + + if remainder > 0: + kv = kv[:, :cutoff] + score = score[:, :cutoff] + + kv = kv.unflatten(1, (-1, ratio)) + score = score.unflatten(1, (-1, ratio)) + self.ape + + if self.overlap: + kv = self.overlap_transform(kv, 0) + score = self.overlap_transform(score, float("-inf")) + + kv = (kv * score.softmax(dim=2)).sum(dim=2) + else: + should_compress = (start_pos + 1) % self.compress_ratio == 0 + if not should_compress: + return None + return None + + if not should_compress: + return None + + kv = self.norm(kv.to(torch.bfloat16)) + + if start_pos == 0: + freqs_cis = self.freqs_cis[:cutoff:ratio] + else: + freqs_cis = self.freqs_cis[start_pos + 1 - self.compress_ratio].unsqueeze(0) + apply_rotary_emb(kv[..., -rd:], freqs_cis) + + if self.rotate: + kv = rotate_activation(kv) + + if start_pos == 0: + self.kv_cache[:bsz, : seqlen // ratio] = kv + + return kv + + +# -------------------------------------------------------------------------- +# Indexer (for compress_ratio=4 layers) +# -------------------------------------------------------------------------- +class Indexer(nn.Module): + def __init__(self, config: CSAConfig): + super().__init__() + self.dim = config.dim + self.n_heads = config.index_n_heads + self.n_local_heads = config.index_n_heads + self.head_dim = config.index_head_dim + self.rope_head_dim = config.rope_head_dim + self.index_topk = config.index_topk + self.q_lora_rank = config.q_lora_rank + self.compress_ratio = config.compress_ratio + self.softmax_scale = self.head_dim**-0.5 + + self.wq_b = nn.Linear(self.q_lora_rank, self.n_heads * self.head_dim, bias=False, dtype=torch.bfloat16) + self.weights_proj = nn.Linear(self.dim, self.n_heads, bias=False, dtype=torch.bfloat16) + + self.compressor = Compressor(config, self.head_dim, rotate=True) + + self.kv_cache = None + self.freqs_cis = None + + def forward(self, x: torch.Tensor, qr: torch.Tensor, start_pos: int, offset: int): + bsz, seqlen, _ = x.size() + freqs_cis = self.freqs_cis[start_pos : start_pos + seqlen] + ratio = self.compress_ratio + rd = self.rope_head_dim + end_pos = start_pos + seqlen + + if self.compressor.kv_cache is None: + self.compressor.kv_cache = self.kv_cache + self.compressor.freqs_cis = self.freqs_cis + + q = self.wq_b(qr) + q = q.unflatten(-1, (self.n_local_heads, self.head_dim)) + apply_rotary_emb(q[..., -rd:], freqs_cis) + q = rotate_activation(q) + + self.compressor(x, start_pos) + + weights = self.weights_proj(x) * (self.softmax_scale * self.n_heads**-0.5) + + index_score = torch.einsum("bshd,btd->bsht", q, self.kv_cache[:bsz, : end_pos // ratio]) + index_score = (index_score.relu_() * weights.unsqueeze(-1)).sum(dim=2) + + if start_pos == 0: + mask = torch.arange(seqlen // ratio).repeat(seqlen, 1) >= torch.arange(1, seqlen + 1).unsqueeze(1) // ratio + index_score = index_score + torch.where( + mask, float("-inf"), torch.zeros_like(mask, dtype=index_score.dtype) + ) + + k = min(self.index_topk, end_pos // ratio) + topk_idxs = index_score.topk(k, dim=-1)[1] + + if start_pos == 0: + mask = topk_idxs >= torch.arange(1, seqlen + 1).unsqueeze(1) // ratio + topk_idxs = torch.where(mask, torch.tensor(-1), topk_idxs + offset) + else: + topk_idxs = topk_idxs + offset + + return topk_idxs + + +# -------------------------------------------------------------------------- +# Core Attention Module (Steps 11-16) -- Decode +# -------------------------------------------------------------------------- +class CSAAttentionCore(nn.Module): + """ + Core sparse attention for CSA (compress_ratio=4). + + Owns: Compressor, Indexer, attn_sink, freqs_cis, kv_cache. + Does NOT own: wq_a, wq_b, wkv, kv_norm, wo_a, wo_b (projection weights). + """ + + def __init__(self, config: CSAConfig): + super().__init__() + self.config = config + self.n_heads = config.n_heads + self.head_dim = config.head_dim + self.rope_head_dim = config.rope_head_dim + self.window_size = config.window_size + self.compress_ratio = config.compress_ratio + self.softmax_scale = config.head_dim**-0.5 + + self.attn_sink = nn.Parameter(torch.empty(config.n_heads, dtype=torch.float32)) + self.compressor = Compressor(config, config.head_dim, rotate=False) + self.indexer = Indexer(config) + + max_seq_len = config.seq_len + freqs_cis = precompute_freqs_cis( + self.rope_head_dim, + max_seq_len + 1, + config.original_seq_len, + config.compress_rope_theta, + config.rope_factor, + config.beta_fast, + config.beta_slow, + ) + self.register_buffer("freqs_cis", freqs_cis, persistent=False) + + kv_cache_size = config.window_size + max_seq_len // self.compress_ratio + self.register_buffer( + "kv_cache", torch.zeros(config.batch_size, kv_cache_size, self.head_dim, dtype=torch.bfloat16) + ) + + indexer_cache_size = max_seq_len // self.compress_ratio + self.indexer.kv_cache = torch.zeros( + config.batch_size, indexer_cache_size, config.index_head_dim, dtype=torch.bfloat16 + ) + + def prefill(self, q, kv, x, qr): + """ + Run prefill to populate KV caches. Identical to the prefill path in + deepseek_v4_csa_attn_core.py. + + Args: + q: [B, S, n_heads, head_dim] + kv: [B, S, head_dim] + x: [B, S, dim] + qr: [B, S, q_lora_rank] + + Returns: + o: [B, S, n_heads, head_dim] -- prefill attention output + """ + bsz, seqlen, _ = x.size() + win = self.window_size + start_pos = 0 + + if self.compressor.kv_cache is None: + self.compressor.kv_cache = self.kv_cache[:, win:] + self.compressor.freqs_cis = self.freqs_cis + self.indexer.freqs_cis = self.freqs_cis + + # Step 11: Window indices + topk_idxs = get_window_topk_idxs(win, bsz, seqlen, start_pos) + + # Step 13: Indexer + offset = kv.size(1) + compress_topk_idxs = self.indexer(x, qr, start_pos, offset) + + # Step 14: Merge + topk_idxs = torch.cat([topk_idxs, compress_topk_idxs], dim=-1) + topk_idxs = topk_idxs.int() + + # Step 12 & 15: Compress KV and prepare + if seqlen <= win: + self.kv_cache[:bsz, :seqlen] = kv + else: + cutoff = seqlen % win + self.kv_cache[:bsz, cutoff:win], self.kv_cache[:bsz, :cutoff] = kv[:, -win:].split( + [win - cutoff, cutoff], dim=1 + ) + + kv_compress = self.compressor(x, start_pos) + if kv_compress is not None: + kv_full = torch.cat([kv, kv_compress], dim=1) + else: + kv_full = kv + + # Step 16: Sparse attention + o = sparse_attn_cpu(q, kv_full, self.attn_sink, topk_idxs, self.softmax_scale) + return o + + def forward(self, q, kv, x, qr, start_pos): + """ + Decode step: compute attention for a single new token. + + Args: + q: [B, 1, n_heads, head_dim] -- projected query for new token + kv: [B, 1, head_dim] -- projected KV for new token + x: [B, 1, dim] -- raw hidden states for new token + qr: [B, 1, q_lora_rank] -- normalized query latent for new token + start_pos: int -- position of the new token (== prefill seq_len) + + Returns: + o: [B, 1, n_heads, head_dim] -- attention output + """ + bsz, seqlen, _ = x.size() + assert seqlen == 1, f"Decode expects seqlen=1, got {seqlen}" + win = self.window_size + + if self.compressor.kv_cache is None: + self.compressor.kv_cache = self.kv_cache[:, win:] + self.compressor.freqs_cis = self.freqs_cis + self.indexer.freqs_cis = self.freqs_cis + + # Step 11: Window indices + topk_idxs = get_window_topk_idxs(win, bsz, seqlen, start_pos) + + # Step 13: Indexer + offset = win + compress_topk_idxs = self.indexer(x, qr, start_pos, offset) + + # Step 14: Merge + topk_idxs = torch.cat([topk_idxs, compress_topk_idxs], dim=-1) + topk_idxs = topk_idxs.int() + + # Step 12 & 15: Place new KV in window cache and try compress + self.kv_cache[:bsz, start_pos % win] = kv.squeeze(1) + self.compressor(x, start_pos) + + # Step 16: Sparse attention against full cache (window + compressed) + o = sparse_attn_cpu(q, self.kv_cache[:bsz], self.attn_sink, topk_idxs, self.softmax_scale) + return o + + +def _init_block_weights( + model, seed: int = 42, weight_gain: float = 1.0, norm_init: float = 1.0, sink_scale: float = 1.0 +): + """Deterministic init. + + 2-D weights: xavier_uniform(gain=weight_gain). 1-D params default to a + small uniform, EXCEPT RMSNorm weights (initialized near `norm_init`, i.e. + ~identity scaling) and attn_sink (scaled by `sink_scale`). With gain=0.1 and + norm_init~=0 the output decays to ~1e-6 (bf16 noise floor); gain=1.0 + + norm_init=1.0 gives an O(1) output where real numerical error is visible. + """ + torch.manual_seed(seed) + for name, param in model.named_parameters(): + if param.dim() >= 2: + nn.init.xavier_uniform_(param, gain=weight_gain) + elif param.dim() == 1: + if name.endswith("attn_sink"): + nn.init.uniform_(param, -0.1 * sink_scale, 0.1 * sink_scale) + elif "norm" in name: + # RMSNorm gamma near 1.0 (identity) with a small spread. + nn.init.uniform_(param, norm_init - 0.05, norm_init + 0.05) + else: + nn.init.uniform_(param, -0.1, 0.1) + + +# -------------------------------------------------------------------------- +# Full CSA Attention Block (projections + core) -- Prefill +# -------------------------------------------------------------------------- +class CSAAttentionBlockPrefill(nn.Module): + """Complete CSA attention block (compress_ratio=4), prefill phase. + + Owns the projection weights (wq_a, q_norm, wq_b, wkv, kv_norm, wo_a, wo_b) + plus a CSAAttentionCore (attn_sink, Compressor, Indexer, sparse attention, + KV caches). Runs the whole sequence through the prefill path in one forward. + + The output projection (wo_a, wo_b) is sharded for `tp_size`-way tensor + parallelism; this module holds ONLY rank `tp_rank`'s shard and returns that + rank's partial [B, S, dim] contribution (the all-reduce over ranks is the + caller's responsibility). Structurally identical to the decode block's + projection wiring -- only the phase (start_pos=0, full sequence) differs. + """ + + def __init__(self, config: CSAConfig, tp_size: int = 4, tp_rank: int = 0): + super().__init__() + self.config = config + self.dim = config.dim + self.n_heads = config.n_heads + self.q_lora_rank = config.q_lora_rank + self.o_lora_rank = config.o_lora_rank + self.head_dim = config.head_dim + self.rope_head_dim = config.rope_head_dim + self.n_groups = config.o_groups + self.eps = config.norm_eps + + assert self.n_groups % tp_size == 0, "o_groups must be divisible by tp_size" + self.tp_size = tp_size + self.tp_rank = tp_rank + self.n_local_groups = self.n_groups // tp_size # groups owned by this rank + self.group_in = self.n_heads * self.head_dim // self.n_groups # per-group wo_a input width + + pdt = torch.bfloat16 + + # ----- Query projection ----- + self.wq_a = nn.Linear(self.dim, self.q_lora_rank, bias=False, dtype=pdt) + self.q_norm = RMSNorm(self.q_lora_rank, self.eps) + self.wq_b = nn.Linear(self.q_lora_rank, self.n_heads * self.head_dim, bias=False, dtype=pdt) + + # ----- KV projection ----- + self.wkv = nn.Linear(self.dim, self.head_dim, bias=False, dtype=pdt) + self.kv_norm = RMSNorm(self.head_dim, self.eps) + + # ----- Output projection (grouped low-rank), rank-local shard ----- + # Full (world_size=1) wo_a: [n_groups * o_lora_rank, group_in] + # wo_b: [dim, n_groups * o_lora_rank] + # Rank r owns groups [r*n_local_groups : (r+1)*n_local_groups]: + # wo_a shard rows = those groups' o_lora_rank outputs -> [n_local_groups*o_lora_rank, group_in] + # wo_b shard cols = those groups' flattened inputs -> [dim, n_local_groups*o_lora_rank] + self.wo_a = nn.Linear(self.group_in, self.n_local_groups * self.o_lora_rank, bias=False, dtype=pdt) + self.wo_b = nn.Linear(self.n_local_groups * self.o_lora_rank, self.dim, bias=False, dtype=pdt) + + # ----- Core sparse attention (owns caches, attn_sink, compressor, indexer) ----- + self.core = CSAAttentionCore(config) + + # RoPE frequencies (shared with core; used here for q/kv RoPE). + self.register_buffer("freqs_cis", self.core.freqs_cis, persistent=False) + + # ---- shared projection helpers ----------------------------------------- + def _project_q(self, x, freqs_cis): + rd = self.rope_head_dim + qr = self.q_norm(self.wq_a(x)) + q = self.wq_b(qr) + q = q.unflatten(-1, (self.n_heads, self.head_dim)) + q = q * torch.rsqrt(q.square().mean(-1, keepdim=True) + self.eps) + apply_rotary_emb(q[..., -rd:], freqs_cis) + return q, qr + + def _project_kv(self, x, freqs_cis): + rd = self.rope_head_dim + kv = self.wkv(x) + kv = self.kv_norm(kv) + apply_rotary_emb(kv[..., -rd:], freqs_cis) + return kv + + def _output_projection(self, o, bsz, seqlen): + """Rank-local grouped low-rank output projection. + + o: [B, S, n_heads, head_dim] (post de-RoPE). This rank only consumes the + n_local_groups groups it owns. + """ + o = o.reshape(bsz, seqlen, self.n_groups, self.group_in) + g0 = self.tp_rank * self.n_local_groups + o_local = o[:, :, g0 : g0 + self.n_local_groups, :] # [B, S, n_local_groups, group_in] + + wo_a = self.wo_a.weight.view(self.n_local_groups, self.o_lora_rank, self.group_in) + o_local = torch.einsum("bsgd,grd->bsgr", o_local, wo_a) # [B, S, n_local_groups, o_lora_rank] + out_partial = self.wo_b(o_local.flatten(2)) # [B, S, dim] -- rank partial + + return out_partial + + # ---- prefill forward (whole sequence) ---------------------------------- + @torch.no_grad() + def forward(self, x, start_pos: int = 0): + bsz, seqlen, _ = x.size() + rd = self.rope_head_dim + freqs_cis = self.freqs_cis[start_pos : start_pos + seqlen] + + q, qr = self._project_q(x, freqs_cis) + kv = self._project_kv(x, freqs_cis) + + o = self.core.prefill(q, kv, x, qr) # [B, S, n_heads, head_dim] + + # De-rotate RoPE on the output's rope channels. + apply_rotary_emb(o[..., -rd:], freqs_cis, inverse=True) + + return self._output_projection(o, bsz, seqlen) # [B, S, dim] rank partial + + +# -------------------------------------------------------------------------- +# Tensor-parallel reference-data generation (head-parallel over `tp_size` ranks) +# -------------------------------------------------------------------------- +def generate_prefill_block_reference_tp( + full_config, + tp_size: int = 4, + weight_gain: float = 0.2, + norm_init: float = 1.0, + sink_scale: float = 1.0, + input_scale: float = 1.0, +): + """Head-parallel TP decomposition of the FULL prefill model over `tp_size` ranks. + + The production model has full_config.n_heads (=128) query heads and + full_config.o_groups (=16) output-projection groups. Sharded HEAD-PARALLEL: + rank r owns query heads [r*Hl : (r+1)*Hl] (Hl = n_heads/tp_size = 32) which, + since a group tiles n_heads/o_groups (=8) heads, coincide EXACTLY with output + groups [r*Gl : (r+1)*Gl] (Gl = o_groups/tp_size = 4). Each rank is a + self-contained Hl-head / Gl-group block with its OWN weights -- no redundant + compute. + + Builds ONE full 128-head model, runs the prefill forward to the true full + output [B,S,dim] (the all-reduce target), then for each rank derives (a) its + sharded weight state_dict for an Hl-head/Gl-group block and (b) its golden + output-projection PARTIAL. Because prefill attention is per-head (the + indexer's top-k is head-independent and the KV is shared), heads + [r*Hl:(r+1)*Hl] of the full model equal an independent Hl-head block's + output; and because the grouped projection sums over disjoint group blocks, + sum_r partial_r == full output. + + Returns the shared prefill input `x`, ref_output_full, per-rank partials, + and per-rank shard state_dicts. + """ + full = CSAAttentionBlockPrefill(full_config, tp_size=1, tp_rank=0) + _init_block_weights(full, weight_gain=weight_gain, norm_init=norm_init, sink_scale=sink_scale) + full.eval() + + B, S = full_config.batch_size, full_config.seq_len + Hl = full_config.n_heads // tp_size # 32 query heads per rank + Gl = full_config.o_groups // tp_size # 4 output groups per rank + hd = full_config.head_dim + R = full_config.o_lora_rank + group_in = full_config.n_heads * hd // full_config.o_groups # 4096 (full) + + torch.manual_seed(99) + x = (torch.randn(B, S, full_config.dim) * input_scale).to(torch.bfloat16) + rd = full_config.rope_head_dim + with torch.no_grad(): + # Full prefill up to the per-head attention output o (post de-RoPE), then + # the full grouped output projection = the golden all-reduce target. + freqs_cis = full.freqs_cis[0:S] + q, qr = full._project_q(x, freqs_cis) + kv = full._project_kv(x, freqs_cis) + o_full = full.core.prefill(q, kv, x, qr) # [B,S,128,head_dim] + apply_rotary_emb(o_full[..., -rd:], freqs_cis, inverse=True) + ref_output_full = full._output_projection(o_full.clone(), B, S) # [B,S,dim] + + full_sd = { + k: v + for k, v in full.state_dict().items() + if not k.startswith("core.kv_cache") + and not k.startswith("core.freqs_cis") + and not k.startswith("freqs_cis") + and not k.startswith("core.indexer.kv_cache") + } + + per_rank_weights, ref_partials = [], [] + for r in range(tp_size): + h0, h1 = r * Hl, (r + 1) * Hl # this rank's query heads + g0r, g1r = r * Gl * R, (r + 1) * Gl * R # this rank's wo_a rows / wo_b cols + sd_r = {} + for k, v in full_sd.items(): + if k == "wq_b.weight": # [n_heads*hd, q_lora] -> this rank's heads + sd_r[k] = v[h0 * hd : h1 * hd, :].clone() + elif k == "core.attn_sink": # [n_heads] -> this rank's heads + sd_r[k] = v[h0:h1].clone() + elif k == "wo_a.weight": # [o_groups*R, group_in] -> this rank's groups + sd_r[k] = v[g0r:g1r, :].clone() + elif k == "wo_b.weight": # [dim, o_groups*R] -> this rank's group cols + sd_r[k] = v[:, g0r:g1r].clone() + else: # wq_a/q_norm/wkv/kv_norm/indexer/compressor replicated + sd_r[k] = v.clone() + per_rank_weights.append(sd_r) + with torch.no_grad(): + og = o_full[:, :, h0:h1, :].reshape(B, S, Gl, group_in) # [B,S,Gl,group_in] + wo_a_r = sd_r["wo_a.weight"].view(Gl, R, group_in) + lat = torch.einsum("bsgd,grd->bsgr", og, wo_a_r) # [B,S,Gl,R] + partial_r = torch.matmul(lat.reshape(B, S, Gl * R), sd_r["wo_b.weight"].t()) # [B,S,dim] + ref_partials.append(partial_r) + + summed = torch.stack([p.float() for p in ref_partials], 0).sum(0) + max_sum_err = (summed - ref_output_full.float()).abs().max().item() + + return { + "x": x, # prefill input [B,S,dim] (the trace input) + "ref_output_full": ref_output_full, # all-reduce target [B,S,dim] + "ref_partials": ref_partials, # list of tp_size [B,S,dim] partials + "per_rank_weights": per_rank_weights, # list of tp_size shard state_dicts + "max_sum_err": max_sum_err, + "tp_size": tp_size, + } + + +# -------------------------------------------------------------------------- +# Full CSA Attention Block (projections + core) -- Decode +# -------------------------------------------------------------------------- +class CSAAttentionBlockDecode(nn.Module): + """Complete CSA attention block (compress_ratio=4), decode phase. + + Owns the projection weights (wq_a, q_norm, wq_b, wkv, kv_norm, wo_a, wo_b) + plus a CSAAttentionCore (attn_sink, Compressor, Indexer, sparse attention, + KV caches). Runs prefill to populate the caches, then a single decode step. + + The output projection (wo_a, wo_b) is sharded for `tp_size`-way tensor + parallelism; this module holds ONLY rank `tp_rank`'s shard and returns that + rank's partial [B, 1, dim] contribution (the all-reduce over ranks is the + caller's responsibility). + """ + + def __init__(self, config: CSAConfig, tp_size: int = 4, tp_rank: int = 0): + super().__init__() + self.config = config + self.dim = config.dim + self.n_heads = config.n_heads + self.q_lora_rank = config.q_lora_rank + self.o_lora_rank = config.o_lora_rank + self.head_dim = config.head_dim + self.rope_head_dim = config.rope_head_dim + self.n_groups = config.o_groups + self.eps = config.norm_eps + + assert self.n_groups % tp_size == 0, "o_groups must be divisible by tp_size" + self.tp_size = tp_size + self.tp_rank = tp_rank + self.n_local_groups = self.n_groups // tp_size # groups owned by this rank + self.group_in = self.n_heads * self.head_dim // self.n_groups # per-group wo_a input width + + pdt = torch.bfloat16 + + # ----- Query projection ----- + self.wq_a = nn.Linear(self.dim, self.q_lora_rank, bias=False, dtype=pdt) + self.q_norm = RMSNorm(self.q_lora_rank, self.eps) + self.wq_b = nn.Linear(self.q_lora_rank, self.n_heads * self.head_dim, bias=False, dtype=pdt) + + # ----- KV projection (for the new decode token) ----- + self.wkv = nn.Linear(self.dim, self.head_dim, bias=False, dtype=pdt) + self.kv_norm = RMSNorm(self.head_dim, self.eps) + + # ----- Output projection (grouped low-rank), rank-local shard ----- + # Full (world_size=1) wo_a: [n_groups * o_lora_rank, group_in] + # wo_b: [dim, n_groups * o_lora_rank] + # Rank r owns groups [r*n_local_groups : (r+1)*n_local_groups]: + # wo_a shard rows = those groups' o_lora_rank outputs -> [n_local_groups*o_lora_rank, group_in] + # wo_b shard cols = those groups' flattened inputs -> [dim, n_local_groups*o_lora_rank] + self.wo_a = nn.Linear(self.group_in, self.n_local_groups * self.o_lora_rank, bias=False, dtype=pdt) + self.wo_b = nn.Linear(self.n_local_groups * self.o_lora_rank, self.dim, bias=False, dtype=pdt) + + # ----- Core sparse attention (owns caches, attn_sink, compressor, indexer) ----- + self.core = CSAAttentionCore(config) + + # RoPE frequencies (shared with core; used here for q/kv RoPE). + self.register_buffer("freqs_cis", self.core.freqs_cis, persistent=False) + + # ---- shared projection helpers ----------------------------------------- + def _project_q(self, x, freqs_cis): + rd = self.rope_head_dim + qr = self.q_norm(self.wq_a(x)) + q = self.wq_b(qr) + q = q.unflatten(-1, (self.n_heads, self.head_dim)) + q = q * torch.rsqrt(q.square().mean(-1, keepdim=True) + self.eps) + apply_rotary_emb(q[..., -rd:], freqs_cis) + return q, qr + + def _project_kv(self, x, freqs_cis): + rd = self.rope_head_dim + kv = self.wkv(x) + kv = self.kv_norm(kv) + apply_rotary_emb(kv[..., -rd:], freqs_cis) + return kv + + def _output_projection(self, o, bsz, seqlen): + """Rank-local grouped low-rank output projection. + + o: [B, S, n_heads, head_dim] (post de-RoPE). This rank only consumes the + n_local_groups groups it owns. + """ + # Flatten heads then view as groups; keep only this rank's groups. + o = o.reshape(bsz, seqlen, self.n_groups, self.group_in) + g0 = self.tp_rank * self.n_local_groups + o_local = o[:, :, g0 : g0 + self.n_local_groups, :] # [B, S, n_local_groups, group_in] + + wo_a = self.wo_a.weight.view(self.n_local_groups, self.o_lora_rank, self.group_in) + o_local = torch.einsum("bsgd,grd->bsgr", o_local, wo_a) # [B, S, n_local_groups, o_lora_rank] + out_partial = self.wo_b(o_local.flatten(2)) # [B, S, dim] -- rank partial + + return out_partial + + # ---- prefill (populate caches) ----------------------------------------- + @torch.no_grad() + def prefill(self, x): + bsz, seqlen, _ = x.size() + freqs_cis = self.freqs_cis[0:seqlen] + q, qr = self._project_q(x, freqs_cis) + kv = self._project_kv(x, freqs_cis) + # core.prefill fills self.core.kv_cache and self.core.indexer.kv_cache. + self.core.prefill(q, kv, x, qr) + + # ---- decode (single new token) ----------------------------------------- + @torch.no_grad() + def forward(self, x, start_pos): + bsz, seqlen, _ = x.size() + assert seqlen == 1, f"Decode expects seqlen=1, got {seqlen}" + rd = self.rope_head_dim + freqs_cis = self.freqs_cis[start_pos : start_pos + 1] + + q, qr = self._project_q(x, freqs_cis) + kv = self._project_kv(x, freqs_cis) + + # Core sparse attention (inserts the new kv into the window cache and + # attends over window + top-k compressed positions). + o = self.core(q, kv, x, qr, start_pos) # [B, 1, n_heads, head_dim] + + # De-rotate RoPE on the output's rope channels. + apply_rotary_emb(o[..., -rd:], freqs_cis, inverse=True) + + return self._output_projection(o, bsz, seqlen) # [B, 1, dim] rank partial + + +# -------------------------------------------------------------------------- +# Reference data generation (prefill -> extract caches -> decode) +# -------------------------------------------------------------------------- +def generate_decode_block_reference_tp( + full_config, + tp_size: int = 4, + weight_gain: float = 0.46, + norm_init: float = 1.0, + sink_scale: float = 1.0, + input_scale: float = 1.0, +): + """Head-parallel tensor-parallel decomposition of the FULL model over `tp_size` ranks. + + The production model has full_config.n_heads (=128) query heads and + full_config.o_groups (=16) output-projection groups. We shard it HEAD-PARALLEL: + rank r owns query heads [r*Hl : (r+1)*Hl] (Hl = n_heads/tp_size = 32) which, + since a group tiles n_heads/o_groups (=8) heads, coincide EXACTLY with output + groups [r*Gl : (r+1)*Gl] (Gl = o_groups/tp_size = 4). So each rank is a + self-contained Hl-head / Gl-group block with its OWN weights, computing a + DIFFERENT slice of the attention — no redundant compute. + + Builds ONE full model, runs prefill+decode to the true full output (the + all-reduce target), then for each rank derives (a) its sharded weight + state_dict for an Hl-head/Gl-group block and (b) its golden output-projection + PARTIAL. Because decode attention is per-head (the indexer's top-k is + head-independent and the KV is shared), heads [r*Hl:(r+1)*Hl] of the full + model equal an independent Hl-head block's output; and because the grouped + projection sums over disjoint group blocks, sum_r partial_r == full output. + + Returns shared decode inputs + caches (head-independent, shared by all ranks), + ref_output_full, the per-rank partials, and the per-rank shard state_dicts. + """ + full = CSAAttentionBlockDecode(full_config, tp_size=1, tp_rank=0) + _init_block_weights(full, weight_gain=weight_gain, norm_init=norm_init, sink_scale=sink_scale) + full.eval() + + B, S = full_config.batch_size, full_config.seq_len + W = full_config.window_size + T_c = S // full_config.compress_ratio + Hl = full_config.n_heads // tp_size # 32 query heads per rank + Gl = full_config.o_groups // tp_size # 4 output groups per rank + hd = full_config.head_dim + R = full_config.o_lora_rank + + torch.manual_seed(99) + x_prefill = (torch.randn(B, S, full_config.dim) * input_scale).to(torch.bfloat16) + with torch.no_grad(): + full.prefill(x_prefill) + + # Head-independent caches, shared by every rank. + kv_window = full.core.kv_cache[:B, :W].clone() + kv_compress = full.core.kv_cache[:B, W : W + T_c].clone() + indexer_kv_cache = full.core.indexer.kv_cache[:B, :T_c].clone() + + torch.manual_seed(200) + x_dec = (torch.randn(B, 1, full_config.dim) * input_scale).to(torch.bfloat16) + rd = full_config.rope_head_dim + with torch.no_grad(): + # Full decode up to the per-head attention output o (post de-RoPE), then + # the full grouped output projection = the golden all-reduce target. + freqs_cis = full.freqs_cis[S : S + 1] + q, qr = full._project_q(x_dec, freqs_cis) + kv = full._project_kv(x_dec, freqs_cis) + o_full = full.core(q, kv, x_dec, qr, start_pos=S) # [B,1,128,head_dim] + apply_rotary_emb(o_full[..., -rd:], freqs_cis, inverse=True) + ref_output_full = full._output_projection(o_full.clone(), B, 1) # [B,1,dim] + + full_sd = { + k: v + for k, v in full.state_dict().items() + if not k.startswith("core.kv_cache") + and not k.startswith("core.freqs_cis") + and not k.startswith("freqs_cis") + and not k.startswith("core.indexer.kv_cache") + } + + per_rank_weights, ref_partials = [], [] + for r in range(tp_size): + h0, h1 = r * Hl, (r + 1) * Hl # this rank's query heads + g0r, g1r = r * Gl * R, (r + 1) * Gl * R # this rank's wo_a rows / wo_b cols + sd_r = {} + for k, v in full_sd.items(): + if k == "wq_b.weight": # [n_heads*hd, q_lora] -> this rank's heads + sd_r[k] = v[h0 * hd : h1 * hd, :].clone() + elif k == "core.attn_sink": # [n_heads] -> this rank's heads + sd_r[k] = v[h0:h1].clone() + elif k == "wo_a.weight": # [o_groups*R, group_in] -> this rank's groups + sd_r[k] = v[g0r:g1r, :].clone() + elif k == "wo_b.weight": # [dim, o_groups*R] -> this rank's group cols + sd_r[k] = v[:, g0r:g1r].clone() + else: # wq_a/q_norm/wkv/kv_norm/indexer/compressor replicated + sd_r[k] = v.clone() + per_rank_weights.append(sd_r) + # Golden partial: an Hl-head/Gl-group block over this rank's head slice of + # o_full (== what the NKI rank block computes independently). + with torch.no_grad(): + og = o_full[:, :, h0:h1, :].reshape(B, 1, Gl, Hl * hd // Gl) # [B,1,Gl,group_in] + wo_a_r = sd_r["wo_a.weight"].view(Gl, R, Hl * hd // Gl) + lat = torch.einsum("bsgd,grd->bsgr", og, wo_a_r) # [B,1,Gl,R] + partial_r = torch.matmul(lat.reshape(B, 1, Gl * R), sd_r["wo_b.weight"].t()) # [B,1,dim] + ref_partials.append(partial_r) + + summed = torch.stack([p.float() for p in ref_partials], 0).sum(0) + max_sum_err = (summed - ref_output_full.float()).abs().max().item() + + return { + "x_dec": x_dec, + "kv_window": kv_window, + "kv_compress": kv_compress, + "indexer_kv_cache": indexer_kv_cache, + "ref_output_full": ref_output_full, # all-reduce target [B,1,dim] + "ref_partials": ref_partials, # list of tp_size [B,1,dim] partials + "per_rank_weights": per_rank_weights, # list of tp_size shard state_dicts + "max_sum_err": max_sum_err, + "tp_size": tp_size, + } diff --git a/src/nkilib_src/nkilib/experimental/deepseek_v4_csa/csa_common.py b/src/nkilib_src/nkilib/experimental/deepseek_v4_csa/csa_common.py new file mode 100644 index 0000000..280bff0 --- /dev/null +++ b/src/nkilib_src/nkilib/experimental/deepseek_v4_csa/csa_common.py @@ -0,0 +1,276 @@ +# 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. + +"""Configuration and host-side helpers shared across the DeepSeek-V4 CSA kernels.""" + +import math +from dataclasses import dataclass, replace + +import torch +from torch import nn + + +@dataclass +class CSAConfig: + """One tensor-parallel rank's view of the DeepSeek-V4 CSA attention block. + + ``n_heads`` and ``o_groups`` are the RANK-LOCAL counts. The compressed cache + holds ``seq_len // compress_ratio`` positions (``T_c``), which is what the + indexer scores and what the top-k selects from. + """ + + dim: int = 7168 + """Model hidden size.""" + + n_heads: int = 32 + """Rank-local query heads (128 in the full model, over 4 ranks).""" + + head_dim: int = 512 + """Attention head dimension.""" + + rope_head_dim: int = 64 + """Rotated channels of each head; the leading ``head_dim - rope_head_dim`` pass through.""" + + q_lora_rank: int = 1536 + """Query latent rank, shared between the attention q-path and the indexer.""" + + o_groups: int = 16 + """Rank-local output projection groups (16 in the full model, over 4 ranks).""" + + o_lora_rank: int = 1024 + """Output projection latent rank -- the low-rank bottleneck between wo_a and wo_b.""" + + window_size: int = 128 + """Sliding window positions attended in full, on top of the selected compressed ones.""" + + compress_ratio: int = 4 + """Raw tokens folded into one compressed cache position.""" + + norm_eps: float = 1e-6 + """RMSNorm epsilon.""" + + index_n_heads: int = 64 + """Lightning indexer query heads.""" + + index_head_dim: int = 128 + """Lightning indexer head dimension.""" + + index_topk: int = 1024 + """Compressed positions the sparse attention selects (``k``).""" + + compress_rope_theta: float = 160000.0 + """RoPE base.""" + + original_seq_len: int = 65536 + """RoPE reference length; 0 disables the NTK correction.""" + + rope_factor: float = 16.0 + """RoPE scaling factor.""" + + beta_fast: int = 32 + """High end of the RoPE correction range.""" + + beta_slow: int = 1 + """Low end of the RoPE correction range.""" + + batch_size: int = 1 + """Batch size.""" + + seq_len: int = 8192 + """Context length. The compressed cache holds ``seq_len // compress_ratio`` positions.""" + + @property + def compressed_len(self) -> int: + """``T_c`` -- compressed positions in the cache.""" + return self.seq_len // self.compress_ratio + + @property + def group_in(self) -> int: + """Per-group input width of the output projection's wo_a.""" + return self.n_heads * self.head_dim // self.o_groups + + +@dataclass +class CSAConfigFull(CSAConfig): + """The full unsharded shape: 128 query heads, 16 output groups. + + A whole chip holds this configuration as ``tp_size`` head-parallel ranks; each + rank runs ``shard_for_tp(config, tp_size)`` on 2 logical NeuronCores. + """ + + n_heads: int = 128 + o_groups: int = 16 + + +def shard_for_tp(config: CSAConfig, tp_size: int) -> CSAConfig: + """Return one head-parallel rank's config: ``n_heads`` and ``o_groups`` divided by ``tp_size``. + + Head-parallel sharding leaves ``group_in`` unchanged, because both the head + count and the group count divide by the same factor. Every other field is + replicated, since the indexer, the compressor and the sliding window are + shared: each rank scores all ``T_c`` compressed positions and selects the same + ``index_topk`` of them. + """ + if config.n_heads % tp_size != 0: + raise ValueError(f"n_heads={config.n_heads} must be divisible by tp_size={tp_size}") + if config.o_groups % tp_size != 0: + raise ValueError(f"o_groups={config.o_groups} must be divisible by tp_size={tp_size}") + return replace(config, n_heads=config.n_heads // tp_size, o_groups=config.o_groups // tp_size) + + +def precompute_freqs_cos_sin(dim, seqlen, original_seq_len, base, factor, beta_fast, beta_slow): + """Build the RoPE rotation tables as real ``(cos, sin)`` of shape ``[seqlen, dim // 2]``. + + Real tables rather than complex ones, because the kernels rotate with + multiplies and adds on the Vector engine and there is no complex dtype on + device. ``original_seq_len > 0`` enables the YaRN-style NTK correction: the + low-frequency channels are divided by ``factor`` and the high-frequency ones + are left alone, with ``beta_fast``/``beta_slow`` setting the ramp between. + """ + + def find_correction_dim(num_rotations, dim, base, max_seq_len): + return dim * math.log(max_seq_len / (num_rotations * 2 * math.pi)) / (2 * math.log(base)) + + def find_correction_range(low_rot, high_rot, dim, base, max_seq_len): + low = math.floor(find_correction_dim(low_rot, dim, base, max_seq_len)) + high = math.ceil(find_correction_dim(high_rot, dim, base, max_seq_len)) + return max(low, 0), min(high, dim - 1) + + def linear_ramp_factor(min_val, max_val, dim): + if min_val == max_val: + max_val += 0.001 + linear_func = (torch.arange(dim, dtype=torch.float32) - min_val) / (max_val - min_val) + return torch.clamp(linear_func, 0, 1) + + freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + if original_seq_len > 0: + low, high = find_correction_range(beta_fast, beta_slow, dim, base, original_seq_len) + smooth = 1 - linear_ramp_factor(low, high, dim // 2) + freqs = freqs / factor * (1 - smooth) + freqs * smooth + + t = torch.arange(seqlen) + freqs = torch.outer(t, freqs) + return torch.cos(freqs), torch.sin(freqs) + + +def apply_rotary_emb_functional(x_rope, freqs_cos_sin, inverse=False): + """Functional RoPE using real cos/sin. Computes in fp32, returns input dtype.""" + dtype = x_rope.dtype + cos_f, sin_f = freqs_cos_sin + + x_pairs = x_rope.float().unflatten(-1, (-1, 2)) + x1 = x_pairs[..., 0] + x2 = x_pairs[..., 1] + + if inverse: + sin_f = -sin_f + + if x1.ndim == 3: + cos_f = cos_f.unsqueeze(0) + sin_f = sin_f.unsqueeze(0) + else: + cos_f = cos_f.unsqueeze(0).unsqueeze(2) + sin_f = sin_f.unsqueeze(0).unsqueeze(2) + + y1 = x1 * cos_f - x2 * sin_f + y2 = x1 * sin_f + x2 * cos_f + + return torch.stack([y1, y2], dim=-1).flatten(-2).to(dtype) + + +class RMSNorm(nn.Module): + """RMSNorm with a learnable fp32 gain, normalizing in fp32 and casting back on the way out. + + The fp32 accumulation and the cast at the output boundary are load-bearing: + the fused kernels reproduce exactly this dtype flow, so a kernel result can be + compared against this module bit-for-bit. + """ + + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim, dtype=torch.float32)) + + def forward(self, x): + dtype = x.dtype + x = x.float() + var = x.square().mean(-1, keepdim=True) + return (self.weight * (x * torch.rsqrt(var + self.eps))).to(dtype) + + +_hadamard_cache: dict = {} + + +def get_hadamard_matrix(n, device, dtype): + """Normalized ``[n, n]`` Sylvester Hadamard matrix, cached per (n, device, dtype). + + ``n`` must be a power of two. Scaled by ``n ** -0.5`` so the transform is + orthonormal and does not change the magnitude of what it rotates. + """ + key = (n, device, dtype) + if key not in _hadamard_cache: + H = torch.tensor([[1.0]]) + while H.shape[0] < n: + H = torch.cat( + [ + torch.cat([H, H], dim=1), + torch.cat([H, -H], dim=1), + ], + dim=0, + ) + _hadamard_cache[key] = (H * (n**-0.5)).to(dtype=dtype, device=device) + return _hadamard_cache[key] + + +def hadamard_transform(x): + """Rotate the last axis of ``x`` by the orthonormal Hadamard matrix, as one matmul. + + The indexer applies this to its query so the scored channels are decorrelated; + being orthonormal it leaves the dot products the indexer takes unchanged in + aggregate while spreading each channel's contribution. + """ + n = x.shape[-1] + H = get_hadamard_matrix(n, x.device, x.dtype) + return x @ H + + +def precompute_win_bias_parts(S, W): + """Static sliding-window bias parts for prefill, both ``[S, 2 * W]``. + + Returns ``(base, sink_indicator)``: + + * ``base`` is ``0`` at window positions a query may attend and ``-1e9`` + elsewhere, so adding it before ``exp`` masks the rest out. + * ``sink_indicator`` is ``1.0`` at each query's attention-sink slot and ``0.0`` + elsewhere, so the caller scales it by the per-head sink scalar and adds it. + + Splitting the bias this way keeps both halves independent of the sink weights, + which are learned -- so these two tables are constant for a given ``(S, W)`` + and are built once at module construction rather than per call. + """ + TILE_Q = W + WIN_SIZE = 2 * W + + q_pos = torch.arange(S) + i_idx = (q_pos % TILE_Q).unsqueeze(1) + q_start = ((q_pos // TILE_Q) * TILE_Q).unsqueeze(1) + j_idx = torch.arange(WIN_SIZE).unsqueeze(0) + + valid = (j_idx >= i_idx + 1) & (j_idx <= i_idx + W) & (j_idx >= (W - q_start)) + sink_j = W - q_start + is_sink = (j_idx == sink_j) & valid + + base = torch.where(valid, torch.zeros(1), torch.tensor(-1e9)) + sink_indicator = is_sink.float() + return base, sink_indicator diff --git a/src/nkilib_src/nkilib/experimental/deepseek_v4_csa/csa_decode_attention.py b/src/nkilib_src/nkilib/experimental/deepseek_v4_csa/csa_decode_attention.py new file mode 100644 index 0000000..96ad462 --- /dev/null +++ b/src/nkilib_src/nkilib/experimental/deepseek_v4_csa/csa_decode_attention.py @@ -0,0 +1,1746 @@ +# 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 CSA decode attention kernels. + +Single-token decode over the compressed sparse attention block. The headline +kernel is ``nki_indexer_score_topk_gather_2core``, which runs the lightning +indexer's scoring, the top-k selection and the O(k) sparse attention in ONE +launch on a ``[2]``-grid (two logical NeuronCores), so neither the score row nor +the selected-index array returns to the host. + +Batching: decode is a SINGLE token, ``batch_size = 1``. Every kernel here takes one +query position per head (``S = 1``) and the ``[2]``-grid is spent on splitting the +sequence, not the batch. Prefill handles the multi-position case. + +""" + +import nki +import nki.isa as nisa +import nki.language as nl + +from ...core.utils.kernel_assert import kernel_assert + + +# -------------------------------------------------------------------------- +# NKI Kernel: fused RMS(+optional learnable gain) + RoPE for BOTH the q-path +# and the kv-path, merged into ONE kernel over a [n_heads+1, head_dim] tile. +# +# Replaces the two torch tails +# q = q * rsqrt(q.square().mean(-1) + eps); q[...,-rd:] = RoPE(q[...,-rd:]) +# kv = kv_norm(wkv(x)); kv[...,-rd:] = RoPE(kv[...,-rd:]) +# with q = [n_heads, head_dim] and kv = [1, head_dim]. +# +# Both do the same per-partition work -- RMS over the free axis in fp32, round to bf16 +# at the RMSNorm boundary, then RoPE in fp32 on the trailing rope channels with the +# same cos/sin. The one difference (q has no learnable gain, kv scales by +# kv_norm.weight) is unified by a per-partition gain tile: 1.0 on the q rows, the +# weight on the kv row. `x * 1.0 == x` is exact in fp32, so the q rows are unchanged. +# +# Packing both onto one partition tile turns two kernels into one launch, one +# HBM->SBUF load and one store. +# -------------------------------------------------------------------------- +@nki.jit +def nki_qkv_rms_rope_kernel( + q_in: nl.NkiTensor, + kv_in: nl.NkiTensor, + weight_in: nl.NkiTensor, + cos_in: nl.NkiTensor, + sin_in: nl.NkiTensor, + eps_val: float, +) -> nl.NkiTensor: + """Fused RMS(+per-partition gain) + RoPE for the merged q/kv tile. + + q_in: [n_heads, head_dim] bf16 — the query heads. + kv_in: [1, head_dim] bf16 — the single decode-token KV latent. + weight_in:[1, head_dim] fp32 — kv_norm.weight (learnable gain, kv row only). + cos_in: [1, half_rope] fp32 (NOT pre-repeated). sin_in: [1, half_rope] fp32. + Returns: [n_heads+1, head_dim] bf16 — rows 0..n_heads-1 = q heads (RMS+RoPE), + row n_heads = kv (learnable RMSNorm+RoPE). + + The q rows and the kv row are ASSEMBLED onto one partition tile inside the + kernel (two DMAs), so the host does no concatenation. The per-partition gain + tile is built on-chip: memset 1.0 (q rows -> x*1.0==x exact in fp32), DMA + kv_norm.weight into the kv row. Single-launch / single-core (shared_hbm). + """ + n_heads = q_in.shape[0] + head_dim = q_in.shape[1] + n_rows = n_heads + 1 # q heads + the single kv row + half_rope = cos_in.shape[1] + rope_head_dim = 2 * half_rope + nope_dim = head_dim - rope_head_dim + TILE = 128 # SBUF partitions + + kernel_assert( + n_rows <= TILE, + f"n_heads + 1 must fit the {TILE} SBUF partitions, got n_heads={n_heads}; " + f"shard the heads (tp_size >= 2) or split the kv row out", + ) + + out = nl.ndarray((n_rows, head_dim), dtype=nl.bfloat16, buffer=nl.shared_hbm) + + # ---- Assemble q rows + kv row onto ONE [n_rows, head_dim] tile (2 DMAs) ---- + x_sb = nl.ndarray((TILE, head_dim), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.dma_copy(dst=x_sb[0:n_heads, 0:head_dim], src=q_in[0:n_heads, 0:head_dim], priority=0) + nisa.dma_copy(dst=x_sb[n_heads:n_rows, 0:head_dim], src=kv_in[0:1, 0:head_dim], priority=0) + + # ---- RMS statistic ------------------------------------------------------ + x_sq = nl.ndarray((TILE, head_dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=x_sq[0:n_rows, 0:head_dim], + data1=x_sb[0:n_rows, 0:head_dim], + data2=x_sb[0:n_rows, 0:head_dim], + op=nl.multiply, + ) + msq = nl.ndarray((TILE, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_reduce(dst=msq[0:n_rows, 0:1], data=x_sq[0:n_rows, 0:head_dim], op=nl.add, axis=1) + # mean = sum / head_dim, then + eps (fused), then rsqrt. + nisa.tensor_scalar( + dst=msq[0:n_rows, 0:1], + data=msq[0:n_rows, 0:1], + op0=nl.multiply, + operand0=1.0 / head_dim, + op1=nl.add, + operand1=eps_val, + ) + rms = nl.ndarray((TILE, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation(dst=rms[0:n_rows, 0:1], op=nl.rsqrt, data=msq[0:n_rows, 0:1]) + + # Per-partition learnable gain built ON-CHIP: memset 1.0 (q rows -> x*1.0==x + # exact in fp32), DMA kv_norm.weight into the kv row (row n_heads) -> reproduces + # the learnable kv RMSNorm exactly. Host does no concatenation. + gain = nl.ndarray((TILE, head_dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.memset(dst=gain[0:n_rows, 0:head_dim], value=1.0) + nisa.dma_copy(dst=gain[n_heads:n_rows, 0:head_dim], src=weight_in[0:1, 0:head_dim], priority=1) + + # ---- normed = (x * rms) * gain, rounded to bf16, in ONE pass per region ---- + normed_nope = nl.ndarray((TILE, nope_dim), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.scalar_tensor_tensor( + dst=normed_nope[0:n_rows, 0:nope_dim], + data=x_sb[0:n_rows, 0:nope_dim], + op0=nl.multiply, + operand0=rms[0:n_rows, 0:1], + op1=nl.multiply, + operand1=gain[0:n_rows, 0:nope_dim], + ) + normed_pairs = nl.ndarray((TILE, half_rope, 2), dtype=nl.bfloat16, buffer=nl.sbuf) + normed_rope = normed_pairs.reshape((TILE, rope_head_dim)) + nisa.scalar_tensor_tensor( + dst=normed_rope[0:n_rows, 0:rope_head_dim], + data=x_sb[0:n_rows, nope_dim:head_dim], + op0=nl.multiply, + operand0=rms[0:n_rows, 0:1], + op1=nl.multiply, + operand1=gain[0:n_rows, nope_dim:head_dim], + ) + + # ---- Write the nope channels (0..nope_dim-1) straight to output ---- + nisa.dma_copy(dst=out[0:n_rows, 0:nope_dim], src=normed_nope[0:n_rows, 0:nope_dim]) + + # ---- RoPE on the last rope_head_dim channels (fp32 math) ---- + # normed_pairs is [.., half_rope, 2], so [...,0]=even (x1), [...,1]=odd (x2). + x1 = nl.ndarray((TILE, half_rope), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=x1[0:n_rows, 0:half_rope], src=normed_pairs[0:n_rows, 0:half_rope, 0]) + x2 = nl.ndarray((TILE, half_rope), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=x2[0:n_rows, 0:half_rope], src=normed_pairs[0:n_rows, 0:half_rope, 1]) + + cos_h = nl.ndarray((TILE, half_rope), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=cos_h[0:n_rows, 0:half_rope], src=cos_in.ap(pattern=[[0, n_rows], [1, half_rope]]), priority=2) + sin_h = nl.ndarray((TILE, half_rope), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=sin_h[0:n_rows, 0:half_rope], src=sin_in.ap(pattern=[[0, n_rows], [1, half_rope]]), priority=2) + + rope_bf16 = nl.ndarray((TILE, half_rope, 2), dtype=nl.bfloat16, buffer=nl.sbuf) + tmp_a = nl.ndarray((TILE, half_rope), dtype=nl.float32, buffer=nl.sbuf) + tmp_b = nl.ndarray((TILE, half_rope), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=tmp_a[0:n_rows, 0:half_rope], + data1=x1[0:n_rows, 0:half_rope], + data2=cos_h[0:n_rows, 0:half_rope], + op=nl.multiply, + ) + nisa.tensor_tensor( + dst=tmp_b[0:n_rows, 0:half_rope], + data1=x2[0:n_rows, 0:half_rope], + data2=sin_h[0:n_rows, 0:half_rope], + op=nl.multiply, + ) + nisa.tensor_tensor( + dst=rope_bf16[0:n_rows, 0:half_rope, 0], + data1=tmp_a[0:n_rows, 0:half_rope], + data2=tmp_b[0:n_rows, 0:half_rope], + op=nl.subtract, + ) + nisa.tensor_tensor( + dst=tmp_a[0:n_rows, 0:half_rope], + data1=x1[0:n_rows, 0:half_rope], + data2=sin_h[0:n_rows, 0:half_rope], + op=nl.multiply, + ) + nisa.tensor_tensor( + dst=tmp_b[0:n_rows, 0:half_rope], + data1=x2[0:n_rows, 0:half_rope], + data2=cos_h[0:n_rows, 0:half_rope], + op=nl.multiply, + ) + nisa.tensor_tensor( + dst=rope_bf16[0:n_rows, 0:half_rope, 1], + data1=tmp_a[0:n_rows, 0:half_rope], + data2=tmp_b[0:n_rows, 0:half_rope], + op=nl.add, + ) + + rope_bf16_flat = rope_bf16.reshape((TILE, rope_head_dim)) + nisa.dma_copy(dst=out[0:n_rows, nope_dim:head_dim], src=rope_bf16_flat[0:n_rows, 0:rope_head_dim]) + + return out + + +# -------------------------------------------------------------------------- +# NKI Kernel: decode indexer q projection + RoPE + Hadamard + Q^T broadcast +# -------------------------------------------------------------------------- +@nki.jit +def nki_indexer_qproj_rope_had_gemv( + wT: nl.NkiTensor, # [n_ktiles, 128, N] bf16 — wq_b with wT[t, kk, n] == w[n, t*128+kk] + qr_in: nl.NkiTensor, # [1, K] bf16 — the decode q-latent, K = n_ktiles * 128 + cos_in: nl.NkiTensor, # [1, half_rope] fp32 — this position's cos (NOT pre-repeated) + sin_in: nl.NkiTensor, # [1, half_rope] fp32 + hadamard: nl.NkiTensor, # [head_dim, head_dim] bf16 — normalized Hadamard + s_q: int, # query-tile width the scorer expects (the decode row is replicated s_q times) +) -> nl.NkiTensor: + """The whole decode indexer q path, in the launch that already did its projection. + + Extends ``nki_indexer_qproj_gemv`` (which this leaves untouched, since it has a torch + oracle and unit tests) to also do the RoPE, the Hadamard rotation and the Q^T + replication. That removes FIVE host ops -- the ``.t().contiguous().reshape``, + ``apply_rotary_emb_functional``, the ``cat`` rejoining the roped tail, + ``hadamard_transform``, and the ``permute/unsqueeze/expand/reshape/contiguous`` -- plus + the two intermediates they fed, without adding a launch. + + Returns [head_dim, n_heads * s_q] bf16, the ``q_T_all`` the scorer consumes, where every + one of a head's s_q columns is the same decode row (that replication is what the host + ``expand`` was doing). + + """ + core_id = nl.program_id(0) + n_cores = nl.num_programs() + + n_ktiles = wT.shape[0] + K_TILE = wT.shape[1] + N = wT.shape[2] + head_dim = hadamard.shape[0] + half_rope = cos_in.shape[1] + rope_head_dim = 2 * half_rope + nope_dim = head_dim - rope_head_dim + N_TILE = 128 + + kernel_assert(K_TILE == 128, "k-tile must be the 128-row nc_matmul contraction dim") + kernel_assert(head_dim == N_TILE, f"head_dim={head_dim} must be {N_TILE} (one n-tile per head)") + kernel_assert(N % N_TILE == 0, "N must tile evenly by 128 (one head's channels per tile)") + n_ntiles = N // N_TILE # == n_heads, since head_dim == N_TILE + kernel_assert(n_ntiles % n_cores == 0, "n-tiles must split evenly across the grid") + + nt_core = n_ntiles // n_cores + j0 = core_id * nt_core + c0 = j0 * N_TILE + N_core = nt_core * N_TILE + + out = nl.ndarray((head_dim, n_ntiles * s_q), dtype=nl.bfloat16, buffer=nl.shared_hbm, name="qproj_qT_all") + + qr_sb = nl.ndarray((K_TILE, n_ktiles), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.dma_copy(dst=qr_sb, src=qr_in.ap(pattern=[[1, K_TILE], [K_TILE, n_ktiles]], offset=0), priority=1) + + h_sb = nl.ndarray((head_dim, head_dim), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.dma_copy(dst=h_sb, src=hadamard[0:head_dim, 0:head_dim], priority=1) + + # ---- projection: unchanged from nki_indexer_qproj_gemv ---- + acc = nl.ndarray((K_TILE, nt_core), dtype=nl.float32, buffer=nl.psum) + w_sb = nl.ndarray((K_TILE, N_core), dtype=nl.bfloat16, buffer=nl.sbuf) + for t in nl.sequential_range(n_ktiles): + nisa.dma_copy(dst=w_sb, src=wT[t, 0:K_TILE, c0 : c0 + N_core], priority=0) + for j in nl.affine_range(nt_core): + nisa.nc_matmul( + dst=acc[0:K_TILE, j : j + 1], + stationary=w_sb[0:K_TILE, j * N_TILE : (j + 1) * N_TILE], + moving=qr_sb[0:K_TILE, t : t + 1], + accumulate=(t > 0), + ) + # Single fp32 -> bf16 rounding, where the bf16 nn.Linear rounded. + qc = nl.ndarray((K_TILE, nt_core), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.tensor_copy(dst=qc, src=acc) + + # ---- channel-major -> head-major so RoPE sees channels on the free axis ---- + hm_ps = nl.ndarray((nt_core, head_dim), dtype=nl.bfloat16, buffer=nl.psum) + nisa.nc_transpose(dst=hm_ps, data=qc) + hm = nl.ndarray((nt_core, head_dim), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.tensor_copy(dst=hm, src=hm_ps) + + # ---- RoPE on the trailing rope_head_dim channels, fp32 math ---- + rope_f = nl.ndarray((nt_core, rope_head_dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=rope_f, src=hm[0:nt_core, nope_dim:head_dim]) + pairs = rope_f.reshape((nt_core, half_rope, 2)) + x1 = nl.ndarray((nt_core, half_rope), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=x1, src=pairs[0:nt_core, 0:half_rope, 0]) + x2 = nl.ndarray((nt_core, half_rope), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=x2, src=pairs[0:nt_core, 0:half_rope, 1]) + + cos_h = nl.ndarray((nt_core, half_rope), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=cos_h, src=cos_in.ap(pattern=[[0, nt_core], [1, half_rope]]), priority=2) + sin_h = nl.ndarray((nt_core, half_rope), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=sin_h, src=sin_in.ap(pattern=[[0, nt_core], [1, half_rope]]), priority=2) + + ta = nl.ndarray((nt_core, half_rope), dtype=nl.float32, buffer=nl.sbuf) + tb = nl.ndarray((nt_core, half_rope), dtype=nl.float32, buffer=nl.sbuf) + y = nl.ndarray((nt_core, half_rope, 2), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=ta, data1=x1, data2=cos_h, op=nl.multiply) + nisa.tensor_tensor(dst=tb, data1=x2, data2=sin_h, op=nl.multiply) + nisa.tensor_tensor(dst=y[0:nt_core, 0:half_rope, 0], data1=ta, data2=tb, op=nl.subtract) + nisa.tensor_tensor(dst=ta, data1=x1, data2=sin_h, op=nl.multiply) + nisa.tensor_tensor(dst=tb, data1=x2, data2=cos_h, op=nl.multiply) + nisa.tensor_tensor(dst=y[0:nt_core, 0:half_rope, 1], data1=ta, data2=tb, op=nl.add) + # Rejoin in place: this is the `cat` the host was doing. + nisa.tensor_copy(dst=hm[0:nt_core, nope_dim:head_dim], src=y.reshape((nt_core, rope_head_dim))) + + # ---- Hadamard, which also restores the scorer's channel-major layout ---- + rt_ps = nl.ndarray((head_dim, nt_core), dtype=nl.bfloat16, buffer=nl.psum) + nisa.nc_transpose(dst=rt_ps, data=hm) + rt = nl.ndarray((head_dim, nt_core), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.tensor_copy(dst=rt, src=rt_ps) + had = nl.ndarray((head_dim, nt_core), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=had, stationary=h_sb, moving=rt) + qT = nl.ndarray((head_dim, nt_core), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=qT, src=had) + + # ---- replicate each head's column s_q times, then ONE DMA for the whole range ---- + ones = nl.ndarray((head_dim, s_q), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.memset(dst=ones, value=1.0) + out_sb = nl.ndarray((head_dim, nt_core * s_q), dtype=nl.bfloat16, buffer=nl.sbuf) + for j in nl.affine_range(nt_core): + nisa.tensor_scalar( + dst=out_sb[0:head_dim, j * s_q : (j + 1) * s_q], + data=ones, + op0=nl.multiply, + operand0=qT[0:head_dim, j : j + 1], + ) + nisa.dma_copy(dst=out[0:head_dim, j0 * s_q : (j0 + nt_core) * s_q], src=out_sb, priority=1) + + return out + +# -------------------------------------------------------------------------- +# nisa.topk batched kernel + encode/decode helpers +# -------------------------------------------------------------------------- +NISA_TOPK_GROUP_SIZE = 16 +NISA_TOPK_PARTITIONS = 128 +NISA_TOPK_GROUPS_PER_CALL = NISA_TOPK_PARTITIONS // NISA_TOPK_GROUP_SIZE # 8 + +# Sentinel used to pad a score row up to a proven-safe nisa.topk `n`. Every real +# indexer score is >= 0 (post-relu), so a padded slot can never enter the top-k. +_TOPK_PAD_SENTINEL = -1e9 + + +@nki.jit +def nisa_topk_snake_kernel(in_tensor: nl.NkiTensor, k_val: int, n_val: int) -> tuple[nl.NkiTensor, nl.NkiTensor]: + """Batched nisa.topk on snake-encoded input. + + in_tensor: [num_batches * 128, src_x] bf16 — snake-encoded scores. + Returns (values [num_batches * 128, k_val], indices [num_batches * 128, k_val]). + """ + total_rows = in_tensor.shape[0] + src_x = in_tensor.shape[1] + num_batches = total_rows // 128 + par_dim = 128 + + out_values = nl.ndarray((total_rows, k_val), dtype=nl.bfloat16, buffer=nl.shared_hbm) + out_indices = nl.ndarray((total_rows, k_val), dtype=nl.uint32, buffer=nl.shared_hbm) + + for b in nl.affine_range(num_batches): + src = nl.ndarray((par_dim, src_x), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.dma_copy(dst=src, src=in_tensor[b * 128 : (b + 1) * 128, 0:src_x]) + + val_dst = nl.ndarray((par_dim, k_val), dtype=nl.bfloat16, buffer=nl.sbuf) + idx_dst = nl.ndarray((par_dim, k_val), dtype=nl.uint32, buffer=nl.sbuf) + nisa.topk(val_dst=val_dst, idx_dst=idx_dst, src=src, n=n_val) + + nisa.dma_copy(dst=out_values[b * 128 : (b + 1) * 128, 0:k_val], src=val_dst) + nisa.dma_copy(dst=out_indices[b * 128 : (b + 1) * 128, 0:k_val], src=idx_dst) + + return out_values, out_indices + + +# -------------------------------------------------------------------------- +# NKI Kernel: Indexer per-head scoring (outputs raw scores for external topk) +# -------------------------------------------------------------------------- +@nki.jit +def nki_indexer_score_kernel( + q_T_all: nl.NkiTensor, # [head_dim, n_heads * S_q] — all heads' Q^T stacked (bf16) + kv_t: nl.NkiTensor, # [head_dim, T_c] — indexer_kv transposed (bf16), shared + weights: nl.NkiTensor, # [S_q, n_heads] — per-row per-head weights * weight_scale (fp32) + causal_bias: nl.NkiTensor, # [S_q, T_c] — causal bias (0 / -1e9) (fp32) +) -> nl.NkiTensor: + """Indexer scoring kernel — computes per-row scores for nkilib topk. + + Computes, per query row s and compressed kv position t: + index_score[s, t] = sum_h relu(q[s, h, :] . kv[t, :]) * weights[s, h] + + causal_bias[s, t] + + Returns: + scores: [S_q, T_c] fp32 — raw index scores (higher = more relevant). + """ + head_dim = q_T_all.shape[0] + total_q_free = q_T_all.shape[1] + T_c = kv_t.shape[1] + n_heads = weights.shape[1] + S_q = total_q_free // n_heads + + TILE_Q = 128 + SCORE_CHUNK = 512 if T_c >= 512 else T_c + num_score_chunks = (T_c + SCORE_CHUNK - 1) // SCORE_CHUNK + num_q_tiles = (S_q + TILE_Q - 1) // TILE_Q + + scores_out = nl.ndarray((S_q, T_c), dtype=nl.float32, buffer=nl.shared_hbm) + + kv_t_sb = nl.ndarray((head_dim, T_c), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.dma_copy(dst=kv_t_sb, src=kv_t[0:head_dim, 0:T_c]) + + core_id = nl.program_id(0) + n_cores = nl.num_programs() + tiles_per_core = num_q_tiles // n_cores + + for q_local in nl.affine_range(tiles_per_core): + q_idx = core_id * tiles_per_core + q_local + q_start = q_idx * TILE_Q + + w_tile = nl.ndarray((TILE_Q, n_heads), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=w_tile, src=weights[q_start : q_start + TILE_Q, 0:n_heads]) + + index_score_bf16 = nl.ndarray((TILE_Q, T_c), dtype=nl.bfloat16, buffer=nl.sbuf) + + for h in nl.affine_range(n_heads): + q_global = h * S_q + q_start + q_T = nl.ndarray((head_dim, TILE_Q), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.dma_copy(dst=q_T, src=q_T_all[0:head_dim, q_global : q_global + TILE_Q]) + + w_h = nl.ndarray((TILE_Q, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=w_h, src=w_tile[0:TILE_Q, h : h + 1]) + + for m_idx in nl.affine_range(num_score_chunks): + m_start = m_idx * SCORE_CHUNK + kt_slice = kv_t_sb[0:head_dim, m_start : m_start + SCORE_CHUNK] + s_psum = nl.ndarray((TILE_Q, SCORE_CHUNK), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=s_psum, stationary=q_T, moving=kt_slice) + s_relu = nl.ndarray((TILE_Q, SCORE_CHUNK), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.activation(dst=s_relu, op=nl.relu, data=s_psum) + if h == 0: + nisa.tensor_scalar( + dst=index_score_bf16[0:TILE_Q, m_start : m_start + SCORE_CHUNK], + data=s_relu, + op0=nl.multiply, + operand0=w_h, + ) + else: + nisa.scalar_tensor_tensor( + dst=index_score_bf16[0:TILE_Q, m_start : m_start + SCORE_CHUNK], + data=s_relu, + op0=nl.multiply, + operand0=w_h, + op1=nl.add, + operand1=index_score_bf16[0:TILE_Q, m_start : m_start + SCORE_CHUNK], + ) + + index_score = nl.ndarray((TILE_Q, T_c), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=index_score, src=index_score_bf16) + + cbias = nl.ndarray((TILE_Q, T_c), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=cbias, src=causal_bias[q_start : q_start + TILE_Q, 0:T_c]) + nisa.tensor_tensor(dst=index_score, data1=index_score, data2=cbias, op=nl.add) + + nisa.dma_copy(dst=scores_out[q_start : q_start + TILE_Q, 0:T_c], src=index_score) + + return scores_out + + +# -------------------------------------------------------------------------- +# NKI Kernel: FUSED indexer scoring + nisa.topk (single-chunk decode path) +# +# Fuses the two-kernel (nki_indexer_score_kernel -> HBM fp32 scores -> torch +# encode_snake -> nisa_topk_snake_kernel) pipeline into ONE @nki.jit kernel: +# 1. Score all T_c positions (reuses the proven per-head relu*weight matmul +# producing score[128, T_c] with the query on the partition dim; only +# row 0 is meaningful in decode since all query rows are identical). +# 2. Build the nisa.topk SNAKE src [128, T_c/16]: snake[r, c] = score[16c+r] for +# r in [0,16), c in [0, T_c/16). SBUF cannot stride its partition dim, so the +# free->partition fold routes row 0 through a tiny HBM scratch (T_c bf16 = +# 4-16KB) that `_snake_fill` reads back contiguously. +# 3. Run nisa.topk(n=T_c) -> local indices in snake layout. +# 4. Write group-0's k indices (== global indices for the single chunk) to a +# [TOPK_ROWS, k] HBM tensor. Row 0 receives ALL k indices as an unordered +# SET, which is exactly what the downstream permutation-invariant softmax +# over gathered positions needs (== what Pass-1 candidate_indices[0] gave). +# +# This replaces the two-kernel pipeline's fp32 [S_q, T_c] HBM write (1MB @ +# T_c=2048), the torch encode_snake/decode_snake host glue, and the separate +# topk kernel launch with one fused kernel + a 4-16KB scratch round-trip. +# -------------------------------------------------------------------------- +@nki.jit +def nki_indexer_score_topk_kernel( + q_T_all: nl.NkiTensor, # [head_dim, n_heads * S_q] — all heads' Q^T stacked (bf16) + kv_t: nl.NkiTensor, # [head_dim, T_c] — indexer_kv transposed (bf16) + weights: nl.NkiTensor, # [S_q, n_heads] — per-row per-head weights * weight_scale (fp32) + k_val: int, # number of top-k indices to return (== T_c-capped topk) +) -> nl.NkiTensor: + """Fused indexer scoring + top-k. Returns [TOPK_ROWS, k_val] uint32 indices. + + Computes, per compressed kv position t (row 0 of the identical decode query): + score[t] = sum_h relu(q[0, h, :] . kv[t, :]) * weights[0, h] + then returns the GLOBAL indices of the k_val largest scores (order-agnostic). + + Requirements (single-chunk decode path): + T_c divisible by 128 and by 16; k_val divisible by 16; S_q >= 128. + """ + head_dim = q_T_all.shape[0] + total_q_free = q_T_all.shape[1] + T_c = kv_t.shape[1] + n_heads = weights.shape[1] + S_q = total_q_free // n_heads + + TILE_Q = 128 + GROUP = 16 # nisa.topk snake group size + SNAKE_X = T_c // GROUP # snake free dim (T_c/16): 128 (T_c=2048) or 512 (T_c=8192) + SCORE_CHUNK = 512 if T_c >= 512 else T_c + num_score_chunks = (T_c + SCORE_CHUNK - 1) // SCORE_CHUNK + TOPK_ROWS = 8 # minimum rows nisa.topk operates on (128/16) + PAR = 128 + + out_indices = nl.ndarray((TOPK_ROWS, k_val), dtype=nl.uint32, buffer=nl.shared_hbm) + + # --- Preload shared KV^T (query on partition scoring matmul) --- + kv_t_sb = nl.ndarray((head_dim, T_c), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.dma_copy(dst=kv_t_sb, src=kv_t[0:head_dim, 0:T_c]) + + # Per-head weights for query row 0 (all rows identical in decode). + w_tile = nl.ndarray((TILE_Q, n_heads), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=w_tile, src=weights[0:TILE_Q, 0:n_heads]) + + # --- Stage 1: score all T_c, query on partition (bf16 accumulation) --- + index_score_bf16 = nl.ndarray((TILE_Q, T_c), dtype=nl.bfloat16, buffer=nl.sbuf) + for h in nl.affine_range(n_heads): + q_global = h * S_q # q_start = 0 (single tile) + q_T = nl.ndarray((head_dim, TILE_Q), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.dma_copy(dst=q_T, src=q_T_all[0:head_dim, q_global : q_global + TILE_Q]) + + w_h = nl.ndarray((TILE_Q, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=w_h, src=w_tile[0:TILE_Q, h : h + 1]) + + for m_idx in nl.affine_range(num_score_chunks): + m_start = m_idx * SCORE_CHUNK + kt_slice = kv_t_sb[0:head_dim, m_start : m_start + SCORE_CHUNK] + s_psum = nl.ndarray((TILE_Q, SCORE_CHUNK), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=s_psum, stationary=q_T, moving=kt_slice) + s_relu = nl.ndarray((TILE_Q, SCORE_CHUNK), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.activation(dst=s_relu, op=nl.relu, data=s_psum) + if h == 0: + nisa.tensor_scalar( + dst=index_score_bf16[0:TILE_Q, m_start : m_start + SCORE_CHUNK], + data=s_relu, + op0=nl.multiply, + operand0=w_h, + ) + else: + nisa.scalar_tensor_tensor( + dst=index_score_bf16[0:TILE_Q, m_start : m_start + SCORE_CHUNK], + data=s_relu, + op0=nl.multiply, + operand0=w_h, + op1=nl.add, + operand1=index_score_bf16[0:TILE_Q, m_start : m_start + SCORE_CHUNK], + ) + + # --- Stage 2: build the nisa.topk snake src [128, SNAKE_X] --- + scratch = nl.ndarray((1, T_c), dtype=nl.bfloat16, buffer=nl.shared_hbm) + nisa.dma_copy(dst=scratch, src=index_score_bf16[0:1, 0:T_c]) + + snake_src = nl.ndarray((PAR, SNAKE_X), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.memset(dst=snake_src, value=0) + _snake_fill(scratch, snake_src, 0, SNAKE_X) + + # --- Stage 3: nisa.topk (snake layout) -> local indices --- + val_dst = nl.ndarray((PAR, k_val), dtype=nl.bfloat16, buffer=nl.sbuf) + idx_dst = nl.ndarray((PAR, k_val), dtype=nl.uint32, buffer=nl.sbuf) + nisa.topk(val_dst=val_dst, idx_dst=idx_dst, src=snake_src, n=T_c) + + # --- Stage 4: write group-0 indices (== global, single chunk) as a SET --- + # idx_dst group 0 = partitions 0..15, columns 0..k/16-1 (k_val local positions). + # Flatten into out row 0: out[0, p*(k/16) + c] = idx_dst[p, c]. Order-agnostic. + k_cols = k_val // GROUP + idx_grp0 = nl.ndarray((GROUP, k_cols), dtype=nl.uint32, buffer=nl.sbuf) + nisa.tensor_copy(dst=idx_grp0, src=idx_dst[0:GROUP, 0:k_cols]) + nisa.dma_copy(dst=out_indices[0:1, :].ap(pattern=[[k_cols, GROUP], [1, k_cols]], offset=0), src=idx_grp0) + + return out_indices + + +# -------------------------------------------------------------------------- +# NKI Kernel: 2-LNC indexer scoring (disjoint T_c halves) -> shared bf16 scores +# +# Splits the Stage-1 scoring across 2 LNC cores; the attention kernel that +# follows already uses [2], so the 2nd core would otherwise idle through the +# whole indexing phase. Each core scores a contiguous T_c/n_cores slice over ALL +# heads and writes only row 0 (all decode query rows are identical) to its +# DISJOINT half of a shared [1, W] bf16 buffer -> no cross-core reduction. The +# cross-core barrier that guarantees both halves land before the top-k reads them +# is `nisa.core_barrier(cores=(0,1))` INSIDE the merged kernel below +# (nki_indexer_score_topk_2core), so the score->topk hand-off costs no @nki.jit +# launch boundary; the top-k itself still runs on ONE core. +# +# HEAD-BATCHED scoring: the compressed index-KV is SHARED across all n_heads +# (only the query differs per head), so instead of the old sequential +# `for h in range(n_heads)` loop that re-streamed the same KV through the PE +# array 64x (64 matmuls + 64 relu + 64 bf16 weight-accumulates per chunk), this +# scores all heads with TWO matmuls per SCORE_CHUNK: +# allh[h, t] = q_compact[:, h] . kv[:, t] (matmul-1, heads on free) +# score[t] = sum_h relu(allh[h, t]) * w[h] (relu; matmul-2 reduces h) +# matmul-2's contraction over the 64 head-partitions accumulates in fp32 (more +# accurate than the old bf16 running sum), so max_abs_diff should stay <= the +# bit-identical baseline. q_compact [head_dim, n_heads] and w_compact [n_heads,1] +# are derived in-kernel via strided DMAs from the existing q_T_all / weights +# inputs (q_T_all[d, h*S_q] == q[h,d]; weights[0, h]), so the kernel signature, +# both call sites, and all host code are unchanged. +# -------------------------------------------------------------------------- +def _score_2core_stage( + q_T_all: nl.NkiTensor, # [head_dim, n_heads * S_q] — all heads' Q^T stacked (bf16) + kv_t: nl.NkiTensor, # [head_dim, T_c] — indexer_kv transposed (bf16) + weights: nl.NkiTensor, # [S_q, n_heads] — per-row per-head weights * weight_scale (fp32) + scores_dst: nl.NkiTensor, # [1, W] bf16 shared_hbm (W >= T_c) — destination score row +) -> None: + """Score this LNC core's disjoint T_c slice into scores_dst[0:1, t_base:...].""" + head_dim = q_T_all.shape[0] + total_q_free = q_T_all.shape[1] + T_c = kv_t.shape[1] + n_heads = weights.shape[1] + S_q = total_q_free // n_heads + + SCORE_CHUNK = 512 if T_c >= 512 else T_c + + core_id = nl.program_id(0) + n_cores = nl.num_programs() + Tc_per_core = T_c // n_cores # 1024 (s8192, [2]) or 4096 (s32768, [2]) + num_score_chunks = (Tc_per_core + SCORE_CHUNK - 1) // SCORE_CHUNK + t_base = core_id * Tc_per_core # this core's global T_c offset (register) + + # --- Preload this core's KV^T slice (shared across all heads) --- + kv_t_sb = nl.ndarray((head_dim, Tc_per_core), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.dma_copy(dst=kv_t_sb, src=kv_t[0:head_dim, t_base : t_base + Tc_per_core], priority=0) + + # --- Compact query [head_dim, n_heads]: q_compact[d, h] = q_T_all[d, h*S_q] + q_compact = nl.ndarray((head_dim, n_heads), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.dma_copy( + dst=q_compact, src=q_T_all.ap(pattern=[[n_heads * S_q, head_dim], [S_q, n_heads]], offset=0), priority=1 + ) + + w_f32 = nl.ndarray((n_heads, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=w_f32, src=weights.ap(pattern=[[1, n_heads], [1, 1]], offset=0), priority=1) + w_bf16 = nl.ndarray((n_heads, 1), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.tensor_copy(dst=w_bf16, src=w_f32) + + # --- Score this core's T_c slice: 2 matmuls per chunk (heads batched) --- + for m_idx in nl.affine_range(num_score_chunks): + m_start = m_idx * SCORE_CHUNK + kt_slice = kv_t_sb[0:head_dim, m_start : m_start + SCORE_CHUNK] + + # matmul-1: allh[h, t] = sum_d q_compact[d, h] * kt_slice[d, t] (fp32 psum) + allh = nl.ndarray((n_heads, SCORE_CHUNK), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=allh, stationary=q_compact, moving=kt_slice) + + # relu (matches reference: relu THEN weight) -> bf16 for matmul-2 + R = nl.ndarray((n_heads, SCORE_CHUNK), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.activation(dst=R, op=nl.relu, data=allh) + + # matmul-2: score[t] = sum_h w[h] * R[h, t] (fp32 accumulation over heads) + sc = nl.ndarray((1, SCORE_CHUNK), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=sc, stationary=w_bf16, moving=R) + + sc_bf = nl.ndarray((1, SCORE_CHUNK), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.tensor_copy(dst=sc_bf, src=sc) + nisa.dma_copy(dst=scores_dst[0:1, t_base + m_start : t_base + m_start + SCORE_CHUNK], src=sc_bf, priority=0) + + +@nki.jit +def nki_indexer_score_2core( + q_T_all: nl.NkiTensor, # [head_dim, n_heads * S_q] — all heads' Q^T stacked (bf16) + kv_t: nl.NkiTensor, # [head_dim, T_c] — indexer_kv transposed (bf16) + weights: nl.NkiTensor, # [S_q, n_heads] — per-row per-head weights * weight_scale (fp32) +) -> nl.NkiTensor: + """Score all T_c across n_cores LNC cores (disjoint T_c halves). Returns [1, T_c] bf16. + + Scores-only entry point, used by the MULTI-chunk path (s131072), whose per-chunk + top-k + Pass-2 merge is host-orchestrated over the assembled row. The + single-chunk decode path instead calls `nki_indexer_score_topk_2core`, which + runs this same scoring stage and the top-k inside ONE kernel launch. + """ + T_c = kv_t.shape[1] + scores_out = nl.ndarray((1, T_c), dtype=nl.bfloat16, buffer=nl.shared_hbm, name="indexer_scores_2core_out") + _score_2core_stage(q_T_all, kv_t, weights, scores_out) + return scores_out + + +# -------------------------------------------------------------------------- +# nisa.topk reads a [128, SNAKE_X] tile as 8 INDEPENDENT groups of 16 partitions, and +# within a group the element it calls logical index `j` lives at partition `j % 16`, +# column `j // 16` -- the "snake" layout. Only group 0 is filled and read, so the tile +# holds n_val = 16 * SNAKE_X scores. +# -------------------------------------------------------------------------- +_SNAKE_GROUP = 16 + + +def _snake_fill( + scores: nl.NkiTensor, # [1, n_val] bf16 — the score row, contiguous by position + dst: nl.NkiTensor, # [>=16, count] bf16 sbuf — snake tile (columns 0..count-1) + c0: int, # first snake column to produce + count: int, # number of snake columns to produce + priority: int = 0, +) -> None: + """dst[r, j] = scores[16 * (c0 + j) + r] for r in [0,16), j in [0,count). + + Reads `scores` as its natural [n_val/16, 16] row-major view (partition stride 16, + 16 contiguous elements per partition) and folds free->partition with nc_transpose + rather than with DMA descriptors. + """ + GROUP = _SNAKE_GROUP + NC = 128 if count % 128 == 0 else count + kernel_assert(NC <= 128, "snake column count must be <= 128 or a multiple of 128") + for b in nl.affine_range(count // NC): + blk = nl.ndarray((NC, GROUP), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.dma_copy( + dst=blk, + src=scores.ap(pattern=[[GROUP, NC], [1, GROUP]], offset=(c0 + b * NC) * GROUP), + priority=priority, + ) + tp = nl.ndarray((GROUP, NC), dtype=nl.bfloat16, buffer=nl.psum) + nisa.nc_transpose(dst=tp, data=blk) + nisa.tensor_copy(dst=dst[0:GROUP, b * NC : (b + 1) * NC], src=tp) + + +# -------------------------------------------------------------------------- +# Stage helper: snake-encode + nisa.topk on the assembled score row (ONE core) +# +# Stages 2-4 of nki_indexer_score_topk_kernel, factored into a plain Python +# helper so the merged kernel below can run it on core 0 only (after the +# cross-core barrier) with no extra @nki.jit launch. Consumes the [1, n_val] bf16 +# score row assembled by `_score_2core_stage` (+ its -1e9 top-k padding tail) and +# writes group-0's k GLOBAL indices as an unordered set in out row 0. +# -------------------------------------------------------------------------- +def _snake_topk_stage( + scores: nl.NkiTensor, # [1, n_val] bf16 shared_hbm — assembled + padded score row + out_indices: nl.NkiTensor, # [TOPK_ROWS, k_val] uint32 shared_hbm — destination + k_val: int, # number of top-k indices to return + n_val: int, # nisa.topk n (the proven-safe padded width) +) -> None: + """nisa.topk over snake-encoded scores -> out_indices row 0 (unordered set).""" + GROUP = _SNAKE_GROUP + SNAKE_X = n_val // GROUP + PAR = 128 + + snake_src = nl.ndarray((PAR, SNAKE_X), dtype=nl.bfloat16, buffer=nl.sbuf) + _snake_fill(scores, snake_src, 0, SNAKE_X, priority=0) + + val_dst = nl.ndarray((PAR, k_val), dtype=nl.bfloat16, buffer=nl.sbuf) + idx_dst = nl.ndarray((PAR, k_val), dtype=nl.uint32, buffer=nl.sbuf) + nisa.topk(val_dst=val_dst, idx_dst=idx_dst, src=snake_src, n=n_val) + + # The strided fill makes the returned index a global position already. Write group + # 0's indices as a SET into out row 0: out[0, p*(k/16) + c]. Order-agnostic. + k_cols = k_val // GROUP + idx_grp0 = nl.ndarray((GROUP, k_cols), dtype=nl.uint32, buffer=nl.sbuf) + nisa.tensor_copy(dst=idx_grp0, src=idx_dst[0:GROUP, 0:k_cols]) + nisa.dma_copy(dst=out_indices[0:1, :].ap(pattern=[[k_cols, GROUP], [1, k_cols]], offset=0), src=idx_grp0) + + +def _snake_topk_stage_2core( + scores: nl.NkiTensor, # [1, n_val] bf16 shared_hbm — assembled + padded score row + out_indices: nl.NkiTensor, # [TOPK_ROWS, k_val] uint32 shared_hbm — destination + k_val: int, + n_val: int, # nisa.topk n (proven-safe padded width) + core_id: int, # this LNC's program id +) -> None: + """2-LNC version of `_snake_topk_stage`: split the descriptor-bound snake + reformat DMA across both cores, exchange via nisa.sendrecv, top-k on core 0. + + The single-core stage runs entirely on core 0 while core 1 sits at the barrier with + its 16 DMA engines idle. Splitting the snake free axis `c` in half puts half the + fill on each core's own engines; one nisa.sendrecv then swaps the halves + SBUF<->SBUF (no HBM round-trip) so core 0 can reassemble the whole tile. + + Core 0 keeps its half at columns [0, HALF) and places the received half at + [HALF, SNAKE_X), reproducing exactly the tile the single-core path builds -- so the + snake-to-global remap is the same, and only core 0 runs the top-k, leaving the + n-safety argument untouched. + """ + GROUP = _SNAKE_GROUP + SNAKE_X = n_val // GROUP + HALF = SNAKE_X // 2 + PAR = 128 + peer = 1 - core_id + + my_half = nl.ndarray((PAR, HALF), dtype=nl.bfloat16, buffer=nl.sbuf) + _snake_fill(scores, my_half, core_id * HALF, HALF, priority=0) + + peer_half = nl.ndarray((PAR, HALF), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.sendrecv(src=my_half, dst=peer_half, send_to_rank=peer, recv_from_rank=peer, pipe_id=0) + + if core_id == 0: + snake_src = nl.ndarray((PAR, SNAKE_X), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.tensor_copy(dst=snake_src[0:GROUP, 0:HALF], src=my_half[0:GROUP, 0:HALF]) + nisa.tensor_copy(dst=snake_src[0:GROUP, HALF:SNAKE_X], src=peer_half[0:GROUP, 0:HALF]) + + val_dst = nl.ndarray((PAR, k_val), dtype=nl.bfloat16, buffer=nl.sbuf) + idx_dst = nl.ndarray((PAR, k_val), dtype=nl.uint32, buffer=nl.sbuf) + nisa.topk(val_dst=val_dst, idx_dst=idx_dst, src=snake_src, n=n_val) + + k_cols = k_val // GROUP + idx_grp0 = nl.ndarray((GROUP, k_cols), dtype=nl.uint32, buffer=nl.sbuf) + nisa.tensor_copy(dst=idx_grp0, src=idx_dst[0:GROUP, 0:k_cols]) + nisa.dma_copy(dst=out_indices[0:1, :].ap(pattern=[[k_cols, GROUP], [1, k_cols]], offset=0), src=idx_grp0) + + +# -------------------------------------------------------------------------- +# NKI Kernel: MERGED 2-LNC indexer scoring + nisa.topk in ONE launch +# +# Removes one @nki.jit launch boundary from the decode critical path by running the +# 2-core scoring stage and the single-core top-k inside one `[2]`-grid kernel. The +# scoring matmuls are a tiny share of PE work, so this is entirely about the boundary +# (kernel-boundary DMA staging is the largest sync-engine opcode in the profile) plus +# the [1, T_c] HBM score round-trip. +# +# The hand-off the launch boundary used to provide becomes an intra-kernel barrier: +# 1. both cores write their DISJOINT T_c halves into the shared_hbm score row; +# 2. `nisa.core_barrier(data=scores_pad, cores=(0, 1))` establishes visibility; +# 3. core 0 alone runs the snake read + nisa.topk + index write-back. +# +# Per-core gating: the kernel is traced ONCE PER LOGICAL CORE and `nl.program_id(0)` +# folds to a Python int during that trace, so `if core_id == 0:` is real code +# specialization -- core 1's NEFF contains no top-k. (The "no device-if on a +# register" hazard applies to values that really are registers, such as +# nisa.register_load results, which need nl.dynamic_range / nl.while_loop.) +# +# nisa.topk n-safety moves ON-CHIP: the score row is allocated `n_val` wide and its +# [T_c, n_val) tail is memset to a -1e9 sentinel before the barrier, so the top-k +# runs at a validated n. On this heavily tied score distribution the selection depends +# on the width, so the width is pinned rather than tracking T_c. Every real score +# is >= 0 after the relu, so the sentinel can never win. +# -------------------------------------------------------------------------- +@nki.jit +def nki_indexer_score_topk_2core( + q_T_all: nl.NkiTensor, # [head_dim, n_heads * S_q] — all heads' Q^T stacked (bf16) + kv_t: nl.NkiTensor, # [head_dim, T_c] — indexer_kv transposed (bf16) + weights: nl.NkiTensor, # [S_q, n_heads] — per-row per-head weights * weight_scale (fp32) + k_val: int, # number of top-k indices to return + n_val: int, # nisa.topk n (proven-safe padded width, >= T_c) +) -> nl.NkiTensor: + """2-core scoring + core-0 nisa.topk in ONE launch. Returns [TOPK_ROWS, k_val] uint32. + + Requirements: launched on the [2] grid; n_val >= T_c and n_val % 16 == 0; + k_val % 16 == 0. Bit-identical to the score_2core -> host F.pad -> snake_topk + pipeline it replaces. + """ + T_c = kv_t.shape[1] + GROUP = 16 + TOPK_ROWS = 8 + + core_id = nl.program_id(0) + n_cores = nl.num_programs() + kernel_assert(n_cores == 2, "nki_indexer_score_topk_2core must be launched on the [2] grid") + kernel_assert(n_val >= T_c and n_val % GROUP == 0, "n_val must be >= T_c and a multiple of 16") + kernel_assert(k_val % GROUP == 0, "k_val must be a multiple of 16") + out_indices = nl.ndarray((TOPK_ROWS, k_val), dtype=nl.uint32, buffer=nl.shared_hbm, name="indexer_topk_out") + + # Score row padded out to the proven-safe nisa.topk width, and the cross-core + # exchange buffer the barrier synchronizes on. + + scores_pad = nl.ndarray((1, n_val), dtype=nl.bfloat16, buffer=nl.shared_hbm, name="indexer_scores_shared") + + if core_id == 0 and n_val > T_c: + pad_sb = nl.ndarray((1, n_val - T_c), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.memset(dst=pad_sb, value=_TOPK_PAD_SENTINEL) + nisa.dma_copy(dst=scores_pad[0:1, T_c:n_val], src=pad_sb) + + # Stage 1 (BOTH cores): disjoint T_c halves -> scores_pad[0:1, 0:T_c]. + _score_2core_stage(q_T_all, kv_t, weights, scores_pad) + + # Cross-core barrier: replaces the old @nki.jit launch boundary. Both cores' + # score halves (and the pad tail) are visible to every core after this point. + nisa.core_barrier(data=scores_pad, cores=(0, 1)) + + # Stages 2-4 on core 0 ONLY (nisa.topk must see the WHOLE assembled row, and + # runs on exactly one core). core 1's trace ends at the barrier. + if core_id == 0: + _snake_topk_stage(scores_pad, out_indices, k_val, n_val) + + return out_indices + + +# -------------------------------------------------------------------------- +# Stage helper: the WHOLE O(k) decode-attention body (gather -> score -> softmax +# -> V accumulate -> output de-RoPE -> write-back). +# +# A plain Python helper, not a @nki.jit kernel, so `nki_decode_gather_ok_kernel` (the +# standalone [1]-grid attention) and `nki_indexer_score_topk_gather_2core` (the fused +# [2]-grid kernel, which runs this inside its `core_id == 0` branch) trace the SAME +# instruction sequence rather than two copies of it. +# +# `idx_chunks` comes from the CALLER because the two callers read the same index bytes +# from differently-shaped sources -- the standalone kernel from a host-supplied [k, S] +# tensor, the fused kernel from row 0 of the top-k output it just wrote. Both are k +# contiguous uint32 with partition stride 1. +# +# `h_base` / `H_BATCH` / `output` are parameters rather than derived from +# nl.num_programs() because the fused kernel runs this body on one core of a [2] grid +# with all n_heads batched. +# -------------------------------------------------------------------------- +def _gather_attn_stage( + idx_chunks: list[nl.NkiTensor], # list[num_k_chunks] of [COMP_CHUNK, 1] uint32 SBUF row offsets + all_q_T: nl.NkiTensor, # [head_dim, n_heads * S] — replicated query per head + win_K_T: nl.NkiTensor, # [head_dim, W] — window K^T (real window only) + win_V: nl.NkiTensor, # [W, head_dim] — window V (real window only) + compress_kv: nl.NkiTensor, # [T_c, head_dim] bf16 — compressed KV row-major (K=V in CSA) + attn_sink_in: nl.NkiTensor, # [1, n_heads] — per-head sink scalars + derope_cos: nl.NkiTensor, # [1, half_rope] fp32 — inverse-RoPE cos (output de-RoPE) + derope_sin: nl.NkiTensor, # [1, half_rope] fp32 — inverse-RoPE sin + output: nl.NkiTensor, # [n_heads * S, head_dim] bf16 shared_hbm — destination + k: int, # number of gathered compressed positions + S: int, # query rows per head (1 in decode) + n_heads: int, # total heads spanned by `output` / all_q_T + h_base: int, # global head offset of this call's head batch + H_BATCH: int, # heads processed by this call +) -> None: + """O(k) decode attention: gather K and V in chunks, score+accumulate via matmul.""" + head_dim = all_q_T.shape[0] + KV_CHUNK = 128 + WIN_SIZE = KV_CHUNK + COMP_CHUNK = 128 + num_k_chunks = k // COMP_CHUNK + + HD_CHUNK = 128 + HD_TILES = (head_dim + HD_CHUNK - 1) // HD_CHUNK + + # Per-head attn-sink scalar for this core's 16 heads, on the partition dim: + # sink_hb[h_local, 0] = attn_sink_in[0, h_base + h_local]. + sink_hb = nl.ndarray((H_BATCH, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=sink_hb, src=attn_sink_in.ap(pattern=[[1, H_BATCH], [1, 1]], offset=h_base), priority=1) + + q_start = 0 + + # idx_chunks (the [COMP_CHUNK, 1] uint32 gather offsets) were loaded by the + # caller — see the header note on why the load lives there. + + # ---- Prefetch: gather all k compressed-KV chunks ONCE, up front ---------- + kv_chunks = [None] * num_k_chunks + for c_idx in nl.affine_range(num_k_chunks): + kv_bf = nl.ndarray((COMP_CHUNK, head_dim), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.dma_copy( + dst=kv_bf, + src=compress_kv.ap( + pattern=[[head_dim, COMP_CHUNK], [1, head_dim]], + vector_offset=idx_chunks[c_idx], + indirect_dim=0, + ), + dge_mode=nisa.dge_mode.swdge, + priority=0, # highest — the swdge gather gates all downstream compute + ) + kv_chunks[c_idx] = nl.ndarray((COMP_CHUNK, head_dim), dtype=nl.float16, buffer=nl.sbuf) + nisa.tensor_copy(dst=kv_chunks[c_idx], src=kv_bf) + + # Load window K^T (shared across all heads). + kv_t_win = [None] * HD_TILES + for hd in nl.affine_range(HD_TILES): + hd_start = hd * HD_CHUNK + hd_sz = min(HD_CHUNK, head_dim - hd_start) + kv_t_win[hd] = nl.ndarray((hd_sz, WIN_SIZE), dtype=nl.float16, buffer=nl.sbuf) + nisa.dma_copy( + dst=kv_t_win[hd], src=win_K_T[hd_start : hd_start + hd_sz, q_start : q_start + WIN_SIZE], priority=2 + ) # window K^T — lower priority than the gated gather + + # Load window V (shared across all heads). WIN_SIZE=KV_CHUNK now, so the + # whole real window is one chunk (the dead zero-pad half was dropped). + win_v_0 = nl.ndarray((KV_CHUNK, head_dim), dtype=nl.float16, buffer=nl.sbuf) + nisa.dma_copy( + dst=win_v_0, src=win_V[q_start : q_start + KV_CHUNK, 0:head_dim], priority=2 + ) # window V — lower priority than the gated gather + + # ---- Head-batched query [head_dim, H_BATCH] (head on the free/moving dim) ---- + q_hb = [None] * HD_TILES + for hd in nl.affine_range(HD_TILES): + hd_start = hd * HD_CHUNK + hd_sz = min(HD_CHUNK, head_dim - hd_start) + q_hb[hd] = nl.ndarray((hd_sz, H_BATCH), dtype=nl.float16, buffer=nl.sbuf) + nisa.dma_copy( + dst=q_hb[hd], + src=all_q_T.ap(pattern=[[n_heads * S, hd_sz], [S, H_BATCH]], offset=hd_start * (n_heads * S) + h_base * S), + priority=1, + ) # query — mid priority (needed for scoring, after the gather) + + # ---- Window scores [H_BATCH, WIN_SIZE] via one matmul over HD_TILES ---- + win_scores_psum = nl.ndarray((H_BATCH, WIN_SIZE), dtype=nl.float32, buffer=nl.psum) + for hd in nl.affine_range(HD_TILES): + nisa.nc_matmul(dst=win_scores_psum, stationary=q_hb[hd], moving=kv_t_win[hd]) + # ---- Window bias collapsed to the sink scalar (bit-identical dead-weight cut) ---- + win_scores = nl.ndarray((H_BATCH, WIN_SIZE), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=win_scores, src=win_scores_psum) + nisa.tensor_scalar( + dst=win_scores[0:H_BATCH, 0:1], + data=win_scores_psum[0:H_BATCH, 0:1], + op0=nl.add, + operand0=sink_hb[0:H_BATCH, 0:1], + ) + + # ---- Gathered compressed scoring [H_BATCH, k] ---- + # Gather+transpose all k positions into wide K^T tiles, then score with WIDE + # matmuls to amortize LDWEIGHTS + PE pipeline fill. + + SCORE_W = 512 # max moving free dim + num_score_groups = (k + SCORE_W - 1) // SCORE_W + + # Build wide K^T tiles kt_all[hd] = [HD_CHUNK, k], filled per gathered chunk. + kt_all = [None] * HD_TILES + for hd in range(HD_TILES): + hd_start = hd * HD_CHUNK + hd_sz = min(HD_CHUNK, head_dim - hd_start) + kt_all[hd] = nl.ndarray((hd_sz, k), dtype=nl.float16, buffer=nl.sbuf) + + for c_idx in nl.affine_range(num_k_chunks): + c_start = c_idx * COMP_CHUNK + + # K chunk was already gathered+cast into kv_chunks[c_idx] above (K=V in CSA, + # shared with the V matmul). Transpose K -> K^T [head_dim, COMP_CHUNK] into + # this chunk's slice of kt_all. + k_chunk = kv_chunks[c_idx] + for hd in range(HD_TILES): + hd_start = hd * HD_CHUNK + hd_sz = min(HD_CHUNK, head_dim - hd_start) + kt_psum = nl.ndarray((hd_sz, COMP_CHUNK), dtype=nl.float16, buffer=nl.psum) + nisa.nc_transpose(dst=kt_psum, data=k_chunk[0:COMP_CHUNK, hd_start : hd_start + hd_sz]) + nisa.tensor_copy(dst=kt_all[hd][0:hd_sz, c_start : c_start + COMP_CHUNK], src=kt_psum) + + # Wide scoring: SCORE_W positions per matmul group, HD_TILES accumulation. + comp_scores = nl.ndarray((H_BATCH, k), dtype=nl.float32, buffer=nl.sbuf) + for g in nl.affine_range(num_score_groups): + g_start = g * SCORE_W + g_sz = min(SCORE_W, k - g_start) + scores_psum = nl.ndarray((H_BATCH, g_sz), dtype=nl.float32, buffer=nl.psum) + for hd in nl.affine_range(HD_TILES): + hd_start = hd * HD_CHUNK + hd_sz = min(HD_CHUNK, head_dim - hd_start) + nisa.nc_matmul(dst=scores_psum, stationary=q_hb[hd], moving=kt_all[hd][0:hd_sz, g_start : g_start + g_sz]) + nisa.tensor_copy(dst=comp_scores[0:H_BATCH, g_start : g_start + g_sz], src=scores_psum) + + # ---- Softmax over window + k gathered scores (all heads, one pass each) ---- + win_max = nl.ndarray((H_BATCH, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_reduce(dst=win_max, data=win_scores, op=nl.maximum, axis=1) + comp_max = nl.ndarray((H_BATCH, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_reduce(dst=comp_max, data=comp_scores, op=nl.maximum, axis=1) + neg_max = nl.ndarray((H_BATCH, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=neg_max, data1=win_max, data2=comp_max, op=nl.maximum) + nisa.tensor_scalar(dst=neg_max, data=neg_max, op0=nl.multiply, operand0=-1.0) + + win_sum = nl.ndarray((H_BATCH, 1), dtype=nl.float32, buffer=nl.sbuf) + win_exp = nl.ndarray((H_BATCH, WIN_SIZE), dtype=nl.float16, buffer=nl.sbuf) + nisa.activation( + dst=win_exp, + op=nl.exp, + data=win_scores, + bias=neg_max, + reduce_op=nl.add, + reduce_res=win_sum, + reduce_cmd=nisa.reduce_cmd.reset_reduce, + ) + comp_sum = nl.ndarray((H_BATCH, 1), dtype=nl.float32, buffer=nl.sbuf) + comp_exp = nl.ndarray((H_BATCH, k), dtype=nl.float16, buffer=nl.sbuf) + nisa.activation( + dst=comp_exp, + op=nl.exp, + data=comp_scores, + bias=neg_max, + reduce_op=nl.add, + reduce_res=comp_sum, + reduce_cmd=nisa.reduce_cmd.reset_reduce, + ) + total_sum = nl.ndarray((H_BATCH, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=total_sum, data1=win_sum, data2=comp_sum, op=nl.add) + + # ---- Output accumulation [H_BATCH, head_dim] ---- + out_psum = nl.ndarray((H_BATCH, head_dim), dtype=nl.float32, buffer=nl.psum) + + we_T0 = nl.ndarray((KV_CHUNK, H_BATCH), dtype=nl.float16, buffer=nl.psum) + nisa.nc_transpose(dst=we_T0, data=win_exp[0:H_BATCH, 0:KV_CHUNK]) + we_T0_sb = nl.ndarray((KV_CHUNK, H_BATCH), dtype=nl.float16, buffer=nl.sbuf) + nisa.tensor_copy(dst=we_T0_sb, src=we_T0) + nisa.nc_matmul(dst=out_psum, stationary=we_T0_sb, moving=win_v_0) + + for c_idx in nl.affine_range(num_k_chunks): + c_start = c_idx * COMP_CHUNK + + v_chunk = kv_chunks[c_idx] + + ce_T = nl.ndarray((COMP_CHUNK, H_BATCH), dtype=nl.float16, buffer=nl.psum) + nisa.nc_transpose(dst=ce_T, data=comp_exp[0:H_BATCH, c_start : c_start + COMP_CHUNK]) + ce_T_sb = nl.ndarray((COMP_CHUNK, H_BATCH), dtype=nl.float16, buffer=nl.sbuf) + nisa.tensor_copy(dst=ce_T_sb, src=ce_T) + nisa.nc_matmul(dst=out_psum, stationary=ce_T_sb, moving=v_chunk) + + # ---- Finalize: normalize by total_sum, write this core's H_BATCH rows ---- + out_sbuf = nl.ndarray((H_BATCH, head_dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=out_sbuf, src=out_psum) + inv_sum = nl.ndarray((H_BATCH, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation(dst=inv_sum, op=nl.reciprocal, data=total_sum) + nisa.tensor_scalar(dst=out_sbuf, data=out_sbuf, op0=nl.multiply, operand0=inv_sum[0:H_BATCH, 0:1]) + out_bf16 = nl.ndarray((H_BATCH, head_dim), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.tensor_copy(dst=out_bf16, src=out_sbuf) + + # ---- Fused output de-RoPE (inverse rotation on the last rope channels) ---- + half_rope = derope_cos.shape[1] + rope_head_dim = 2 * half_rope + nope_dim = head_dim - rope_head_dim + + # Widen the rope channels bf16 -> fp32 and split into even (x1) / odd (x2). + d_rope_f = nl.ndarray((H_BATCH, rope_head_dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=d_rope_f[0:H_BATCH, 0:rope_head_dim], src=out_bf16[0:H_BATCH, nope_dim:head_dim]) + d_pairs = d_rope_f.reshape((H_BATCH, half_rope, 2)) + dx1 = nl.ndarray((H_BATCH, half_rope), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=dx1[0:H_BATCH, 0:half_rope], src=d_pairs[0:H_BATCH, 0:half_rope, 0]) + dx2 = nl.ndarray((H_BATCH, half_rope), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=dx2[0:H_BATCH, 0:half_rope], src=d_pairs[0:H_BATCH, 0:half_rope, 1]) + + dcos = nl.ndarray((H_BATCH, half_rope), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy( + dst=dcos[0:H_BATCH, 0:half_rope], src=derope_cos.ap(pattern=[[0, H_BATCH], [1, half_rope]]), priority=3 + ) + dsin = nl.ndarray((H_BATCH, half_rope), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy( + dst=dsin[0:H_BATCH, 0:half_rope], src=derope_sin.ap(pattern=[[0, H_BATCH], [1, half_rope]]), priority=3 + ) + + # y1 = x1*cos + x2*sin ; y2 = x2*cos - x1*sin + dta = nl.ndarray((H_BATCH, half_rope), dtype=nl.float32, buffer=nl.sbuf) + dtb = nl.ndarray((H_BATCH, half_rope), dtype=nl.float32, buffer=nl.sbuf) + dy1 = nl.ndarray((H_BATCH, half_rope), dtype=nl.float32, buffer=nl.sbuf) + dy2 = nl.ndarray((H_BATCH, half_rope), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=dta[0:H_BATCH, 0:half_rope], + data1=dx1[0:H_BATCH, 0:half_rope], + data2=dcos[0:H_BATCH, 0:half_rope], + op=nl.multiply, + ) + nisa.tensor_tensor( + dst=dtb[0:H_BATCH, 0:half_rope], + data1=dx2[0:H_BATCH, 0:half_rope], + data2=dsin[0:H_BATCH, 0:half_rope], + op=nl.multiply, + ) + nisa.tensor_tensor( + dst=dy1[0:H_BATCH, 0:half_rope], data1=dta[0:H_BATCH, 0:half_rope], data2=dtb[0:H_BATCH, 0:half_rope], op=nl.add + ) + nisa.tensor_tensor( + dst=dta[0:H_BATCH, 0:half_rope], + data1=dx2[0:H_BATCH, 0:half_rope], + data2=dcos[0:H_BATCH, 0:half_rope], + op=nl.multiply, + ) + nisa.tensor_tensor( + dst=dtb[0:H_BATCH, 0:half_rope], + data1=dx1[0:H_BATCH, 0:half_rope], + data2=dsin[0:H_BATCH, 0:half_rope], + op=nl.multiply, + ) + nisa.tensor_tensor( + dst=dy2[0:H_BATCH, 0:half_rope], + data1=dta[0:H_BATCH, 0:half_rope], + data2=dtb[0:H_BATCH, 0:half_rope], + op=nl.subtract, + ) + + # Re-interleave y1 (even) / y2 (odd), cast bf16, overwrite the rope channels. + d_out = nl.ndarray((H_BATCH, half_rope, 2), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=d_out[0:H_BATCH, 0:half_rope, 0], src=dy1[0:H_BATCH, 0:half_rope]) + nisa.tensor_copy(dst=d_out[0:H_BATCH, 0:half_rope, 1], src=dy2[0:H_BATCH, 0:half_rope]) + d_out_flat = d_out.reshape((H_BATCH, rope_head_dim)) + nisa.tensor_copy(dst=out_bf16[0:H_BATCH, nope_dim:head_dim], src=d_out_flat[0:H_BATCH, 0:rope_head_dim]) + + nisa.dma_copy( + dst=output.ap(pattern=[[S * head_dim, H_BATCH], [1, head_dim]], offset=h_base * S * head_dim), + src=out_bf16, + priority=1, + ) + + +def _gather_attn_stage_ksplit( + idx_chunks: list[nl.NkiTensor], # list[k_half/COMP_CHUNK] of [COMP_CHUNK,1] uint32 — THIS core's half + all_q_T: nl.NkiTensor, + win_K_T: nl.NkiTensor, + win_V: nl.NkiTensor, + compress_kv: nl.NkiTensor, + attn_sink_in: nl.NkiTensor, + derope_cos: nl.NkiTensor, + derope_sin: nl.NkiTensor, + output: nl.NkiTensor, + k: int, # TOTAL gathered positions (both halves) + S: int, + n_heads: int, + core_id: int, # 0 or 1 + k_half: int, # k // 2, this core's share of gathered positions +) -> None: + """O(k) decode attention with the K DIMENSION split across both LNCs. + + WHY THIS EXISTS (and why the HEAD split could not work). Splitting the head + batch was measured to change core 0's tensor time by 0.05% (383.1 -> 382.9 us): + the heads live on the matmul OUTPUT-PARTITION dim (M), and with M = 32 or 16 of + 128 the cost is set by the MOVING free dim (the k gathered positions) and the + head_dim contraction, not by M. Worse, everything expensive in this body is + k-driven and HEAD-INDEPENDENT — the swdge gather, the K^T transpose build, the + window loads — so a head split DUPLICATES all of it (8 -> 16 DMA_INDIRECT, + +1.25 MB HBM) while removing nothing from the critical path. + + Splitting along K fixes exactly that: core c owns gathered positions + [c*k_half, (c+1)*k_half), so per core the gather drops 8 -> 4 chunks, the + transpose build drops to k_half columns, and the scoring matmul's MOVING dim + drops k -> k_half. Total HBM traffic is UNCHANGED (each row is gathered once, + by exactly one core) — unlike the head split, which read every row twice. + """ + head_dim = all_q_T.shape[0] + KV_CHUNK = 128 + WIN_SIZE = KV_CHUNK + COMP_CHUNK = 128 + HD_CHUNK = 128 + HD_TILES = (head_dim + HD_CHUNK - 1) // HD_CHUNK + PAR = 128 + H_BATCH = n_heads # ALL heads on BOTH cores (the split is over k, not heads) + h_base = 0 + num_half_chunks = k_half // COMP_CHUNK + q_start = 0 + peer = 1 - core_id + + sink_hb = nl.ndarray((H_BATCH, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=sink_hb, src=attn_sink_in.ap(pattern=[[1, H_BATCH], [1, 1]], offset=h_base), priority=1) + + # ---- Gather ONLY this core's k_half positions (4 chunks, not 8) ---------- + kv_chunks = [None] * num_half_chunks + for c_idx in nl.affine_range(num_half_chunks): + kv_bf = nl.ndarray((COMP_CHUNK, head_dim), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.dma_copy( + dst=kv_bf, + src=compress_kv.ap( + pattern=[[head_dim, COMP_CHUNK], [1, head_dim]], + vector_offset=idx_chunks[c_idx], + indirect_dim=0, + ), + dge_mode=nisa.dge_mode.swdge, + priority=0, + ) + kv_chunks[c_idx] = nl.ndarray((COMP_CHUNK, head_dim), dtype=nl.float16, buffer=nl.sbuf) + nisa.tensor_copy(dst=kv_chunks[c_idx], src=kv_bf) + + # Window K^T / V. BOTH cores load it: the window max participates in the global + # softmax shift, so both need win_scores to form a bit-identical global max. Only + # core 0 accumulates the window V contribution (below), so it is counted once. + kv_t_win = [None] * HD_TILES + for hd in nl.affine_range(HD_TILES): + hd_start = hd * HD_CHUNK + hd_sz = min(HD_CHUNK, head_dim - hd_start) + kv_t_win[hd] = nl.ndarray((hd_sz, WIN_SIZE), dtype=nl.float16, buffer=nl.sbuf) + nisa.dma_copy( + dst=kv_t_win[hd], src=win_K_T[hd_start : hd_start + hd_sz, q_start : q_start + WIN_SIZE], priority=2 + ) + win_v_0 = nl.ndarray((KV_CHUNK, head_dim), dtype=nl.float16, buffer=nl.sbuf) + nisa.dma_copy(dst=win_v_0, src=win_V[q_start : q_start + KV_CHUNK, 0:head_dim], priority=2) + + q_hb = [None] * HD_TILES + for hd in nl.affine_range(HD_TILES): + hd_start = hd * HD_CHUNK + hd_sz = min(HD_CHUNK, head_dim - hd_start) + q_hb[hd] = nl.ndarray((hd_sz, H_BATCH), dtype=nl.float16, buffer=nl.sbuf) + nisa.dma_copy( + dst=q_hb[hd], + src=all_q_T.ap(pattern=[[n_heads * S, hd_sz], [S, H_BATCH]], offset=hd_start * (n_heads * S) + h_base * S), + priority=1, + ) + + # ---- Window scores + window max (both cores, identical values) ----------- + win_scores_psum = nl.ndarray((H_BATCH, WIN_SIZE), dtype=nl.float32, buffer=nl.psum) + for hd in nl.affine_range(HD_TILES): + nisa.nc_matmul(dst=win_scores_psum, stationary=q_hb[hd], moving=kv_t_win[hd]) + win_scores = nl.ndarray((H_BATCH, WIN_SIZE), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=win_scores, src=win_scores_psum) + nisa.tensor_scalar( + dst=win_scores[0:H_BATCH, 0:1], + data=win_scores_psum[0:H_BATCH, 0:1], + op0=nl.add, + operand0=sink_hb[0:H_BATCH, 0:1], + ) + + # ---- This core's half of the gathered scoring --------------------------- + SCORE_W = 512 + num_score_groups = (k_half + SCORE_W - 1) // SCORE_W + kt_all = [None] * HD_TILES + for hd in range(HD_TILES): + hd_start = hd * HD_CHUNK + hd_sz = min(HD_CHUNK, head_dim - hd_start) + kt_all[hd] = nl.ndarray((hd_sz, k_half), dtype=nl.float16, buffer=nl.sbuf) + for c_idx in nl.affine_range(num_half_chunks): + c_start = c_idx * COMP_CHUNK + k_chunk = kv_chunks[c_idx] + for hd in range(HD_TILES): + hd_start = hd * HD_CHUNK + hd_sz = min(HD_CHUNK, head_dim - hd_start) + kt_psum = nl.ndarray((hd_sz, COMP_CHUNK), dtype=nl.float16, buffer=nl.psum) + nisa.nc_transpose(dst=kt_psum, data=k_chunk[0:COMP_CHUNK, hd_start : hd_start + hd_sz]) + nisa.tensor_copy(dst=kt_all[hd][0:hd_sz, c_start : c_start + COMP_CHUNK], src=kt_psum) + + comp_scores = nl.ndarray((H_BATCH, k_half), dtype=nl.float32, buffer=nl.sbuf) + for g in nl.affine_range(num_score_groups): + g_start = g * SCORE_W + g_sz = min(SCORE_W, k_half - g_start) + scores_psum = nl.ndarray((H_BATCH, g_sz), dtype=nl.float32, buffer=nl.psum) + for hd in nl.affine_range(HD_TILES): + hd_start = hd * HD_CHUNK + hd_sz = min(HD_CHUNK, head_dim - hd_start) + nisa.nc_matmul(dst=scores_psum, stationary=q_hb[hd], moving=kt_all[hd][0:hd_sz, g_start : g_start + g_sz]) + nisa.tensor_copy(dst=comp_scores[0:H_BATCH, g_start : g_start + g_sz], src=scores_psum) + + # ---- PHASE 1 sendrecv: exchange local comp max -> identical global max ---- + # Packed into a 128-partition tile (a <128-partition sendrecv does NOT preserve + # layout on this HW — HW-verified earlier this iteration). + my_max = nl.ndarray((PAR, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_reduce(dst=my_max[0:H_BATCH, 0:1], data=comp_scores, op=nl.maximum, axis=1) + peer_max = nl.ndarray((PAR, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.sendrecv(src=my_max, dst=peer_max, send_to_rank=peer, recv_from_rank=peer, pipe_id=0) + + comp_max = nl.ndarray((H_BATCH, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=comp_max, data1=my_max[0:H_BATCH, 0:1], data2=peer_max[0:H_BATCH, 0:1], op=nl.maximum) + win_max = nl.ndarray((H_BATCH, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_reduce(dst=win_max, data=win_scores, op=nl.maximum, axis=1) + neg_max = nl.ndarray((H_BATCH, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=neg_max, data1=win_max, data2=comp_max, op=nl.maximum) + nisa.tensor_scalar(dst=neg_max, data=neg_max, op0=nl.multiply, operand0=-1.0) + + # ---- exp with the GLOBAL max (identical shift to the unsplit body) ------- + comp_sum = nl.ndarray((H_BATCH, 1), dtype=nl.float32, buffer=nl.sbuf) + comp_exp = nl.ndarray((H_BATCH, k_half), dtype=nl.float16, buffer=nl.sbuf) + nisa.activation( + dst=comp_exp, + op=nl.exp, + data=comp_scores, + bias=neg_max, + reduce_op=nl.add, + reduce_res=comp_sum, + reduce_cmd=nisa.reduce_cmd.reset_reduce, + ) + + out_psum = nl.ndarray((H_BATCH, head_dim), dtype=nl.float32, buffer=nl.psum) + win_sum = nl.ndarray((H_BATCH, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.memset(dst=win_sum, value=0.0) + if core_id == 0: + win_exp = nl.ndarray((H_BATCH, WIN_SIZE), dtype=nl.float16, buffer=nl.sbuf) + nisa.activation( + dst=win_exp, + op=nl.exp, + data=win_scores, + bias=neg_max, + reduce_op=nl.add, + reduce_res=win_sum, + reduce_cmd=nisa.reduce_cmd.reset_reduce, + ) + we_T0 = nl.ndarray((KV_CHUNK, H_BATCH), dtype=nl.float16, buffer=nl.psum) + nisa.nc_transpose(dst=we_T0, data=win_exp[0:H_BATCH, 0:KV_CHUNK]) + we_T0_sb = nl.ndarray((KV_CHUNK, H_BATCH), dtype=nl.float16, buffer=nl.sbuf) + nisa.tensor_copy(dst=we_T0_sb, src=we_T0) + # First write into out_psum on core 0 (overwrites). + nisa.nc_matmul(dst=out_psum, stationary=we_T0_sb, moving=win_v_0) + + # This core's half of the V accumulation (no explicit accumulate= — see above). + for c_idx in nl.affine_range(num_half_chunks): + c_start = c_idx * COMP_CHUNK + v_chunk = kv_chunks[c_idx] + ce_T = nl.ndarray((COMP_CHUNK, H_BATCH), dtype=nl.float16, buffer=nl.psum) + nisa.nc_transpose(dst=ce_T, data=comp_exp[0:H_BATCH, c_start : c_start + COMP_CHUNK]) + ce_T_sb = nl.ndarray((COMP_CHUNK, H_BATCH), dtype=nl.float16, buffer=nl.sbuf) + nisa.tensor_copy(dst=ce_T_sb, src=ce_T) + nisa.nc_matmul(dst=out_psum, stationary=ce_T_sb, moving=v_chunk) + + # ---- PHASE 2 sendrecv: exchange (partial acc | partial sums) ------------- + # One 128-partition fp32 tile: cols [0, head_dim) = partial V accumulator, + # col head_dim = this core's comp exp-sum, col head_dim+1 = its window sum. + PACK = head_dim + 2 + my_pack = nl.ndarray((PAR, PACK), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=my_pack[0:H_BATCH, 0:head_dim], src=out_psum) + nisa.tensor_copy(dst=my_pack[0:H_BATCH, head_dim : head_dim + 1], src=comp_sum) + nisa.tensor_copy(dst=my_pack[0:H_BATCH, head_dim + 1 : head_dim + 2], src=win_sum) + peer_pack = nl.ndarray((PAR, PACK), dtype=nl.float32, buffer=nl.sbuf) + nisa.sendrecv(src=my_pack, dst=peer_pack, send_to_rank=peer, recv_from_rank=peer, pipe_id=1) + + if core_id != 0: + return + + # ---- Core 0: combine the two partials, normalize, de-RoPE, write -------- + out_sbuf = nl.ndarray((H_BATCH, head_dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=out_sbuf, data1=my_pack[0:H_BATCH, 0:head_dim], data2=peer_pack[0:H_BATCH, 0:head_dim], op=nl.add + ) + total_sum = nl.ndarray((H_BATCH, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=total_sum, + data1=my_pack[0:H_BATCH, head_dim : head_dim + 1], + data2=peer_pack[0:H_BATCH, head_dim : head_dim + 1], + op=nl.add, + ) + # window sums: exactly one core wrote a nonzero value, so adding both is exact. + win_tot = nl.ndarray((H_BATCH, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=win_tot, + data1=my_pack[0:H_BATCH, head_dim + 1 : head_dim + 2], + data2=peer_pack[0:H_BATCH, head_dim + 1 : head_dim + 2], + op=nl.add, + ) + nisa.tensor_tensor(dst=total_sum, data1=total_sum, data2=win_tot, op=nl.add) + + inv_sum = nl.ndarray((H_BATCH, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation(dst=inv_sum, op=nl.reciprocal, data=total_sum) + nisa.tensor_scalar(dst=out_sbuf, data=out_sbuf, op0=nl.multiply, operand0=inv_sum[0:H_BATCH, 0:1]) + out_bf16 = nl.ndarray((H_BATCH, head_dim), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.tensor_copy(dst=out_bf16, src=out_sbuf) + + # Fused output de-RoPE (identical algebra to `_gather_attn_stage`). + half_rope = derope_cos.shape[1] + rope_head_dim = 2 * half_rope + nope_dim = head_dim - rope_head_dim + d_rope_f = nl.ndarray((H_BATCH, rope_head_dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=d_rope_f[0:H_BATCH, 0:rope_head_dim], src=out_bf16[0:H_BATCH, nope_dim:head_dim]) + d_pairs = d_rope_f.reshape((H_BATCH, half_rope, 2)) + dx1 = nl.ndarray((H_BATCH, half_rope), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=dx1[0:H_BATCH, 0:half_rope], src=d_pairs[0:H_BATCH, 0:half_rope, 0]) + dx2 = nl.ndarray((H_BATCH, half_rope), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=dx2[0:H_BATCH, 0:half_rope], src=d_pairs[0:H_BATCH, 0:half_rope, 1]) + dcos = nl.ndarray((H_BATCH, half_rope), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy( + dst=dcos[0:H_BATCH, 0:half_rope], src=derope_cos.ap(pattern=[[0, H_BATCH], [1, half_rope]]), priority=3 + ) + dsin = nl.ndarray((H_BATCH, half_rope), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy( + dst=dsin[0:H_BATCH, 0:half_rope], src=derope_sin.ap(pattern=[[0, H_BATCH], [1, half_rope]]), priority=3 + ) + dta = nl.ndarray((H_BATCH, half_rope), dtype=nl.float32, buffer=nl.sbuf) + dtb = nl.ndarray((H_BATCH, half_rope), dtype=nl.float32, buffer=nl.sbuf) + dy1 = nl.ndarray((H_BATCH, half_rope), dtype=nl.float32, buffer=nl.sbuf) + dy2 = nl.ndarray((H_BATCH, half_rope), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=dta, data1=dx1, data2=dcos, op=nl.multiply) + nisa.tensor_tensor(dst=dtb, data1=dx2, data2=dsin, op=nl.multiply) + nisa.tensor_tensor(dst=dy1, data1=dta, data2=dtb, op=nl.add) + nisa.tensor_tensor(dst=dta, data1=dx2, data2=dcos, op=nl.multiply) + nisa.tensor_tensor(dst=dtb, data1=dx1, data2=dsin, op=nl.multiply) + nisa.tensor_tensor(dst=dy2, data1=dta, data2=dtb, op=nl.subtract) + d_out = nl.ndarray((H_BATCH, half_rope, 2), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=d_out[0:H_BATCH, 0:half_rope, 0], src=dy1) + nisa.tensor_copy(dst=d_out[0:H_BATCH, 0:half_rope, 1], src=dy2) + d_out_flat = d_out.reshape((H_BATCH, rope_head_dim)) + nisa.tensor_copy(dst=out_bf16[0:H_BATCH, nope_dim:head_dim], src=d_out_flat[0:H_BATCH, 0:rope_head_dim]) + + nisa.dma_copy( + dst=output.ap(pattern=[[S * head_dim, H_BATCH], [1, head_dim]], offset=h_base * S * head_dim), + src=out_bf16, + priority=1, + ) + + +def _split_head_fraction(T_c: int) -> tuple[int, int]: + if T_c <= 4096: + return 1, 2 # even split: both cores take half the heads + return 0, 1 # don't split — measured to lose at every ratio for T_c>4096 + + +# -------------------------------------------------------------------------- +# NKI Kernel: O(k) Decode Attention — gathered K scoring + gathered V matmul +# +# Standalone `[1]`-grid entry point, kept for the MULTI-chunk indexer path (whose +# per-chunk top-k + Pass-2 merge is host-orchestrated, so the indices genuinely +# have to come back as a tensor). The single-chunk decode path instead calls +# `nki_indexer_score_topk_gather_2core`, which runs this same body inside the +# indexer's own launch. Both share `_gather_attn_stage` verbatim. +# -------------------------------------------------------------------------- +@nki.jit +def nki_decode_gather_ok_kernel( + topk_indices_T: nl.NkiTensor, # [k, S] uint32 — top-k indices transposed (partition=k) + all_q_T: nl.NkiTensor, # [head_dim, n_heads * S] — replicated query per head + win_K_T: nl.NkiTensor, # [head_dim, W] — window K^T (real window only) + win_V: nl.NkiTensor, # [W, head_dim] — window V (real window only) + compress_kv: nl.NkiTensor, # [T_c, head_dim] bf16 — compressed KV row-major (K=V in CSA) + attn_sink_in: nl.NkiTensor, # [1, n_heads] — per-head sink scalars + derope_cos: nl.NkiTensor, # [1, half_rope] fp32 — inverse-RoPE cos (output de-RoPE) + derope_sin: nl.NkiTensor, # [1, half_rope] fp32 — inverse-RoPE sin +) -> nl.NkiTensor: + """O(k) decode attention on a `[1]` (or `[2]`) grid. Returns [n_heads*S, head_dim]. + + topk_indices_T is [k, S] (transposed) so that partition-dim slicing gives + COMP_CHUNK distinct indices per chunk for the swdge gather. + """ + k, S = topk_indices_T.shape + head_dim = all_q_T.shape[0] + n_heads = all_q_T.shape[1] // S + COMP_CHUNK = 128 + num_k_chunks = k // COMP_CHUNK + + n_cores = nl.num_programs() + H_BATCH = n_heads // n_cores # per-core head batch (n_heads/2 @ [2], n_heads @ [1]) + core_id = nl.program_id(0) + h_base = core_id * H_BATCH # global head offset for this core + + output = nl.ndarray((n_heads * S, head_dim), dtype=nl.bfloat16, buffer=nl.shared_hbm) + + idx_chunks = [None] * num_k_chunks + for c_idx in nl.affine_range(num_k_chunks): + c_start = c_idx * COMP_CHUNK + idx_chunks[c_idx] = nl.ndarray((COMP_CHUNK, 1), dtype=nl.uint32, buffer=nl.sbuf) + nisa.dma_copy(dst=idx_chunks[c_idx], src=topk_indices_T[c_start : c_start + COMP_CHUNK, 0:1], priority=0) + + _gather_attn_stage( + idx_chunks, + all_q_T, + win_K_T, + win_V, + compress_kv, + attn_sink_in, + derope_cos, + derope_sin, + output, + k, + S, + n_heads, + h_base, + H_BATCH, + ) + return output + + +# -------------------------------------------------------------------------- +# NKI Kernel: FUSED 2-LNC indexer scoring + top-k + O(k) attention in ONE launch +# +# Replaces `nki_indexer_score_topk_2core[2]` -> `nki_decode_gather_ok_kernel[1]` on +# the single-chunk decode path by folding the attention body into the indexer's +# `[2]`-grid launch. Two things go away: one @nki.jit launch boundary, and the [k, S] +# index array's trip out to the host graph between the two launches. The indices now +# stay on chip -- the top-k writes them and the gather's row-offset loads read them +# straight back from the same buffer. +# +# Structure: +# 1. both cores score their DISJOINT T_c halves into the named shared_hbm row; +# 2. `nisa.core_barrier(data=..., cores=(0, 1))` so core 0 sees core 1's half; +# 3. core 0 only: the top-k, then the whole gather+attention body. The gather needs +# all n_heads on one core, which is what the standalone [1]-grid launch does +# anyway. Core 1's trace ends at the barrier (or takes a share of the work -- +# see the K-split and head-split branches below). +# +# The index hand-off is byte-for-byte what the host chain delivered: k contiguous +# uint32 with partition stride 1, read from row 0 of the same top-k output buffer. +# -------------------------------------------------------------------------- +@nki.jit +def nki_indexer_score_topk_gather_2core( + q_T_all: nl.NkiTensor, # [idx_head_dim, idx_n_heads * S_q] — indexer Q^T (bf16) + kv_t: nl.NkiTensor, # [idx_head_dim, T_c] — indexer_kv transposed (bf16) + weights: nl.NkiTensor, # [S_q, idx_n_heads] — per-head weights * weight_scale (fp32) + k_val: int, # number of top-k indices / gathered positions + n_val: int, # nisa.topk n (proven-safe padded width, >= T_c) + all_q_T: nl.NkiTensor, # [head_dim, n_heads * S] f16 — attention query per head + win_K_T: nl.NkiTensor, # [head_dim, W] f16 — window K^T (real window only) + win_V: nl.NkiTensor, # [W, head_dim] f16 — window V (real window only) + compress_kv: nl.NkiTensor, # [T_c, head_dim] bf16 — compressed KV row-major (K=V in CSA) + attn_sink_in: nl.NkiTensor, # [1, n_heads] fp32 — per-head sink scalars + derope_cos: nl.NkiTensor, # [1, half_rope] fp32 — inverse-RoPE cos (output de-RoPE) + derope_sin: nl.NkiTensor, # [1, half_rope] fp32 — inverse-RoPE sin +) -> nl.NkiTensor: + """2-core indexer score + core-0 top-k + core-0 O(k) attention, ONE launch. + + Returns the attention output [n_heads * S, head_dim] bf16 — the same tensor + `nki_decode_gather_ok_kernel` returns today. + + Requirements: launched on the [2] grid; n_val >= T_c and n_val % 16 == 0; + k_val % 128 == 0. Bit-identical to the + score_topk_2core[2] -> nki_decode_gather_ok_kernel[1] pipeline it replaces. + """ + T_c = kv_t.shape[1] + GROUP = 16 + TOPK_ROWS = 8 + COMP_CHUNK = 128 + + core_id = nl.program_id(0) + n_cores = nl.num_programs() + kernel_assert(n_cores == 2, "nki_indexer_score_topk_gather_2core needs the [2] grid") + kernel_assert(n_val >= T_c and n_val % GROUP == 0, "n_val must be >= T_c and a multiple of 16") + kernel_assert(k_val % COMP_CHUNK == 0, "k_val must be a multiple of the gather chunk (128)") + n_heads = attn_sink_in.shape[1] + S = all_q_T.shape[1] // n_heads + head_dim = all_q_T.shape[0] + output = nl.ndarray((n_heads * S, head_dim), dtype=nl.bfloat16, buffer=nl.shared_hbm, name="gather_attn_out") + + topk_idx = nl.ndarray((TOPK_ROWS, k_val), dtype=nl.uint32, buffer=nl.shared_hbm, name="indexer_topk_out") + scores_pad = nl.ndarray((1, n_val), dtype=nl.bfloat16, buffer=nl.shared_hbm, name="indexer_scores_shared") + + if core_id == 0 and n_val > T_c: + pad_sb = nl.ndarray((1, n_val - T_c), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.memset(dst=pad_sb, value=_TOPK_PAD_SENTINEL) + nisa.dma_copy(dst=scores_pad[0:1, T_c:n_val], src=pad_sb) + + # Stage 1 (BOTH cores): disjoint T_c halves -> scores_pad[0:1, 0:T_c]. + _score_2core_stage(q_T_all, kv_t, weights, scores_pad) + + # Cross-core barrier: both cores' score halves (and the pad tail) are visible + # to every core after this point. + nisa.core_barrier(data=scores_pad, cores=(0, 1)) + + # Stage 2: top-k. + if (n_val // 2) % GROUP == 0: + _snake_topk_stage_2core(scores_pad, topk_idx, k_val, n_val, core_id) + elif core_id == 0: + _snake_topk_stage(scores_pad, topk_idx, k_val, n_val) + + # ---- Stage 3: split the attention head batch over both cores, or not -------- + CORE1_HEAD_FRAC_NUM, CORE1_HEAD_FRAC_DEN = _split_head_fraction(T_c) + c1_heads = (n_heads * CORE1_HEAD_FRAC_NUM) // CORE1_HEAD_FRAC_DEN + split_heads = (c1_heads > 0) and (c1_heads < n_heads) + + # ---- K-SPLIT: divide the GATHERED POSITIONS (not the heads) over both cores ---- + use_ksplit = k_val % (2 * COMP_CHUNK) == 0 + if use_ksplit: + # Publish the top-k winners so BOTH cores can read their own half of the offsets. + nisa.core_barrier(data=topk_idx, cores=(0, 1)) + k_half = k_val // 2 + num_half_chunks = k_half // COMP_CHUNK + base = core_id * k_half # this core's first gathered position + idx_chunks = [None] * num_half_chunks + for c_idx in nl.affine_range(num_half_chunks): + c_start = base + c_idx * COMP_CHUNK + idx_chunks[c_idx] = nl.ndarray((COMP_CHUNK, 1), dtype=nl.uint32, buffer=nl.sbuf) + nisa.dma_copy( + dst=idx_chunks[c_idx], src=topk_idx.ap(pattern=[[1, COMP_CHUNK], [1, 1]], offset=c_start), priority=0 + ) + _gather_attn_stage_ksplit( + idx_chunks, + all_q_T, + win_K_T, + win_V, + compress_kv, + attn_sink_in, + derope_cos, + derope_sin, + output, + k_val, + S, + n_heads, + core_id, + k_half, + ) + return output + + if split_heads: + nisa.core_barrier(data=topk_idx, cores=(0, 1)) + + if split_heads or core_id == 0: + # core 0 takes the first (n_heads - c1_heads), core 1 the trailing c1_heads. + if split_heads: + H_BATCH = c1_heads if core_id == 1 else n_heads - c1_heads + h_base = (n_heads - c1_heads) if core_id == 1 else 0 + else: + H_BATCH = n_heads + h_base = 0 + + num_k_chunks = k_val // COMP_CHUNK + idx_chunks = [None] * num_k_chunks + for c_idx in nl.affine_range(num_k_chunks): + c_start = c_idx * COMP_CHUNK + idx_chunks[c_idx] = nl.ndarray((COMP_CHUNK, 1), dtype=nl.uint32, buffer=nl.sbuf) + nisa.dma_copy( + dst=idx_chunks[c_idx], src=topk_idx.ap(pattern=[[1, COMP_CHUNK], [1, 1]], offset=c_start), priority=0 + ) + + _gather_attn_stage( + idx_chunks, + all_q_T, + win_K_T, + win_V, + compress_kv, + attn_sink_in, + derope_cos, + derope_sin, + output, + k_val, + S, + n_heads, + h_base, + H_BATCH, + ) + + return output + + +@nki.jit +def nki_indexer_qproj_gemv(wT: nl.NkiTensor, qr_in: nl.NkiTensor) -> nl.NkiTensor: + """Indexer q-projection q = wq_b @ qr, returned TRANSPOSED as [head_dim, n_heads]. + + wT: [n_ktiles, 128, N] bf16 — the frozen wq_b weight pre-tiled so that + wT[t, kk, n] == wq_b.weight[n, t*128 + kk]. A pure host-side transform of + a frozen constant, so neuronx-cc constant-folds it and the ONLY weight + materialized is this one, at its true size. + qr_in: [1, K] bf16 — the decode q-latent, K = n_ktiles * 128. + Returns [head_dim, n_heads] bf16 where out[c, j] == q[head j, channel c]. + + Requires head_dim == 128 (the indexer's index_head_dim), so that an N-tile of + 128 columns is exactly one head's channel block and the PSUM tile is q^T. + """ + core_id = nl.program_id(0) + n_cores = nl.num_programs() + + n_ktiles = wT.shape[0] + K_TILE = wT.shape[1] # 128 — the nc_matmul contraction dim + N = wT.shape[2] # n_heads * head_dim + N_TILE = 128 # stationary free dim: 128*128*2 = 32 KB, matches + # the compiler's observed 1:1 MATMUL:LDWEIGHTS tile + n_ntiles = N // N_TILE + kernel_assert(K_TILE == 128, "k-tile must be the 128-row nc_matmul contraction dim") + kernel_assert(N % N_TILE == 0, "N must tile evenly by 128 (one head's channels per tile)") + kernel_assert(n_ntiles % n_cores == 0, "n-tiles must split evenly across the grid") + + nt_core = n_ntiles // n_cores + j0 = core_id * nt_core # first n-tile this core owns + c0 = j0 * N_TILE # first weight column this core owns + N_core = nt_core * N_TILE # weight columns this core streams + + out = nl.ndarray((K_TILE, n_ntiles), dtype=nl.bfloat16, buffer=nl.shared_hbm, name="qproj_qT") + + qr_sb = nl.ndarray((K_TILE, n_ktiles), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.dma_copy(dst=qr_sb, src=qr_in.ap(pattern=[[1, K_TILE], [K_TILE, n_ktiles]], offset=0), priority=1) + + # fp32 PSUM accumulator over THIS CORE's n-tiles only, + # [head_dim=128, n_heads/n_cores] = 128 B/partition at [2] (one bank). + acc = nl.ndarray((K_TILE, nt_core), dtype=nl.float32, buffer=nl.psum) + + w_sb = nl.ndarray((K_TILE, N_core), dtype=nl.bfloat16, buffer=nl.sbuf) + for t in nl.sequential_range(n_ktiles): + nisa.dma_copy(dst=w_sb, src=wT[t, 0:K_TILE, c0 : c0 + N_core], priority=0) + for j in nl.affine_range(nt_core): + nisa.nc_matmul( + dst=acc[0:K_TILE, j : j + 1], + stationary=w_sb[0:K_TILE, j * N_TILE : (j + 1) * N_TILE], + moving=qr_sb[0:K_TILE, t : t + 1], + accumulate=(t > 0), + ) + + # Single fp32 -> bf16 rounding at the end (matches the nn.Linear's bf16 output). + out_bf = nl.ndarray((K_TILE, nt_core), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.tensor_copy(dst=out_bf, src=acc) + nisa.dma_copy(dst=out[0:K_TILE, j0 : j0 + nt_core], src=out_bf, priority=1) + return out diff --git a/src/nkilib_src/nkilib/experimental/deepseek_v4_csa/csa_decode_attention_torch.py b/src/nkilib_src/nkilib/experimental/deepseek_v4_csa/csa_decode_attention_torch.py new file mode 100644 index 0000000..453546b --- /dev/null +++ b/src/nkilib_src/nkilib/experimental/deepseek_v4_csa/csa_decode_attention_torch.py @@ -0,0 +1,298 @@ +# 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. + +"""CPU references for the DeepSeek-V4 CSA decode kernels. + +One reference per tested kernel, each taking the SAME parameter names as its +kernel so the test framework can pair them. They are written as the plain +definition of what the kernel computes -- dense torch ops in fp32, no tiling, no +layout tricks -- so a disagreement points at the kernel rather than at a shared +mistake. + +Two conventions worth knowing when reading these: + +* The kernels take the query already TRANSPOSED and head-major, so + ``q_T_all[d, h * S_q + s] == q[s, h, d]``. Every reference undoes that + indexing explicitly rather than reshaping, since getting it wrong silently + is exactly the bug these tests are for. +* Decode is a single token, so ``S == 1``: each head's query is one column and + the ``[n_heads * S, head_dim]`` output has one row per head. The references + keep ``S`` general in the indexing anyway, and read only column 0, matching the + kernels. +""" + +import torch +import torch.nn.functional as F + + +def _rms_rope_rows(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, gain, eps: float) -> torch.Tensor: + """RMSNorm over the free axis then RoPE on the trailing ``2 * cos.shape[-1]`` channels. + + Mirrors the kernels' dtype flow exactly, which is what lets a kernel be graded + tightly rather than loosely: normalize in fp32, round to bf16 at the RMSNorm + output boundary (as the model does), then widen those rope channels back to + fp32 for the rotation and round once more at the end. + """ + half_rope = cos.shape[-1] + rope_dim = 2 * half_rope + nope_dim = x.shape[-1] - rope_dim + + xf = x.float() + rms = torch.rsqrt(xf.square().mean(-1, keepdim=True) + eps) + normed = xf * rms + if gain is not None: + normed = normed * gain + normed = normed.to(torch.bfloat16) + + rope = normed[:, nope_dim:].float().unflatten(-1, (half_rope, 2)) + x1, x2 = rope[..., 0], rope[..., 1] + y1 = x1 * cos - x2 * sin + y2 = x1 * sin + x2 * cos + rotated = torch.stack([y1, y2], dim=-1).flatten(-2).to(torch.bfloat16) + return torch.cat([normed[:, :nope_dim], rotated], dim=-1) + + +def _inverse_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + """Undo the query rotation on the trailing rope channels of an attention output. + + The inverse of the forward rotation is the forward formula with ``sin`` + negated, which is what the kernels fuse into their finalize step: + ``y1 = x1 * cos + x2 * sin``, ``y2 = x2 * cos - x1 * sin``. + """ + half_rope = cos.shape[-1] + rope_dim = 2 * half_rope + nope_dim = x.shape[-1] - rope_dim + + rope = x[:, nope_dim:].float().unflatten(-1, (half_rope, 2)) + x1, x2 = rope[..., 0], rope[..., 1] + y1 = x1 * cos + x2 * sin + y2 = x2 * cos - x1 * sin + rotated = torch.stack([y1, y2], dim=-1).flatten(-2).to(torch.bfloat16) + return torch.cat([x[:, :nope_dim], rotated], dim=-1) + + +def nki_qkv_rms_rope_torch_ref( + q_in: torch.Tensor, + kv_in: torch.Tensor, + weight_in: torch.Tensor, + cos_in: torch.Tensor, + sin_in: torch.Tensor, + eps_val: float, +) -> dict[str, torch.Tensor]: + """Oracle for ``nki_qkv_rms_rope_kernel``: the q heads and the kv row on one tile. + + The kernel packs both projection tails onto one partition tile and unifies them + with a per-partition gain of ``1.0`` on the q rows and ``kv_norm.weight`` on the + kv row. This reference instead runs the two paths SEPARATELY -- q with no gain, + kv with its learnable gain -- and concatenates. That is the point: if the + kernel's gain trick were not exact (``x * 1.0 == x`` in fp32), the two would + disagree. + """ + q_out = _rms_rope_rows(q_in, cos_in, sin_in, None, eps_val) + kv_out = _rms_rope_rows(kv_in, cos_in, sin_in, weight_in.float(), eps_val) + return {"output_0": torch.cat([q_out, kv_out], dim=0)} + + +def nki_indexer_qproj_gemv_torch_ref(wT: torch.Tensor, qr_in: torch.Tensor) -> dict[str, torch.Tensor]: + """Oracle for ``nki_indexer_qproj_gemv``: the indexer q-projection, returned transposed. + + ``wT`` is the frozen ``wq_b`` weight pre-tiled on the host so that + ``wT[t, kk, n] == wq_b.weight[n, t * 128 + kk]``. The reference rebuilds the + plain ``[N, K]`` weight from that tiling and does one dense matvec, then lays + the result out as the kernel's PSUM tile is: ``out[c, j] == q[head j, channel c]`` + with ``j`` an N-tile of 128 columns, i.e. one head's channel block. + """ + n_ktiles, k_tile, n = wT.shape + weight = wT.float().permute(2, 0, 1).reshape(n, n_ktiles * k_tile) + q = weight @ qr_in.float().reshape(-1) + return {"output_0": q.reshape(n // k_tile, k_tile).t().to(torch.bfloat16)} + + +def _indexer_scores(q_T_all: torch.Tensor, kv_t: torch.Tensor, weights: torch.Tensor) -> torch.Tensor: + """``score[s, t] = sum_h relu(q[s, h, :] . kv[t, :]) * weights[s, h]``, as ``[S_q, T_c]``. + + The relu comes BEFORE the per-head weight, so a head whose dot product is + negative contributes nothing at all rather than contributing negatively. That + is what makes every real indexer score non-negative, which in turn is what lets + the kernels pad a score row with a negative sentinel and know the padding can + never win the top-k. + """ + head_dim, total_q_free = q_T_all.shape + n_heads = weights.shape[1] + s_q = total_q_free // n_heads + + # q_T_all[d, h * S_q + s] == q[s, h, d] -> [S_q, n_heads, head_dim] + q = q_T_all.float().reshape(head_dim, n_heads, s_q).permute(2, 1, 0) + per_head = torch.einsum("shd,dt->sht", q, kv_t.float()) + return torch.einsum("sht,sh->st", F.relu(per_head), weights.float()) + + +def nki_indexer_score_torch_ref( + q_T_all: torch.Tensor, + kv_t: torch.Tensor, + weights: torch.Tensor, + causal_bias: torch.Tensor, +) -> dict[str, torch.Tensor]: + """Oracle for ``nki_indexer_score_kernel``: raw indexer scores plus the causal bias.""" + return {"output_0": _indexer_scores(q_T_all, kv_t, weights) + causal_bias.float()} + + +def nki_indexer_score_2core_torch_ref( + q_T_all: torch.Tensor, + kv_t: torch.Tensor, + weights: torch.Tensor, +) -> dict[str, torch.Tensor]: + """Oracle for ``nki_indexer_score_2core``: the assembled ``[1, T_c]`` score row. + + The kernel splits ``T_c`` in half across the two logical cores and each writes + its own half of a shared buffer, so a reference that scores the WHOLE range is + what catches a half that never landed -- the failure mode a per-core allocation + produces, where core 1's half reads back as zeros. + + Decode's query rows are all identical, so the kernel writes only row 0 and this + returns only row 0. + """ + scores = _indexer_scores(q_T_all, kv_t, weights) + return {"output_0": scores[0:1].to(torch.bfloat16)} + + +def _gather_attention( + idx: torch.Tensor, + all_q_T: torch.Tensor, + win_K_T: torch.Tensor, + win_V: torch.Tensor, + compress_kv: torch.Tensor, + attn_sink_in: torch.Tensor, + derope_cos: torch.Tensor, + derope_sin: torch.Tensor, + n_heads: int, + s_len: int, +) -> torch.Tensor: + """Dense O(W + k) sparse attention over the sliding window plus ``idx``'s rows. + + This is the whole point of CSA: the softmax runs over ``window_size + k`` + positions and nothing else, so its cost is independent of how long the context + is. The window is a full valid permutation in decode with no intra-window mask, + and the attention sink is a per-head bias on window column 0 only. + + ``all_q_T`` arrives ALREADY scaled by ``softmax_scale`` (the host folds it in), + so no scaling happens here. + """ + head_dim = all_q_T.shape[0] + + # Column 0 of each head: [head_dim, n_heads]. All S columns per head are + # identical in decode, and the kernels read only this one. + q_hb = all_q_T.float()[:, 0 : n_heads * s_len : s_len] + + win_scores = q_hb.t() @ win_K_T.float() + win_scores[:, 0] = win_scores[:, 0] + attn_sink_in.float().reshape(-1) + + gathered = compress_kv.float()[idx.long()] # [k, head_dim] + comp_scores = q_hb.t() @ gathered.t() # [n_heads, k] + + both = torch.cat([win_scores, comp_scores], dim=-1) + shift = both.max(dim=-1, keepdim=True).values + win_exp = torch.exp(win_scores - shift) + comp_exp = torch.exp(comp_scores - shift) + total = win_exp.sum(-1, keepdim=True) + comp_exp.sum(-1, keepdim=True) + + out = (win_exp @ win_V.float() + comp_exp @ gathered) / total + out = _inverse_rope(out.to(torch.bfloat16), derope_cos.float(), derope_sin.float()) + + # The kernel writes head h to output row h * S, leaving the other rows + # untouched; with S == 1 that is every row. + full = torch.zeros((n_heads * s_len, head_dim), dtype=torch.bfloat16) + full[0 : n_heads * s_len : s_len] = out + return full + + +def nki_decode_gather_ok_torch_ref( + topk_indices_T: torch.Tensor, + all_q_T: torch.Tensor, + win_K_T: torch.Tensor, + win_V: torch.Tensor, + compress_kv: torch.Tensor, + attn_sink_in: torch.Tensor, + derope_cos: torch.Tensor, + derope_sin: torch.Tensor, +) -> dict[str, torch.Tensor]: + """Oracle for ``nki_decode_gather_ok_kernel``: O(k) attention on caller-supplied indices. + + ``topk_indices_T`` is ``[k, S]``, transposed so the kernel can slice ``k`` onto + the partition axis for the gather; every column holds the same indices in + decode, so column 0 is what is read. + """ + n_heads = attn_sink_in.shape[1] + s_len = all_q_T.shape[1] // n_heads + out = _gather_attention( + topk_indices_T[:, 0], + all_q_T, + win_K_T, + win_V, + compress_kv, + attn_sink_in, + derope_cos, + derope_sin, + n_heads, + s_len, + ) + return {"output_0": out} + + +def nki_indexer_score_topk_gather_2core_torch_ref( + q_T_all: torch.Tensor, + kv_t: torch.Tensor, + weights: torch.Tensor, + k_val: int, + n_val: int, + all_q_T: torch.Tensor, + win_K_T: torch.Tensor, + win_V: torch.Tensor, + compress_kv: torch.Tensor, + attn_sink_in: torch.Tensor, + derope_cos: torch.Tensor, + derope_sin: torch.Tensor, +) -> dict[str, torch.Tensor]: + """Oracle for the fused decode kernel: indexer score, top-k, and O(k) attention. + + The kernel does all three inside one launch, so nothing intermediate is + observable and the only way to grade it is end-to-end. This reference scores + with ``torch``, selects with ``torch.topk``, and attends densely. + + That makes the test sensitive to top-k ties: if the k-th and (k+1)-th scores are + equal, the kernel and ``torch.topk`` may legitimately pick different positions + and the outputs will differ. The test's input generator is therefore built to + give the scores real separation -- which is also the realistic regime, since a + tied indexer means the selection carries no information. + """ + del n_val + n_heads = attn_sink_in.shape[1] + s_len = all_q_T.shape[1] // n_heads + + scores = _indexer_scores(q_T_all, kv_t, weights) + idx = torch.topk(scores[0].float(), k_val).indices + + out = _gather_attention( + idx, + all_q_T, + win_K_T, + win_V, + compress_kv, + attn_sink_in, + derope_cos, + derope_sin, + n_heads, + s_len, + ) + return {"output_0": out} diff --git a/src/nkilib_src/nkilib/experimental/deepseek_v4_csa/csa_prefill_attention.py b/src/nkilib_src/nkilib/experimental/deepseek_v4_csa/csa_prefill_attention.py new file mode 100644 index 0000000..c4a18fb --- /dev/null +++ b/src/nkilib_src/nkilib/experimental/deepseek_v4_csa/csa_prefill_attention.py @@ -0,0 +1,1499 @@ +# 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 CSA prefill kernels. + +The prefill counterpart of ``csa_decode_attention``. Prefill processes ``S`` +query positions at once rather than one, so the partition axis carries the +SEQUENCE and every kernel here tiles over it; the decode kernels instead put +heads or compressed positions on partitions, because their query is a single +token. + +Five kernels, in the order the block calls them: + +``nki_rms_rope_kernel`` + Fused RMS(+optional learnable gain) + RoPE over a ``[S_rows, head_dim]`` tile. + One kernel covers three call sites via trace-time flags: the q-path + (``gain_in=None``), the kv-path (``gain_in=kv_norm.weight``) and the output + de-RoPE (``do_rms=0, inverse=1``, rotation only). + +``nki_compressor_core_kernel`` + Gated pooling over the ``2 * compress_ratio`` overlapped slots, then RMSNorm + over ``head_dim``, then RoPE -- the compressor that folds raw tokens into + compressed cache positions. + +``nki_indexer_score_mask_kernel`` + Lightning-indexer scoring plus top-k selection, emitting a ``0 / -1e9`` + additive mask. The threshold comes from a fixed number of bisection rounds + rather than a sort, so the whole selection stays on-chip and unrolled. + +``nki_fused_csa_attn_kernel`` + Mask-predicated sparse attention over ``[window | compressed]``, used for the + leading ``split_pos`` positions where every compressed position is still + within the causal frontier. + +``nki_gather_csa_attn_kernel`` + The same math with a per-tile COMPILE-TIME causal bound on the compressed + loop, used for the trailing positions. The static bound is what removes the + sequential dynamic-range device loop, the online-softmax rescaling and all + indirect DMA, leaving every loop unrolled and pipelinable. + +As in the decode file, ``priority=`` is a NeuronCore-v4 (trn3) DMA +class-of-service hint that changes no byte and no MAC. +""" + +import nki +import nki.isa as nisa +import nki.language as nl +from nki.isa.constants import oob_mode + +from ...core.utils.kernel_assert import kernel_assert + + +# -------------------------------------------------------------------------- +# The prefill analogue of csa_nki_model_decode_block.nki_qkv_rms_rope_kernel. +# That kernel is decode-shaped: it packs the n_heads q rows + the 1 kv row of a +# SINGLE token onto one [n_heads+1, head_dim] partition tile. Prefill has S rows +# per head, so the partition dim is the SEQUENCE and the kernel tiles over it, +# processing [TILE_S, head_dim] blocks with cos/sin sliced per tile (decode +# broadcasts one position with a stride-0 .ap()). +# +# -------------------------------------------------------------------------- +@nki.jit +def nki_rms_rope_kernel( + x_in: nl.NkiTensor, + cos_in: nl.NkiTensor, + sin_in: nl.NkiTensor, + gain_in: nl.NkiTensor | None, + eps_val: float, + do_rms: int = 1, + inverse: int = 0, + heads: int = 1, + in_head_major: int = 1, + out_head_major: int = 1, + do_rope: int = 1, +) -> nl.NkiTensor: + """Fused RMS(+optional gain) + RoPE over a [S_rows, head_dim] tile. + + x_in: [S_rows, head_dim] bf16 — rows are (head-major) sequence positions. + cos_in/sin_in: [S, half_rope] fp32 — per-POSITION rotation; with ``heads > 1`` + the table is indexed by position alone, so the caller does not repeat + it per head. + gain_in:[1, head_dim] fp32 or None — learnable RMSNorm gain, broadcast. + Returns:[S_rows, head_dim] bf16 — nope channels passthrough, rope rotated. + + ``heads``/``in_head_major``/``out_head_major`` select the LAYOUT of the multi-head + q and de-RoPE tensors, independently on each side: + + * head-major is ``[heads * S, head_dim]``, row ``h * S + s``; + * query-major (``*_head_major=0``) is ``[S, heads * head_dim]``, row ``s`` and + column block ``h``. + + """ + if in_head_major: + head_dim = x_in.shape[1] + S = x_in.shape[0] // heads + else: + head_dim = x_in.shape[1] // heads + S = x_in.shape[0] + S_rows = S * heads + # do_rope=0 is the norm-only variant (the q-latent RMSNorm, which has no rotation): + # every channel passes through the nope path and cos_in/sin_in are unused, so the + # caller passes None rather than a dummy table. + half_rope = cos_in.shape[1] if do_rope else 0 + rope_head_dim = 2 * half_rope + nope_dim = head_dim - rope_head_dim + TILE = 128 # partition tile (v4 hard cap) + kernel_assert(S % TILE == 0, f"S={S} must be a multiple of {TILE}") + n_tiles = S // TILE + rows = TILE + + if out_head_major: + out = nl.ndarray((S_rows, head_dim), dtype=nl.bfloat16, buffer=nl.shared_hbm) + else: + out = nl.ndarray((S, heads * head_dim), dtype=nl.bfloat16, buffer=nl.shared_hbm) + + if gain_in is not None: + gain = nl.ndarray((TILE, head_dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=gain[0:TILE, 0:head_dim], src=gain_in.ap(pattern=[[0, TILE], [1, head_dim]]), priority=1) + + core_id = nl.program_id(0) + n_cores = nl.num_programs() + if heads % n_cores == 0: + heads_per_core = heads // n_cores + tiles_lo = 0 + tiles_hi = n_tiles + else: + kernel_assert( + n_tiles % n_cores == 0, + f"neither heads={heads} nor n_tiles={n_tiles} divides the {n_cores}-core grid", + ) + heads_per_core = heads + tiles_lo = 0 + tiles_hi = n_tiles // n_cores + + for h_local in nl.affine_range(heads_per_core): + h = core_id * heads_per_core + h_local if heads % n_cores == 0 else h_local + for ts_local in nl.affine_range(tiles_hi - tiles_lo): + ts = ts_local if heads % n_cores == 0 else core_id * tiles_hi + ts_local + s0 = ts * TILE + cos_row = s0 + src_row = h * S + s0 if in_head_major else s0 + src_col = 0 if in_head_major else h * head_dim + dst_row = h * S + s0 if out_head_major else s0 + dst_col = 0 if out_head_major else h * head_dim + # priority=0: this load gates the whole RMS+RoPE chain below. + x_sb = nl.ndarray((TILE, head_dim), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.dma_copy( + dst=x_sb[0:rows, 0:head_dim], + src=x_in[src_row : src_row + rows, src_col : src_col + head_dim], + priority=0, + ) + x_f32 = nl.ndarray((TILE, head_dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=x_f32[0:rows, 0:head_dim], src=x_sb[0:rows, 0:head_dim]) + + if do_rms: + # mean(x^2) over the free axis -> *1/head_dim + eps (fused) -> rsqrt. + x_sq = nl.ndarray((TILE, head_dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=x_sq[0:rows, 0:head_dim], + data1=x_f32[0:rows, 0:head_dim], + data2=x_f32[0:rows, 0:head_dim], + op=nl.multiply, + ) + msq = nl.ndarray((TILE, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_reduce(dst=msq[0:rows, 0:1], data=x_sq[0:rows, 0:head_dim], op=nl.add, axis=1) + nisa.tensor_scalar( + dst=msq[0:rows, 0:1], + data=msq[0:rows, 0:1], + op0=nl.multiply, + operand0=1.0 / head_dim, + op1=nl.add, + operand1=eps_val, + ) + rms = nl.ndarray((TILE, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation(dst=rms[0:rows, 0:1], op=nl.rsqrt, data=msq[0:rows, 0:1]) + nisa.tensor_scalar( + dst=x_f32[0:rows, 0:head_dim], + data=x_f32[0:rows, 0:head_dim], + op0=nl.multiply, + operand0=rms[0:rows, 0:1], + ) + if gain_in is not None: + nisa.tensor_tensor( + dst=x_f32[0:rows, 0:head_dim], + data1=x_f32[0:rows, 0:head_dim], + data2=gain[0:rows, 0:head_dim], + op=nl.multiply, + ) + + # Cast at the RMSNorm output boundary (the reference casts back here). + normed = nl.ndarray((TILE, head_dim), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.tensor_copy(dst=normed[0:rows, 0:head_dim], src=x_f32[0:rows, 0:head_dim]) + + # nope channels pass straight through. + if nope_dim > 0: + nisa.dma_copy( + dst=out[dst_row : dst_row + rows, dst_col : dst_col + nope_dim], src=normed[0:rows, 0:nope_dim] + ) + + if do_rope: + # ---- RoPE on the trailing rope_head_dim channels (fp32 math) ---- + rope_f = nl.ndarray((TILE, rope_head_dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=rope_f[0:rows, 0:rope_head_dim], src=normed[0:rows, nope_dim:head_dim]) + # View as [.., half_rope, 2]: [...,0]=even (x1), [...,1]=odd (x2). + rope_pairs = rope_f.reshape((TILE, half_rope, 2)) + x1 = nl.ndarray((TILE, half_rope), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=x1[0:rows, 0:half_rope], src=rope_pairs[0:rows, 0:half_rope, 0]) + x2 = nl.ndarray((TILE, half_rope), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=x2[0:rows, 0:half_rope], src=rope_pairs[0:rows, 0:half_rope, 1]) + + # priority=2: cos/sin are consumed LAST (only by the rotation), so they + # yield DMA bandwidth to the loads the pipeline stalls on first. + cos_h = nl.ndarray((TILE, half_rope), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=cos_h[0:rows, 0:half_rope], src=cos_in[cos_row : cos_row + rows, 0:half_rope], priority=2) + sin_h = nl.ndarray((TILE, half_rope), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=sin_h[0:rows, 0:half_rope], src=sin_in[cos_row : cos_row + rows, 0:half_rope], priority=2) + + # y1 = x1*cos - x2*sin ; y2 = x1*sin + x2*cos (inverse negates sin, so + # the signs swap: y1 = x1*cos + x2*sin ; y2 = -x1*sin + x2*cos) + tmp_a = nl.ndarray((TILE, half_rope), dtype=nl.float32, buffer=nl.sbuf) + tmp_b = nl.ndarray((TILE, half_rope), dtype=nl.float32, buffer=nl.sbuf) + y1 = nl.ndarray((TILE, half_rope), dtype=nl.float32, buffer=nl.sbuf) + y2 = nl.ndarray((TILE, half_rope), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=tmp_a[0:rows, 0:half_rope], + data1=x1[0:rows, 0:half_rope], + data2=cos_h[0:rows, 0:half_rope], + op=nl.multiply, + ) + nisa.tensor_tensor( + dst=tmp_b[0:rows, 0:half_rope], + data1=x2[0:rows, 0:half_rope], + data2=sin_h[0:rows, 0:half_rope], + op=nl.multiply, + ) + if inverse: + nisa.tensor_tensor( + dst=y1[0:rows, 0:half_rope], + data1=tmp_a[0:rows, 0:half_rope], + data2=tmp_b[0:rows, 0:half_rope], + op=nl.add, + ) + else: + nisa.tensor_tensor( + dst=y1[0:rows, 0:half_rope], + data1=tmp_a[0:rows, 0:half_rope], + data2=tmp_b[0:rows, 0:half_rope], + op=nl.subtract, + ) + nisa.tensor_tensor( + dst=tmp_a[0:rows, 0:half_rope], + data1=x1[0:rows, 0:half_rope], + data2=sin_h[0:rows, 0:half_rope], + op=nl.multiply, + ) + nisa.tensor_tensor( + dst=tmp_b[0:rows, 0:half_rope], + data1=x2[0:rows, 0:half_rope], + data2=cos_h[0:rows, 0:half_rope], + op=nl.multiply, + ) + if inverse: + nisa.tensor_tensor( + dst=y2[0:rows, 0:half_rope], + data1=tmp_b[0:rows, 0:half_rope], + data2=tmp_a[0:rows, 0:half_rope], + op=nl.subtract, + ) + else: + nisa.tensor_tensor( + dst=y2[0:rows, 0:half_rope], + data1=tmp_a[0:rows, 0:half_rope], + data2=tmp_b[0:rows, 0:half_rope], + op=nl.add, + ) + + # Re-interleave y1 (even) / y2 (odd), cast bf16, write out. + rope_out = nl.ndarray((TILE, half_rope, 2), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=rope_out[0:rows, 0:half_rope, 0], src=y1[0:rows, 0:half_rope]) + nisa.tensor_copy(dst=rope_out[0:rows, 0:half_rope, 1], src=y2[0:rows, 0:half_rope]) + rope_flat = rope_out.reshape((TILE, rope_head_dim)) + rope_bf16 = nl.ndarray((TILE, rope_head_dim), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.tensor_copy(dst=rope_bf16[0:rows, 0:rope_head_dim], src=rope_flat[0:rows, 0:rope_head_dim]) + nisa.dma_copy( + dst=out[dst_row : dst_row + rows, dst_col + nope_dim : dst_col + head_dim], + src=rope_bf16[0:rows, 0:rope_head_dim], + ) + + return out + + +# -------------------------------------------------------------------------- +# NKI Kernel: q projection + per-head RMS + RoPE, in ONE existing launch +# -------------------------------------------------------------------------- +@nki.jit +def nki_qb_rms_rope_kernel( + latent_in: nl.NkiTensor, # [S, R] — the normalized q latent (bf16) + w_packed: nl.NkiTensor, # [heads, 128, R/128, head_dim] — wq_b packed per head (bf16) + cos_in: nl.NkiTensor, # [S, half_rope] — per-position cos (fp32) + sin_in: nl.NkiTensor, # [S, half_rope] — per-position sin (fp32) + eps: float, + heads_per_pass: int = 8, +) -> nl.NkiTensor: + """``wq_b`` then per-head RMSNorm then RoPE, emitted query-major. + + This replaces a host ``self.wq_b(qr)`` whose output was already being handed straight + to ``nki_rms_rope_kernel``. Folding the projection into that consumer is the point: the + q path already issues two NKI launches with an XLA GEMM wedged between them, so doing + the GEMM here removes the XLA op and adds NO launch. + + """ + S, R = latent_in.shape + heads, pmax_w, n_r_tiles, head_dim = w_packed.shape + half_rope = cos_in.shape[1] + rope_head_dim = 2 * half_rope + nope_dim = head_dim - rope_head_dim + TILE = 128 + kernel_assert(pmax_w == TILE, f"w_packed partition dim {pmax_w} must be {TILE}") + kernel_assert(R == n_r_tiles * TILE, f"latent R={R} vs w_packed R={n_r_tiles * TILE}") + kernel_assert(S % TILE == 0, f"S={S} must be a multiple of {TILE}") + kernel_assert(heads % heads_per_pass == 0, f"heads={heads} must be a multiple of {heads_per_pass}") + n_tiles = S // TILE + n_passes = heads // heads_per_pass + + out = nl.ndarray((S, heads * head_dim), dtype=nl.bfloat16, buffer=nl.shared_hbm) + + # SPMD over position tiles: each core owns disjoint output rows, no reduction. + core_id = nl.program_id(0) + n_cores = nl.num_programs() + tiles_per_core = (n_tiles + n_cores - 1) // n_cores + + for hp in nl.affine_range(n_passes): + # Weights for this pass's heads, one contiguous DMA each, reused by every tile. + w_sb = [None] * heads_per_pass + for hl in nl.affine_range(heads_per_pass): + h = hp * heads_per_pass + hl + w_sb[hl] = nl.ndarray((TILE, n_r_tiles, head_dim), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.dma_copy(dst=w_sb[hl], src=w_packed[h], priority=1) + + for tl in nl.affine_range(tiles_per_core): + ts = core_id * tiles_per_core + tl + if ts >= n_tiles: + continue + s0 = ts * TILE + + # Latent tile once per (pass, tile), transposed to put R on partitions. + lat_sb = nl.ndarray((TILE, R), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.dma_copy(dst=lat_sb, src=latent_in[s0 : s0 + TILE, 0:R], priority=0) + lat_t = [None] * n_r_tiles + for rt in nl.affine_range(n_r_tiles): + tp = nl.ndarray((TILE, TILE), dtype=nl.bfloat16, buffer=nl.psum) + nisa.nc_transpose(dst=tp, data=lat_sb[0:TILE, rt * TILE : (rt + 1) * TILE]) + lat_t[rt] = nl.ndarray((TILE, TILE), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.tensor_copy(dst=lat_t[rt], src=tp) + + # cos/sin for these positions, shared by every head. + cos_h = nl.ndarray((TILE, half_rope), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=cos_h, src=cos_in[s0 : s0 + TILE, 0:half_rope], priority=2) + sin_h = nl.ndarray((TILE, half_rope), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=sin_h, src=sin_in[s0 : s0 + TILE, 0:half_rope], priority=2) + + for hl in nl.affine_range(heads_per_pass): + h = hp * heads_per_pass + hl + + # ---- projection: [TILE, head_dim] accumulated over R in fp32 PSUM ---- + acc = nl.ndarray((TILE, head_dim), dtype=nl.float32, buffer=nl.psum) + for rt in nl.affine_range(n_r_tiles): + nisa.nc_matmul( + dst=acc, + stationary=lat_t[rt], + moving=w_sb[hl][0:TILE, rt, 0:head_dim], + accumulate=(rt > 0), + ) + # Round to bf16 exactly where the bf16 nn.Linear did, then widen for the + # fp32 norm -- keeps this bit-comparable rather than more precise. + proj_bf = nl.ndarray((TILE, head_dim), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.tensor_copy(dst=proj_bf, src=acc) + x = nl.ndarray((TILE, head_dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=x, src=proj_bf) + + # ---- per-head RMSNorm over head_dim, no learnable gain ---- + sq = nl.ndarray((TILE, head_dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=sq, data1=x, data2=x, op=nl.multiply) + msq = nl.ndarray((TILE, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_reduce(dst=msq, data=sq, op=nl.add, axis=1) + nisa.tensor_scalar( + dst=msq, data=msq, op0=nl.multiply, operand0=1.0 / head_dim, op1=nl.add, operand1=eps + ) + rms = nl.ndarray((TILE, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation(dst=rms, op=nl.rsqrt, data=msq) + normed = nl.ndarray((TILE, head_dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar(dst=normed, data=x, op0=nl.multiply, operand0=rms) + normed_bf = nl.ndarray((TILE, head_dim), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.tensor_copy(dst=normed_bf, src=normed) + + # ---- RoPE on the trailing rope_head_dim channels, fp32 math ---- + row = nl.ndarray((TILE, head_dim), dtype=nl.bfloat16, buffer=nl.sbuf) + if nope_dim > 0: + nisa.tensor_copy(dst=row[0:TILE, 0:nope_dim], src=normed_bf[0:TILE, 0:nope_dim]) + rope_f = nl.ndarray((TILE, rope_head_dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=rope_f, src=normed_bf[0:TILE, nope_dim:head_dim]) + pairs = rope_f.reshape((TILE, half_rope, 2)) + x1 = nl.ndarray((TILE, half_rope), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=x1, src=pairs[0:TILE, 0:half_rope, 0]) + x2 = nl.ndarray((TILE, half_rope), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=x2, src=pairs[0:TILE, 0:half_rope, 1]) + + ta = nl.ndarray((TILE, half_rope), dtype=nl.float32, buffer=nl.sbuf) + tb = nl.ndarray((TILE, half_rope), dtype=nl.float32, buffer=nl.sbuf) + y1 = nl.ndarray((TILE, half_rope), dtype=nl.float32, buffer=nl.sbuf) + y2 = nl.ndarray((TILE, half_rope), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=ta, data1=x1, data2=cos_h, op=nl.multiply) + nisa.tensor_tensor(dst=tb, data1=x2, data2=sin_h, op=nl.multiply) + nisa.tensor_tensor(dst=y1, data1=ta, data2=tb, op=nl.subtract) + nisa.tensor_tensor(dst=ta, data1=x1, data2=sin_h, op=nl.multiply) + nisa.tensor_tensor(dst=tb, data1=x2, data2=cos_h, op=nl.multiply) + nisa.tensor_tensor(dst=y2, data1=ta, data2=tb, op=nl.add) + + rope_out = nl.ndarray((TILE, half_rope, 2), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=rope_out[0:TILE, 0:half_rope, 0], src=y1) + nisa.tensor_copy(dst=rope_out[0:TILE, 0:half_rope, 1], src=y2) + nisa.tensor_copy( + dst=row[0:TILE, nope_dim:head_dim], src=rope_out.reshape((TILE, rope_head_dim)) + ) + + nisa.dma_copy(dst=out[s0 : s0 + TILE, h * head_dim : (h + 1) * head_dim], src=row) + + return out + + +# -------------------------------------------------------------------------- +# NKI Kernel: Compressor gated-pooling + RMSNorm + RoPE core +# -------------------------------------------------------------------------- +@nki.jit +def nki_compressor_core_kernel( + kv8: nl.NkiTensor, # [T_c, ratio2, head_dim] — overlapped kv slots (bf16) + score8: nl.NkiTensor, # [T_c, ratio2, head_dim] — overlapped gate scores, NO ape (bf16) + norm_weight: nl.NkiTensor, # [1, head_dim] — RMSNorm gain (fp32) + cos_rep: nl.NkiTensor, # [T_c, rope_head_dim] — per-pair cos, each value repeated (fp32) + sin_rep: nl.NkiTensor, # [T_c, rope_head_dim] — per-pair sin, each value repeated (fp32) + eps: float, # RMSNorm epsilon + ape: nl.NkiTensor | None = None, # [ratio2, head_dim] — per-slot gate bias (fp32) + hadamard: nl.NkiTensor | None = None, # [head_dim, head_dim] — orthonormal rotation (bf16) +) -> nl.NkiTensor: + """Gated pooling over the size-(2*ratio) axis, RMSNorm over head_dim, then RoPE. + + Per compressed position t and channel c: + w[t, j, c] = softmax_j(score8[t, j, c]) + pooled[t, c] = sum_j kv8[t, j, c] * w[t, j, c] + Then RMSNorm(pooled.to(bf16)) over c, and RoPE on the last rope_head_dim dims. + + Layout: partition = compressed positions (tiled by 128), free = head_dim channels. + The softmax over the slot axis is computed independently per (position, channel). + + When ``hadamard`` is given the whole row is finally rotated by it (the indexer's + compressor does this so the channels it scores are decorrelated). Being orthonormal + it needs the nope and roped halves together, so they are assembled in SBUF and the + row is written once after the rotation instead of half at a time. + + Returns: + out: [T_c, head_dim] bf16 — normalized, roped, optionally rotated compressed kv. + """ + T_c, ratio2, head_dim = kv8.shape + rope_head_dim = cos_rep.shape[1] + nope_dim = head_dim - rope_head_dim + half_rope = rope_head_dim // 2 + + TILE_P = 128 + num_tiles = (T_c + TILE_P - 1) // TILE_P + + out = nl.ndarray((T_c, head_dim), dtype=nl.bfloat16, buffer=nl.shared_hbm) + + # The rotation operand is the same for every position tile, so it is loaded once for + # the whole kernel rather than per tile. It has to fit on partitions because the row + # is transposed to put the contracted channel axis there. + h_sb = None + if hadamard is not None: + kernel_assert(head_dim <= 128, "hadamard rotation requires head_dim <= 128") + h_sb = nl.ndarray((head_dim, head_dim), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.dma_copy(dst=h_sb, src=hadamard[0:head_dim, 0:head_dim]) + + core_id = nl.program_id(0) + n_cores = nl.num_programs() + tiles_per_core = num_tiles // n_cores + + for t_local in nl.affine_range(tiles_per_core): + t_idx = core_id * tiles_per_core + t_local + p_start = t_idx * TILE_P + p_sz = min(TILE_P, T_c - p_start) + + # RMSNorm gain replicated to all partitions of this tile via a + # partition-stride-0 DMA from the single HBM gain row. + gain = nl.ndarray((p_sz, head_dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=gain, src=norm_weight.ap(pattern=[[0, p_sz], [1, head_dim]])) + + kv_slots = [None] * ratio2 + score_slots = [None] * ratio2 + for j in nl.affine_range(ratio2): + kv_bf = nl.ndarray((p_sz, head_dim), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.dma_copy(dst=kv_bf, src=kv8[p_start : p_start + p_sz, j, 0:head_dim]) + kv_slots[j] = nl.ndarray((p_sz, head_dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=kv_slots[j], src=kv_bf) + score_bf = nl.ndarray((p_sz, head_dim), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.dma_copy(dst=score_bf, src=score8[p_start : p_start + p_sz, j, 0:head_dim]) + score_slots[j] = nl.ndarray((p_sz, head_dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=score_slots[j], src=score_bf) + if ape is not None: + # The per-slot gate bias is added HERE, in fp32, so the bf16 handoff above + # stays exact: the host would otherwise add an fp32 parameter to a bf16 + # projection output, promoting the whole tensor to fp32 before the kernel + # ever sees it. Row-invariant, so one stride-0 partition broadcast. + ape_sb = nl.ndarray((p_sz, head_dim), dtype=nl.float32, buffer=nl.sbuf) + ape_row = ape[j : j + 1, 0:head_dim] + nisa.dma_copy(dst=ape_sb, src=ape_row.ap(pattern=[[0, p_sz], [1, head_dim]])) + nisa.tensor_tensor(dst=score_slots[j], data1=score_slots[j], data2=ape_sb, op=nl.add) + + # --- Softmax over the slot axis (per position & channel) --- + # Elementwise max across the ratio2 slots. + slot_max = nl.ndarray((p_sz, head_dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=slot_max, src=score_slots[0]) + for j in range(1, ratio2): + nisa.tensor_tensor(dst=slot_max, data1=slot_max, data2=score_slots[j], op=nl.maximum) + + # exp(score - max) per slot, and accumulate the denominator. + neg_max = nl.ndarray((p_sz, head_dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar(dst=neg_max, data=slot_max, op0=nl.multiply, operand0=-1.0) + + exp_slots = [None] * ratio2 + denom = nl.ndarray((p_sz, head_dim), dtype=nl.float32, buffer=nl.sbuf) + for j in range(ratio2): + exp_slots[j] = nl.ndarray((p_sz, head_dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=exp_slots[j], data1=score_slots[j], data2=neg_max, op=nl.add) + nisa.activation(dst=exp_slots[j], op=nl.exp, data=exp_slots[j]) + if j == 0: + nisa.tensor_copy(dst=denom, src=exp_slots[0]) + else: + nisa.tensor_tensor(dst=denom, data1=denom, data2=exp_slots[j], op=nl.add) + + inv_denom = nl.ndarray((p_sz, head_dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.reciprocal(dst=inv_denom, data=denom) + + # --- Weighted sum: pooled = sum_j kv_j * (exp_j / denom) --- + pooled = nl.ndarray((p_sz, head_dim), dtype=nl.float32, buffer=nl.sbuf) + prod = nl.ndarray((p_sz, head_dim), dtype=nl.float32, buffer=nl.sbuf) + for j in range(ratio2): + # weight_j = exp_j * inv_denom (softmax); pooled += kv_j * weight_j + nisa.tensor_tensor(dst=prod, data1=exp_slots[j], data2=inv_denom, op=nl.multiply) + nisa.tensor_tensor(dst=prod, data1=prod, data2=kv_slots[j], op=nl.multiply) + if j == 0: + nisa.tensor_copy(dst=pooled, src=prod) + else: + nisa.tensor_tensor(dst=pooled, data1=pooled, data2=prod, op=nl.add) + + # --- Cast pooled to bf16 (matches reference kv.to(bf16) before norm) --- + pooled_bf16 = nl.ndarray((p_sz, head_dim), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.tensor_copy(dst=pooled_bf16, src=pooled) + # Re-widen to fp32 for the norm compute (reference RMSNorm computes in fp32). + x_norm = nl.ndarray((p_sz, head_dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=x_norm, src=pooled_bf16) + + # --- RMSNorm over head_dim (free axis) --- + sq = nl.ndarray((p_sz, head_dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=sq, data1=x_norm, data2=x_norm, op=nl.multiply) + msq = nl.ndarray((p_sz, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_reduce(dst=msq, data=sq, op=nl.add, axis=1) + # mean = sum / head_dim, then rsqrt(mean + eps). + nisa.tensor_scalar(dst=msq, data=msq, op0=nl.multiply, operand0=1.0 / head_dim, op1=nl.add, operand1=eps) + rms = nl.ndarray((p_sz, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation(dst=rms, op=nl.rsqrt, data=msq) + # normed = x * rms (broadcast over free) * gain (broadcast over partition). + normed = nl.ndarray((p_sz, head_dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar(dst=normed, data=x_norm, op0=nl.multiply, operand0=rms) + nisa.tensor_tensor(dst=normed, data1=normed, data2=gain, op=nl.multiply) + + # Cast to bf16 (RMSNorm output dtype), then RoPE reads bf16 -> fp32. + normed_bf16 = nl.ndarray((p_sz, head_dim), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.tensor_copy(dst=normed_bf16, src=normed) + + # --- The nope part (channels 0..nope_dim-1) --- + # Unrotated it goes straight to HBM. Rotated it must meet the roped half first, + # so it is staged in an SBUF row that the matmul below consumes whole. + full_bf16 = None + if hadamard is None: + nisa.dma_copy(dst=out[p_start : p_start + p_sz, 0:nope_dim], src=normed_bf16[0:p_sz, 0:nope_dim]) + else: + full_bf16 = nl.ndarray((p_sz, head_dim), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.tensor_copy(dst=full_bf16[0:p_sz, 0:nope_dim], src=normed_bf16[0:p_sz, 0:nope_dim]) + + # --- RoPE on the last rope_head_dim channels --- + # Load cos/sin (already repeated per pair) for these positions. + cos_sb = nl.ndarray((p_sz, rope_head_dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=cos_sb, src=cos_rep[p_start : p_start + p_sz, 0:rope_head_dim]) + sin_sb = nl.ndarray((p_sz, rope_head_dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=sin_sb, src=sin_rep[p_start : p_start + p_sz, 0:rope_head_dim]) + + # Widen the rope channels back to fp32 (reference computes RoPE in fp32). + rope_f = nl.ndarray((p_sz, rope_head_dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=rope_f, src=normed_bf16[0:p_sz, nope_dim:head_dim]) + # View as [p_sz, half_rope, 2] so [...,0]=even (x1), [...,1]=odd (x2). + rope_pairs = rope_f.reshape((p_sz, half_rope, 2)) + x1 = nl.ndarray((p_sz, half_rope), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=x1, src=rope_pairs[0:p_sz, 0:half_rope, 0]) + x2 = nl.ndarray((p_sz, half_rope), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=x2, src=rope_pairs[0:p_sz, 0:half_rope, 1]) + + cos_pairs = cos_sb.reshape((p_sz, half_rope, 2)) + sin_pairs = sin_sb.reshape((p_sz, half_rope, 2)) + cos_h = nl.ndarray((p_sz, half_rope), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=cos_h, src=cos_pairs[0:p_sz, 0:half_rope, 0]) + sin_h = nl.ndarray((p_sz, half_rope), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=sin_h, src=sin_pairs[0:p_sz, 0:half_rope, 0]) + + # y1 = x1*cos - x2*sin ; y2 = x1*sin + x2*cos + tmp_a = nl.ndarray((p_sz, half_rope), dtype=nl.float32, buffer=nl.sbuf) + tmp_b = nl.ndarray((p_sz, half_rope), dtype=nl.float32, buffer=nl.sbuf) + y1 = nl.ndarray((p_sz, half_rope), dtype=nl.float32, buffer=nl.sbuf) + y2 = nl.ndarray((p_sz, half_rope), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=tmp_a, data1=x1, data2=cos_h, op=nl.multiply) + nisa.tensor_tensor(dst=tmp_b, data1=x2, data2=sin_h, op=nl.multiply) + nisa.tensor_tensor(dst=y1, data1=tmp_a, data2=tmp_b, op=nl.subtract) + nisa.tensor_tensor(dst=tmp_a, data1=x1, data2=sin_h, op=nl.multiply) + nisa.tensor_tensor(dst=tmp_b, data1=x2, data2=cos_h, op=nl.multiply) + nisa.tensor_tensor(dst=y2, data1=tmp_a, data2=tmp_b, op=nl.add) + + # Re-interleave y1 (even) and y2 (odd) into [p_sz, half_rope, 2] then bf16. + rope_out = nl.ndarray((p_sz, half_rope, 2), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=rope_out[0:p_sz, 0:half_rope, 0], src=y1) + nisa.tensor_copy(dst=rope_out[0:p_sz, 0:half_rope, 1], src=y2) + rope_out_flat = rope_out.reshape((p_sz, rope_head_dim)) + rope_out_bf16 = nl.ndarray((p_sz, rope_head_dim), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.tensor_copy(dst=rope_out_bf16, src=rope_out_flat) + + if hadamard is None: + nisa.dma_copy(dst=out[p_start : p_start + p_sz, nope_dim:head_dim], src=rope_out_bf16) + else: + nisa.tensor_copy(dst=full_bf16[0:p_sz, nope_dim:head_dim], src=rope_out_bf16) + row_t_psum = nl.ndarray((head_dim, p_sz), dtype=nl.bfloat16, buffer=nl.psum) + nisa.nc_transpose(dst=row_t_psum, data=full_bf16[0:p_sz, 0:head_dim]) + row_t = nl.ndarray((head_dim, p_sz), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.tensor_copy(dst=row_t, src=row_t_psum) + rot_psum = nl.ndarray((p_sz, head_dim), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=rot_psum, stationary=row_t, moving=h_sb) + rot_bf16 = nl.ndarray((p_sz, head_dim), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.tensor_copy(dst=rot_bf16, src=rot_psum) + nisa.dma_copy(dst=out[p_start : p_start + p_sz, 0:head_dim], src=rot_bf16) + + return out + + +# -------------------------------------------------------------------------- +# NKI Kernel: Indexer per-head scoring + binary-search top-k mask +# -------------------------------------------------------------------------- +@nki.jit +def nki_indexer_score_mask_kernel( + q_T_all: nl.NkiTensor, # [head_dim, n_heads * S_q] — all heads' Q^T stacked (bf16) + kv_t: nl.NkiTensor, # [head_dim, T_c] — indexer_kv transposed (bf16), shared + weights: nl.NkiTensor, # [S_q, n_heads] — per-row per-head weights * weight_scale (fp32) + causal_bias: nl.NkiTensor, # [S_q, T_c] — causal bias (0 / -1e9) (fp32) + k: int, # top-k count +) -> nl.NkiTensor: + """Indexer scoring + top-k selection mask in a single NKI kernel. + + Computes, per query row s and compressed kv position t: + index_score[s, t] = sum_h relu(q[s, h, :] . kv[t, :]) * weights[s, h] + + causal_bias[s, t] + then finds, per row, a threshold via 10 iterations of bisection, and builds: + sel_mask[s, t] = 0 if index_score[s, t] >= threshold[s] + -1e9 otherwise + + Returns: + sel_mask: [S_q, T_c] fp32 — selection mask (0 / -1e9). + """ + head_dim = q_T_all.shape[0] + total_q_free = q_T_all.shape[1] + T_c = kv_t.shape[1] + n_heads = weights.shape[1] + S_q = total_q_free // n_heads + + TILE_Q = 128 + SCORE_CHUNK = 512 if T_c >= 512 else T_c + num_score_chunks = (T_c + SCORE_CHUNK - 1) // SCORE_CHUNK + num_q_tiles = (S_q + TILE_Q - 1) // TILE_Q + + NEG_INF = -1e9 + + sel_mask = nl.ndarray((S_q, T_c), dtype=nl.float32, buffer=nl.shared_hbm) + + kv_t_sb = nl.ndarray((head_dim, T_c), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.dma_copy(dst=kv_t_sb, src=kv_t[0:head_dim, 0:T_c]) + + core_id = nl.program_id(0) + n_cores = nl.num_programs() + tiles_per_core = num_q_tiles // n_cores + + for q_local in nl.affine_range(tiles_per_core): + q_idx = core_id * tiles_per_core + q_local + q_start = q_idx * TILE_Q + + w_tile = nl.ndarray((TILE_Q, n_heads), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=w_tile, src=weights[q_start : q_start + TILE_Q, 0:n_heads]) + + index_score_bf16 = nl.ndarray((TILE_Q, T_c), dtype=nl.bfloat16, buffer=nl.sbuf) + + nisa.memset(dst=index_score_bf16, value=0) + for h in nl.affine_range(n_heads): + q_global = h * S_q + q_start + q_T = nl.ndarray((head_dim, TILE_Q), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.dma_copy(dst=q_T, src=q_T_all[0:head_dim, q_global : q_global + TILE_Q]) + + w_h = nl.ndarray((TILE_Q, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=w_h, src=w_tile[0:TILE_Q, h : h + 1]) + + for m_idx in nl.affine_range(num_score_chunks): + m_start = m_idx * SCORE_CHUNK + kt_slice = kv_t_sb[0:head_dim, m_start : m_start + SCORE_CHUNK] + s_psum = nl.ndarray((TILE_Q, SCORE_CHUNK), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=s_psum, stationary=q_T, moving=kt_slice) + s_relu = nl.ndarray((TILE_Q, SCORE_CHUNK), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.activation(dst=s_relu, op=nl.relu, data=s_psum) + nisa.scalar_tensor_tensor( + dst=index_score_bf16[0:TILE_Q, m_start : m_start + SCORE_CHUNK], + data=s_relu, + op0=nl.multiply, + operand0=w_h, + op1=nl.add, + operand1=index_score_bf16[0:TILE_Q, m_start : m_start + SCORE_CHUNK], + ) + + index_score = nl.ndarray((TILE_Q, T_c), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=index_score, src=index_score_bf16) + + cbias = nl.ndarray((TILE_Q, T_c), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=cbias, src=causal_bias[q_start : q_start + TILE_Q, 0:T_c]) + nisa.tensor_tensor(dst=index_score, data1=index_score, data2=cbias, op=nl.add) + + # Binary search for the per-row threshold. + hi = nl.ndarray((TILE_Q, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_reduce(dst=hi, op=nl.maximum, data=index_score, axis=1) + is_valid = nl.ndarray((TILE_Q, T_c), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar(dst=is_valid, data=index_score, op0=nl.greater, operand0=-1e8) + score_minus_hi = nl.ndarray((TILE_Q, T_c), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar(dst=score_minus_hi, data=index_score, op0=nl.subtract, operand0=hi) + masked_for_min = nl.ndarray((TILE_Q, T_c), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=masked_for_min, data1=score_minus_hi, data2=is_valid, op=nl.multiply) + nisa.tensor_scalar(dst=masked_for_min, data=masked_for_min, op0=nl.add, operand0=hi) + lo = nl.ndarray((TILE_Q, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_reduce(dst=lo, op=nl.minimum, data=masked_for_min, axis=1) + + mid = nl.ndarray((TILE_Q, 1), dtype=nl.float32, buffer=nl.sbuf) + count = nl.ndarray((TILE_Q, 1), dtype=nl.float32, buffer=nl.sbuf) + ge_mid = nl.ndarray((TILE_Q, T_c), dtype=nl.float32, buffer=nl.sbuf) + ge_count_flag = nl.ndarray((TILE_Q, 1), dtype=nl.float32, buffer=nl.sbuf) + for _it in range(9): + nisa.tensor_tensor(dst=mid, data1=lo, data2=hi, op=nl.add) + nisa.tensor_scalar(dst=mid, data=mid, op0=nl.multiply, operand0=0.5) + nisa.tensor_scalar_reduce( + dst=ge_mid, data=index_score, op0=nl.greater_equal, operand0=mid, reduce_op=nl.add, reduce_res=count + ) + nisa.tensor_scalar(dst=ge_count_flag, data=count, op0=nl.greater_equal, operand0=float(k)) + mid_minus_lo = nl.ndarray((TILE_Q, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=mid_minus_lo, data1=mid, data2=lo, op=nl.subtract) + nisa.tensor_tensor(dst=mid_minus_lo, data1=mid_minus_lo, data2=ge_count_flag, op=nl.multiply) + nisa.tensor_tensor(dst=lo, data1=lo, data2=mid_minus_lo, op=nl.add) + lt_count_flag = nl.ndarray((TILE_Q, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=lt_count_flag, data=ge_count_flag, op0=nl.multiply, operand0=-1.0, op1=nl.add, operand1=1.0 + ) + mid_minus_hi = nl.ndarray((TILE_Q, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=mid_minus_hi, data1=mid, data2=hi, op=nl.subtract) + nisa.tensor_tensor(dst=mid_minus_hi, data1=mid_minus_hi, data2=lt_count_flag, op=nl.multiply) + nisa.tensor_tensor(dst=hi, data1=hi, data2=mid_minus_hi, op=nl.add) + + # Build mask: sel = (score >= lo) ? 0 : -1e9 + sel = nl.ndarray((TILE_Q, T_c), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar(dst=sel, data=index_score, op0=nl.greater_equal, operand0=lo) + nisa.tensor_scalar(dst=sel, data=sel, op0=nl.multiply, operand0=-NEG_INF, op1=nl.add, operand1=NEG_INF) + nisa.dma_copy(dst=sel_mask[q_start : q_start + TILE_Q, 0:T_c], src=sel) + + return sel_mask + + +# -------------------------------------------------------------------------- +# NKI Kernel: Fused CSA Attention with PSum Accumulation + V Preloading +# -------------------------------------------------------------------------- + + +@nki.jit +def nki_fused_csa_attn_kernel( + compress_sel: nl.NkiTensor, # [S, T_c] — selection mask (0/-inf), shared across heads + all_q_T: nl.NkiTensor, # [head_dim, n_heads * S] — all heads' Q^T stacked + all_K_T: nl.NkiTensor, # [head_dim, S + W + T_c] — concatenated K^T (shared) + all_V: nl.NkiTensor, # [S + W + T_c, head_dim] — concatenated V (shared) + win_bias_base_in: nl.NkiTensor, # [S, 256] — base window bias (0/-1e9), shared across heads + win_bias_sink_in: nl.NkiTensor, # [S, 256] — sink indicator (0/1), shared across heads + attn_sink_in: nl.NkiTensor, # [1, n_heads] — per-head sink scalars +) -> nl.NkiTensor: + S, T_c = compress_sel.shape + head_dim = all_q_T.shape[0] + total_q_free = all_q_T.shape[1] + n_heads = total_q_free // S + W = 128 + TILE_Q = 128 + KV_CHUNK = 128 + COMP_V_CHUNK = min(KV_CHUNK, T_c) + SCORE_CHUNK = min(512, T_c) + WIN_SIZE = 2 * KV_CHUNK + num_q_tiles = S // TILE_Q + num_c_chunks = T_c // COMP_V_CHUNK + num_score_chunks = T_c // SCORE_CHUNK + H_BATCH = 16 if T_c >= 512 else 8 + num_h_batches = n_heads // H_BATCH + + HD_CHUNK = 128 + HD_TILES = (head_dim + HD_CHUNK - 1) // HD_CHUNK + + output = nl.ndarray((n_heads * S, head_dim), dtype=nl.bfloat16, buffer=nl.shared_hbm) + + kv_comp_offset = S + W + + all_comp_kt = [None] * HD_TILES + for hd in range(HD_TILES): + hd_start = hd * HD_CHUNK + hd_sz = min(HD_CHUNK, head_dim - hd_start) + all_comp_kt[hd] = nl.ndarray((hd_sz, T_c), dtype=nl.float16, buffer=nl.sbuf) + nisa.dma_copy( + dst=all_comp_kt[hd], src=all_K_T[hd_start : hd_start + hd_sz, kv_comp_offset : kv_comp_offset + T_c] + ) + + comp_v = [] + for i in range(num_c_chunks): + c_start = i * COMP_V_CHUNK + v_tile = nl.ndarray((COMP_V_CHUNK, head_dim), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.dma_copy( + dst=v_tile, src=all_V[kv_comp_offset + c_start : kv_comp_offset + c_start + COMP_V_CHUNK, 0:head_dim] + ) + comp_v.append(v_tile) + + # Preload per-head attn_sink scalars, replicated to all TILE_Q partitions + # using partition-stride-0 DMA so each partition has the full n_heads vector. + attn_sink_sb = nl.ndarray((TILE_Q, n_heads), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=attn_sink_sb, src=attn_sink_in.ap(pattern=[[0, TILE_Q], [1, n_heads]])) + + core_id = nl.program_id(0) + n_cores = nl.num_programs() + tiles_per_core = num_q_tiles // n_cores + + for q_local in nl.affine_range(tiles_per_core): + q_idx = core_id * tiles_per_core + q_local + q_start = q_idx * TILE_Q + + comp_mask = nl.ndarray((TILE_Q, T_c), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.dma_copy(dst=comp_mask, src=compress_sel[q_start : q_start + TILE_Q, 0:T_c]) + + kv_t_win = [None] * HD_TILES + for hd in range(HD_TILES): + hd_start = hd * HD_CHUNK + hd_sz = min(HD_CHUNK, head_dim - hd_start) + kv_t_win[hd] = nl.ndarray((hd_sz, WIN_SIZE), dtype=nl.float16, buffer=nl.sbuf) + nisa.dma_copy(dst=kv_t_win[hd], src=all_K_T[hd_start : hd_start + hd_sz, q_start : q_start + WIN_SIZE]) + + win_v_0 = nl.ndarray((KV_CHUNK, head_dim), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.dma_copy(dst=win_v_0, src=all_V[q_start : q_start + KV_CHUNK, 0:head_dim]) + win_v_1 = nl.ndarray((KV_CHUNK, head_dim), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.dma_copy(dst=win_v_1, src=all_V[q_start + KV_CHUNK : q_start + WIN_SIZE, 0:head_dim]) + + # Load base and sink_ind ONCE per tile (shared across all heads). + bias_base_tile = nl.ndarray((TILE_Q, WIN_SIZE), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=bias_base_tile, src=win_bias_base_in[q_start : q_start + TILE_Q, 0:WIN_SIZE]) + bias_sink_tile = nl.ndarray((TILE_Q, WIN_SIZE), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=bias_sink_tile, src=win_bias_sink_in[q_start : q_start + TILE_Q, 0:WIN_SIZE]) + + for hb_idx in nl.affine_range(num_h_batches): + q_T = [None] * H_BATCH + for h_local in nl.affine_range(H_BATCH): + h = hb_idx * H_BATCH + h_local + q_global = h * S + q_start + q_T[h_local] = [None] * HD_TILES + for hd in nl.affine_range(HD_TILES): + hd_start = hd * HD_CHUNK + hd_sz = min(HD_CHUNK, head_dim - hd_start) + q_T[h_local][hd] = nl.ndarray((hd_sz, TILE_Q), dtype=nl.float16, buffer=nl.sbuf) + nisa.dma_copy( + dst=q_T[h_local][hd], src=all_q_T[hd_start : hd_start + hd_sz, q_global : q_global + TILE_Q] + ) + + # === Compute window biases inline for H_BATCH heads === + win_bias = [None] * H_BATCH + for h_local in nl.affine_range(H_BATCH): + h = hb_idx * H_BATCH + h_local + win_bias[h_local] = nl.ndarray((TILE_Q, WIN_SIZE), dtype=nl.float32, buffer=nl.sbuf) + nisa.scalar_tensor_tensor( + dst=win_bias[h_local], + data=bias_sink_tile, + op0=nl.multiply, + operand0=attn_sink_sb[0:TILE_Q, h : h + 1], + op1=nl.add, + operand1=bias_base_tile, + ) + + # === Compute window scores for H_BATCH heads === + win_scores = [None] * H_BATCH + for h_local in nl.affine_range(H_BATCH): + win_scores_psum = nl.ndarray((TILE_Q, WIN_SIZE), dtype=nl.float32, buffer=nl.psum) + for hd in nl.affine_range(HD_TILES): + nisa.nc_matmul(dst=win_scores_psum, stationary=q_T[h_local][hd], moving=kv_t_win[hd]) + win_scores[h_local] = nl.ndarray((TILE_Q, WIN_SIZE), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=win_scores[h_local], data1=win_scores_psum, data2=win_bias[h_local], op=nl.add) + + # === Compute compressed scores for H_BATCH heads === + comp_scores = [None] * H_BATCH + for h_local in nl.affine_range(H_BATCH): + comp_scores[h_local] = nl.ndarray((TILE_Q, T_c), dtype=nl.float32, buffer=nl.sbuf) + for m_idx in nl.affine_range(num_score_chunks): + m_start = m_idx * SCORE_CHUNK + scores_chunk_psum = nl.ndarray((TILE_Q, SCORE_CHUNK), dtype=nl.float32, buffer=nl.psum) + for hd in nl.affine_range(HD_TILES): + kt_slice = all_comp_kt[hd][0 : all_comp_kt[hd].shape[0], m_start : m_start + SCORE_CHUNK] + nisa.nc_matmul(dst=scores_chunk_psum, stationary=q_T[h_local][hd], moving=kt_slice) + nisa.tensor_tensor( + dst=comp_scores[h_local][0:TILE_Q, m_start : m_start + SCORE_CHUNK], + data1=scores_chunk_psum, + data2=comp_mask[0:TILE_Q, m_start : m_start + SCORE_CHUNK], + op=nl.add, + ) + + # === Unified global max softmax for H_BATCH heads === + total_sum = [None] * H_BATCH + win_exp = [None] * H_BATCH + comp_exp = [None] * H_BATCH + for h_local in nl.affine_range(H_BATCH): + win_max = nl.ndarray((TILE_Q, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_reduce(dst=win_max, data=win_scores[h_local], op=nl.maximum, axis=1) + comp_max = nl.ndarray((TILE_Q, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_reduce(dst=comp_max, data=comp_scores[h_local], op=nl.maximum, axis=1) + neg_max = nl.ndarray((TILE_Q, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=neg_max, data1=win_max, data2=comp_max, op=nl.maximum) + nisa.tensor_scalar(dst=neg_max, data=neg_max, op0=nl.multiply, operand0=-1.0) + + win_sum = nl.ndarray((TILE_Q, 1), dtype=nl.float32, buffer=nl.sbuf) + win_exp[h_local] = nl.ndarray((TILE_Q, WIN_SIZE), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.activation( + dst=win_exp[h_local], + op=nl.exp, + data=win_scores[h_local], + bias=neg_max, + reduce_op=nl.add, + reduce_res=win_sum, + reduce_cmd=nisa.reduce_cmd.reset_reduce, + ) + comp_sum = nl.ndarray((TILE_Q, 1), dtype=nl.float32, buffer=nl.sbuf) + comp_exp[h_local] = nl.ndarray((TILE_Q, T_c), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.activation( + dst=comp_exp[h_local], + op=nl.exp, + data=comp_scores[h_local], + bias=neg_max, + reduce_op=nl.add, + reduce_res=comp_sum, + reduce_cmd=nisa.reduce_cmd.reset_reduce, + ) + total_sum[h_local] = nl.ndarray((TILE_Q, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=total_sum[h_local], data1=win_sum, data2=comp_sum, op=nl.add) + + # === V multiply with PSum accumulation for H_BATCH heads === + for h_local in nl.affine_range(H_BATCH): + out_psum = nl.ndarray((TILE_Q, head_dim), dtype=nl.float32, buffer=nl.psum) + + scores_T_psum = nl.ndarray((KV_CHUNK, TILE_Q), dtype=nl.bfloat16, buffer=nl.psum) + nisa.nc_transpose(dst=scores_T_psum, data=win_exp[h_local][0:TILE_Q, 0:KV_CHUNK]) + scores_T_sb = nl.ndarray((KV_CHUNK, TILE_Q), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.tensor_copy(dst=scores_T_sb, src=scores_T_psum) + nisa.nc_matmul(dst=out_psum, stationary=scores_T_sb, moving=win_v_0) + + scores_T_psum = nl.ndarray((KV_CHUNK, TILE_Q), dtype=nl.bfloat16, buffer=nl.psum) + nisa.nc_transpose(dst=scores_T_psum, data=win_exp[h_local][0:TILE_Q, KV_CHUNK:WIN_SIZE]) + scores_T_sb = nl.ndarray((KV_CHUNK, TILE_Q), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.tensor_copy(dst=scores_T_sb, src=scores_T_psum) + nisa.nc_matmul(dst=out_psum, stationary=scores_T_sb, moving=win_v_1) + + for c_idx in nl.affine_range(num_c_chunks): + c_start = c_idx * COMP_V_CHUNK + scores_T_psum = nl.ndarray((COMP_V_CHUNK, TILE_Q), dtype=nl.bfloat16, buffer=nl.psum) + nisa.nc_transpose( + dst=scores_T_psum, data=comp_exp[h_local][0:TILE_Q, c_start : c_start + COMP_V_CHUNK] + ) + scores_T_sb = nl.ndarray((COMP_V_CHUNK, TILE_Q), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.tensor_copy(dst=scores_T_sb, src=scores_T_psum) + nisa.nc_matmul(dst=out_psum, stationary=scores_T_sb, moving=comp_v[c_idx]) + + # Finalize: single copy from PSum, divide by sum, write output + out_sbuf = nl.ndarray((TILE_Q, head_dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=out_sbuf, src=out_psum) + + inv_sum = nl.ndarray((TILE_Q, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation(dst=inv_sum, op=nl.reciprocal, data=total_sum[h_local]) + nisa.tensor_scalar(dst=out_sbuf, data=out_sbuf, op0=nl.multiply, operand0=inv_sum) + + h = hb_idx * H_BATCH + h_local + q_global = h * S + q_start + out_bf16 = nl.ndarray((TILE_Q, head_dim), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.tensor_copy(dst=out_bf16, src=out_sbuf) + nisa.dma_copy(dst=output[q_global : q_global + TILE_Q, 0:head_dim], src=out_bf16) + + return output + + +# -------------------------------------------------------------------------- +# NKI Kernel: Full gather-based sparse attention (indirect DMA) +# -------------------------------------------------------------------------- + + +@nki.jit +def nki_gather_csa_attn_kernel( + topk_sel_bias: nl.NkiTensor, # [S, T_c] bfloat16 — selection bias (0/-inf), causal-masked + all_q_T: nl.NkiTensor, # [head_dim, n_heads * S] — all heads' Q^T stacked (bf16) + all_K_T_win: nl.NkiTensor, # [head_dim, S + W] — window K^T (padded) + all_V_win: nl.NkiTensor, # [S + W, head_dim] — window V (padded) + compress_kv_T: nl.NkiTensor, # [head_dim, T_c] bf16 — compressed KV transposed (for scoring) + compress_kv: nl.NkiTensor, # [T_c, head_dim] bf16 — compressed KV row-major (for V) + win_bias_base_in: nl.NkiTensor, # [S, 256] — base window bias + win_bias_sink_in: nl.NkiTensor, # [S, 256] — sink indicator + attn_sink_in: nl.NkiTensor, # [1, n_heads] — per-head sink scalars + split_pos: int, # global position offset of this second half + ratio: int, # compression ratio +) -> nl.NkiTensor: + """Sparse attention with static causal-bound + global-max softmax + mask predication.""" + S = topk_sel_bias.shape[0] + head_dim = all_q_T.shape[0] + T_c = compress_kv_T.shape[1] + n_heads = all_q_T.shape[1] // S + TILE_Q = 128 + KV_CHUNK = 128 + COMP_V_CHUNK = min(KV_CHUNK, T_c) + SCORE_CHUNK = min(512, T_c) + WIN_SIZE = 2 * KV_CHUNK + num_q_tiles = S // TILE_Q + num_c_chunks = T_c // COMP_V_CHUNK + H_BATCH = 16 + num_h_batches = n_heads // H_BATCH + HD_CHUNK = 128 + HD_TILES = (head_dim + HD_CHUNK - 1) // HD_CHUNK + + output = nl.ndarray((n_heads * S, head_dim), dtype=nl.bfloat16, buffer=nl.shared_hbm) + + attn_sink_sb = nl.ndarray((TILE_Q, n_heads), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=attn_sink_sb, src=attn_sink_in.ap(pattern=[[0, TILE_Q], [1, n_heads]])) + + # Preload full compressed K^T (shared across all Q tiles and heads). + all_comp_kt = [None] * HD_TILES + for hd in range(HD_TILES): + hd_start = hd * HD_CHUNK + hd_sz = min(HD_CHUNK, head_dim - hd_start) + all_comp_kt[hd] = nl.ndarray((hd_sz, T_c), dtype=nl.float16, buffer=nl.sbuf) + nisa.dma_copy(dst=all_comp_kt[hd], src=compress_kv_T[hd_start : hd_start + hd_sz, 0:T_c]) + + # Preload full compressed V chunks (shared across all Q tiles and heads). + comp_v = [None] * num_c_chunks + for i in range(num_c_chunks): + c_start = i * COMP_V_CHUNK + comp_v[i] = nl.ndarray((COMP_V_CHUNK, head_dim), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.dma_copy(dst=comp_v[i], src=compress_kv[c_start : c_start + COMP_V_CHUNK, 0:head_dim]) + + core_id = nl.program_id(0) + n_cores = nl.num_programs() + # Split work on head-batch dimension so both cores process the same Q tile. + hb_per_core = num_h_batches // n_cores + + for q_idx in nl.static_range(num_q_tiles): + q_start = q_idx * TILE_Q + + # Per-tile compile-time causal bound (CEIL-DIV — rounding up is always safe; + # extra chunks are fully -1e9-masked in sel_bias so they contribute 0). + global_q_end = split_pos + (q_idx + 1) * TILE_Q + causal_chunks = min(num_c_chunks, (global_q_end + (ratio * COMP_V_CHUNK) - 1) // (ratio * COMP_V_CHUNK)) + comp_cols = causal_chunks * COMP_V_CHUNK + num_score_chunks = (comp_cols + SCORE_CHUNK - 1) // SCORE_CHUNK + + # Load sel_bias for this tile's causal columns only (plain contiguous DMA). + comp_mask = nl.ndarray((TILE_Q, comp_cols), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.dma_copy(dst=comp_mask, src=topk_sel_bias[q_start : q_start + TILE_Q, 0:comp_cols]) + + # Load window data (shared across heads for this Q tile). + kv_t_win = [None] * HD_TILES + for hd in nl.affine_range(HD_TILES): + hd_start = hd * HD_CHUNK + hd_sz = min(HD_CHUNK, head_dim - hd_start) + kv_t_win[hd] = nl.ndarray((hd_sz, WIN_SIZE), dtype=nl.float16, buffer=nl.sbuf) + nisa.dma_copy(dst=kv_t_win[hd], src=all_K_T_win[hd_start : hd_start + hd_sz, q_start : q_start + WIN_SIZE]) + + win_v_0 = nl.ndarray((KV_CHUNK, head_dim), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.dma_copy(dst=win_v_0, src=all_V_win[q_start : q_start + KV_CHUNK, 0:head_dim]) + win_v_1 = nl.ndarray((KV_CHUNK, head_dim), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.dma_copy(dst=win_v_1, src=all_V_win[q_start + KV_CHUNK : q_start + WIN_SIZE, 0:head_dim]) + + bias_base_tile = nl.ndarray((TILE_Q, WIN_SIZE), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=bias_base_tile, src=win_bias_base_in[q_start : q_start + TILE_Q, 0:WIN_SIZE]) + bias_sink_tile = nl.ndarray((TILE_Q, WIN_SIZE), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=bias_sink_tile, src=win_bias_sink_in[q_start : q_start + TILE_Q, 0:WIN_SIZE]) + + for hb_local in nl.affine_range(hb_per_core): + hb_idx = core_id * hb_per_core + hb_local + q_T = [None] * H_BATCH + for h_local in nl.affine_range(H_BATCH): + h = hb_idx * H_BATCH + h_local + q_global = h * S + q_start + q_T[h_local] = [None] * HD_TILES + for hd in nl.affine_range(HD_TILES): + hd_start = hd * HD_CHUNK + hd_sz = min(HD_CHUNK, head_dim - hd_start) + q_T[h_local][hd] = nl.ndarray((hd_sz, TILE_Q), dtype=nl.float16, buffer=nl.sbuf) + nisa.dma_copy( + dst=q_T[h_local][hd], src=all_q_T[hd_start : hd_start + hd_sz, q_global : q_global + TILE_Q] + ) + + win_exp_all = [None] * H_BATCH + comp_exp_all = [None] * H_BATCH + total_sum_all = [None] * H_BATCH + + # --- Stage A: scores + unified global-max softmax per head --- + for h_local in nl.affine_range(H_BATCH): + h = hb_idx * H_BATCH + h_local + + win_bias_h = nl.ndarray((TILE_Q, WIN_SIZE), dtype=nl.float32, buffer=nl.sbuf) + nisa.scalar_tensor_tensor( + dst=win_bias_h, + data=bias_sink_tile, + op0=nl.multiply, + operand0=attn_sink_sb[0:TILE_Q, h : h + 1], + op1=nl.add, + operand1=bias_base_tile, + ) + win_scores_psum = nl.ndarray((TILE_Q, WIN_SIZE), dtype=nl.float32, buffer=nl.psum) + for hd in nl.affine_range(HD_TILES): + nisa.nc_matmul(dst=win_scores_psum, stationary=q_T[h_local][hd], moving=kv_t_win[hd]) + win_scores_h = nl.ndarray((TILE_Q, WIN_SIZE), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=win_scores_h, data1=win_scores_psum, data2=win_bias_h, op=nl.add) + win_max = nl.ndarray((TILE_Q, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_reduce(dst=win_max, data=win_scores_h, op=nl.maximum, axis=1) + + comp_scores = nl.ndarray((TILE_Q, comp_cols), dtype=nl.float32, buffer=nl.sbuf) + for m_idx in nl.affine_range(num_score_chunks): + m_start = m_idx * SCORE_CHUNK + m_sz = min(SCORE_CHUNK, comp_cols - m_start) + scores_chunk_psum = nl.ndarray((TILE_Q, m_sz), dtype=nl.float32, buffer=nl.psum) + for hd in nl.affine_range(HD_TILES): + kt_slice = all_comp_kt[hd][0 : all_comp_kt[hd].shape[0], m_start : m_start + m_sz] + nisa.nc_matmul(dst=scores_chunk_psum, stationary=q_T[h_local][hd], moving=kt_slice) + nisa.tensor_tensor( + dst=comp_scores[0:TILE_Q, m_start : m_start + m_sz], + data1=scores_chunk_psum, + data2=comp_mask[0:TILE_Q, m_start : m_start + m_sz], + op=nl.add, + ) + comp_max = nl.ndarray((TILE_Q, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_reduce(dst=comp_max, data=comp_scores, op=nl.maximum, axis=1) + + neg_max = nl.ndarray((TILE_Q, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=neg_max, data1=win_max, data2=comp_max, op=nl.maximum) + nisa.tensor_scalar(dst=neg_max, data=neg_max, op0=nl.multiply, operand0=-1.0) + + win_sum = nl.ndarray((TILE_Q, 1), dtype=nl.float32, buffer=nl.sbuf) + win_exp_all[h_local] = nl.ndarray((TILE_Q, WIN_SIZE), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.activation( + dst=win_exp_all[h_local], + op=nl.exp, + data=win_scores_h, + bias=neg_max, + reduce_op=nl.add, + reduce_res=win_sum, + reduce_cmd=nisa.reduce_cmd.reset_reduce, + ) + comp_sum = nl.ndarray((TILE_Q, 1), dtype=nl.float32, buffer=nl.sbuf) + comp_exp_all[h_local] = nl.ndarray((TILE_Q, comp_cols), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.activation( + dst=comp_exp_all[h_local], + op=nl.exp, + data=comp_scores, + bias=neg_max, + reduce_op=nl.add, + reduce_res=comp_sum, + reduce_cmd=nisa.reduce_cmd.reset_reduce, + ) + total_sum_all[h_local] = nl.ndarray((TILE_Q, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=total_sum_all[h_local], data1=win_sum, data2=comp_sum, op=nl.add) + + # --- Stage B: V multiply + normalize + write per head --- + for h_local in nl.affine_range(H_BATCH): + h = hb_idx * H_BATCH + h_local + out_psum = nl.ndarray((TILE_Q, head_dim), dtype=nl.float32, buffer=nl.psum) + + scores_T_psum = nl.ndarray((KV_CHUNK, TILE_Q), dtype=nl.bfloat16, buffer=nl.psum) + nisa.nc_transpose(dst=scores_T_psum, data=win_exp_all[h_local][0:TILE_Q, 0:KV_CHUNK]) + scores_T_sb = nl.ndarray((KV_CHUNK, TILE_Q), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.tensor_copy(dst=scores_T_sb, src=scores_T_psum) + nisa.nc_matmul(dst=out_psum, stationary=scores_T_sb, moving=win_v_0) + + scores_T_psum2 = nl.ndarray((KV_CHUNK, TILE_Q), dtype=nl.bfloat16, buffer=nl.psum) + nisa.nc_transpose(dst=scores_T_psum2, data=win_exp_all[h_local][0:TILE_Q, KV_CHUNK:WIN_SIZE]) + scores_T_sb2 = nl.ndarray((KV_CHUNK, TILE_Q), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.tensor_copy(dst=scores_T_sb2, src=scores_T_psum2) + nisa.nc_matmul(dst=out_psum, stationary=scores_T_sb2, moving=win_v_1) + + for c_idx in nl.affine_range(causal_chunks): + c_start = c_idx * COMP_V_CHUNK + exp_T_psum = nl.ndarray((COMP_V_CHUNK, TILE_Q), dtype=nl.bfloat16, buffer=nl.psum) + nisa.nc_transpose( + dst=exp_T_psum, data=comp_exp_all[h_local][0:TILE_Q, c_start : c_start + COMP_V_CHUNK] + ) + exp_T_sb = nl.ndarray((COMP_V_CHUNK, TILE_Q), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.tensor_copy(dst=exp_T_sb, src=exp_T_psum) + nisa.nc_matmul(dst=out_psum, stationary=exp_T_sb, moving=comp_v[c_idx]) + + out_sbuf = nl.ndarray((TILE_Q, head_dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=out_sbuf, src=out_psum) + inv_sum = nl.ndarray((TILE_Q, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation(dst=inv_sum, op=nl.reciprocal, data=total_sum_all[h_local]) + nisa.tensor_scalar(dst=out_sbuf, data=out_sbuf, op0=nl.multiply, operand0=inv_sum) + + q_global = h * S + q_start + out_bf16 = nl.ndarray((TILE_Q, head_dim), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.tensor_copy(dst=out_bf16, src=out_sbuf) + nisa.dma_copy(dst=output[q_global : q_global + TILE_Q, 0:head_dim], src=out_bf16) + + return output + + +# -------------------------------------------------------------------------- +# NKI Kernel: TRUE sparse prefill attention (per-query indirect gather) +# +# The other two prefill attention kernels are dense-plus-mask: they score every +# causal compressed column and predicate the unselected ones to -1e9. That is +# correct but computes `causal_cols / k` more score positions than the model needs +# (4.3x at seq_len=32768, 128x at 1M). +# +# This kernel instead gathers each query's `k` SELECTED compressed rows with an +# indirect DMA and scores only those, so its compressed cost is O(k) and +# independent of context length. +# +# The window is a per-query causal SLICE rather than a masked 256-column block, so +# no additive window bias is needed: query p reads window columns [p, p+W) and the +# window is the causal W-column slice ENDING AT the query's own position. +# -------------------------------------------------------------------------- +@nki.jit +def nki_prefill_sparse_attn_kernel( + topk_idx_T: nl.NkiTensor, # [k, S] uint32 — per-query selected compressed positions + all_q: nl.NkiTensor, # [S * n_heads, head_dim] f16 — QUERY-MAJOR: one query's heads contiguous + all_K_T_win: nl.NkiTensor, # [head_dim, S + W] f16 — window K^T (padded) + all_V_win: nl.NkiTensor, # [S + W, head_dim] f16 — window V (padded) + compress_kv: nl.NkiTensor, # [T_c, head_dim] f16 — FULL compressed KV, replicated per rank + attn_sink_in: nl.NkiTensor, # [1, n_heads] f32 — read for n_heads only; see the window note +) -> nl.NkiTensor: + """O(k) sparse prefill attention. Returns [n_heads * S, head_dim] bf16. + + Requires n_heads == 128 (one full matmul output-partition tile) and + k % 128 == 0. `S` here is the number of queries THIS launch covers. + """ + head_dim = all_q.shape[1] + k_val = topk_idx_T.shape[0] + S = topk_idx_T.shape[1] + n_heads = attn_sink_in.shape[1] + W = 128 + COMP_CHUNK = 128 + HD_CHUNK = 128 + HD_TILES = head_dim // HD_CHUNK + num_chunks = k_val // COMP_CHUNK + + kernel_assert(n_heads == 128, "sparse prefill needs all 128 heads on one rank (sequence-parallel sharding)") + kernel_assert(k_val % COMP_CHUNK == 0, "k must be a multiple of the gather chunk (128)") + kernel_assert(head_dim % HD_CHUNK == 0, "head_dim must be a multiple of 128") + + core_id = nl.program_id(0) + n_cores = nl.num_programs() + kernel_assert(S % n_cores == 0, "query count must divide across the launch grid") + s_per_core = S // n_cores + s_start = core_id * s_per_core + + output = nl.ndarray((n_heads * S, head_dim), dtype=nl.bfloat16, buffer=nl.shared_hbm, name="sparse_prefill_out") + + for q_local in nl.static_range(s_per_core): + q = s_start + q_local # this core's query, a trace-time constant + p = q # window slice is pre-positioned by the caller, so p is tile-relative + + q_hb = [None] * HD_TILES + for hd in range(HD_TILES): + q_hb[hd] = nl.ndarray((HD_CHUNK, n_heads), dtype=nl.float16, buffer=nl.sbuf) + nisa.dma_transpose( + dst=q_hb[hd], + src=all_q.ap( + pattern=[[head_dim, n_heads], [1, HD_CHUNK]], offset=q * n_heads * head_dim + hd * HD_CHUNK + ), + ) + + # ---- indirect gather of this query's k selected compressed rows ---- + kv_chunks = [None] * num_chunks + for c in nl.affine_range(num_chunks): + idx = nl.ndarray((COMP_CHUNK, 1), dtype=nl.uint32, buffer=nl.sbuf) + nisa.dma_copy(dst=idx, src=topk_idx_T.ap(pattern=[[S, COMP_CHUNK], [1, 1]], offset=c * COMP_CHUNK * S + q)) + kv_chunks[c] = nl.ndarray((COMP_CHUNK, head_dim), dtype=nl.float16, buffer=nl.sbuf) + nisa.memset(dst=kv_chunks[c], value=0) + nisa.dma_copy( + dst=kv_chunks[c], + src=compress_kv.ap(pattern=[[head_dim, COMP_CHUNK], [1, head_dim]], vector_offset=idx, indirect_dim=0), + dge_mode=nisa.dge_mode.swdge, + oob_mode=oob_mode.skip, + priority=0, + ) + + # ---- window K^T / V: the causal W-column slice for THIS query ---- + win_kt = [None] * HD_TILES + for hd in range(HD_TILES): + hd_start = hd * HD_CHUNK + win_kt[hd] = nl.ndarray((HD_CHUNK, W), dtype=nl.float16, buffer=nl.sbuf) + nisa.dma_copy( + dst=win_kt[hd], src=all_K_T_win[hd_start : hd_start + HD_CHUNK, p + 1 : p + 1 + W], priority=2 + ) + win_v = nl.ndarray((W, head_dim), dtype=nl.float16, buffer=nl.sbuf) + nisa.dma_copy(dst=win_v, src=all_V_win[p + 1 : p + 1 + W, 0:head_dim], priority=2) + + # ---- window scores ---- + win_ps = nl.ndarray((n_heads, W), dtype=nl.float32, buffer=nl.psum) + for hd in nl.affine_range(HD_TILES): + nisa.nc_matmul(dst=win_ps, stationary=q_hb[hd], moving=win_kt[hd]) + win_scores = nl.ndarray((n_heads, W), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=win_scores, src=win_ps) + + # ---- compressed scores over the k gathered positions ONLY ---- + comp_scores = nl.ndarray((n_heads, k_val), dtype=nl.float32, buffer=nl.sbuf) + for c in nl.affine_range(num_chunks): + c0 = c * COMP_CHUNK + ps = nl.ndarray((n_heads, COMP_CHUNK), dtype=nl.float32, buffer=nl.psum) + for hd in nl.affine_range(HD_TILES): + hd_start = hd * HD_CHUNK + tp = nl.ndarray((HD_CHUNK, COMP_CHUNK), dtype=nl.float16, buffer=nl.psum) + nisa.nc_transpose(dst=tp, data=kv_chunks[c][0:COMP_CHUNK, hd_start : hd_start + HD_CHUNK]) + kt_c = nl.ndarray((HD_CHUNK, COMP_CHUNK), dtype=nl.float16, buffer=nl.sbuf) + nisa.tensor_copy(dst=kt_c, src=tp) + nisa.nc_matmul(dst=ps, stationary=q_hb[hd], moving=kt_c) + nisa.tensor_copy(dst=comp_scores[0:n_heads, c0 : c0 + COMP_CHUNK], src=ps) + + # ---- one global-max softmax over [window | gathered] ---- + win_max = nl.ndarray((n_heads, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_reduce(dst=win_max, data=win_scores, op=nl.maximum, axis=1) + comp_max = nl.ndarray((n_heads, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_reduce(dst=comp_max, data=comp_scores, op=nl.maximum, axis=1) + neg_max = nl.ndarray((n_heads, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=neg_max, data1=win_max, data2=comp_max, op=nl.maximum) + nisa.tensor_scalar(dst=neg_max, data=neg_max, op0=nl.multiply, operand0=-1.0) + + win_sum = nl.ndarray((n_heads, 1), dtype=nl.float32, buffer=nl.sbuf) + win_exp = nl.ndarray((n_heads, W), dtype=nl.float16, buffer=nl.sbuf) + nisa.activation( + dst=win_exp, + op=nl.exp, + data=win_scores, + bias=neg_max, + reduce_op=nl.add, + reduce_res=win_sum, + reduce_cmd=nisa.reduce_cmd.reset_reduce, + ) + comp_sum = nl.ndarray((n_heads, 1), dtype=nl.float32, buffer=nl.sbuf) + comp_exp = nl.ndarray((n_heads, k_val), dtype=nl.float16, buffer=nl.sbuf) + nisa.activation( + dst=comp_exp, + op=nl.exp, + data=comp_scores, + bias=neg_max, + reduce_op=nl.add, + reduce_res=comp_sum, + reduce_cmd=nisa.reduce_cmd.reset_reduce, + ) + + # ---- V accumulation: window slice first, then the gathered chunks ---- + out_psum = nl.ndarray((n_heads, head_dim), dtype=nl.float32, buffer=nl.psum) + we_T = nl.ndarray((W, n_heads), dtype=nl.float16, buffer=nl.psum) + nisa.nc_transpose(dst=we_T, data=win_exp) + we_T_sb = nl.ndarray((W, n_heads), dtype=nl.float16, buffer=nl.sbuf) + nisa.tensor_copy(dst=we_T_sb, src=we_T) + nisa.nc_matmul(dst=out_psum, stationary=we_T_sb, moving=win_v) + for c in nl.affine_range(num_chunks): + c0 = c * COMP_CHUNK + ce_T = nl.ndarray((COMP_CHUNK, n_heads), dtype=nl.float16, buffer=nl.psum) + nisa.nc_transpose(dst=ce_T, data=comp_exp[0:n_heads, c0 : c0 + COMP_CHUNK]) + ce_T_sb = nl.ndarray((COMP_CHUNK, n_heads), dtype=nl.float16, buffer=nl.sbuf) + nisa.tensor_copy(dst=ce_T_sb, src=ce_T) + nisa.nc_matmul(dst=out_psum, stationary=ce_T_sb, moving=kv_chunks[c]) + + # ---- normalize by the shared denominator and write out head-major ---- + total = nl.ndarray((n_heads, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=total, data1=win_sum, data2=comp_sum, op=nl.add) + inv = nl.ndarray((n_heads, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation(dst=inv, op=nl.reciprocal, data=total) + o_f32 = nl.ndarray((n_heads, head_dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=o_f32, src=out_psum) + nisa.tensor_scalar(dst=o_f32, data=o_f32, op0=nl.multiply, operand0=inv) + o_bf = nl.ndarray((n_heads, head_dim), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.tensor_copy(dst=o_bf, src=o_f32) + nisa.dma_copy(dst=output.ap(pattern=[[S * head_dim, n_heads], [1, head_dim]], offset=q * head_dim), src=o_bf) + + return output + + +# -------------------------------------------------------------------------- +# NKI Kernel: per-query top-k over the indexer scores, entirely on-chip +# +# The block used to build nisa.topk's snake layout on the HOST: reshape/transpose +# [S_q, T_c] into [S_q * 128, T_c/16] and hand it back to a topk kernel that +# returned [S_q * 128, k] values AND indices. Only 1/128 of that output is ever +# read (snake group 0 = 16 partitions x k/16 columns), so at S_q = 7168, T_c = 8192 +# the stage moved ~6.5 GB through HBM to deliver 29 MB of indices. +# +# Here the snake tile is built in SBUF from the score rows directly, with the same +# nc_transpose fold `_snake_fill` uses in the decode indexer. HBM traffic becomes +# read S_q * T_c bf16 + write S_q * k uint32 -- 117 MB + 29 MB at those shapes. +# +# -------------------------------------------------------------------------- +_SNAKE_GROUP = 16 +_SNAKE_GROUPS = 8 +_SNAKE_PAR = 128 +_SNAKE_NEG = -1.0e30 +"""Sentinel for snake positions that hold no score. Must be strictly below every real +score, including the indexer's -1e9 causal mask, so padding can never be selected.""" + + +@nki.jit +def nki_prefill_topk_kernel( + scores: nl.NkiTensor, # [S_q, T_c] bf16 — indexer scores, causal bias already folded in + k_val: int, # top-k count; multiple of 16 + n_val: int, # nisa.topk width; T_c padded up to a proven-safe width +) -> nl.NkiTensor: + """Per-query top-k positions. Returns [S_q, k_val] uint32 GLOBAL compressed positions. + + The returned index is a global position because the snake fill satisfies + ``snake[16 * g + r, c] == scores[base + g, 16 * c + r]``, which is exactly + nisa.topk's per-group index encoding. + + The k winners of a row are an UNORDERED SET: nisa.topk emits each snake partition's + winners in ascending position order, not by value. The gather that consumes them is + order-agnostic. + """ + S_q = scores.shape[0] + T_c = scores.shape[1] + GROUP = _SNAKE_GROUP + GROUPS = _SNAKE_GROUPS + PAR = _SNAKE_PAR + snake_x = n_val // GROUP + live_x = T_c // GROUP + k_cols = k_val // GROUP + + kernel_assert(T_c % (GROUP * 128) == 0, "T_c must be a multiple of 16*128 for the snake fold") + kernel_assert(k_val % GROUP == 0, "k must be a multiple of the snake group size") + kernel_assert(n_val >= T_c, "topk width must cover T_c") + + out = nl.ndarray((S_q, k_val), dtype=nl.uint32, buffer=nl.shared_hbm, name="prefill_topk_idx") + + core_id = nl.program_id(0) + n_cores = nl.num_programs() + kernel_assert(S_q % (GROUPS * n_cores) == 0, "score rows must divide into 8-row tiles across the grid") + tiles_per_core = S_q // (GROUPS * n_cores) + tile_base = core_id * tiles_per_core + + snake = nl.ndarray((PAR, snake_x), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.memset(dst=snake, value=_SNAKE_NEG) + + val = nl.ndarray((PAR, k_val), dtype=nl.bfloat16, buffer=nl.sbuf) + idx = nl.ndarray((PAR, k_val), dtype=nl.uint32, buffer=nl.sbuf) + + for t_local in nl.sequential_range(tiles_per_core): + base = (tile_base + t_local) * GROUPS + + for b in nl.static_range(live_x // 128): + blk = nl.ndarray((128, PAR), dtype=nl.bfloat16, buffer=nl.sbuf) + for g in nl.static_range(GROUPS): + nisa.dma_copy( + dst=blk[0:128, g * GROUP : (g + 1) * GROUP], + src=scores.ap(pattern=[[GROUP, 128], [1, GROUP]], offset=(base + g) * T_c + b * 128 * GROUP), + ) + tp = nl.ndarray((PAR, 128), dtype=nl.bfloat16, buffer=nl.psum) + nisa.nc_transpose(dst=tp, data=blk) + nisa.tensor_copy(dst=snake[0:PAR, b * 128 : (b + 1) * 128], src=tp) + + nisa.topk(val_dst=val, idx_dst=idx, src=snake, n=n_val) + + nisa.dma_copy( + dst=out.ap(pattern=[[k_val, GROUPS], [k_cols, GROUP], [1, k_cols]], offset=base * k_val), + src=idx[0:PAR, 0:k_cols], + ) + + return out diff --git a/src/nkilib_src/nkilib/experimental/deepseek_v4_csa/csa_prefill_attention_torch.py b/src/nkilib_src/nkilib/experimental/deepseek_v4_csa/csa_prefill_attention_torch.py new file mode 100644 index 0000000..83a541e --- /dev/null +++ b/src/nkilib_src/nkilib/experimental/deepseek_v4_csa/csa_prefill_attention_torch.py @@ -0,0 +1,419 @@ +# 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. + +"""CPU references for the DeepSeek-V4 CSA prefill kernels. + +One reference per tested kernel, taking the SAME parameter names as its kernel so +the test framework can pair them. + +The one structural thing to know: the sliding window is tiled by QUERY TILE, not +by query row. A tile of 128 queries starting at ``q_start`` all read the same 256 +window key columns ``[q_start, q_start + 256)``, and which of those a given row may +actually attend to is decided by ``win_bias_base`` rather than by the slice. So the +references loop over 128-row tiles exactly as the kernels do -- a reference that +sliced per row would disagree with a correct kernel. + +The window key/value tensors arrive left-padded by ``W``, so window column ``j`` of +tile ``q_start`` is sequence position ``q_start + j - W``. That padding is what lets +the first tile use the same slice arithmetic as every other one. +""" + +import torch +import torch.nn.functional as F + +_TILE_Q = 128 +_W = 128 +_WIN_SIZE = 2 * _TILE_Q +_NEG_INF = -1e9 + + +def _rope_pairs(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, inverse: bool) -> torch.Tensor: + """Rotate interleaved (even, odd) channel pairs of ``x`` by ``cos``/``sin``, in fp32. + + ``inverse`` negates ``sin``, which turns the forward rotation into its inverse -- + the form the output de-RoPE uses. + """ + half_rope = cos.shape[-1] + pairs = x.float().unflatten(-1, (half_rope, 2)) + x1, x2 = pairs[..., 0], pairs[..., 1] + s = -sin if inverse else sin + y1 = x1 * cos - x2 * s + y2 = x1 * s + x2 * cos + return torch.stack([y1, y2], dim=-1).flatten(-2) + + +def nki_rms_rope_torch_ref( + x_in: torch.Tensor, + cos_in: torch.Tensor, + sin_in: torch.Tensor, + gain_in, + eps_val: float, + do_rms: int = 1, + inverse: int = 0, + heads: int = 1, + in_head_major: int = 1, + out_head_major: int = 1, + do_rope: int = 1, +) -> dict[str, torch.Tensor]: + """Oracle for ``nki_rms_rope_kernel``: RMSNorm(+gain) then RoPE over ``[S_rows, head_dim]``. + """ + # do_rope=0 is the norm-only variant; cos_in/sin_in are then unused and may be None. + half_rope = cos_in.shape[1] if do_rope else 0 + rope_dim = 2 * half_rope + if in_head_major: + head_dim, s_len = x_in.shape[1], x_in.shape[0] // heads + rows = x_in + else: + head_dim, s_len = x_in.shape[1] // heads, x_in.shape[0] + rows = x_in.reshape(s_len, heads, head_dim).permute(1, 0, 2).reshape(heads * s_len, head_dim) + nope_dim = head_dim - rope_dim + if heads > 1: + cos_in = cos_in.repeat(heads, 1) + sin_in = sin_in.repeat(heads, 1) + + x = rows.float() + if do_rms: + x = x * torch.rsqrt(x.square().mean(-1, keepdim=True) + eps_val) + if gain_in is not None: + x = x * gain_in.float() + normed = x.to(torch.bfloat16) + + if do_rope: + rotated = _rope_pairs(normed[:, nope_dim:], cos_in.float(), sin_in.float(), bool(inverse)) + out = torch.cat([normed[:, :nope_dim], rotated.to(torch.bfloat16)], dim=-1) + else: + out = normed + if not out_head_major: + out = out.reshape(heads, s_len, head_dim).permute(1, 0, 2).reshape(s_len, heads * head_dim) + return {"output_0": out} + + +def nki_compressor_core_torch_ref( + kv8: torch.Tensor, + score8: torch.Tensor, + norm_weight: torch.Tensor, + cos_rep: torch.Tensor, + sin_rep: torch.Tensor, + eps: float, + ape=None, + hadamard=None, +) -> dict[str, torch.Tensor]: + """Oracle for ``nki_compressor_core_kernel``: gated pooling, then RMSNorm, then RoPE. + + The gate is a softmax over the ``ratio2`` overlapped slots taken INDEPENDENTLY + PER CHANNEL -- not per position -- so each channel of a compressed position + pools its ``ratio2`` candidates with its own weights. That per-channel + independence is why the kernel can keep positions on partitions and channels on + the free axis and never reduce across cores. + + ``cos_rep``/``sin_rep`` arrive with each pair's angle DUPLICATED across the two + channels of the pair, so the kernel can load them with the same access pattern + as the data; only the even entries are read, which is what this mirrors. + """ + rope_dim = cos_rep.shape[1] + head_dim = kv8.shape[2] + nope_dim = head_dim - rope_dim + + # ape, when given, is the per-slot gate bias the kernel adds in fp32 after widening + # the bf16 score; shaped [ratio2, head_dim] and broadcast over positions. + scores = score8.float() + if ape is not None: + scores = scores + ape.float().unsqueeze(0) + weights = torch.softmax(scores, dim=1) + pooled = (kv8.float() * weights).sum(dim=1) + + # bf16 at the pooling output, matching the model casting before its RMSNorm. + x = pooled.to(torch.bfloat16).float() + normed = x * torch.rsqrt(x.square().mean(-1, keepdim=True) + eps) * norm_weight.float() + normed = normed.to(torch.bfloat16) + + # Each pair's angle is duplicated across its two channels; take one per pair. + cos = cos_rep.float()[:, 0::2] + sin = sin_rep.float()[:, 0::2] + rotated = _rope_pairs(normed[:, nope_dim:], cos, sin, inverse=False) + out = torch.cat([normed[:, :nope_dim], rotated.to(torch.bfloat16)], dim=-1) + + # The indexer's compressor rotates the finished row by an orthonormal Hadamard. The + # reference takes the bf16 row into the matmul exactly as the kernel does, so the + # single rounding before the rotation is shared rather than being a kernel artifact. + if hadamard is not None: + out = (out.float() @ hadamard.float()).to(torch.bfloat16) + return {"output_0": out} + + +def _indexer_scores(q_T_all: torch.Tensor, kv_t: torch.Tensor, weights: torch.Tensor) -> torch.Tensor: + """``score[s, t] = sum_h relu(q[s, h, :] . kv[t, :]) * weights[s, h]``, as ``[S_q, T_c]``. + + The per-head sum is accumulated in **bf16**, one head at a time, because that is + what the kernel does: its running total lives in a bf16 SBUF tile and each head's + contribution is folded in with a single `scalar_tensor_tensor`. Accumulating in + fp32 here instead would make the reference systematically more precise, which + matters more than usual downstream -- ``nki_indexer_score_mask_kernel`` turns this + score into a THRESHOLDED mask, so a score difference of one bf16 ulp near the + threshold flips a position and shows up as a full ``1e9`` mask error rather than + as a small numeric one. + """ + head_dim, total_q_free = q_T_all.shape + n_heads = weights.shape[1] + s_q = total_q_free // n_heads + + q = q_T_all.float().reshape(head_dim, n_heads, s_q).permute(2, 1, 0) + per_head = F.relu(torch.einsum("shd,dt->sht", q, kv_t.float())).to(torch.bfloat16) + + acc = torch.zeros((s_q, kv_t.shape[1]), dtype=torch.bfloat16) + for h in range(n_heads): + acc = (per_head[:, h] * weights[:, h : h + 1].float()).to(torch.bfloat16).add(acc).to(torch.bfloat16) + return acc.float() + + +def nki_indexer_score_mask_torch_ref( + q_T_all: torch.Tensor, + kv_t: torch.Tensor, + weights: torch.Tensor, + causal_bias: torch.Tensor, + k: int, +) -> dict[str, torch.Tensor]: + """Oracle for ``nki_indexer_score_mask_kernel``: the ``0 / -1e9`` top-k selection mask. + + The threshold comes from 9 rounds of BISECTION on the score range, not from a + sort, so this reference runs the same bisection rather than calling + ``torch.topk``. That is the honest comparison: bisection to a fixed depth does + not always select exactly ``k`` positions (ties, and a residual interval), and a + ``topk``-based reference would report those as kernel errors. + + ``lo`` starts at the smallest score that is not causally masked (found by + folding the masked entries up to ``hi``), ``hi`` at the largest. Each round keeps + the half that still holds at least ``k`` positions, and the final mask admits + every score at or above ``lo``. + """ + scores = _indexer_scores(q_T_all, kv_t, weights) + causal_bias.float() + + hi = scores.max(dim=-1, keepdim=True).values + # Fold causally-masked entries (score <= -1e8) up to `hi` so they cannot become + # the minimum, then take the min over what is left. + is_valid = (scores > -1e8).float() + lo = ((scores - hi) * is_valid + hi).min(dim=-1, keepdim=True).values + + for _ in range(9): + mid = (lo + hi) * 0.5 + enough = ((scores >= mid).float().sum(dim=-1, keepdim=True) >= float(k)).float() + lo = lo + (mid - lo) * enough + hi = hi + (mid - hi) * (1.0 - enough) + + sel = (scores >= lo).float() * (-_NEG_INF) + _NEG_INF + return {"output_0": sel} + + +def _tiled_sparse_attention( + all_q_T: torch.Tensor, + win_K_T: torch.Tensor, + win_V: torch.Tensor, + comp_K_T: torch.Tensor, + comp_V: torch.Tensor, + comp_sel: torch.Tensor, + win_bias_base: torch.Tensor, + win_bias_sink: torch.Tensor, + attn_sink: torch.Tensor, +) -> torch.Tensor: + """Global-max softmax over [window | selected compressed], tiled by query tile. + + Both prefill attention kernels compute exactly this. The window and the + compressed positions share ONE softmax normalization (a single global max over + the two score sets), which is what makes the two contributions directly + comparable and removes any need for online rescaling. + + Masking is additive and happens before ``exp``: ``win_bias_base`` is ``-1e9`` + outside a row's causal window and ``comp_sel`` is ``-1e9`` at unselected + compressed positions, so those terms underflow to exactly zero and contribute + nothing to either the denominator or the value sum. That is also why a + reference over ALL ``T_c`` columns matches a kernel that truncates the + compressed loop at its causal bound. + + Returns ``[n_heads * S, head_dim]``, head-major as the kernels write it. + """ + head_dim, total_q_free = all_q_T.shape + s_len = comp_sel.shape[0] + n_heads = total_q_free // s_len + t_c = comp_sel.shape[1] + + # all_q_T[d, h * S + s] == q[s, h, d] -> [S, n_heads, head_dim] + q = all_q_T.float().reshape(head_dim, n_heads, s_len).permute(2, 1, 0) + out = torch.zeros((n_heads * s_len, head_dim), dtype=torch.bfloat16) + + comp_scores_all = torch.einsum("shd,dt->sht", q, comp_K_T.float()) + comp_sel.float().unsqueeze(1) + + for q_start in range(0, s_len, _TILE_Q): + rows = slice(q_start, q_start + _TILE_Q) + q_tile = q[rows] # [tile, n_heads, head_dim] + + # Every row of the tile reads the same 256 window columns; win_bias_base + # decides which of them the row may actually attend to. + k_win = win_K_T.float()[:, q_start : q_start + _WIN_SIZE] # [head_dim, 256] + v_win = win_V.float()[q_start : q_start + _WIN_SIZE] # [256, head_dim] + + bias = win_bias_sink.float()[rows].unsqueeze(1) * attn_sink.float().reshape( + 1, n_heads, 1 + ) + win_bias_base.float()[rows].unsqueeze(1) + win_scores = torch.einsum("shd,dj->shj", q_tile, k_win) + bias + comp_scores = comp_scores_all[rows] + + shift = torch.maximum( + win_scores.max(dim=-1, keepdim=True).values, + comp_scores.max(dim=-1, keepdim=True).values, + ) + win_exp = torch.exp(win_scores - shift) + comp_exp = torch.exp(comp_scores - shift) + total = win_exp.sum(-1, keepdim=True) + comp_exp.sum(-1, keepdim=True) + + tile_out = ( + torch.einsum("shj,jd->shd", win_exp, v_win) + torch.einsum("sht,td->shd", comp_exp, comp_V.float()) + ) / total + + for h in range(n_heads): + out[h * s_len + q_start : h * s_len + q_start + q_tile.shape[0]] = tile_out[:, h].to(torch.bfloat16) + + del t_c + return out + + +def nki_fused_csa_attn_torch_ref( + compress_sel: torch.Tensor, + all_q_T: torch.Tensor, + all_K_T: torch.Tensor, + all_V: torch.Tensor, + win_bias_base_in: torch.Tensor, + win_bias_sink_in: torch.Tensor, + attn_sink_in: torch.Tensor, +) -> dict[str, torch.Tensor]: + """Oracle for ``nki_fused_csa_attn_kernel``: mask-predicated attention on concatenated KV. + + This kernel takes the window and compressed keys/values in ONE concatenated + tensor laid out ``[padded window (S + W) | compressed (T_c)]``, so the reference + splits them back apart at ``S + W`` before attending. + """ + s_len, t_c = compress_sel.shape + split = s_len + _W + out = _tiled_sparse_attention( + all_q_T, + all_K_T[:, 0:split], + all_V[0:split], + all_K_T[:, split : split + t_c], + all_V[split : split + t_c], + compress_sel, + win_bias_base_in, + win_bias_sink_in, + attn_sink_in, + ) + return {"output_0": out} + + +def nki_gather_csa_attn_torch_ref( + topk_sel_bias: torch.Tensor, + all_q_T: torch.Tensor, + all_K_T_win: torch.Tensor, + all_V_win: torch.Tensor, + compress_kv_T: torch.Tensor, + compress_kv: torch.Tensor, + win_bias_base_in: torch.Tensor, + win_bias_sink_in: torch.Tensor, + attn_sink_in: torch.Tensor, + split_pos: int, + ratio: int, +) -> dict[str, torch.Tensor]: + """Oracle for ``nki_gather_csa_attn_kernel``: the same attention with a static causal bound. + + ``split_pos`` and ``ratio`` only set how far the kernel's compressed loop runs: + it stops at the chunk holding the tile's last causally-reachable compressed + position. This reference deliberately attends over ALL ``T_c`` compressed + columns and relies on ``topk_sel_bias`` being ``-1e9`` past the frontier, so if + the kernel's bound were ever too tight -- dropping a column that mattered -- the + two would disagree. + """ + del split_pos, ratio + out = _tiled_sparse_attention( + all_q_T, + all_K_T_win, + all_V_win, + compress_kv_T, + compress_kv, + topk_sel_bias, + win_bias_base_in, + win_bias_sink_in, + attn_sink_in, + ) + return {"output_0": out} + + +def nki_prefill_sparse_attn_torch_ref( + topk_idx_T: torch.Tensor, + all_q: torch.Tensor, + all_K_T_win: torch.Tensor, + all_V_win: torch.Tensor, + compress_kv: torch.Tensor, + attn_sink_in: torch.Tensor, +) -> dict[str, torch.Tensor]: + """Oracle for ``nki_prefill_sparse_attn_kernel``. + + Same global-max softmax over ``[window | compressed]`` as the dense prefill + kernels. Two differences follow from the sparse formulation: the window is read as + the causal ``W``-column slice ending at the query's own position rather than via an + additive bias over a shared ``2W`` tile, and the compressed set is the ``k`` GATHERED + positions rather than all ``T_c`` behind a ``-1e9`` predicate. Both are equivalent to + the dense form -- an unselected position contributes ``exp(-1e9 - shift) == 0``. + + Written against ``csa_block_torch``, NOT against the kernel: an oracle that copies + the kernel's indexing cannot detect an indexing bug, and that is exactly how a + one-position window shift and a misplaced attention sink previously passed here. + """ + k_val, s_len = topk_idx_T.shape + head_dim = all_q.shape[1] + n_heads = all_q.shape[0] // s_len + w = _W + + # all_q is [S * n_heads, head_dim] query-major -> [S, H, D] + q = all_q.float().reshape(s_len, n_heads, head_dim) + del attn_sink_in # sink applies to absolute position 0 only, which is outside this window + out = torch.zeros((n_heads * s_len, head_dim), dtype=torch.bfloat16) + + for s in range(s_len): + p = s # window slice is pre-positioned by the caller + qs = q[s] # [H, D] + + # [p + 1, p + 1 + W) on a buffer front-padded by W == original positions + # [g - W + 1, g] for the query at global position g -- the window INCLUDING the + # query's own key, which is what csa_block_torch.get_window_topk_idxs attends. + # No attention sink: that bias belongs to ABSOLUTE position 0 only, and this + # kernel runs solely on the scored region where position 0 is far outside the + # W = 128 window. + win_k = all_K_T_win.float()[:, p + 1 : p + 1 + w] # [D, W] + win_v = all_V_win.float()[p + 1 : p + 1 + w] # [W, D] + win_scores = qs @ win_k # [H, W] + + sel = topk_idx_T[:, s].long() # [k] + comp_k = compress_kv.float()[sel] # [k, D] + comp_scores = qs @ comp_k.T # [H, k] + + shift = torch.maximum( + win_scores.max(dim=-1, keepdim=True).values, + comp_scores.max(dim=-1, keepdim=True).values, + ) + we = torch.exp(win_scores - shift) + ce = torch.exp(comp_scores - shift) + total = we.sum(-1, keepdim=True) + ce.sum(-1, keepdim=True) + o = (we @ win_v + ce @ comp_k) / total # [H, D] + + for h in range(n_heads): + out[h * s_len + s] = o[h].to(torch.bfloat16) + + return {"output_0": out} diff --git a/src/nkilib_src/nkilib/experimental/deepseek_v4_csa/csa_tp_all_reduce.py b/src/nkilib_src/nkilib/experimental/deepseek_v4_csa/csa_tp_all_reduce.py new file mode 100644 index 0000000..a981979 --- /dev/null +++ b/src/nkilib_src/nkilib/experimental/deepseek_v4_csa/csa_tp_all_reduce.py @@ -0,0 +1,127 @@ +# 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. + +"""Tensor-parallel output all-reduce for the CSA attention blocks. + +""" + +import nki +import nki.collectives as ncc +import nki.isa as nisa +import nki.language as nl +from nki.collectives import ReplicaGroup +from torch import nn + + +@nki.jit +def nki_tp_all_reduce_kernel(input: nl.NkiTensor, replica_group: ReplicaGroup) -> nl.NkiTensor: + """Sum `input` across the ranks of `replica_group` (op=nl.add) at lnc=2. + + `input` is a 2D [P, F] tile. This is the canonical nki-library + all_reduce_hbm_kernel over the WHOLE tensor — collective src/dst must be + freshly-allocated nl.shared_hbm WITH name=, and a collective cannot read/write + IO tensors directly, so the input is staged in via dma_copy and the result + copied back out. + + """ + src = nl.ndarray(input.shape, dtype=input.dtype, buffer=nl.shared_hbm, name="src") + dst = nl.ndarray(input.shape, dtype=input.dtype, buffer=nl.shared_hbm, name="dst") + out = nl.ndarray(input.shape, dtype=input.dtype, buffer=nl.shared_hbm) + + nisa.dma_copy(dst=src, src=input, priority=0) + + ncc.all_reduce(dsts=[dst], srcs=[src], op=nl.add, replica_group=replica_group, priority=0) + nisa.dma_copy(dst=out, src=dst, priority=1) + return out + + +@nki.jit +def nki_tp_all_gather_kernel( + input: nl.NkiTensor, replica_group: ReplicaGroup, world: int, rows: int, free: int +) -> nl.NkiTensor: + """Concatenate this rank's ``input`` shard with every other rank's, along dim 0. + + ``input`` is a 2D ``[rows, free]`` shard; the result is ``[world * rows, free]`` in + RANK ORDER, which is what sequence-parallel prefill needs: rank r computes compressed + positions ``[r * T_c / world, (r + 1) * T_c / world)`` and every rank then needs all + ``T_c`` of them, because a query's top-k may select any compressed position. + + ``rows``/``free`` are passed as compile-time ints rather than read off + ``input.shape``: a traced kernel cannot tuple-unpack the shape of an IO tensor + (``error: failed to resolve name 'input.shape'``), and the gathered ``dst`` needs + ``world * rows``, so the extent cannot be forwarded whole the way ``all_reduce`` + forwards ``input.shape`` into a same-shape allocation. + + """ + src = nl.ndarray((rows, free), dtype=input.dtype, buffer=nl.shared_hbm, name="ag_src") + dst = nl.ndarray((world * rows, free), dtype=input.dtype, buffer=nl.shared_hbm, name="ag_dst") + out = nl.ndarray((world * rows, free), dtype=input.dtype, buffer=nl.shared_hbm) + + nisa.dma_copy(dst=src, src=input, priority=0) + ncc.all_gather(dsts=[dst], srcs=[src], replica_group=replica_group, collective_dim=0, priority=0) + nisa.dma_copy(dst=out, src=dst, priority=1) + return out + + +def tp_all_gather_rows(shard, replica_ranks): + """All-gather a ``[n, F]`` row-shard into ``[world * n, F]`` in rank order. + + Returns ``shard`` unchanged for a single-rank group, so the same call site works on + the sequential (no-peer) harness and on a real multi-worker launch. + """ + ranks = list(replica_ranks) + if len(ranks) == 1: + return shard + rows, free = shard.shape + replica_group = ReplicaGroup([ranks]) + return nki_tp_all_gather_kernel[2](shard.contiguous(), replica_group, len(ranks), rows, free) + + +def tp_all_reduce(partial, replica_ranks): + """All-reduce (sum) a [B, 1, dim] partial across `replica_ranks` on 2 LNC. + + Reshapes the partial to a balanced 2D [P, F] tile and launches the collective + on the `[2]` grid (so it runs in the block's lnc=2 context). The all_reduce + sum is element-wise (layout-invariant, value-EXACT), so the reshape does not + affect the result; a [1, dim] single-partition input would instead leave a + logical core empty. P is chosen so total is P*F with P even. Returns the full + summed [B, 1, dim]. + """ + bsz, seqlen, dim = partial.shape + total = bsz * seqlen * dim + P = 128 + while P > 1 and (total % P != 0 or P % 2 != 0): + P -= 2 + flat = partial.reshape(P, total // P).contiguous() + replica_group = ReplicaGroup([list(replica_ranks)]) + summed = nki_tp_all_reduce_kernel[2](flat, replica_group) # [P, total//P] + return summed.reshape(bsz, seqlen, dim) + + +class TPAllReduceNKI(nn.Module): + """Standalone module wrapping ONLY the 2-LNC ncc.all_reduce collective, so + torch_neuronx.trace(..., compiler_args=["--logical-nc-config=2"]) produces a + NEFF that is JUST the collective (for isolated profiling). + + replica_ranks defines the ReplicaGroup ([[0,1,2,3]] for a 4-rank world); on a + genuine multi-worker launch every rank returns the full summed [B,1,dim], and + with a single-rank group it is an identity passthrough. + """ + + def __init__(self, replica_ranks): + super().__init__() + self.replica_ranks = list(replica_ranks) + + def forward(self, partial): + return tp_all_reduce(partial, self.replica_ranks) diff --git a/src/nkilib_src/nkilib/experimental/deepseek_v4_csa/csa_tp_all_reduce_torch.py b/src/nkilib_src/nkilib/experimental/deepseek_v4_csa/csa_tp_all_reduce_torch.py new file mode 100644 index 0000000..ebc9816 --- /dev/null +++ b/src/nkilib_src/nkilib/experimental/deepseek_v4_csa/csa_tp_all_reduce_torch.py @@ -0,0 +1,55 @@ +# 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. + +"""CPU reference for the CSA tensor-parallel output all-reduce. + +The reduction itself is an ordinary ``dist.all_reduce`` sum. What the kernel adds +around it -- staging the input into a NAMED ``shared_hbm`` buffer, launching on the +``[2]`` grid so the ``lnc=2`` lowering distributes one whole-tensor collective +across the rank's two logical cores, and copying the result back out -- is exactly +what the test is for, since each of those is a way to get a silently HALF-reduced +answer rather than an error. + +Uses the same ``get_pg`` adapter as the other collective references, so it works +under both the simulated and the real distributed runner. +""" + +import numpy as np +import torch +import torch.distributed as dist +from nki.collectives import ReplicaGroup + +from ..collectives.distributed_adapter import get_pg + + +def nki_tp_all_reduce_torch_ref(input: np.ndarray, replica_group: ReplicaGroup) -> dict: + """Sum ``input`` elementwise across the ranks of ``replica_group``. + + An all-reduce sum is elementwise and therefore layout-invariant, which is why + the kernel is free to reshape a ``[B, S, dim]`` partial into a balanced ``[P, F]`` + tile before reducing -- the reshape cannot change the result, and it keeps both + logical cores occupied where a single-partition input would leave one empty. + """ + dtype = input.dtype + try: + tensor = torch.from_numpy(input.copy()) + except TypeError: + tensor = torch.from_numpy(input.astype(np.float32)) + + dist.all_reduce(tensor, op=dist.ReduceOp.SUM, group=get_pg(replica_group)) + + out = tensor.numpy() + if out.dtype != dtype: + out = out.astype(dtype) + return {"out": out} diff --git a/test/integration/nkilib/experimental/deepseek_v4_csa/__init__.py b/test/integration/nkilib/experimental/deepseek_v4_csa/__init__.py new file mode 100644 index 0000000..ce63f8f --- /dev/null +++ b/test/integration/nkilib/experimental/deepseek_v4_csa/__init__.py @@ -0,0 +1,13 @@ +# 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. diff --git a/test/integration/nkilib/experimental/deepseek_v4_csa/test_csa_block.py b/test/integration/nkilib/experimental/deepseek_v4_csa/test_csa_block.py new file mode 100644 index 0000000..4a18d24 --- /dev/null +++ b/test/integration/nkilib/experimental/deepseek_v4_csa/test_csa_block.py @@ -0,0 +1,82 @@ +# 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. +"""End-to-end tests for the whole CSA attention block. + +The per-kernel tests in this directory grade each ``@nki.jit`` kernel against a CPU +reference in isolation. That leaves the block itself untested, and the block is where +a distinct class of bug lives: the torch projections around the kernels, the +trace-time dispatch that picks WHICH kernel a given shape uses, the head-parallel +weight sharding, and whether the kernels compose at all once ``torch_neuronx.trace`` +compiles them together rather than one at a time. + +That last one is not hypothetical. The XLA trace path rejects Python constructs the +standalone kernel tests happily compile, so a block-level compile failure can sit +behind a fully green per-kernel suite. + +Each test traces every head-parallel rank, sums the partials on the host, and grades +the result against a 128-head CPU golden -- so a sharding mistake surfaces as a +numeric failure rather than as a plausible-looking number. The collective is covered +separately by ``test_csa_tp_all_reduce``; here ``replica_ranks=None``, which isolates +the block compute. +""" + +import os +from typing import final + +import pytest + +from test.utils.common_dataclasses import Platforms +from test.utils.pytest_parametrize import pytest_parametrize +from test.utils.pytest_test_metadata import pytest_marks, pytest_test_metadata + +pytestmark = pytest.mark.platforms(exclude=list(set(Platforms) - {Platforms.TRN3, Platforms.TRN3_A0})) + +# Unlike the per-kernel tests, these do not go through the Orchestrator -- the block is +# an nn.Module, so it is traced with torch_neuronx IN THIS PROCESS. That means they only +# run where pytest itself is on a Trainium host, not when the run ships kernels to a +# remote fleet host, so skip rather than fail when there is no local device. +_HAS_LOCAL_DEVICE = os.path.exists("/dev/neuron0") + + +@final +@pytest_test_metadata(name="DeepSeek-V4 CSA Attention Block") +@pytest_marks(["deepseek_v4_csa"]) +@pytest.mark.skip_simulation +@pytest.mark.high_rank +@pytest.mark.slow +class TestCsaBlock: + """The full block: torch projections + every CSA kernel, graded end to end.""" + + _PARAMS = "phase, seq_len, tp_size" + # tp_size=4 is the production sharding (128 heads -> 32/rank). seq_len picks the + # dispatch path: 8192 gives T_c=2048 (fused single-chunk indexer), 32768 gives + # T_c=8192, which is where n_val == T_c and the sentinel pad is skipped. + _CASES = [ + ("decode", 8192, 4), + ("decode", 32768, 4), + ("prefill", 8192, 4), + ] + _ABBREVS = {"seq_len": "s", "tp_size": "tp"} + + @pytest.mark.skipif(not _HAS_LOCAL_DEVICE, reason="in-process tracing needs a local Neuron device") + @pytest_parametrize(_PARAMS, _CASES, abbrevs=_ABBREVS) + def test_block_matches_cpu_golden(self, phase: str, seq_len: int, tp_size: int): + """Trace each rank's block, host-sum the partials, grade against the golden.""" + pytest.importorskip("torch_neuronx", reason="block tracing needs torch_neuronx") + + from nkilib_src.nkilib.experimental.deepseek_v4_csa.csa_block import run_sequential + from nkilib_src.nkilib.experimental.deepseek_v4_csa.csa_common import CSAConfigFull + + passed = run_sequential(phase, CSAConfigFull(seq_len=seq_len), tp_size) + assert passed, f"{phase} block at seq_len={seq_len}, tp_size={tp_size} exceeded its tolerance" diff --git a/test/integration/nkilib/experimental/deepseek_v4_csa/test_csa_decode_attention.py b/test/integration/nkilib/experimental/deepseek_v4_csa/test_csa_decode_attention.py new file mode 100644 index 0000000..72250ef --- /dev/null +++ b/test/integration/nkilib/experimental/deepseek_v4_csa/test_csa_decode_attention.py @@ -0,0 +1,590 @@ +# 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. +"""Integration tests for the DeepSeek-V4 CSA decode kernels. + +The headline test is ``test_score_topk_gather_fused``, which grades the fused +megakernel end to end. That is the only way to grade it -- the indexer score, the +GpSimd top-k and the O(k) attention all happen inside one launch, so nothing in +between is observable -- and it is also what transitively covers the snake +reformat, the cross-core barrier and the ``sendrecv`` softmax merge, none of which +has an observable output of its own. + +Three kernels are covered compile-only rather than numerically, and deliberately: +``nisa_topk_snake_kernel``, ``nki_indexer_score_topk_kernel`` and +``nki_indexer_score_topk_2core`` all return SELECTED POSITIONS. The kernels emit +their winners as an unordered set into row 0 of an 8-row buffer, leaving rows 1-7 +undefined, and ``nisa.topk`` gives no ordering guarantee within the set -- so an +elementwise comparison against ``torch.topk`` would report differences that are not +errors. Their numerics are covered through the fused kernel, whose softmax over +gathered positions is permutation-invariant and therefore insensitive to exactly +the freedom that makes a direct comparison meaningless. + +These kernels use ``priority=`` DMA class-of-service hints, which exist only on +NeuronCore-v4, so every test here is trn3-only. +""" + +from typing import Any, final + +import ml_dtypes +import neuron_dtypes as dt +import nki.language as nl +import numpy as np +import pytest +from nkilib_src.nkilib.experimental.deepseek_v4_csa.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 nkilib_src.nkilib.experimental.deepseek_v4_csa.csa_decode_attention_torch import ( + nki_decode_gather_ok_torch_ref, + nki_indexer_qproj_gemv_torch_ref, + nki_indexer_score_2core_torch_ref, + nki_indexer_score_topk_gather_2core_torch_ref, + nki_indexer_score_torch_ref, + nki_qkv_rms_rope_torch_ref, +) + +from test.utils.common_dataclasses import CompilerArgs, InferenceArgs, Platforms +from test.utils.pytest_parametrize import pytest_parametrize +from test.utils.pytest_test_metadata import pytest_marks, pytest_test_metadata +from test.utils.test_orchestrator import Orchestrator +from test.utils.unit_test_framework import UnitTestFramework, torch_ref_wrapper + +# The `priority=` DMA hints these kernels carry are NeuronCore-v4 only. +pytestmark = pytest.mark.platforms(exclude=list(set(Platforms) - {Platforms.TRN3, Platforms.TRN3_A0})) + +_INDEX_HEAD_DIM = 128 +_WINDOW = 128 +_S_Q = 128 + +_BF16 = ml_dtypes.bfloat16 +_WARMUP_RUNS = 2 + + +def _rng(seed: int = 42) -> np.random.Generator: + """A seeded generator, so every test case is reproducible.""" + return np.random.default_rng(seed) + + +def _bf16(x: np.ndarray) -> np.ndarray: + return dt.static_cast(x.astype(np.float32), nl.bfloat16) + + +def _f16(x: np.ndarray) -> np.ndarray: + return x.astype(np.float16) + + +def _indexer_inputs(t_c: int, n_index_heads: int, k: int | None = None, seed: int = 42) -> dict[str, Any]: + """Indexer scoring operands. With ``k`` given, the top-k boundary is UNAMBIGUOUS. + + Why that matters: the kernels select with ``nisa.topk`` and the reference with + ``torch.topk``. If the k-th and (k+1)-th scores are merely *close*, the two can + legitimately pick different positions, and a test built on smoothly-varying + scores fails on a difference that is not an error -- which is exactly what a + ``linspace`` of per-position gains produced here at ``T_c = 8192``, where 8192 + gains inside one bounded range sit far closer together than a bf16 ulp. + + So when ``k`` is supplied the scores are built as TWO WIDELY SEPARATED CLUSTERS + of DISTINCT values: exactly ``k`` positions land in a high band and the rest in a + band ~50x below it, all sharing one direction so the score is strictly + proportional to a per-position gain. The 50x gap puts the k-th boundary far + beyond anything rounding can cross, so the top-k is precisely the high cluster. + + Distinct *within* each band is the other half of the requirement, and it is not + optional. Giving a whole band one shared value instead makes ``k`` positions score + *exactly* the same, and on that degenerate tied distribution which ``k`` of the + equal scores come back is not pinned down -- so the comparison against + ``torch.topk`` fails on the tie, not on the kernel. Real indexer scores are a + projection of activations and are never tied like that, so a tied input tests a + regime the kernel never sees. + + Order within the high cluster is left arbitrary on purpose: all ``k`` of them are + selected regardless, and the softmax over gathered positions is + permutation-invariant. The high positions are shuffled, so a kernel that quietly + returned "the first k" still fails. + + ``kv_t`` is the INDEXER cache and feeds scoring only; the attention gathers from + a separate ``compress_kv`` whose rows stay distinct, so sharing the scoring + direction costs the attention no coverage. + + With ``k`` omitted (the score-only tests, which have no selection step) each + position gets its own gain instead, which exercises a wider score range. + """ + rng = _rng(seed) + q = rng.standard_normal((_INDEX_HEAD_DIM, n_index_heads)).astype(np.float32) * 0.5 + + if k is None: + gains = np.linspace(0.2, 1.8, t_c, dtype=np.float32) + rng.shuffle(gains) + kv = rng.standard_normal((_INDEX_HEAD_DIM, t_c)).astype(np.float32) * 0.3 * gains + else: + gains = rng.uniform(0.01, 0.02, t_c).astype(np.float32) + high = rng.permutation(t_c)[:k] + gains[high] = rng.uniform(1.0, 2.0, k).astype(np.float32) + base = rng.standard_normal((_INDEX_HEAD_DIM, 1)).astype(np.float32) * 0.3 + kv = base * gains + + # q_T_all[d, h * S_q + s] == q[s, h, d]; all S_q rows are the same query. + q_T_all = np.repeat(q, _S_Q, axis=1) + weights = np.tile(rng.uniform(0.5, 1.5, (1, n_index_heads)).astype(np.float32), (_S_Q, 1)) + return {"q_T_all": _bf16(q_T_all), "kv_t": _bf16(kv), "weights": weights} + + +@final +@pytest_test_metadata(name="DeepSeek-V4 CSA Decode") +@pytest_marks(["attention", "deepseek_v4_csa"]) +class TestCsaDecodeAttention: + """Decode-side CSA kernels: projection tail, indexer, top-k, and O(k) attention.""" + + # ---------------- fused RMS + RoPE projection tail ---------------- + + _RMS_ROPE_PARAMS = "n_heads, head_dim, rope_head_dim" + _RMS_ROPE_CASES = [ + (32, 512, 64), # the production per-rank shape + (8, 128, 32), + ] + _RMS_ROPE_ABBREVS = {"n_heads": "h", "head_dim": "d", "rope_head_dim": "rd"} + + @pytest.mark.fast + @pytest_parametrize(_RMS_ROPE_PARAMS, _RMS_ROPE_CASES, abbrevs=_RMS_ROPE_ABBREVS) + def test_qkv_rms_rope( + self, + test_manager: Orchestrator, + platform_target: Platforms, + n_heads: int, + head_dim: int, + rope_head_dim: int, + ): + """The q heads and the kv row normalized and rotated on one packed tile. + + The kernel unifies the two projection tails with a per-partition gain of + ``1.0`` on the q rows, relying on ``x * 1.0 == x`` being exact in fp32. The + reference runs the two paths separately, so an inexact gain would show up + here as a q-row mismatch. + """ + rng = _rng() + half_rope = rope_head_dim // 2 + eps = 1e-6 + + def input_generator(test_config): + return { + "q_in": _bf16(rng.standard_normal((n_heads, head_dim)) * 0.5), + "kv_in": _bf16(rng.standard_normal((1, head_dim)) * 0.5), + # A gain near 1 with real spread: an all-ones gain would not + # distinguish the kv row's learnable scaling from the q rows'. + "weight_in": rng.uniform(0.8, 1.2, (1, head_dim)).astype(np.float32), + "cos_in": np.cos(rng.uniform(0, 2 * np.pi, (1, half_rope))).astype(np.float32), + "sin_in": np.sin(rng.uniform(0, 2 * np.pi, (1, half_rope))).astype(np.float32), + "eps_val": eps, + } + + def output_tensors(kernel_input: dict[str, Any]) -> dict[str, Any]: + return {"output_0": np.zeros((n_heads + 1, head_dim), dtype=_BF16)} + + UnitTestFramework( + test_manager=test_manager, + kernel_entry=nki_qkv_rms_rope_kernel, + torch_ref=torch_ref_wrapper(nki_qkv_rms_rope_torch_ref), + kernel_input_generator=input_generator, + output_tensor_descriptor=output_tensors, + ).run_test( + test_config=None, + compiler_args=CompilerArgs(logical_nc_config=1, platform_target=platform_target), + inference_args=InferenceArgs(num_runs=_WARMUP_RUNS), + atol=1e-2, + rtol=3e-2, + ) + + # ---------------- indexer q-projection GEMV ---------------- + + _QPROJ_PARAMS = "q_lora_rank, n_index_heads, lnc" + _QPROJ_CASES = [ + (1536, 64, 2), # production: 12 k-tiles x 64 n-tiles, sharded over both cores + (1536, 64, 1), # the same weight on one core -- must agree with the sharded run + (256, 8, 2), + ] + _QPROJ_ABBREVS = {"q_lora_rank": "r", "n_index_heads": "h", "lnc": "lnc"} + + @pytest.mark.fast + @pytest_parametrize(_QPROJ_PARAMS, _QPROJ_CASES, abbrevs=_QPROJ_ABBREVS) + def test_indexer_qproj_gemv( + self, + test_manager: Orchestrator, + platform_target: Platforms, + q_lora_rank: int, + n_index_heads: int, + lnc: int, + ): + """The indexer query projection as a hand-tiled GEMV, returned transposed. + + Run at both ``lnc=1`` and ``lnc=2``. The two-core launch shards the N axis, + with each core writing a disjoint half of a NAMED ``shared_hbm`` output; the + one-core launch writes all of it. Both are graded against the same + reference, so a half that never landed fails rather than passing as + plausible numbers. + """ + rng = _rng() + n_ktiles = q_lora_rank // 128 + n = n_index_heads * _INDEX_HEAD_DIM + + def input_generator(test_config): + # wT[t, kk, n] == wq_b.weight[n, t * 128 + kk], the host-side pre-tiling. + weight = rng.standard_normal((n, q_lora_rank)).astype(np.float32) / np.sqrt(q_lora_rank) + w_t = weight.reshape(n, n_ktiles, 128).transpose(1, 2, 0) + return { + "wT": _bf16(np.ascontiguousarray(w_t)), + "qr_in": _bf16(rng.standard_normal((1, q_lora_rank)) * 0.5), + } + + def output_tensors(kernel_input: dict[str, Any]) -> dict[str, Any]: + return {"output_0": np.zeros((128, n // 128), dtype=_BF16)} + + UnitTestFramework( + test_manager=test_manager, + kernel_entry=nki_indexer_qproj_gemv, + torch_ref=torch_ref_wrapper(nki_indexer_qproj_gemv_torch_ref), + kernel_input_generator=input_generator, + output_tensor_descriptor=output_tensors, + ).run_test( + test_config=None, + compiler_args=CompilerArgs(logical_nc_config=lnc, platform_target=platform_target), + inference_args=InferenceArgs(num_runs=_WARMUP_RUNS), + atol=2e-2, + rtol=5e-2, + ) + + # ---------------- indexer scoring ---------------- + + _SCORE_PARAMS = "t_c, n_index_heads" + _SCORE_CASES = [(1024, 64), (512, 8)] + _SCORE_ABBREVS = {"t_c": "tc", "n_index_heads": "h"} + + @pytest.mark.fast + @pytest_parametrize(_SCORE_PARAMS, _SCORE_CASES, abbrevs=_SCORE_ABBREVS) + def test_indexer_score(self, test_manager: Orchestrator, platform_target: Platforms, t_c: int, n_index_heads: int): + """Raw indexer scores plus the causal bias, one core, ``[S_q, T_c]`` out. + + The relu is applied per head BEFORE the per-head weight, so a head with a + negative dot product contributes nothing rather than contributing + negatively. That ordering is what makes every real score non-negative, which + the top-k padding downstream depends on. + """ + + def input_generator(test_config): + inputs = _indexer_inputs(t_c, n_index_heads) + # Zero bias: decode has no causal masking to apply within the cache, and + # a zero table is what the block passes. + inputs["causal_bias"] = np.zeros((_S_Q, t_c), dtype=np.float32) + return inputs + + def output_tensors(kernel_input: dict[str, Any]) -> dict[str, Any]: + return {"output_0": np.zeros((_S_Q, t_c), dtype=np.float32)} + + UnitTestFramework( + test_manager=test_manager, + kernel_entry=nki_indexer_score_kernel, + torch_ref=torch_ref_wrapper(nki_indexer_score_torch_ref), + kernel_input_generator=input_generator, + output_tensor_descriptor=output_tensors, + ).run_test( + test_config=None, + compiler_args=CompilerArgs(logical_nc_config=1, platform_target=platform_target), + inference_args=InferenceArgs(num_runs=_WARMUP_RUNS), + atol=5e-2, + rtol=5e-2, + ) + + @pytest.mark.fast + @pytest_parametrize(_SCORE_PARAMS, [(1024, 64)], abbrevs=_SCORE_ABBREVS) + def test_indexer_score_2core( + self, test_manager: Orchestrator, platform_target: Platforms, t_c: int, n_index_heads: int + ): + """Two-core scoring of disjoint ``T_c`` halves into one shared score row. + + This is the test for the shared-buffer hand-off: both cores write disjoint + halves of a NAMED ``shared_hbm`` row, and if that allocation were anonymous + each core would get a private copy and the returned row would carry core 0's + half with core 1's left as zeros. The reference scores the whole range, so + that failure shows up as a mismatch over the second half rather than as + plausible output. + """ + + def input_generator(test_config): + return _indexer_inputs(t_c, n_index_heads) + + def output_tensors(kernel_input: dict[str, Any]) -> dict[str, Any]: + return {"output_0": np.zeros((1, t_c), dtype=_BF16)} + + UnitTestFramework( + test_manager=test_manager, + kernel_entry=nki_indexer_score_2core, + torch_ref=torch_ref_wrapper(nki_indexer_score_2core_torch_ref), + kernel_input_generator=input_generator, + output_tensor_descriptor=output_tensors, + ).run_test( + test_config=None, + compiler_args=CompilerArgs(logical_nc_config=2, platform_target=platform_target), + inference_args=InferenceArgs(num_runs=_WARMUP_RUNS), + atol=5e-2, + rtol=5e-2, + ) + + # ---------------- O(k) gathered attention ---------------- + + _GATHER_PARAMS = "n_heads, head_dim, rope_head_dim, t_c, k" + _GATHER_CASES = [ + (32, 512, 64, 1024, 256), # production head/dim shape, 2 gather chunks + (32, 512, 64, 2048, 512), # 4 chunks and a full-width score group + (8, 256, 64, 512, 128), # one chunk, and a head count well under 128 + ] + _GATHER_ABBREVS = {"n_heads": "h", "head_dim": "d", "rope_head_dim": "rd", "t_c": "tc", "k": "k"} + + @pytest.mark.fast + @pytest_parametrize(_GATHER_PARAMS, _GATHER_CASES, abbrevs=_GATHER_ABBREVS) + def test_decode_gather_attention( + self, + test_manager: Orchestrator, + platform_target: Platforms, + n_heads: int, + head_dim: int, + rope_head_dim: int, + t_c: int, + k: int, + ): + """O(k) attention on caller-supplied indices, with the output de-RoPE fused in. + + The selected positions are SHUFFLED rather than a contiguous or monotone + run: the indirect gather reads them in whatever order the index row holds, + so a monotone list would hide an off-by-one in the gather's address + arithmetic. It also passes the same position more than once nowhere, so each + gathered row is distinguishable. + """ + rng = _rng() + half_rope = rope_head_dim // 2 + s_len = 1 + + def input_generator(test_config): + idx = rng.permutation(t_c)[:k].astype(np.uint32).reshape(k, s_len) + return { + "topk_indices_T": idx, + # Already scaled by softmax_scale, as the block hands it over. + "all_q_T": _f16(rng.standard_normal((head_dim, n_heads * s_len)) * (head_dim**-0.5)), + "win_K_T": _f16(rng.standard_normal((head_dim, _WINDOW)) * 0.3), + "win_V": _f16(rng.standard_normal((_WINDOW, head_dim)) * 0.3), + "compress_kv": _bf16(rng.standard_normal((t_c, head_dim)) * 0.3), + "attn_sink_in": rng.standard_normal((1, n_heads)).astype(np.float32) * 0.5, + "derope_cos": np.cos(rng.uniform(0, 2 * np.pi, (1, half_rope))).astype(np.float32), + "derope_sin": np.sin(rng.uniform(0, 2 * np.pi, (1, half_rope))).astype(np.float32), + } + + def output_tensors(kernel_input: dict[str, Any]) -> dict[str, Any]: + return {"output_0": np.zeros((n_heads * s_len, head_dim), dtype=_BF16)} + + UnitTestFramework( + test_manager=test_manager, + kernel_entry=nki_decode_gather_ok_kernel, + torch_ref=torch_ref_wrapper(nki_decode_gather_ok_torch_ref), + kernel_input_generator=input_generator, + output_tensor_descriptor=output_tensors, + ).run_test( + test_config=None, + compiler_args=CompilerArgs(logical_nc_config=1, platform_target=platform_target), + inference_args=InferenceArgs(num_runs=_WARMUP_RUNS), + atol=2e-2, + rtol=5e-2, + ) + + # ---------------- the fused megakernel ---------------- + + _FUSED_PARAMS = "n_heads, head_dim, rope_head_dim, n_index_heads, t_c, k, n_val" + _FUSED_CASES = [ + # k divides 2 * COMP_CHUNK, so the kernel takes the K-SPLIT path: the + # gathered positions are halved across the cores and the softmax is + # recombined with two sendrecv exchanges. + (32, 512, 64, 64, 1024, 256, 2048), + # k does NOT divide 256, so the K-split is rejected and T_c <= 4096 selects + # the HEAD-SPLIT path instead -- the two cores take disjoint head blocks + # after a second core_barrier publishes the winners. + (32, 512, 64, 64, 1024, 384, 2048), + # The production 32K-context shape. k still divides 2 * COMP_CHUNK so this is + # also the K-split path, but at 8x the cache: the score stage runs 8 chunks per + # core and the top-k runs at its full pinned width with no padding tail. + (32, 512, 64, 64, 8192, 1024, 8192), + # 8K context, but the top-k still runs at the pinned n_val=8192 -- so three + # QUARTERS of the score row is sentinel padding. This is what the block + # actually traces at seq_len=8192, and it is the case where the padding + # dominates: a top-k that mishandles a mostly-padded row is wrong here and + # right at both of the shapes above (one has no padding, the other half). + (32, 512, 64, 64, 2048, 1024, 8192), + ] + _FUSED_ABBREVS = { + "n_heads": "h", + "head_dim": "d", + "rope_head_dim": "rd", + "n_index_heads": "ih", + "t_c": "tc", + "k": "k", + "n_val": "n", + } + + @pytest.mark.fast + @pytest_parametrize(_FUSED_PARAMS, _FUSED_CASES[:2], abbrevs=_FUSED_ABBREVS) + def test_score_topk_gather_fused( + self, + test_manager: Orchestrator, + platform_target: Platforms, + n_heads: int, + head_dim: int, + rope_head_dim: int, + n_index_heads: int, + t_c: int, + k: int, + n_val: int, + ): + """Indexer score, GpSimd top-k and O(k) attention in ONE launch, graded end to end. + + Nothing between the three stages is observable, so this single comparison is + what covers the snake reformat, the cross-core barrier, the ``nisa.topk`` + call and the gather that reads its winners back. The parametrization picks + which cross-core strategy the kernel compiles to -- K-split or head-split -- + via ``k`` and ``T_c``, both trace-time constants. + """ + self._run_fused(test_manager, platform_target, n_heads, head_dim, rope_head_dim, n_index_heads, t_c, k, n_val) + + @pytest_parametrize(_FUSED_PARAMS, _FUSED_CASES[2:], abbrevs=_FUSED_ABBREVS) + def test_score_topk_gather_fused_large( + self, + test_manager: Orchestrator, + platform_target: Platforms, + n_heads: int, + head_dim: int, + rope_head_dim: int, + n_index_heads: int, + t_c: int, + k: int, + n_val: int, + ): + """The fused kernel at the production 32K-context cache size.""" + self._run_fused(test_manager, platform_target, n_heads, head_dim, rope_head_dim, n_index_heads, t_c, k, n_val) + + def _run_fused(self, test_manager, platform_target, n_heads, head_dim, rope_head_dim, n_index_heads, t_c, k, n_val): + rng = _rng() + half_rope = rope_head_dim // 2 + s_len = 1 + + def input_generator(test_config): + inputs = _indexer_inputs(t_c, n_index_heads, k=k) + inputs.update( + { + "k_val": k, + "n_val": n_val, + "all_q_T": _f16(rng.standard_normal((head_dim, n_heads * s_len)) * (head_dim**-0.5)), + "win_K_T": _f16(rng.standard_normal((head_dim, _WINDOW)) * 0.3), + "win_V": _f16(rng.standard_normal((_WINDOW, head_dim)) * 0.3), + "compress_kv": _bf16(rng.standard_normal((t_c, head_dim)) * 0.3), + "attn_sink_in": rng.standard_normal((1, n_heads)).astype(np.float32) * 0.5, + "derope_cos": np.cos(rng.uniform(0, 2 * np.pi, (1, half_rope))).astype(np.float32), + "derope_sin": np.sin(rng.uniform(0, 2 * np.pi, (1, half_rope))).astype(np.float32), + } + ) + return inputs + + def output_tensors(kernel_input: dict[str, Any]) -> dict[str, Any]: + return {"output_0": np.zeros((n_heads * s_len, head_dim), dtype=_BF16)} + + UnitTestFramework( + test_manager=test_manager, + kernel_entry=nki_indexer_score_topk_gather_2core, + torch_ref=torch_ref_wrapper(nki_indexer_score_topk_gather_2core_torch_ref), + kernel_input_generator=input_generator, + output_tensor_descriptor=output_tensors, + ).run_test( + test_config=None, + compiler_args=CompilerArgs(logical_nc_config=2, platform_target=platform_target), + inference_args=InferenceArgs(num_runs=_WARMUP_RUNS), + atol=2e-2, + rtol=5e-2, + ) + + # ---------------- compile-only coverage of the index-returning kernels ---------------- + + @pytest.mark.fast + def test_topk_snake_compiles(self, test_manager: Orchestrator, platform_target: Platforms): + """``nisa_topk_snake_kernel`` traces and compiles on the snake-encoded layout. + + Compile-only: the kernel returns the top-k values AND their positions, and + ``nisa.topk`` orders neither, so an elementwise comparison against + ``torch.topk`` would flag orderings that are equally correct. + """ + rng = _rng() + n_val, k_val = 2048, 64 + src_x = n_val // 16 + kernel_input = { + "in_tensor": _bf16(rng.standard_normal((128, src_x))), + "k_val": k_val, + "n_val": n_val, + } + UnitTestFramework( + test_manager=test_manager, + kernel_entry=nisa_topk_snake_kernel, + kernel_input_generator=lambda _: kernel_input, + trace_only=True, + ).run_test( + test_config=None, + compiler_args=CompilerArgs(logical_nc_config=1, platform_target=platform_target), + ) + + @pytest.mark.fast + def test_score_topk_single_core_compiles(self, test_manager: Orchestrator, platform_target: Platforms): + """``nki_indexer_score_topk_kernel`` traces and compiles (fallback path). + + Compile-only for the same reason as above: the returned winners are an + unordered set written into row 0 of an 8-row buffer whose other rows are + never written. + """ + kernel_input = dict(_indexer_inputs(2048, 64, k=128), k_val=128) + UnitTestFramework( + test_manager=test_manager, + kernel_entry=nki_indexer_score_topk_kernel, + kernel_input_generator=lambda _: kernel_input, + trace_only=True, + ).run_test( + test_config=None, + compiler_args=CompilerArgs(logical_nc_config=1, platform_target=platform_target), + ) + + @pytest.mark.fast + def test_score_topk_2core_compiles(self, test_manager: Orchestrator, platform_target: Platforms): + """``nki_indexer_score_topk_2core`` traces and compiles on the ``[2]`` grid. + + Compile-only for the unordered-set reason; its scoring stage is graded + numerically by ``test_indexer_score_2core`` and its top-k by the fused + kernel, both of which share this code verbatim. + """ + kernel_input = dict(_indexer_inputs(1024, 64, k=128), k_val=128, n_val=2048) + UnitTestFramework( + test_manager=test_manager, + kernel_entry=nki_indexer_score_topk_2core, + kernel_input_generator=lambda _: kernel_input, + trace_only=True, + ).run_test( + test_config=None, + compiler_args=CompilerArgs(logical_nc_config=2, platform_target=platform_target), + ) diff --git a/test/integration/nkilib/experimental/deepseek_v4_csa/test_csa_prefill_attention.py b/test/integration/nkilib/experimental/deepseek_v4_csa/test_csa_prefill_attention.py new file mode 100644 index 0000000..737652b --- /dev/null +++ b/test/integration/nkilib/experimental/deepseek_v4_csa/test_csa_prefill_attention.py @@ -0,0 +1,663 @@ +# 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. +"""Integration tests for the DeepSeek-V4 CSA prefill kernels. + +All five prefill kernels are graded numerically. Two things about the inputs are +load-bearing rather than incidental: + +* The selection masks are built with CAUSAL MASKING ALREADY BAKED IN, as the + indexer produces them: a compressed position is selectable only once every raw + token it covers is at or before the query. ``nki_gather_csa_attn_kernel`` + truncates its compressed loop at a per-tile compile-time causal bound and relies + on everything past that bound being ``-1e9``, while the reference attends over all + ``T_c`` columns -- so a bound that was ever too TIGHT would show up as a mismatch. +* The window bias tables come from the library's own + ``precompute_win_bias_parts``, so the kernel and the reference agree on which + window position each row may attend and where its attention sink sits. + +These kernels use ``priority=`` DMA class-of-service hints, which exist only on +NeuronCore-v4, so every test here is trn3-only. +""" + +from typing import Any, final + +import ml_dtypes +import neuron_dtypes as dt +import nki.language as nl +import numpy as np +import pytest +from nkilib_src.nkilib.experimental.deepseek_v4_csa.csa_common import precompute_win_bias_parts +from nkilib_src.nkilib.experimental.deepseek_v4_csa.csa_prefill_attention import ( + nki_compressor_core_kernel, + nki_fused_csa_attn_kernel, + nki_gather_csa_attn_kernel, + nki_indexer_score_mask_kernel, + nki_prefill_sparse_attn_kernel, + nki_prefill_topk_kernel, + nki_rms_rope_kernel, +) +from nkilib_src.nkilib.experimental.deepseek_v4_csa.csa_prefill_attention_torch import ( + nki_compressor_core_torch_ref, + nki_fused_csa_attn_torch_ref, + nki_gather_csa_attn_torch_ref, + nki_indexer_score_mask_torch_ref, + nki_prefill_sparse_attn_torch_ref, + nki_rms_rope_torch_ref, +) + +from test.utils.common_dataclasses import CompilerArgs, InferenceArgs, Platforms +from test.utils.pytest_parametrize import pytest_parametrize +from test.utils.pytest_test_metadata import pytest_marks, pytest_test_metadata +from test.utils.test_orchestrator import Orchestrator +from test.utils.unit_test_framework import UnitTestFramework, torch_ref_wrapper + +pytestmark = pytest.mark.platforms(exclude=list(set(Platforms) - {Platforms.TRN3, Platforms.TRN3_A0})) + +_INDEX_HEAD_DIM = 128 +_WINDOW = 128 +_WIN_SIZE = 2 * _WINDOW +_NEG_INF = -1e9 +_BF16 = ml_dtypes.bfloat16 + +# Every device test runs the kernel twice and grades the SECOND execution +# (`--save-nth-output` tracks `num_runs`), i.e. one warmup then one measured run. +# +# This is load-bearing, not boilerplate: the warmup makes the tests measure steady +# state, which is what they are for. It does mean they do NOT cover +# first-execution behaviour -- grading run 0 is a separate exercise, and dropping +# `inference_args` to save a run changes what these tests assert. +_WARMUP_RUNS = 2 + + +def _rng(seed: int = 42) -> np.random.Generator: + return np.random.default_rng(seed) + + +def _bf16(x: np.ndarray) -> np.ndarray: + return dt.static_cast(x.astype(np.float32), nl.bfloat16) + + +def _f16(x: np.ndarray) -> np.ndarray: + return x.astype(np.float16) + + +def _hadamard(n: int) -> np.ndarray: + """Orthonormal ``[n, n]`` Sylvester Hadamard matrix, matching the model's own.""" + h = np.ones((1, 1), dtype=np.float32) + while h.shape[0] < n: + h = np.block([[h, h], [h, -h]]) + return h * n**-0.5 + + +def _window_bias(s_len: int) -> tuple[np.ndarray, np.ndarray]: + """The library's own ``[S, 2W]`` window bias tables, as numpy fp32. + + Taking them from the library rather than rebuilding them here is deliberate: the + kernel and the reference must agree on which window column each row may attend + and which one carries the attention sink, and rebuilding the rule twice is how + that agreement quietly breaks. + """ + base, sink = precompute_win_bias_parts(s_len, _WINDOW) + return base.numpy().astype(np.float32), sink.numpy().astype(np.float32) + + +def _causal_selection_bias(s_len: int, t_c: int, ratio: int, split_pos: int, density: float) -> np.ndarray: + """A ``0 / -1e9`` compressed-selection mask with causal masking baked in. + + Compressed position ``t`` pools raw tokens ``[t * ratio, (t + 1) * ratio)``, so a + query at global position ``p`` may attend it only once ``(t + 1) * ratio - 1 <= p``. + Among the positions that pass that test, a random ``density`` fraction is marked + selected, mimicking the indexer's top-k. Everything else is ``-1e9``, which is + what makes the additive mask both the selection and the causal predicate. + """ + rng = _rng(7) + positions = split_pos + np.arange(s_len) + covered_by = (np.arange(t_c) + 1) * ratio - 1 + causal = covered_by[None, :] <= positions[:, None] + chosen = rng.random((s_len, t_c)) < density + return np.where(causal & chosen, 0.0, _NEG_INF).astype(np.float32) + + +@final +@pytest_test_metadata(name="DeepSeek-V4 CSA Prefill") +@pytest_marks(["attention", "deepseek_v4_csa"]) +class TestCsaPrefillAttention: + """Prefill-side CSA kernels: projection tail, compressor, indexer mask, attention.""" + + # ---------------- fused RMS + RoPE projection tail ---------------- + + _RMS_ROPE_PARAMS = "s_rows, head_dim, rope_head_dim, do_rms, inverse, with_gain" + _RMS_ROPE_CASES = [ + # kv-path: learnable RMSNorm gain then forward rotation. + (256, 512, 64, 1, 0, True), + # q-path: per-head RMS with no gain, forward rotation. + (256, 512, 64, 1, 0, False), + # output de-RoPE: rotation only, inverted, no norm. + (256, 512, 64, 0, 1, False), + (128, 128, 32, 1, 0, True), + ] + _RMS_ROPE_ABBREVS = { + "s_rows": "s", + "head_dim": "d", + "rope_head_dim": "rd", + "do_rms": "rms", + "inverse": "inv", + "with_gain": "gain", + } + + @pytest.mark.fast + @pytest_parametrize(_RMS_ROPE_PARAMS, _RMS_ROPE_CASES, abbrevs=_RMS_ROPE_ABBREVS) + def test_rms_rope( + self, + test_manager: Orchestrator, + platform_target: Platforms, + s_rows: int, + head_dim: int, + rope_head_dim: int, + do_rms: int, + inverse: int, + with_gain: bool, + ): + """One kernel covering all three of prefill's projection tails. + + ``do_rms``, ``inverse`` and whether ``gain_in`` is supplied are trace-time + constants, so each combination compiles to its own specialization -- which is + why they are test parameters rather than separate tests. + """ + rng = _rng() + half_rope = rope_head_dim // 2 + + def input_generator(test_config): + angles = rng.uniform(0, 2 * np.pi, (s_rows, half_rope)) + return { + "x_in": _bf16(rng.standard_normal((s_rows, head_dim)) * 0.5), + "cos_in": np.cos(angles).astype(np.float32), + "sin_in": np.sin(angles).astype(np.float32), + "gain_in": rng.uniform(0.8, 1.2, (1, head_dim)).astype(np.float32) if with_gain else None, + "eps_val": 1e-6, + "do_rms": do_rms, + "inverse": inverse, + } + + def output_tensors(kernel_input: dict[str, Any]) -> dict[str, Any]: + return {"output_0": np.zeros((s_rows, head_dim), dtype=_BF16)} + + UnitTestFramework( + test_manager=test_manager, + kernel_entry=nki_rms_rope_kernel, + torch_ref=torch_ref_wrapper(nki_rms_rope_torch_ref), + kernel_input_generator=input_generator, + output_tensor_descriptor=output_tensors, + ).run_test( + test_config=None, + compiler_args=CompilerArgs(logical_nc_config=1, platform_target=platform_target), + inference_args=InferenceArgs(num_runs=_WARMUP_RUNS), + atol=1e-2, + rtol=3e-2, + ) + + # ---------------- compressor ---------------- + + _COMPRESSOR_PARAMS = "t_c, head_dim, rope_head_dim, compress_ratio, lnc, rotate" + _COMPRESSOR_CASES = [ + (256, 512, 64, 4, 1, 0), + (256, 512, 64, 4, 2, 0), + (128, 256, 64, 4, 1, 0), + # rotate=1 is the INDEXER's compressor: head_dim 128, result rotated by an + # orthonormal Hadamard. Both lnc values, because the rotation is per position + # tile and must not depend on which core owns the tile. + (256, 128, 64, 4, 1, 1), + (256, 128, 64, 4, 2, 1), + ] + _COMPRESSOR_ABBREVS = { + "t_c": "tc", + "head_dim": "d", + "rope_head_dim": "rd", + "compress_ratio": "r", + "lnc": "lnc", + "rotate": "rot", + } + + @pytest.mark.fast + @pytest_parametrize(_COMPRESSOR_PARAMS, _COMPRESSOR_CASES, abbrevs=_COMPRESSOR_ABBREVS) + def test_compressor_core( + self, + test_manager: Orchestrator, + platform_target: Platforms, + t_c: int, + head_dim: int, + rope_head_dim: int, + compress_ratio: int, + lnc: int, + rotate: int, + ): + """Gated pooling over the overlapped slots, then RMSNorm, then RoPE. + + The gate softmax runs over the ``2 * compress_ratio`` slots INDEPENDENTLY PER + CHANNEL, so the gate scores are drawn with real spread across both the slot + and the channel axis -- a per-position-only gate would let a kernel that + collapsed the channel axis still pass. + + Both reductions are per-position, so nothing is reduced across cores; running + at ``lnc=2`` checks that the position tiles really are partitioned rather than + duplicated. + """ + rng = _rng() + ratio2 = 2 * compress_ratio + half_rope = rope_head_dim // 2 + + def input_generator(test_config): + angles = rng.uniform(0, 2 * np.pi, (t_c, half_rope)) + # cos/sin arrive with each pair's angle duplicated across its two channels. + cos_rep = np.repeat(np.cos(angles), 2, axis=1).astype(np.float32) + sin_rep = np.repeat(np.sin(angles), 2, axis=1).astype(np.float32) + inputs = { + # kv8/score8 are BF16: the block hands the kernel its bf16 projection + # output directly and the kernel widens on load, so ape (the fp32 gate + # bias) is added inside the kernel rather than by the caller. + "kv8": _bf16(rng.standard_normal((t_c, ratio2, head_dim)) * 0.5), + "score8": _bf16(rng.standard_normal((t_c, ratio2, head_dim)) * 1.5), + "norm_weight": rng.uniform(0.8, 1.2, (1, head_dim)).astype(np.float32), + "cos_rep": cos_rep, + "sin_rep": sin_rep, + "eps": 1e-6, + "ape": (rng.standard_normal((ratio2, head_dim)) * 0.3).astype(np.float32), + } + if rotate: + inputs["hadamard"] = _bf16(_hadamard(head_dim)) + return inputs + + def output_tensors(kernel_input: dict[str, Any]) -> dict[str, Any]: + return {"output_0": np.zeros((t_c, head_dim), dtype=_BF16)} + + UnitTestFramework( + test_manager=test_manager, + kernel_entry=nki_compressor_core_kernel, + torch_ref=torch_ref_wrapper(nki_compressor_core_torch_ref), + kernel_input_generator=input_generator, + output_tensor_descriptor=output_tensors, + ).run_test( + test_config=None, + compiler_args=CompilerArgs(logical_nc_config=lnc, platform_target=platform_target), + inference_args=InferenceArgs(num_runs=_WARMUP_RUNS), + atol=1e-2, + rtol=3e-2, + ) + + # ---------------- indexer scoring + bisection selection mask ---------------- + + _MASK_PARAMS = "s_q, t_c, n_index_heads, k, lnc" + _MASK_CASES = [ + (256, 1024, 64, 256, 1), + (256, 1024, 64, 256, 2), + (128, 512, 8, 64, 1), + ] + _MASK_ABBREVS = {"s_q": "sq", "t_c": "tc", "n_index_heads": "h", "k": "k", "lnc": "lnc"} + + @pytest.mark.fast + @pytest_parametrize(_MASK_PARAMS, _MASK_CASES, abbrevs=_MASK_ABBREVS) + def test_indexer_score_mask( + self, + test_manager: Orchestrator, + platform_target: Platforms, + s_q: int, + t_c: int, + n_index_heads: int, + k: int, + lnc: int, + ): + """Indexer scoring plus the bisection threshold that becomes the selection mask. + + The threshold comes from a fixed 9 rounds of bisection, so it does not + necessarily admit exactly ``k`` positions -- the reference runs the same + bisection rather than ``torch.topk``, which is what makes the comparison a + test of the kernel instead of a test of the algorithm's approximation. + + Each query row gets its own causal bias, so the rows have genuinely different + valid ranges and therefore different thresholds; a kernel that computed one + threshold for the whole tile would fail. + """ + rng = _rng() + + def input_generator(test_config): + q = rng.standard_normal((_INDEX_HEAD_DIM, n_index_heads, s_q)).astype(np.float32) * 0.5 + q_T_all = q.transpose(0, 1, 2).reshape(_INDEX_HEAD_DIM, n_index_heads * s_q) + + # Two widely separated score clusters, k positions high and the rest at + # 1/20th scale, sharing one direction so scores WITHIN a cluster are + # exactly equal. The mask is a step function of the score, so smoothly + # varying scores make the threshold land between near-identical values and + # a single bf16 ulp of disagreement flips a position -- a full 1e9 mask + # error rather than a small numeric one. Equal-within-cluster scores make + # `score >= threshold` give the same answer in kernel and reference + # wherever the threshold lands, including when the causal mask leaves a row + # with fewer than k selectable positions. + gains = np.full(t_c, 0.05, dtype=np.float32) + gains[rng.permutation(t_c)[:k]] = 1.0 + kv = (rng.standard_normal((_INDEX_HEAD_DIM, 1)).astype(np.float32) * 0.3) * gains + + # Causal bias: query row s may attend compressed position t only once + # every token t pools sits at or before it. Rows therefore differ. + covered_by = (np.arange(t_c) + 1) * 4 - 1 + causal = covered_by[None, :] <= np.arange(s_q)[:, None] * 4 + causal_bias = np.where(causal, 0.0, _NEG_INF).astype(np.float32) + + return { + "q_T_all": _bf16(np.ascontiguousarray(q_T_all)), + "kv_t": _bf16(kv), + "weights": rng.uniform(0.5, 1.5, (s_q, n_index_heads)).astype(np.float32), + "causal_bias": causal_bias, + "k": k, + } + + def output_tensors(kernel_input: dict[str, Any]) -> dict[str, Any]: + return {"output_0": np.zeros((s_q, t_c), dtype=np.float32)} + + UnitTestFramework( + test_manager=test_manager, + kernel_entry=nki_indexer_score_mask_kernel, + torch_ref=torch_ref_wrapper(nki_indexer_score_mask_torch_ref), + kernel_input_generator=input_generator, + output_tensor_descriptor=output_tensors, + ).run_test( + test_config=None, + compiler_args=CompilerArgs(logical_nc_config=lnc, platform_target=platform_target), + # The mask is exactly 0 or -1e9, so any threshold disagreement is a huge + # absolute difference -- these tolerances admit no wrong bit, they only + # allow the -1e9 magnitude itself to compare equal. + atol=1e-3, + rtol=1e-3, + ) + + # ---------------- mask-predicated sparse attention ---------------- + + _ATTN_PARAMS = "s_len, t_c, n_heads, head_dim, lnc" + _ATTN_CASES = [ + (256, 512, 16, 512, 1), + (256, 512, 32, 256, 2), + ] + _ATTN_ABBREVS = {"s_len": "s", "t_c": "tc", "n_heads": "h", "head_dim": "d", "lnc": "lnc"} + + @pytest.mark.fast + @pytest_parametrize(_ATTN_PARAMS, _ATTN_CASES, abbrevs=_ATTN_ABBREVS) + def test_fused_csa_attention( + self, + test_manager: Orchestrator, + platform_target: Platforms, + s_len: int, + t_c: int, + n_heads: int, + head_dim: int, + lnc: int, + ): + """Sparse attention over ``[window | compressed]`` with one global-max softmax. + + The window and the compressed positions share a single normalization, so the + two contributions are directly comparable and no online rescaling is needed. + Masking is additive before ``exp``, which is what makes an unselected position + contribute exactly zero to both the denominator and the value sum. + """ + rng = _rng() + split = s_len + _WINDOW + base, sink = _window_bias(s_len) + + def input_generator(test_config): + total = split + t_c + return { + "compress_sel": _bf16(_causal_selection_bias(s_len, t_c, 4, 0, 0.3)), + # Already scaled by softmax_scale, as the core hands it over. + "all_q_T": _f16(rng.standard_normal((head_dim, n_heads * s_len)) * (head_dim**-0.5)), + "all_K_T": _f16(rng.standard_normal((head_dim, total)) * 0.3), + "all_V": _bf16(rng.standard_normal((total, head_dim)) * 0.3), + "win_bias_base_in": base, + "win_bias_sink_in": sink, + "attn_sink_in": (rng.standard_normal((1, n_heads)) * 0.5).astype(np.float32), + } + + def output_tensors(kernel_input: dict[str, Any]) -> dict[str, Any]: + return {"output_0": np.zeros((n_heads * s_len, head_dim), dtype=_BF16)} + + UnitTestFramework( + test_manager=test_manager, + kernel_entry=nki_fused_csa_attn_kernel, + torch_ref=torch_ref_wrapper(nki_fused_csa_attn_torch_ref), + kernel_input_generator=input_generator, + output_tensor_descriptor=output_tensors, + ).run_test( + test_config=None, + compiler_args=CompilerArgs(logical_nc_config=lnc, platform_target=platform_target), + inference_args=InferenceArgs(num_runs=_WARMUP_RUNS), + atol=2e-2, + rtol=5e-2, + ) + + # ---------------- causal-bound sparse attention ---------------- + + _GATHER_PARAMS = "s_len, t_c, n_heads, head_dim, split_pos, ratio, lnc" + _GATHER_CASES = [ + (256, 512, 32, 512, 1024, 4, 1), + (256, 512, 32, 256, 1024, 4, 2), + # split_pos = 0 makes the causal bound at its tightest, so the first query + # tile reaches only the leading compressed chunk. + (256, 512, 32, 256, 0, 4, 1), + # A/B shape vs the sparse kernel above: matched s_len / t_c / n_heads / + # head_dim, split_pos chosen so the causal bound spans ALL of t_c. + (128, 2048, 128, 512, 8064, 4, 1), + # Establish the dense kernel's slope in t_c by MEASUREMENT rather than assuming + # linearity, at both the sequence-parallel head count (128) and the current + # head-parallel one (32). split_pos is set so the causal bound spans all of t_c. + (128, 4096, 128, 512, 16256, 4, 1), + (128, 8192, 128, 512, 32640, 4, 1), + (128, 2048, 32, 512, 8064, 4, 1), + (128, 4096, 32, 512, 16256, 4, 1), + (128, 8192, 32, 512, 32640, 4, 1), + # lnc=2 so the dense A/B baseline uses BOTH LNC cores, matching the sparse + # kernel's [2] grid. Comparing a 2-core sparse kernel against a 1-core dense + # one would overstate the win. + (128, 2048, 128, 512, 8064, 4, 2), + (128, 4096, 128, 512, 16256, 4, 2), + (128, 8192, 128, 512, 32640, 4, 2), + ] + _GATHER_ABBREVS = { + "s_len": "s", + "t_c": "tc", + "n_heads": "h", + "head_dim": "d", + "split_pos": "sp", + "ratio": "r", + "lnc": "lnc", + } + + @pytest.mark.fast + @pytest_parametrize(_GATHER_PARAMS, _GATHER_CASES, abbrevs=_GATHER_ABBREVS) + def test_gather_csa_attention( + self, + test_manager: Orchestrator, + platform_target: Platforms, + s_len: int, + t_c: int, + n_heads: int, + head_dim: int, + split_pos: int, + ratio: int, + lnc: int, + ): + """The same attention with the compressed loop capped at a compile-time causal bound. + + The reference deliberately attends over ALL ``T_c`` compressed columns while + the kernel stops at its per-tile bound. That only agrees because every + column past the bound is ``-1e9`` in the selection mask and so contributes + exactly zero -- so if the bound were ever too tight, dropping a column that + mattered, the two would disagree. ``split_pos=0`` is included because it makes + the bound tightest, where an off-by-one is most likely. + """ + rng = _rng() + base, sink = _window_bias(s_len) + + def input_generator(test_config): + return { + "topk_sel_bias": _bf16(_causal_selection_bias(s_len, t_c, ratio, split_pos, 0.3)), + "all_q_T": _f16(rng.standard_normal((head_dim, n_heads * s_len)) * (head_dim**-0.5)), + "all_K_T_win": _f16(rng.standard_normal((head_dim, s_len + _WINDOW)) * 0.3), + "all_V_win": _bf16(rng.standard_normal((s_len + _WINDOW, head_dim)) * 0.3), + "compress_kv_T": _f16(rng.standard_normal((head_dim, t_c)) * 0.3), + "compress_kv": _bf16(rng.standard_normal((t_c, head_dim)) * 0.3), + "win_bias_base_in": base, + "win_bias_sink_in": sink, + "attn_sink_in": (rng.standard_normal((1, n_heads)) * 0.5).astype(np.float32), + "split_pos": split_pos, + "ratio": ratio, + } + + def output_tensors(kernel_input: dict[str, Any]) -> dict[str, Any]: + return {"output_0": np.zeros((n_heads * s_len, head_dim), dtype=_BF16)} + + UnitTestFramework( + test_manager=test_manager, + kernel_entry=nki_gather_csa_attn_kernel, + torch_ref=torch_ref_wrapper(nki_gather_csa_attn_torch_ref), + kernel_input_generator=input_generator, + output_tensor_descriptor=output_tensors, + ).run_test( + test_config=None, + compiler_args=CompilerArgs(logical_nc_config=lnc, platform_target=platform_target), + inference_args=InferenceArgs(num_runs=_WARMUP_RUNS), + atol=2e-2, + rtol=5e-2, + ) + + # ---------------- TRUE sparse prefill (per-query indirect gather) ---------------- + + _SPARSE_PARAMS = "s_len, t_c, k_val, n_heads, head_dim, q_base" + _SPARSE_CASES = [ + (8, 2048, 1024, 128, 512, 4096), + (8, 8192, 1024, 128, 512, 16384), + (16, 8192, 1024, 128, 512, 16384), + # A/B shapes vs the dense kernel below: one query tile, all 128 heads. + # t_c is varied at fixed k to show the cost is FLAT in context length. + (128, 2048, 1024, 128, 512, 8192), + (128, 8192, 1024, 128, 512, 8192), + # 256 is the tile size the block actually launches with. + (256, 4096, 1024, 128, 512, 8192), + (256, 8192, 1024, 128, 512, 8192), + ] + _SPARSE_ABBREVS = { + "s_len": "s", + "t_c": "tc", + "k_val": "k", + "n_heads": "h", + "head_dim": "d", + "q_base": "qb", + } + + @pytest.mark.fast + @pytest_parametrize(_SPARSE_PARAMS, _SPARSE_CASES, abbrevs=_SPARSE_ABBREVS) + def test_prefill_sparse_attention( + self, + test_manager: Orchestrator, + platform_target: Platforms, + s_len: int, + t_c: int, + k_val: int, + n_heads: int, + head_dim: int, + q_base: int, + ): + """O(k) sparse prefill: gather each query's selected rows instead of masking all T_c. + + Graded against a reference that gathers the same rows, so a wrong index, a + wrong window slice or a mis-shared softmax denominator all show up as a + numeric failure. ``t_c`` is varied at fixed ``k`` because the kernel's work + must be independent of context length -- only the gather addresses change. + """ + rng = _rng() + + def input_generator(test_config): + idx = np.stack([rng.permutation(t_c)[:k_val] for _ in range(s_len)], axis=1).astype(np.uint32) + return { + "topk_idx_T": idx, + "all_q": _f16(rng.standard_normal((s_len * n_heads, head_dim)) * (head_dim**-0.5)), + "all_K_T_win": _f16(rng.standard_normal((head_dim, s_len + _WINDOW)) * 0.3), + "all_V_win": _f16(rng.standard_normal((s_len + _WINDOW, head_dim)) * 0.3), + "compress_kv": _f16(rng.standard_normal((t_c, head_dim)) * 0.3), + "attn_sink_in": (rng.standard_normal((1, n_heads)) * 0.5).astype(np.float32), + } + + def output_tensors(kernel_input: dict[str, Any]) -> dict[str, Any]: + return {"output_0": np.zeros((n_heads * s_len, head_dim), dtype=_BF16)} + + UnitTestFramework( + test_manager=test_manager, + kernel_entry=nki_prefill_sparse_attn_kernel, + torch_ref=torch_ref_wrapper(nki_prefill_sparse_attn_torch_ref), + kernel_input_generator=input_generator, + output_tensor_descriptor=output_tensors, + ).run_test( + test_config=None, + compiler_args=CompilerArgs(logical_nc_config=2, platform_target=platform_target), + inference_args=InferenceArgs(num_runs=_WARMUP_RUNS), + atol=2e-2, + rtol=5e-2, + ) + + # ---------------- fused on-chip per-query top-k ---------------- + + _TOPK_PARAMS = "s_q, t_c, k_val, n_val" + _TOPK_CASES = [ + (256, 4096, 1024, 8192), + (256, 8192, 1024, 8192), + (2048, 8192, 1024, 8192), + ] + _TOPK_ABBREVS = {"s_q": "sq", "t_c": "tc", "k_val": "k", "n_val": "n"} + + @pytest.mark.fast + @pytest_parametrize(_TOPK_PARAMS, _TOPK_CASES, abbrevs=_TOPK_ABBREVS) + def test_prefill_topk( + self, + test_manager: Orchestrator, + platform_target: Platforms, + s_q: int, + t_c: int, + k_val: int, + n_val: int, + ): + """Per-query top-k positions for the sparse prefill's gather. + + ``trace_only``, because the kernel returns the k winners as an UNORDERED SET + (nisa.topk emits each snake partition's winners in ascending POSITION order, + not by value), so there is no elementwise oracle: a correct result is a + permutation of ``torch.topk``'s indices, and in bf16 -- where the k-th and + (k+1)-th scores are frequently exact ties -- not even the same set. What this + guards is that every shape the block dispatches still compiles and allocates; + the numeric contract on the indices is graded end-to-end by + ``test_csa_block``'s prefill cases, which fail if a selected position is wrong. + """ + rng = _rng() + + def input_generator(test_config): + # Indexer-shaped scores: non-negative after the relu, with a -1e9 causal tail. + sc = np.abs(rng.standard_normal((s_q, t_c))).astype(np.float32) * 0.5 + frontier = np.minimum((np.arange(s_q) + k_val * 4 + 1) // 4, t_c) + sc[np.arange(t_c)[None, :] >= frontier[:, None]] = -1e9 + return {"scores": sc.astype(_BF16), "k_val": k_val, "n_val": n_val} + + UnitTestFramework( + test_manager=test_manager, + kernel_entry=nki_prefill_topk_kernel, + kernel_input_generator=input_generator, + trace_only=True, + ).run_test( + test_config=None, + compiler_args=CompilerArgs(logical_nc_config=2, platform_target=platform_target), + ) diff --git a/test/integration/nkilib/experimental/deepseek_v4_csa/test_csa_tp_all_reduce.py b/test/integration/nkilib/experimental/deepseek_v4_csa/test_csa_tp_all_reduce.py new file mode 100644 index 0000000..7db3ec5 --- /dev/null +++ b/test/integration/nkilib/experimental/deepseek_v4_csa/test_csa_tp_all_reduce.py @@ -0,0 +1,101 @@ +# 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. +"""Integration tests for the CSA tensor-parallel output all-reduce. + +The reduction is an ordinary sum; what these tests are actually for is everything +the kernel does around it, because each of those has a failure mode that produces a +plausible-looking but HALF-REDUCED answer rather than an error: + +* The collective's ``src``/``dst`` must be freshly allocated ``shared_hbm`` buffers + with an explicit ``name=``, and a collective cannot read or write an IO tensor + directly -- hence the staging copies in and out. +* It must launch on the ``[2]`` grid to match the block's ``lnc=2`` context. A + ``[1]``-grid collective inside that graph fails outright, and hand-splitting it + into two ``program_id``-sliced collectives instead wires only one slice across the + ranks and silently leaves the other unreduced. + +Each rank contributes DIFFERENT data, so a rank whose contribution never arrived +changes the sum. Ranks are given equal-magnitude values rather than a distinguishing +scale, so no single rank dominates and a dropped contribution cannot hide inside +rounding. +""" + +from typing import final + +import numpy as np +import pytest +from nki.collectives import ReplicaGroup +from nkilib_src.nkilib.experimental.deepseek_v4_csa.csa_tp_all_reduce import nki_tp_all_reduce_kernel +from nkilib_src.nkilib.experimental.deepseek_v4_csa.csa_tp_all_reduce_torch import nki_tp_all_reduce_torch_ref + +from test.utils.common_dataclasses import CompilerArgs, InferenceArgs, Platforms +from test.utils.pytest_parametrize import pytest_parametrize +from test.utils.pytest_test_metadata import pytest_marks, pytest_test_metadata +from test.utils.test_orchestrator import Orchestrator +from test.utils.unit_test_collective_framework import CollectiveUnitTestFramework + +pytestmark = pytest.mark.platforms(exclude=list(set(Platforms) - {Platforms.TRN3, Platforms.TRN3_A0})) + +# The block reshapes its [B, S, dim] partial to a balanced [P, F] tile before +# reducing. P = 128 with dim = 7168 gives F = 56, the shape the block emits. +_PARTITIONS = 128 +_FREE = 56 + + +@final +@pytest_test_metadata(name="DeepSeek-V4 CSA TP All-Reduce") +@pytest_marks(["collectives", "deepseek_v4_csa"]) +@pytest.mark.skip_simulation +@pytest.mark.high_rank +class TestCsaTpAllReduce: + """The 2-LNC collective that sums the head-parallel output partials across ranks.""" + + _PARAMS = "collective_ranks, logical_nc_config" + # 4 ranks x 2 LNC is the production topology; 2 ranks is the cheap smoke case. + _CASES = [(4, 2), (2, 2)] + _ABBREVS = {"collective_ranks": "ranks", "logical_nc_config": "lnc"} + + @pytest.mark.fast + @pytest_parametrize(_PARAMS, _CASES, abbrevs=_ABBREVS) + def test_tp_all_reduce( + self, + test_manager: Orchestrator, + platform_target: Platforms, + collective_ranks: int, + logical_nc_config: int, + ): + """Sum distinct per-rank partials and check every rank returns the full sum.""" + rng = np.random.default_rng(42) + # Distinct data per rank, at comparable magnitudes, so the sum depends on + # every rank without any one of them dominating. + per_rank = rng.standard_normal((collective_ranks, _PARTITIONS, _FREE)).astype(np.float32) + replica_group = ReplicaGroup([list(range(collective_ranks))]) + + def create_inputs(rank_id: int): + return {"input": per_rank[rank_id], "replica_group": replica_group} + + CollectiveUnitTestFramework( + test_manager=test_manager, + kernel_entry=nki_tp_all_reduce_kernel, + torch_ref=nki_tp_all_reduce_torch_ref, + per_rank_input_generator=create_inputs, + collective_ranks=collective_ranks, + ).run_test( + test_config=None, + compiler_args=CompilerArgs(logical_nc_config=logical_nc_config, platform_target=platform_target), + output_keys=["out"], + rtol=1e-3, + atol=1e-3, + inference_args=InferenceArgs(collective_ranks=collective_ranks, enable_determinism_check=True, num_runs=10), + )