@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 78e84261ab

[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>

JooHo Lee committed Sep 14, 2026 at 16:42 UTC 78e84261ab2ffb691a8079d698d1f3f96ecb7635
15 files changed +555 -21
benchmarks/kernels/benchmark_gelu_and_mul_sparse.py new
+90
@@ -0,0 +1,90 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import itertools
5 +
6 +import torch
7 +
8 +import vllm.kernels # noqa: F401
9 +from vllm import ir
10 +from vllm.benchmarks.lib.utils import default_vllm_config
11 +from vllm.triton_utils import triton
12 +from vllm.utils.argparse_utils import FlexibleArgumentParser
13 +from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed
14 +
15 +NUM_TOKENS = [1, 2, 8, 32, 128, 512, 2048]
16 +INTERMEDIATE_SIZES = [2048, 4096, 8192, 13824, 16384, 32768]
17 +CONFIGS = list(itertools.product(NUM_TOKENS, INTERMEDIATE_SIZES))
18 +STD_MULTIPLIER = 1.6448536269514722
19 +
20 +
21 +@default_vllm_config()
22 +def benchmark_gelu_and_mul_sparse(
23 + num_tokens: int,
24 + intermediate_size: int,
25 + provider: str,
26 + dtype: torch.dtype,
27 +):
28 + set_random_seed(0)
29 + torch.set_default_device("cuda")
30 + x = torch.randn(num_tokens, 2 * intermediate_size, dtype=dtype)
31 + native = ir.ops.gelu_and_mul_sparse.impls["native"].impl_fn
32 + triton_impl = ir.ops.gelu_and_mul_sparse.impls["triton"].impl_fn
33 + args = (x, STD_MULTIPLIER, "tanh")
34 +
35 + expected = native(*args)
36 + actual = triton_impl(*args)
37 + torch.testing.assert_close(
38 + actual,
39 + expected,
40 + **ir.ops.gelu_and_mul_sparse.get_tolerance(actual.dtype),
41 + )
42 +
43 + if provider == "triton":
44 + fn = lambda: triton_impl(*args)
45 + elif provider == "native_compiled":
46 + compiled = torch.compile(native, fullgraph=True)
47 + compiled(*args)
48 + fn = lambda: compiled(*args)
49 + else:
50 + fn = lambda: native(*args)
51 +
52 + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
53 + fn, quantiles=[0.5, 0.2, 0.8]
54 + )
55 + return ms, max_ms, min_ms
56 +
57 +
58 +if __name__ == "__main__":
59 + parser = FlexibleArgumentParser(
60 + description="Benchmark the Gemma3n sparse GELU activation."
61 + )
62 + parser.add_argument(
63 + "--dtype",
64 + type=str,
65 + choices=["half", "bfloat16", "float"],
66 + default="bfloat16",
67 + )
68 + args = parser.parse_args()
69 + dtype = STR_DTYPE_TO_TORCH_DTYPE[args.dtype]
70 +
71 + perf_report = triton.testing.perf_report(
72 + triton.testing.Benchmark(
73 + x_names=["num_tokens", "intermediate_size"],
74 + x_vals=CONFIGS,
75 + line_arg="provider",
76 + line_vals=["triton", "native_compiled", "native_eager"],
77 + line_names=["Triton", "Native (torch.compile)", "Native (eager)"],
78 + styles=[("blue", "-"), ("green", "-"), ("red", "-")],
79 + ylabel="ms",
80 + plot_name="gelu-and-mul-sparse-performance",
81 + args={},
82 + )
83 + )
84 + perf_report(
85 + lambda num_tokens, intermediate_size, provider: (
86 + benchmark_gelu_and_mul_sparse(
87 + num_tokens, intermediate_size, provider, dtype
88 + )
89 + )
90 + ).run(print_data=True)
tests/compile/passes/ir/test_lowering.py
+34
@@ -11,6 +11,7 @@ from vllm.compilation.passes.ir.lowering_pass import (
11 )
12 from vllm.config import get_current_vllm_config
13 from vllm.ir import ops
14 +from vllm.model_executor.layers.activation import GeluAndMulSparse
15 from vllm.platforms import current_platform
16
17 from ...backend import TestBackend
@@ -34,6 +35,15 @@ class Model(nn.Module):
35 return x6 + 3.0
36
37
38 +class SparseActivationModel(nn.Module):
39 + def __init__(self):
40 + super().__init__()
41 + self.activation = GeluAndMulSparse(0.95, "tanh")
42 +
43 + def forward(self, x):
44 + return self.activation(x)
45 +
46 +
47 @pytest.mark.parametrize("rms_provider", ops.rms_norm.supported_providers())
48 def test_lowering_rms_norm(rms_provider, default_vllm_config):
49 torch.set_default_device(current_platform.device_type)
@@ -67,3 +77,27 @@ def test_lowering_rms_norm(rms_provider, default_vllm_config):
77
78 torch.testing.assert_close(output_unlowered, output)
79 torch.testing.assert_close(output_unlowered, output2)
80 +
81 +
82 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider")
83 +def test_lowering_gelu_and_mul_sparse(default_vllm_config):
84 + torch.set_default_device("cuda")
85 + lowering_pass = VllmIRLoweringPass(get_current_vllm_config())
86 + backend = TestBackend(lowering_pass)
87 + model = SparseActivationModel()
88 + x = torch.randn(8, 32768, dtype=torch.bfloat16)
89 +
90 + with (
91 + ops.gelu_and_mul_sparse.set_priority(["triton", "native"]),
92 + ir.enable_torch_wrap(True),
93 + ):
94 + compiled_model = torch.compile(model, backend=backend, fullgraph=True)
95 + output = compiled_model(x)
96 +
97 + selected = lowering_pass.selected_impls["gelu_and_mul_sparse"]
98 + assert selected == {"gelu_and_mul_sparse": "triton"}
99 +
100 + expected = ops.gelu_and_mul_sparse.impls["native"].impl_fn(
101 + x, 1.6448536269514722, "tanh"
102 + )
103 + torch.testing.assert_close(output, expected, rtol=1.6e-2, atol=1e-3)
tests/compile/test_aot_compile.py
+65
@@ -13,6 +13,8 @@ import pytest
13 import torch
14
15 import vllm.envs as envs
16 +import vllm.kernels # noqa: F401 to register kernels
17 +from vllm import ir
18 from vllm.compilation.backends import VllmBackend
19 from vllm.compilation.caching import (
20 StandaloneCompiledArtifacts,
@@ -28,6 +30,9 @@ from vllm.config import (
30 )
31 from vllm.envs import disable_envs_cache
32 from vllm.forward_context import set_forward_context
33 +from vllm.ir import ops
34 +from vllm.model_executor.layers.activation import GeluAndMulSparse
35 +from vllm.platforms import current_platform
36 from vllm.utils.torch_utils import is_torch_equal_or_newer
37
38 from ..utils import create_new_process_for_each_test
@@ -77,6 +82,18 @@ class CompiledModTuple(torch.nn.Module):
82 return reference_fn_tuple(x)
83
84
85 +@support_torch_compile(dynamic_arg_dims={"x": 0})
86 +class SparseActivationAotModel(torch.nn.Module):
87 + def __init__(self, **kwargs):
88 + super().__init__()
89 + self.up_proj = torch.nn.Linear(32, 256, bias=False)
90 + self.activation = GeluAndMulSparse(0.95, "tanh")
91 + self.down_proj = torch.nn.Linear(128, 16, bias=False)
92 +
93 + def forward(self, x: torch.Tensor):
94 + return self.down_proj(self.activation(self.up_proj(x)))
95 +
96 +
97 def make_vllm_config() -> VllmConfig:
98 return VllmConfig(
99 compilation_config=CompilationConfig(
@@ -175,6 +192,54 @@ def test_save_and_load(monkeypatch: pytest.MonkeyPatch):
192 assert torch.allclose(ret, expected)
193
194
195 +@pytest.mark.skipif(
196 + not current_platform.is_cuda() or not is_torch_equal_or_newer("2.10.0"),
197 + reason="requires CUDA and torch 2.10",
198 +)
199 +def test_sparse_activation_standalone_aot_preserves_dependencies(
200 + monkeypatch: pytest.MonkeyPatch, vllm_tmp_cache: Path
201 +):
202 + def native_reference(model: SparseActivationAotModel, x: torch.Tensor):
203 + projected = torch.nn.functional.linear(x, model.up_proj.weight)
204 + activated = ops.gelu_and_mul_sparse.impls["native"].impl_fn(
205 + projected, model.activation.std_multiplier, "tanh"
206 + )
207 + return torch.nn.functional.linear(activated, model.down_proj.weight)
208 +
209 + with monkeypatch.context() as m:
210 + m.setenv("VLLM_USE_AOT_COMPILE", "1")
211 + m.setenv("VLLM_USE_MEGA_AOT_ARTIFACT", "1")
212 + m.setenv("VLLM_USE_STANDALONE_COMPILE", "1")
213 + m.setenv("VLLM_CACHE_ROOT", str(vllm_tmp_cache / "vllm_cache"))
214 +
215 + vllm_config = make_vllm_config()
216 + torch.manual_seed(0)
217 + inputs = [
218 + torch.randn(4, 32, device="cuda", dtype=torch.bfloat16),
219 + torch.randn(7, 32, device="cuda", dtype=torch.bfloat16),
220 + ]
221 +
222 + with (
223 + use_vllm_config(vllm_config),
224 + vllm_config.kernel_config.ir_op_priority.set_priority(),
225 + ir.enable_torch_wrap(True),
226 + ):
227 + compiled = SparseActivationAotModel(vllm_config=vllm_config).to(
228 + device="cuda", dtype=torch.bfloat16
229 + )
230 + compiled.requires_grad_(False)
231 + expected = [native_reference(compiled, x) for x in inputs]
232 + actual = [compiled(x) for x in inputs]
233 +
234 + for result, reference in zip(actual, expected):
235 + assert torch.isfinite(result).all()
236 + torch.testing.assert_close(
237 + result,
238 + reference,
239 + **ops.gelu_and_mul_sparse.get_tolerance(result.dtype),
240 + )
241 +
242 +
243 @pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
244 def test_save_and_load_slice(monkeypatch: pytest.MonkeyPatch):
245 from torch._subclasses import FakeTensorMode
tests/kernels/ir/test_activation.py new
+183
@@ -0,0 +1,183 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +import torch
6 +
7 +import vllm.kernels # noqa: F401
8 +from tests.ir.ir_test_utils import assert_close, clone_args
9 +from vllm import ir
10 +from vllm.platforms import current_platform
11 +from vllm.triton_utils import HAS_TRITON
12 +
13 +gelu_and_mul_sparse_native = ir.ops.gelu_and_mul_sparse.impls["native"].impl_fn
14 +
15 +
16 +def test_gelu_and_mul_sparse_registration():
17 + expected = {
18 + "native": True,
19 + "triton": HAS_TRITON and current_platform.is_cuda(),
20 + }
21 + actual = {
22 + provider: impl.supported
23 + for provider, impl in ir.ops.gelu_and_mul_sparse.impls.items()
24 + }
25 + assert actual == expected
26 +
27 +
28 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider")
29 +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
30 +@pytest.mark.parametrize("num_tokens", [1, 8, 64])
31 +@pytest.mark.parametrize("intermediate_size", [2048, 4096, 8192, 13824, 16384, 32768])
32 +def test_gelu_and_mul_sparse_triton(
33 + dtype: torch.dtype, num_tokens: int, intermediate_size: int
34 +):
35 + torch.manual_seed(0)
36 + torch.set_default_device("cuda")
37 + args = ir.ops.gelu_and_mul_sparse.generate_inputs(
38 + num_tokens=num_tokens,
39 + hidden_size=intermediate_size,
40 + dtype=dtype,
41 + )
42 + impl = ir.ops.gelu_and_mul_sparse.impls["triton"]
43 + assert impl.supports_args(*args)
44 +
45 + expected = gelu_and_mul_sparse_native(*clone_args(args))
46 + actual = impl.impl_fn(*clone_args(args))
47 + assert_close(ir.ops.gelu_and_mul_sparse, actual, expected)
48 +
49 + mask_mismatch = torch.count_nonzero((actual == 0) != (expected == 0))
50 + assert mask_mismatch / actual.numel() < 1e-3
51 +
52 +
53 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider")
54 +def test_gelu_and_mul_sparse_triton_second_seed():
55 + torch.manual_seed(1)
56 + torch.set_default_device("cuda")
57 + args = ir.ops.gelu_and_mul_sparse.generate_inputs(
58 + num_tokens=32,
59 + hidden_size=8192,
60 + dtype=torch.bfloat16,
61 + )
62 +
63 + actual = ir.ops.gelu_and_mul_sparse.impls["triton"].impl_fn(*clone_args(args))
64 + expected = gelu_and_mul_sparse_native(*clone_args(args))
65 + assert_close(ir.ops.gelu_and_mul_sparse, actual, expected)
66 +
67 +
68 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider")
69 +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
70 +@pytest.mark.parametrize("intermediate_size", [1, 31, 128])
71 +def test_gelu_and_mul_sparse_triton_small_sizes(
72 + dtype: torch.dtype, intermediate_size: int
73 +):
74 + torch.manual_seed(0)
75 + torch.set_default_device("cuda")
76 + args = ir.ops.gelu_and_mul_sparse.generate_inputs(
77 + num_tokens=5,
78 + hidden_size=intermediate_size,
79 + dtype=dtype,
80 + )
81 +
82 + actual = ir.ops.gelu_and_mul_sparse.impls["triton"].impl_fn(*clone_args(args))
83 + expected = gelu_and_mul_sparse_native(*clone_args(args))
84 + assert_close(ir.ops.gelu_and_mul_sparse, actual, expected)
85 + assert torch.equal(actual == 0, expected == 0)
86 +
87 +
88 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider")
89 +def test_gelu_and_mul_sparse_triton_nd_and_empty():
90 + torch.set_default_device("cuda")
91 + impl = ir.ops.gelu_and_mul_sparse.impls["triton"]
92 +
93 + x = torch.randn(2, 3, 16384, dtype=torch.bfloat16)
94 + args = (x, 1.6448536269514722, "tanh")
95 + actual = impl.impl_fn(*args)
96 + expected = gelu_and_mul_sparse_native(*args)
97 + assert actual.shape == (2, 3, 8192)
98 + assert_close(ir.ops.gelu_and_mul_sparse, actual, expected)
99 +
100 + empty = torch.empty(0, 16384, dtype=torch.bfloat16)
101 + empty_out = impl.impl_fn(empty, 1.6448536269514722, "tanh")
102 + assert empty_out.shape == (0, 8192)
103 + assert empty_out.dtype == empty.dtype
104 +
105 +
106 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider")
107 +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
108 +def test_gelu_and_mul_sparse_triton_constant_rows(dtype: torch.dtype):
109 + torch.set_default_device("cuda")
110 + gate = torch.full((4, 8192), 3.0, dtype=dtype)
111 + up = torch.randn_like(gate)
112 + x = torch.cat((gate, up), dim=-1)
113 + args = (x, 1.6448536269514722, "tanh")
114 +
115 + actual = ir.ops.gelu_and_mul_sparse.impls["triton"].impl_fn(*args)
116 + expected = gelu_and_mul_sparse_native(*args)
117 + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0)
118 + assert torch.count_nonzero(actual) == 0
119 +
120 +
121 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider")
122 +@pytest.mark.parametrize("nonfinite", [float("nan"), float("inf")])
123 +def test_gelu_and_mul_sparse_triton_nonfinite_gate(nonfinite: float):
124 + torch.set_default_device("cuda")
125 + x = torch.randn(2, 16384, dtype=torch.bfloat16)
126 + x[0, 0] = nonfinite
127 + args = (x, 1.6448536269514722, "tanh")
128 +
129 + actual = ir.ops.gelu_and_mul_sparse.impls["triton"].impl_fn(*args)
130 + expected = gelu_and_mul_sparse_native(*args)
131 + torch.testing.assert_close(actual, expected, equal_nan=True)
132 +
133 +
134 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider")
135 +def test_gelu_and_mul_sparse_triton_fallback_conditions():
136 + torch.set_default_device("cuda")
137 + impl = ir.ops.gelu_and_mul_sparse.impls["triton"]
138 + x = torch.randn(4, 16384, dtype=torch.bfloat16)
139 +
140 + assert impl.supports_args(x, 1.0, "tanh")
141 + assert not impl.supports_args(x, 1.0, "none")
142 + assert not impl.supports_args(x[:, ::2], 1.0, "tanh")
143 + assert not impl.supports_args(x[:, :-1], 1.0, "tanh")
144 + assert not impl.supports_args(
145 + torch.randn(1, 65538, dtype=torch.bfloat16), 1.0, "tanh"
146 + )
147 +
148 +
149 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider")
150 +def test_gelu_and_mul_sparse_triton_dispatch_and_opcheck():
151 + torch.set_default_device("cuda")
152 + args = ir.ops.gelu_and_mul_sparse.generate_inputs(
153 + num_tokens=8,
154 + hidden_size=4096,
155 + dtype=torch.bfloat16,
156 + )
157 + impl = ir.ops.gelu_and_mul_sparse.impls["triton"]
158 +
159 + with ir.ops.gelu_and_mul_sparse.set_priority(["triton", "native"]):
160 + dispatched = ir.ops.gelu_and_mul_sparse(*args)
161 + torch.library.opcheck(torch.ops.vllm_ir.gelu_and_mul_sparse, args)
162 +
163 + direct = impl.impl_fn(*args)
164 + torch.testing.assert_close(dispatched, direct, rtol=0.0, atol=0.0)
165 +
166 +
167 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider")
168 +def test_gelu_and_mul_sparse_triton_cudagraph():
169 + torch.set_default_device("cuda")
170 + x = torch.randn(16, 32768, dtype=torch.bfloat16)
171 + impl = ir.ops.gelu_and_mul_sparse.impls["triton"]
172 +
173 + for _ in range(3):
174 + expected = impl.impl_fn(x, 1.6448536269514722, "tanh")
175 + torch.accelerator.synchronize()
176 +
177 + graph = torch.cuda.CUDAGraph()
178 + with torch.cuda.graph(graph):
179 + actual = impl.impl_fn(x, 1.6448536269514722, "tanh")
180 + graph.replay()
181 + torch.accelerator.synchronize()
182 +
183 + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0)
tests/test_config.py
+16
@@ -3322,6 +3322,22 @@ def test_ir_op_priority_default():
3322 assert priority_config.fused_add_rms_norm == ["native"]
3323
3324
3325 +@pytest.mark.parametrize("mode", [CompilationMode.NONE, CompilationMode.VLLM_COMPILE])
3326 +@pytest.mark.parametrize("backend", ["inductor", "eager"])
3327 +def test_ir_op_platform_defaults_support_sparse_gelu(mode, backend):
3328 + """Worker initialization must not select an unregistered sparse GELU provider."""
3329 + from vllm import ir
3330 +
3331 + config = SimpleNamespace(
3332 + compilation_config=CompilationConfig(mode=mode, backend=backend)
3333 + )
3334 + priority = current_platform.get_default_ir_op_priority(config)
3335 + expected = ["triton", "native"] if current_platform.is_cuda() else ["native"]
3336 +
3337 + with priority.set_priority():
3338 + assert ir.ops.gelu_and_mul_sparse.get_priority() == expected
3339 +
3340 +
3341 def test_ir_op_priority_str():
3342 """Test that passing a comma-delimited string works"""
3343 from vllm.config.kernel import IrOpPriorityConfig
vllm/config/kernel.py
+3
@@ -34,6 +34,9 @@ class IrOpPriorityConfig:
34 fused_add_rms_norm: list[str] = Field(default_factory=list)
35 """Priority list for vllm.ir.ops.fused_add_rms_norm"""
36
37 + gelu_and_mul_sparse: list[str] = Field(default_factory=list)
38 + """Priority list for vllm.ir.ops.gelu_and_mul_sparse"""
39 +
40 def compute_hash(self) -> str:
41 """
42 Produces a hash unique to the pass configuration.
vllm/ir/ops/__init__.py
+2 -1
@@ -1,5 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from .activation import gelu_and_mul_sparse
4 from .layernorm import fused_add_rms_norm, rms_norm
5
5 -__all__ = ["rms_norm", "fused_add_rms_norm"]
6 +__all__ = ["rms_norm", "fused_add_rms_norm", "gelu_and_mul_sparse"]
vllm/ir/ops/activation.py new
+34
@@ -0,0 +1,34 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +import torch.nn.functional as F
6 +from torch import Tensor
7 +
8 +from ..op import register_op
9 +
10 +
11 +@register_op
12 +def gelu_and_mul_sparse(
13 + x: Tensor, std_multiplier: float, approximate: str = "none"
14 +) -> Tensor:
15 + """Apply Gaussian sparsification, GELU, and gated multiplication."""
16 + d = x.shape[-1] // 2
17 + gate = x[..., :d]
18 + # Statistics intentionally remain local to each tensor-parallel shard.
19 + mean = torch.mean(gate, dim=-1, keepdim=True)
20 + std = torch.std(gate, dim=-1, keepdim=True, unbiased=False)
21 + sparse_gate = F.relu(gate - (mean + std * std_multiplier))
22 + return F.gelu(sparse_gate, approximate=approximate) * x[..., d:]
23 +
24 +
25 +@gelu_and_mul_sparse.register_input_generator
26 +def _gelu_and_mul_sparse_input_generator(
27 + num_tokens: int,
28 + hidden_size: int,
29 + dtype: torch.dtype,
30 + std_multiplier: float = 1.6448536269514722,
31 + approximate: str = "tanh",
32 +) -> tuple:
33 + x = torch.randn(num_tokens, hidden_size * 2, dtype=dtype)
34 + return x, std_multiplier, approximate
vllm/kernels/__init__.py
+2 -2
@@ -2,6 +2,6 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Kernel implementations for vLLM."""
4
5 -from . import aiter_ops, oink_ops, vllm_c
5 +from . import aiter_ops, oink_ops, triton, vllm_c
6
7 -__all__ = ["vllm_c", "aiter_ops", "oink_ops"]
7 +__all__ = ["vllm_c", "aiter_ops", "oink_ops", "triton"]
vllm/kernels/triton/__init__.py
+4
@@ -1,3 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Triton kernel implementations."""
4 +
5 +from . import activation
6 +
7 +__all__ = ["activation"]
vllm/kernels/triton/activation.py new
+110
@@ -0,0 +1,110 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +from torch import Tensor
6 +from torch.library import triton_op, wrap_triton
7 +
8 +from vllm import ir
9 +from vllm.platforms import current_platform
10 +from vllm.triton_utils import HAS_TRITON, tl, triton
11 +
12 +_MAX_INTERMEDIATE_SIZE = 32768
13 +_SUPPORTED_DTYPES = {torch.float16, torch.bfloat16, torch.float32}
14 +
15 +
16 +@triton.jit
17 +def _gelu_and_mul_sparse_kernel(
18 + x_ptr,
19 + output_ptr,
20 + d: tl.constexpr,
21 + std_multiplier,
22 + BLOCK_SIZE: tl.constexpr,
23 +):
24 + row = tl.program_id(0).to(tl.int64)
25 + offsets = tl.arange(0, BLOCK_SIZE)
26 + mask = offsets < d
27 + row_start = row * 2 * d
28 +
29 + gate = tl.load(x_ptr + row_start + offsets, mask=mask, other=0.0).to(tl.float32)
30 + # TODO: Investigate delaying this load to reduce spills for large reductions
31 + # without regressing smaller shapes.
32 + up = tl.load(x_ptr + row_start + d + offsets, mask=mask, other=0.0).to(tl.float32)
33 +
34 + input_dtype = x_ptr.dtype.element_ty
35 + mean = tl.sum(gate, axis=0) / d
36 + centered = tl.where(mask, gate - mean, 0.0)
37 + variance = tl.sum(centered * centered, axis=0) / d
38 + # Match the input-dtype rounding at each native PyTorch op boundary.
39 + mean = mean.to(input_dtype).to(tl.float32)
40 + std = tl.sqrt(variance).to(input_dtype).to(tl.float32)
41 + scaled_std = (std * std_multiplier).to(input_dtype).to(tl.float32)
42 + cutoff = (mean + scaled_std).to(input_dtype).to(tl.float32)
43 + sparse_gate = (gate - cutoff).to(input_dtype).to(tl.float32)
44 + sparse_gate = tl.where(sparse_gate < 0.0, 0.0, sparse_gate)
45 +
46 + inner = 0.7978845608028654 * (
47 + sparse_gate + 0.044715 * sparse_gate * sparse_gate * sparse_gate
48 + )
49 + tanh = 2.0 * tl.sigmoid(2.0 * inner) - 1.0
50 + activated = 0.5 * sparse_gate * (1.0 + tanh)
51 + activated = activated.to(input_dtype).to(tl.float32)
52 + tl.store(output_ptr + row * d + offsets, activated * up, mask=mask)
53 +
54 +
55 +def _supports_gelu_and_mul_sparse(
56 + x: Tensor, std_multiplier: float, approximate: str = "none"
57 +) -> bool:
58 + del std_multiplier
59 + if (
60 + x.device.type != "cuda"
61 + or x.dtype not in _SUPPORTED_DTYPES
62 + or not x.is_contiguous()
63 + or x.ndim == 0
64 + or approximate != "tanh"
65 + ):
66 + return False
67 + last_dim = x.shape[-1]
68 + return (
69 + last_dim > 0 and last_dim % 2 == 0 and last_dim // 2 <= _MAX_INTERMEDIATE_SIZE
70 + )
71 +
72 +
73 +@triton_op(
74 + "vllm::gelu_and_mul_sparse_triton",
75 + mutates_args=(),
76 +)
77 +def _gelu_and_mul_sparse_triton_op(x: Tensor, std_multiplier: float) -> Tensor:
78 + d = x.shape[-1] // 2
79 + output = torch.empty((*x.shape[:-1], d), dtype=x.dtype, device=x.device)
80 + if output.numel() == 0:
81 + return output
82 +
83 + n_rows = x.numel() // x.shape[-1]
84 +
85 + def grid(meta):
86 + return (n_rows,)
87 +
88 + block_size = triton.next_power_of_2(d)
89 + num_warps = min(max(block_size // 256, 1), 8)
90 + wrap_triton(_gelu_and_mul_sparse_kernel)[grid](
91 + x,
92 + output,
93 + d=d,
94 + std_multiplier=std_multiplier,
95 + BLOCK_SIZE=block_size,
96 + num_warps=num_warps,
97 + )
98 + return output
99 +
100 +
101 +@ir.ops.gelu_and_mul_sparse.register_impl(
102 + "triton",
103 + supported=HAS_TRITON and current_platform.is_cuda(),
104 + supports_args=_supports_gelu_and_mul_sparse,
105 +)
106 +def gelu_and_mul_sparse_triton(
107 + x: Tensor, std_multiplier: float, approximate: str = "none"
108 +) -> Tensor:
109 + assert approximate == "tanh"
110 + return _gelu_and_mul_sparse_triton_op(x, std_multiplier)
vllm/model_executor/layers/activation.py
+3 -15
@@ -8,6 +8,7 @@ import torch
8 import torch.nn as nn
9 import torch.nn.functional as F
10
11 +from vllm import ir
12 from vllm.logger import init_logger
13 from vllm.model_executor.custom_op import CustomOp
14 from vllm.platforms import CpuArchEnum, current_platform
@@ -332,24 +333,11 @@ class GeluAndMulSparse(CustomOp):
333 raise ValueError("activation_sparsity is 0.0. Please use GeluAndMul.")
334 target_sparsity_tensor = torch.tensor(activation_sparsity, dtype=torch.float32)
335 normal_dist = torch.distributions.normal.Normal(0, 1)
335 - self.std_multiplier = normal_dist.icdf(target_sparsity_tensor)
336 -
337 - def _gaussian_topk(self, x: torch.Tensor) -> torch.Tensor:
338 - """Get % sparse percentile of the Gaussian distribution."""
339 - # NOTE(rob): for TP>1, we could all-gather to get the means/std.
340 - # But we do not do this because in expectation they are the same
341 - # and in practice the eval scores are good without gathering.
342 - mean = torch.mean(x, dim=-1, keepdim=True)
343 - std = torch.std(x, dim=-1, keepdim=True, unbiased=False)
344 - cutoff_x = mean + std * self.std_multiplier
345 - return nn.functional.relu(x - cutoff_x)
336 + self.std_multiplier = normal_dist.icdf(target_sparsity_tensor).item()
337
338 def forward_native(self, x: torch.Tensor) -> torch.Tensor:
339 """PyTorch-native implementation equivalent to forward()."""
349 - d = x.shape[-1] // 2
350 - out = self._gaussian_topk(x[..., :d])
351 - out = F.gelu(out, approximate=self.approximate)
352 - return out * x[..., d:]
340 + return ir.ops.gelu_and_mul_sparse(x, self.std_multiplier, self.approximate)
341
342 def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
343 return self.forward_native(x)
vllm/platforms/cuda.py
+4 -1
@@ -755,7 +755,10 @@ class CudaPlatformBase(Platform):
755 rms_norm = ["oink"] + default
756
757 return IrOpPriorityConfig.with_default(
758 - default, rms_norm=rms_norm, fused_add_rms_norm=rms_norm
758 + default,
759 + rms_norm=rms_norm,
760 + fused_add_rms_norm=rms_norm,
761 + gelu_and_mul_sparse=["triton", "native"],
762 )
763
764 @classmethod
vllm/platforms/rocm.py
+4 -1
@@ -1161,7 +1161,10 @@ class RocmPlatform(Platform):
1161 rms_norm = default
1162
1163 return IrOpPriorityConfig.with_default(
1164 - default, rms_norm=rms_norm, fused_add_rms_norm=rms_norm
1164 + default,
1165 + rms_norm=rms_norm,
1166 + fused_add_rms_norm=rms_norm,
1167 + gelu_and_mul_sparse=["native"],
1168 )
1169
1170 @classmethod
vllm/platforms/xpu.py
+1 -1
@@ -491,7 +491,7 @@ class XPUPlatform(Platform):
491 using_inductor = cc.backend == "inductor" and cc.mode != CompilationMode.NONE
492 default = ["native"] if using_inductor else ["vllm_c", "native"]
493
494 - return IrOpPriorityConfig.with_default(default)
494 + return IrOpPriorityConfig.with_default(default, gelu_and_mul_sparse=["native"])
495
496 @classmethod
497 def device_count(cls) -> int: