[WS1][CUDA][Qwen3-Next] MoE route/combine contract with a batch-invariant grouped expert GEMM (RFC #428 stack 5/8)

.github/workflows/ci.yml

@@ -5,9 +5,9 @@ name: CI-Pipeline
 
 on:
   push:
-    branches: [ main, test ]
+    branches: [ main, test, test-qwennext ]
   pull_request:
-    branches: [ main, test ]
+    branches: [ main, test, test-qwennext ]
 
 permissions:
   contents: read
@@ -75,7 +75,7 @@ jobs:
         run: |
           python -m pytest tests/runtime/test_dispatch.py -v
           PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 python -m pytest tests/ops/attention/test_attention_correctness.py -q -rs
-          python -m pytest tests/validation/operators/test_forward_invariance.py tests/contracts/test_tolerance_contract.py tests/validation/operators/test_ws1_workload.py tests/validation/operators/test_gradient_invariance.py tests/validation/operators/test_elementwise_inventory.py tests/validation/operators/test_four_judgment_matrix.py tests/validation/operators/test_op_checks.py tests/validation/operators/test_operator_inputs.py tests/benchmarks/test_profiler.py tests/validation/operators/test_kv_consistency.py tests/models/qwen3/test_ws1_qwen3_dense.py tests/validation/models/test_ws1_chain_integration.py -q
+          python -m pytest tests/validation/operators/test_forward_invariance.py tests/contracts/test_tolerance_contract.py tests/validation/operators/test_ws1_workload.py tests/validation/operators/test_gradient_invariance.py tests/validation/operators/test_elementwise_inventory.py tests/validation/operators/test_four_judgment_matrix.py tests/validation/operators/test_op_checks.py tests/validation/operators/test_operator_inputs.py tests/benchmarks/test_profiler.py tests/validation/operators/test_kv_consistency.py tests/models/qwen3/test_ws1_qwen3_dense.py tests/validation/models/test_ws1_chain_integration.py tests/models/qwen3_next/test_qwen3_next_norm.py tests/models/qwen3_next/test_gdn_state_contract.py tests/validation/common/test_tensor_identity.py tests/models/qwen3_next/test_qwen3_next_forward_contract.py tests/models/qwen3_next/test_qwen3_next_tp_blocks.py -q
 
       - name: Run Cross-Configuration Contract Tests (CPU-safe)
         run: |

.github/workflows/qwen3-next-provider-gpu.yml

@@ -0,0 +1,90 @@
+# SPDX-License-Identifier: Apache-2.0
+# RFC #428 Qwen3-Next provider comparisons: the C1/C3/C4 norm providers, the C6
+# GDN decode-step goldens and the shared GEMM/MoE route-combine provider, each
+# checked against the real vLLM 0.30.0 provider.
+#
+# These check_ files import vLLM, so the default pytest collection never runs them
+# (tests/integrations/common/test_framework_operator_integrations.py asserts vLLM is not imported). This
+# workflow runs them in their own process on a self-hosted runner that a maintainer
+# registers with the labels below and keeps on the audited runtime (CUDA 13,
+# torch 2.13.0+cu130, vLLM 0.30.0). Without such a runner the job queues rather than
+# reporting a false pass.
+#
+# Security: do not use pull_request_target. Fork PRs never reach the self-hosted
+# runner; a maintainer dispatches the reviewed commit from a trusted branch.
+
+name: Qwen3-Next-provider-GPU
+
+on:
+  pull_request:
+    branches: [main, test-qwennext]
+    paths:
+      - "rl_engine/backends/**"
+      - "rl_engine/reference/**"
+      - "rl_engine/models/qwen3_next/**"
+      - "rl_engine/integrations/engines/train/vllm/qwen3_next*"
+      - "rl_engine/validation/models/qwen3_next*"
+      - "tools/validation/models/*qwen3_next*"
+      - "tests/models/qwen3_next/**"
+      - "tests/models/qwen3_next/check_gdn_recurrent_golden.py"
+      - ".github/workflows/qwen3-next-provider-gpu.yml"
+  workflow_dispatch:
+
+concurrency:
+  group: qwen3-next-provider-gpu-${{ github.ref }}
+  cancel-in-progress: false
+
+permissions:
+  contents: read
+
+jobs:
+  fork-pr-notice:
+    if: github.event_name == 'pull_request' && github.event.pull_request.head.repo.full_name != github.repository
+    runs-on: ubuntu-latest
+    steps:
+      - name: Report required trusted execution
+        run: |
+          echo "Fork code does not run on the self-hosted Qwen3-Next provider runner."
+          echo "A maintainer must dispatch this workflow from a trusted upstream branch."
+          echo "source_repository=${{ github.event.pull_request.head.repo.full_name }}"
+          echo "source_sha=${{ github.event.pull_request.head.sha }}"
+
+  provider:
+    if: github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository
+    runs-on: [self-hosted, linux, x64, rl-kernel-qwen3-next]
+    timeout-minutes: 30
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          persist-credentials: false
+      - name: Require the audited runtime
+        run: |
+          python3 - <<'PY'
+          import torch, vllm
+          assert torch.__version__ == "2.13.0+cu130"
+          assert vllm.__version__ == "0.30.0"
+          assert torch.cuda.is_available()
+          PY
+      - name: Build the extension against that runtime
+        run: python3 setup.py build_ext --inplace
+      - name: Require the extension built from this checkout
+        # A stale editable install elsewhere on the runner would otherwise satisfy the
+        # import and test some other tree's _C.
+        run: |
+          python3 - <<'PY'
+          import os
+          from pathlib import Path
+
+          from rl_engine import _C
+
+          workspace = Path(os.environ["GITHUB_WORKSPACE"]).resolve()
+          loaded = Path(_C.__file__).resolve()
+          assert workspace in loaded.parents, f"rl_engine._C loaded from {loaded}, outside {workspace}"
+          print("rl_engine._C:", loaded)
+          PY
+      - name: Check pinned providers in their own process
+        # The FP32 router needs IEEE Triton dots from process start.
+        env:
+          TRITON_F32_DEFAULT: ieee
+        run: |
+          python3 -m pytest -q tests/models/qwen3_next/check_qwen3_next_norm_providers.py tests/models/qwen3_next/check_gdn_recurrent_golden.py tests/models/qwen3_next/check_qwen3_next_forward.py

.github/workflows/ws1-gtest-gpu.yml

@@ -10,7 +10,7 @@ name: WS1-gtest-GPU
 
 on:
   pull_request:
-    branches: [ main ]
+    branches: [ main, test-qwennext ]
     paths:
       - "rl_engine/validation/**"
       - "rl_engine/backends/**"
@@ -28,6 +28,7 @@ on:
       - "tools/validation/operators/ws1_candidate_evidence.py"
       - "tests/validation/**"
       - "tests/models/**"
+      - "tests/ops/norm/**"
       - "tests/validation/operators/test_forward_invariance.py"
       - "tests/validation/operators/test_gradient_invariance.py"
       - "tests/validation/operators/test_four_judgment_matrix.py"
@@ -36,7 +37,7 @@ on:
       - "ci/scripts/run_gpu_ci.sh"
       - ".github/workflows/ws1-gtest-gpu.yml"
   push:
-    branches: [ main ]
+    branches: [ main, test-qwennext ]
     paths:
       - "rl_engine/validation/**"
       - "rl_engine/backends/**"
@@ -54,6 +55,7 @@ on:
       - "tools/validation/operators/ws1_candidate_evidence.py"
       - "tests/validation/**"
       - "tests/models/**"
+      - "tests/ops/norm/**"
       - "tests/validation/operators/test_forward_invariance.py"
       - "tests/validation/operators/test_gradient_invariance.py"
       - "tests/validation/operators/test_four_judgment_matrix.py"

ci/scripts/run_ws1_gtest.sh

@@ -19,6 +19,7 @@ echo "[ws1-gtest] interpreter=$PY out=$OUT"
 
 "$PY" -m pytest -q \
   tests/validation/operators/test_ws1_gtest_gpu.py \
+  tests/models/qwen3_next/test_qwen3_next_norm.py \
   tests/ops/attention/test_triton_batch_invariant_attention.py \
   tests/validation/operators/test_four_judgment_matrix.py \
   tests/validation/operators/test_ws1_candidate_evidence.py \
@@ -75,3 +76,14 @@ for cell in required:
         )
 print("[ws1-gtest] C8 gate passed")
 PY
