Module compress 

DeepSeek-V4.1 compressed attention (compress_ratio > 0): the Compressor (softmax pooling of ratio tokens into one latent, or a plain projection at ratio 1), the Indexer (fp4 query heads against one shared fp4 key per compressed position, rectified scores combined by weights_proj), select_candidate_blocks (level one of the two-level top-k), the SharedAttentionRuntime slots layers hand down the stack, YaRN RoPE with compress_rope_theta, the two fp4 quantisers, and the full Attention.forward for any ratio. CPU reference against the golden; the ratio-0 pieces are shared with attn.rs.

Precision follows the reference: the ratio-2 compressor is f32 end to end until its RMSNorm casts to bf16; the indexer’s q and k go through the fp4 e2m1 round trip (e8m0 scales, block 32) and its scores are bf16 tensors; the compressed KV latent goes through e2m1 with e4m3 scales (block 16). Index selections are compared exactly, tensors at bf16 tolerance.

Structs

CompAttnCfg

CompAttnRun

CompressorState

CompressorWeights

IndexerCfg

IndexerWeights

LayerAttnState

SharedRuntime

SharedAttentionRuntime: what source layers publish and later layers read.

Constants

FP4_BLOCK

fp4_block_size (model.py:28).

LATENT_BLOCK

Block size the compressed KV latent is quantised with (fp4_act_quant(latent, 16, ...)).

Functions

attention_any

Attention.forward for any compress_ratio, mutating the layer state and the shared slots.

compressor

Compressor.forward: the pre-RoPE latent [groups][hd] (bf16 values), or None while a group is still filling up in decode.

fp4_quant_e4m3_inplace

fp4_act_quant(x, 16, inplace=True, scale_dtype=e4m3): amax floored at 6 * 2^-9, scale = amax / 6 rounded through e4m3, values to e2m1, dequantised back to bf16.

fp4_quant_e8m0_inplace

fp4_act_quant(x, block, inplace=True) with e8m0 scales: scale = pow2 ceiling of amax / 6 (amax floored at 6 * 2^-126), values to e2m1, dequantised back to bf16.

indexer

Indexer.forward. Returns [queries][topk] compressed-position indices, offset by offset, -1 where unreachable; also publishes index keys when this layer owns them.

linear_f32

f32 F.linear (the ratio-2 compressor’s wkv / wgate are fp32 weights on an fp32 input).

select_candidate_blocks

select_candidate_blocks: logits[queries][width] with unreachable positions at -inf; compress_lens[q]. Returns the keep mask [queries][width].

to_e2m1_rne

Round onto the e2m1 grid with ties to the even code, keeping the sign (_to_e2m1_rne).

torch_cpu_topk_set

The index SET torch’s CPU topk(k, largest=True) returns for vals, ties included. For the row lengths here (k * 64 > n) torch runs std::nth_element on (value, index) pairs with a descending comparator and returns the first k pairs, so among equal values (the indexer’s many -inf entries) the picks are whatever libstdc++’s introselect leaves in front. That order is deterministic and this is a line-for-line emulation of it (__introselect, __move_median_to_first, __unguarded_partition, __insertion_sort). Pinned by topk_tie_order_matches_torch_cpu against this box’s torch.

yarn_freqs_cis

precompute_freqs_cis WITH YaRN (original_seq_len > 0), as the compress layers build it from compress_rope_theta. [seqlen][dim/2] of (cos, sin).