[6/8] feat(models): port Qwen3.5 / Qwen3-Next GDN to the head-sharded layer

[Zhichenzzz]

What

Qwen3.5 / 3.6 / 3.8 and Qwen3-Next GDN move from the replicated HF shim onto the shared layer. GatedDeltaNet joins linear_attn.py as the second family. Qwen3_5GatedDeltaNet and Qwen3NextGatedDeltaNet only declare their projections.

• Only q/k/v and the conv change layout (group-major). The bridges and exporters permute exactly those two; the FP8 export is unchanged.
• Old-layout checkpoints are rejected, not loaded scrambled: Megatron's default load keeps init values for missing tensors. The layer saves a weight_layout_version buffer, load_checkpoint checks for it, and the mbridge bridges keep it out of HF weights. CONVERSION_VERSION goes to 4.
• Qwen3.8-Next keeps the replicated GDN (qwen3_8_next/gdn.py) and HF's flat layout. Its bridge and exporter skip the permutation; without that, R3 replay mismatches 283 of 384 tokens.

vs current main (GB300, 2026-10-08)

The Qwen3.5 GDN self_attention from main (replicated HF shim) vs this PR, forward + backward per layer, SP when TP > 1, real model configs:

Model	TP	16k	64k	peak memory at 64k
Qwen3.5-9B (16 k / 32 v heads)	1	1.12x	1.12x	11.7 -> 10.5 GiB
	2	1.50x	1.92x	11.5 -> 5.8 GiB
	4	1.51x	3.06x	11.3 -> 3.75 GiB
Qwen3.8-27B (16 k / 48 v heads)	1	1.14x	1.13x	16.3 -> 14.1 GiB
	2	1.90x	1.99x	16.0 -> 7.7 GiB
	4	2.12x	3.16x	15.8 -> 4.7 GiB

At 4k tokens the layer is launch-bound (~7 ms either way) and roughly even with main.

Summary of an earlier run (H200, before the rebase)

Precision against the replicated layer:

• TP=1: output and every parameter gradient are bitwise equal, except those downstream of fla's grouped-value accumulation. For Qwen3.5 and Qwen3-Next, those are as close or closer to an fp32 reference.
• TP=2/4/8: 2.9e-3 to 4.5e-3 relative, which is bf16 reduction order.
• CP=2/4, with a sequence boundary inside a shard: GDN 1e-3, the same distance from fp32 as without CP.

Qwen3.8-Flash GDN (2560 hidden, 16 k-heads, 48 v-heads), forward + backward per layer. Speedup and memory saving are against the replicated layer:

Tokens	TP=2	TP=4	TP=8	TP=8 memory saving
4k	1.13x	1.18x	1.45x	2.9x
16k	2.01x	3.10x	3.41x	4.1x
128k	2.14x	3.60x	5.32x	4.75x

RL end to end, Qwen3.5-4B with TP=2 + SP: train / rollout logprob absolute difference is 0.007–0.010, and step time matches the replicated baseline.

Note on CP. GDN scales less well with CP (3.05x at CP=8, 128k) than KDA does. fla dispatches GDN to FlashQLA without CP and to its Triton kernel with CP, so prefer TP over CP for GDN.

This branch (0d33e13a1)

Test	H200	GB300
Head-sharded vs replicated, TP=2: Qwen3.5 / Qwen3-Next / KDA	out 2.85e-3 / 2.84e-3 / 2.86e-3, dx 4.9e-3 / 4.8e-3 / 4.2e-3	pass
Head-sharded vs replicated, TP=4: Qwen3.5 / Qwen3-Next / KDA	out 3.5e-3 / 3.5e-3 / 3.5e-3, dx 5.4e-3 / 5.2e-3 / 4.7e-3	pass
CP correctness vs no CP, CP=2: GDN / KDA	fwd 4.9e-4 / 1.8e-3, dx 1.0e-3 / 3.6e-3	pass
CP correctness vs no CP, CP=4: GDN / KDA	fwd 5.3e-4 / 1.7e-3, dx 1.2e-3 / 3.4e-3	pass
CP relayout, CP=2 / CP=4	pass / pass	pass / pass
test_linear_attn_layout.py, test_linear_attn_short_conv.py	4 passed, 2 passed	–, 2 passed
test_qwen3_5_mtp_bridge_mapping.py, test_fp32_param_utils.py	6 passed, 13 passed (same as main)	–