+
+echo "[ws1-gtest] Qwen3-Next C3/C4 norms (operator scope)"
+QWEN3_NEXT_NORM_MANIFEST=rl_engine/validation/models/qwen3_next_norm_manifest.json
+"$PY" tools/validation/operators/check_forward_invariance.py --manifest "$QWEN3_NEXT_NORM_MANIFEST" \
+  --op qwen3_next_rms_norm --candidate cuda --backend-profile cuda_bf16 --hidden 2048 --head-dim 128
+"$PY" tools/validation/operators/check_gradient_invariance.py --manifest "$QWEN3_NEXT_NORM_MANIFEST" \
+  --op qwen3_next_rms_norm --candidate cuda --backend-profile cuda_bf16 --hidden 2048 --head-dim 128
+"$PY" tools/validation/operators/check_forward_invariance.py --manifest "$QWEN3_NEXT_NORM_MANIFEST" \
+  --op rms_norm_gated --candidate cuda --backend-profile cuda_bf16 --hidden 2048 --head-dim 128
+"$PY" tools/validation/operators/check_gradient_invariance.py --manifest "$QWEN3_NEXT_NORM_MANIFEST" \
+  --op rms_norm_gated --candidate cuda --backend-profile cuda_bf16 --hidden 2048 --head-dim 128

