Add GatedDeltaNet evaluation export for Qwen hybrid models - #2442
Conversation
There was a problem hiding this comment.
Pull request overview
Adds an evaluation-oriented export path for Qwen hybrid recurrent layers that replaces the packed varlen linear-attention recurrence with ONNX Runtime’s com.microsoft::GatedDeltaNet, while keeping the upstream normalization/gating behavior and preserving the FP32 V-major recurrent-state contract for hybrid paging/Engine usage.
Changes:
- Add builder helpers for packed varlen conv/state, packed varlen linear attention, and a GatedDeltaNet “evaluation” emission path (including FP32 gate casts and BF16 rejection).
- Update Qwen hybrid builder to support packed leading dims and to emit
VarlenCausalConvWithState+GatedDeltaNetwhen using the packed paged-attention layout. - Add/extend Python builder tests validating node wiring, attributes, dtype/state contracts, and packed reshapes.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| test/python/builder/test_decoder_state_groups.py | Adds unit tests covering varlen conv/linear attention state contracts and the new GatedDeltaNet evaluation emission behavior. |
| src/python/py/models/README.md | Updates paged-attention documentation, including hybrid packed-layout behavior and state_groups/CUDA-graph implications. |
| src/python/py/models/builders/qwen.py | Implements packed-vs-dense leading-dim handling, packed mRoPE path, and swaps packed recurrent op emission to GatedDeltaNet with FP32 V-major state. |
| src/python/py/models/builders/base.py | Adds paged_attention_uses_packed_layout() and new builder helpers for varlen conv/state, varlen linear attention, and GatedDeltaNet evaluation emission; disables CUDA graph for fixed-state hybrid builds. |
Suppressed comments (2)
src/python/py/models/builders/qwen.py:1960
- This comment says state shapes are identical between dense and packed layouts, but the recurrent state layout differs (dense uses K-major [B, heads, key_dim, value_dim] while packed uses V-major [B, heads, value_dim, key_dim]). Keeping the note about the batch dimension is helpful, but the shape claim is incorrect.
Dense and packed builds share every op except the two above: state batch dimension is
always the request count (never the packed token count), so state shapes are identical
between the two layouts.
src/python/py/models/builders/qwen.py:1956
- The _make_linear_attention docstring lists LinearAttention/VarlenLinearAttention as the recurrent ops, but the packed path now uses GatedDeltaNet (and the dense path uses LinearAttention). The operator list should match the actual op selection to avoid confusion when debugging exported graphs.
Uses com.microsoft contrib ops:
- CausalConvWithState / VarlenCausalConvWithState: fused depthwise conv1d + SiLU + carry
state, dense channel-first [B,C,S] or packed token-major [num_tokens,C].
- LinearAttention / VarlenLinearAttention: fused linear attention with GQA, dense
[B,S,...] or packed token-major [num_tokens,...].
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
9c13ab9 to
9b3d165
Compare
Tianlei Wu (tianleiwu)
left a comment
There was a problem hiding this comment.
Thanks — I built this stack end-to-end and it works. Everything below is about how
it lands next to #2453, plus one rebase hazard that lives lower in the stack.
What I validated (this branch + #2453, on ORT #32182 + #32168; Qwen3.8-27B-NVFP4, H200):
- 64-layer paged export: 16
PagedAttention, 0GroupQueryAttention, 48GatedDeltaNet,
48VarlenCausalConvWithState;input_ids [num_tokens], noattention_mask;
state_groups= 1paged_kv(16 layers) + 2 fixed (48 each). - Ragged-batch correctness: a batch of 17 / 512 / 1993-token prompts produced
bit-identical output to the same prompts run alone. That is the property packed
conv/recurrent state is most likely to get wrong, and it holds. - Quality is statistically indistinguishable from the dense export under an identical greedy
recipe: MMLU-Pro 800 83.25% vs 83.00% (McNemar exact p=0.88), GPQA-diamond
75.25% vs 76.26% (p=0.86), 0 unextracted answers. - Throughput: 5,969 tok/s prefill at 2k (vs 914 dense), decode 343 tok/s at batch 8,
587 tok/s at batch 64, 24 tok/s at a 256k context.
Three inline comments cover the #2453 reconciliation. Three further notes that do not map to
a line in this PR's diff:
1. A rebase hazard that belongs to #2427, not here. _make_config_model in
test/python/builder/test_decoder_state_groups.py:61 builds its model with __new__ and
never sets use_windowed_paged_kv_cache. That attribute did not exist when #2427 was
branched, but #2358 has since landed on main and made make_genai_config call
has_windowed_paged_layers() for every paged build. I rebased this stack onto current
main and got four failures:
AttributeError: 'Qwen35TextModel' object has no attribute 'use_windowed_paged_kv_cache'
test_common_paged_builder_preserves_legacy_manifest_absence
test_common_nonpaged_builder_preserves_manifest_absence
test_qwen_all_attention_builder_preserves_legacy_manifest_absence
test_qwen38_official_geometry_emits_exact_sparse_groups
Adding model.use_windowed_paged_kv_cache = False beside model.window_size = None fixes
all four; False is right because Qwen3.8 has no sliding window. Flagging it here because
this is the PR whose export it blocks, but the one-line change belongs at the bottom of the
stack.
2. _setup_hybrid_cache_io needs its state-dtype rule consolidated with #2453. This PR
keys V-major FP32 recurrent state off use_paged_attention; #2453 keys it off
linear_attn_op == "gated_delta_net". Merging both naively leaves two competing assignments
to recurrent_state_dtype and a stale state_dtype reference. The single rule that worked
for me was v_major_state = self.use_paged_attention or self.linear_attn_op == "gated_delta_net",
feeding both the dtype and the trailing-extent swap.
3. Unrelated but it affects every paged build. The native FP8/NVFP4 MatMul emitters in
base.py hardcode ["batch_size", seq_dim, out_features] rather than going through
hidden_state_shape(). With use_paged_attention=true prune_lm_head=true the LM head then
declares logits as [batch_size, batch_size, vocab]. Already fixed in #2449, so it
resolves itself if that merges first — noting it so it is not diagnosed twice.
I have the full reconciliation applied and tested in a combined tree, and am happy to hand
over the patch or open a PR against this branch.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: a8d618bb-0feb-4d05-a918-93b6accde1b1
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: a8d618bb-0feb-4d05-a918-93b6accde1b1
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: a8d618bb-0feb-4d05-a918-93b6accde1b1
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: a8d618bb-0feb-4d05-a918-93b6accde1b1
beb17d8 to
3951b7b
Compare
Rename the packed evaluation helper to avoid the dense GDN contract collision and remove the unused VarlenLinearAttention helper for the dropped ORT operator. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: a8d618bb-0feb-4d05-a918-93b6accde1b1
3951b7b to
266f08a
Compare
|
Superseded by #2453. Closing this PR. |
Summary
This PR adds a GatedDeltaNet export path for Qwen hybrid recurrent layers.
It replaces
VarlenLinearAttentionwith the ONNX RuntimeGatedDeltaNetcontrib operator while preserving the same upstream normalization, Q/K/V layout, causal convolution, effective gates, recurrent-state layout, and model weights. This keeps model-level comparisons focused on the recurrent operator.What changed
com.microsoft::GatedDeltaNet.gated_deltaupdate rule.none.Evaluation contract
This path is intentionally conservative. It reuses the same precomputed gates as the varlen graph so the comparison isolates GatedDeltaNet, but casting FP16 gates to FP32 cannot recover precision already lost in the shared FP16 path.
A production exporter should use GatedDeltaNet's raw Qwen decay and sigmoid-beta fusion inputs directly. That would avoid the precision round-trip and remove separate gate kernels.
Dependency
This branch depends on ONNX Runtime PR #32182 and the reviewed local fixes for:
dt_biasvalidation;The exporter selects GatedDeltaNet unconditionally, so this should remain an evaluation PR until the ONNX Runtime dependency is available. A production version should add an explicit backend option or use a coordinated ONNX Runtime version boundary.