[fusheng-ji]

RFC #428 stack 3/4 (base test-qwennext, merge bottom-up): #466 (merged) → #467 → #468 → #469. This PR contains #467's changes, so GitHub also shows them until #467 merges.

Claim levels (what this branch's tests and gates assert; GDN gated RMSNorm row of #428 §7.2, CUDA track):

• L1 for Qwen3NextRMSNormGatedCudaOp: prefix slices in pytest (bf16, H = 128, batch ∈ {1, 2, 8, 16, 32, 48, 64, 512} of 512 rows, one seed) and for dweight (fp32 inputs, batch ∈ {1, 8, 64, 512}); batch size, chunking, padding and permutation in the C3/C4 gates (cuda_bf16, hidden 2048 / head-dim 128, one seed). The same gates cover [WS1][CUDA][Qwen3-Next] C1: zero-centred RMSNorm references and CUDA kernel (RFC #428 stack 2/4) #467's zero-centred CUDA op.
• The gated forward's rstd is bitwise identical to the plain kernel's for the same x, asserted over 36 configurations.
• L0 is not separately tested for the gated CUDA op. L2 is not claimed.
• Not bitwise vs vLLM. The CUDA op was not compared against vLLM; the vLLM-convention PyTorch reference differs from vLLM's two eager paths on 6/40 and 18/40 seeds (committed check), and vLLM's paths are not bitwise equal to each other. The CUDA op and the PyTorch reference use different fixed-order reductions and are compared only under the contract's reduction row.

Evidence location: measurement outputs (JSON, figures, logs) are not part of this PR's diff. They are on the fork's pr_assets branch (links fixed to one commit). At the head: B200, H100. Older: pr468/ (pre-rewrite B200 charts and reports), b200-pr468-b200 and h100-pr468-h100 (1e2eaec0). Paths such as reports/experiments/... below name where a runner writes its output.

What

y = x * rstd * (weight_offset + weight) * act(gate)

• The gated kernel does every multiply in fp32 with one cast at the store. This matches vLLM's RMSNormGated with norm_before_gate=True and group_size=None, which is how the GDN block constructs it. The op has no parameter for either, so other values are not implemented rather than rejected.
• The activation is silu (swish is an alias) or sigmoid, fixed at construction and validated before the extension check. The registry constructs silu; vLLM takes it from config.output_gate_type (default silu), so the registered op matches the released config, not every config. The kernel uses expf, not __expf; a KERNEL_ALIGN_USE_FAST_MATH=1 build (off by default) turns it back into the intrinsic.
• rstd copies the plain kernel's sum-of-squares loop with the same block_reduce_sum and choose_threads(H) launch shape, so for a fixed H the reduction tree does not depend on the row count (what L1 rests on). The 36-case identity test spans fp32/fp16/bf16 × H ∈ {128, 2048, 5120} × {silu, sigmoid} × offset ∈ {0, 1}.
• The backward is deterministic. dx is a new kernel (the plain dx with (w + offset) replaced by (w + offset) * act(z)); dweight reuses rmsnorm_dweight_rows_fp32 and the ascending-row fp32 left fold through one shared _fold_dweight_rows; dgate is row-local, reduction-free fp32. parameter_vjp_contributions_fp32 takes rstd from _C.rmsnorm_gated_forward, the statistic the backward uses (pinned by test_cuda_gated_parameter_vjp_contributions_match_the_fold).
• A stale extension is rejected. The constructor and the autograd forward share _require_cuda_gated(), which checks both gated symbols and rmsnorm_api_version == 2 and raises RuntimeError (not an assert), so the registry falls through to the PyTorch reference on a build without them.
• Registration. rms_norm_gated and qwen3_next_rms_norm are added to OP_SPECS (each with a cuda-sm90 candidate), operator_inputs.py, the registry (CUDA first, PyTorch reference elsewhere) and the WS1 registered-ops set. This is the first branch where check_operator --op accepts their names.
• Changes to existing code. csrc/bindings/ops.cpp now TORCH_CHECKs the plain kernels' inputs; fp64 is now rejected by rmsnorm_forward and rmsnorm_backward_dx (on main they dispatched it). The USE_ROCM guard around the c10/cuda includes in rmsnorm.cu is removed and the launchers use at::cuda::OptionalCUDAGuard as activation.cu does; the ROCm build has not been run. The shared C3/C4 gate scripts gain --manifest and adapters can be model-scoped; the Qwen3-Next norm manifest is a norm-only workload and its validator raises WorkloadError for any malformed manifest.
• Question for the maintainers (also asked in [WS1][CUDA][Qwen3-Next] C1: zero-centred RMSNorm references and CUDA kernel (RFC #428 stack 2/4) #467): which source of truth should the CUDA strict profile use for forward arithmetic? See Scope.

Prior art & reuse decision

Measured with tools/validation/models/qwen3_next_norm_reuse_check.py, committed in this PR; rms_norm_gated.json was written from a clean tree on an otherwise idle B200 (recorded as abb56c2, now 6ac74cd8). Batch invariance is bitwise over every row alone vs inside full batches of three sizes, the full 262144-row batch vs sub-batches of 4097, 65536 and 131072 rows that cover every row and of 1 and 7 rows sampled at a stride of 512 rows, and a dense batch-size sweep. The reuse reports re-run at the head (B200 and H100, see Results) cover every row for every sub-batch size and also find the rl-kernel op batch-invariant. Accuracy is against FP64; the C3/C4 column runs this repository's gates with the candidate swapped for each library, judging the library's own dweight (rerun at 339bc3ff, rms_norm_gated_gates.json).

Implementation	Batch-invariant	Forward correctly rounded	Worst grad err	Forward / forward + backward	C3/C4 gates
rl-kernel Qwen3NextRMSNormGatedCudaOp	yes	99.9989%	2.4e-03	235 µs / 115.3 ms	pass
rl-kernel PyTorch reference	yes	99.9990%	2.4e-03	772 µs / 1991 µs	—
transformers 5.17.0 Qwen3NextRMSNormGated	yes	65.4748%	5.6e-03	695 µs / 2093 µs	C3 pass, C4 fail
FLA 0.5.2 layernorm_gated.rmsnorm_fn	yes	99.9987%	2.4e-03	177 µs / 882 µs	C3 pass, C4 fail
FLA 0.5.2 fused_norm_gate.rms_norm_gated	no (247 sub-batches differ)	99.9987%	2.4e-03	93 µs / 865 µs	C3 pass, C4 fail
vLLM 0.30.0 RMSNormGated.forward_cuda	no (3 rows differ; 3 sub-batches differ)	99.9988%	—	87 µs / —	—

Decision: write our own; no existing implementation passes C4. FLA's layernorm_gated.rmsnorm_fn is the closest: batch-invariant on every check, vLLM's arithmetic (99.9987% correctly rounded, like this kernel), a backward, and 882 µs vs 115 ms forward + backward. It passes C3 but fails C4 on dweight under chunked prefill and singleton-aggregate, as do transformers and FLA's fused_norm_gate. Those configurations add dweight across calls, and each library returns it in BF16. This kernel passes because it declares FP32 per-row dweight contributions that are folded in a fixed order. That fold is also why its backward is slow. vLLM's forward_cuda is not batch-invariant and has no backward. An earlier version of this section said FLA "passes the same gates"; that run judged rl-kernel's dweight for every swapped library (review r4379475184).

Comparison charts

B200, head dim 128, BF16	rl-kernel CUDA	PyTorch reference	transformers (cast-first)	vLLM RMSNormGated*
rows differing alone vs in a batch, forward / dx,dgate (of 768)	0 / 0	0 / 0	0 / 0	0 / n/a
forward elements correctly rounded vs the golden	99.999%	99.999%	65.6%	99.999%
forward, 262144 rows	235 µs	774 µs	697 µs	81 µs
backward, 262144 rows	114 ms	1.4 ms	1.5 ms	n/a

* forward only. Written by qwen3_next_norm_evidence.py on an otherwise idle B200 (recorded as 822b085, now ff271d5d); the same report re-measures the zero-centred op. Known cost: the CUDA backward is about 75× slower than transformers at 262144 rows, for the same reason as #467 (a serial ascending dweight fold that meets the singleton-aggregate check bitwise); FLA's rmsnorm_fn passes the same gate without one.

Results

B200 and H100 are separate data sets; per RFC #428 §2.2 no claim is made across cards. Both cards were run at the head 5aa9aa1c on an otherwise idle GPU from a clean tree. 1e2eaec0..5aa9aa1c changed only the Python guard _require_cuda_gated, an rstd reshape in the gated backward, a WorkloadError wrapper for malformed norm manifests, comments, and a comment in the WS1 gate script; no kernel source.

B200 (sm_100)

One NVIDIA B200 (driver 580.126.20), torch 2.13.0+cu130, vLLM 0.30.0, transformers 5.17.0, FlashInfer 0.6.18.post1; rl_engine._C checked to load from the tree under test; Transformer Engine, Megatron-core, Liger and FLA not installed, so those comparison rows are absent. Files: b200-pr468-head-5aa9aa1c-idle-b200.

check	result at 5aa9aa1c
unit suite (see Tests)	307 passed, 0 failed, 73 skipped (72 Ascend rmsnorm_ascend not compiled, 1 needs two CUDA devices)
check_operator rms_norm / rms_norm_gated / qwen3_next_rms_norm	pass_rate=1.0000 each
C3/C4 gates, four commands	68 invariance pairs (12 + 12 + 16 + 28), all passed=True and bitwise (max_abs=0)
norm evidence report, both comparison reports	ran, git_dirty: false; qwen3_next_rms_norm: rl-kernel CUDA batch-invariant, reference / transformers / vLLM not, FlashInfer yes; rms_norm_gated: rl-kernel CUDA, reference and transformers batch-invariant, vLLM not. Against the 1e2eaec0 reports only timings, commit ids and paths differ

H100 (sm90)

One NVIDIA H100-94GB (95830 MiB, driver 580.173.02), torch 2.13.0+cu130, vLLM 0.30.0, transformers 5.13.1, FlashInfer 0.6.18.post1; extension loaded from the tree under test; Liger and FLA not installed; Transformer Engine 2.18.0 ran for qwen3_next_rms_norm (not batch-invariant); megatron-core 0.16.0rc0 was excluded by the harness's zero-weight probe, and neither library applies to the gated op. Files: h100-pr468-head-5aa9aa1c-h100; the norm report and both reuse checks were run a second time at the head into h100-pr468-h100 on 2026-10-26 UTC, with the same probe and reuse verdicts.

check	result at 5aa9aa1c
unit suite (see Tests)	307 passed, 0 failed, 73 skipped (72 Ascend not compiled, 1 needs two CUDA devices)
check_operator rms_norm / rms_norm_gated / qwen3_next_rms_norm	pass_rate=1.0000 each
C3/C4 gates, four commands	68 invariance pairs (12 + 12 + 16 + 28), all bitwise (max_abs=0)
norm evidence report	ran, git_dirty: false. qwen3_next_rms_norm: rl-kernel CUDA 0 differing forward rows and 0 differing dx rows; PyTorch reference 8 differing dx rows; transformers 1 forward and 9 dx; vLLM forward 1; FlashInfer forward 0. rms_norm_gated: rl-kernel CUDA, reference and transformers 0 differing rows; vLLM forward 0 on the probe and 3 differing rows at batch 4097 in the reuse sweep

Earlier measurements (historical, superseded by the tables above)

• B200, shared GPU, at 5aa9aa1c (b200-pr468-head-5aa9aa1c-b200): unit suite 307 / 0 / 73, check_operator and the four C3/C4 gates as above; the GPU was shared with other jobs, so the idle-GPU run replaces it. Its C3/C4 invariance lines are identical to the idle run.
• B200, idle, at 1e2eaec0 (b200-pr468-b200, five commits before the head): 300 / 0 / 73; norm report and comparisons ran; check_operator 3 of 3 pass_rate=1.0000.
• H100, at 1e2eaec0 (in h100-pr468-h100): unit suite and check_operator 300 / 0 / 73, three ops pass_rate=1.0000.
• Comparison tables above: recorded under the pre-rewrite ids abb56c2, a66493c, 88f59f7, 822b085; now 6ac74cd8, 4f92dfc0, da6aa8db, ff271d5d (same code).

Pre-migration (main 32b765e vs this branch 0cceb84, same B200 type, driver 580.126.20, triton 3.7.1, Python 3.12.14, extension built for sm_100; failure sets compared by test ID):

check	main	0cceb84
full suite	17 failed, 2942 passed, 923 skipped	15 failed, 3103 passed, 923 skipped: no new failures, the same 2 fixed
C3/C4 gates (4 commands)	—	all passed=True; 72 invariance and singleton-aggregate comparisons, all bitwise (max_abs=0) across batch size, chunk, padding and permutation
device guard on a second GPU	—	all six launchers (plain fwd/dx/dw, gated fwd/dx) on cuda:1 inputs write to cuda:1, bitwise equal to cuda:0; the current device is unchanged (one-off script, not in this PR)

Status

• Head a2037197 on test-qwennext, which now carries the merged [WS1][CUDA][Qwen3-Next] C1: zero-centred RMSNorm references and CUDA kernel (RFC #428 stack 2/4) #467 (merged here in 6c63f0b7) and, through it, [WS1][CUDA][Qwen3-Next] Restore CUDA dispatch for rms_norm (RFC #428 stack 1/4) #466.
• CI at the head: lint, unit-tests, docs, DCO and the OpenCodeReview bot job pass; ws1-gtest is skipped (no self-hosted runner).
• Review: the bot's findings on this PR were answered one by one in their threads; those that were valid are fixed (4a3d19f2 API-version guard, f9fa9b85 rstd reshape, d634d0bd WorkloadError for malformed manifests, 08eee506 and 5aa9aa1c comments). No human review decision yet. @inaniloquentee's C4 finding (swapped libraries were judged on rl-kernel's dweight) is fixed in 339bc3ff; the rerun and the changed reuse decision are in 789a2854.
• No known failures.

History

• 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: caedea59 rejects inputs past the kernels' 32-bit row offsets and drops the unused gate_activation_grad; a2037197 makes the norm manifest require the Dense top-level keys and valid chain semantics and reject a duplicated op, names both tensors in the gate dtype error, and adds tests for the offset-aware gated backward and N-D parameter-VJP rows. Kernel arithmetic is unchanged, so the measured results stand.
• 2026-10-26 (later): merged test-qwennext after [WS1][CUDA][Qwen3-Next] C1: zero-centred RMSNorm references and CUDA kernel (RFC #428 stack 2/4) #467 landed (6c63f0b7, no tree change). Review r4379475184: the reuse check's swapped wrappers stop inheriting parameter_vjp_contributions_fp32 (339bc3ff); C3/C4 rerun on a B200 at 339bc3ff: every existing implementation now fails C4 on dweight, so the reuse decision changes from FLA layernorm_gated.rmsnorm_fn to this kernel (789a2854).
• 2026-10-26: B200 re-run at the head 5aa9aa1c on an idle GPU (unit suite, check_operator, C3/C4, norm report and both comparison reports); Results restructured into one table per card with the older measurements collapsed; evidence links retargeted to the current pr_assets commit.
• 2026-10-25: bot review rounds (see above); the fixes touch no kernel source. B200 / H100 reruns at the then-head (1e2eaec0) are in Results; 2026-10-26 H100 norm report and reuse checks updated to 5aa9aa1c in h100-pr468-h100.
• 2026-10-10: repository refactor. Merged test-qwennext at a5a92ecf (notice); code moved to rl_engine/backends, reference, runtime, validation; tests to tests/models/qwen3_next; tools to tools/validation/{models,operators}. After the migration: 295 passed, 72 skipped in the relevant suites (B200).
• 2026-10-08: added the existing-implementation comparison (qwen3_next_norm_reuse_check.py) and the comparison charts; merged [WS1][CUDA][Qwen3-Next] C1: zero-centred RMSNorm references and CUDA kernel (RFC #428 stack 2/4) #467's latest state (1edc9d4, with bec7e7a), where the plain op uses [WS1][CUDA][Qwen3-Next] C1: zero-centred RMSNorm references and CUDA kernel (RFC #428 stack 2/4) #467's _require_cuda_rmsnorm and the gated op kept the generic symbol check, since superseded by _require_cuda_gated. At 1edc9d4 (two B200s): norm suites 263 passed (72 skipped, all Ascend), providers check 9 passed, the four C3/C4 gates pass.
• 2026-10-04: the maintainers retargeted the stack to test-qwennext; re-stacked linearly onto [WS1][CUDA][Qwen3-Next] C1: zero-centred RMSNorm references and CUDA kernel (RFC #428 stack 2/4) #467 (d03c852) with every commit signed off; patch-id and final tree identical to the pre-rebase tip. [WS1][CUDA][Qwen3-Next] Restore CUDA dispatch for rms_norm (RFC #428 stack 1/4) #466 added the device guard in the plain launchers; the duplicate guards here were reduced to the T == 0 early returns, kernel launch checks and the guards on the two gated launchers. 0cceb84 switched the launchers' guard to at::cuda::OptionalCUDAGuard, the form the ROCm build already compiles.
• Mutation checks (when the tests were added): adding +1e-4f to each square in the gated sum of squares fails the rstd test in all 36 configurations; making dgate add the offset unconditionally fails exactly test_cuda_gated_dgate_keeps_signed_zero_weight (3 dtypes); mapping swish to sigmoid fails exactly test_cuda_gated_swish_is_an_alias_for_silu.
• Retracted statements: an earlier version said no training-capable batch-invariant implementation exists, that vLLM's RMSNormGated.forward_cuda is batch-invariant, and that FLA's fused_norm_gate is batch-invariant at the workload size; the Prior art section above replaces them.
• Earlier results at 4ddbaa0 / aaf0986 (full suite 15 failed, 3098 passed, 923 skipped, no new failures; first run of the C3/C4 gates, all 72 comparisons bitwise) are superseded by the 0cceb84 row.

Tests

At head 5aa9aa1c (B200; the layout below is the post-migration one):

pip install --no-build-isolation --no-deps -e .    # MAX_JOBS=32 TORCH_CUDA_ARCH_LIST=10.0 RL_KERNEL_REQUIRE_EXT=1
pytest tests/models/qwen3_next tests/ops/norm/test_rms_norm.py tests/ops/autograd/test_vjp_fp32.py tests/runtime/test_dispatch.py tests/validation/operators/test_ws1_gtest_gpu.py -q -p no:randomly -rfs
python tools/validation/operators/check_operator.py --op rms_norm --candidate cuda --device cuda --dtype bf16 --batch 2 --seq 16 --normalized-dim 4096 --seed 123 --check-grad
python tools/validation/operators/check_operator.py --op rms_norm_gated --candidate cuda --device cuda --dtype bf16 --batch 2 --seq 16 --head-dim 128 --seed 123 --check-grad
python tools/validation/operators/check_operator.py --op qwen3_next_rms_norm --candidate cuda --device cuda --dtype bf16 --batch 2 --seq 16 --normalized-dim 4096 --seed 123 --check-grad
# C3/C4 gates, as in ci/scripts/run_ws1_gtest.sh (same flags for both ops; rms_norm_gated runs at --head-dim 128)
python tools/validation/operators/check_forward_invariance.py  --manifest rl_engine/validation/models/qwen3_next_norm_manifest.json --op qwen3_next_rms_norm --candidate cuda --backend-profile cuda_bf16 --hidden 2048 --head-dim 128
python tools/validation/operators/check_gradient_invariance.py --manifest rl_engine/validation/models/qwen3_next_norm_manifest.json --op qwen3_next_rms_norm --candidate cuda --backend-profile cuda_bf16 --hidden 2048 --head-dim 128
python tools/validation/operators/check_forward_invariance.py  --manifest rl_engine/validation/models/qwen3_next_norm_manifest.json --op rms_norm_gated --candidate cuda --backend-profile cuda_bf16 --hidden 2048 --head-dim 128
python tools/validation/operators/check_gradient_invariance.py --manifest rl_engine/validation/models/qwen3_next_norm_manifest.json --op rms_norm_gated --candidate cuda --backend-profile cuda_bf16 --hidden 2048 --head-dim 128

Check	Result
pytest (five paths above)	307 passed, 0 failed, 73 skipped
check_operator rms_norm / rms_norm_gated / qwen3_next_rms_norm	pass_rate=1.0000 each
C3/C4 gates (forward and gradient, both ops, flags above)	B200 at 5aa9aa1c: 68 invariance pairs (12 + 12 + 16 + 28), all passed=True and bitwise (max_abs=0)

The same commands at this head on one idle H100-94GB: unit suite 307 passed, 0 failed, 73 skipped; check_operator pass_rate=1.0000 for each of the three ops; C3/C4 68 invariance pairs, all bitwise (max_abs=0). Skips are environment only: 72 × rmsnorm_ascend kernel not compiled, 1 × needs two CUDA devices. The repository's GPU workflows (ws1-gtest, provider) are skipped without a self-hosted runner, so the B200 / H100 runs above are the GPU test evidence.

Scope / Known limitations

• CUDA only; the ROCm build has not been run. Static check only: PyTorch's hipify translates rmsnorm.cu without error, and the headers and launch-check macro are already used unguarded by sources the ROCm build compiles.
• Question for the maintainers (RFC [RFC][Qwen3-Next-80B-A3B-Instruct][CUDA/ROCm] VIME/vLLM operator-level train-rollout consistency roadmap, ablation matrix, and integration plan #428 §0 item 1): which source of truth should the CUDA strict profile use for forward arithmetic? vLLM has three numerical paths for this norm and none is bitwise equal to another: eager forward_native, FLA forward_cuda (enforce_eager), and forward_native compiled by Inductor (default mode). The RFC's profile table points to a shared vLLM forward; the Qwen3 dense strict profile instead patches vLLM's RMSNorm to PyTorch eager arithmetic. The answer also decides which gated convention is the strict default (vLLM's fp32 weight multiply vs transformers' cast-first). VLLM_BATCH_INVARIANT=1 does not reach this norm in either mode (one-off probe, not committed).
• Norm manifest provenance. qwen3_next_norm_manifest.json pins Qwen/Qwen3-Next-80B-A3B-Instruct at 9c7f2fbe… with 41 shard SHA-256s. The validator checks internal consistency only (hex format, total size, content_hash); the values are reproducible from the public Hugging Face LFS metadata at that revision. The gates never load weights (synthetic inputs at the checkpoint's dimensions, full_model_evidence: false).
• Left as they are (none changes behavior): the gated forward duplicates the plain kernel's sum-of-squares loop, so the bitwise rstd guarantee rests on the 36-case test rather than shared code (a shared device function is a follow-up); qwen3_next_workload.py imports six private validators from ws1_workload and back; validate_norm_dimensions hard-codes 2048 / 128 already present in NORM_OPS; the Triton op's parameter-VJP hook still recomputes rstd with mean(-1) as on main.

Summary by CodeRabbit

• New Features

  • Added Qwen3-Next zero-centred and gated RMSNorm operations, with CUDA and PyTorch support, configurable weight offsets, and SiLU or sigmoid gating.
  • Added fallback to the PyTorch implementation when CUDA support is unavailable.
  • Added operator documentation and workload-specific validation.
  • Added benchmark and comparison reports for accuracy, invariance, and performance.

• Bug Fixes

  • Strengthened input validation and handling of empty batches for RMSNorm operations.

• Tests

  • Expanded coverage for forward and gradient behavior, backend dispatch, invariance, and compatibility with reference implementations.

[coderabbitai]

Walkthrough

📝 Walkthrough

Walkthrough

The change adds zero-centred and gated RMSNorm operators for Qwen3-Next, with CUDA and PyTorch implementations and backend dispatch. It also adds a pinned norm-only workload manifest, validation, invariance checks, tests, benchmark tooling, evidence reports, and operator documentation.

Changes

Qwen3-Next RMSNorm

Layer / File(s)	Summary
CUDA API and kernels csrc/cuda/rmsnorm.cu, csrc/ops.cpp, rl_engine/_C.pyi, rl_engine/kernels/ops/cuda/norm/rmsnorm.py, tests/test_rms_norm.py	RMSNorm forward and input-gradient APIs accept a weight offset. Gated forward and input-gradient APIs support SiLU and sigmoid. CUDA kernels add offset-aware scaling, empty-row returns, and launch checks. The bindings declare API version 2.
Operator implementations and dispatch rl_engine/kernels/ops/pytorch/norm/*, rl_engine/kernels/registry.py, rl_engine/kernels/gtest/operator_inputs.py, rl_engine/kernels/gtest/operator_specs.py, rl_engine/tests/test_dispatch.py, tests/test_qwen3_next_norm.py, docs/operators/*, docs/.nav.yml, docs/operators/README.md	Adds PyTorch and CUDA operators for zero-centred and gated RMSNorm, plus CUDA-first dispatch with PyTorch candidates on other backends. Tests cover formulas, gradients, invariance, validation, and signed-zero behavior. Operator pages describe conventions, constraints, dispatch, and documented evidence.
Workload manifest and validation rl_engine/testing/qwen3_next_norm_manifest.json, rl_engine/testing/qwen3_next_workload.py, rl_engine/testing/ws1_workload.py, tests/test_qwen3_next_workload.py	Adds a pinned Qwen3-Next norm-only manifest and validation for model identity, fixtures, required operators, and dimensions. Manifest loading selects the scope-specific validator and can report workload metadata.
Invariance checks and gradient adapters rl_engine/kernels/gtest/gradient_adapters.py, scripts/check_forward_invariance.py, scripts/check_gradient_invariance.py, ci/run_ws1_gtest.sh, tests/test_ws1_ascend_closeout.py, tests/test_ws1_gtest_gpu.py, .github/workflows/ci.yml, .github/workflows/ws1-gtest-gpu.yml	Adds model-scoped gradient adapters and operator inputs. The invariance scripts accept workload manifests and include workload details in JSON output. CI adds the norm tests and manifest-based forward and gradient checks.
Provider checks and benchmark evidence tests/check_qwen3_next_norm_providers.py, scripts/qwen3_next_norm_evidence.py, scripts/qwen3_next_norm_reuse_check.py, scripts/plot_qwen3_next_norm_evidence.py, docs/usage/evidence/*	Adds provider comparison checks and CUDA tools that record accuracy, row invariance, gradient, and latency results. Adds benchmark reports and a plotting script.

Priority: Low

Estimated code review effort: 4 (Complex) | ~60 minutes

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant OperatorRegistry
  participant Qwen3NextRMSNormCudaOp
  participant PyTorchReference
  OperatorRegistry->>Qwen3NextRMSNormCudaOp: try CUDA candidate
  Qwen3NextRMSNormCudaOp->>OperatorRegistry: report unavailable or incompatible CUDA bindings
  OperatorRegistry->>PyTorchReference: resolve fallback candidate

Loading

Merge Risk: Low · up to eeb8b

The new kernels are not shown to be broken. The comparison tooling overstates its every-row batch-invariance coverage, and a few docs and tooling gaps remain. Fix these before relying on the published evidence.

Security Architecture Review

Security architecture risk: 🔵 Low · up to eeb8b

The expanded native-call contract retains strong tensor validation and per-call state ownership. No introduced security issue was established, but deployment exposure and device-fault recovery remain incompletely evidenced.

Retained concerns
No architecture-level concerns identified.

Security review details

Security Blast Radius

• inferred — The demonstrated new reachability is from in-process callers supplying tensors and activation choices into native GPU execution. Its evidenced exposure is the invoking process and selected GPU; the supplied relationships do not establish a remote, tenant, credential, or cross-service authority transition.

Trust Boundaries and Controls

• observed — New caller-controlled gate tensors and activation values pass through shared native checks in both forward and backward. Backward additionally checks gradient and saved-statistic compatibility. The full comparison shows stronger plain tensor checks rather than removal of an existing control.

Resilience and Maintainability Implications

• inferred — Fresh native outputs, per-invocation saved statistics, and guarded stream selection limit interference between repeated or concurrent numerical calls. The inspected transition does not expose a persistent reservation or rollback state; device-fault recovery remains outside the established source guarantees.
• observed — Reuse validation replaces an adapter loader inside a separately launched gate-check process, rather than modifying the production registry. The parent records the child's exit status, containing the replacement during normal execution or child failure. This process separation is not a credential sandbox: the child inherits the environment.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 27.16% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 243 functions across 25 files. (8 skipped…	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Title check		The title clearly and concisely identifies the main change: a CUDA gated RMSNorm kernel for the Qwen3-Next GDN block. It matches the pull request objectives and changeset.
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.

Full details: Docstring Coverage

Explanation

Docstring coverage is 27.16% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 243 functions across 25 files. (8 skipped: 8 unsupported.)

Resolution

Write docstrings for the functions missing them to satisfy the coverage threshold.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

• Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

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

[Flink-ddd]

please resolve the code conflicts first, Thank you.

[fusheng-ji]

@coderabbitai review

[coderabbitai]

✅ Action performed

Review finished.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[fusheng-ji]

OpenCodeReview · csrc/bindings/ops.cpp (L336-L336): rmsnorm_check_backward now enforces rstd.scalar_type() == torch::kFloat32, and rmsnorm_forward always allocates rstd as fp32 (torch::empty({T}, x.options().dtype(torch::kFloat32))), so the public forward/backward pair is consistent.

Not changed; this is not a behaviour change. Before, the launch already read dy.data_ptr<scalar_t>() with scalar_t dispatched from x, and rstd.data_ptr<float>(). data_ptr<T> throws on a dtype mismatch, so the new TORCH_CHECKs give the same rejection with a clearer message. rstd is always allocated as float32 by the forward bindings, and no Python code calls _C.rmsnorm_backward_dw (RMSNormCuda.backward uses rmsnorm_dweight_rows_fp32).

[fusheng-ji]

OpenCodeReview · csrc/bindings/ops.cpp (L433-L434): rmsnorm_gated_backward_dx has no mask argument, but the plain rmsnorm_backward_dw does.

Not changed: a float32 weight with a bf16 or fp16 x is supported on purpose. rmsnorm_check_weight (shared with the plain op) allows weight in x's dtype or float32, the kernels dispatch on the weight dtype separately, and the Python backward uses weight.float() and casts dweight back to weight.dtype. gate and dy are checked against x's dtype in the binding.

[fusheng-ji]

OpenCodeReview · rl_engine/backends/cuda/norm/rmsnorm.py (L299-L299): RMSNormGatedCuda.backward saves rstd via save_for_backward and passes it straight to _C.rmsnorm_gated_backward_dx, whose binding now requires rstd to be contiguous 1-D fp32 of length T.

Not changed. Autograd casts the incoming gradient to the dtype of the forward output before backward runs, and y is allocated as empty_like(x), so grad_out always arrives in x's dtype and passes the binding's dy check. RMSNormCuda.backward relies on the same.

[fusheng-ji]

OpenCodeReview · rl_engine/backends/cuda/norm/rmsnorm.py (L284-L284): RMSNormGatedCuda.forward guards with bare assert (CUDA-ness, contiguity, rank, gate shape).

Not changed. The extension check is _require_cuda_gated(), which raises RuntimeError (r4372677163). The asserts that remain are not the only guard: the binding repeats each check with TORCH_CHECK (rmsnorm_gated_check in csrc/bindings/ops.cpp: CUDA, contiguous, 2-D x, weight [H], device and dtype of weight and gate), so under python -O a bad tensor still raises before any pointer is read. The plain RMSNormCuda.forward uses the same pattern.

[fusheng-ji]

OpenCodeReview · rl_engine/reference/norm/qwen3_next_rms_norm.py (L117-L118): Qwen3NextRMSNormGatedOp._normalized now rejects a gate whose dtype differs from x.

Not changed: the reference now has the CUDA contract, which rejects gate.dtype != x.dtype (test_cuda_gated_rejects_mismatched_gate). The only model caller passes the GDN z in BF16 like x, and the goldens go through forward_fp32 with matching dtypes.

[fusheng-ji]

OpenCodeReview · rl_engine/backends/cuda/norm/rmsnorm.py (L361-L363): Qwen3NextRMSNormGatedCudaOp declares weight_offset = 0.0 (plain weight), but its backward (dweight, dgate) and parameter_vjp_contributions_fp32 honor an arbitrary weight_offset.

Not changed: the offset is applied wherever it enters a derivative. dweight has no offset term, because d/dw (w + c) = 1 (the plain op notes the same); dgate uses weight.float() + weight_offset, and the dx and forward kernels add the offset in fp32. test_cuda_gated_backward_is_offset_aware (added in a203719 on #468) checks that dx and dgate with weight_offset=1 equal those with an explicit 1 + w weight, and that dweight equals the plain-weight dweight.

[fusheng-ji]

OpenCodeReview · csrc/bindings/ops.cpp (L437-L440): rmsnorm_gated_backward_dx validates dy/rstd with rmsnorm_check_input (contiguity) but the saved gate and x are only checked in rmsnorm_gated_check, which re-uses rmsnorm_check_input(gate, ...); however rmsnorm_backward_dw (changed in this diff) also gained a mask.device() and positive-H check, while rmsnorm_gated_backward_dx never checks dy/gate contiguity after the Python backward reuses ctx.saved_tensors (raw saved gate/x, not the contiguous copies made in the wrapper's forward).

Not changed: rmsnorm_gated_backward_dx calls rmsnorm_gated_check, which runs rmsnorm_check_input (CUDA and contiguous) on x, weight and gate before dy and rstd are checked, so a non-contiguous saved tensor raises before any pointer is read.

[fusheng-ji]

OpenCodeReview · rl_engine/validation/models/qwen3_next_workload.py (L52-L52): _validate_model_identity is called before _validate_stochastic_policy, but the model-identity check is only reached after the bare subscripts identity["model_id"] / identity["revision"].

Not changed. _validate_model_identity checks every pinned FINGERPRINT key against the manifest, as the Dense path does, and the committed fingerprint has exactly those 13 keys. A key beyond them is not left unchecked either: the runtime's verify_checkpoint compares every manifest fingerprint key with the checkpoint's config.json.

[fusheng-ji]

OpenCodeReview · rl_engine/config/workload.py (L289-L289): workload_report reads manifest.raw["fixture_identity_sha256"] with a subscript.

Same point as 543-I4381854144, 543-I4381854154; not changed. validate_norm_manifest reads fixture_identity_sha256 and compares it with the recomputed identity hash, and raises WorkloadError when it is missing, so every manifest that reaches workload_report through the norm branch has it.

[fusheng-ji]

OpenCodeReview · rl_engine/validation/operators/operator_specs.py (L61-L62): For qwen3_next_rms_norm the gold_path/pytorch candidate is Qwen3NextRMSNormOp, which is a NativeRMSNormOp subclass applying 1 + weight.

Not changed. The C3/C4 runners build the norm parameters with _row_parameters, not the shared _make_rms_norm_inputs, and the convention does not need near-zero weights to show: (1 + w)·n and w·n differ by n for any w. It is pinned by test_zero_weight_is_identity_scaling, test_weight_offset_is_one, test_cuda_offset_equals_explicit_fp32_weight and test_cuda_zero_centred_backward_is_offset_aware.

[fusheng-ji]

OpenCodeReview · rl_engine/validation/operators/operator_specs.py (L68-L69): The cuda entry here is a bare string while the sibling cuda-sm90 entry is written as a parenthesized string ("..."), which is also just a string, not a 1-tuple.

Not changed: every candidate_paths value is a str. ("...") without a trailing comma is a parenthesised string, not a tuple, and OperatorSpec.candidate_paths is typed dict[str, str]. Every consumer treats the value as a dotted path (_load_object does path.rsplit(".", 1)); the parentheses only wrap long paths, as for other entries in the same table.

[fusheng-ji]

OpenCodeReview · rl_engine/reference/norm/qwen3_next_rms_norm.py (L117-L118): The dtype/device guard raises the same ValueError type as the preceding shape checks but omits the offending values, unlike the shape error which prints both shapes.

Fixed in a203719 on #468: the error now names both tensors' dtype and device; test_gated_dtype_mismatch_names_both_dtypes checks it.

[fusheng-ji]

OpenCodeReview · rl_engine/reference/norm/qwen3_next_rms_norm.py (L117-L122): This new strict gate dtype/device check tightens the reference contract, but the CUDA candidate this reference is compared against does not enforce the same thing at the op level: Qwen3NextRMSNormGatedCudaOp.forward only validates gate.shape == x.shape and delegates to rmsnorm_gated_cuda, so a dtype-mismatched gate surfaces as an AssertionError from RMSNormGatedCuda.forward (or, in parameter_vjp_contributions_fp32, may reach the kernel with no Python-side dtype check at all).

Same point as 468-S5, 469-S1; not changed. Both sides reject a gate whose dtype differs from x: the reference raises ValueError, and the CUDA op reaches the binding's TORCH_CHECK on the gate dtype, which raises RuntimeError (test_cuda_gated_rejects_mismatched_gate). There is no dtype assert in RMSNormGatedCuda.forward that would behave differently.

[fusheng-ji]

OpenCodeReview · rl_engine/backends/cuda/norm/rmsnorm.py (L308-L309): dgate/dweight use the kernel rstd, which is bit-identical to the plain kernel's for the same x (as the forward comment claims), yet this code computes act(gate) in fp32 via torch.nn.functional.silu/sigmoid while the kernel uses 1.0f / (1.0f + expf(-z)).

Agreed, no change: the gated kernel computes rstd with the plain kernel's loop, and test_cuda_gated_rstd_is_bitwise_identical_to_plain_kernel pins the equality across dtype, width, activation and offset.

[fusheng-ji]

OpenCodeReview · rl_engine/validation/models/qwen3_next_workload.py (L59-L59): _validate_model_identity is called with a caller-supplied fingerprint mapping.

Same point as 468-S3; not changed. Both callers pass module constants: the norm validator passes the 13-key FINGERPRINT and the Dense path uses the official default, so no caller can pass an empty mapping today.

[fusheng-ji]

OpenCodeReview · rl_engine/config/workload.py (L289-L290): workload_report dereferences manifest.raw["fixture_identity_sha256"] and manifest.model_identity["model_id"] with bare subscripting.

Same point as 468-S4, 543-I4381854144; not changed. Every manifest that workload_report sees was accepted by load_manifest, and both validators now require the same top-level keys (_REQUIRED_TOP_LEVEL, which includes workload_id, model_identity and fixture_identity_sha256), so these lookups cannot fail.

[fusheng-ji]

OpenCodeReview · rl_engine/validation/models/qwen3_next_workload.py (L70-L71): caps["required_chain_ops"] comparison only sorts e["op"], so a typo in another capability field would surface as a KeyError wrapped into a generic "malformed" message.

Not changed: the wrapper re-raises with {exc!r}, so the WorkloadError names the missing key (for example KeyError('op')) rather than a generic "malformed".

[Flink-ddd]

LGTM. The review findings have been addressed or explained, and the corrected reuse comparison and Hopper validation look good. Ready to merge into test-qwennext.