[Kernel] Implement CUDA kernel for ReLUSquaredActivation (relu^2) (#39058)

Signed-off-by: Tanish Malekar <tanishmalekar32@gmail.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>

benchmarks/kernels/benchmark_relu_squared.py

@@ -0,0 +1,108 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Benchmark ReLUSquaredActivation: custom CUDA kernel vs forward_native, both
+# eager and under torch.compile (Inductor fuses relu+square into one kernel).
+
+import itertools
+
+import torch
+import torch.nn.functional as F
+
+import vllm.model_executor.layers.activation  # noqa: F401
+from vllm.benchmarks.lib.utils import default_vllm_config
+from vllm.triton_utils import triton
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed
+
+# Capped so the largest tensor stays under 2**31 elements: the shared activation
+# kernel computes the per-token pointer offset (blockIdx.x * d) in 32-bit, which
+# overflows for tensors with >2**32 elements. Realistic token counts are well
+# below this; the kernel-vs-native gap is already clear at these sizes.
+batch_size_range = [1, 16, 128]
+seq_len_range = [1, 16, 64, 1024]
+intermediate_size = [3072, 9728, 12288]
+configs = list(itertools.product(batch_size_range, seq_len_range, intermediate_size))
+
+
+@default_vllm_config()
+def benchmark_relu_squared(
+    batch_size: int,
+    seq_len: int,
+    intermediate_size: int,
+    provider: str,
+    dtype: torch.dtype,
+):
+    device = "cuda"
+    num_tokens = batch_size * seq_len
+    set_random_seed(42)
+    torch.set_default_device(device)
+
+    x = torch.randn(num_tokens, intermediate_size, dtype=dtype, device=device)
+    out = torch.empty_like(x)
+
+    def native(x: torch.Tensor) -> torch.Tensor:
+        return torch.square(F.relu(x))
+
+    # Verify the custom kernel matches the native implementation before timing.
+    ref = native(x)
+    torch.ops._C.relu_squared(out, x)
+    torch.testing.assert_close(out, ref)
+
+    if provider == "custom":
+        # Custom CUDA kernel — single fused kernel.
+        fn = lambda: torch.ops._C.relu_squared(out, x)
+    elif provider == "native":
+        # forward_native, eager — relu and square as separate ops.
+        fn = lambda: native(x)
+    elif provider == "native_compiled":
+        # forward_native under torch.compile — Inductor fuses relu+square.
+        # This is the real production baseline (custom ops are off when
+        # Inductor is enabled), so it is the comparison reviewers care about.
+        compiled = torch.compile(native)
+        compiled(x)  # warm up / trigger compilation before timing
+        fn = lambda: compiled(x)
+
+    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
+        fn, quantiles=[0.5, 0.2, 0.8]
+    )
+    return ms, max_ms, min_ms
+
+
+if __name__ == "__main__":
+    parser = FlexibleArgumentParser(
+        description="Benchmark ReLUSquaredActivation: custom kernel vs native."
+    )
+    parser.add_argument(
+        "--dtype",
+        type=str,
+        choices=["half", "bfloat16", "float"],
+        default="bfloat16",
+    )
+    args = parser.parse_args()
+
+    dtype = STR_DTYPE_TO_TORCH_DTYPE[args.dtype]
+
+    perf_report = triton.testing.perf_report(
+        triton.testing.Benchmark(
+            x_names=["batch_size", "seq_len", "intermediate_size"],
+            x_vals=configs,
+            line_arg="provider",
+            line_vals=["custom", "native_compiled", "native"],
+            line_names=[
+                "Custom Kernel",
+                "Native (torch.compile)",
+                "Native (eager)",
+            ],
+            styles=[("blue", "-"), ("green", "-"), ("red", "-")],
+            ylabel="ms",
+            plot_name="relu_squared-eager-performance",
+            args={},
+        )
+    )
+
+    perf_report(
+        lambda batch_size, seq_len, intermediate_size, provider: benchmark_relu_squared(
+            batch_size, seq_len, intermediate_size, provider, dtype
+        )
+    ).run(print_data=True)

csrc/libtorch_stable/activation_kernels.cu

@@ -669,6 +669,14 @@ __device__ __forceinline__ T gelu_quick_kernel(const T& x) {
   return (T)(((float)x) / (1.0f + expf(-1.702f * (float)x)));
 }
 
