Module glm5next_kda_ref

Module glm5next_kda_ref 

Source
Expand description

GLM-5.3-Flash KDA CPU reference (Slice 2 design artifact – not a production forward path). GLM-5.3-Flash KDA (Kimi Delta Attention) CPU reference — Slice 2B.

Design artifact, not a production path. Nothing here runs on GPU, nothing here is wired into a forward pass, and no checkpoint tensor is bound. Its only job is to prove the GLM KDA equations in Atlas-shaped code against golden vectors produced by HuggingFace itself, before a single CUDA kernel is written.

§Why this exists

GLM KDA is not a thin parameterization of Atlas’s Qwen GDN (layers::qwen3_ssm). Four structural items are genuinely new:

  1. Decay is per (head, key-channel), not scalar-per-head. Atlas’s compute_gdn_gates writes gate_out[num_tokens, num_v_heads]; KDA needs [T, H, head_dim].
  2. The decay source is a low-rank projection f_a: hidden -> head_dim, f_b: head_dim -> H*head_dim. Atlas derives its scalar from the fused BA projection.
  3. The gate law is bounded: lower_bound * sigmoid(exp(A_log) * (g + dt_bias)), versus Atlas’s unbounded exp(-exp(A_log) * softplus(a + dt_bias)).
  4. The output gate is low-rank g_a/g_b. Atlas takes a full-rank Z out of its fused QKVZ projection; the KDA checkpoint has no Z tensor at all.

See docs/glm5next/KDA-VS-QWEN-GDN.md for the full REUSE/ADAPT/NEW table.

§Provenance of the goldens

kda_golden.json is generated by gen_kda_golden.py (kept beside it), which calls the real transformers 5.16.1 glm5_next module — Glm5NextTextForgetGate, Glm5NextTextRMSNormGated, l2norm, recurrent_kimi_delta_attention, chunk_kimi_delta_attention. No equation is re-derived on the Python side either.

transformers 5.16.0 contains no glm5_next at all even though the checkpoint declares that version; 5.16.1 is the first source-bearing release.

§Config values are read, never defaulted

gate_lower_bound, rms_norm_eps and hidden_act must all come from the checkpoint config. vLLM happens to agree with HF on this checkpoint only by coincidence: it looks up the legacy key lower_bound while the checkpoint stores gate_lower_bound, and falls back to a default of -5.0 that happens to match. Atlas must not inherit that.

§Layout conventions

  • q/k/v/gate: [T, H, D], row-major, index ((t * H) + h) * D + d.
  • beta: [T, H], index t * H + h.
  • dt_bias: [H * D] — per channel.
  • a_log: [H] — per head.
  • state: [H, D_k, D_v], index (h * D_k + kd) * D_v + vd.
  • Weights follow the torch Linear convention [out, in], row-major.

Structs§

KdaDims
KDA geometry. Production is hidden 4096 / heads 64 / head_dim 128; the low-rank width for both f_a/f_b and g_a/g_b equals head_dim.
KdaWeights
Weights for one KDA layer, reference-only. Torch Linear layout [out, in] throughout.

Functions§

bounded_gate
The bounded GLM forget gate, item 3 of the four NEW structural items.
kda_chunked
Prefill formulation: chunked (WY-style) delta rule, mirroring HF’s chunk_kimi_delta_attention. Included at reference level because the per-channel decay mask is exactly the part that a scalar-decay GDN kernel cannot express — proving it here is cheaper than proving it in CUDA.
kda_chunked_prenorm
Same chunked formulation, but q/k are already L2-normalised — Atlas’s contract, where the conv path (fused on decode, l2_norm_bf16 on prefill) has already normalised them.
kda_recurrent
Decode formulation: one token at a time, carrying state.
kda_recurrent_prenorm
Same recurrence, but q/k are already L2-normalised.
kda_reference_layer
End-to-end reference for one KDA layer, from post-conv q/k/v to o_proj output.
l2norm_rows
x / sqrt(sum(x^2) + eps) over the trailing dimension.
linear
y = x @ w^T for x: [m, k], w: [n, k] (torch Linear weight layout), no bias.
rms_norm_gated
Gated RMSNorm over the trailing d, strict FP32, sigmoid gate.
unbounded_gdn_gate
Unbounded Qwen-GDN gate, kept only so the microtest can show the two laws diverge.