[Performance] Add Triton kernel for Gemma3n sparse GELU (#48498)

Signed-off-by: BWAAEEEK <jooho414@gmail.com>
Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>

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)

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)

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

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)

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

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.

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"]

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

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"]

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"]

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)

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)

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

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

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: