Module compress 

2026-09-25: DeepSeek-V4.1 compressed attention (compress_ratio > 0), CPU reference against the golden: YaRN RoPE with compress_rope_theta, the two fp4 quantisers, linear_f32, and attention_any, the attention forward for any ratio. The compressor and the indexer are in the compressor and indexer submodules; the ratio-0 pieces are shared with attn.rs.

Precision: the ratio > 1 compressor is f32 until its pooled result is rounded to bf16 for the RMSNorm; the indexer’s q and k go through the e2m1 round trip (power-of-two scales, block 32) and its scores are bf16 values; the compressed KV latent goes through e2m1 with e4m3 scales (block 16). The tests compare index selections exactly and tensors at bf16 tolerance.

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

Structs

CompAttnCfg

CompAttnRun

CompressorState

CompressorWeights

2026-09-25: wkv holds f32 values at ratio > 1 and bf16 values at ratio 1; wgate is required at ratio > 1 only.

IndexerCfg

IndexerWeights

LayerAttnState

SharedRuntime

2026-09-25: What source layers publish and later layers read: the latest kv source’s whole latent cache [max_groups][hd] and index-key cache [max_groups][index_hd], the latest index selection [queries][topk] (offset past the window rows), and the candidate mask [queries][width].

Constants

FP4_BLOCK

2026-09-25: Values per power-of-two scale in the indexer’s fp4 round trip.

LATENT_BLOCK

2026-09-25: Values per e4m3 scale in the compressed KV latent’s fp4 round trip.

Functions

attention_any

2026-09-25: The attention forward for any compress_ratio, updating the layer state and the shared slots.

compressor

2026-09-25: The pre-RoPE latent [groups][hd] (bf16 values), or None while a group is still filling (prefill shorter than ratio, or a decode step that does not complete a group).

fp4_quant_e4m3_inplace

2026-09-25: The shims’ fp4_act_quant with e4m3 scales, in place: amax floored at 6 * 2^-9, scale = amax / 6 rounded through e4m3, values to e2m1, dequantised and rounded to bf16.

fp4_quant_e8m0_inplace

2026-09-25: The shims’ fp4_act_quant with e8m0 scales, in place: scale = power-of-two ceiling of amax / 6 (amax floored at 6 * 2^-126), values to e2m1, dequantised and rounded to bf16.

indexer

2026-09-25: The indexer forward. Returns ([queries][topk], topk) compressed-position indices, offset by offset, -1 where unreachable; also publishes index keys when latent and k_cache are given (a kv source).

linear_f32

2026-09-25: f32 x · w^T (the ratio > 1 compressor’s wkv / wgate are f32 weights). w is [out_dim, in_dim].

select_candidate_blocks

2026-09-25: Candidate blocks: per query, score each block-wide block of logits[queries][width] (unreachable positions at -inf) by its max, pin the block holding position compress_lens[q] - 1, keep the top topk_blocks blocks that are not all -inf. Returns the keep mask [queries][width].

to_e2m1_rne

2026-09-25: Round onto the e2m1 grid with ties to the even code, keeping the sign (as _to_e2m1_rne in ds41_ref_shims.py).

torch_cpu_topk_set

2026-09-25: The index set torch’s CPU topk(k, largest=True) returns for vals, ties included, emulated as libstdc++’s introselect (__introselect, __move_median_to_first, __unguarded_partition, __insertion_sort) over (value, index) pairs with a descending comparator, keeping the first k pairs. Among equal values (the indexer’s many -inf entries) the set is whatever that order leaves in front. Only k * 64 > n is emulated (asserted). Pinned by topk_tie_order_matches_torch_cpu.

yarn_freqs_cis

2026-09-25: The RoPE table with YaRN frequency correction, which the compressed layers build from compress_rope_theta: [seqlen][dim/2] of (cos, sin).