Skip to content

[Feature] Add GLM-5.3-Flash F3: Kimi Delta Attention (KDA) - #2106

Open
jayhenry wants to merge 1 commit into
feat/glm53flash-materialize-full-f0from
feat/glm53flash-f3-kda
Open

jayhenry wants to merge 1 commit into
feat/glm53flash-materialize-full-f0from
feat/glm53flash-f3-kda

Conversation

@jayhenry

@jayhenry jayhenry commented Sep 23, 2026

Copy link
Copy Markdown
Collaborator

Stack (bottom to top):

  1. [Feature] Add GLM-5.3-Flash F0: 25B cropped reference checkpoint builder #2105 feat/glm53flash-materialize-full-f0main
  2. [Feature] Add GLM-5.3-Flash F3: Kimi Delta Attention (KDA) #2106 feat/glm53flash-f3-kdafeat/glm53flash-materialize-full-f0 ← you are here
  3. [Feature] Add GLM-5.3-Flash F4: mHC four-stream residual #2107 feat/glm53flash-f4-mhcfeat/glm53flash-f3-kda
  4. [Feature] Add GLM-5.3-Flash F5: NoPE DSA + KPool indexer + clamped SwiGLU #2108 feat/glm53flash-f5-nope-dsafeat/glm53flash-f4-mhc
  5. [Feature] Add GLM-5.3-Flash F1: VL data preprocessing pipeline #2109 feat/glm53flash-f1-vl-datafeat/glm53flash-f5-nope-dsa
  6. [Feature] Add GLM-5.3-Flash F2: vision tower + projector (eager) #2110 feat/glm53flash-f2-vision-towerfeat/glm53flash-f1-vl-data
  7. [Feature] Add GLM-5.3-Flash F6 core: text model + MTP + compose model #2111 feat/glm53flash-f6-text-moefeat/glm53flash-f2-vision-tower

Base is #2105's branch (layer 1). Review only this PR's own diff.


Summary

Stack layer 2/7 of GLM-5.3-Flash support (base: layer 1, F0 + materialize_full).

Adds xtuner/v1/module/attention/kda.py (KDAConfig / KimiDeltaAttention) per doc/xtuner_glm5p3flash_design.md F3, with module structure and parameter names matching the published checkpoint layout (separate q/k/v_conv1d, low-rank f_a_proj->f_b_proj forget gate with fp32 A_log/dt_bias, low-rank g_a_proj->g_b_proj output gate).

Root cause of a real correctness bug found while implementing: the installed fla 0.4.2's chunk_kda/fused_recurrent_kda call convention differs from the xtuner-ncp-k3 reference — it doesn't accept A_log/dt_bias/use_beta_sigmoid_in_kernel, and silently swallows them into **kwargs without raising, producing a KDA that trains but computes the wrong gate. Fixed by computing the gate externally via fla.ops.kda.gate.fused_kda_gate before calling the kernel, matching HF's own Glm5NextTextLinearAttention reference forward. A regression test pins these three kwargs' absence from chunk_kda's signature so this can't silently regress.

Also relaxes TransformerConfig.linear_attention to GatedDeltaNetConfig | KDAConfig | None (mypy-strict plumbing only, no behavior change for existing GatedDeltaNet layers).

Test Plan

tests/model/test_glm53_kda.py (5 tests, single/2-GPU): gate parity against fla's naive reference, module-level parity against HF's Glm5NextTextLinearAttention (weights bridged per the published<->HF key mapping), packed multi-document isolation, and Ulysses SP parity.

Adds xtuner/v1/module/attention/kda.py (KDAConfig / KimiDeltaAttention)
following doc/xtuner_glm5p3flash_design.md F3, with module structure and
parameter names matching the published checkpoint layout (separate
q/k/v_conv1d, low-rank f_a_proj->f_b_proj forget gate with fp32
A_log/dt_bias, low-rank g_a_proj->g_b_proj output gate).

Critically, the chunk_kda/fused_recurrent_kda call convention differs from
the ~/github/xtuner-ncp-k3 reference: the installed fla 0.4.2's chunk_kda
does not accept A_log/dt_bias/use_beta_sigmoid_in_kernel, so those get
silently swallowed into **kwargs without raising if passed the ncp-k3 way,
producing a KDA that trains but computes the wrong gate. The gate is now
computed externally via fla.ops.kda.gate.fused_kda_gate before calling the
kernel, matching HF's own Glm5NextTextLinearAttention reference forward. A
regression test asserts these three kwargs stay absent from chunk_kda's
signature.

Relaxes TransformerConfig.linear_attention to GatedDeltaNetConfig |
KDAConfig | None and the corresponding attention_config type annotations in
dense/moe decoder layers and model builders (mypy-strict plumbing only, no
behavior change for existing GatedDeltaNet layers).

Verified with tests/model/test_glm53_kda.py (5 tests, single/2-GPU): gate
parity against fla's naive reference, module-level parity against HF's
Glm5NextTextLinearAttention (weights bridged per the published<->HF key
mapping), packed multi-document isolation, and Ulysses SP parity.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant