feat(compile): traceable GatedDeltaNet decoder loop (Qwen3.5 + Qwen3-Next) + model-agnostic FA2 varlen collator

[thad0ctor]

Description

Under torch_compile, one graph break inside the GatedDeltaNet decoder for-loop made dynamo skip the whole text-model forward frame (a break inside a loop is unresumable), so every layer (attention included) ran eagerly.

This makes the loop traceable for every GatedDeltaNet model (Qwen3.5 dense + MoE, Qwen3-Next), plus a model-agnostic collator path that fixes the same break for any flash-attention + packed model under compile:

• gated_delta_net_ops.py (new, shared) — opaque torch.library ops (axolotl_gdn::gdn_conv/gdn_chunk + backward) wrapping the same FLA kernels with identical saved tensors. They derive cu_seqlens from position_ids inside the op, so no aten.nonzero enters the graph. A cast_g flag matches each architecture's eager call (qwen3_5 casts g; qwen3_next keeps f32) for bitwise parity; packed input enforces batch-1 with an explicit ValueError. Also hosts the shared FusedRMSNormGated compile boundary.
• qwen3_5/ + qwen3_next/modeling.py — route the GDN training path through the ops under compile, and install the FusedRMSNormGated opaque wrapper (its backward recomputes via autograd.grad, since the FLA backward isn't meta-traceable). Self-attention stays in-graph under GC+training, behind a dynamo.disable boundary otherwise (guards an Inductor FA2-backward fusion that corrupts packed grads; pinned by test_fa2_compiled_matches_eager_grads). The qwen3_next wiring also fixes a pre-existing crash (its decoder passed cache_position to a GatedDeltaNet.forward that rejects it).
• collators/batching.py + builders/causal.py — under torch_compile, the multipack collator precomputes the FA2 varlen kwargs (cu_seq_lens_q/k, max_length_q/k) once per batch, so transformers skips its per-layer nonzero derivation. Model-agnostic — gated on SUPPORTED_MULTIPACK_MODEL_TYPES, so any packed FA2 model compiles clean.
• patch_manager.py — threads torch_compile into the qwen3_5/qwen3_5_moe/qwen3_next packing patches.

Opt-in via torch_compile; inert otherwise.

Version coupling — the ==0.4.1 pin is load-bearing. The ops wrap FLA internal functions by fixed arity and the fakes hardcode 0.4.1 shapes/dtypes; FLA 0.4.2/0.5.0 already changed chunk_gated_delta_rule_fwd's arity (crash at first forward). So a bump needs revalidating the ops — guarded by test_opcheck_custom_ops + test_eager_parity_ops_vs_legacy_bitwise (the call-time crash isn't caught by the build-time warning). The upstream fix fla-org/flash-linear-attention#909 removes only the index-helper breaks, not the nonzero/@torch.compiler.disable ones these ops target, so the wrapper can't just be dropped on a newer FLA.

Known limit: auto-enabled LoRA triton kernels + a GDN packing patch + torch_compile + GC NaN gradients on torch ≤ 2.10 (pre-existing on main; clean on 2.11, and with any one ingredient removed). PatchManager force-disables those kernels with a warning on torch < 2.11.

Motivation and Context

The decoder loop is the bulk of training compute. Skipping its compilation forfeits the largest fusion payoff: the eager RMSNorm fp32 round-trips (~870 GB/step of traffic at 9B/32k, interposer-measured) + per-kernel Python launch overhead. Once the loop traces, Inductor fuses those.

Scope: the GDN ops cover Qwen3.5/Qwen3.6 (model_type qwen3_5/qwen3_5_moe, dense + MoE, text + VL — every shipped checkpoint is *ForConditionalGeneration) and Qwen3-Next (qwen3_next, same FLA GatedDeltaNet kernels). The collator change is model-agnostic. Dense Qwen3/Qwen2.5/Qwen3-VL have no GatedDeltaNet, so the loop-internal breaks don't exist for them — but they still benefit from the collator path under FA2+packing+compile.

How has this been tested?

Stack: torch 2.11.0+cu130, transformers 5.9.0, FLA 0.4.1, liger-kernel 0.8.0, real flash_attn 2.8.3 (the validated stack). RTX 3090 / 3090 Ti / 5090.

Functional — 12 tests, all passing:

tests/monkeypatch/test_qwen3_5_compiled_loop.py (10): zero-breaks-with-GC (1 graph), FA2-loop-with-varlen-kwargs, VL ForConditionalGeneration 3-D MRoPE, aten.nonzero-break-gone, bitwise eager parity (ops vs legacy), compiled-vs-eager, FA2+GC compiled grad parity, toy MoE (0 breaks), B>1-packed-raises, opcheck (packed/MRoPE/dense × bias × both cast_g).
tests/monkeypatch/test_qwen3_next_compiled_loop.py (2): qwen3_next decoder loop 0 breaks + bitwise eager parity (cast_g=False, g stays f32).
Existing test_qwen3_5_fused_attn.py: 12/12 (no regression). Suite also passes one major version down (torch 2.9.1 / 2.10.0 + transformers 5.5.4).

