[RFC][Qwen3-Next-80B-A3B-Instruct][CUDA/ROCm] VIME/vLLM operator-level train-rollout consistency roadmap, ablation matrix, and integration plan

[inaniloquentee]

Qwen3-Next-80B-A3B-Instruct

Status: Proposed
Target: post-v0.1.0 community roadmap

---

1. Motivation

Qwen3-Next combines Gated DeltaNet (GDN), full attention, and 512-expert top-10 MoE. The main train/rollout mismatch is not that every operator needs a new kernel: training usually runs chunked sequences, while rollout runs prompt prefill followed by recurrent decode. Strict mode therefore needs the same provider, state boundary, reduction order, and runtime manifest on both sides.

Goal

Given the same logical tokens, weights, positions, masks, and sampling context, VIME replay and vLLM rollout must produce exactly equal selected-token log-probabilities in strict mode.

Reuse-first decision

The following are intentionally not claim rows when an existing implementation passes shape and numerical qualification: causal convolution, full attention, basic GDN forward, RMSNorm, dense GEMM, LM-head/logprob, standard MoE grouped GEMM/dispatch, and NCCL/RCCL collectives. Candidate providers are FlashInfer, FlashQLA/FLA, causal-conv1d, vLLM, AITER, and the native BLAS/collective libraries. RL-Kernel only owns the missing contract, bridge, ordering, and validation work.

2. Checkpoint fingerprint

The initial claim is pinned to the exact Qwen3-Next-80B-A3B-Instruct checkpoint and runtime manifest used by each validation run:

• Qwen3NextForCausalLM, BF16, hidden size 2048, vocabulary size 151936, untied embeddings.
• 48 decoder layers: [GDN, GDN, GDN, FullAttention] × 12; 36 GDN layers and 12 full-attention layers.
• Full attention: Q/KV heads 16/2, head dimension 256, partial RoPE 0.25, RoPE theta 10,000,000.
• GDN: key/value heads 16/32, head dimensions 128/128, depthwise causal-conv kernel 4.
• MoE: 512 routed experts, top-10 activation, expert and shared-expert intermediate size 512.
• RMSNorm epsilon 1e-6; top-k probabilities normalized; maximum position 262144.

A different checkpoint revision, accelerator generation, library revision, or kernel manifest is a separate claim.

3. Architecture map

3.1 Platform provider map

CUDA and ROCm share the ABI, trace schema, test generator, and numerical contract, but use separate provider manifests and kernel goldens.

3.2 Strict replay boundary

Concatenating prompt and completion into one chunked GDN call is mathematically equivalent, but is not generally bitwise equivalent to prompt prefill followed by recurrent decode. The prompt-derived GDN, convolution, and KV state must therefore cross the same explicit boundary in trainer replay and rollout.

4. Current stack and strict contract

• VIME already contains a Qwen3-Next model path and exposes train/rollout logprob comparison. The remaining work is provider alignment, state handoff, stable MoE ordering, and platform manifests.
• vLLM exposes batch-invariant execution, but an environment variable alone is not acceptance evidence; the selected backend, kernel, split policy, and provenance must be checked.
• ROCm is a separate claim. AITER, FLA, ROCm attention, hipBLASLt/rocBLAS, and RCCL may be reused, but CUDA goldens cannot be reused as ROCm evidence.

Every strict row must satisfy:

1. Fixed accumulator dtype, reduction order, casts, and provider configuration.
2. No dynamic Split-K, split-KV, cross-CTA atomics, or load-dependent autotuning unless the merge tree is locked in the manifest.
3. Target row/token/state is unchanged by unrelated batch members, padding, packing, sequence order, or concurrency.
4. Prompt/response state layout, dtype, padding, and ownership are identical between replay and rollout.
5. Unsupported shape, dtype, topology, accelerator, workspace, or provider fails closed.
6. Runtime provenance records backend/library revision, binary or build ID, kernel/algorithm ID, tiles, split policy, accumulator dtype, compile key, graph state, collective algorithm, and topology.
7. Strict acceptance is torch.equal / zero absolute difference for the declared output; allclose is diagnostic only.

