wehub-resource-sync
109 lines
3.7 KiB
Python
109 lines
3.7 KiB
Python
# 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)
|