This document is relevant for: Trn2, Trn3

Sparse Attention Indexer MX BF16Score Kernel API Reference#

DeepSeek Sparse Attention Indexer — MX projections + BF16 score.

Supported/optimal usage: head_dim == 128; dim and q_lora_rank multiples of 512; batch_size == 1 and S a multiple of P_MAX (=128) on the validated v3_long/h64 configs; LNC sharding degree up to 2.

Background#

The sparse_attention_indexer_mx_bf16score kernel implements the DeepSeek Sparse Attention Indexer with MX projections and a BF16 score path.

API Reference#

Source code for this kernel API can be found at: sparse_attention_indexer_mx_bf16score.py

sparse_attention_indexer_mx_bf16score#

nkilib.experimental.sparse_attention_indexer.sparse_attention_indexer_mx_bf16score(x: nl.NkiTensor, wq_b: nl.NkiTensor, wk: nl.NkiTensor, k_norm_gamma: nl.NkiTensor, k_norm_beta: nl.NkiTensor, weights_proj: nl.NkiTensor, cos: nl.NkiTensor, sin: nl.NkiTensor, k_cache: nl.NkiTensor, mask: nl.NkiTensor, n_heads: int, head_dim: int, rope_head_dim: int, index_topk: int, start_pos: int, use_hadamard: bool = False, batch_size: int = 1, wq_b_scale: Optional[nl.NkiTensor] = None, wk_scale: Optional[nl.NkiTensor] = None, k_scale_cache: Optional[nl.NkiTensor] = None, x_non_mx: Optional[nl.NkiTensor] = None, x_mx_data: Optional[nl.NkiTensor] = None, x_mx_scale: Optional[nl.NkiTensor] = None, qr_qtz_hbm: Optional[nl.NkiTensor] = None, qr_scale_hbm: Optional[nl.NkiTensor] = None, phase: str = 'all', k_seq_out_hbm: Optional[nl.NkiTensor] = None, end_pos_arg: Optional[int] = None, emit_flat_topk: bool = False, emit_tiled_topk: bool = False) → tuple[nl.NkiTensor, nl.NkiTensor]#

DeepSeek Sparse Attention Indexer — MX projections + BF16 score.

Parameters:

  • x (nl.NkiTensor) – [B*S, dim] hidden states.

  • qr_qtz_hbm (Optional[nl.NkiTensor]) – [num_s_tiles, P_MAX, q_lora_rank // (P_MAX*H_PACK), P_MAX] uint32 — pre-quantized (fp8x4-as-uint32) compressed query from the upstream QKV kernel.

  • qr_scale_hbm (Optional[nl.NkiTensor]) – matching uint8 block-32 e8m0 scales for qr_qtz_hbm.

  • wq_b (nl.NkiTensor) – [q_lora_rank // 4, n_heads * head_dim] fp8_e4m3fn_x4.

  • wq_b_scale (Optional[nl.NkiTensor]) – [q_lora_rank // 32, n_heads * head_dim] uint8 e8m0.

  • wk (nl.NkiTensor) – [dim // 4, head_dim] fp8_e4m3fn_x4 (contraction dim on partition).

  • wk_scale (Optional[nl.NkiTensor]) – [dim // 32, head_dim] uint8 e8m0. k_norm_gamma, k_norm_beta: [head_dim] LayerNorm gamma/beta.

  • weights_proj (nl.NkiTensor) – [n_heads, dim] per-head weights projection (always bf16). cos, sin: [S, rope_head_dim // 2] RoPE caches.

  • k_cache (nl.NkiTensor) – [B, head_dim, max_seq_len] bf16 (mutable).

  • k_scale_cache (Optional[nl.NkiTensor]) – ignored (no MX scales in bf16 path).

  • mask (nl.NkiTensor) – [B*S, end_pos] attention mask.

  • index_topk (int) – number of top positions to select per query (topk k).

  • x_non_mx (Optional[nl.NkiTensor]) – pre-cast bf16 view of x for W-projection (skips the f32->bf16 cast). x_mx_data, x_mx_scale: pre-quantized x for K-projection (skips the in-kernel swizzle + quantize_mx). qr_qtz_hbm, qr_scale_hbm: pre-quantized qr latent from an upstream QKV kernel (indexer skips its own qr transpose+norm+quantize). This is the only Q input path — qr is never quantized in-kernel.

  • phase (str) – context-parallel phased split. “all” (default): fused self-attention (queries == keys == S). “kproj”: project K for this rank’s S shard only; write seq-major [S, head_dim] to k_seq_out_hbm, return (that, None), no score. “score”: skip K-proj; score this rank’s S queries vs the pre-gathered full k_cache [B, head_dim, end_pos_arg=S_kv] over [0, S_kv).

  • k_seq_out_hbm (Optional[nl.NkiTensor]) – [S, head_dim] bf16 K output for phase=”kproj”.

  • end_pos_arg (Optional[int]) – key range for phase=”score” (defaults to start_pos + S).

  • emit_flat_topk (bool) – decode the snake topk to per-query positions in-kernel and return flat [B*S, index_topk] int32 (the sparse_mla_latent_attn topk_indices contract). Fuses the out-of-kernel snake->flat decode.

  • emit_tiled_topk (bool) – split-fix output — skip the per-query scatter and return positions in the natural partition-tiled layout [num_s_tiles, NUM_TOPK_BATCHES, P_MAX, index_topk // 16] int32, read per-query by the paired attention kernel. Mutually exclusive with emit_flat_topk; both default False (snake output, decoded by the consumer, e.g. the fused mega kernel).

Returns:

[B*S, end_pos] f32 — pre-topk per-position scores.

Return type:

nl.ndarray

Returns:

[num_S_tiles_total, NUM_TOPK_BATCHES, P_MAX, index_topk] uint32 — hardware-topk output. topk_idx[s_tile, t, 16g + j%16, j//16] holds the position (0-based in [0, end_pos)) of query 8t + g’s j-th largest score (ascending). num_S_tiles_total = batch_size * ceil(S / P_MAX).

Return type:

nl.ndarray

Notes:

  • Q/K/W projections still use nc_matmul_mx for speed, but the score matmul uses bf16 nc_matmul (no Q/K quantization, no Hadamard rotation). This removes the Q swizzle + quantize_mx + Hadamard chain from the per-S-tile critical path.

  • K cache stores bf16 instead of fp8x4 (no need for MX format if downstream consumers use this same kernel; otherwise downstream must coordinate).

  • Algorithm equivalent to the MX variant via Hadamard’s orthogonality: q_H @ (k_H).T = q @ H @ H.T @ k.T = q @ k.T so dropping Hadamard is mathematically a no-op. The MX quantization that Hadamard exists to support is also dropped on the score path. Score precision improves (bf16 score matmul vs MX-rounded score matmul).

Dimensions:

  • B: batch size (= batch_size).

  • S: query sequence length per batch.

  • M: B * S (total query rows, = x.shape[0]).

  • dim: hidden size (= x.shape[1]).

  • q_lora_rank: compressed-query rank (derived from qr_qtz_hbm layout).

  • n_heads: number of indexer heads.

  • head_dim: per-head dimension (== 128).

  • rope_head_dim: RoPE-rotated slice of head_dim.

  • end_pos: start_pos + S (KV positions scored per query).

  • P_MAX: partition tile size (= 128).

This document is relevant for: Trn2, Trn3