5. Claim table

Status: 🙋 open → ⏳ in progress → 👀 in review → ✅ merged

One row × one platform = one PR. Reused kernels are dependencies, not claim rows. A row is complete only when its implementation and validation land together.

Work item	What it does	Reuse / dependency	CUDA owner	CUDA PR	CUDA status	ROCm owner	ROCm PR	ROCm status
strict_provider_manifest	Pin checkpoint, provider, kernel/config fingerprints, reduction settings, and fail-closed checks for replay and rollout	Extend VIME/vLLM tracing; new integration glue	@fusheng-ji		🙋			🙋
gdn_state_bridge	Make prompt-prefill → decode GDN/convolution state, layout, dtype, padding, and gradient handoff identical	Reuse FlashQLA/FLA/causal-conv1d/AITER where qualified; no new GDN math unless a gap is proven	@fusheng-ji		🙋			🙋
moe_route_combine_contract	Lock top-10 routing order, FP32 weight normalization, permutation, and stable atomic-free combine	Reuse vLLM/AITER dispatch and grouped GEMM where semantics match; implement only ordering/reduction gaps	@fusheng-ji		🙋			🙋
distributed_logprob_parity	Lock TP/EP reduction order and exact selected-token logprob between VIME replay and vLLM rollout	Reuse NCCL/RCCL and existing logprob paths; add no collective kernel unless required			🙋			🙋
full_model_chain	End-to-end train/prefill/decode parity with first-drift localization by layer family	Compose qualified providers and the rows above; extend existing model-chain tests	@fusheng-ji		🙋			🙋
batch_invariance_ci	Validate packing, chunk boundaries, unrelated-sequence interference, and 1–64 concurrent requests	Extend the vLLM batch-invariant/metamorphic harness; no new operator	@fusheng-ji		🙋			🙋

6. Recommended claim order

1. strict_provider_manifest
2. gdn_state_bridge, moe_route_combine_contract, and distributed_logprob_parity in parallel
3. full_model_chain
4. batch_invariance_ci
5. TP/EP topologies and additional accelerator-specific manifests after the TP1 strict claim is green

7. Contribution notes

• CUDA and ROCm rows implement the same versioned contract and test suite; only providers and native performance baselines differ.
• A PR that covers only GDN chunked prefill does not close gdn_state_bridge; prompt state, decode state, and the trainer gradient path belong to the same claim.
• MoE PRs must state expert ordering, combine ordering, payload dtype, and atomic policy explicitly.
• Strict mode fails closed instead of silently falling back to an unqualified backend or schedule.
• Performance numbers are reported against the qualified native path for training, prefill, and decode; correctness remains the primary merge gate.

If you are interested, just ping below!

[Bignonia7]

can i take this issue? @Flink-ddd

[Flink-ddd]

Hi @Bignonia7 , Thanks for your interest! This RFC covers the full Qwen3-Next roadmap across multiple operator families and CUDA and ROCm tracks. Are you planning to coordinate the full-model effort, or would you like to take a specific operator or milestone? If it’s a specific part, please let us know which one so we can scope a separate issue for it.

[fusheng-ji]

Hi @Flink-ddd — following up on the claim question above with specific rows. I would like to take the CUDA column of the following rows of the claim table:

Row	Status I propose	What is ready
strict_provider_manifest	⏳ in progress → 👀 in review	Registry dispatch fix, Qwen3-Next RMSNorm references with a zero-centred CUDA kernel, and the gated RMSNorm kernel, each with contract tolerances and reproducible checks. Opened: #466 (registry dispatch fix, bottom of the stack), #467 (RMSNorm references + zero-centred CUDA kernel, on #466) and #468 (gated RMSNorm kernel, on #467).
gdn_state_bridge	⏳ in progress → 👀 in review	GDN decode-step PyTorch goldens and the explicit convolution/recurrent state contract (RFC C6), with the measured conv mismatch causes documented. Opened: #469 (GDN decode-step PyTorch goldens and the explicit convolution/recurrent state contract, on #468).
moe_route_combine_contract	⏳ in progress	Top-10 routing with fixed-order FP32 softmax/normalisation (shape-independent) and ordered atomic-free combine; validated in a TP4 native-scheduler batch-invariance matrix (batch 1/4 × prompt 8/64/256/1024 × full/chunked, 16/16 raw-bit). PR follows once the full chain lands.
full_model_chain	⏳ in progress	Full-model L2 on the pinned Qwen3-Next-80B-A3B-Instruct checkpoint, TP4 BF16 weights / FP32 recurrent state: vLLM rollout vs VIME (Megatron) recompute, 16 cases × 48 layers raw-bit equal (tokens, BF16 logits, FP32 logprobs, per-layer state digests), with first-drift localisation by layer. Training/reload acceptance in progress. PR follows once it passes on one frozen commit.
batch_invariance_ci	⏳ in progress	The native-scheduler matrix above plus packing/chunk-boundary cases; the metamorphic harness extension comes with the chain PR.

distributed_logprob_parity is exercised by the chain (TP4 selected-token logprob equality across engines) but I am not claiming it separately yet; happy to hand it to someone or fold it in later. ROCm columns stay open.

Hardware: a single 4×B200 node per run; CUDA 13.0, vLLM 0.30.0, TE 2.16.1. All evidence is pinned to the exact checkpoint revision and a runtime manifest as the RFC requires. The first four PRs reference each other and must merge bottom-up (registry → RMSNorm → gated RMSNorm → GDN goldens).

PRs opened so far (CUDA column):

• [WS1][CUDA][Qwen3-Next] Restore CUDA dispatch for rms_norm (RFC #428 stack 1/4) #466 — [WS1][CUDA][Qwen3-Next] Restore CUDA dispatch for rms_norm (RFC #428 stack 1/4)
• [WS1][CUDA][Qwen3-Next] C1: zero-centred RMSNorm references and CUDA kernel (RFC #428 stack 2/4) #467 — [WS1][CUDA][Qwen3-Next] C1: zero-centred RMSNorm references and CUDA kernel (RFC #428 stack 2/4) (stacked on [WS1][CUDA][Qwen3-Next] Restore CUDA dispatch for rms_norm (RFC #428 stack 1/4) #466)
• [WS1][CUDA][Qwen3-Next] Gated RMSNorm CUDA kernel for the GDN block (RFC #428 stack 3/4) #468 — [WS1][CUDA][Qwen3-Next] Gated RMSNorm CUDA kernel for the GDN block (RFC #428 stack 3/4) (stacked on [WS1][CUDA][Qwen3-Next] C1: zero-centred RMSNorm references and CUDA kernel (RFC #428 stack 2/4) #467)
• [WS1][CUDA][Qwen3-Next] C6: GDN decode-step goldens and recurrent/conv state contract (RFC #428 stack 4/4) #469 — [WS1][CUDA][Qwen3-Next] C6: GDN decode-step goldens and recurrent/conv state contract (RFC #428 stack 4/4) (stacked on [WS1][CUDA][Qwen3-Next] Gated RMSNorm CUDA kernel for the GDN block (RFC #428 stack 3/4) #468)

The four PRs merge bottom-up; each one's description states what must merge first. Per the note above,
they now target test-qwennext and have been rebased onto it (conflicts resolved); #470 is an independent
CI fix and stays on main.

[Flink-ddd]

cc @inaniloquentee PTAL

[Flink-ddd]

Hi everyone! Please note that all PRs should target the test-qwennext branch instead of main. Once the CI and validation tests pass successfully, the changes will be merged into main. Thanks for your contribution!