Skip to content

refactor(layers): share one GatedDeltaNet op in layers/linear_attn - #841

Merged
jason-fxz merged 2 commits into
mainfrom
refactor/shared-gdn-layer
Oct 29, 2026
Merged

jason-fxz merged 2 commits into
mainfrom
refactor/shared-gdn-layer

Conversation

@jason-fxz

@jason-fxz jason-fxz commented Oct 28, 2026 •

Copy link
Copy Markdown
Collaborator

Refs #828

The two GatedDeltaNet ops, Qwen3_5GatedDeltaNet (qwen3_5_moe) and Qwen4ExpGatedDeltaNet (qwen4_exp), become one GatedDeltaNet in freetoken/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's output_gate to GatedRMSNorm, while qwen3_5_moe used GatedRMSNorm's default, silu, which is also what its group config sets. qwen4_exp also imported its kernel wrappers from models/qwen3_5_moe/gdn_kernels.py, one model family importing another.

  • layers/linear_attn/gdn.py: GatedDeltaNet, with an output_gate argument (default silu). Both families build it with output_gate=g.output_gate, so each keeps the gate it had.
  • The layer calls the fla kernels in kernel/fla/ (chunk_gated_delta_rule, fused_sigmoid_gating_delta_rule_update) directly, as the KDA layer does. The two wrappers in models/qwen3_5_moe/gdn_kernels.py passed fixed arguments, cast the slot indices and offsets to int32 / int64 and took o[0]; that now sits in the layer and the wrappers are removed; the prefill branch now takes the per-chunk states h from the kernel's return value instead of a return_h flag.
  • 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.py and both family gdn_reference.py files are removed; comments that pointed at the old paths now point at the new one.
  • KDA (models/glm5_next/kda.py) has one user and stays in its family.
  • qwen4_exp still imports other pieces from qwen3_5_moe (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 (QwenGatedDeltaNetAttention in vllm/model_executor/layers/mamba/gdn/).

Tests

tests/models/qwen4_exp/test_gdn.py moves to tests/layers/test_gdn.py (the README's layers/ row now lists GatedDeltaNet; the test sets the TP info itself and swaps its context into core._GLOBAL_CTX with monkeypatch, which the qwen4_exp conftest.py did 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:

  • the single 1000-token prefill, which the ragged prefill covers;
  • the gate test, which the two parameter sets cover: each gate is checked against a reference with that gate, so a layer that ignored the config fails one of them;
  • the test that the reference's chunked and recurrent delta rules agree, which tested the reference, not the layer.

Test environment

  • GPU: 1x NVIDIA H100 80GB HBM3, driver 580.95.05
  • CPU: Intel Xeon Platinum 8480+ (26 vCPUs), 221 GB RAM, Ubuntu 22.04.5
  • CUDA 13.0, torch 2.14.1+cu130, Python 3.12.12
  • This branch at 84cbd87, A/B against main a2ce636. The unit tests ran on 84cbd87; the end-to-end runs ran on 531d855, and the later commit changes only the layer test and comments.

Unit tests

FREETOKEN_QWEN36_MODEL=/path/to/Qwen3.6-27B PYTHONPATH=python pytest tests/layers tests/scheduler tests/engine tests/kvcache tests/attention tests/moe/test_offload.py tests/server tests/tokenizer tests/models tests/kernels/test_kda.py

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 with quant_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:

Checkpoint main this branch Token by token
Qwen/Qwen3.6-35B-A3B (qwen3_5_moe, silu gate) 16384 tokens (hit the limit) 16384 tokens (hit the limit) identical
Qwen/Qwen3.6-27B (qwen3_5_moe, dense) 15676 tokens, answer 70 15676 tokens, answer 70 identical
nvidia/Qwen3.8-Flash-Next-NVFP4 (qwen4_exp, sigmoid gate) 727 tokens, answer 70 727 tokens, answer 70 identical
Qwen/Qwen3.6-35B-A3B-FP8 (qwen3_5_moe, split in_proj_qkvz / in_proj_ba) 15188 tokens, answer 70 15188 tokens, answer 70 identical
ISTA-DASLab/Qwen3.8-Flash-Next-GSQ-RCO-GGUF IQ3_XXS (qwen4_exp, GGUF tiled ssm_out heads) 850 tokens, answer 70 850 tokens, answer 70 identical

AIME25 problems 1-3 submitted together on Qwen/Qwen3.6-35B-A3B (cache_type=hybrid_radix, greedy, 1024 tokens each with ignore_eos): all three identical to main token by token.

Smoke test on this branch, passed:

FREETOKEN_TEST_MODEL=/path/to/Qwen3.6-35B-A3B FREETOKEN_AIME25_JSONL=/path/to/aime25.jsonl PYTHONPATH=python pytest tests/e2e/test_aime.py

@jason-fxz
jason-fxz merged commit d043fb6 into main Oct 29, 2026
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant