Conversation
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>
jayhenry
added this pull request to stack #2112
September 23, 2026 06:54
This was referenced Sep 23, 2026
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stack (bottom to top):
feat/glm53flash-materialize-full-f0→mainfeat/glm53flash-f3-kda→feat/glm53flash-materialize-full-f0← you are herefeat/glm53flash-f4-mhc→feat/glm53flash-f3-kdafeat/glm53flash-f5-nope-dsa→feat/glm53flash-f4-mhcfeat/glm53flash-f1-vl-data→feat/glm53flash-f5-nope-dsafeat/glm53flash-f2-vision-tower→feat/glm53flash-f1-vl-datafeat/glm53flash-f6-text-moe→feat/glm53flash-f2-vision-towerSummary
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) perdoc/xtuner_glm5p3flash_design.mdF3, with module structure and parameter names matching the published checkpoint layout (separateq/k/v_conv1d, low-rankf_a_proj->f_b_projforget gate with fp32A_log/dt_bias, low-rankg_a_proj->g_b_projoutput gate).Root cause of a real correctness bug found while implementing: the installed
fla0.4.2'schunk_kda/fused_recurrent_kdacall convention differs from thextuner-ncp-k3reference — it doesn't acceptA_log/dt_bias/use_beta_sigmoid_in_kernel, and silently swallows them into**kwargswithout raising, producing a KDA that trains but computes the wrong gate. Fixed by computing the gate externally viafla.ops.kda.gate.fused_kda_gatebefore calling the kernel, matching HF's ownGlm5NextTextLinearAttentionreference forward. A regression test pins these three kwargs' absence fromchunk_kda's signature so this can't silently regress.Also relaxes
TransformerConfig.linear_attentiontoGatedDeltaNetConfig | 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'sGlm5NextTextLinearAttention(weights bridged per the published<->HF key mapping), packed multi-document isolation, and Ulysses SP parity.