Add context parallelism to the Gated Delta Net (Qwen3-Next / Qwen3.5)

[NuojCheng]

Description

Adds context parallelism to the Gated Delta Net used by Qwen3-Next and Qwen3.5, with a cross-rank payload that does not grow with sequence length.

Why

GDN currently forces the full sequence onto every rank. The sharding constraint before the causal conv replicates the sequence axis, and jax_chunk_gated_delta_rule is a lax.scan over chunks that assumes it owns the whole sequence. With ici_context_parallelism > 1 the GDN layers all-gather everything back, so the context axis buys nothing on the linear-attention half of these hybrid models.

Softmax-attention context parallelism (ring, all-gather, Ulysses) exchanges O(S) key/value tensors. Linear attention does not have to. The chunked gated delta rule is affine in the incoming recurrent state:

h_out = (gamma * I - k_tilde^T @ w) @ h + k_tilde^T @ u   ==  M @ h + U
o     = (q_g - attn_i @ w) @ h + attn_i @ u

Affine maps compose associatively, so a rank's entire shard collapses to a single pair (M_r, U_r), and stitching the ranks back together is a prefix scan over those pairs. M_r is (B, H, K, K) and U_r is (B, H, K, V) — both independent of sequence length. For qwen3.5-35b-a3b the crossover against softmax-attention CP is around 1024 tokens per shard; past that, GDN's exchange is strictly cheaper and stays flat while everything else grows.

Note this is not the Mamba-2/GLA situation. Those have a scalar decay, so M degenerates to a scalar and the combination is trivial. GDN's delta rule makes M a genuine K x K matrix, which is why it needs the explicit transition-composition machinery below.

Implementation

jax_chunk_gated_delta_rule_cp runs three phases:

1. Local pass from a zero state, carrying the composed transition M_r alongside the state. Outputs are skipped, since the true incoming state is not known yet.
2. all_gather the (M_r, U_r) pairs and run an exclusive prefix scan, giving every rank the state that enters its shard.
3. Local pass again with the correct incoming state, reusing the WY factors from phase 1, now emitting outputs.

M is composed in its low-rank form (M_new = gamma * M - k_tilde^T (w @ M)), costing 2 * K^2 * C per chunk instead of a dense K^3 product.

Phase 1 makes the GDN core about 50% more expensive, but the core is only ~8% of the GDN block (the projections dominate: ~67 MFLOP/token vs ~5.9 MFLOP/token). AOT confirms this: +0.20% total model FLOPs at cp=4.

To share numerics between the two paths, the monolithic jax_chunk_gated_delta_rule is split into _gdn_chunk_prepare / _gdn_scan_body / _gdn_finalize. Its public signature and behavior are unchanged; the single-device path is the same computation, just reassembled.

On the layer side:

• The pre-conv sharding constraint keeps the sequence sharded when the context axis is active. Replicating there was what forced the full-sequence all-gather.
• No manual halo code for the depthwise conv. GSPMD already handles the seq-sharded causal conv1d: the compiled HLO shows a collective-permute of bf16[1,3,512] — exactly gdn_conv_kernel_dim - 1 tokens — and zero all-gathers.
• The recurrent state is never sequence-sharded, so state_pspec is untouched and stays replicated across the context axis.
• Chunk padding is provably harmless per-rank: padded positions get g=0 (identity transition) and k=v=q=0 (zero contribution), so they are a no-op for the recurrence and get sliced off the output.

Limitation: load balancing must be off

context_parallel_load_balance (default true) applies a zigzag reorder so rank r holds global chunks r and 2*cp-1-r. That gives each rank two non-contiguous pieces of the sequence, which a sequential recurrence cannot consume — the prefix scan assumes rank order equals sequence order. Rather than silently computing the wrong result, the layer raises:

ValueError: GDN context parallelism requires context_parallel_load_balance=False. The gated
delta rule is a sequential recurrence, so each rank must own a contiguous, in-order slice of
the sequence; the load-balancing reorder interleaves chunks across ranks and would silently
compute the wrong result.

Since GDN has no causal-mask triangle, its per-rank cost is uniform anyway, so it gains nothing from load balancing. The full-attention layers of these hybrid models do lose the balancing, which is a real cost and a candidate for future work (the reorder is applied at the data-iterator level in train_utils.py, so making it layer-local would be the fix).

Tests

New tests/unit/gdn_context_parallelism_test.py (5 tests, all passing on a v5 TPU VM):

• test_forward_matches_single_device / test_backward_matches_single_device — kernel parity against the single-device recurrence under shard_map.
• test_cross_rank_payload_is_sequence_length_independent — parses the compiled HLO, sums all-gather volume at S=512/1024/2048, asserts it is constant and equals exactly cp * B * H * K * (K + V) elements.
• test_layer_output_matches_without_context_parallelism — the full Qwen3NextGatedDeltaNet layer is invariant to CP degree.
• test_load_balancing_is_rejected — the guard above fires.

python -m pytest tests/unit/gdn_context_parallelism_test.py -q
# 5 passed

Measured results:

Check	Result
Kernel forward parity (cp=4, f32)	output rel err 9.9e-8, final state 1.4e-7
Kernel backward parity	q/k/v/g/beta rel err 2.3e-7 … 4.3e-7
Cross-rank HLO payload	f32[8,4,128,128], identical at S = 512 / 1024 / 2048
Conv halo	collective-permute of bf16[1,3,512], zero all-gathers

AOT compile, qwen3-next-80b-a3b on v5p-64 at 8192 context:

	baseline	cp=4	delta
FLOPs	21,874,845,155,328	21,919,038,439,424	+0.20%
HBM bytes	559,759,491,072	579,720,052,736	+3.6%
peak temp	64,807,110,912	70,465,474,464	+8.7%

Real training runs, 4x TPU v5, 4096 context, 6 steps, synthetic data, seeded:

step	cp=1	cp=2	cp=4
0	8.119	8.119	8.119
1	8.104	8.105	8.105
2	8.091	8.092	8.092
3	8.082	8.082	8.082
4	8.076	8.076	8.076
5	8.073	8.073	8.074

Matching to float reduction-order noise. Also verified cp=1 vs cp=2 at 1024 context (identical: 8.126/8.098/8.073/8.055/8.044/8.039).

One note for reviewers reproducing this: cp=4 at 1024 context fails with q_block_size=512 should divide q_seq_len_per_shard=256. That is the pre-existing splash-attention kernel constraint on the hybrid model's full-attention layers, not GDN. Hence 4096 context for the cp=4 comparison.

Also ran pyink and pylint (10.00/10) over both changed files.

Checklist

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Code Review

This pull request introduces context parallelism for the Gated Delta Net (GDN) in Qwen3. It refactors the jax_chunk_gated_delta_rule implementation into modular helper functions and adds jax_chunk_gated_delta_rule_cp to support context-parallel execution. Additionally, a new unit test suite tests/unit/gdn_context_parallelism_test.py is added to verify correctness, sequence-length independence of cross-rank communication, and proper error handling. I have no feedback to provide as there are no review comments.

[codecov]

Codecov Report

❌ Patch coverage is 91.30435% with 6 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/models/qwen3.py	91.30%	3 Missing and 3 partials ⚠️

📢 Thoughts on this report? Let us know!