Module layer

Module layer 

Source
Expand description

Glm5NextDsaLayer โ€” the DSA block: NoPE MLA attention over indexer-selected tokens.

Decode, end to end:

hidden โ”€โ”ฌโ”€ q_a_proj โ”€ RMSNorm โ”€โ”ฌโ”€ q_absorb โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ Q (latent space)
        โ”‚                      โ””โ”€ indexer.wq_b โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ q_idx  โ”€โ”
        โ”œโ”€ indexer.wk โ”€ LayerNorm(w,b) โ”€ state.k_normed โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ค
        โ”œโ”€ compress_gate โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ state.gate โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ผโ”€ select_tokens
        โ”œโ”€ weights_proj โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ head weights โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜        โ”‚
        โ””โ”€ kv_a_proj โ”€ RMSNorm โ”€ FP8 โ”€โ”€โ”€ paged latent cache                โ”‚
                                                                           โ–ผ
                                           glm5next_dsa_mla_decode_fp8 (gather)

ยง๐Ÿชค Four silent-wrong-answer traps this file exists to hold

  • Two RMSNorm kernels differ only by a +1. rms_norm computes x * rms * (1 + w); rms_norm_vanilla computes x * rms * w. Same signature, same shapes. GLM is plain, so every norm here takes the vanilla entry point.
  • indexer.k_norm is an nn.LayerNorm with a bias, not an RMSNorm at all โ€” mean subtraction plus a bias term. It takes nllb_layernorm_bf16(x, w, b, โ€ฆ).
  • weights_proj output must already carry index_heads^-0.5. dsa_index_scores does not apply it. Folded into the weight at load โ€” see Glm5NextDsaWeights.
  • Q must be absorbed into latent space before it reaches the decode kernel. The kernel dots Q against the 512-dim latent directly, so q_absorb is q_b_proj pre-multiplied by kv_b_projโ€™s K half. A raw q_b_proj is the right shape per head (256 vs 512 is not) but the wrong space.

Structsยง

Glm5NextDsaLayer
Glm5NextDsaLayerKernels
Every kernel a DSA block launches, beyond the selection set.
Glm5NextDsaWeights
One DSA blockโ€™s weights, already sharded for this rank.
Glm5NextDsaWorkspace
Scratch reused across decode steps. Allocated once per layer.