Merge branch 'main' into feat/issue-145

Author: Flink-ddd

benchmarks/benchmark_rmsnorm.py

@@ -3,14 +3,15 @@
 
 import torch
 
-from rl_engine.kernels.ops.pytorch.norm.rmsnorm_ref import rmsnorm_ref_custom
+from rl_engine.kernels.ops.pytorch.norm.rms_norm import NativeRMSNormOp
 from rl_engine.kernels.ops.triton.rmsnorm_triton import rmsnorm_triton
 
 try:
+    from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE
     from rl_engine.kernels.ops.cuda.norm.rmsnorm import rmsnorm_cuda
 
-    HAS_CUDA_EXT = True
-except Exception:
+    HAS_CUDA_EXT = _EXT_AVAILABLE and hasattr(_C, "rmsnorm_forward")
+except ImportError:
     HAS_CUDA_EXT = False
 
 
@@ -54,8 +55,10 @@ def make_inputs():
             w_base.detach().clone().requires_grad_(True),
         )
 
+    native = NativeRMSNormOp()
+
     x, w = make_inputs()
-    t_ref = bench(lambda a, b: rmsnorm_ref_custom(a, b), x, w, dy)
+    t_ref = bench(lambda a, b: native.forward(a, b), x, w, dy)
     print(f"pytorch ref : {t_ref:.4f} ms")
 
     x, w = make_inputs()

ci/run_gpu_ci.sh

@@ -186,7 +186,7 @@ TORCH_INDEX_URL="${TORCH_INDEX_URL:-https://download.pytorch.org/whl/cu124}"
 # --no-build-isolation: torch must be visible to setup.py, else the extension is silently skipped.
 # --no-deps: keep the pinned torch; do not let the editable install re-resolve it.
 "$PY" -m pip install --no-build-isolation --no-deps -e .
-"$PY" -m pip install --no-cache-dir numpy tabulate accelerate transformers pytest
+"$PY" -m pip install --no-cache-dir numpy tabulate accelerate "transformers==5.13.1" pytest
 nvidia-smi
 # Fail fast if _C did not build or cannot launch, instead of silently using native fallbacks.
 "$PY" scripts/ci_smoke.py

csrc/cuda/rmsnorm.cu

@@ -66,6 +66,9 @@ static int choose_threads(int H) {
     return 512;
 }
 
