Skip to content

Add GatedDeltaNet evaluation export for Qwen hybrid models - #2442

Closed
Baiju Meswani (baijumeswani) wants to merge 14 commits into
baijumeswani/packed-hybrid-io-publishfrom
baijumeswani/qwen38-gdn-eval-publish-ready
Closed

Add GatedDeltaNet evaluation export for Qwen hybrid models#2442
Baiju Meswani (baijumeswani) wants to merge 14 commits into
baijumeswani/packed-hybrid-io-publishfrom
baijumeswani/qwen38-gdn-eval-publish-ready

Conversation

@baijumeswani

Copy link
Copy Markdown
Collaborator

Summary

This PR adds a GatedDeltaNet export path for Qwen hybrid recurrent layers.

It replaces VarlenLinearAttention with the ONNX Runtime GatedDeltaNet contrib 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

  • Added a builder helper that emits com.microsoft::GatedDeltaNet.
  • Uses packed THD Q/K/V inputs and cumulative sequence lengths from the existing varlen export.
  • Preserves the FP32 V-major recurrent-state contract.
  • Uses the reviewed gated_delta update rule.
  • Passes precomputed effective decay and beta gates with activation attributes set to none.
  • Casts effective gates to the operator's FP32 gate ABI.
  • Disables state-checkpoint outputs because the Engine commits final state transactionally.
  • Rejects BF16 model I/O because the current ONNX Runtime kernel supports FP16 and FP32 only.
  • Adds builder tests for node inputs, outputs, attributes, gate casts, state geometry, update-rule selection, and supported dtypes.

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:

  • scalar-per-head dt_bias validation;
  • dynamic shared-memory setup;
  • forced-plan behavior;
  • supported dtype advertising;
  • update-rule input validation;
  • state alias declaration.

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.

Copilot AI lite review requested due to automatic review settings August 21, 2026 02:06
@baijumeswani
Baiju Meswani (baijumeswani) requested a review from a team as a code owner August 21, 2026 02:06

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

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 + GatedDeltaNet when 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.

Comment thread src/python/py/models/builders/qwen.py
Comment thread src/python/py/models/README.md Outdated
@baijumeswani
Baiju Meswani (baijumeswani) force-pushed the baijumeswani/qwen38-gdn-eval-publish-ready branch from 9c13ab9 to 9b3d165 Compare August 22, 2026 05:01
Comment thread src/python/py/models/builders/qwen.py Fixed
Comment thread src/python/py/models/builders/qwen.py Fixed
Comment thread src/python/py/models/builders/qwen.py Fixed

@tianleiwu Tianlei Wu (tianleiwu) left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

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, 0 GroupQueryAttention, 48 GatedDeltaNet,
    48 VarlenCausalConvWithState; input_ids [num_tokens], no attention_mask;
    state_groups = 1 paged_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.

Comment thread src/python/py/models/builders/base.py Outdated
Comment thread src/python/py/models/builders/qwen.py Outdated
Comment thread src/python/py/models/builders/base.py Outdated
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
@baijumeswani
Baiju Meswani (baijumeswani) force-pushed the baijumeswani/qwen38-gdn-eval-publish-ready branch from beb17d8 to 3951b7b Compare August 23, 2026 07:21
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
@baijumeswani
Baiju Meswani (baijumeswani) force-pushed the baijumeswani/qwen38-gdn-eval-publish-ready branch from 3951b7b to 266f08a Compare August 23, 2026 07:23
@baijumeswani

Copy link
Copy Markdown
Collaborator Author

Superseded by #2453. Closing this PR.

@baijumeswani
Baiju Meswani (baijumeswani) deleted the baijumeswani/qwen38-gdn-eval-publish-ready branch August 28, 2026 19:38
An error occurred while trying to automatically change base from baijumeswani/packed-hybrid-io-publish to main August 29, 2026 07:37
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.

4 participants