+template <typename T>
+__device__ __forceinline__ T relu_squared_kernel(const T& x) {
+  // relu(x)^2 — introduced in https://arxiv.org/abs/2109.08668v2
+  const float f = (float)x;
+  const float val = f > 0.0f ? f : 0.0f;
+  return (T)(val * val);
+}
+
 }  // namespace vllm
 
 void gelu_new(torch::stable::Tensor& out,    // [..., d]
@@ -688,3 +696,9 @@ void gelu_quick(torch::stable::Tensor& out,    // [..., d]
 {
   LAUNCH_ACTIVATION_KERNEL(vllm::gelu_quick_kernel);
 }
+
+void relu_squared(torch::stable::Tensor& out,    // [..., d]
+                  torch::stable::Tensor& input)  // [..., d]
+{
+  LAUNCH_ACTIVATION_KERNEL(vllm::relu_squared_kernel);
+}

csrc/libtorch_stable/ops.h

@@ -414,6 +414,8 @@ void gelu_new(torch::stable::Tensor& out, torch::stable::Tensor& input);
 void gelu_fast(torch::stable::Tensor& out, torch::stable::Tensor& input);
 void gelu_quick(torch::stable::Tensor& out, torch::stable::Tensor& input);
 
+void relu_squared(torch::stable::Tensor& out, torch::stable::Tensor& input);
+
 // INT8 quantization kernels (shared CUDA/ROCm)
 void static_scaled_int8_quant(torch::stable::Tensor& out,
                               torch::stable::Tensor const& input,

csrc/libtorch_stable/torch_bindings.cpp

@@ -539,6 +539,9 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
   // Quick GELU implementation.
   ops.def("gelu_quick(Tensor! out, Tensor input) -> ()");
 
+  // relu(x)^2 activation from https://arxiv.org/abs/2109.08668v2
+  ops.def("relu_squared(Tensor! out, Tensor input) -> ()");
+
   // Compute int8 quantized tensor for given scaling factor.
   ops.def(
       "static_scaled_int8_quant(Tensor! result, Tensor input, Tensor scale,"
@@ -715,6 +718,7 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
   ops.impl("gelu_new", TORCH_BOX(&gelu_new));
   ops.impl("gelu_fast", TORCH_BOX(&gelu_fast));
   ops.impl("gelu_quick", TORCH_BOX(&gelu_quick));
+  ops.impl("relu_squared", TORCH_BOX(&relu_squared));
   ops.impl("silu_and_mul_with_clamp", TORCH_BOX(&silu_and_mul_clamp));
 
   // INT8 quantization kernels

csrc/ops.h

@@ -43,6 +43,8 @@ void gelu_fast(torch::Tensor& out, torch::Tensor& input);
 
 void gelu_quick(torch::Tensor& out, torch::Tensor& input);
 
+void relu_squared(torch::Tensor& out, torch::Tensor& input);
+
 void static_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input,
                               torch::Tensor const& scale,
                               std::optional<torch::Tensor> const& azp);

tests/kernels/core/test_activation.py

@@ -15,6 +15,7 @@ from vllm.model_executor.layers.activation import (
     MulAndSilu,
     NewGELU,
     QuickGELU,
+    ReLUSquaredActivation,
     SiluAndMul,
     SiluAndMulWithClamp,
     SwigluOAIAndMul,
@@ -202,6 +203,7 @@ def test_silu_and_mul_with_clamp(
         (FastGELU, torch.ops._C.gelu_fast),
         (NewGELU, torch.ops._C.gelu_new),
         (QuickGELU, torch.ops._C.gelu_quick),
+        (ReLUSquaredActivation, torch.ops._C.relu_squared),
     ],
 )
 @pytest.mark.parametrize("num_tokens", NUM_TOKENS)

vllm/model_executor/layers/activation.py

@@ -613,13 +613,19 @@ class ReLUSquaredActivation(CustomOp):
 
     # --8<-- [end:relu2]
 
+    def __init__(self):
+        super().__init__()
+        if current_platform.is_cuda_alike():
+            self.op = torch.ops._C.relu_squared
+
     def forward_native(self, x: torch.Tensor) -> torch.Tensor:
         """PyTorch-native implementation equivalent to forward()."""
         return torch.square(F.relu(x))
 
     def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
-        # TODO : implement cuda kernels
-        return self.forward_native(x)
+        out = torch.empty_like(x)
+        self.op(out, x)
+        return out
 
 
 # --8<-- [start:xielu]