CI's GLM-5 end-to-end jobs on this PR stop before training because the runner's wandb key is rejected. Both were run on H200 at this stack (dd2cc6f6b) without wandb and pass: test_glm5_744b_a40b_4layer_r3.py (train/rollout logprob abs diff 0.047, 0.049) and test_glm5_744b_a40b_4layer_ci.py (0.009). They caught two GLM-5 call-site bugs in #3607, fixed there; see its description.

H200 results

8×H200 on the current radixark/miles:dev image (pulled 2026-09-23): torch 2.13 + CUDA 13.0, sglang 0.5.21.dev55, tilelang 0.1.14, fla 0.5.2, FlashQLA 0.1.2, FlashMLA; Megatron-LM 7a1f7fe3f (the image's); this stack at dd2cc6f6b. Sequence parallel on; CP uses the zigzag relayout as in training; error is against the replicated layer at CP=1 and against no CP otherwise.

Speedups are against the replicated layer this stack replaces. In multi-GPU runs that reference repeatedly measured 30–900× slow while the head-sharded timings stayed put (they match the earlier runs to a few percent), so the TP=2/4/8 rows carry no speedup here. Clean measurements of the same code: TP=4 at 32k, 3.37× (Qwen3.8 GDN), 3.39× (GLM-5.3 KDA), 3.58× (Kimi-K3 KDA), from two separate reruns on two H200 boxes; TP=8, 1.45–5.32× (GDN), 2.2–5.5× (GLM-5.3 KDA), 3.8–6.2× (Kimi-K3 KDA) with 2.9–5.6× lower activation memory, from the first full run.

qwen3.8-flash, forward + backward ms per layer (speedup vs replicated at CP=1)

TP×CP	4k	8k	16k	32k	64k	128k	max rel err out / dx
1×1	4.4 (1.11x)	7.0 (1.14x)	13.3 (1.17x)	26.3 (1.17x)	52.6 (1.15x)	110.5 (1.14x)	0.0e+00 / 3.7e-03
1×2	7.9	7.8	12.8	22.6	44.4	89.4	4.2e-03 / 5.0e-03
2×1	4.7	4.8	8.5	15.7	29.7	58.5	2.9e-03 / 4.8e-03
1×4	7.0	6.7	7.3	13.4	25.8	50.7	4.2e-03 / 5.0e-03
2×2	7.8	8.0	8.5	13.9	29.0	51.6	4.8e-03 / 5.5e-03
4×1	4.7	4.7	6.2	10.5	18.3	35.3	3.5e-03 / 5.2e-03
1×8	6.6	6.8	6.9	9.0	16.8	32.3	4.2e-03 / 5.0e-03
2×4	7.8	7.5	7.3	8.4	15.7	30.4	4.8e-03 / 5.5e-03
4×2	9.0	8.2	8.2	10.3	18.7	35.7	5.5e-03 / 6.2e-03
8×1	4.8	4.8	4.7	6.9	13.7	24.4	4.5e-03 / 6.0e-03

GB300 results

8×GB300 (sm_103) on two 4-GPU nodes, GPUs otherwise idle; torch 2.13 + CUDA 13.0, tilelang 0.1.12, fla 0.5.2, FlashMLA b7643bd built for sm_100f; commits as pushed to this stack. Sequence parallel on; CP runs use the zigzag relayout as in training. FlashQLA is not installed on GB300, so fla runs GDN on its Triton kernel with and without CP; that is why TP=1 differs from the replicated layer by 5e-3 here (native GVA vs repeated q/k take different kernels) while it is bitwise on H200.

qwen3.8-flash, forward + backward ms per layer (speedup vs replicated at CP=1)

TP×CP	4k	8k	16k	32k	64k	128k	max rel err out / dx
1×1	5.2 (0.98x)	7.5 (0.95x)	11.1 (1.16x)	21.9 (1.17x)	43.5 (1.16x)	86.9 (1.16x)	5.1e-03 / 5.5e-03
1×2	10.3	8.8	8.9	15.9	30.8	61.1	0.0e+00 / 1.4e-04
2×1	7.9 (1.20x)	7.8 (0.95x)	8.5 (1.54x)	14.3 (1.81x)	26.7 (1.91x)	51.1 (1.98x)	5.4e-03 / 6.3e-03
1×4	10.7	9.4	10.1	13.7	18.4	39.2	8.7e-03 / 5.6e-03
2×2	10.8	10.0	11.2	12.3	19.3	37.2	0.0e+00 / 1.6e-04
4×1	7.1 (1.26x)	7.7 (0.94x)	6.4 (2.10x)	9.7 (2.67x)	18.5 (2.91x)	37.0 (2.75x)	6.2e-03 / 6.6e-03

Local end to end, tier 2 (8×H200, stack tip dd2cc6f6b, Megatron 8a5dbe5ff, no wandb)

These are the non-model-specific megatron e2e tests CI runs, rerun locally because the CI wandb key is rejected on PR runs. All 19 pass:

Test	Result	Time
test_qwen3_0.6B_sft_snapshot_eval	pass	185 s
test_mimo_7B_mtp_only_grad	pass	372 s
test_qwen3_4B_muon_offload_disk	pass	581 s
test_qwen3_4B_offload_disk	pass	592 s
test_qwen3_4B_offload_disk_stream	pass	460 s
test_qwen3_30B_A3B/test_baseline	pass	1228 s
test_qwen3_30B_A3B/test_dp_attention	pass	1023 s
test_qwen3_30B_A3B/test_r3_baseline	pass	1024 s
test_qwen3_4B_variable_gbs	pass	343 s
test_qwen3_4B_ppo	pass	531 s
test_gpt_oss_20b_moe_lora_ci	pass	729 s
test_qwen3_4B_disk_delta	pass	331 s
test_qwen3_30B_A3B_p2p	pass	660 s
test_glm47_flash/test_r3_mtp	pass	946 s
test_qwen3_30B_A3B/test_disagg_broadcast	pass	1087 s
test_qwen3_30B_A3B/test_fully_async	pass	906 s
test_qwen3_4B_offload_disaggregated	pass	293 s
test_nemotron_3_ultra_4layer_ci	pass	590 s
test_qwen3_4b_fully_async_eval	pass	1129 s

CI's H100 timeout on test_mtp1_spec_v2_r3.py

CI's stage-c-8-gpu-h100 (1) timed out this test at its 2000 s budget while it was still making progress (R3 replay checks passing), and its sibling test_mtp0_spec_v2.py passed at 2000 s in the same run. Rerun on one 8×H200 box, main (a6b72d31f) and this stack (6893c4e88) both pass: 2111 s on main, 2030 s here. The timeout is the test's budget, not a slowdown from this PR.

Local end to end, tier 1b (8×H200, this stack 6893c4e88, Megatron 8a5dbe5ff, no wandb)

Test	Result	Time
test_glm5_744b_a40b_4layer_r3	pass	1289 s
test_glm5_2_744b_a40b_5layer_ci	pass	803 s
test_deepseek_v32_5layer_fp8	pass	1311 s
test_deepseek_v4_flash_4layer_ci	pass	1008 s
test_deepseek_v4_flash_4layer_megatron_impl_ci	pass	663 s
test_dsv4_bshd_thd_rollout_parity	pass	870 s
test_glm5_1_744b_a40b_6layer_lora_ci	fails, on main too	–

test_glm5_1_744b_a40b_6layer_lora_ci stops at actor init with dsa_kernel_backend=tilelang requires fused DSA kernels ... TileLang DSA kernels. Megatron's check imports megatron.core.transformer.experimental_attention_variant.ops.tilelang_dsa, which is not in the image's Megatron (7a1f7fe3f), in 8a5dbe5ff, or in the current miles-main. The same test on main (a6b72d31f) on the same box fails with the same error, so it is not from this stack.

Stack

	PR	Change
1	#3605	non-attention kernels → miles/kernels
2	#3606	attention kernels → miles/kernels/attention
3	#3607	one DSA kernel pair (GLM-5, DeepSeek-V4, V4.1)
4	#3608	per-arch DSA sparse attention dispatch
5	#3609	shared head-sharded linear-attention layer (miles_plugins/models/linear_attn.py)
6	#3610	Qwen3.5 / Qwen3-Next GDN on it
7	#3632	Kimi-K3 KDA on it
8	#3634	short-conv and KDA backward backends

Rules for miles/kernels are in its README: only kernels we write, one directory per op, names say what is computed (never the backend), callers import the entry, a torch-reference test per op.

[claude]

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.

_{Tip: disable this comment in your organization's Code Review settings.}

[yueming-yuan]

[P2] Exclude the layout marker from offline HF export
This persistent buffer is included in distributed checkpoints, but tools/convert_torch_dist_to_hf.py passes every loaded model tensor to convert_to_hf(). The Qwen converters reject this key with ValueError: Unknown parameter name: module.module.decoder.layers.0.self_attention.linear_attn.weight_layout_version, so checkpoints produced with this layer can no longer be exported through that tool. The bridge's filtering does not cover the offline exporter. Please exclude this metadata from HF export while retaining it in distributed checkpoints.