Multimodal: Multimodal training is non-packed in axolotl (SFT vision examples + MM-CPT (pending PR axolotl-ai-cloud#3629) ship sample_packing: false), so the GDN ops are inert there by construction this was validated with real images (llava-instruct through the Qwen3.5-2B vision tower: packing patches absent from logs, eager-vs-compiled step-1 loss Δ0.004, eval clean).

Multiple test training runs and benches (below).

---

Benchmarks (all on the validated stack above)

Qwen3.5-2B/9B = Qwen3_5ForConditionalGeneration, LoRA r=16 q/k/v/o, sample packing, bf16, gradient checkpointing, attn_implementation: flash_attention_2. 60 steps, steady-state window after warmup, TORCH_LOGS=recompiles. Throughput = tok/s; VRAM = reserved (GiB). Zero NaN and step-1/final-10 loss parity in every cell.

A. Core compile speedup On real FA2, torch_compile without this PR is near-neutral (the loop graph-breaks); this PR is what makes compile worth it (RTX 3090):

model	eager	compile (main, no PR)	compile + PR	PR vs eager	PR vs compile-main
Qwen3.5-2B (seq 4096)	3593 tok/s	3788 (+5.4%)	4617	+28.5%	+21.9%
Qwen3.5-9B (seq 4096, CCE)	2724 tok/s	2732 (+0.3%)	3330	+22.2%	+21.9%

B. Feature compatibility / composition (Qwen3.5-2B, compile + PR, varying feature): (RTX 3090)

config	tok/s	reserved VRAM	note
compile + PR (LoRA kernels on — default)	4617	14.66 GiB	baseline
+ LoRA kernels off	4585	14.40 GiB	kernels ≈ neutral under compile
+ fused_attn_kernel (axolotl-ai-cloud#3680)	4394	14.67 GiB	−4.8% — redundant under compile
+ axolotl CCE (cut_cross_entropy)	4781	6.03 GiB	fastest and ~2.4× less memory
+ liger fused_linear_cross_entropy	4593	14.66 GiB	no-op on Qwen3.5 (= no memory benefit) — requires overlooked gating bug fix following liger 0.8.0 bump
+ liger full stack (rms_norm/gated/swiglu/rope)	4137	10.60 GiB	−10% — kernels redundant under compile
(eager + fused_attn_kernel)	3648	18.07 GiB	eager: fused_attn only +1.5%

LoRA triton kernels and fused_attn_kernel are subsumed by Inductor once the loop compiles (consistent with axolotl-ai-cloud#3680's own "fused+compile −18%" dense-Qwen3 numbers); liger composes cleanly but adds no value under compile.

C. Generalization — the model-agnostic collator (compile on/off, FA2 + packing): (RTX 5090)

model	eager	compile + PR	speedup	reserved (eager → compile)
Llama-3.2-1B (seq 2048)	5181 tok/s	7572	+46.1%	6.09 → 4.54 GiB
Qwen3-0.6B (seq 2048)	5445 tok/s	11375	+108.9%	5.54 → 5.01 GiB

Neither model has a GatedDeltaNet so speedups come from the collator removing the FA2 per-layer break. Isolation (toy Llama, FA2 + packed, under compile): 3 graph breaks (4 graphs) without the varlen kwargs → 0 breaks (1 graph) with them. 0 recompiles in steady state; loss parity holds.

---

Convergence: loss is bit-stable across every cell above — step-1 identical to eager and final-10 within float noise (e.g. 2B eager 1.1275 vs compile+PR 1.1275; 9B 0.9264 vs 0.9266). The opaque-op path is bitwise identical to legacy eager (regression-tested) for both qwen3_5 and qwen3_next.

Composition with axolotl-ai-cloud#3732 (fused-LoRA GDN routing): merges cleanly, both touch the GDN forward, the projection call sites route through _la_proj_fwd while the compiled-ops gating stays intact, and the fused LoRA_O autograd Function traces inside the compiled loop with 0 graph breaks (earlier 2B bench: both-merged+compile +44% vs eager, axolotl-ai-cloud#3732's +10% marginal surviving compilation because it substitutes bf16 adapter GEMMs Inductor can't infer).

(Orthogonal note: any LoRA + GC + compile run carries one pre-loop graph break from transformers' enable_input_require_grads embedding hook; on torch 2.11 it fires before the decoder loop and doesn't affect it.)

AI Usage Disclaimer

Opus 4.8 / Fable 5 used throughout.

Types of changes

• New feature (torch_compile support for the GatedDeltaNet decoder loop — Qwen3.5 dense/MoE + Qwen3-Next)
• New feature (model-agnostic FA2 varlen collator kwargs for any packed model under compile)
• Performance improvement (non-breaking; numerics-neutral, bitwise eager fallback)
• New tests

Summary by CodeRabbit

• New Features

  • Added torch.compile optimization support for Qwen3.5 and Qwen3Next models with improved kernel handling.
  • Enhanced Flash Attention 2 compatibility with varlen metadata precomputation for better performance.

• Bug Fixes

  • Fixed Qwen3.5 LoRA training compatibility when using sample packing with gradient checkpointing on PyTorch < 2.11.

• Tests

  • Added regression test suites validating torch.compile compatibility and performance parity for Qwen model architectures.

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 9b1e5889-d874-4ad8-9550-f1c4f59c52b8

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 Walkthrough

Walkthrough

Adds opaque torch.library.custom_op wrappers (axolotl_gdn.gdn_conv, gdn_chunk) around FLA GatedDeltaNet kernels to prevent aten.nonzero from entering torch.compile graphs. Threads torch_compile flags through Qwen3.5 and Qwen3-Next monkeypatches to route through these ops. Precomputes FlashAttention varlen metadata in DataCollatorForSeq2Seq when enabled. Wires everything via PatchManager and adds CUDA regression tests for zero graph breaks and eager/compiled parity.

Changes

torch.compile GatedDeltaNet Opaque Op Integration

Layer / File(s)	Summary
GatedDeltaNet opaque custom-op library src/axolotl/monkeypatch/models/gated_delta_net_ops.py	New module: lazily registers torch.library.custom_op wrappers (gdn_conv, gdn_chunk) that move cu_seqlens computation from position_ids inside the op boundary, preventing aten.nonzero from entering compile graphs. Also registers a FusedRMSNormGated compile boundary with recompute-based backward, and exposes fla_ops_available() / fla_ops_build_error() with cached build state.
Qwen3.5 modeling patches for compiled-op routing src/axolotl/monkeypatch/models/qwen3_5/modeling.py	Adds _FLA_COMPILED_OPS module state and _call_self_attn Dynamo-disabled wrapper. Adds compile_boundary flag to _inject_fla_kernels to install FusedRMSNormGated boundary. Updates decoder forward to choose between torch.ops.axolotl_gdn.gdn_conv/gdn_chunk (compiled path) vs cu_seqlens+chunk_gated_delta_rule (eager path). Extends _apply_packing_patches, patch_qwen3_5_modeling_packing, and patch_qwen3_5_moe_modeling_packing with torch_compile parameter.
Qwen3-Next modeling patches for compiled-op routing src/axolotl/monkeypatch/models/qwen3_next/modeling.py	Mirrors Qwen3.5 changes: adds _FLA_COMPILED_OPS, Dynamo-disabled self-attn wrapper, and torch_compile parameter to all four public patch functions. Decoder layer threads position_ids into linear_attention and conditionally wraps self_attn with the Dynamo-disable boundary. GatedDeltaNet layer selects axolotl_gdn.gdn_conv/gdn_chunk or falls back to FLA eager paths.
FA varlen kwargs precomputation in collator and builder src/axolotl/utils/collators/batching.py, src/axolotl/core/builders/causal.py	DataCollatorForSeq2Seq gains emit_fa_varlen_kwargs: bool = False; when set, precomputes cu_seq_lens_q/k and max_length_q/k from position_ids once per batch with a warning_once fallback. HFCausalTrainerBuilder.build_collator sets this flag when flash_attention_2, packed mode, torch_compile, and a supported multipack model type are all active.
PatchManager wiring and LoRA kernel guard src/axolotl/loaders/patch_manager.py	Passes torch_compile=bool(self.cfg.torch_compile) to all three Qwen3-family packing patch calls. Adds a pre-model-load guard that disables Qwen3.5 LoRA triton kernels on PyTorch < 2.11 when sample packing + torch_compile + gradient checkpointing + adapter training are all enabled.
CUDA regression tests tests/monkeypatch/test_qwen3_5_compiled_loop.py, tests/monkeypatch/test_qwen3_next_compiled_loop.py	CUDA-only test suites with packing_patched fixtures that snapshot/restore patched methods. Qwen3.5 tests cover zero Dynamo graph breaks across gradient-checkpointing, FA2 varlen, VL, and MoE scenarios; bitwise and tolerance parity between compiled-op and eager paths; batch-gt-1 raise guard; and torch.library.opcheck for gdn_chunk/gdn_conv. Qwen3-Next tests cover zero graph breaks with gradient checkpointing and bitwise parity toggling _FLA_COMPILED_OPS.

Estimated code review effort

🎯 5 (Critical) | ⏱️ ~120 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 40.79% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and accurately summarizes the main changes: adding torch.compile support for GatedDeltaNet decoder loops in Qwen3.5/Qwen3-Next models and a model-agnostic FA2 varlen collator.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch feat/gdn-compiled-decoder-loop

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[github-actions]

📖 Documentation Preview:

Deployed on Netlify from commit 9953bce