[fusheng-ji]

Stack: #466 (merged) → #467 → #468 → #469 → #503 → #507 → #543 → #544 → this PR → #557 → benchmarks (#551), base test-qwennext; merges after #544. This PR's own diff. RFC #428 row distributed_logprob_parity, CUDA, TP part (EP = 1); expert parallelism is the next PR (#557).

What

• Contract (docs/operators/qwen3-next-distributed-logprob.md): every TP4 point where ranks combine values -- embedding, row-parallel GDN/attention/MoE outputs, column-parallel input gradients, the vocabulary-parallel selected logprob / LSE / entropy, sampling logits -- with its operation, order, dtype and transport (NCCL all-gather for bytes only; RL-Kernel's fixed-order collective cuda_ipc_fixed_tree for sums).
• Parity runner (tools/validation/models/qwen3_next/qwen3_next_logprob_parity.py, 4 GPUs, real weights): SHA-256 of every point's output and gradient per rank, a TP1 reference, and a raw-NCCL negative control; qwen3_next_logprob_parity_compare.py checks one digest per rank across NCCL settings.
• CI guard (tests/models/qwen3_next/test_qwen3_next_tp_reductions.py): no raw torch.distributed reduction in Qwen3-Next model, engine or VIME modules; the TP reduction must use the fixed-order collective.
• Evidence publisher with the chain publisher's redaction rules.
• Layout as reviewed on this PR: validation code and JSON under rl_engine/validation/models/qwen3_next/, scripts under tools/validation/models/qwen3_next/, comparison benchmarks and plots in the benchmark PR ([WS1][CUDA][Qwen3-Next] Comparison benchmarks and plots for the Qwen3-Next stack: MoE, attention, logprob, EP (RFC #428 stack 11/11) #551), raw JSON as release ZIPs with summaries in reports/experiments/qwen3_next/. The earlier PRs of the stack now follow the same layout themselves.

Prior art & reuse decision

Same official-LM-head logits for every measured candidate (TP4 shards are slices of one full-vocabulary product), 1024 tokens:

Library	Implementation	TP4 == TP1	Invariant over 13 NCCL settings	Batch invariant	Max abs vs FP64	TP4 latency	Verdict
RL-Kernel WS2	shared_logprob (both engines here)	yes	yes	yes	7.8e-7	2.55 ms
vLLM 0.30.0	sampler: all-gather + log_softmax	yes	yes	yes	8.1e-7	1.51 ms	exact, but rollout only
SGLang 0.5.21	logits processor: all-gather + log_softmax (srt/layers/logits_processor.py)	—	—	—	—	—	same structure as vLLM (source); rollout only
Megatron-core 0.16	vocab_parallel_cross_entropy	no	no (2 results)	yes	8.1e-7	0.57 ms	NCCL all-reduce of max and sum-exp
Vime 0.3.2	default calculate_log_probs_and_entropy	no	no (2 results)	yes	8.1e-7	0.56 ms	as Megatron
Miles (radixark/miles@23d3d38)	fused_log_probs_and_entropy, imported unmodified	no	no (2 results)	yes	8.1e-7	0.33 ms	NCCL all-reduce of max and sums; repeat-stable at a fixed setting; backward store-only, reproducible and batch invariant
verl (volcengine/verl@e5a627c)	linear_cross_entropy, imported unmodified (fuses the LM-head GEMM; own logits, FP64 reference of its own product)	no	no (2 results)	yes	1.2e-5	0.63 ms	NCCL all-reduce; forward repeat-stable (its forward atomic_add runs only for reduction sum/mean); backward d_hidden/d_weight use tl.atomic_add but are repeat-stable (20/20), not batch invariant (probe rows 1/10 cells)
CUTLASS / FlashInfer / FlashAttention	—						N/A: no vocabulary-parallel logprob

The measured implementations differ from each other by up to 9.5e-7 (Megatron and VIME agree with each other). vLLM's path is exact for rollout, but the training side has no counterpart that reproduces it; Megatron/VIME, Miles and verl (all measured) reduce max and sum-exp with NCCL all-reduce, so they move with NCCL settings. Decision: both engines keep RL-Kernel's WS2 provider -- the only measured way to equal bits on both sides. Its cost is latency (4.5× Megatron at this shape); a faster tile-stats kernel can keep the same merge order and is follow-up work.

Results (4 × B200, official checkpoint)

Top: digest of every TP4 point under each NCCL setting vs NCCL defaults; the last row is the raw-NCCL negative control. Bottom: prior-art invariance, largest pairwise difference, latency.

Check	Result
17 TP4 points × 4 ranks × 13 NCCL settings (defaults; Ring, all-reduce Tree × LL / LL128 / Simple × 1 / default channels)	one digest per point and rank
TP4 vs TP1 on the same logits	selected logprob and LSE bitwise equal; entropy differs (documented)
negative control: real o_proj partials, NCCL all_reduce	5 different results
full-vocabulary GEMM vs gathered TP4 logits	bitwise equal
end to end (#543): vLLM rollout vs VIME recompute selected logprobs	16 cases × 4 weight versions bitwise

NCCL 2.29 has no Tree all-gather, so the Tree settings restrict all-reduce only (NCCL_ALGO=allreduce:tree). Summary, tested commits and commands: reports/experiments/qwen3_next/qwen3-next-distributed-logprob-b200.md; raw reports ZIP (SHA-256 a48bf29d…2cd9), parity measured at 04db726, prior art at 33495ad. Post-review GPU rerun at b10b6616 (ZIP): all 233 mandatory operator cases and four norm gates passed; all 17 points invariant across the 13 settings; selected logprob / LSE equal TP1; the negative control again gave five digests per rank.

Status

• Head 3f7cf6ec, on top of [WS1][CUDA][Qwen3-Next] Batch-invariance matrix for the full model, 97/97 cells bitwise (RFC #428 stack 8/11) #544 (db1eb37f). CI is green on this head: lint, unit-tests, docs, DCO and the bot review pass; provider and ws1-gtest skip because no self-hosted runner is attached.
• Ready for review. No new kernel or collective: the row documents and proves the order RL-Kernel's existing fixed-order collective and WS2 vocabulary-parallel logprob already give, and adds a CI guard against raw NCCL reductions.
• Known: entropy merges per-shard contributions in TP-rank order, so it is not TP-degree invariant; rollout and replay run the same TP degree and agree.
• Review changes since the first version: layout as requested; benchmarks and plots moved to [WS1][CUDA][Qwen3-Next] Comparison benchmarks and plots for the Qwen3-Next stack: MoE, attention, logprob, EP (RFC #428 stack 11/11) #551; evidence externalised; Python 3.10 hashing fix (3927fad7); stdlib-only launcher runtime.
• Rerun after aligning the stack (clean clone of the stack top 90c09164): the EP1 chain passed through l2, train1, train2 and train (L2 artifacts bitwise identical to the c67ec14 run) and the full BI suite passed 97 / 97 cells (see [WS1][CUDA][Qwen3-Next] Full-model chain: vLLM rollout vs VIME recompute, raw-bit through two updates and a reload (RFC #428 stack 7/11) #543 and [WS1][CUDA][Qwen3-Next] Batch-invariance matrix for the full model, 97/97 cells bitwise (RFC #428 stack 8/11) #544).

History

• 2026-10-27 (provider path filter): the provider workflow's pull_request paths were missing Qwen3-Next modules that its checks import; merged the additions of [WS1][CUDA][Qwen3-Next] Full-model chain: vLLM rollout vs VIME recompute, raw-bit through two updates and a reload (RFC #428 stack 7/11) #543 and [WS1][CUDA][Qwen3-Next] Batch-invariance matrix for the full model, 97/97 cells bitwise (RFC #428 stack 8/11) #544 (3f7cf6ec). Workflow-only change; shared infrastructure such as rl_engine/ops and rl_engine/platforms stays out of the filter.
• 2026-10-26 (review): two OpenCodeReview rounds on the stack (the bot re-reviews every PR's whole diff after each push); every finding of these two rounds is answered in its own thread or comment, most as repeats of earlier answers or as not applicable with the code reference. Fixes in this PR: 5ffd4c31 runs test_file_digest.py once in the CPU job. The fixes of the PRs below it are merged in; new head 5ffd4c31. Kernel arithmetic is unchanged, so the measured results stand.
• 2026-10-26: merged test-qwennext after [WS1][CUDA][Qwen3-Next] C1: zero-centred RMSNorm references and CUDA kernel (RFC #428 stack 2/4) #467 landed (c382dd4b), through [WS1][CUDA][Qwen3-Next] Batch-invariance matrix for the full model, 97/97 cells bitwise (RFC #428 stack 8/11) #544's new head eb04b15f, in fe6ed386. This PR moves the norm workload, manifest and evidence scripts under qwen3_next/, so [WS1][CUDA][Qwen3-Next] Gated RMSNorm CUDA kernel for the GDN block (RFC #428 stack 3/4) #468's later changes to them were taken at those paths, and [WS1][CUDA][Qwen3-Next] Gated RMSNorm CUDA kernel for the GDN block (RFC #428 stack 3/4) #468's reuse check (new here) reads the moved manifest.

Tests

python -m pytest -q tests/models/qwen3_next/test_qwen3_next_tp_reductions.py
NCCL_ALGO=<algo> NCCL_PROTO=<proto> [NCCL_MAX_NCHANNELS=1] torchrun --nproc-per-node 4 \
    tools/validation/models/qwen3_next/qwen3_next_logprob_parity.py --checkpoint <ckpt> --output <dir>/<setting>
python tools/validation/models/qwen3_next/qwen3_next_logprob_parity_compare.py --runs <dir>/* --out report.json

Check	Result
Qwen3-Next CPU list in ci.yml + digest tests	439 passed (99 CUDA-only skips)
CI guard: a planted raw all_reduce / reduce_scatter_tensor is reported	passed
evidence redaction test, pre-commit on all files, layout check	clean
remote CI (lint, unit tests, docs, DCO)	passed at the current head (0b88a839)

Scope

CUDA, TP4 × EP1 × CP1, one node of four B200. Expert parallelism is #557; ROCm is a separate claim.

Results (H100 / Hopper, 4 GPUs, official checkpoint)

Run at this PR's head 8e88cd82. 4 × NVIDIA H100 80GB, torch 2.13.0+cu130, vLLM 0.30.0, NCCL 2.29.7, checkpoint revision 9c7f2fbe. Raw data: hopper-qwen3-next/pr549.

Check	Result
tests/models/qwen3_next/test_*.py + check_*.py (4 GPUs)	787 passed
adapter, TP4 attention, TP4 GDN, TP4 head, TP4 MoE gates	all passed
TP4 logprob parity over 13 NCCL settings (default; Ring, all-reduce Tree × LL / LL128 / Simple × 1 / default channels), 1024 tokens	passed

• All 18 recorded TP4 points × 4 ranks give one digest per point and rank across the 13 settings (deterministic_points_invariant: true).
• TP4 equals TP1 for the selected logprob and LSE bitwise (tp4_equals_tp1: true); entropy differs between TP4 and TP1 as documented (entropy_tp4_equals_tp1: false).
• The control behaves as on B200: the NCCL all_reduce of the o_proj partials gives 8 distinct results over the 13 settings, the fixed-order collective gives 1.
• The full logits match the gathered logits.

Logprob prior-art on H100 (1024 tokens, TP4 vs TP1, 10 repeats; runner at 77d59a59):

Candidate	TP4 == TP1	batch bitwise
rl_kernel_ws2	yes	yes
vllm_gather	yes	yes
megatron_ce	no	yes
vime_default	no	yes
miles_fused, verl_fused	not available in env	–

Report: pr549_logprob/report.json (also parity/report.json for the gate parity).

[fusheng-ji]

OpenCodeReview · ci/scripts/run_qwen3_next_chain_gate.sh (L10-L10): This gate relies on pytest's exit code alone.

Not changed. test_qwen3_next_chain_report.py has no skip or importorskip, runs on CPU tensors, and imports torch and the contracts at module level, so an import failure is a collection error (exit 2) that stops the script under set -e; the provider workflow needs its skip check because its CUDA and vLLM modules can skip, this module cannot.

[fusheng-ji]

OpenCodeReview · ci/scripts/run_qwen3_next_chain_gate.sh (L22-L22): The report gate only asserts passed, which build_report can return true for a structurally valid report with zero cells (empty matrices/case_ids yields first is None, vacuous all(...) over empty detections, passed=True).

Same point as 544-S1, 549-S7, 557-S3; not changed. Under set -euo pipefail, a missing key raises KeyError and the script exits non-zero before sys.exit(0), and build_report cannot produce an empty report (it raises on empty matrices, and every cell adds invariance entries).

[fusheng-ji]

OpenCodeReview · rl_engine/backends/cuda/norm/rmsnorm.py (L373-L376): No validation that weight.shape == (x.shape[-1],) (nor that x/weight/gate are CUDA and share a dtype).

Same point as 468-S4, 503-S9, 543-S5, 549-S1; not changed. RMSNormCuda.forward uses the same asserts, and the binding repeats every check with TORCH_CHECK (weight [H], device and dtype; gate device, shape and dtype), so under python -O a bad tensor still raises before any pointer is read.

[fusheng-ji]

OpenCodeReview · rl_engine/backends/cuda/norm/rmsnorm.py (L389-L391): The ungated parameter_vjp_contributions_fp32 accepts positional x, weight, grad_output, eps=... and applies the mask, while this gated version is keyword-only (*, x, weight, gate, grad_output, eps).

Not changed: the ungated hook is also keyword-only and takes no mask, and the only caller, the C4 harness, passes keywords.

[fusheng-ji]

OpenCodeReview · rl_engine/backends/cuda/norm/rmsnorm.py (L365-L368): _require_cuda_gated reuses the module-wide _RMSNORM_API_VERSION that guards the already-shipped ungated RMSNorm surface.

Not changed. The plain and gated bindings ship in the same .so with one rmsnorm_api_version attribute, so any bump covers both. The gated symbols exist only in builds of this code, _require_cuda_gated checks both symbols and the version, and a signature change would also raise a pybind TypeError.

[fusheng-ji]

OpenCodeReview · rl_engine/backends/cuda/norm/rmsnorm.py (L280-L285): The forward saves only x, weight, gate, rstd, but backward calls rstd.reshape(x.shape[:-1]) and dgate uses rstd.reshape(x.shape[:-1]).unsqueeze(-1).

Same point as 503-S7, 549-S2, 557-I4383087967; not changed. RMSNormGatedCuda.forward and the binding both require a 2-D x, so backward always sees a [T] rstd; the op flattens N-D input first, and the harness passes 2-D x. test_cuda_gated_parameter_vjp_contributions_match_the_fold covers N-D x for the hook.

[fusheng-ji]

OpenCodeReview · rl_engine/backends/cuda/norm/rmsnorm.py (L389-L390): parameter_vjp_contributions_fp32 returns unreduced per-row contributions (rows), not the folded dweight, unlike the plain RMSNormCudaOp.parameter_vjp_contributions_fp32 which does the same.

Same point as 469-S4, 507-S1, 544-S4; not changed. Both hooks return per-row contributions on purpose: the C4 harness keys them per logical token, checks the row count and folds them itself in fixed fp32 order. It always passes 2-D x; an N-D row count would raise, not pass. The plain hook is covered for (2, 3, 2048) and the gated one for N-D x.

[fusheng-ji]

OpenCodeReview · rl_engine/backends/cuda/norm/rmsnorm.py (L404-L408): _gate_activation_fp32 computes act(gate) with torch.nn.functional.silu/torch.sigmoid in fp32, while the CUDA kernel's gate_activation is 1/(1+expf(-z)) (and z*sigma for silu).

Same point as 507-S3, 507-S2; not changed. The hook and the backward share _gate_activation_fp32, so their dweight rows agree bit for bit (checked by the fold test); the kernel's activation only feeds y and dx, and it uses expf, not __expf.

[fusheng-ji]

OpenCodeReview · pyproject.toml (L21-L21): The base dependency was loosened to transformers>=5.13.1,<=5.17.0, while the new qwen3-next extra pins transformers==5.17.0.

Same point as 549-S3, 543-S3; not changed. pip merges the base and extra specifiers, so .[qwen3-next] installs exactly transformers 5.17.0, and the preflight enforces the audited runtime.

[fusheng-ji]

OpenCodeReview · pyproject.toml (L37-L40): torch==2.13.0 / vllm==0.30.0 / triton==3.7.1 are hard-pinned, but the pre-existing vllm = ["vllm>=0.6.0"] extra still allows an unbounded vllm version.

Same point as 543-S4; not changed. >=0.6.0 and ==0.30.0 intersect at 0.30.0, so requesting both extras resolves to the audited version.

[fusheng-ji]

One OpenCodeReview summary item:

1. run_qwen3_next_chain_gate.sh L26, unreadable report path: fixed in 4e2c33af, which reaches this PR through the merge 30c0e5a7; a missing or unreadable file now exits 2 with invalid chain report: cannot read .... No other change in this PR.

[github-actions]

This 28-module qwen3_next block was appended to the "Run Mocked Hardware Discovery Tests" step, which does not build the rl_engine._C extension and is CPU-only. Several of the newly listed modules (e.g. test_qwen3_next_imported_binary.py, test_qwen3_next_installed_provider_identity.py, test_qwen3_next_storage.py, test_qwen3_next_training_checkpoint.py, test_qwen3_next_l2.py) exercise compiled-extension, vLLM-import or on-disk binary/provider evidence paths. Any module that imports those at module scope without pytest.importorskip/skip markers will fail collection here, turning the required CI job red instead of skipping. Please confirm each newly added module is import-safe on a runner without the extension/vLLM (or split these into their own step with the appropriate skip guard).

[Flink-ddd]

LGTM, Thank for you push it. also cc @inaniloquentee