[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>
Tanish Malekar committed
Jul 12, 2026 at 19:18 UTC
e26264f3ef5987a172162bbf497b8ae7ecbe0108
7 files changed
+140
-2
benchmarks/kernels/benchmark_relu_squared.py
new
+108
@@ -0,0 +1,108 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+# Benchmark ReLUSquaredActivation: custom CUDA kernel vs forward_native, both
5
+# eager and under torch.compile (Inductor fuses relu+square into one kernel).
6
+
7
+import itertools
8
+
9
+import torch
10
+import torch.nn.functional as F
11
+
12
+import vllm.model_executor.layers.activation # noqa: F401
13
+from vllm.benchmarks.lib.utils import default_vllm_config
14
+from vllm.triton_utils import triton
15
+from vllm.utils.argparse_utils import FlexibleArgumentParser
16
+from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed
17
+
18
+# Capped so the largest tensor stays under 2**31 elements: the shared activation
19
+# kernel computes the per-token pointer offset (blockIdx.x * d) in 32-bit, which
20
+# overflows for tensors with >2**32 elements. Realistic token counts are well
21
+# below this; the kernel-vs-native gap is already clear at these sizes.
22
+batch_size_range = [1, 16, 128]
23
+seq_len_range = [1, 16, 64, 1024]
24
+intermediate_size = [3072, 9728, 12288]
25
+configs = list(itertools.product(batch_size_range, seq_len_range, intermediate_size))
26
+
27
+
28
+@default_vllm_config()
29
+def benchmark_relu_squared(
30
+ batch_size: int,
31
+ seq_len: int,
32
+ intermediate_size: int,
33
+ provider: str,
34
+ dtype: torch.dtype,
35
+):
36
+ device = "cuda"
37
+ num_tokens = batch_size * seq_len
38
+ set_random_seed(42)
39
+ torch.set_default_device(device)
40
+
41
+ x = torch.randn(num_tokens, intermediate_size, dtype=dtype, device=device)
42
+ out = torch.empty_like(x)
43
+
44
+ def native(x: torch.Tensor) -> torch.Tensor:
45
+ return torch.square(F.relu(x))
46
+
47
+ # Verify the custom kernel matches the native implementation before timing.
48
+ ref = native(x)
49
+ torch.ops._C.relu_squared(out, x)
50
+ torch.testing.assert_close(out, ref)
51
+
52
+ if provider == "custom":
53
+ # Custom CUDA kernel — single fused kernel.
54
+ fn = lambda: torch.ops._C.relu_squared(out, x)
55
+ elif provider == "native":
56
+ # forward_native, eager — relu and square as separate ops.
57
+ fn = lambda: native(x)
58
+ elif provider == "native_compiled":
59
+ # forward_native under torch.compile — Inductor fuses relu+square.
60
+ # This is the real production baseline (custom ops are off when
61
+ # Inductor is enabled), so it is the comparison reviewers care about.
62
+ compiled = torch.compile(native)
63
+ compiled(x) # warm up / trigger compilation before timing
64
+ fn = lambda: compiled(x)
65
+
66
+ ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
67
+ fn, quantiles=[0.5, 0.2, 0.8]
68
+ )
69
+ return ms, max_ms, min_ms
70
+
71
+
72
+if __name__ == "__main__":
73
+ parser = FlexibleArgumentParser(
74
+ description="Benchmark ReLUSquaredActivation: custom kernel vs native."
75
+ )
76
+ parser.add_argument(
77
+ "--dtype",
78
+ type=str,
79
+ choices=["half", "bfloat16", "float"],
80
+ default="bfloat16",
81
+ )
82
+ args = parser.parse_args()
83
+
84
+ dtype = STR_DTYPE_TO_TORCH_DTYPE[args.dtype]
85
+
86
+ perf_report = triton.testing.perf_report(
87
+ triton.testing.Benchmark(
88
+ x_names=["batch_size", "seq_len", "intermediate_size"],
89
+ x_vals=configs,
90
+ line_arg="provider",
91
+ line_vals=["custom", "native_compiled", "native"],
92
+ line_names=[
93
+ "Custom Kernel",
94
+ "Native (torch.compile)",
95
+ "Native (eager)",
96
+ ],
97
+ styles=[("blue", "-"), ("green", "-"), ("red", "-")],
98
+ ylabel="ms",
99
+ plot_name="relu_squared-eager-performance",
100
+ args={},
101
+ )
102
+ )
103
+
104
+ perf_report(
105
+ lambda batch_size, seq_len, intermediate_size, provider: benchmark_relu_squared(
106
+ batch_size, seq_len, intermediate_size, provider, dtype
107
+ )
108
+ ).run(print_data=True)
csrc/libtorch_stable/activation_kernels.cu
+14
@@ -669,6 +669,14 @@ __device__ __forceinline__ T gelu_quick_kernel(const T& x) {
669
return (T)(((float)x) / (1.0f + expf(-1.702f * (float)x)));
670
}
671
672
+template <typename T>
673
+__device__ __forceinline__ T relu_squared_kernel(const T& x) {
674
+ // relu(x)^2 — introduced in https://arxiv.org/abs/2109.08668v2
675
+ const float f = (float)x;
676
+ const float val = f > 0.0f ? f : 0.0f;
677
+ return (T)(val * val);
678
+}
679
+
680
} // namespace vllm
681
682
void gelu_new(torch::stable::Tensor& out, // [..., d]
@@ -688,3 +696,9 @@ void gelu_quick(torch::stable::Tensor& out, // [..., d]
696
{
697
LAUNCH_ACTIVATION_KERNEL(vllm::gelu_quick_kernel);
698
}
699
+
700
+void relu_squared(torch::stable::Tensor& out, // [..., d]
701
+ torch::stable::Tensor& input) // [..., d]
702
+{
703
+ LAUNCH_ACTIVATION_KERNEL(vllm::relu_squared_kernel);
704
+}
csrc/libtorch_stable/ops.h
+2
@@ -414,6 +414,8 @@ void gelu_new(torch::stable::Tensor& out, torch::stable::Tensor& input);
414
void gelu_fast(torch::stable::Tensor& out, torch::stable::Tensor& input);
415
void gelu_quick(torch::stable::Tensor& out, torch::stable::Tensor& input);
416
417
+void relu_squared(torch::stable::Tensor& out, torch::stable::Tensor& input);
418
+
419
// INT8 quantization kernels (shared CUDA/ROCm)
420
void static_scaled_int8_quant(torch::stable::Tensor& out,
421
torch::stable::Tensor const& input,
csrc/libtorch_stable/torch_bindings.cpp
+4
@@ -539,6 +539,9 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
539
// Quick GELU implementation.
540
ops.def("gelu_quick(Tensor! out, Tensor input) -> ()");
541
542
+ // relu(x)^2 activation from https://arxiv.org/abs/2109.08668v2
543
+ ops.def("relu_squared(Tensor! out, Tensor input) -> ()");
544
+
545
// Compute int8 quantized tensor for given scaling factor.
546
ops.def(
547
"static_scaled_int8_quant(Tensor! result, Tensor input, Tensor scale,"
@@ -715,6 +718,7 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
718
ops.impl("gelu_new", TORCH_BOX(&gelu_new));
719
ops.impl("gelu_fast", TORCH_BOX(&gelu_fast));
720
ops.impl("gelu_quick", TORCH_BOX(&gelu_quick));
721
+ ops.impl("relu_squared", TORCH_BOX(&relu_squared));
722
ops.impl("silu_and_mul_with_clamp", TORCH_BOX(&silu_and_mul_clamp));
723
724
// INT8 quantization kernels
csrc/ops.h
+2
@@ -43,6 +43,8 @@ void gelu_fast(torch::Tensor& out, torch::Tensor& input);
43
44
void gelu_quick(torch::Tensor& out, torch::Tensor& input);
45
46
+void relu_squared(torch::Tensor& out, torch::Tensor& input);
47
+
48
void static_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input,
49
torch::Tensor const& scale,
50
std::optional<torch::Tensor> const& azp);
tests/kernels/core/test_activation.py
+2
@@ -15,6 +15,7 @@ from vllm.model_executor.layers.activation import (
15
MulAndSilu,
16
NewGELU,
17
QuickGELU,
18
+ ReLUSquaredActivation,
19
SiluAndMul,
20
SiluAndMulWithClamp,
21
SwigluOAIAndMul,
@@ -202,6 +203,7 @@ def test_silu_and_mul_with_clamp(
203
(FastGELU, torch.ops._C.gelu_fast),
204
(NewGELU, torch.ops._C.gelu_new),
205
(QuickGELU, torch.ops._C.gelu_quick),
206
+ (ReLUSquaredActivation, torch.ops._C.relu_squared),
207
],
208
)
209
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
vllm/model_executor/layers/activation.py
+8
-2
@@ -613,13 +613,19 @@ class ReLUSquaredActivation(CustomOp):
613
614
# --8<-- [end:relu2]
615
616
+ def __init__(self):
617
+ super().__init__()
618
+ if current_platform.is_cuda_alike():
619
+ self.op = torch.ops._C.relu_squared
620
+
621
def forward_native(self, x: torch.Tensor) -> torch.Tensor:
622
"""PyTorch-native implementation equivalent to forward()."""
623
return torch.square(F.relu(x))
624
625
def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
621
- # TODO : implement cuda kernels
622
- return self.forward_native(x)
626
+ out = torch.empty_like(x)
627
+ self.op(out, x)
628
+ return out
629
630
631
# --8<-- [start:xielu]