Module layer

Module layer 

Source
Expand description

2026-09-25: Glm5NextDsaLayer, one GLM-5.3 DSA block: the projections, the indexer cache write, token selection and the NoPE MLA gather-attend over the selected tokens.

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

  • decode_k fails before its first launch when the indexer cache is behind the sequence (len < seq_len) or k is 0 or above the workspace’s max_rows. An indexer cache ahead of the sequence is rewound to seq_len first.
  • indexer_forward checks that the indexer cache has room (ensure_room) before it writes.

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)

§Kernel choices the shapes do not check

  • The norms use rms_norm_vanilla, x * rms * w. The rms_norm kernel (module norm) computes x * rms * (1 + w) with the same signature.
  • indexer.k_norm is a LayerNorm with a bias: nllb_layernorm_bf16(x, w, b, …).
  • weights_proj carries index_heads^-0.5 from load (Glm5NextDsaWeights), because dsa_index_scores does not apply it.
  • Q reaches the decode kernel as q_absorb: q_b_proj multiplied through kv_b_proj’s K half, so it is in the kv_lora_rank-wide latent space the kernel dots against. The raw q_b_proj is qk_head_dim (256) per head.

Structs§

Glm5NextDsaLayer
Glm5NextDsaLayerKernels
2026-09-25: The projection, norm and latent-write kernels of a DSA block. The selection kernels are in Glm5NextDsaKernels, the decode kernel in attend.
Glm5NextDsaWeights
2026-09-25: One DSA block’s weights on the device, already sharded for this rank.
Glm5NextDsaWorkspace
2026-09-25: Scratch reused across decode_k calls; each layer owns one.