diff --git a/benchmarks/kernels/benchmark_relu_squared.py b/benchmarks/kernels/benchmark_relu_squared.py new file mode 100644 index 0000000000..00550ca475 --- /dev/null +++ b/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) diff --git a/csrc/libtorch_stable/activation_kernels.cu b/csrc/libtorch_stable/activation_kernels.cu index e1dc013460..60b8ca5f38 100644 --- a/csrc/libtorch_stable/activation_kernels.cu +++ b/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 +__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); +} diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 4520f4a99e..0aeba41471 100644 --- a/csrc/libtorch_stable/ops.h +++ b/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, diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 7581769f4f..67487daf3e 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/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 diff --git a/csrc/ops.h b/csrc/ops.h index 274cd52bea..d7ee2d080c 100644 --- a/csrc/ops.h +++ b/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 const& azp); diff --git a/tests/kernels/core/test_activation.py b/tests/kernels/core/test_activation.py index 3f1d45ba8e..f698c385fe 100644 --- a/tests/kernels/core/test_activation.py +++ b/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) diff --git a/vllm/model_executor/layers/activation.py b/vllm/model_executor/layers/activation.py index 7c3fe4407d..a8c91af9ad 100644 --- a/vllm/model_executor/layers/activation.py +++ b/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]