Module glm5next_dsa

Module glm5next_dsa 

Source
Expand description

2026-09-25: GLM-5.3-Flash DSA (DeepSeek Sparse Attention): the kernel handles, the per-layer geometry and its validation.

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

  • A config returned by Glm5NextDsaConfig::from_config has passed validate: index_topk is a nonzero multiple of index_kpool, index_kpool is at most KERNEL_MAX_KPOOL, kv_lora_rank equals KERNEL_KV_LORA_DIM, and local_heads is nonzero.

The checkpoint is the one kernels/gb10/glm-5.3-flash/MODEL.toml names (LibertAIDAI/GLM-5.3-Flash-NVFP4, revision 9e0d74e3). The selection kernels are in kernels/gb10/common/dsa_indexer.cu, and examples/dsa_indexer_microtest.rs runs them on a GPU. crate::glm5next_dsa_ref is the CPU reference for the equations.

§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

§Checkpoint and kernel facts the code depends on

  • indexer.k_norm is a LayerNorm with a bias. indexer.k_norm.bias is the only norm bias among the checkpoint’s text tensors (crates/model-engine/tests/fixtures/glm53-nvfp4-9e0d74e3-structural.txt).
  • The pool softmax runs over the pool-slot axis, per channel (dsa_kpool_compress).
  • Pooling starts at first_key. A pool counts only if every one of its index_kpool slots is valid, so a trailing partial pool is not a pool.
  • NoPE: qk_rope_head_dim is 0, so there is no rope section and no wkv_a_rope tensor.
  • dsa_expand_selection writes -1 over the whole index row before it stores any index, so every slot is written.

Modules§

attend
2026-09-25: GLM-5.3 DSA attention: the launcher for the selected-index NoPE MLA paged decode.
aux_state
2026-09-25: Snapshot and restore of the DSA indexer cache, the part of a GLM-5.3 sequence that a KV-only prefix-cache hit cannot rebuild.
binding
2026-09-25: The tensor contract of one GLM-5.3 DSA block: the 14 self_attn tensors dsa_tensor_specs lists, their dtype and unsharded shape, and a verifier.
build
2026-09-25: Loading one DSA block for this rank: tensor-parallel sharding plus four load-time weight transforms.
layer
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.
select
2026-09-25: GLM-5.3 DSA token selection: launch geometry, scratch, and the launcher for the selection kernels.
state
2026-09-25: Per-sequence DSA indexer cache: the k_normed, gate and valid rows the selector reads.
tp
2026-09-25: GLM-5.3-Flash DSA tensor-parallel shard plan: MLA heads sharded, the indexer replicated.

Structs§

Glm5NextDsaConfig
2026-09-25: DSA geometry for one layer, from ModelConfig.
Glm5NextDsaKernels
2026-09-25: The DSA selection kernels, the indexer LayerNorm and the two oracle kernels.

Constants§

DSA_MODULE
2026-09-25: Module name the DSA selection kernels resolve from. dsa_indexer.cu has no entry in kernels/gb10/common/KERNEL.toml [modules], so its module is its file stem.
KERNEL_KV_LORA_DIM
2026-09-25: #define GLM_KV_LORA_DIM in kernels/gb10/glm-5.3-flash/nvfp4/glm5next_dsa_mla_decode.cu, the latent width the decode kernel is compiled for. Glm5NextDsaConfig::validate and attend::decode_attention refuse a kv_lora_rank that differs.
KERNEL_MAX_KPOOL
2026-09-25: float lg[8] in dsa_kpool_compress, the most pool slots the compression kernel holds. The kernel loops s < KP && s < 8, so it would drop slots past 8 without an error; config validation refuses a larger index_kpool instead.
LAYERNORM_MODULE
2026-09-25: Module of nllb_layernorm_bf16, the BF16 LayerNorm with a bias that the indexer’s k_norm needs: kernels/gb10/common/nllb_encoder.cu, named by its file stem.
MASKED_ATTN_MAX_KEYS
2026-09-25: The most keys dsa_mla_masked_attn takes: it stages the [S] f32 score row in dynamic shared memory, and 4 · S <= 49,152 B gives 12,288.