refactor: replace warp shuffle with CUB BlockReduce

- Replace hand-written warp_reduce_sum/block_reduce_sum with CUB BlockReduce
- Add BLOCK_SIZE template param to fwd_small for compile-time CUB sizing
- Dispatch fwd_small with matched BlockReduce<float, 32/64/128/256>
- Simpler code, no performance regression (small dims slightly faster)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

activation/rms_norm.cu

@@ -13,62 +13,12 @@ template <typename type, int N> struct alignas(sizeof(type) * N) type_vec_t {
   type data[N];
 };
 
-// Warp shuffle block reduction
-// Uses ~128 bytes shared memory vs ~8KB for CUB
-__device__ __forceinline__ float warp_reduce_sum(float val) {
-#pragma unroll
-  for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2)
-    val += __shfl_down_sync(0xffffffff, val, offset);
-  return val;
-}
-
-__device__ __forceinline__ float2 warp_reduce_sum2(float2 v) {
-#pragma unroll
-  for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
-    v.x += __shfl_down_sync(0xffffffff, v.x, offset);
-    v.y += __shfl_down_sync(0xffffffff, v.y, offset);
-  }
-  return v;
-}
-
-__device__ __forceinline__ float2 block_reduce_sum2(float2 v) {
-  __shared__ float warp_x[32];
-  __shared__ float warp_y[32];
-  int lane = threadIdx.x % WARP_SIZE;
-  int warp = threadIdx.x / WARP_SIZE;
-  int num_warps = (blockDim.x + WARP_SIZE - 1) / WARP_SIZE;
-
-  v = warp_reduce_sum2(v);
-  if (lane == 0) {
-    warp_x[warp] = v.x;
-    warp_y[warp] = v.y;
-  }
-  __syncthreads();
-
-  if (warp == 0) {
-    v.x = (lane < num_warps) ? warp_x[lane] : 0.0f;
-    v.y = (lane < num_warps) ? warp_y[lane] : 0.0f;
-    v = warp_reduce_sum2(v);
+// Float2 sum operator for CUB BlockReduce
+struct Float2SumOp {
+  __device__ float2 operator()(const float2 &a, const float2 &b) const {
+    return make_float2(a.x + b.x, a.y + b.y);
   }
-  return v;
-}
-
-__device__ __forceinline__ float block_reduce_sum(float val) {
-  __shared__ float warp_sums[32];
-  int lane = threadIdx.x % WARP_SIZE;
-  int warp = threadIdx.x / WARP_SIZE;
-  int num_warps = (blockDim.x + WARP_SIZE - 1) / WARP_SIZE;
-
-  val = warp_reduce_sum(val);
-  if (lane == 0)
-    warp_sums[warp] = val;
-  __syncthreads();
-
-  val = (lane < num_warps) ? warp_sums[lane] : 0.0f;
-  if (warp == 0)
-    val = warp_reduce_sum(val);
-  return val;
-}
+};
 
 // ---------------------------------------------------------------------------
 // Forward (dim ≤ 2048): single-pass with register caching
@@ -78,7 +28,8 @@ __device__ __forceinline__ float block_reduce_sum(float val) {
 //   Pass 2: write output from cache (no second global read)
 //   Also writes inv_rms[token] for backward
 // ---------------------------------------------------------------------------
-template <typename scalar_t, typename acc_t, int width, int NVECS>
+template <typename scalar_t, typename acc_t, int width, int NVECS,
+          int BLOCK_SIZE = 256>
 __global__ void rms_norm_fwd_small(scalar_t *__restrict__ out,
                                    acc_t *__restrict__ inv_rms_out,
                                    const scalar_t *__restrict__ input,
@@ -103,7 +54,11 @@ __global__ void rms_norm_fwd_small(scalar_t *__restrict__ out,
     }
   }
 
-  sum_square = block_reduce_sum(sum_square);
+  {
+    using BlockReduce = cub::BlockReduce<float, BLOCK_SIZE>;
+    __shared__ typename BlockReduce::TempStorage reduceStore;
+    sum_square = BlockReduce(reduceStore).Sum(sum_square, blockDim.x);
+  }
 
   __shared__ acc_t s_scale;
   if (threadIdx.x == 0) {
@@ -159,7 +114,11 @@ __global__ void rms_norm_fwd_large(scalar_t *__restrict__ out,
     }
   }
 
-  sum_square = block_reduce_sum(sum_square);
+  {
+    using BlockReduce = cub::BlockReduce<float, 256>;
+    __shared__ typename BlockReduce::TempStorage reduceStore;
+    sum_square = BlockReduce(reduceStore).Sum(sum_square, blockDim.x);
+  }
 
   __shared__ acc_t s_scale;
   if (threadIdx.x == 0) {
@@ -203,7 +162,11 @@ __global__ void rms_norm_fwd_scalar(scalar_t *__restrict__ out,
     sum_square += x * x;
   }
 
-  sum_square = block_reduce_sum(sum_square);
+  {
+    using BlockReduce = cub::BlockReduce<float, 256>;
+    __shared__ typename BlockReduce::TempStorage reduceStore;
+    sum_square = BlockReduce(reduceStore).Sum(sum_square, blockDim.x);
+  }
 
   __shared__ acc_t s_scale;
   if (threadIdx.x == 0) {
@@ -264,7 +227,11 @@ __global__ __launch_bounds__(256, 4) void rms_norm_bwd_large_input_grad(
     }
   }
 
-  d_sum = block_reduce_sum(d_sum);
+  {
+    using BlockReduce = cub::BlockReduce<float, 256>;
+    __shared__ typename BlockReduce::TempStorage reduceStore;
+    d_sum = BlockReduce(reduceStore).Sum(d_sum, blockDim.x);
+  }
 
   __shared__ acc_t s_dxx;
   if (threadIdx.x == 0) {
@@ -331,6 +298,13 @@ __global__ void rms_norm_bwd_fused(
   int64_t token_start = static_cast<int64_t>(blockIdx.x) * tpb;
   int64_t token_end = min(token_start + tpb, num_tokens);
 
+  // Shared TempStorage for CUB block reductions — declared once, reused per
+  // iteration
+  using BlockReduce2 = cub::BlockReduce<float2, 256>;
+  __shared__ typename BlockReduce2::TempStorage reduceStore2;
+  using BlockReduce = cub::BlockReduce<float, 256>;
+  __shared__ typename BlockReduce::TempStorage reduceStore;
+
   // Process tokens in pairs with float2 reduction
   int64_t t = token_start;
   for (; t + 1 < token_end; t += 2) {
@@ -355,7 +329,9 @@ __global__ void rms_norm_bwd_fused(
       }
     }
 
-    float2 sums = block_reduce_sum2(make_float2(dsum0, dsum1));
+    float2 sums =
+        BlockReduce2(reduceStore2)
+            .Reduce(make_float2(dsum0, dsum1), Float2SumOp{}, blockDim.x);
 
     // dxx = d_sum * scale^3 / d
     __shared__ acc_t sd0, sd1;
@@ -405,7 +381,7 @@ __global__ void rms_norm_bwd_fused(
                  static_cast<acc_t>(x_vec.data[i]) *
                  static_cast<acc_t>(w_vec.data[i]);
     }
-    d_sum = block_reduce_sum(d_sum);
+    d_sum = BlockReduce(reduceStore).Sum(d_sum, blockDim.x);
 
     __shared__ acc_t s_dxx;
     if (threadIdx.x == 0)
@@ -520,7 +496,11 @@ __global__ void rms_norm_bwd_scalar(scalar_t *__restrict__ input_grad,
     d_sum += dy * x * w;
   }
 
-  d_sum = block_reduce_sum(d_sum);
+  {
+    using BlockReduce = cub::BlockReduce<float, 256>;
+    __shared__ typename BlockReduce::TempStorage reduceStore;
+    d_sum = BlockReduce(reduceStore).Sum(d_sum, blockDim.x);
+  }
 
   __shared__ acc_t s_dxx;
   if (threadIdx.x == 0) {
@@ -574,15 +554,26 @@ rms_norm(const torch::Tensor &input,  // [..., d]
       // Single-pass: 1 vec per thread (dim <= 2048)
       int block_size = ((vec_d + 31) / 32) * 32;
       block_size = std::max(block_size, 32);
-      dim3 block(block_size);
-      MOTIF_DISPATCH_FLOATING_TYPES(
-          input.scalar_type(), "rms_norm_fwd_small", [&] {
-            motif::rms_norm_fwd_small<scalar_t, float, 8, 1>
-                <<<grid, block, 0, stream>>>(
-                    out.data_ptr<scalar_t>(), inv_rms.data_ptr<float>(),
-                    input.data_ptr<scalar_t>(), weight.data_ptr<scalar_t>(),
-                    eps, d);
-          });
+
+#define LAUNCH_FWD_SMALL(BS)                                                   \
+  MOTIF_DISPATCH_FLOATING_TYPES(                                               \
+      input.scalar_type(), "rms_norm_fwd_small", [&] {                         \
+        motif::rms_norm_fwd_small<scalar_t, float, 8, 1, BS>                   \
+            <<<grid, dim3(BS), 0, stream>>>(                                   \
+                out.data_ptr<scalar_t>(), inv_rms.data_ptr<float>(),           \
+                input.data_ptr<scalar_t>(), weight.data_ptr<scalar_t>(), eps,  \
+                d);                                                            \
+      })
+
+      if (block_size <= 32)
+        LAUNCH_FWD_SMALL(32);
+      else if (block_size <= 64)
+        LAUNCH_FWD_SMALL(64);
+      else if (block_size <= 128)
+        LAUNCH_FWD_SMALL(128);
+      else
+        LAUNCH_FWD_SMALL(256);
+#undef LAUNCH_FWD_SMALL
     } else {
       // Large dims: 2-pass
       dim3 block(block_2pass);
@@ -619,23 +610,23 @@ rms_norm(const torch::Tensor &input,  // [..., d]
 
 std::tuple<torch::Tensor, torch::Tensor>
 rms_norm_backward(const torch::Tensor &output_grad, // [..., d]
-                  const torch::Tensor &output,  // [..., d] — forward output y
-                  const torch::Tensor &weight,  // [d]
-                  const torch::Tensor &inv_rms, // [num_tokens]
+                  const torch::Tensor &input,       // [..., d]
+                  const torch::Tensor &weight,      // [d]
+                  const torch::Tensor &inv_rms,     // [num_tokens]
                   double eps) {
-  torch::Tensor input_grad = torch::empty_like(output);
+  torch::Tensor input_grad = torch::empty_like(input);
   torch::Tensor weight_grad = torch::empty_like(weight);
 
   AssertTensorContiguous(output_grad, "output_grad");
-  AssertTensorContiguous(output, "output");
+  AssertTensorContiguous(input, "input");
   AssertTensorContiguous(weight, "weight");
   AssertTensorContiguous(inv_rms, "inv_rms");
 
-  int d = output.size(-1);
-  int64_t num_tokens = output.numel() / output.size(-1);
+  int d = input.size(-1);
+  int64_t num_tokens = input.numel() / input.size(-1);
 
   dim3 grid(num_tokens);
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(output));
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
   const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
 
   if (d % 8 == 0) {
@@ -648,20 +639,20 @@ rms_norm_backward(const torch::Tensor &output_grad, // [..., d]
       int64_t num_blocks_mt = (num_tokens + tpb - 1) / tpb;
 
       torch::Tensor wg_acc =
-          torch::zeros({d}, output.options().dtype(torch::kFloat));
+          torch::zeros({d}, input.options().dtype(torch::kFloat));
 
       int block_size = std::min(vec_d, 128);
       block_size = ((block_size + 31) / 32) * 32;
       block_size = std::max(block_size, 32);
       size_t smem = d * sizeof(float);
-      // 'output' C++ arg receives input (saved by Python autograd)
+      // Small/medium dims (d <= 8192): multi-token block with shared memory
       MOTIF_DISPATCH_FLOATING_TYPES(
-          output.scalar_type(), "rms_norm_bwd_fused", [&] {
+          input.scalar_type(), "rms_norm_bwd_fused", [&] {
             motif::rms_norm_bwd_fused<scalar_t, float, 8>
                 <<<num_blocks_mt, block_size, smem, stream>>>(
                     input_grad.data_ptr<scalar_t>(), wg_acc.data_ptr<float>(),
                     output_grad.data_ptr<scalar_t>(),
-                    output.data_ptr<scalar_t>(), weight.data_ptr<scalar_t>(),
+                    input.data_ptr<scalar_t>(), weight.data_ptr<scalar_t>(),
                     inv_rms.data_ptr<float>(), d, num_tokens, tpb);
           });
 
@@ -671,20 +662,20 @@ rms_norm_backward(const torch::Tensor &output_grad, // [..., d]
     } else {
       // Large dims (d > 8192): input-based bwd + column-parallel weight grad
       MOTIF_DISPATCH_FLOATING_TYPES(
-          output.scalar_type(), "rms_norm_bwd_large_input_grad", [&] {
+          input.scalar_type(), "rms_norm_bwd_large_input_grad", [&] {
             motif::rms_norm_bwd_large_input_grad<scalar_t, float, 8>
                 <<<grid, dim3(256), 0, stream>>>(
                     input_grad.data_ptr<scalar_t>(),
                     output_grad.data_ptr<scalar_t>(),
-                    output.data_ptr<scalar_t>(), weight.data_ptr<scalar_t>(),
+                    input.data_ptr<scalar_t>(), weight.data_ptr<scalar_t>(),
                     inv_rms.data_ptr<float>(), d);
           });
 
       if (weight_grad.defined()) {
         int64_t chunk_size = 256;
         int64_t num_chunks = (num_tokens + chunk_size - 1) / chunk_size;
-        torch::Tensor partial_wg = torch::empty(
-            {num_chunks, d}, output.options().dtype(torch::kFloat));
+        torch::Tensor partial_wg =
+            torch::empty({num_chunks, d}, input.options().dtype(torch::kFloat));
 
         constexpr int TILE_T = 64;
         constexpr int VEC_W = 8;
@@ -692,13 +683,13 @@ rms_norm_backward(const torch::Tensor &output_grad, // [..., d]
         int cols_per_block = wg_threads * VEC_W;
         dim3 wg_grid((d + cols_per_block - 1) / cols_per_block, num_chunks);
         MOTIF_DISPATCH_FLOATING_TYPES(
-            output.scalar_type(), "rms_norm_bwd_large_weight_grad", [&] {
+            input.scalar_type(), "rms_norm_bwd_large_weight_grad", [&] {
               motif::rms_norm_bwd_large_weight_grad<scalar_t, float, TILE_T,
                                                     VEC_W>
                   <<<wg_grid, wg_threads, 0, stream>>>(
                       partial_wg.data_ptr<float>(),
                       output_grad.data_ptr<scalar_t>(),
-                      output.data_ptr<scalar_t>(), weight.data_ptr<scalar_t>(),
+                      input.data_ptr<scalar_t>(), weight.data_ptr<scalar_t>(),
                       inv_rms.data_ptr<float>(), d, num_tokens, chunk_size);
             });
 
@@ -712,19 +703,19 @@ rms_norm_backward(const torch::Tensor &output_grad, // [..., d]
     // Scalar fallback: temp buffer + at::sum_out (still uses input-based for
     // scalar)
     torch::Tensor temp_weight_grad =
-        torch::empty({num_tokens, d}, output.options().dtype(torch::kFloat));
+        torch::empty({num_tokens, d}, input.options().dtype(torch::kFloat));
 
     int block_size = std::min(d, 256);
     block_size = ((block_size + 31) / 32) * 32;
     block_size = std::max(block_size, 32);
     dim3 block(block_size);
     MOTIF_DISPATCH_FLOATING_TYPES(
-        output.scalar_type(), "rms_norm_bwd_scalar", [&] {
+        input.scalar_type(), "rms_norm_bwd_scalar", [&] {
           motif::rms_norm_bwd_scalar<scalar_t, float>
               <<<grid, block, 0, stream>>>(
                   input_grad.data_ptr<scalar_t>(),
                   temp_weight_grad.data_ptr<float>(),
-                  output_grad.data_ptr<scalar_t>(), output.data_ptr<scalar_t>(),
+                  output_grad.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(),
                   weight.data_ptr<scalar_t>(), inv_rms.data_ptr<float>(), d);
         });
 

torch-ext/torch_binding.cpp

@@ -20,7 +20,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "rms_norm(Tensor input, Tensor weight, float eps) -> (Tensor, Tensor)");
   ops.impl("rms_norm", torch::kCUDA, &rms_norm);
 
-  ops.def("rms_norm_backward(Tensor output_grad, Tensor output, Tensor weight, "
+  ops.def("rms_norm_backward(Tensor output_grad, Tensor input, Tensor weight, "
           "Tensor inv_rms, float eps) -> (Tensor, Tensor)");
   ops.impl("rms_norm_backward", torch::kCUDA, &rms_norm_backward);
 

torch-ext/torch_binding.h

@@ -14,7 +14,7 @@ void poly_norm_backward(torch::Tensor &input_grad, torch::Tensor &weight_grad,
 std::tuple<torch::Tensor, torch::Tensor>
 rms_norm(const torch::Tensor &input, const torch::Tensor &weights, double eps);
 std::tuple<torch::Tensor, torch::Tensor>
-rms_norm_backward(const torch::Tensor &output_grad, const torch::Tensor &output,
+rms_norm_backward(const torch::Tensor &output_grad, const torch::Tensor &input,
                   const torch::Tensor &weight, const torch::Tensor &inv_rms,
                   double eps);