csrc/bindings/ops.cpp

@@ -263,14 +263,36 @@ void rmsnorm_forward_cuda(
   torch::Tensor weight,
   torch::Tensor y,
   torch::Tensor rstd,
-  double eps);
+  double eps,
+  double weight_offset);
 
 void rmsnorm_backward_dx_cuda(
   torch::Tensor dy,
   torch::Tensor x,
   torch::Tensor weight,
   torch::Tensor rstd,
-  torch::Tensor dx);
+  torch::Tensor dx,
+  double weight_offset);
+
+void rmsnorm_gated_forward_cuda(
+  torch::Tensor x,
+  torch::Tensor weight,
+  torch::Tensor gate,
+  torch::Tensor y,
+  torch::Tensor rstd,
+  double eps,
+  double weight_offset,
+  int64_t activation);
+
+void rmsnorm_gated_backward_dx_cuda(
+  torch::Tensor dy,
+  torch::Tensor x,
+  torch::Tensor weight,
+  torch::Tensor gate,
+  torch::Tensor rstd,
+  torch::Tensor dx,
+  double weight_offset,
+  int64_t activation);
 
 void rmsnorm_backward_partial_dw_cuda(
   torch::Tensor dy,
@@ -296,23 +318,39 @@ static void rmsnorm_check_input(const torch::Tensor& x, const char* name) {
   TORCH_CHECK(x.is_contiguous(), name, " must be contiguous");
 }
 
+static void rmsnorm_check_weight(const torch::Tensor& x, const torch::Tensor& weight) {
+  TORCH_CHECK(x.dim() == 2 && x.size(1) > 0, "x must be 2D with positive hidden size");
+  TORCH_CHECK(x.scalar_type() == torch::kFloat32 || x.scalar_type() == torch::kFloat16 ||
+              x.scalar_type() == torch::kBFloat16, "x must be float32, float16 or bfloat16");
+  TORCH_CHECK(weight.dim() == 1 && weight.size(0) == x.size(1), "weight must be [H]");
+  TORCH_CHECK(weight.device() == x.device(), "weight must be on the same device as x");
+  TORCH_CHECK(weight.scalar_type() == x.scalar_type() || weight.scalar_type() == torch::kFloat32,
+              "weight must have x dtype or float32");
+}
+
+static void rmsnorm_check_backward(const torch::Tensor& dy, const torch::Tensor& x,
+                                   const torch::Tensor& rstd) {
+  TORCH_CHECK(dy.device() == x.device() && rstd.device() == x.device(),
+              "dy and rstd must be on the same device as x");
+  TORCH_CHECK(dy.scalar_type() == x.scalar_type(), "dy must have the same dtype as x");
+  TORCH_CHECK(rstd.scalar_type() == torch::kFloat32, "rstd must be float32");
+}
+
 std::vector<torch::Tensor> rmsnorm_forward(
   torch::Tensor x,
   torch::Tensor weight,
-  double eps)
+  double eps,
+  double weight_offset)
 {
   rmsnorm_check_input(x, "x");
   rmsnorm_check_input(weight, "weight");
-
-  TORCH_CHECK(x.dim() == 2, "x must be 2D [T, H]");
-  TORCH_CHECK(weight.dim() == 1, "weight must be 1D [H]");
-  TORCH_CHECK(x.size(1) == weight.size(0), "x.size(1) must equal weight.size(0)");
+  rmsnorm_check_weight(x, weight);
 
   auto T = x.size(0);
   auto y = torch::empty_like(x);
   auto rstd = torch::empty({T}, x.options().dtype(torch::kFloat32));
 
-  rmsnorm_forward_cuda(x, weight, y, rstd, eps);
+  rmsnorm_forward_cuda(x, weight, y, rstd, eps, weight_offset);
 
   return {y, rstd};
 }
