Module attn_v41 

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

V4.1 has three kinds of attention layer, told apart by compress_ratio: ratio 0 attends over a sliding window of its own fp8 latent rows; ratio 1 and ratio 2 attend over that window PLUS a set of compressed positions chosen by an indexer from a latent cache that only the four kv_source_layer_ids produce (the SharedAttentionRuntime: four layers write, forty read). The index selection itself is computed by eight index_source_layer_ids and reused by the layers between them; layer candidate_source_layer also publishes a block-level candidate mask the later index sources respect.

This module is the production form of deepseek_v41_ref::compress::attention_any, stage for stage, and is held to it layer by layer in attn_v41_tests.rs:

• GEMMs on common/dense_gemm_bf16 (f32 accumulate, bf16 out = linear_bf16);
• RMSNorm, RoPE, the fp8/fp4 quantisers, the compressor pooling, the indexer scores, sparse attention with the sink and the grouped output projection on kernels/gb10/deepseek-v4-flash/nvfp4/attn_v41.cu;
• the two top-k selections (candidate blocks, index positions) on the CPU through the reference’s torch_cpu_topk_set, because the reference returns torch’s CPU tie order among equal scores and a kernel that picks any other set diverges from it on exactly those queries. The scores are bf16-valued and few ([tokens, width]), so this costs a download, not a kernel.

Layer state (window ring, compressor partial group, the two caches) and the shared slots live on the device; the shared index selection and candidate mask are host vectors, as the reference keeps them.

Split into submodules:

• init: AttnV41::new (kernels, RoPE tables, workspaces) and free
• primitives: the GEMM/GEMV, RMSNorm, RoPE and quantiser launch wrappers
• sources: the kv-source compressor and the index-source indexer
• forward: one layer’s attention forward

Re-exports

pub use crate::layers::ops::ResidentMat as AttnMat;

Structs

AttnV41

The attention runtime: kernels, RoPE tables, and workspaces for up to max_tokens positions per call.

AttnV41Cfg

Model-wide attention geometry, from ModelConfig.

AttnV41LayerState

One layer’s device state across prefill and decode.

AttnV41LayerWeights

One layer’s resident attention weights. bf16 [out, in] unless noted.

AttnV41Run

The intermediates of one layer’s forward, for oracles and for the caller.

CompressorWeightsGpu

Compressor weights (kv sources only). At ratio > 1 kv and gate are f32 [hd, dim] as the checkpoint stores them; at ratio 1 kv is bf16.

IndexerWeightsGpu

Indexer weights (index sources only); wk / k_norm on kv sources.

LayerRole

What one layer is, in the shared runtime.

SharedV41

The SharedAttentionRuntime: what source layers publish for the layers below them.