Module glm5next_dsa_ref 

2026-09-25: CPU reference for the GLM-5.3-Flash DSA kpool indexer and NoPE MLA.

Owner: model-arch (GLM-5.3-Flash DSA reference). Invariants:

• Nothing here runs on a GPU or reads a checkpoint; production code does not call it.
• Every index row expand_selection returns is written for its full out_width, with INVALID in every slot that holds no token.

examples/dsa_indexer_microtest.rs checks it against goldens from gen_dsa_indexer_golden.py and gen_dsa_mla_golden.py, which run the transformers glm5_next indexer and attention.

• k_norm is a LayerNorm: it subtracts the mean and adds a bias (indexer.k_norm.bias). The DSA block’s other norms (q_a_layernorm, kv_a_layernorm) have no bias.
• The pool softmax runs over the slot axis, separately for each head_dim channel.
• 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].
• A pool is valid only if all kpool slots are valid. A trailing partial pool is not a pool; the tail append covers it, so a 7-token sequence has one pool.
• NoPE: qk_rope_head_dim = 0, so kv_a_proj_with_mqa emits kv_lora_rank values and there is no rope section.

Structs

DsaDims

2026-09-25: Indexer and MLA geometry. No field has a default.

Pools

2026-09-25: The compressed k-pool candidates.

Constants

INVALID

2026-09-25: The invalid-index sentinel, the -1 the golden generator writes.

Functions

expand_kv

2026-09-25: Expand the compressed latent into per-head K and V. NoPE only: it asserts qk_rope_head_dim == 0, so K is exactly qk_nope_head_dim wide.

expand_selection

2026-09-25: Expand selected pools into raw token indices, append the visible tail, and pad with INVALID. Every row is written for all out_width entries: the output starts filled with INVALID, and a padded query row stays that way.

index_scores

2026-09-25: Per-(query, pool) index score: sum_h weights[h] * relu(scale * q_h . key).

kept_pools

2026-09-25: The original ids of the pools pool_states keeps. Computed from padding and sequence length alone.

layer_norm

2026-09-25: LayerNorm over the trailing d: subtract the mean, divide by the standard deviation, then w * x + b.

linear

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

mla_masked_attention

2026-09-25: NoPE MLA over a per-query selected key set.

pool_states

2026-09-25: Build the pools. k/gate are [seq, index_head_dim], valid is [seq]. The indexer golden records the output of transformers’ get_pooled_states for comparison.

rms_norm

2026-09-25: RMSNorm without bias, the contrast to layer_norm in the tests.

sigmoid_f32

2026-09-25: Public sigmoid; nothing in the DSA reference calls it.

topk_pools

2026-09-25: Select up to select_k pools per query, ordered by higher score, then smaller pool index, so ties resolve the same way on every run. Invalid candidates score f32::MIN.

topk_to_mask

2026-09-25: Turn an index row into the boolean visibility mask the attention reads. Duplicate indices collapse, so a repeated token is attended once; out-of-range and INVALID entries are dropped.

visible

2026-09-25: Which keys a query at q_pos may see: causal and not padding.