Repository navigation
refactor(layers): share one GatedDeltaNet op in layers/linear_attn - #841
Merged
Merged
Conversation
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.
Refs #828
The two GatedDeltaNet ops,
Qwen3_5GatedDeltaNet(qwen3_5_moe) andQwen4ExpGatedDeltaNet(qwen4_exp), become oneGatedDeltaNetinfreetoken/layers/linear_attn/. No behavior change.Summary
The two classes were copies of one file: their
forward, conv helpers, gating and parameter names were the same line for line. The only difference in behavior was the output-norm gate: qwen4_exp passed the group config'soutput_gatetoGatedRMSNorm, while qwen3_5_moe usedGatedRMSNorm's default,silu, which is also what its group config sets. qwen4_exp also imported its kernel wrappers frommodels/qwen3_5_moe/gdn_kernels.py, one model family importing another.layers/linear_attn/gdn.py:GatedDeltaNet, with anoutput_gateargument (defaultsilu). Both families build it withoutput_gate=g.output_gate, so each keeps the gate it had.kernel/fla/(chunk_gated_delta_rule,fused_sigmoid_gating_delta_rule_update) directly, as the KDA layer does. The two wrappers inmodels/qwen3_5_moe/gdn_kernels.pypassed fixed arguments, cast the slot indices and offsets to int32 / int64 and tooko[0]; that now sits in the layer and the wrappers are removed; the prefill branch now takes the per-chunk stateshfrom the kernel's return value instead of areturn_hflag.layers/linear_attn/gdn_reference.py: the pure-torch reference the layer test compares against. It is the qwen4_exp reference, which takes the gate as an argument; the qwen3_5_moe reference had no users and is removed, as is the reference's chunked delta rule, which only a test of the reference itself used.models/qwen3_5_moe/gdn.py,models/qwen3_5_moe/gdn_kernels.py,models/qwen4_exp/gdn.pyand both familygdn_reference.pyfiles are removed; comments that pointed at the old paths now point at the new one.models/glm5_next/kda.py) has one user and stays in its family.Qwen3_5MoE,iter_expert_pieces); this PR only moves the GDN op.Why now: the speculative verify step (#836) adds a verify branch to the GDN layer. With one layer it is written once, next to the decode and prefill branches. vLLM also keeps one GDN layer for its Qwen GDN models (
QwenGatedDeltaNetAttentioninvllm/model_executor/layers/mamba/gdn/).Tests
tests/models/qwen4_exp/test_gdn.pymoves totests/layers/test_gdn.py(the README'slayers/row now lists GatedDeltaNet; the test sets the TP info itself and swaps its context intocore._GLOBAL_CTXwithmonkeypatch, which the qwen4_expconftest.pydid before) and keeps one test,test_ragged_prefill_then_decode: a ragged bs=3 prefill (128, 1000 and 37 tokens) and one decode step against the reference, for the 2:1 head ratio with the silu gate and the 3:1 ratio with the sigmoid gate. It replaces three tests:Test environment
Unit tests
1684 passed, 202 skipped, 1 failed. The failure is
tests/models/test_quant_config.py::test_scheme_for_agrees_with_the_stored_tensors[apodex/Apodex-1.1-mini-NVFP4[index-only]](a local checkpoint withquant_method: modelopt_mixed); it fails the same way on main.End-to-end A/B against main
AIME25 problem 1 through
freetoken.llm.LLM(model_path, dtype=torch.bfloat16, max_running_req=4)with the chat template, greedy (temperature=0,max_tokens=16384, natural end), on main a2ce636 and on this branch, both with CUDA graphs, compared token by token:Qwen/Qwen3.6-35B-A3B(qwen3_5_moe, silu gate)Qwen/Qwen3.6-27B(qwen3_5_moe, dense)nvidia/Qwen3.8-Flash-Next-NVFP4(qwen4_exp, sigmoid gate)Qwen/Qwen3.6-35B-A3B-FP8(qwen3_5_moe, splitin_proj_qkvz/in_proj_ba)ISTA-DASLab/Qwen3.8-Flash-Next-GSQ-RCO-GGUFIQ3_XXS (qwen4_exp, GGUF tiledssm_outheads)AIME25 problems 1-3 submitted together on
Qwen/Qwen3.6-35B-A3B(cache_type=hybrid_radix, greedy, 1024 tokens each withignore_eos): all three identical to main token by token.Smoke test on this branch, passed: