[WS1][CUDA][Qwen3-Next] Gated RMSNorm CUDA kernel for the GDN block (RFC #428 stack 3/4)

.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 -q
 
       - name: Run Cross-Configuration Contract Tests (CPU-safe)
         run: |

.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",