@@ -321,22 +359,89 @@ torch::Tensor rmsnorm_backward_dx(
   torch::Tensor dy,
   torch::Tensor x,
   torch::Tensor weight,
-  torch::Tensor rstd)
+  torch::Tensor rstd,
+  double weight_offset)
 {
   rmsnorm_check_input(dy, "dy");
   rmsnorm_check_input(x, "x");
   rmsnorm_check_input(weight, "weight");
+  rmsnorm_check_weight(x, weight);
   rmsnorm_check_input(rstd, "rstd");
+  rmsnorm_check_backward(dy, x, rstd);
 
   TORCH_CHECK(dy.sizes() == x.sizes(), "dy and x must have same shape");
-  TORCH_CHECK(x.dim() == 2, "x must be 2D [T, H]");
-  TORCH_CHECK(weight.dim() == 1, "weight must be 1D [H]");
   TORCH_CHECK(rstd.dim() == 1, "rstd must be 1D [T]");
   TORCH_CHECK(rstd.size(0) == x.size(0), "rstd.size(0) must equal x.size(0)");
 
   auto dx = torch::empty_like(x);
 
-  rmsnorm_backward_dx_cuda(dy, x, weight, rstd, dx);
+  rmsnorm_backward_dx_cuda(dy, x, weight, rstd, dx, weight_offset);
+
+  return dx;
+}
+
+static void rmsnorm_gated_check(
+  const torch::Tensor& x,
+  const torch::Tensor& weight,
+  const torch::Tensor& gate,
+  int64_t activation)
+{
+  rmsnorm_check_input(x, "x");
+  rmsnorm_check_input(weight, "weight");
+  rmsnorm_check_weight(x, weight);
+  rmsnorm_check_input(gate, "gate");
+  TORCH_CHECK(gate.device() == x.device(), "gate must be on the same device as x");
+  TORCH_CHECK(gate.sizes() == x.sizes(), "gate must have the same shape as x");
+  TORCH_CHECK(
+    gate.scalar_type() == x.scalar_type(),
+    "gate must have the same dtype as x");
+  // 0 = silu/swish, 1 = sigmoid. Anything else fails closed rather than
+  // silently computing a different activation (RFC #428 section 6, item 7).
+  TORCH_CHECK(
+    activation == 0 || activation == 1,
+    "activation must be 0 (silu) or 1 (sigmoid), got ", activation);
+}
+
+std::vector<torch::Tensor> rmsnorm_gated_forward(
+  torch::Tensor x,
+  torch::Tensor weight,
+  torch::Tensor gate,
+  double eps,
+  double weight_offset,
+  int64_t activation)
+{
+  rmsnorm_gated_check(x, weight, gate, activation);
+
+  auto T = x.size(0);
+  auto y = torch::empty_like(x);
+  auto rstd = torch::empty({T}, x.options().dtype(torch::kFloat32));
+
+  rmsnorm_gated_forward_cuda(x, weight, gate, y, rstd, eps, weight_offset, activation);
+
+  return {y, rstd};
+}
+
+torch::Tensor rmsnorm_gated_backward_dx(
+  torch::Tensor dy,
+  torch::Tensor x,
+  torch::Tensor weight,
+  torch::Tensor gate,
+  torch::Tensor rstd,
+  double weight_offset,
+  int64_t activation)
+{
+  rmsnorm_gated_check(x, weight, gate, activation);
+  rmsnorm_check_input(dy, "dy");
+  rmsnorm_check_input(rstd, "rstd");
+  rmsnorm_check_backward(dy, x, rstd);
+
+  TORCH_CHECK(dy.sizes() == x.sizes(), "dy must have the same shape as x");
+  TORCH_CHECK(rstd.dim() == 1 && rstd.size(0) == x.size(0), "rstd must be [T]");
+
+  auto dx = torch::empty_like(x);
+
+  rmsnorm_gated_backward_dx_cuda(
+    dy, x, weight, gate, rstd, dx, weight_offset, activation);
 
   return dx;
 }