+constexpr int RMSNORM_DW_ROWS_PER_CHUNK = 256;
+constexpr int RMSNORM_DW_H_TILE = 128;
+
 
 template <typename scalar_t, typename weight_t>
 __global__ void rmsnorm_fwd_kernel(
@@ -164,46 +167,48 @@ __global__ void rmsnorm_partial_dw_kernel(
     int T,
     int H
 ) {
-    int idx = blockIdx.x * blockDim.x + threadIdx.x;
-    int total = T * H;
-
-    if (idx >= total) return;
-
-    int row = idx / H;
-
-    if (!mask[row]) {
-        partial_dw[idx] = 0.0f;
+    int chunk = blockIdx.x;
+    int h = blockIdx.y * RMSNORM_DW_H_TILE + threadIdx.x;
+    if (h >= H) {
         return;
     }
 
-    float dyv = load_as_float<scalar_t>(dy + idx);
-    float xv = load_as_float<scalar_t>(x + idx);
-    float r = rstd[row];
+    int t0 = chunk * RMSNORM_DW_ROWS_PER_CHUNK;
+    float acc = 0.0f;
+
+#pragma unroll
+    for (int i = 0; i < RMSNORM_DW_ROWS_PER_CHUNK; ++i) {
+        int row = t0 + i;
+        float contrib = 0.0f;
+        if (row < T && mask[row]) {
+            int idx = row * H + h;
+            float dyv = load_as_float<scalar_t>(dy + idx);
+            float xv = load_as_float<scalar_t>(x + idx);
+            contrib = dyv * xv * rstd[row];
+        }
+        acc += contrib;
+    }
 
-    partial_dw[idx] = dyv * xv * r;
+    partial_dw[chunk * H + h] = acc;
 }
 
 
 __global__ void rmsnorm_reduce_dw_kernel(
     const float* __restrict__ partial_dw,
     float* __restrict__ dw,
-    int T,
+    int chunks,
     int H
 ) {
-    int h = blockIdx.x;
-    int tid = threadIdx.x;
-
-    float local_sum = 0.0f;
-
-    for (int t = tid; t < T; t += blockDim.x) {
-        local_sum += partial_dw[t * H + h];
+    int h = blockIdx.x * RMSNORM_DW_H_TILE + threadIdx.x;
+    if (h >= H) {
+        return;
     }
 
-    float sum = block_reduce_sum(local_sum);
-
-    if (tid == 0) {
-        dw[h] = sum;
+    float acc = 0.0f;
+    for (int chunk = 0; chunk < chunks; ++chunk) {
+        acc += partial_dw[chunk * H + h];
     }
+    dw[h] = acc;
 }
 
 
@@ -280,10 +285,10 @@ void rmsnorm_backward_partial_dw_cuda(
 ) {
     int T = x.size(0);
     int H = x.size(1);
-    int total = T * H;
 
-    int threads = 256;
-    int blocks = (total + threads - 1) / threads;
+    int chunks = (T + RMSNORM_DW_ROWS_PER_CHUNK - 1) / RMSNORM_DW_ROWS_PER_CHUNK;
+    dim3 blocks(chunks, (H + RMSNORM_DW_H_TILE - 1) / RMSNORM_DW_H_TILE);
+    dim3 threads(RMSNORM_DW_H_TILE);
 
     cudaStream_t stream = at::cuda::getCurrentCUDAStream();
 
@@ -305,18 +310,18 @@ void rmsnorm_backward_reduce_dw_cuda(
     torch::Tensor partial_dw,
     torch::Tensor dw
 ) {
-    int T = partial_dw.size(0);
+    int chunks = partial_dw.size(0);
     int H = partial_dw.size(1);
 
-    int threads = choose_threads(T);
-    size_t smem = threads * sizeof(float);
+    dim3 blocks((H + RMSNORM_DW_H_TILE - 1) / RMSNORM_DW_H_TILE);
+    dim3 threads(RMSNORM_DW_H_TILE);
 
     cudaStream_t stream = at::cuda::getCurrentCUDAStream();
 
-    rmsnorm_reduce_dw_kernel<<<H, threads, smem, stream>>>(
+    rmsnorm_reduce_dw_kernel<<<blocks, threads, 0, stream>>>(
         partial_dw.data_ptr<float>(),
         dw.data_ptr<float>(),
-        T,
+        chunks,
         H
     );
 }

csrc/ops.cpp

@@ -77,6 +77,8 @@ torch::Tensor deterministic_logp_forward_indexed_fp32(torch::Tensor logits, torc
 
 // RMSNorm Declarations & Wrappers
 
+constexpr int RMSNORM_DW_ROWS_PER_CHUNK = 256;
+
 void rmsnorm_forward_cuda(
   torch::Tensor x,
   torch::Tensor weight,
@@ -173,8 +175,9 @@ torch::Tensor rmsnorm_backward_dw(
 
   auto T = x.size(0);
   auto H = x.size(1);
+  auto chunks = (T + RMSNORM_DW_ROWS_PER_CHUNK - 1) / RMSNORM_DW_ROWS_PER_CHUNK;
 
-  auto partial_dw = torch::empty({T, H}, x.options().dtype(torch::kFloat32));
+  auto partial_dw = torch::empty({chunks, H}, x.options().dtype(torch::kFloat32));
   auto dw = torch::empty({H}, x.options().dtype(torch::kFloat32));
 
   rmsnorm_backward_partial_dw_cuda(dy, x, rstd, mask, partial_dw);

pyproject.toml

@@ -17,7 +17,7 @@ dependencies = [
     "tabulate",
     "numpy",
     "accelerate",
-    "transformers",
+    "transformers==5.13.1",
 ]
 
 [project.optional-dependencies]

rl_engine/kernels/gtest/operator_specs.py

@@ -32,6 +32,19 @@ def _load_object(path: str) -> Any:
 
 
 OP_SPECS = {
+    "rms_norm": OperatorSpec(
+        name="rms_norm",
+        op_class="reduction",
+        gold_path="rl_engine.kernels.ops.pytorch.norm.rms_norm.NativeRMSNormOp",
+        gold_method="forward_fp32",
+        candidate_paths={
+            "pytorch": "rl_engine.kernels.ops.pytorch.norm.rms_norm.NativeRMSNormOp",
+            "triton": "rl_engine.kernels.ops.triton.rmsnorm_triton.RMSNormTritonOp",
+            "cuda": "rl_engine.kernels.ops.cuda.norm.rmsnorm.RMSNormCudaOp",
+            "cuda-sm90": "rl_engine.kernels.ops.cuda.norm.rmsnorm.RMSNormCudaOp",
+        },
+        grad_input_names=("x", "weight"),
+    ),
     "logp": OperatorSpec(
         name="logp",
         op_class="logprob",

rl_engine/kernels/ops/cuda/norm/rmsnorm.py

@@ -74,3 +74,16 @@ def rmsnorm_cuda(x, weight, eps=1e-6, mask=None):
         y = rmsnorm_cuda(x, weight, mask=mask)
     """
     return RMSNormCuda.apply(x, weight, mask, eps)
+
+
+class RMSNormCudaOp:
+    """CUDA RMSNorm wrapper compatible with the shared operator harness."""
+
+    def __call__(self, x, weight, *, eps=1e-6):
+        return self.forward(x, weight, eps=eps)
+
+    def forward(self, x, weight, *, eps=1e-6):
+        hidden = x.shape[-1]
+        x_2d = x.contiguous().view(-1, hidden)
+        y_2d = rmsnorm_cuda(x_2d, weight.contiguous(), eps=eps)
+        return y_2d.view_as(x)

rl_engine/kernels/ops/pytorch/norm/rmsnorm_native.py

@@ -96,3 +96,16 @@ def backward(ctx, grad_out):
 
 def rmsnorm_triton(x, weight, eps: float = 1e-6):
     return RMSNormTriton.apply(x, weight, eps)
+
+
+class RMSNormTritonOp:
+    """Triton RMSNorm wrapper compatible with the shared operator harness."""
+
+    def __call__(self, x, weight, *, eps: float = 1e-6):
+        return self.forward(x, weight, eps=eps)
+
+    def forward(self, x, weight, *, eps: float = 1e-6):
+        hidden = x.shape[-1]
+        x_2d = x.contiguous().view(-1, hidden)
+        y_2d = rmsnorm_triton(x_2d, weight.contiguous(), eps=eps)
+        return y_2d.view_as(x)