Skip to content

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

Merged
jayhenry merged 5 commits into
feat/glm53flash-materialize-full-f0from
feat/glm53flash-f3-kda
Oct 15, 2026
Merged

jayhenry merged 5 commits into
feat/glm53flash-materialize-full-f0from
feat/glm53flash-f3-kda

Conversation

@jayhenry

@jayhenry jayhenry commented Sep 23, 2026 •

Copy link
Copy Markdown
Collaborator

Stack #2179 (bottom to top):

  1. [Fix] Preserve DeepEP layout events in synchronous dispatch #2177 fix/deepep-layout-event → main
  2. [Feature] Add GLM-5.3-Flash F0: 25B cropped reference checkpoint builder #2105 feat/glm53flash-materialize-full-f0 → fix/deepep-layout-event
  3. [Feature] Add GLM-5.3-Flash F3: Kimi Delta Attention (KDA) #2106 feat/glm53flash-f3-kda → feat/glm53flash-materialize-full-f0 ← you are here
  4. [Feature] Add GLM-5.3-Flash F4: mHC four-stream residual #2107 feat/glm53flash-f4-mhc → feat/glm53flash-f3-kda
  5. [Feature] Add GLM-5.3-Flash F5: NoPE DSA + KPool indexer + clamped SwiGLU #2108 feat/glm53flash-f5-nope-dsa → feat/glm53flash-f4-mhc
  6. [Feature] Add GLM-5.3-Flash F1: VL data preprocessing pipeline #2109 feat/glm53flash-f1-vl-data → feat/glm53flash-f5-nope-dsa
  7. [Feature] Add GLM-5.3-Flash F2: vision tower + projector (eager) #2110 feat/glm53flash-f2-vision-tower → feat/glm53flash-f1-vl-data
  8. [Fix] Correct generic MTP future embeddings and empty-shard loss #2178 fix/mtp-local-future-embeddings → feat/glm53flash-f2-vision-tower
  9. [Feature] Add GLM-5.3-Flash F6 core: text model + MTP + compose model #2111 feat/glm53flash-f6-text-moe → fix/mtp-local-future-embeddings

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


Summary

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

The KDA compile, precision, kernel-gradient and sequence-parallel fixes previously in F6 are now included in this layer. The shared fp32_linear utility is here because the KDA gate uses it.

Adds xtuner/v1/module/attention/kda.py (KDAConfig / KimiDeltaAttention) 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).

The Read the Docs configs now mock einops, which KDA imports while Sphinx loads XTuner with docs-only dependencies. English and Chinese HTML builds passed locally with Python 3.12 and Sphinx 6.2.1.

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.

from fla.ops.kda import fused_recurrent_kda as _fused_recurrent_kda
from fla.ops.kda.gate import fused_kda_gate as _fused_kda_gate

class FusedRMSNormGated(_FLAFusedRMSNormGated):

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

是否可以和 Qwen3.5 的合并?

batch_size, seq_len, _ = hidden_states.shape
assert batch_size == 1, "KDA currently supports packed batch size 1"

cu_seqlens = seq_ctx.cu_seq_lens_q

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

这样做和使用 mixed_qkv相比,哪个好?

cu_seqlens = seq_ctx.cu_seq_lens_q
projection_size = self.num_heads * self.head_dim

q = self._sp_short_conv(

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

直接让 self.q_conv1d 等支持 SP 是否更好?

jayhenry and others added 4 commits October 15, 2026 03:39
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>
Compare the chunked KDA forward output and every module gradient against
HF's Glm5NextTextLinearAttention on GPU, including the split causal
convolutions that HF fuses into one conv1d.
GLM-5.3 HF parity tests require the glm5_next API introduced in Transformers 5.17. The upgrade also changes Qwen video interpolation baselines. Use the public video processor API, preserve fp32 vision RoPE, and mirror HF gated norm rounding in parity mode.
Co-authored-by: 纪焘 (Tao Ji) <taoji@fudan.edu.cn>
Co-authored-by: Kuang <128125671+Knlife@users.noreply.github.com>
Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
@jayhenry
jayhenry force-pushed the feat/glm53flash-f3-kda branch from 1eec2f5 to d8d52d5 Compare October 15, 2026 03:55
@jayhenry
jayhenry removed this pull request from stack #2112 October 15, 2026 03:58
@jayhenry
jayhenry added this pull request to stack #2179 October 15, 2026 03:59
@jayhenry
jayhenry merged commit e619594 into main Oct 15, 2026
9 of 10 checks passed
@jayhenry
jayhenry deleted the feat/glm53flash-f3-kda branch October 15, 2026 05:28
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.

2 participants