@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 96161fe978
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 8, 08:45 UTC

[Kernel] [Helion] [4/N] Add silu_mul_fp8 Helion kernel (#33373)

Signed-off-by: Yanan Cao <gmagogsfm@gmail.com>

Yanan Cao committed Feb 12, 2026 at 18:13 UTC 96161fe9785814bf1adcce49dfd3c47863a2ecac
7 files changed +1002 -4
tests/kernels/helion/test_register.py
+9 -1
@@ -554,11 +554,19 @@ class TestKernelRegistry:
554 """Test suite for kernel registry functionality."""
555
556 def setup_method(self):
557 - """Clear the registry before each test."""
557 + """Save and clear the registry before each test."""
558 from vllm.kernels.helion.register import _REGISTERED_KERNELS
559
560 + self._saved_registry = dict(_REGISTERED_KERNELS)
561 _REGISTERED_KERNELS.clear()
562
563 + def teardown_method(self):
564 + """Restore the registry after each test."""
565 + from vllm.kernels.helion.register import _REGISTERED_KERNELS
566 +
567 + _REGISTERED_KERNELS.clear()
568 + _REGISTERED_KERNELS.update(self._saved_registry)
569 +
570 def test_get_registered_kernels_returns_copy(self):
571 """Test get_registered_kernels returns copy of registry."""
572 result1 = get_registered_kernels()
tests/kernels/helion/test_silu_mul_fp8.py new
+331
@@ -0,0 +1,331 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +import torch
6 +import torch.nn.functional as F
7 +
8 +from vllm.utils.import_utils import has_helion
9 +
10 +if not has_helion():
11 + pytest.skip(
12 + "Helion is not installed. Install with: pip install vllm[helion]",
13 + allow_module_level=True,
14 + )
15 +
16 +from vllm.kernels.helion.config_manager import ConfigManager
17 +from vllm.kernels.helion.ops.silu_mul_fp8 import (
18 + pick_silu_mul_fp8_config,
19 + silu_mul_fp8,
20 + silu_mul_fp8_baseline,
21 +)
22 +
23 +
24 +def skip_if_platform_unsupported():
25 + try:
26 + from vllm.kernels.helion.utils import get_canonical_gpu_name
27 +
28 + if not torch.cuda.is_available():
29 + pytest.skip("CUDA not available")
30 +
31 + platform = get_canonical_gpu_name()
32 +
33 + try:
34 + config_manager = ConfigManager.get_instance()
35 + except RuntimeError:
36 + config_manager = ConfigManager()
37 +
38 + configs = config_manager.get_platform_configs("silu_mul_fp8", platform)
39 + if len(configs) == 0:
40 + pytest.skip("Current GPU platform not supported for silu_mul_fp8 kernel")
41 +
42 + except (ImportError, RuntimeError, KeyError):
43 + pytest.skip("Error detecting platform support for silu_mul_fp8 kernel")
44 +
45 +
46 +@pytest.fixture(autouse=True)
47 +def reset_config_manager_singleton():
48 + ConfigManager.reset_instance()
49 + ConfigManager()
50 + yield
51 + ConfigManager.reset_instance()
52 +
53 +
54 +class TestSiluMulFp8ConfigPicker:
55 + def test_config_picker_exact_match(self):
56 + config_keys = [
57 + "intermediate_2048_batchsize_256",
58 + "intermediate_4096_batchsize_256",
59 + ]
60 +
61 + input_tensor = torch.randn(32, 4096, dtype=torch.bfloat16, device="cuda")
62 + scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
63 + args = (input_tensor, scale)
64 +
65 + selected_key = pick_silu_mul_fp8_config(args, config_keys)
66 + assert selected_key == "intermediate_2048_batchsize_256"
67 +
68 + def test_config_picker_closest_match(self):
69 + config_keys = [
70 + "intermediate_2048_batchsize_256",
71 + "intermediate_4096_batchsize_256",
72 + ]
73 + # Use 7000 (intermediate_size=3500) which is closer to 4096 than 2048
74 + input_tensor = torch.randn(32, 7000, dtype=torch.bfloat16, device="cuda")
75 + scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
76 + args = (input_tensor, scale)
77 +
78 + selected_key = pick_silu_mul_fp8_config(args, config_keys)
79 + assert selected_key == "intermediate_4096_batchsize_256"
80 +
81 + def test_config_picker_fallback_to_default(self):
82 + config_keys = ["default", "some_other_key"]
83 +
84 + input_tensor = torch.randn(32, 4096, dtype=torch.bfloat16, device="cuda")
85 + scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
86 + args = (input_tensor, scale)
87 +
88 + selected_key = pick_silu_mul_fp8_config(args, config_keys)
89 + assert selected_key == "default"
90 +
91 + def test_config_picker_no_configs(self):
92 + config_keys: list[str] = []
93 +
94 + input_tensor = torch.randn(32, 4096, dtype=torch.bfloat16, device="cuda")
95 + scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
96 + args = (input_tensor, scale)
97 +
98 + selected_key = pick_silu_mul_fp8_config(args, config_keys)
99 + assert selected_key is None
100 +
101 + @pytest.mark.parametrize("intermediate_size", [2048, 4096, 5120])
102 + def test_config_picker_different_sizes(self, intermediate_size):
103 + config_keys = [
104 + "intermediate_2048_batchsize_256",
105 + "intermediate_4096_batchsize_256",
106 + "intermediate_5120_batchsize_256",
107 + ]
108 +
109 + input_tensor = torch.randn(
110 + 32, 2 * intermediate_size, dtype=torch.bfloat16, device="cuda"
111 + )
112 + scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
113 + args = (input_tensor, scale)
114 +
115 + selected_key = pick_silu_mul_fp8_config(args, config_keys)
116 + expected_key = f"intermediate_{intermediate_size}_batchsize_256"
117 + assert selected_key == expected_key
118 +
119 +
120 +class TestSiluMulFp8Correctness:
121 + @pytest.mark.parametrize("batch_size", [1, 8, 32, 128])
122 + @pytest.mark.parametrize("intermediate_size", [2048, 3000, 3500, 4096, 5000])
123 + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
124 + def test_silu_mul_fp8_correctness(self, batch_size, intermediate_size, dtype):
125 + skip_if_platform_unsupported()
126 +
127 + input_size = 2 * intermediate_size
128 + input_tensor = torch.randn(batch_size, input_size, dtype=dtype, device="cuda")
129 + scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
130 +
131 + reference_output = silu_mul_fp8_baseline(input_tensor, scale)
132 + helion_output = silu_mul_fp8(input_tensor, scale)
133 +
134 + assert helion_output.shape == reference_output.shape
135 + assert helion_output.dtype == torch.float8_e4m3fn
136 + assert reference_output.dtype == torch.float8_e4m3fn
137 +
138 + ref_f32 = reference_output.to(torch.float32)
139 + helion_f32 = helion_output.to(torch.float32)
140 + # FP8 E4M3 has limited precision. Values near quantization boundaries
141 + # can round differently due to intermediate precision differences.
142 + torch.testing.assert_close(
143 + helion_f32,
144 + ref_f32,
145 + atol=0.05,
146 + rtol=0.05,
147 + msg=f"Mismatch at batch={batch_size}, size={intermediate_size}",
148 + )
149 +
150 + def test_silu_mul_fp8_shape_inference(self):
151 + skip_if_platform_unsupported()
152 + batch_size, input_size = 32, 8192
153 + intermediate_size = input_size // 2
154 +
155 + input_tensor = torch.randn(
156 + batch_size, input_size, dtype=torch.bfloat16, device="cuda"
157 + )
158 + scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
159 +
160 + output = silu_mul_fp8(input_tensor, scale)
161 +
162 + expected_shape = (batch_size, intermediate_size)
163 + assert output.shape == expected_shape
164 + assert output.dtype == torch.float8_e4m3fn
165 +
166 + def test_silu_mul_fp8_scale_variations(self):
167 + skip_if_platform_unsupported()
168 + batch_size, input_size = 16, 4096
169 +
170 + input_tensor = torch.randn(
171 + batch_size, input_size, dtype=torch.bfloat16, device="cuda"
172 + )
173 +
174 + scales = [0.1, 0.5, 1.0, 2.0, 10.0]
175 +
176 + for scale_val in scales:
177 + scale = torch.tensor([scale_val], dtype=torch.float32, device="cuda")
178 +
179 + reference_output = silu_mul_fp8_baseline(input_tensor, scale)
180 + helion_output = silu_mul_fp8(input_tensor, scale)
181 + ref_f32 = reference_output.to(torch.float32)
182 + helion_f32 = helion_output.to(torch.float32)
183 +
184 + torch.testing.assert_close(
185 + helion_f32,
186 + ref_f32,
187 + atol=0.05,
188 + rtol=0.05,
189 + msg=f"Mismatch for scale={scale_val}",
190 + )
191 +
192 + @pytest.mark.parametrize(
193 + "shape",
194 + [
195 + (1, 4096),
196 + (16, 4096),
197 + (128, 4096),
198 + (1024, 4096),
199 + (1, 8192),
200 + (16, 8192),
201 + (128, 8192),
202 + ],
203 + )
204 + def test_silu_mul_fp8_various_shapes(self, shape):
205 + skip_if_platform_unsupported()
206 +
207 + input_tensor = torch.randn(*shape, dtype=torch.bfloat16, device="cuda")
208 + scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
209 +
210 + reference_output = silu_mul_fp8_baseline(input_tensor, scale)
211 + helion_output = silu_mul_fp8(input_tensor, scale)
212 +
213 + assert helion_output.shape == reference_output.shape
214 +
215 + ref_f32 = reference_output.to(torch.float32)
216 + helion_f32 = helion_output.to(torch.float32)
217 +
218 + torch.testing.assert_close(
219 + helion_f32, ref_f32, atol=0.05, rtol=0.05, msg=f"Mismatch for shape={shape}"
220 + )
221 +
222 +
223 +def silu_mul_fp8_pytorch(input: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
224 + """Pure PyTorch reference using F.silu.
225 +
226 + This matches vLLM's SiluAndMul.forward_native exactly:
227 + F.silu(x[..., :d]) * x[..., d:]
228 + """
229 + d = input.shape[-1] // 2
230 + result = F.silu(input[..., :d]) * input[..., d:]
231 + return (result.to(torch.float32) / scale).to(torch.float8_e4m3fn)
232 +
233 +
234 +class TestSiluMulFp8PytorchReference:
235 + """Tests comparing Helion kernel against pure PyTorch implementation.
236 +
237 + Uses tighter tolerance since both use PyTorch's FP8 conversion
238 + (same rounding mode), unlike the vLLM C++ baseline which uses
239 + NVIDIA's hardware FP8 conversion with different rounding.
240 + """
241 +
242 + @pytest.mark.parametrize("batch_size", [1, 8, 32, 128, 256])
243 + @pytest.mark.parametrize("intermediate_size", [1024, 2048, 4096])
244 + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
245 + def test_silu_mul_fp8_vs_pytorch(self, batch_size, intermediate_size, dtype):
246 + skip_if_platform_unsupported()
247 +
248 + input_tensor = torch.randn(
249 + batch_size, 2 * intermediate_size, dtype=dtype, device="cuda"
250 + )
251 + scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
252 +
253 + pytorch_output = silu_mul_fp8_pytorch(input_tensor, scale)
254 + helion_output = silu_mul_fp8(input_tensor, scale)
255 +
256 + assert helion_output.shape == pytorch_output.shape
257 + assert helion_output.dtype == torch.float8_e4m3fn
258 +
259 + pytorch_f32 = pytorch_output.to(torch.float32)
260 + helion_f32 = helion_output.to(torch.float32)
261 +
262 + # Tolerance accounts for FP8 quantization boundary effects
263 + torch.testing.assert_close(
264 + helion_f32,
265 + pytorch_f32,
266 + atol=0.05,
267 + rtol=0.05,
268 + msg=(
269 + f"Mismatch at batch={batch_size}, size={intermediate_size}, "
270 + f"dtype={dtype}"
271 + ),
272 + )
273 +
274 + @pytest.mark.parametrize(
275 + "shape",
276 + [
277 + (1, 2, 4096), # 3D input
278 + (2, 4, 2048), # 3D input
279 + (1, 1, 1, 8192), # 4D input
280 + ],
281 + )
282 + def test_silu_mul_fp8_multidim_vs_pytorch(self, shape):
283 + skip_if_platform_unsupported()
284 +
285 + input_tensor = torch.randn(*shape, dtype=torch.bfloat16, device="cuda")
286 + scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
287 +
288 + pytorch_output = silu_mul_fp8_pytorch(input_tensor, scale)
289 + helion_output = silu_mul_fp8(input_tensor, scale)
290 +
291 + assert helion_output.shape == pytorch_output.shape
292 +
293 + pytorch_f32 = pytorch_output.to(torch.float32)
294 + helion_f32 = helion_output.to(torch.float32)
295 +
296 + torch.testing.assert_close(
297 + helion_f32,
298 + pytorch_f32,
299 + atol=0.05,
300 + rtol=0.05,
301 + msg=f"Mismatch for shape={shape}",
302 + )
303 +
304 +
305 +class TestSiluMulFp8Integration:
306 + def test_kernel_registration_integration(self):
307 + from vllm.kernels.helion.register import get_registered_kernels
308 +
309 + registered_kernels = get_registered_kernels()
310 + assert "silu_mul_fp8" in registered_kernels
311 +
312 + kernel_wrapper = registered_kernels["silu_mul_fp8"]
313 + assert kernel_wrapper.op_name == "silu_mul_fp8"
314 + assert kernel_wrapper._config_picker is not None
315 +
316 + def test_fake_impl_functionality(self):
317 + skip_if_platform_unsupported()
318 + from vllm.kernels.helion.register import get_registered_kernels
319 +
320 + input_tensor = torch.randn(32, 4096, dtype=torch.bfloat16, device="cuda")
321 + scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
322 + registered_kernels = get_registered_kernels()
323 + kernel_wrapper = registered_kernels["silu_mul_fp8"]
324 + fake_impl = kernel_wrapper._fake_impl
325 +
326 + fake_output = fake_impl(input_tensor, scale)
327 +
328 + expected_shape = (32, 2048)
329 + assert fake_output.shape == expected_shape
330 + assert fake_output.dtype == torch.float8_e4m3fn
331 + assert fake_output.device == input_tensor.device
vllm/kernels/helion/__init__.py
+1
@@ -2,6 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Helion integration for vLLM."""
4
5 +import vllm.kernels.helion.ops # noqa: F401 Auto-register all Helion ops
6 from vllm.kernels.helion.config_manager import (
7 ConfigManager,
8 ConfigSet,
vllm/kernels/helion/config_manager.py
-3
@@ -104,9 +104,6 @@ class ConfigSet:
104 result[platform] = {}
105
106 for config_key, config in config_keys_dict.items():
107 - # Convert helion.Config to dict using to_json() + json.loads()
108 - import json
109 -
107 result[platform][config_key] = json.loads(config.to_json())
108
109 return result
vllm/kernels/helion/configs/silu_mul_fp8.json new
+550
@@ -0,0 +1,550 @@
1 +{
2 + "nvidia_h200": {
3 + "intermediate_2048_batchsize_256": {
4 + "block_sizes": [
5 + 64,
6 + 128
7 + ],
8 + "loop_orders": [
9 + [
10 + 0,
11 + 1
12 + ]
13 + ],
14 + "flatten_loops": [
15 + true
16 + ],
17 + "l2_groupings": [
18 + 2
19 + ],
20 + "range_unroll_factors": [
21 + 0
22 + ],
23 + "range_num_stages": [
24 + 0
25 + ],
26 + "range_multi_buffers": [
27 + null
28 + ],
29 + "range_flattens": [
30 + null
31 + ],
32 + "load_eviction_policies": [
33 + "",
34 + "",
35 + ""
36 + ],
37 + "num_warps": 32,
38 + "num_stages": 1,
39 + "indexing": [
40 + "pointer",
41 + "tensor_descriptor",
42 + "pointer",
43 + "pointer"
44 + ],
45 + "pid_type": "flat",
46 + "range_warp_specializes": []
47 + },
48 + "intermediate_4096_batchsize_256": {
49 + "block_sizes": [
50 + 16,
51 + 64
52 + ],
53 + "loop_orders": [
54 + [
55 + 0,
56 + 1
57 + ]
58 + ],
59 + "flatten_loops": [
60 + false
61 + ],
62 + "l2_groupings": [
63 + 1
64 + ],
65 + "range_unroll_factors": [
66 + 0
67 + ],
68 + "range_num_stages": [
69 + 0
70 + ],
71 + "range_multi_buffers": [
72 + null
73 + ],
74 + "range_flattens": [
75 + null
76 + ],
77 + "load_eviction_policies": [
78 + "",
79 + "",
80 + ""
81 + ],
82 + "num_warps": 2,
83 + "num_stages": 1,
84 + "indexing": [
85 + "pointer",
86 + "pointer",
87 + "pointer",
88 + "pointer"
89 + ],
90 + "pid_type": "flat",
91 + "range_warp_specializes": []
92 + },
93 + "default": {
94 + "block_sizes": [
95 + 1,
96 + 512
97 + ],
98 + "loop_orders": [
99 + [
100 + 1,
101 + 0
102 + ]
103 + ],
104 + "flatten_loops": [
105 + false
106 + ],
107 + "l2_groupings": [
108 + 4
109 + ],
110 + "range_unroll_factors": [
111 + 0
112 + ],
113 + "range_num_stages": [
114 + 0
115 + ],
116 + "range_multi_buffers": [
117 + null
118 + ],
119 + "range_flattens": [
120 + null
121 + ],
122 + "load_eviction_policies": [
123 + "",
124 + "first",
125 + ""
126 + ],
127 + "num_warps": 8,
128 + "num_stages": 2,
129 + "indexing": [
130 + "tensor_descriptor",
131 + "tensor_descriptor",
132 + "tensor_descriptor",
133 + "pointer"
134 + ],
135 + "pid_type": "flat",
136 + "range_warp_specializes": []
137 + }
138 + },
139 + "nvidia_h100_pcie": {
140 + "intermediate_2048_batchsize_256": {
141 + "block_sizes": [
142 + 1,
143 + 512
144 + ],
145 + "loop_orders": [
146 + [
147 + 1,
148 + 0
149 + ]
150 + ],
151 + "flatten_loops": [
152 + false
153 + ],
154 + "l2_groupings": [
155 + 4
156 + ],
157 + "range_unroll_factors": [
158 + 0
159 + ],
160 + "range_num_stages": [
161 + 0
162 + ],
163 + "range_multi_buffers": [
164 + null
165 + ],
166 + "range_flattens": [
167 + null
168 + ],
169 + "load_eviction_policies": [
170 + "",
171 + "first",
172 + ""
173 + ],
174 + "num_warps": 8,
175 + "num_stages": 2,
176 + "indexing": [
177 + "tensor_descriptor",
178 + "tensor_descriptor",
179 + "tensor_descriptor",
180 + "pointer"
181 + ],
182 + "pid_type": "flat",
183 + "range_warp_specializes": []
184 + },
185 + "intermediate_4096_batchsize_256": {
186 + "block_sizes": [
187 + 256,
188 + 128
189 + ],
190 + "loop_orders": [
191 + [
192 + 0,
193 + 1
194 + ]
195 + ],
196 + "flatten_loops": [
197 + true
198 + ],
199 + "l2_groupings": [
200 + 1
201 + ],
202 + "range_unroll_factors": [
203 + 2
204 + ],
205 + "range_num_stages": [
206 + 3
207 + ],
208 + "range_multi_buffers": [
209 + false
210 + ],
211 + "range_flattens": [
212 + true
213 + ],
214 + "load_eviction_policies": [
215 + "last",
216 + "last",
217 + ""
218 + ],
219 + "num_warps": 32,
220 + "num_stages": 3,
221 + "indexing": [
222 + "tensor_descriptor",
223 + "tensor_descriptor",
224 + "tensor_descriptor",
225 + "pointer"
226 + ],
227 + "pid_type": "persistent_blocked",
228 + "range_warp_specializes": []
229 + },
230 + "default": {
231 + "block_sizes": [
232 + 1,
233 + 512
234 + ],
235 + "loop_orders": [
236 + [
237 + 1,
238 + 0
239 + ]
240 + ],
241 + "flatten_loops": [
242 + false
243 + ],
244 + "l2_groupings": [
245 + 4
246 + ],
247 + "range_unroll_factors": [
248 + 0
249 + ],
250 + "range_num_stages": [
251 + 0
252 + ],
253 + "range_multi_buffers": [
254 + null
255 + ],
256 + "range_flattens": [
257 + null
258 + ],
259 + "load_eviction_policies": [
260 + "",
261 + "first",
262 + ""
263 + ],
264 + "num_warps": 8,
265 + "num_stages": 2,
266 + "indexing": [
267 + "tensor_descriptor",
268 + "tensor_descriptor",
269 + "tensor_descriptor",
270 + "pointer"
271 + ],
272 + "pid_type": "flat",
273 + "range_warp_specializes": []
274 + }
275 + },
276 + "nvidia_h100_sxm5": {
277 + "intermediate_2048_batchsize_256": {
278 + "block_sizes": [
279 + 1,
280 + 512
281 + ],
282 + "loop_orders": [
283 + [
284 + 1,
285 + 0
286 + ]
287 + ],
288 + "flatten_loops": [
289 + false
290 + ],
291 + "l2_groupings": [
292 + 4
293 + ],
294 + "range_unroll_factors": [
295 + 0
296 + ],
297 + "range_num_stages": [
298 + 0
299 + ],
300 + "range_multi_buffers": [
301 + null
302 + ],
303 + "range_flattens": [
304 + null
305 + ],
306 + "load_eviction_policies": [
307 + "",
308 + "first",
309 + ""
310 + ],
311 + "num_warps": 8,
312 + "num_stages": 2,
313 + "indexing": [
314 + "tensor_descriptor",
315 + "tensor_descriptor",
316 + "tensor_descriptor",
317 + "pointer"
318 + ],
319 + "pid_type": "flat",
320 + "range_warp_specializes": []
321 + },
322 + "intermediate_4096_batchsize_256": {
323 + "block_sizes": [
324 + 256,
325 + 128
326 + ],
327 + "loop_orders": [
328 + [
329 + 0,
330 + 1
331 + ]
332 + ],
333 + "flatten_loops": [
334 + true
335 + ],
336 + "l2_groupings": [
337 + 1
338 + ],
339 + "range_unroll_factors": [
340 + 2
341 + ],
342 + "range_num_stages": [
343 + 3
344 + ],
345 + "range_multi_buffers": [
346 + false
347 + ],
348 + "range_flattens": [
349 + true
350 + ],
351 + "load_eviction_policies": [
352 + "last",
353 + "last",
354 + ""
355 + ],
356 + "num_warps": 32,
357 + "num_stages": 3,
358 + "indexing": [
359 + "tensor_descriptor",
360 + "tensor_descriptor",
361 + "tensor_descriptor",
362 + "pointer"
363 + ],
364 + "pid_type": "persistent_blocked",
365 + "range_warp_specializes": []
366 + },
367 + "default": {
368 + "block_sizes": [
369 + 1,
370 + 512
371 + ],
372 + "loop_orders": [
373 + [
374 + 1,
375 + 0
376 + ]
377 + ],
378 + "flatten_loops": [
379 + false
380 + ],
381 + "l2_groupings": [
382 + 4
383 + ],
384 + "range_unroll_factors": [
385 + 0
386 + ],
387 + "range_num_stages": [
388 + 0
389 + ],
390 + "range_multi_buffers": [
391 + null
392 + ],
393 + "range_flattens": [
394 + null
395 + ],
396 + "load_eviction_policies": [
397 + "",
398 + "first",
399 + ""
400 + ],
401 + "num_warps": 8,
402 + "num_stages": 2,
403 + "indexing": [
404 + "tensor_descriptor",
405 + "tensor_descriptor",
406 + "tensor_descriptor",
407 + "pointer"
408 + ],
409 + "pid_type": "flat",
410 + "range_warp_specializes": []
411 + }
412 + },
413 + "nvidia_h100": {
414 + "intermediate_2048_batchsize_256": {
415 + "block_sizes": [
416 + 1,
417 + 512
418 + ],
419 + "loop_orders": [
420 + [
421 + 1,
422 + 0
423 + ]
424 + ],
425 + "flatten_loops": [
426 + false
427 + ],
428 + "l2_groupings": [
429 + 4
430 + ],
431 + "range_unroll_factors": [
432 + 0
433 + ],
434 + "range_num_stages": [
435 + 0
436 + ],
437 + "range_multi_buffers": [
438 + null
439 + ],
440 + "range_flattens": [
441 + null
442 + ],
443 + "load_eviction_policies": [
444 + "",
445 + "first",
446 + ""
447 + ],
448 + "num_warps": 8,
449 + "num_stages": 2,
450 + "indexing": [
451 + "tensor_descriptor",
452 + "tensor_descriptor",
453 + "tensor_descriptor",
454 + "pointer"
455 + ],
456 + "pid_type": "flat",
457 + "range_warp_specializes": []
458 + },
459 + "intermediate_4096_batchsize_256": {
460 + "block_sizes": [
461 + 256,
462 + 128
463 + ],
464 + "loop_orders": [
465 + [
466 + 0,
467 + 1
468 + ]
469 + ],
470 + "flatten_loops": [
471 + true
472 + ],
473 + "l2_groupings": [
474 + 1
475 + ],
476 + "range_unroll_factors": [
477 + 2
478 + ],
479 + "range_num_stages": [
480 + 3
481 + ],
482 + "range_multi_buffers": [
483 + false
484 + ],
485 + "range_flattens": [
486 + true
487 + ],
488 + "load_eviction_policies": [
489 + "last",
490 + "last",
491 + ""
492 + ],
493 + "num_warps": 32,
494 + "num_stages": 3,
495 + "indexing": [
496 + "tensor_descriptor",
497 + "tensor_descriptor",
498 + "tensor_descriptor",
499 + "pointer"
500 + ],
501 + "pid_type": "persistent_blocked",
502 + "range_warp_specializes": []
503 + },
504 + "default": {
505 + "block_sizes": [
506 + 1,
507 + 512
508 + ],
509 + "loop_orders": [
510 + [
511 + 1,
512 + 0
513 + ]
514 + ],
515 + "flatten_loops": [
516 + false
517 + ],
518 + "l2_groupings": [
519 + 4
520 + ],
521 + "range_unroll_factors": [
522 + 0
523 + ],
524 + "range_num_stages": [
525 + 0
526 + ],
527 + "range_multi_buffers": [
528 + null
529 + ],
530 + "range_flattens": [
531 + null
532 + ],
533 + "load_eviction_policies": [
534 + "",
535 + "first",
536 + ""
537 + ],
538 + "num_warps": 8,
539 + "num_stages": 2,
540 + "indexing": [
541 + "tensor_descriptor",
542 + "tensor_descriptor",
543 + "tensor_descriptor",
544 + "pointer"
545 + ],
546 + "pid_type": "flat",
547 + "range_warp_specializes": []
548 + }
549 + }
550 +}
\ No newline at end of file
vllm/kernels/helion/ops/__init__.py new
+11
@@ -0,0 +1,11 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Auto-import all Helion op modules to trigger kernel registration."""
4 +
5 +import importlib
6 +import pkgutil
7 +
8 +# Automatically import all submodules so that @register_kernel
9 +# decorators execute and register ops with torch.ops.vllm_helion.
10 +for _module_info in pkgutil.iter_modules(__path__):
11 + importlib.import_module(f"{__name__}.{_module_info.name}")
vllm/kernels/helion/ops/silu_mul_fp8.py new
+100
@@ -0,0 +1,100 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from typing import Any
5 +
6 +import torch
7 +
8 +from vllm.logger import init_logger
9 +from vllm.utils.import_utils import has_helion
10 +
11 +if not has_helion():
12 + raise ImportError(
13 + "silu_mul_fp8 Helion kernel requires helion to be installed. "
14 + "Install it with: pip install helion"
15 + )
16 +
17 +import helion.language as hl
18 +
19 +from vllm.kernels.helion.register import register_kernel
20 +
21 +logger = init_logger(__name__)
22 +
23 +
24 +@register_kernel # type: ignore[misc]
25 +def silu_mul_fp8(input: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
26 + original_shape = input.shape
27 + two_d = hl.specialize(original_shape[-1])
28 + d = two_d // 2
29 + output_shape = original_shape[:-1] + (d,)
30 +
31 + input_2d = input.view(-1, original_shape[-1])
32 + m = input_2d.shape[0]
33 +
34 + # TODO(gmagogsfm): Support for more float8 subtypes (e4m3fnuz, e5m2) coming
35 + out = torch.empty((m, d), device=input.device, dtype=torch.float8_e4m3fn)
36 +
37 + input_part_a = input_2d[:, :d]
38 + input_part_b = input_2d[:, d:]
39 +
40 + assert scale.numel() == 1, "Scale must be a scalar Tensor"
41 +
42 + for tile_m, tile_n in hl.tile([m, d]):
43 + a_vals = input_part_a[tile_m, tile_n]
44 + silu_result = torch.nn.functional.silu(a_vals)
45 + b_vals = input_part_b[tile_m, tile_n]
46 + result = silu_result * b_vals
47 + result_f32 = result.to(torch.float32)
48 + scale_val = hl.load(scale, [0])
49 + inv_scale = 1.0 / scale_val
50 + result_scaled = result_f32 * inv_scale
51 + out[tile_m, tile_n] = result_scaled.to(out.dtype)
52 +
53 + return out.view(output_shape)
54 +
55 +
56 +@silu_mul_fp8.register_config_picker # type: ignore[misc]
57 +def pick_silu_mul_fp8_config(
58 + args: tuple[Any, ...], config_keys: list[str]
59 +) -> str | None:
60 + if not config_keys:
61 + return None
62 +
63 + input_tensor, scale = args
64 + intermediate_size = input_tensor.shape[-1] // 2
65 +
66 + # TODO(gmagosfm): Rerun autotuning to capture config for
67 + # other batch sizes.
68 + target_key = f"intermediate_{intermediate_size}_batchsize_256"
69 + if target_key in config_keys:
70 + return target_key
71 +
72 + intermediate_sizes = []
73 + for key in config_keys:
74 + if key.startswith("intermediate_") and "_batchsize_256" in key:
75 + try:
76 + size_str = key.split("_")[1]
77 + size = int(size_str)
78 + intermediate_sizes.append((abs(size - intermediate_size), key))
79 + except (ValueError, IndexError):
80 + continue
81 +
82 + if intermediate_sizes:
83 + _, best_key = min(intermediate_sizes)
84 + logger.debug(
85 + "No exact config for intermediate_size=%d, using closest match: %s",
86 + intermediate_size,
87 + best_key,
88 + )
89 + return best_key
90 + if "default" in config_keys:
91 + return "default"
92 +
93 + return None
94 +
95 +
96 +def silu_mul_fp8_baseline(input: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
97 + output_shape = input.shape[:-1] + (input.shape[-1] // 2,)
98 + out = torch.empty(output_shape, dtype=torch.float8_e4m3fn, device=input.device)
99 + torch.ops._C.silu_and_mul_quant(out, input, scale)
100 + return out