Module glm5next_dsa

Module glm5next_dsa 

Source
Expand description

GLM-5.3-Flash KDA integrated layer (Slice 6 – one layer, no scheduler/cache wiring). GLM-5.3-Flash DSA (DeepSeek Sparse Attention) production surface.

Scoped to LibertAIDAI/GLM-5.3-Flash-NVFP4@9e0d74e3.

The CUDA kernels already exist and are numerically proven against HF 5.16.1 on real weights — kernels/gb10/common/dsa_indexer.cu, gated by examples/dsa_indexer_microtest.rs (GATE 4 kpool indexer, GATE 5 NoPE MLA over the selected tokens). What was missing, and is what this module adds, is the production surface: kernel resolution and geometry that a real layer can bind, rather than an example wiring pointers by hand.

The CPU reference in crate::layers::glm5next_dsa_ref stays the source of truth for the equations. Nothing here re-derives them.

§Shape of the pipeline

k,gate,valid,ape -> kpool_compress -> pool keys/indices/valid
                 -> index_scores  -> [Q, P] scores + candidate validity
                 -> topk_pools    -> [Q, select_k] pool ids
                 -> expand_selection -> [Q, out_width] token ids (-1 = invalid)
                 -> NoPE MLA restricted to those tokens

§🪤 Traps carried from Slice 8 (do not re-derive)

  • indexer.k_norm is a nn.LayerNorm — mean-subtracting, with a bias — not an RMSNorm. indexer.k_norm.bias existing in the checkpoint is the only tell; every other norm in GLM-5.3 is a bias-free RMSNorm.
  • The pool softmax runs over the pool-slot axis, per channel, not over head_dim and not over pools.
  • Pooling starts at the first valid token, so left padding is skipped rather than pooled.
  • A pool counts only if every kpool slot is valid — a trailing partial pool is not a pool.
  • NoPE: qk_rope_head_dim == 0, so the k_rot slice is zero-width. See the rope > 0 guards in qwen3_attention — a NoPE checkpoint carries no wkv_a_rope and leaves rope_theta unset.
  • The -1 sentinel destination must be fully written. vLLM’s day-0 GLM bug was a torch.empty top-k buffer whose tail was never written, so uninitialised memory became “token indices”.

Modules§

attend
GLM-5.3 DSA attention — the launcher for the selected-index NoPE MLA paged decode.
binding
GLM-5.3 DSA checkpoint binding: the 14 self_attn tensors of a DSA block, their expected dtype and shape, and a verifier that fails loudly.
build
Loading one DSA block: TP sharding plus the three load-time transforms the runtime cannot do per token.
layer
Glm5NextDsaLayer — the DSA block: NoPE MLA attention over indexer-selected tokens.
select
GLM-5.3 DSA token selection — the production launcher for the indexer pipeline.
state
Per-sequence DSA indexer state — the cache the selector reads every decode step.
tp
GLM-5.3 DSA tensor-parallel shard plan — MLA heads sharded, indexer replicated.

Structs§

Glm5NextDsaConfig
DSA geometry for one layer, read from the checkpoint config — never defaulted.
Glm5NextDsaKernels
Every kernel the DSA path launches.

Constants§

DSA_MODULE
Module name the DSA kernels resolve from. Unlisted .cu files take their file stem as the module name, so kernels/gb10/common/dsa_indexer.cu is dsa_indexer.
KERNEL_KV_LORA_DIM
#define KV_LORA_DIM in kernels/gb10/common/mla_paged_decode{,_fp8}.cu.
KERNEL_MAX_KPOOL
float lg[8] in dsa_kpool_compress — the most pool slots the compression kernel can hold. The kernel loops s < KP && s < 8, so a larger index_kpool is silently truncated rather than rejected. Mirrored here so config validation refuses it instead.
LAYERNORM_MODULE
Module carrying the bias-bearing BF16 LayerNorm the indexer’s k_norm needs. Lives in common/, so every target merges it; the name is the .cu file stem.
MASKED_ATTN_MAX_KEYS
🔴 dsa_mla_masked_attn is an oracle, not a serve path. Resolved from source 2026-08-27; do not re-derive.