Module glm5next_dsa_ref 

GLM-5.3-Flash DSA + kpool indexer CPU reference (Slice 8 design artifact). GLM-5.3-Flash DSA (DeepSeek Sparse Attention) + kpool indexer CPU reference — Slice 8.

Design artifact, not a production path. Nothing here runs on GPU and no checkpoint tensor is bound. Its only job is to pin the equations in Atlas-shaped code against goldens produced by HuggingFace transformers 5.16.1 itself, before a CUDA kernel is written.

The indexer is proven before the MLA on purpose: a wrong top-k still produces perfectly plausible attention output, so an MLA test built on a broken selection passes and poisons everything downstream.

Traps this module encodes

• 🪤 k_norm is a nn.LayerNorm, not an RMSNorm — it subtracts the mean and it has a bias. Every other norm in GLM-5.3 is an RMSNorm without bias. The checkpoint carries indexer.k_norm.bias, which is the tell; a loader that binds only .weight silently drops it, and mean-subtraction is invisible in shapes.
• 🪤 The pool softmax runs over the POOL-SLOT axis, per channel — softmax(dim=2) over [pool, slot, head_dim]. It is not a softmax over head_dim and not over pools. Getting the axis wrong still yields a well-formed weighted average.
• 🪤 Pooling starts at the FIRST VALID TOKEN, not at slot 0. With left padding [P, P, A, B, C, D], pool 0 is [A, B, C, D]. Pools are formed from first_key + offset.
• 🪤 A pool is valid only if ALL kpool slots are valid. A trailing partial pool is never a pool; it is handled by the separate tail append. So a 7-token sequence has one pool, not two.
• 🪤 NoPE: qk_rope_head_dim = 0. kv_a_proj_with_mqa emits kv_lora_rank + 0, and the k_rot slice is zero-width — a no-op copy, not a padded RoPE. Do not inherit DeepSeek’s assumption that the rope section exists.
• 🔴 -1 is the invalid sentinel and the destination must be FULLY written. vLLM’s day-0 GLM DSA bug was a torch.empty top-k buffer whose tail was never written when the valid pool count fell below the budget, so uninitialised memory became “token indices”. Every function here writes all out_width entries unconditionally.

Structs

DsaDims

Indexer + MLA geometry, read from the checkpoint config — never defaulted.

Pools

The compressed k-pool candidates.

Constants

INVALID

The invalid-index sentinel. Chosen by the reference implementation, not by us.

Functions

expand_kv

Expand the compressed latent into per-head K and V.

expand_selection

Expand selected pools into raw token indices, append the visible tail, pad with INVALID.

index_scores

Per-(query, pool) index score. Mirrors the matmul → relu → head-weighted sum chain.

kept_pools

The original pool ids that survive HF’s keep = pool_valid.any(0) compaction.

layer_norm

nn.LayerNorm over the trailing d: mean-subtract, variance-normalise, then w * x + b.

linear

y = x @ w^T for x: [m, k], w: [n, k] (torch Linear layout), no bias.

mla_masked_attention

NoPE MLA over a per-query selected key set.

pool_states

Build the pools. k/gate are [seq, index_head_dim], valid is [seq].

rms_norm

RMSNorm without bias — the norm every OTHER GLM module uses.

sigmoid_f32

Sigmoid re-exported for the microtest’s convenience; the DSA path itself does not gate.

topk_pools

Select up to select_k pools per query.

topk_to_mask

Turn an index row into the boolean visibility mask the attention consumes.

visible

Which keys a query at q_pos may see: causal and not padding.