Module attn 

2026-09-25: DeepSeek-V4.1 attention for ratio-0 layers (the reference’s Attention.forward at compress_ratio == 0), CPU reference against the golden: low-rank Q with RMSNorm, one latent KV row per token (wkv + kv_norm), RoPE on the last rope_dim of both, the KV row rounded through fp8 in place (block 32, power-of-two scale), a sliding-window ring cache, sparse_attn with the attention sink, the inverse rotation on the output, then the grouped low-rank output projection (wo_a per group, then wo_b). Also the RoPE, window and fp8 helpers the compressed layers and attn_v41 share.

Precision: every projection is linear_bf16 (f32 accumulate, bf16 result), RMSNorm is f32 back to bf16, RoPE multiplies in f32 and writes bf16, the fp8 round trip is bf16-exact (e4m3 times a power of two), and sparse_attn runs in f32 and returns bf16. The tests compare tensors at bf16 tolerance and the window indices exactly.

The ring cache needs a sequential run to check: prefill of the fixture’s 12 tokens into 8 slots seeds slots[4..8] = tokens 4..8 and slots[0..4] = tokens 8..12; decode step 12 writes slot 4 and reads the whole ring oldest-first; step 13 writes slot 5. The tests run the three regimes in that order through one cache and check sa_kv (the ring) at each decode step.

Owner: model-arch, DeepSeek-V4.1 reference. Invariants: none beyond the types.

Structs

AttnCfg

AttnRun

2026-09-25: The intermediates the tests compare, plus the output. o is sparse_attn’s result before the inverse rotation, the golden’s sa_o.

AttnWeights

2026-09-25: The layer’s parameters, all bf16 values except the f32 sink; projections are [out, in].

WindowCache

2026-09-25: The sliding-window ring cache, [window][head_dim].

Constants

FP8_BLOCK

2026-09-25: Values per fp8 scale block.

Functions

act_quant_inplace

2026-09-25: The shims’ act_quant with a power-of-two scale, in place: per block of FP8_BLOCK, scale = the power-of-two ceiling of amax / 448 (amax floored at 1e-4), values to e4m3 (RNE, clamped), then dequantised and rounded to bf16.

apply_rotary

2026-09-25: Rotate the last rope_dim elements of each row_len-wide row of x, row r at position pos[r] (negated angle when inverse); adjacent pairs are complex numbers, multiplied in f32, written back bf16.

attention

2026-09-25: The ratio-0 attention forward. x: [seqlen][dim] bf16 values; fc from freqs_cis over the maximum sequence length. Writes the new kv rows into cache’s ring.

freqs_cis

2026-09-25: The RoPE table without YaRN, which the ratio-0 layers use: [seqlen][dim/2] of (cos, sin).

pow2_ceil

2026-09-25: 2 ** ceil(log2(v)) for positive normal f32 via the IEEE fields, as _pow2_ceil_log2 in ds41_ref_shims.py.

sparse_attn

2026-09-25: sparse_attn as in ds41_ref_shims.py: per query and head, softmax over the gathered rows plus the sink (denominator only), all in f32; output bf16. q: [s][h][d], kv: [rows][d], idx: [s][topk] with -1 skipped.

window_topk_idxs

2026-09-25: Window selection ([queries][cols], cols) of slot indices, -1 for a slot holding nothing. At start_pos == 0: one causal window per query into the chunk itself. Otherwise: the whole ring, oldest first, for one query.