Module attn_v41

Module attn_v41 

Source
Expand description

2026-09-25: DeepSeek-V4.1 Flash attention on the GPU: one forward for every layer.

A layer’s LayerRole::ratio selects its kind. Ratio 0 attends over a sliding window of its own fp8-rounded kv rows. Ratio > 0 attends over that window plus compressed positions chosen by an indexer from the latent cache the kv-source layers (is_kv_source) publish in SharedV41. The index selection is computed by the index-source layers and reused by the ratio > 0 layers after them; the candidate-source layer also publishes a block-level candidate mask that later index sources apply.

This is the device form of deepseek_v41_ref::compress::attention_any, and attn_v41_tests.rs compares it with the reference layer by layer:

  • projections on dense_gemm_bf16 / dense_gemv_bf16 (bf16 weights) or the Q2_K K-quant kernels;
  • RMSNorm, RoPE, the fp8/fp4 quantisers, the compressor pooling, the indexer scores, sparse attention with the sink and the wo_a column slices on kernels/gb10/deepseek-v4-flash/nvfp4/attn_v41.cu;
  • the candidate-block and index top-k selections on the CPU through the reference’s torch_cpu_topk_set, which reproduces torch’s CPU tie order among equal scores; the scores come back in one stream-ordered download.

Layer state (window ring, compressor partial group, the latent and index-key caches) lives on the device; the shared index selection and candidate mask in SharedV41 are host vectors.

Owner: model-arch, DeepSeek-V4.1. Invariants:

  • A ratio > 0 layer fails (does not attend) when no kv source has published a latent cache (forward, decode_prep_role).
  • forward refuses a selection wider than window + index_topk or 2048, the sparse_attn score buffer.

Structs§

AttnV41
2026-09-25: The attention runtime: kernels, RoPE tables, and workspaces for up to max_tokens positions per call.
AttnV41Cfg
2026-09-25: Model-wide attention geometry, built from ModelConfig and the METRALE_DS41_* limits by the DeepSeek-V4.1 weight loader.
AttnV41LayerState
2026-09-25: One sequence’s attention state for one layer, across prefill and decode: the bf16 window ring [window, hd]; on kv sources also the compressor’s partial group (f32 kv and score, [ratio, hd] each, used at ratio > 1), the bf16 latent cache [max_seq / ratio, hd] and the bf16 index-key cache [max_seq / ratio, index_hd].
AttnV41LayerWeights
2026-09-25: One layer’s resident attention weights. Projections are [out, in]: wo_a is [groups * o_rank, gw], wo_b [dim, groups * o_rank]. sink ([n_heads]), q_norm ([q_rank]) and kv_norm ([hd]) are f32.
AttnV41Run
2026-09-25: The intermediates of one layer’s forward. q and o are bf16 [tokens, n_heads, hd] (q after RoPE, o before the inverse rotation); rows_a are the window rows (the chunk’s own rows at start_pos == 0, the ring otherwise); rows_b the kv source’s latent cache, of which rows_b_len rows are attended; idx is [tokens][topk], window slots then compressed positions offset by rows_a_len; out is bf16 [tokens, dim].
CompressorWeightsGpu
2026-09-25: Compressor weights (kv sources only). At ratio > 1 kv and gate are f32 [hd, dim]; at ratio 1 kv is bf16 and gate is None. norm is f32 [hd].
IndexerWeightsGpu
2026-09-25: Indexer weights (index sources only): bf16 wq_b [index_heads * index_hd, q_rank] and weights_proj [index_heads, dim]; on a layer that is also a kv source, bf16 wk [index_hd, hd] and f32 k_norm [index_hd].
LayerRole
2026-09-25: A layer’s compress ratio and which shared slots it writes or reads.
SharedV41
2026-09-25: What source layers publish for the layers after them: the latest kv source’s latent cache (compress_kv) and index-key cache (index_k); the latest index source’s selection topk_idxs ([queries][topk], offset past the window rows, -1 = absent); the candidate source’s mask candidates ([queries][cand_width]). compress_len is the largest group count written so far and only grows.

Enums§

AttnMat
2026-09-25: An attention projection on the device: bf16, or raw Q2_K blocks. Q3_K is refused by AttnV41::gemm. 2026-09-25: A resident [N, K] row-major projection from the GGUF path: expanded bf16, or the raw Q2_K / Q3_K blocks (WeightDtype::Q2K / Q3K) that the K-quant GEMV and MMQ kernels read.