[WS1][CUDA][Qwen3-Next] C1: zero-centred RMSNorm references and CUDA kernel (RFC #428 stack 2/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 \

csrc/bindings/ops.cpp

@@ -263,14 +263,16 @@ 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_backward_partial_dw_cuda(
   torch::Tensor dy,
@@ -299,7 +301,8 @@ static void rmsnorm_check_input(const torch::Tensor& x, const char* name) {
 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");
@@ -312,7 +315,7 @@ std::vector<torch::Tensor> rmsnorm_forward(
   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,7 +324,8 @@ 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");
@@ -336,7 +340,7 @@ torch::Tensor rmsnorm_backward_dx(
 
   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;
 }
@@ -718,8 +722,14 @@ 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");
 #if !defined(USE_ROCM)
     m.def(

csrc/cuda/norm/rmsnorm.cu

@@ -108,7 +108,8 @@ __global__ void rmsnorm_fwd_kernel(
     float* __restrict__ rstd,
     int T,
     int H,
-    float eps
+    float eps,
+    float weight_offset
 ) {
     int row = blockIdx.x;
     int tid = threadIdx.x;
@@ -135,10 +136,17 @@ __global__ void rmsnorm_fwd_kernel(
 
     __syncthreads();
 
-    // Write y = x * rstd * weight.
+    // Write y = x * rstd * (weight_offset + weight). The offset is added in
+    // fp32 after the upcast: a zero-centred weight (Qwen3-Next, Gemma) must not
+    // have its "+1" folded into the low-precision weight beforehand, which would
+    // round the offset and break the bitwise contract.
     for (int col = tid; col < H; col += blockDim.x) {
         float xv = load_as_float<scalar_t>(x_row + col);
         float wv = load_as_float<weight_t>(weight + col);
+        // Guarded: `-0.0f + 0.0f` is +0.0f, so an unconditional add would flip the
+        // sign bit of -0.0 weights on the plain path. Pinned by
+        // tests/models/qwen3_next/test_qwen3_next_norm.py::test_cuda_plain_signed_zero_is_preserved.
+        if (weight_offset != 0.0f) wv += weight_offset;
         float out = xv * row_rstd * wv;
         store_from_float<scalar_t>(y_row + col, out);
     }
@@ -153,7 +161,8 @@ __global__ void rmsnorm_bwd_dx_kernel(
     const float* __restrict__ rstd,
     scalar_t* __restrict__ dx,
     int T,
-    int H
+    int H,
+    float weight_offset
 ) {
     int row = blockIdx.x;
     int tid = threadIdx.x;
@@ -168,6 +177,7 @@ __global__ void rmsnorm_bwd_dx_kernel(
         float dyv = load_as_float<scalar_t>(dy_row + col);
         float xv = load_as_float<scalar_t>(x_row + col);
         float wv = load_as_float<weight_t>(weight + col);
+        if (weight_offset != 0.0f) wv += weight_offset;
         local_dot += dyv * wv * xv;
     }
 
@@ -180,6 +190,7 @@ __global__ void rmsnorm_bwd_dx_kernel(
         float dyv = load_as_float<scalar_t>(dy_row + col);
         float xv = load_as_float<scalar_t>(x_row + col);
         float wv = load_as_float<weight_t>(weight + col);
+        if (weight_offset != 0.0f) wv += weight_offset;
 
         float out = r * dyv * wv - xv * coeff;
         store_from_float<scalar_t>(dx_row + col, out);
@@ -247,7 +258,8 @@ void rmsnorm_forward_cuda(
     torch::Tensor weight,
     torch::Tensor y,
     torch::Tensor rstd,
-    double eps
+    double eps,
+    double weight_offset
 ) {
     // Launch on x's device: the current CUDA stream belongs to the current
     // device, which need not be x's.
@@ -270,7 +282,8 @@ void rmsnorm_forward_cuda(
                 rstd.data_ptr<float>(),
                 T,
                 H,
-                static_cast<float>(eps)
+                static_cast<float>(eps),
+                static_cast<float>(weight_offset)
             );
         });
     });
@@ -282,7 +295,8 @@ void rmsnorm_backward_dx_cuda(
     torch::Tensor x,
     torch::Tensor weight,
     torch::Tensor rstd,
-    torch::Tensor dx
+    torch::Tensor dx,
+    double weight_offset
 ) {
     const at::cuda::OptionalCUDAGuard device_guard(device_of(x));
     int T = x.size(0);
@@ -303,7 +317,8 @@ void rmsnorm_backward_dx_cuda(
                 rstd.data_ptr<float>(),
                 dx.data_ptr<x_t>(),
                 T,
-                H
+                H,
+                static_cast<float>(weight_offset)
             );
         });
     });

docs/.nav.yml

@@ -29,6 +29,7 @@ nav:
       - operators/sampling.md
       - operators/det-gemm.md
       - operators/embedding.md
+      - operators/qwen3-next-rms-norm.md
   - Developer Guide:
     - contributing/README.md
     - Contributor Guide: contributing/contributor-guide.md

docs/operators/README.md

@@ -32,4 +32,5 @@ Every operator page should include:
 - [Matmul](matmul.md)
 - [Sampling](sampling.md)
 - [Token Embedding](embedding.md)
+- [Qwen3-Next RMSNorm (zero-centred)](qwen3-next-rms-norm.md)
 - [Operator Doc Template](../contributing/operator-doc-template.md)