@@ -350,13 +455,16 @@ torch::Tensor rmsnorm_backward_dw(
   rmsnorm_check_input(dy, "dy");
   rmsnorm_check_input(x, "x");
   rmsnorm_check_input(rstd, "rstd");
+  rmsnorm_check_backward(dy, x, rstd);
   rmsnorm_check_input(mask, "mask");
 
   TORCH_CHECK(dy.sizes() == x.sizes(), "dy and x must have same shape");
   TORCH_CHECK(x.dim() == 2, "x must be 2D [T, H]");
   TORCH_CHECK(rstd.dim() == 1, "rstd must be 1D [T]");
   TORCH_CHECK(mask.dim() == 1, "mask must be 1D [T]");
   TORCH_CHECK(mask.scalar_type() == torch::kBool, "mask must be bool");
+  TORCH_CHECK(mask.device() == x.device(), "mask must be on the same device as x");
+  TORCH_CHECK(x.size(1) > 0, "hidden size must be positive");
   TORCH_CHECK(rstd.size(0) == x.size(0), "rstd.size(0) must equal x.size(0)");
   TORCH_CHECK(mask.size(0) == x.size(0), "mask.size(0) must equal x.size(0)");
 
@@ -718,9 +826,24 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
         &det_gemm_db_transposed,
         "Batch-invariant deterministic GEMM backward in canonical [N,K] layout");
     // registry RMSNorm
-    m.def("rmsnorm_forward", &rmsnorm_forward, "Batch-invariant RMSNorm forward CUDA");
-    m.def("rmsnorm_backward_dx", &rmsnorm_backward_dx, "Batch-invariant RMSNorm backward dx CUDA");
+    // API version 2 includes weight_offset in both entry points.
+    m.attr("rmsnorm_api_version") = 2;
+    m.def("rmsnorm_forward", &rmsnorm_forward, "Batch-invariant RMSNorm forward CUDA",
+          py::arg("x"), py::arg("weight"), py::arg("eps"),
+          py::arg("weight_offset") = 0.0);
+    m.def("rmsnorm_backward_dx", &rmsnorm_backward_dx, "Batch-invariant RMSNorm backward dx CUDA",
+          py::arg("dy"), py::arg("x"), py::arg("weight"), py::arg("rstd"),
+          py::arg("weight_offset") = 0.0);
     m.def("rmsnorm_backward_dw", &rmsnorm_backward_dw, "Deterministic RMSNorm backward dweight CUDA");
+    m.def("rmsnorm_gated_forward", &rmsnorm_gated_forward,
+          "Batch-invariant gated RMSNorm forward CUDA",
+          py::arg("x"), py::arg("weight"), py::arg("gate"), py::arg("eps"),
+          py::arg("weight_offset") = 0.0, py::arg("activation") = 0);
+    m.def("rmsnorm_gated_backward_dx", &rmsnorm_gated_backward_dx,
+          "Batch-invariant gated RMSNorm backward dx CUDA",
+          py::arg("dy"), py::arg("x"), py::arg("weight"), py::arg("gate"),
+          py::arg("rstd"), py::arg("weight_offset") = 0.0,
+          py::arg("activation") = 0);
 #if !defined(USE_ROCM)
     m.def(
         "reduce_rows_fp32_left_fold",