From 78e84261ab2ffb691a8079d698d1f3f96ecb7635 Mon Sep 17 00:00:00 2001 From: JooHo Lee <96564470+BWAAEEEK@users.noreply.github.com> Date: Mon, 14 Sep 2026 16:42:15 +0900 Subject: [PATCH] [Performance] Add Triton kernel for Gemma3n sparse GELU (#48498) Signed-off-by: BWAAEEEK Co-authored-by: Andreas Karatzas Co-authored-by: Cyrus Leung --- .../kernels/benchmark_gelu_and_mul_sparse.py | 90 +++++++++ tests/compile/passes/ir/test_lowering.py | 34 ++++ tests/compile/test_aot_compile.py | 65 +++++++ tests/kernels/ir/test_activation.py | 183 ++++++++++++++++++ tests/test_config.py | 16 ++ vllm/config/kernel.py | 3 + vllm/ir/ops/__init__.py | 3 +- vllm/ir/ops/activation.py | 34 ++++ vllm/kernels/__init__.py | 4 +- vllm/kernels/triton/__init__.py | 4 + vllm/kernels/triton/activation.py | 110 +++++++++++ vllm/model_executor/layers/activation.py | 18 +- vllm/platforms/cuda.py | 5 +- vllm/platforms/rocm.py | 5 +- vllm/platforms/xpu.py | 2 +- 15 files changed, 555 insertions(+), 21 deletions(-) create mode 100644 benchmarks/kernels/benchmark_gelu_and_mul_sparse.py create mode 100644 tests/kernels/ir/test_activation.py create mode 100644 vllm/ir/ops/activation.py create mode 100644 vllm/kernels/triton/activation.py diff --git a/benchmarks/kernels/benchmark_gelu_and_mul_sparse.py b/benchmarks/kernels/benchmark_gelu_and_mul_sparse.py new file mode 100644 index 0000000000..9f8d976121 --- /dev/null +++ b/benchmarks/kernels/benchmark_gelu_and_mul_sparse.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import itertools + +import torch + +import vllm.kernels # noqa: F401 +from vllm import ir +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 + +NUM_TOKENS = [1, 2, 8, 32, 128, 512, 2048] +INTERMEDIATE_SIZES = [2048, 4096, 8192, 13824, 16384, 32768] +CONFIGS = list(itertools.product(NUM_TOKENS, INTERMEDIATE_SIZES)) +STD_MULTIPLIER = 1.6448536269514722 + + +@default_vllm_config() +def benchmark_gelu_and_mul_sparse( + num_tokens: int, + intermediate_size: int, + provider: str, + dtype: torch.dtype, +): + set_random_seed(0) + torch.set_default_device("cuda") + x = torch.randn(num_tokens, 2 * intermediate_size, dtype=dtype) + native = ir.ops.gelu_and_mul_sparse.impls["native"].impl_fn + triton_impl = ir.ops.gelu_and_mul_sparse.impls["triton"].impl_fn + args = (x, STD_MULTIPLIER, "tanh") + + expected = native(*args) + actual = triton_impl(*args) + torch.testing.assert_close( + actual, + expected, + **ir.ops.gelu_and_mul_sparse.get_tolerance(actual.dtype), + ) + + if provider == "triton": + fn = lambda: triton_impl(*args) + elif provider == "native_compiled": + compiled = torch.compile(native, fullgraph=True) + compiled(*args) + fn = lambda: compiled(*args) + else: + fn = lambda: native(*args) + + 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 the Gemma3n sparse GELU activation." + ) + 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=["num_tokens", "intermediate_size"], + x_vals=CONFIGS, + line_arg="provider", + line_vals=["triton", "native_compiled", "native_eager"], + line_names=["Triton", "Native (torch.compile)", "Native (eager)"], + styles=[("blue", "-"), ("green", "-"), ("red", "-")], + ylabel="ms", + plot_name="gelu-and-mul-sparse-performance", + args={}, + ) + ) + perf_report( + lambda num_tokens, intermediate_size, provider: ( + benchmark_gelu_and_mul_sparse( + num_tokens, intermediate_size, provider, dtype + ) + ) + ).run(print_data=True) diff --git a/tests/compile/passes/ir/test_lowering.py b/tests/compile/passes/ir/test_lowering.py index b7ca55e7d1..0ac605e796 100644 --- a/tests/compile/passes/ir/test_lowering.py +++ b/tests/compile/passes/ir/test_lowering.py @@ -11,6 +11,7 @@ from vllm.compilation.passes.ir.lowering_pass import ( ) from vllm.config import get_current_vllm_config from vllm.ir import ops +from vllm.model_executor.layers.activation import GeluAndMulSparse from vllm.platforms import current_platform from ...backend import TestBackend @@ -34,6 +35,15 @@ class Model(nn.Module): return x6 + 3.0 +class SparseActivationModel(nn.Module): + def __init__(self): + super().__init__() + self.activation = GeluAndMulSparse(0.95, "tanh") + + def forward(self, x): + return self.activation(x) + + @pytest.mark.parametrize("rms_provider", ops.rms_norm.supported_providers()) def test_lowering_rms_norm(rms_provider, default_vllm_config): torch.set_default_device(current_platform.device_type) @@ -67,3 +77,27 @@ def test_lowering_rms_norm(rms_provider, default_vllm_config): torch.testing.assert_close(output_unlowered, output) torch.testing.assert_close(output_unlowered, output2) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider") +def test_lowering_gelu_and_mul_sparse(default_vllm_config): + torch.set_default_device("cuda") + lowering_pass = VllmIRLoweringPass(get_current_vllm_config()) + backend = TestBackend(lowering_pass) + model = SparseActivationModel() + x = torch.randn(8, 32768, dtype=torch.bfloat16) + + with ( + ops.gelu_and_mul_sparse.set_priority(["triton", "native"]), + ir.enable_torch_wrap(True), + ): + compiled_model = torch.compile(model, backend=backend, fullgraph=True) + output = compiled_model(x) + + selected = lowering_pass.selected_impls["gelu_and_mul_sparse"] + assert selected == {"gelu_and_mul_sparse": "triton"} + + expected = ops.gelu_and_mul_sparse.impls["native"].impl_fn( + x, 1.6448536269514722, "tanh" + ) + torch.testing.assert_close(output, expected, rtol=1.6e-2, atol=1e-3) diff --git a/tests/compile/test_aot_compile.py b/tests/compile/test_aot_compile.py index b3a71b4c30..90c132bee1 100644 --- a/tests/compile/test_aot_compile.py +++ b/tests/compile/test_aot_compile.py @@ -13,6 +13,8 @@ import pytest import torch import vllm.envs as envs +import vllm.kernels # noqa: F401 to register kernels +from vllm import ir from vllm.compilation.backends import VllmBackend from vllm.compilation.caching import ( StandaloneCompiledArtifacts, @@ -28,6 +30,9 @@ from vllm.config import ( ) from vllm.envs import disable_envs_cache from vllm.forward_context import set_forward_context +from vllm.ir import ops +from vllm.model_executor.layers.activation import GeluAndMulSparse +from vllm.platforms import current_platform from vllm.utils.torch_utils import is_torch_equal_or_newer from ..utils import create_new_process_for_each_test @@ -77,6 +82,18 @@ class CompiledModTuple(torch.nn.Module): return reference_fn_tuple(x) +@support_torch_compile(dynamic_arg_dims={"x": 0}) +class SparseActivationAotModel(torch.nn.Module): + def __init__(self, **kwargs): + super().__init__() + self.up_proj = torch.nn.Linear(32, 256, bias=False) + self.activation = GeluAndMulSparse(0.95, "tanh") + self.down_proj = torch.nn.Linear(128, 16, bias=False) + + def forward(self, x: torch.Tensor): + return self.down_proj(self.activation(self.up_proj(x))) + + def make_vllm_config() -> VllmConfig: return VllmConfig( compilation_config=CompilationConfig( @@ -175,6 +192,54 @@ def test_save_and_load(monkeypatch: pytest.MonkeyPatch): assert torch.allclose(ret, expected) +@pytest.mark.skipif( + not current_platform.is_cuda() or not is_torch_equal_or_newer("2.10.0"), + reason="requires CUDA and torch 2.10", +) +def test_sparse_activation_standalone_aot_preserves_dependencies( + monkeypatch: pytest.MonkeyPatch, vllm_tmp_cache: Path +): + def native_reference(model: SparseActivationAotModel, x: torch.Tensor): + projected = torch.nn.functional.linear(x, model.up_proj.weight) + activated = ops.gelu_and_mul_sparse.impls["native"].impl_fn( + projected, model.activation.std_multiplier, "tanh" + ) + return torch.nn.functional.linear(activated, model.down_proj.weight) + + with monkeypatch.context() as m: + m.setenv("VLLM_USE_AOT_COMPILE", "1") + m.setenv("VLLM_USE_MEGA_AOT_ARTIFACT", "1") + m.setenv("VLLM_USE_STANDALONE_COMPILE", "1") + m.setenv("VLLM_CACHE_ROOT", str(vllm_tmp_cache / "vllm_cache")) + + vllm_config = make_vllm_config() + torch.manual_seed(0) + inputs = [ + torch.randn(4, 32, device="cuda", dtype=torch.bfloat16), + torch.randn(7, 32, device="cuda", dtype=torch.bfloat16), + ] + + with ( + use_vllm_config(vllm_config), + vllm_config.kernel_config.ir_op_priority.set_priority(), + ir.enable_torch_wrap(True), + ): + compiled = SparseActivationAotModel(vllm_config=vllm_config).to( + device="cuda", dtype=torch.bfloat16 + ) + compiled.requires_grad_(False) + expected = [native_reference(compiled, x) for x in inputs] + actual = [compiled(x) for x in inputs] + + for result, reference in zip(actual, expected): + assert torch.isfinite(result).all() + torch.testing.assert_close( + result, + reference, + **ops.gelu_and_mul_sparse.get_tolerance(result.dtype), + ) + + @pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10") def test_save_and_load_slice(monkeypatch: pytest.MonkeyPatch): from torch._subclasses import FakeTensorMode diff --git a/tests/kernels/ir/test_activation.py b/tests/kernels/ir/test_activation.py new file mode 100644 index 0000000000..a36ff8ee33 --- /dev/null +++ b/tests/kernels/ir/test_activation.py @@ -0,0 +1,183 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +import vllm.kernels # noqa: F401 +from tests.ir.ir_test_utils import assert_close, clone_args +from vllm import ir +from vllm.platforms import current_platform +from vllm.triton_utils import HAS_TRITON + +gelu_and_mul_sparse_native = ir.ops.gelu_and_mul_sparse.impls["native"].impl_fn + + +def test_gelu_and_mul_sparse_registration(): + expected = { + "native": True, + "triton": HAS_TRITON and current_platform.is_cuda(), + } + actual = { + provider: impl.supported + for provider, impl in ir.ops.gelu_and_mul_sparse.impls.items() + } + assert actual == expected + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider") +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("num_tokens", [1, 8, 64]) +@pytest.mark.parametrize("intermediate_size", [2048, 4096, 8192, 13824, 16384, 32768]) +def test_gelu_and_mul_sparse_triton( + dtype: torch.dtype, num_tokens: int, intermediate_size: int +): + torch.manual_seed(0) + torch.set_default_device("cuda") + args = ir.ops.gelu_and_mul_sparse.generate_inputs( + num_tokens=num_tokens, + hidden_size=intermediate_size, + dtype=dtype, + ) + impl = ir.ops.gelu_and_mul_sparse.impls["triton"] + assert impl.supports_args(*args) + + expected = gelu_and_mul_sparse_native(*clone_args(args)) + actual = impl.impl_fn(*clone_args(args)) + assert_close(ir.ops.gelu_and_mul_sparse, actual, expected) + + mask_mismatch = torch.count_nonzero((actual == 0) != (expected == 0)) + assert mask_mismatch / actual.numel() < 1e-3 + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider") +def test_gelu_and_mul_sparse_triton_second_seed(): + torch.manual_seed(1) + torch.set_default_device("cuda") + args = ir.ops.gelu_and_mul_sparse.generate_inputs( + num_tokens=32, + hidden_size=8192, + dtype=torch.bfloat16, + ) + + actual = ir.ops.gelu_and_mul_sparse.impls["triton"].impl_fn(*clone_args(args)) + expected = gelu_and_mul_sparse_native(*clone_args(args)) + assert_close(ir.ops.gelu_and_mul_sparse, actual, expected) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider") +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("intermediate_size", [1, 31, 128]) +def test_gelu_and_mul_sparse_triton_small_sizes( + dtype: torch.dtype, intermediate_size: int +): + torch.manual_seed(0) + torch.set_default_device("cuda") + args = ir.ops.gelu_and_mul_sparse.generate_inputs( + num_tokens=5, + hidden_size=intermediate_size, + dtype=dtype, + ) + + actual = ir.ops.gelu_and_mul_sparse.impls["triton"].impl_fn(*clone_args(args)) + expected = gelu_and_mul_sparse_native(*clone_args(args)) + assert_close(ir.ops.gelu_and_mul_sparse, actual, expected) + assert torch.equal(actual == 0, expected == 0) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider") +def test_gelu_and_mul_sparse_triton_nd_and_empty(): + torch.set_default_device("cuda") + impl = ir.ops.gelu_and_mul_sparse.impls["triton"] + + x = torch.randn(2, 3, 16384, dtype=torch.bfloat16) + args = (x, 1.6448536269514722, "tanh") + actual = impl.impl_fn(*args) + expected = gelu_and_mul_sparse_native(*args) + assert actual.shape == (2, 3, 8192) + assert_close(ir.ops.gelu_and_mul_sparse, actual, expected) + + empty = torch.empty(0, 16384, dtype=torch.bfloat16) + empty_out = impl.impl_fn(empty, 1.6448536269514722, "tanh") + assert empty_out.shape == (0, 8192) + assert empty_out.dtype == empty.dtype + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider") +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +def test_gelu_and_mul_sparse_triton_constant_rows(dtype: torch.dtype): + torch.set_default_device("cuda") + gate = torch.full((4, 8192), 3.0, dtype=dtype) + up = torch.randn_like(gate) + x = torch.cat((gate, up), dim=-1) + args = (x, 1.6448536269514722, "tanh") + + actual = ir.ops.gelu_and_mul_sparse.impls["triton"].impl_fn(*args) + expected = gelu_and_mul_sparse_native(*args) + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0) + assert torch.count_nonzero(actual) == 0 + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider") +@pytest.mark.parametrize("nonfinite", [float("nan"), float("inf")]) +def test_gelu_and_mul_sparse_triton_nonfinite_gate(nonfinite: float): + torch.set_default_device("cuda") + x = torch.randn(2, 16384, dtype=torch.bfloat16) + x[0, 0] = nonfinite + args = (x, 1.6448536269514722, "tanh") + + actual = ir.ops.gelu_and_mul_sparse.impls["triton"].impl_fn(*args) + expected = gelu_and_mul_sparse_native(*args) + torch.testing.assert_close(actual, expected, equal_nan=True) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider") +def test_gelu_and_mul_sparse_triton_fallback_conditions(): + torch.set_default_device("cuda") + impl = ir.ops.gelu_and_mul_sparse.impls["triton"] + x = torch.randn(4, 16384, dtype=torch.bfloat16) + + assert impl.supports_args(x, 1.0, "tanh") + assert not impl.supports_args(x, 1.0, "none") + assert not impl.supports_args(x[:, ::2], 1.0, "tanh") + assert not impl.supports_args(x[:, :-1], 1.0, "tanh") + assert not impl.supports_args( + torch.randn(1, 65538, dtype=torch.bfloat16), 1.0, "tanh" + ) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider") +def test_gelu_and_mul_sparse_triton_dispatch_and_opcheck(): + torch.set_default_device("cuda") + args = ir.ops.gelu_and_mul_sparse.generate_inputs( + num_tokens=8, + hidden_size=4096, + dtype=torch.bfloat16, + ) + impl = ir.ops.gelu_and_mul_sparse.impls["triton"] + + with ir.ops.gelu_and_mul_sparse.set_priority(["triton", "native"]): + dispatched = ir.ops.gelu_and_mul_sparse(*args) + torch.library.opcheck(torch.ops.vllm_ir.gelu_and_mul_sparse, args) + + direct = impl.impl_fn(*args) + torch.testing.assert_close(dispatched, direct, rtol=0.0, atol=0.0) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider") +def test_gelu_and_mul_sparse_triton_cudagraph(): + torch.set_default_device("cuda") + x = torch.randn(16, 32768, dtype=torch.bfloat16) + impl = ir.ops.gelu_and_mul_sparse.impls["triton"] + + for _ in range(3): + expected = impl.impl_fn(x, 1.6448536269514722, "tanh") + torch.accelerator.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = impl.impl_fn(x, 1.6448536269514722, "tanh") + graph.replay() + torch.accelerator.synchronize() + + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0) diff --git a/tests/test_config.py b/tests/test_config.py index c2629a75ab..4d39b56d26 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -3322,6 +3322,22 @@ def test_ir_op_priority_default(): assert priority_config.fused_add_rms_norm == ["native"] +@pytest.mark.parametrize("mode", [CompilationMode.NONE, CompilationMode.VLLM_COMPILE]) +@pytest.mark.parametrize("backend", ["inductor", "eager"]) +def test_ir_op_platform_defaults_support_sparse_gelu(mode, backend): + """Worker initialization must not select an unregistered sparse GELU provider.""" + from vllm import ir + + config = SimpleNamespace( + compilation_config=CompilationConfig(mode=mode, backend=backend) + ) + priority = current_platform.get_default_ir_op_priority(config) + expected = ["triton", "native"] if current_platform.is_cuda() else ["native"] + + with priority.set_priority(): + assert ir.ops.gelu_and_mul_sparse.get_priority() == expected + + def test_ir_op_priority_str(): """Test that passing a comma-delimited string works""" from vllm.config.kernel import IrOpPriorityConfig diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index d44e83bf67..4505e0cfd9 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -34,6 +34,9 @@ class IrOpPriorityConfig: fused_add_rms_norm: list[str] = Field(default_factory=list) """Priority list for vllm.ir.ops.fused_add_rms_norm""" + gelu_and_mul_sparse: list[str] = Field(default_factory=list) + """Priority list for vllm.ir.ops.gelu_and_mul_sparse""" + def compute_hash(self) -> str: """ Produces a hash unique to the pass configuration. diff --git a/vllm/ir/ops/__init__.py b/vllm/ir/ops/__init__.py index d4d71afef7..13e2b7dda9 100644 --- a/vllm/ir/ops/__init__.py +++ b/vllm/ir/ops/__init__.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from .activation import gelu_and_mul_sparse from .layernorm import fused_add_rms_norm, rms_norm -__all__ = ["rms_norm", "fused_add_rms_norm"] +__all__ = ["rms_norm", "fused_add_rms_norm", "gelu_and_mul_sparse"] diff --git a/vllm/ir/ops/activation.py b/vllm/ir/ops/activation.py new file mode 100644 index 0000000000..149197d735 --- /dev/null +++ b/vllm/ir/ops/activation.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +import torch.nn.functional as F +from torch import Tensor + +from ..op import register_op + + +@register_op +def gelu_and_mul_sparse( + x: Tensor, std_multiplier: float, approximate: str = "none" +) -> Tensor: + """Apply Gaussian sparsification, GELU, and gated multiplication.""" + d = x.shape[-1] // 2 + gate = x[..., :d] + # Statistics intentionally remain local to each tensor-parallel shard. + mean = torch.mean(gate, dim=-1, keepdim=True) + std = torch.std(gate, dim=-1, keepdim=True, unbiased=False) + sparse_gate = F.relu(gate - (mean + std * std_multiplier)) + return F.gelu(sparse_gate, approximate=approximate) * x[..., d:] + + +@gelu_and_mul_sparse.register_input_generator +def _gelu_and_mul_sparse_input_generator( + num_tokens: int, + hidden_size: int, + dtype: torch.dtype, + std_multiplier: float = 1.6448536269514722, + approximate: str = "tanh", +) -> tuple: + x = torch.randn(num_tokens, hidden_size * 2, dtype=dtype) + return x, std_multiplier, approximate diff --git a/vllm/kernels/__init__.py b/vllm/kernels/__init__.py index e83e2772f0..4821ec43ec 100644 --- a/vllm/kernels/__init__.py +++ b/vllm/kernels/__init__.py @@ -2,6 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Kernel implementations for vLLM.""" -from . import aiter_ops, oink_ops, vllm_c +from . import aiter_ops, oink_ops, triton, vllm_c -__all__ = ["vllm_c", "aiter_ops", "oink_ops"] +__all__ = ["vllm_c", "aiter_ops", "oink_ops", "triton"] diff --git a/vllm/kernels/triton/__init__.py b/vllm/kernels/triton/__init__.py index 6626213145..3a38335050 100644 --- a/vllm/kernels/triton/__init__.py +++ b/vllm/kernels/triton/__init__.py @@ -1,3 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Triton kernel implementations.""" + +from . import activation + +__all__ = ["activation"] diff --git a/vllm/kernels/triton/activation.py b/vllm/kernels/triton/activation.py new file mode 100644 index 0000000000..2bda1f51b4 --- /dev/null +++ b/vllm/kernels/triton/activation.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +from torch import Tensor +from torch.library import triton_op, wrap_triton + +from vllm import ir +from vllm.platforms import current_platform +from vllm.triton_utils import HAS_TRITON, tl, triton + +_MAX_INTERMEDIATE_SIZE = 32768 +_SUPPORTED_DTYPES = {torch.float16, torch.bfloat16, torch.float32} + + +@triton.jit +def _gelu_and_mul_sparse_kernel( + x_ptr, + output_ptr, + d: tl.constexpr, + std_multiplier, + BLOCK_SIZE: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + offsets = tl.arange(0, BLOCK_SIZE) + mask = offsets < d + row_start = row * 2 * d + + gate = tl.load(x_ptr + row_start + offsets, mask=mask, other=0.0).to(tl.float32) + # TODO: Investigate delaying this load to reduce spills for large reductions + # without regressing smaller shapes. + up = tl.load(x_ptr + row_start + d + offsets, mask=mask, other=0.0).to(tl.float32) + + input_dtype = x_ptr.dtype.element_ty + mean = tl.sum(gate, axis=0) / d + centered = tl.where(mask, gate - mean, 0.0) + variance = tl.sum(centered * centered, axis=0) / d + # Match the input-dtype rounding at each native PyTorch op boundary. + mean = mean.to(input_dtype).to(tl.float32) + std = tl.sqrt(variance).to(input_dtype).to(tl.float32) + scaled_std = (std * std_multiplier).to(input_dtype).to(tl.float32) + cutoff = (mean + scaled_std).to(input_dtype).to(tl.float32) + sparse_gate = (gate - cutoff).to(input_dtype).to(tl.float32) + sparse_gate = tl.where(sparse_gate < 0.0, 0.0, sparse_gate) + + inner = 0.7978845608028654 * ( + sparse_gate + 0.044715 * sparse_gate * sparse_gate * sparse_gate + ) + tanh = 2.0 * tl.sigmoid(2.0 * inner) - 1.0 + activated = 0.5 * sparse_gate * (1.0 + tanh) + activated = activated.to(input_dtype).to(tl.float32) + tl.store(output_ptr + row * d + offsets, activated * up, mask=mask) + + +def _supports_gelu_and_mul_sparse( + x: Tensor, std_multiplier: float, approximate: str = "none" +) -> bool: + del std_multiplier + if ( + x.device.type != "cuda" + or x.dtype not in _SUPPORTED_DTYPES + or not x.is_contiguous() + or x.ndim == 0 + or approximate != "tanh" + ): + return False + last_dim = x.shape[-1] + return ( + last_dim > 0 and last_dim % 2 == 0 and last_dim // 2 <= _MAX_INTERMEDIATE_SIZE + ) + + +@triton_op( + "vllm::gelu_and_mul_sparse_triton", + mutates_args=(), +) +def _gelu_and_mul_sparse_triton_op(x: Tensor, std_multiplier: float) -> Tensor: + d = x.shape[-1] // 2 + output = torch.empty((*x.shape[:-1], d), dtype=x.dtype, device=x.device) + if output.numel() == 0: + return output + + n_rows = x.numel() // x.shape[-1] + + def grid(meta): + return (n_rows,) + + block_size = triton.next_power_of_2(d) + num_warps = min(max(block_size // 256, 1), 8) + wrap_triton(_gelu_and_mul_sparse_kernel)[grid]( + x, + output, + d=d, + std_multiplier=std_multiplier, + BLOCK_SIZE=block_size, + num_warps=num_warps, + ) + return output + + +@ir.ops.gelu_and_mul_sparse.register_impl( + "triton", + supported=HAS_TRITON and current_platform.is_cuda(), + supports_args=_supports_gelu_and_mul_sparse, +) +def gelu_and_mul_sparse_triton( + x: Tensor, std_multiplier: float, approximate: str = "none" +) -> Tensor: + assert approximate == "tanh" + return _gelu_and_mul_sparse_triton_op(x, std_multiplier) diff --git a/vllm/model_executor/layers/activation.py b/vllm/model_executor/layers/activation.py index 159c521cee..f159d2dff8 100644 --- a/vllm/model_executor/layers/activation.py +++ b/vllm/model_executor/layers/activation.py @@ -8,6 +8,7 @@ import torch import torch.nn as nn import torch.nn.functional as F +from vllm import ir from vllm.logger import init_logger from vllm.model_executor.custom_op import CustomOp from vllm.platforms import CpuArchEnum, current_platform @@ -332,24 +333,11 @@ class GeluAndMulSparse(CustomOp): raise ValueError("activation_sparsity is 0.0. Please use GeluAndMul.") target_sparsity_tensor = torch.tensor(activation_sparsity, dtype=torch.float32) normal_dist = torch.distributions.normal.Normal(0, 1) - self.std_multiplier = normal_dist.icdf(target_sparsity_tensor) - - def _gaussian_topk(self, x: torch.Tensor) -> torch.Tensor: - """Get % sparse percentile of the Gaussian distribution.""" - # NOTE(rob): for TP>1, we could all-gather to get the means/std. - # But we do not do this because in expectation they are the same - # and in practice the eval scores are good without gathering. - mean = torch.mean(x, dim=-1, keepdim=True) - std = torch.std(x, dim=-1, keepdim=True, unbiased=False) - cutoff_x = mean + std * self.std_multiplier - return nn.functional.relu(x - cutoff_x) + self.std_multiplier = normal_dist.icdf(target_sparsity_tensor).item() def forward_native(self, x: torch.Tensor) -> torch.Tensor: """PyTorch-native implementation equivalent to forward().""" - d = x.shape[-1] // 2 - out = self._gaussian_topk(x[..., :d]) - out = F.gelu(out, approximate=self.approximate) - return out * x[..., d:] + return ir.ops.gelu_and_mul_sparse(x, self.std_multiplier, self.approximate) def forward_cuda(self, x: torch.Tensor) -> torch.Tensor: return self.forward_native(x) diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 6d9444567a..98bddbde29 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -755,7 +755,10 @@ class CudaPlatformBase(Platform): rms_norm = ["oink"] + default return IrOpPriorityConfig.with_default( - default, rms_norm=rms_norm, fused_add_rms_norm=rms_norm + default, + rms_norm=rms_norm, + fused_add_rms_norm=rms_norm, + gelu_and_mul_sparse=["triton", "native"], ) @classmethod diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index a9cf534716..b3b7fe8b08 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -1161,7 +1161,10 @@ class RocmPlatform(Platform): rms_norm = default return IrOpPriorityConfig.with_default( - default, rms_norm=rms_norm, fused_add_rms_norm=rms_norm + default, + rms_norm=rms_norm, + fused_add_rms_norm=rms_norm, + gelu_and_mul_sparse=["native"], ) @classmethod diff --git a/vllm/platforms/xpu.py b/vllm/platforms/xpu.py index 8cb29fbdbd..6920c55e71 100644 --- a/vllm/platforms/xpu.py +++ b/vllm/platforms/xpu.py @@ -491,7 +491,7 @@ class XPUPlatform(Platform): using_inductor = cc.backend == "inductor" and cc.mode != CompilationMode.NONE default = ["native"] if using_inductor else ["vllm_c", "native"] - return IrOpPriorityConfig.with_default(default) + return IrOpPriorityConfig.with_default(default, gelu_and_mul_sparse=["native"]) @classmethod def device_count(cls) -> int: