Module attn 

DeepSeek-V4.1 attention, ratio-0 layers (Attention.forward with compress_ratio == 0): low-rank Q with RMSNorm, a single latent KV row per token (wkv + kv_norm), RoPE on the last rope_head_dim of both, the KV row quantised to fp8 in place (block 32, ue8m0 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 block-diagonal over groups, wo_b). CPU reference against the golden.

Precision follows the reference: every Linear is a bf16 GEMM (f32 accumulate, bf16 result), RMSNorm is f32 back to bf16, RoPE is a complex multiply in f32 written back into the bf16 tensor, the fp8 round trip is exact (e4m3 times a power of two is bf16-exact), and sparse_attn runs in f32 and returns bf16. Accumulation order inside a GEMM and inside the attention sums is the only thing that differs from torch, so tensors are compared at bf16 tolerance; the window indices are compared exactly.

The ring cache is the part that only a sequential run can check: prefill of 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 with one cache and check sa_kv (the ring) at each decode step.

Structs

AttnCfg

AttnRun

Every intermediate the golden captured, plus the output. o is sparse_attn’s result BEFORE the inverse rotation, which is where the golden samples it.

AttnWeights

The layer’s parameters, all bf16 values except the f32 sink. Shapes as nn.Linear stores them ([out, in]).

WindowCache

The sliding-window ring cache, [window][head_dim].

Constants

FP8_BLOCK

fp8_block_size in the reference (model.py:27).

Functions

act_quant_inplace

act_quant(x, 32, "ue8m0", e8m0, inplace=True): 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 stored back as bf16.

apply_rotary

apply_rotary_emb on the last rope_dim elements of each row_len-wide row of x, row r at position pos[r]; adjacent pairs are complex numbers, multiplied in f32, written back bf16.

attention

Attention.forward for a ratio-0 layer. x: [seqlen][dim] bf16 values; fc from freqs_cis over max_seq_len. Mutates cache exactly as the reference’s ring buffer.

freqs_cis

precompute_freqs_cis(dim, seqlen, original_seq_len=0, base, ...), the no-YaRN form the ratio-0 layers use: [seqlen][dim/2] of (cos, sin).

pow2_ceil

2 ** ceil(log2(v)) for positive normal f32 via the IEEE fields, as fast_log2_ceil / fast_pow2 in kernel.py:22-35 (and _pow2_ceil_log2 in the shims).

sparse_attn

sparse_attn (kernel.py:310-405, shim form): 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].

window_topk_idxs

get_window_topk_idxs: [queries][cols] of slot indices, -1 for a slot holding nothing. Prefill: one causal window per query into the chunk itself. Decode: the whole ring, oldest first.