Module glm5next_kda

Module glm5next_kda 

Source
Expand description

GLM-5.3-Flash KDA attention block — the reusable production component.

One Glm5NextKdaLayer is a fully bound KDA self_attn block. Any of the checkpoint’s 34 KDA layers instantiates through the same path; the family is structurally uniform (one distinct name/shape/dtype signature across all 34 — see binding).

q|k|v_proj -> pack -> conv1d + SiLU -> L2(q,k only) -> kda_gate / sigmoid(b_proj)
           -> kda_chunk (prefill) | kda_recurrent (decode)
           -> sigmoid-gated RMSNorm(o_norm, g_b(g_a(h))) -> o_proj

Scope: the attention block only. No DSA/MLA, no MoE/dense FFN, no mHC hyper-connection, no scheduler or cache integration. The model is not loadable on this alone.

§Facts this component encodes (proven, do not re-derive)

  • Nothing in a KDA block is quantised. All 15 self_attn tensors are BF16 except A_log and dt_bias, which are F32 — verified across all 34 blocks, 0 artefacts. There is no NVFP4 dequant and no NVFP4 GEMM anywhere on this path, so a real-checkpoint KDA oracle is the production numerics.
  • 🪤 The checkpoint SPLITS the conv; HF FUSES it. HF holds one depthwise nn.Conv1d(conv_dim); the checkpoint stores q_conv1d/k_conv1d/v_conv1d, each [qkv, 1, kernel]. Binding is concat([q, k, v], dim=0) in that order — the same order as mixed_qkv = cat([q_proj, k_proj, v_proj]). Reordering is silent.
  • 🪤 squeeze(1) is a shape-only fix. [dim, 1, ks] and [dim, ks] have identical row-major bytes, so nothing moves — but a loader trusting shape.len() == 2 rejects the tensor outright. binding asserts rank 3 and squeezes exactly once.
  • 🪤 o_norm is ADAPT, not REUSE. Every Atlas gated RMSNorm applies SiLU to the gate (kernels/gb10/common/rms_norm.cu); GLM’s Glm5NextTextRMSNormGated sets activation = "sigmoid". Shapes, dtypes and launch geometry all agree, which is why it was first mis-classified. Hence kda_o_norm_gated_* in kernels/gb10/common/kda_layer_ops.cu.
  • 🪤 dense_gemm_bf16 writes C[row * N + col] — its output row stride is N and there is no caller-supplied output stride, so the three q/k/v projections cannot be aimed at offsets inside one [T, 3*qkv] buffer. They would overwrite each other for T > 1, while being silently correct at T = 1. Hence the separate parts buffer plus kda_pack_qkv_bf16.
  • Conv state widths differ. HF keeps kernel - 1 slots, Atlas keeps kernel and shifts left before convolving, so HF[0..k-1] == Atlas[1..k] and Atlas slot 0 is a don’t-care.
  • Decode conv fuses L2; prefill does not and needs a separate l2_norm_bf16 over q|k. q/k are normalised exactly once; V never.
  • The recurrent state is FP32 by reference semantics, not Atlas policy — HF stores it via .to(torch.float32) and vLLM’s kda_state_dtype hardcodes fp32.

Modules§

binding
Typed weight binding for the GLM-5.3-Flash KDA attention family.
tp
GLM-5.3 KDA tensor-parallel shard plan — head-parallel, one all-reduce.
tp_bind
Applying the TP plan: the shard copies, as an upstream adapter so the binder is untouched. Applying super::tp::KdaTpPlan — the shard COPIES, at last.

Structs§

Glm5NextKdaConfig
Geometry and the config values that MUST be read from the checkpoint.
Glm5NextKdaKernels
Glm5NextKdaLayer
One bound KDA attention block.
Glm5NextKdaWeights
One KDA block’s device weights. Torch Linear layout [out, in], BF16, except the two F32 gate parameters. There is no Z tensor — the output gate is low-rank g_a/g_b.
Glm5NextKdaWorkspace
Scratch owned by the forward path. Sized once for max_tokens and reused across layers — the whole KDA family shares one workspace because every block has identical geometry.
KdaSeqState
The per-sequence state a KDA layer carries. Both buffers are read-modify-write.

Constants§

SMEM_CEILING
Dynamic shared memory available with no cuFuncSetAttribute opt-in in AtlasCudaBackend. GB10 reports sharedMemPerBlockOptin = 101376, real but unreachable from Atlas today.