Module glm5next_kda_ref 

2026-09-25: CPU reference for GLM-5.3-Flash KDA (Kimi Delta Attention).

Owner: model-arch (GLM-5.3-Flash KDA reference). Invariants:

• Nothing here runs on a GPU or reads a checkpoint; production code does not call it.
• lower_bound and rms_eps are arguments; the reference has no built-in value for either.

It is checked against kda_golden.json, which gen_kda_golden.py produces by calling the transformers glm5_next functions (Glm5NextTextForgetGate, Glm5NextTextRMSNormGated, l2norm, recurrent_kimi_delta_attention, chunk_kimi_delta_attention).

KDA differs from the Qwen GDN path (compute_gdn_gates in kernels/gb10/common/ssm_preprocess.cu) in four ways:

1. The decay is per (head, key channel), [T, H, head_dim]; GDN writes one value per head.
2. The decay comes from a low-rank projection, f_a: hidden -> head_dim, f_b: head_dim -> H*head_dim; GDN reads it from the fused BA projection.
3. The gate is bounded, lower_bound * sigmoid(exp(A_log) * (g + dt_bias)); GDN computes exp(-exp(A_log) * softplus(a + dt_bias)).
4. The output gate is low-rank g_a/g_b; GDN takes a full-rank Z from its fused QKVZ projection, and a KDA checkpoint has no Z tensor.

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

2026-09-25: KDA geometry. The low-rank width of both f_a/f_b and g_a/g_b is head_dim.

KdaWeights

2026-09-25: Weights for one KDA layer, reference only. Torch Linear layout [out, in] throughout. There is no Z tensor: the output gate is g_a/g_b and the decay source is f_a/f_b.

Functions

bounded_gate

2026-09-25: The bounded GLM forget gate:

kda_chunked

2026-09-25: Prefill formulation: the chunked (WY-style) delta rule with a per-channel decay mask, checked against the golden’s chunk_kimi_delta_attention output. tokens % chunk != 0 is handled by zero-padding.

kda_chunked_prenorm

2026-09-25: Same chunked formulation, but q/k are already L2-normalised, as the GPU prefill’s l2_norm_bf16 leaves them (in bf16). Separate from kda_chunked for the reason given on kda_recurrent_prenorm.

kda_recurrent

2026-09-25: Decode formulation: one token at a time, carrying state.

kda_recurrent_prenorm

2026-09-25: Same recurrence, but q/k are already L2-normalised, the input the kda_recurrent GPU kernel gets from causal_conv1d_update_l2norm.

kda_reference_layer

2026-09-25: Reference for one KDA layer from post-conv q/k/v to the o_proj output. The short conv is not included; the golden’s inputs are post-conv.

l2norm_rows

2026-09-25: x / sqrt(sum(x^2) + eps) over the trailing dimension, not x / max(norm, eps); the two differ for small-norm rows. The GPU l2_norm_bf16 uses the same form.

linear

2026-09-25: y = x @ w^T for x: [m, k], w: [n, k] (torch Linear weight layout), no bias.

rms_norm_gated

2026-09-25: Gated RMSNorm over the trailing d, FP32, with a sigmoid gate. eps is the caller’s rms_norm_eps.

unbounded_gdn_gate

2026-09-25: The unbounded Qwen GDN gate, used by the tests to show the two laws diverge.