@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 8ecd213c0b
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 9, 06:15 UTC

[Kernel] [Helion] [2/N] Helion kernel wrapper (#32964)

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

Yanan Cao committed Jan 30, 2026 at 20:53 UTC 8ecd213c0b63d012db3b003120d7d78edafe785f
5 files changed +898
tests/kernels/helion/test_register.py new
+547
@@ -0,0 +1,547 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Unit tests for Helion kernel registration.
5 +
6 +Tests ConfiguredHelionKernel, HelionKernelWrapper, and PresetConfigSearch
7 +including config picker registration, custom autotuner integration, and
8 +PyTorch op registration.
9 +"""
10 +
11 +from unittest.mock import Mock, patch
12 +
13 +import pytest
14 +import torch
15 +
16 +from vllm.utils.import_utils import has_helion
17 +
18 +if not has_helion():
19 + pytest.skip(
20 + "Helion is not installed. Install with: pip install vllm[helion]",
21 + allow_module_level=True,
22 + )
23 +
24 +import helion
25 +
26 +from vllm.kernels.helion.config_manager import ConfigManager
27 +from vllm.kernels.helion.register import (
28 + ConfiguredHelionKernel,
29 + HelionKernelWrapper,
30 + validate_helion_settings,
31 +)
32 +
33 +
34 +@pytest.fixture
35 +def sample_configs():
36 + """Create real Helion config objects for testing."""
37 + return {
38 + "hiddensize_4096_batchsize_32": helion.Config(
39 + block_sizes=[128],
40 + num_warps=4,
41 + num_stages=3,
42 + ),
43 + "hiddensize_4096_batchsize_64": helion.Config(
44 + block_sizes=[256],
45 + num_warps=8,
46 + num_stages=4,
47 + ),
48 + "hiddensize_4096_batchsize_128": helion.Config(
49 + block_sizes=[512],
50 + num_warps=16,
51 + num_stages=2,
52 + ),
53 + "default": helion.Config(
54 + block_sizes=[64],
55 + num_warps=2,
56 + num_stages=2,
57 + ),
58 + }
59 +
60 +
61 +@pytest.fixture
62 +def sample_kernel():
63 + """Create a simple test kernel function."""
64 +
65 + def test_kernel(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
66 + """Simple test kernel that adds two tensors."""
67 + return x + y
68 +
69 + return test_kernel
70 +
71 +
72 +@pytest.fixture
73 +def config_manager_with_test_configs(sample_configs):
74 + """Set up ConfigManager with test configs for nvidia_h200 platform."""
75 + mock_config_manager = Mock(spec=ConfigManager)
76 + mock_config_manager.get_platform_configs = Mock(return_value=sample_configs)
77 + return mock_config_manager
78 +
79 +
80 +@pytest.fixture
81 +def configured_kernel(sample_kernel, sample_configs, config_manager_with_test_configs):
82 + """Create a ConfiguredHelionKernel for testing."""
83 +
84 + def test_config_picker(args, config_keys):
85 + """Simple config picker that returns default."""
86 + return "default"
87 +
88 + with (
89 + patch(
90 + "vllm.kernels.helion.config_manager.ConfigManager.get_instance",
91 + return_value=config_manager_with_test_configs,
92 + ),
93 + patch(
94 + "vllm.kernels.helion.utils.get_canonical_gpu_name",
95 + return_value="nvidia_h200",
96 + ),
97 + patch("vllm.kernels.helion.register.helion.kernel") as mock_kernel,
98 + ):
99 + # Mock just the helion.kernel decorator to avoid actual kernel compilation
100 + mock_decorated = Mock()
101 + mock_kernel.return_value = Mock(return_value=mock_decorated)
102 +
103 + return ConfiguredHelionKernel(
104 + op_name="test_kernel",
105 + config_picker=test_config_picker,
106 + raw_kernel_func=sample_kernel,
107 + helion_settings=None,
108 + )
109 +
110 +
111 +class TestValidateHelionSettings:
112 + """Test suite for validate_helion_settings utility function."""
113 +
114 + def test_accepts_none_settings(self):
115 + """Test that None settings are accepted without error."""
116 + validate_helion_settings(None, "test_kernel") # Should not raise
117 +
118 + def test_accepts_valid_settings(self):
119 + """Test that valid settings without conflicts are accepted."""
120 + settings = helion.Settings()
121 + settings.static_shapes = False
122 + settings.print_output_code = True
123 + validate_helion_settings(settings, "test_kernel") # Should not raise
124 +
125 + def test_rejects_autotuner_fn(self):
126 + """Test that settings with custom autotuner_fn raise ValueError."""
127 + settings = helion.Settings()
128 + settings.autotuner_fn = lambda *args: None # Set custom autotuner function
129 +
130 + with pytest.raises(ValueError, match="uses a custom autotuner"):
131 + validate_helion_settings(settings, "test_kernel")
132 +
133 + def test_warns_on_static_shapes_true(self):
134 + """Test that static_shapes=True emits a warning."""
135 + settings = helion.Settings()
136 + settings.static_shapes = True
137 +
138 + with patch("vllm.kernels.helion.register.logger") as mock_logger:
139 + validate_helion_settings(settings, "test_kernel")
140 + mock_logger.warning.assert_called_once()
141 + assert "static_shapes=True" in mock_logger.warning.call_args[0][0]
142 +
143 +
144 +def create_configured_kernel_with_configs(
145 + op_name,
146 + config_picker,
147 + kernel_func,
148 + configs,
149 + platform="nvidia_h200",
150 + helion_settings=None,
151 +):
152 + """Helper to create ConfiguredHelionKernel with real config objects."""
153 + mock_config_manager = Mock(spec=ConfigManager)
154 + mock_config_manager.get_platform_configs = Mock(return_value=configs)
155 +
156 + with (
157 + patch(
158 + "vllm.kernels.helion.config_manager.ConfigManager.get_instance",
159 + return_value=mock_config_manager,
160 + ),
161 + patch(
162 + "vllm.kernels.helion.utils.get_canonical_gpu_name",
163 + return_value=platform,
164 + ),
165 + patch("vllm.kernels.helion.register.helion.kernel") as mock_kernel,
166 + ):
167 + mock_decorated = Mock()
168 + mock_kernel.return_value = Mock(return_value=mock_decorated)
169 +
170 + return ConfiguredHelionKernel(
171 + op_name=op_name,
172 + config_picker=config_picker,
173 + raw_kernel_func=kernel_func,
174 + helion_settings=helion_settings,
175 + )
176 +
177 +
178 +class TestConfiguredHelionKernel:
179 + """Test suite for ConfiguredHelionKernel."""
180 +
181 + def test_init_raises_without_picker(self, sample_kernel, sample_configs):
182 + """Test that __init__ raises when no picker registered."""
183 + configs = {"default": sample_configs["default"]}
184 + mock_config_manager = Mock(spec=ConfigManager)
185 + mock_config_manager.get_platform_configs = Mock(return_value=configs)
186 +
187 + with (
188 + patch(
189 + "vllm.kernels.helion.config_manager.ConfigManager.get_instance",
190 + return_value=mock_config_manager,
191 + ),
192 + patch(
193 + "vllm.kernels.helion.utils.get_canonical_gpu_name",
194 + return_value="nvidia_h200",
195 + ),
196 + pytest.raises(RuntimeError, match="No config picker registered"),
197 + ):
198 + ConfiguredHelionKernel(
199 + op_name="test_kernel",
200 + config_picker=None, # No picker registered
201 + raw_kernel_func=sample_kernel,
202 + helion_settings=None,
203 + )
204 +
205 + def test_config_selector_validates_picker_result(
206 + self, sample_kernel, sample_configs
207 + ):
208 + """Test that config selector validates picker returns valid key."""
209 +
210 + def invalid_picker(args, config_keys):
211 + return "invalid_key"
212 +
213 + kernel = create_configured_kernel_with_configs(
214 + op_name="test_kernel",
215 + config_picker=invalid_picker,
216 + kernel_func=sample_kernel,
217 + configs=sample_configs,
218 + )
219 +
220 + key_computer = kernel._create_key_computer()
221 + selector = kernel._create_config_selector(key_computer)
222 +
223 + with pytest.raises(
224 + ValueError, match="Config picker returned invalid config key"
225 + ):
226 + selector((torch.randn(32, 4096),))
227 +
228 + def test_config_selector_handles_none_from_picker(
229 + self, sample_kernel, sample_configs
230 + ):
231 + """Test that config selector falls back to 'default' on None."""
232 +
233 + def none_picker(args, config_keys):
234 + return None
235 +
236 + kernel = create_configured_kernel_with_configs(
237 + op_name="test_kernel",
238 + config_picker=none_picker,
239 + kernel_func=sample_kernel,
240 + configs=sample_configs,
241 + )
242 +
243 + key_computer = kernel._create_key_computer()
244 + selector = kernel._create_config_selector(key_computer)
245 +
246 + result = selector((torch.randn(32, 4096),))
247 + assert result is kernel.configs["default"]
248 +
249 + def test_create_decorated_kernel_passes_helion_settings(
250 + self, sample_kernel, sample_configs
251 + ):
252 + """Test that _create_decorated_kernel passes helion_settings."""
253 +
254 + def default_picker(args, config_keys):
255 + return "default"
256 +
257 + settings = helion.Settings()
258 + settings.print_output_code = True
259 + # Note: helion.Settings() defaults static_shapes to True
260 +
261 + mock_config_manager = Mock(spec=ConfigManager)
262 + mock_config_manager.get_platform_configs = Mock(return_value=sample_configs)
263 +
264 + with (
265 + patch("vllm.kernels.helion.register.helion.kernel") as mock_kernel,
266 + patch(
267 + "vllm.kernels.helion.config_manager.ConfigManager.get_instance",
268 + return_value=mock_config_manager,
269 + ),
270 + patch(
271 + "vllm.kernels.helion.utils.get_canonical_gpu_name",
272 + return_value="nvidia_h200",
273 + ),
274 + ):
275 + mock_decorated = Mock()
276 + mock_kernel.return_value = Mock(return_value=mock_decorated)
277 +
278 + ConfiguredHelionKernel(
279 + op_name="test_kernel",
280 + config_picker=default_picker,
281 + raw_kernel_func=sample_kernel,
282 + helion_settings=settings,
283 + )
284 +
285 + call_kwargs = mock_kernel.call_args[1]
286 + assert "print_output_code" in call_kwargs
287 + assert call_kwargs["print_output_code"] is True
288 + # helion.Settings() defaults to static_shapes=True, so it should remain True
289 + assert call_kwargs["static_shapes"] is True
290 +
291 + def test_create_decorated_kernel_preserves_static_shapes_true(
292 + self, sample_kernel, sample_configs
293 + ):
294 + """Test that explicit static_shapes=True is preserved."""
295 +
296 + def default_picker(args, config_keys):
297 + return "default"
298 +
299 + settings = helion.Settings()
300 + settings.static_shapes = True
301 +
302 + mock_config_manager = Mock(spec=ConfigManager)
303 + mock_config_manager.get_platform_configs = Mock(return_value=sample_configs)
304 +
305 + with (
306 + patch("vllm.kernels.helion.register.helion.kernel") as mock_kernel,
307 + patch(
308 + "vllm.kernels.helion.config_manager.ConfigManager.get_instance",
309 + return_value=mock_config_manager,
310 + ),
311 + patch(
312 + "vllm.kernels.helion.utils.get_canonical_gpu_name",
313 + return_value="nvidia_h200",
314 + ),
315 + ):
316 + mock_decorated = Mock()
317 + mock_kernel.return_value = Mock(return_value=mock_decorated)
318 +
319 + ConfiguredHelionKernel(
320 + op_name="test_kernel",
321 + config_picker=default_picker,
322 + raw_kernel_func=sample_kernel,
323 + helion_settings=settings,
324 + )
325 +
326 + call_kwargs = mock_kernel.call_args[1]
327 + assert call_kwargs["static_shapes"] is True
328 +
329 + def test_key_and_config_selector_use_same_logic(
330 + self, sample_kernel, sample_configs
331 + ):
332 + """Test that key and config_selector produce identical results."""
333 +
334 + def tracking_picker(args, config_keys):
335 + x = args[0]
336 + batch_size = x.shape[0]
337 + if batch_size <= 32:
338 + return "hiddensize_4096_batchsize_32"
339 + elif batch_size <= 64:
340 + return "hiddensize_4096_batchsize_64"
341 + return "hiddensize_4096_batchsize_128"
342 +
343 + mock_config_manager = Mock(spec=ConfigManager)
344 + mock_config_manager.get_platform_configs = Mock(return_value=sample_configs)
345 +
346 + with (
347 + patch("vllm.kernels.helion.register.helion.kernel") as mock_helion_kernel,
348 + patch(
349 + "vllm.kernels.helion.config_manager.ConfigManager.get_instance",
350 + return_value=mock_config_manager,
351 + ),
352 + patch(
353 + "vllm.kernels.helion.utils.get_canonical_gpu_name",
354 + return_value="nvidia_h200",
355 + ),
356 + ):
357 + mock_decorated = Mock()
358 + mock_helion_kernel.return_value = Mock(return_value=mock_decorated)
359 +
360 + kernel = ConfiguredHelionKernel(
361 + op_name="test_kernel",
362 + config_picker=tracking_picker,
363 + raw_kernel_func=sample_kernel,
364 + helion_settings=None,
365 + )
366 +
367 + call_kwargs = mock_helion_kernel.call_args[1]
368 + key_fn = call_kwargs["key"]
369 + autotuner_fn = call_kwargs["autotuner_fn"]
370 +
371 + tensor = torch.randn(50, 4096) # batch=50, should select batchsize_64
372 +
373 + # key receives unpacked args, autotuner receives args as tuple
374 + key_result = key_fn(tensor)
375 + autotuner = autotuner_fn(None, (tensor,))
376 + config = autotuner.autotune()
377 +
378 + assert key_result == "hiddensize_4096_batchsize_64"
379 + assert config is kernel.configs["hiddensize_4096_batchsize_64"]
380 +
381 +
382 +class TestHelionKernelWrapper:
383 + """Test suite for HelionKernelWrapper."""
384 +
385 + def test_get_configured_op_validates_configs_available(self, sample_kernel):
386 + """Test get_configured_op validates configs are available."""
387 +
388 + def fake_impl(*args, **kwargs):
389 + return torch.zeros_like(args[0])
390 +
391 + wrapper = HelionKernelWrapper(
392 + raw_kernel_func=sample_kernel,
393 + op_name="test_kernel",
394 + fake_impl=fake_impl,
395 + )
396 +
397 + def default_picker(args, config_keys):
398 + return "default"
399 +
400 + wrapper._config_picker = default_picker
401 +
402 + mock_config_manager = Mock(spec=ConfigManager)
403 + mock_config_manager.get_platform_configs = Mock(
404 + return_value={}
405 + ) # Empty configs
406 +
407 + with (
408 + patch(
409 + "vllm.kernels.helion.config_manager.ConfigManager.get_instance",
410 + return_value=mock_config_manager,
411 + ),
412 + patch(
413 + "vllm.kernels.helion.utils.get_canonical_gpu_name",
414 + return_value="nvidia_h200",
415 + ),
416 + pytest.raises(ValueError, match="No configs available"),
417 + ):
418 + wrapper.get_configured_op()
419 +
420 + def test_get_configured_op_validates_config_picker(
421 + self, sample_kernel, sample_configs
422 + ):
423 + """Test get_configured_op validates config picker."""
424 +
425 + def fake_impl(*args, **kwargs):
426 + return torch.zeros_like(args[0])
427 +
428 + wrapper = HelionKernelWrapper(
429 + raw_kernel_func=sample_kernel,
430 + op_name="test_kernel",
431 + fake_impl=fake_impl,
432 + )
433 + # Don't set config picker - should raise assertion error
434 +
435 + mock_config_manager = Mock(spec=ConfigManager)
436 + mock_config_manager.get_platform_configs = Mock(return_value=sample_configs)
437 +
438 + with (
439 + patch(
440 + "vllm.kernels.helion.config_manager.ConfigManager.get_instance",
441 + return_value=mock_config_manager,
442 + ),
443 + patch(
444 + "vllm.kernels.helion.utils.get_canonical_gpu_name",
445 + return_value="nvidia_h200",
446 + ),
447 + pytest.raises(AssertionError, match="No config picker registered"),
448 + ):
449 + wrapper.get_configured_op()
450 +
451 + def test_get_configured_op_returns_cached_op(self, sample_kernel, sample_configs):
452 + """Test get_configured_op returns cached op when already registered."""
453 +
454 + def fake_impl(*args, **kwargs):
455 + return torch.zeros_like(args[0])
456 +
457 + def default_picker(args, config_keys):
458 + return "default"
459 +
460 + wrapper = HelionKernelWrapper(
461 + raw_kernel_func=sample_kernel,
462 + op_name="test_kernel",
463 + fake_impl=fake_impl,
464 + )
465 + wrapper._config_picker = default_picker
466 +
467 + mock_config_manager = Mock(spec=ConfigManager)
468 + mock_config_manager.get_platform_configs = Mock(return_value=sample_configs)
469 +
470 + existing_op = Mock()
471 + mock_namespace = Mock()
472 + mock_namespace.test_kernel = existing_op
473 +
474 + with (
475 + patch(
476 + "vllm.kernels.helion.config_manager.ConfigManager.get_instance",
477 + return_value=mock_config_manager,
478 + ),
479 + patch(
480 + "vllm.kernels.helion.utils.get_canonical_gpu_name",
481 + return_value="nvidia_h200",
482 + ),
483 + patch.object(torch.ops, "vllm_helion", mock_namespace),
484 + patch("vllm.kernels.helion.register.helion.kernel") as mock_kernel,
485 + ):
486 + mock_decorated = Mock()
487 + mock_kernel.return_value = Mock(return_value=mock_decorated)
488 + result = wrapper.get_configured_op()
489 + assert result is existing_op
490 +
491 + def test_get_configured_op_registers_new_op(self, sample_kernel, sample_configs):
492 + """Test get_configured_op creates and registers new op."""
493 +
494 + def fake_impl(*args, **kwargs):
495 + return torch.zeros_like(args[0])
496 +
497 + def default_picker(args, config_keys):
498 + return "default"
499 +
500 + wrapper = HelionKernelWrapper(
501 + raw_kernel_func=sample_kernel,
502 + op_name="test_kernel",
503 + fake_impl=fake_impl,
504 + )
505 + wrapper._config_picker = default_picker
506 +
507 + mock_config_manager = Mock(spec=ConfigManager)
508 + mock_config_manager.get_platform_configs = Mock(return_value=sample_configs)
509 +
510 + new_op = Mock()
511 + registered_ops: dict[str, Mock] = {}
512 +
513 + class MockNamespace:
514 + def __getattr__(self, name):
515 + if name in registered_ops:
516 + return registered_ops[name]
517 + raise AttributeError(name)
518 +
519 + mock_namespace = MockNamespace()
520 +
521 + def register_side_effect(op_name, op_func, **kwargs):
522 + registered_ops[op_name] = new_op
523 +
524 + with (
525 + patch(
526 + "vllm.kernels.helion.config_manager.ConfigManager.get_instance",
527 + return_value=mock_config_manager,
528 + ),
529 + patch(
530 + "vllm.kernels.helion.utils.get_canonical_gpu_name",
531 + return_value="nvidia_h200",
532 + ),
533 + patch.object(torch.ops, "vllm_helion", mock_namespace),
534 + patch(
535 + "vllm.kernels.helion.register.direct_register_custom_op",
536 + side_effect=register_side_effect,
537 + ) as mock_register,
538 + patch("vllm.kernels.helion.register.helion.kernel") as mock_kernel,
539 + ):
540 + mock_decorated = Mock()
541 + mock_kernel.return_value = Mock(return_value=mock_decorated)
542 + result = wrapper.get_configured_op()
543 +
544 + mock_register.assert_called_once()
545 + assert result is new_op
546 + # Check that op_func is the decorated kernel, not ConfiguredHelionKernel
547 + assert mock_register.call_args[1]["op_func"] is mock_decorated
tests/kernels/helion/test_utils.py new
+30
@@ -0,0 +1,30 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for Helion utility functions."""
4 +
5 +import pytest
6 +
7 +from vllm.kernels.helion.utils import canonicalize_gpu_name
8 +
9 +
10 +@pytest.mark.parametrize(
11 + "driver_reported_name,expected",
12 + [
13 + ("NVIDIA H200", "nvidia_h200"),
14 + ("NVIDIA A100-SXM4-80GB", "nvidia_a100_sxm4_80gb"),
15 + ("NVIDIA H100 80GB HBM3", "nvidia_h100_80gb_hbm3"),
16 + ("NVIDIA GeForce RTX 4090", "nvidia_geforce_rtx_4090"),
17 + ("AMD Instinct MI300X", "amd_instinct_mi300x"),
18 + ("Tesla V100-SXM2-32GB", "tesla_v100_sxm2_32gb"),
19 + ],
20 +)
21 +def test_canonicalize_gpu_name(driver_reported_name, expected):
22 + """Test GPU name canonicalization."""
23 + assert canonicalize_gpu_name(driver_reported_name) == expected
24 +
25 +
26 +@pytest.mark.parametrize("invalid_name", ["", " ", "\t", "\n"])
27 +def test_canonicalize_gpu_name_rejects_empty(invalid_name):
28 + """Test that empty or whitespace-only names are rejected."""
29 + with pytest.raises(ValueError, match="cannot be empty"):
30 + canonicalize_gpu_name(invalid_name)
vllm/kernels/helion/__init__.py
+14
@@ -6,8 +6,22 @@ from vllm.kernels.helion.config_manager import (
6 ConfigManager,
7 ConfigSet,
8 )
9 +from vllm.kernels.helion.register import (
10 + ConfiguredHelionKernel,
11 + HelionKernelWrapper,
12 + vllm_helion_lib,
13 +)
14 +from vllm.kernels.helion.utils import canonicalize_gpu_name, get_canonical_gpu_name
15
16 __all__ = [
17 + # Config management
18 "ConfigManager",
19 "ConfigSet",
20 + # Kernel registration
21 + "ConfiguredHelionKernel",
22 + "HelionKernelWrapper",
23 + "vllm_helion_lib",
24 + # Utilities
25 + "canonicalize_gpu_name",
26 + "get_canonical_gpu_name",
27 ]
vllm/kernels/helion/register.py new
+274
@@ -0,0 +1,274 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +vLLM Helion kernel registration with pre-tuned config selection.
5 +
6 +This module leverages Helion's internal config selection infrastructure to use
7 +pre-tuned configs instead of runtime autotuning.
8 +
9 +How Helion Normally Works
10 +-------------------------
11 +For each kernel invocation, Helion:
12 +1. Computes a cache key from input arguments
13 +2. Looks up the key in its internal compilation cache
14 +3. On cache miss, runs autotuning to find the best config
15 +4. Compiles and caches the kernel with that config
16 +
17 +How We Override It
18 +------------------
19 +We override two Helion hooks to use pre-tuned configs:
20 +
21 +1. **key**: We provide a key function (derived from config_picker) that
22 + computes cache keys matching our pre-tuned config keys. This ensures Helion's
23 + internal cache uses keys that correspond to configs we've prepared.
24 +
25 +2. **autotuner_fn**: We provide PresetConfigSearch which, instead of autotuning,
26 + simply returns the pre-tuned config for the computed key. On cache miss,
27 + Helion calls our autotuner which returns the author-prepared config.
28 +
29 +Both hooks use the same config_picker logic to ensure the cache key computed
30 +by key matches the config returned by the autotuner.
31 +
32 +Key Classes
33 +-----------
34 +- HelionKernelWrapper: Wraps raw kernel + config_picker, creates configured ops
35 +- ConfiguredHelionKernel: Platform-specific kernel registered as PyTorch custom op
36 +- PresetConfigSearch: Custom autotuner that returns pre-tuned configs
37 +"""
38 +
39 +from collections.abc import Callable
40 +from typing import Any
41 +
42 +import torch
43 +from torch.library import Library
44 +
45 +from vllm.logger import init_logger
46 +from vllm.utils.import_utils import has_helion
47 +from vllm.utils.torch_utils import direct_register_custom_op
48 +
49 +if not has_helion():
50 + raise ImportError(
51 + "register module requires helion to be installed. "
52 + "Install it with: pip install helion"
53 + )
54 +
55 +import helion
56 +from helion.autotuner.base_search import BaseAutotuner
57 +from helion.runtime.config import Config
58 +from helion.runtime.settings import default_autotuner_fn
59 +
60 +logger = init_logger(__name__)
61 +
62 +vllm_helion_lib = Library("vllm_helion", "FRAGMENT") # noqa
63 +
64 +
65 +def validate_helion_settings(
66 + helion_settings: "helion.Settings | None", op_name: str
67 +) -> None:
68 + """Validate that helion_settings doesn't contain conflicting options."""
69 + if helion_settings is None:
70 + return
71 +
72 + settings_dict = helion_settings.to_dict()
73 +
74 + if (
75 + "autotuner_fn" in settings_dict
76 + and settings_dict["autotuner_fn"] is not None
77 + and settings_dict["autotuner_fn"] is not default_autotuner_fn
78 + ):
79 + raise ValueError(
80 + f"HelionKernelWrapper for '{op_name}' uses a custom autotuner via "
81 + f"config picker. Remove 'autotuner_fn' from helion_settings and use "
82 + f"@{op_name}.register_config_picker instead."
83 + )
84 +
85 + # Warn if static_shapes is explicitly set to True since most vLLM ops need
86 + # dynamic shapes for variable batch sizes and sequence lengths
87 + if settings_dict.get("static_shapes") is True:
88 + logger.warning(
89 + "Kernel '%s' has static_shapes=True in helion_settings. "
90 + "Most vLLM ops require dynamic shapes for variable batch sizes "
91 + "and sequence lengths. Consider removing this setting.",
92 + op_name,
93 + )
94 +
95 +
96 +class PresetConfigSearch(BaseAutotuner):
97 + """Custom autotuner that uses a preset config selector instead of autotuning."""
98 +
99 + def __init__(
100 + self,
101 + args: tuple[Any, ...],
102 + config_selector: Callable[[tuple[Any, ...]], Config],
103 + ):
104 + self.args = args
105 + self.config_selector = config_selector
106 +
107 + def autotune(self, *, skip_cache: bool = False) -> Config:
108 + return self.config_selector(self.args)
109 +
110 +
111 +class ConfiguredHelionKernel:
112 + """A configured Helion kernel bound to a specific platform."""
113 +
114 + def __init__(
115 + self,
116 + op_name: str,
117 + config_picker: Callable[[tuple[Any, ...], list[str]], str | None],
118 + raw_kernel_func: Callable,
119 + helion_settings: "helion.Settings | None" = None,
120 + ):
121 + self.op_name = op_name
122 + self.config_picker = config_picker
123 + self.raw_kernel_func = raw_kernel_func
124 + self.helion_settings = helion_settings
125 + self._decorated_kernel = self._create_decorated_kernel()
126 +
127 + def __call__(self, *args, **kwargs):
128 + return self._decorated_kernel(*args, **kwargs)
129 +
130 + def _create_key_computer(self):
131 + """
132 + Create a key computer function derived from the config picker.
133 +
134 + The returned function receives kernel arguments unpacked (*args) to match
135 + Helion's key signature (called as self._key_fn(*args)).
136 + """
137 + if self.config_picker is None:
138 + raise RuntimeError(
139 + f"No config picker registered for kernel '{self.op_name}'. "
140 + f"Use @{self.op_name}.register_config_picker to register one."
141 + )
142 +
143 + def key_computer(*args):
144 + config_keys = list(self.configs.keys())
145 + selected_key = self.config_picker(args, config_keys)
146 + if selected_key:
147 + return selected_key
148 + return "default" if "default" in self.configs else None
149 +
150 + return key_computer
151 +
152 + def _create_config_selector(self, key_computer):
153 + def config_selector(args):
154 + # args is a tuple; key_computer expects unpacked args
155 + selected_config_key = key_computer(*args)
156 +
157 + if selected_config_key is None:
158 + raise ValueError(
159 + f"Config picker returned None for kernel '{self.op_name}' "
160 + f"with available config keys: {list(self.configs.keys())}"
161 + )
162 +
163 + if selected_config_key not in self.configs:
164 + raise ValueError(
165 + f"Config picker returned invalid config key "
166 + f"'{selected_config_key}' for kernel '{self.op_name}'. "
167 + f"Available keys: {list(self.configs.keys())}"
168 + )
169 +
170 + return self.configs[selected_config_key]
171 +
172 + return config_selector
173 +
174 + def _load_platform_configs(self) -> None:
175 + from vllm.kernels.helion.config_manager import ConfigManager
176 + from vllm.kernels.helion.utils import get_canonical_gpu_name
177 +
178 + self.platform = get_canonical_gpu_name()
179 + config_manager = ConfigManager.get_instance()
180 + self.configs = config_manager.get_platform_configs(self.op_name, self.platform)
181 +
182 + if not self.configs:
183 + raise ValueError(
184 + f"No configs available for kernel '{self.op_name}' "
185 + f"on platform '{self.platform}'"
186 + )
187 +
188 + def _create_decorated_kernel(self) -> Callable[..., Any]:
189 + self._load_platform_configs()
190 +
191 + key_computer = self._create_key_computer()
192 + config_selector = self._create_config_selector(key_computer)
193 +
194 + kernel_kwargs = {}
195 + if self.helion_settings:
196 + kernel_kwargs.update(self.helion_settings.to_dict())
197 +
198 + # Set static_shapes=False by default if user didn't explicitly set it to True
199 + # This is needed for dynamic batch sizes and sequence lengths in vLLM
200 + if kernel_kwargs.get("static_shapes") is not True:
201 + kernel_kwargs["static_shapes"] = False
202 +
203 + kernel_kwargs["autotuner_fn"] = lambda _, args: PresetConfigSearch(
204 + args, config_selector
205 + )
206 + kernel_kwargs["key"] = key_computer
207 +
208 + logger.debug(
209 + "Creating decorated kernel %s with custom autotuner on platform %s",
210 + self.op_name,
211 + self.platform,
212 + )
213 + return helion.kernel(**kernel_kwargs)(self.raw_kernel_func)
214 +
215 +
216 +class HelionKernelWrapper:
217 + """Wrapper for Helion kernels that creates config-specific PyTorch custom ops."""
218 +
219 + def __init__(
220 + self,
221 + raw_kernel_func: Callable,
222 + op_name: str,
223 + fake_impl: Callable,
224 + helion_settings: "helion.Settings | None" = None,
225 + ):
226 + # Validate helion_settings doesn't conflict with our custom autotuner
227 + validate_helion_settings(helion_settings, op_name)
228 +
229 + self.raw_kernel_func = raw_kernel_func
230 + self.op_name = op_name
231 + self._fake_impl = fake_impl
232 + self.helion_settings = helion_settings
233 + self._config_picker: (
234 + Callable[[tuple[Any, ...], list[str]], str | None] | None
235 + ) = None
236 +
237 + def __call__(self, *args, **kwargs):
238 + configured_op = self.get_configured_op()
239 + return configured_op(*args, **kwargs)
240 +
241 + def register_config_picker(
242 + self, picker_func: Callable[[tuple[Any, ...], list[str]], str | None]
243 + ) -> Callable[[tuple[Any, ...], list[str]], str | None]:
244 + self._config_picker = picker_func
245 + return picker_func
246 +
247 + def get_configured_op(self) -> Any:
248 + assert self._config_picker is not None, (
249 + f"No config picker registered for kernel '{self.op_name}'. "
250 + f"Use @{self.op_name}.register_config_picker to register one."
251 + )
252 +
253 + if hasattr(torch.ops.vllm_helion, self.op_name):
254 + logger.debug("Op vllm_helion::%s already registered", self.op_name)
255 + return getattr(torch.ops.vllm_helion, self.op_name)
256 +
257 + configured_kernel = ConfiguredHelionKernel(
258 + op_name=self.op_name,
259 + config_picker=self._config_picker,
260 + raw_kernel_func=self.raw_kernel_func,
261 + helion_settings=self.helion_settings,
262 + )
263 +
264 + logger.info("Registering op: vllm_helion::%s", self.op_name)
265 + direct_register_custom_op(
266 + op_name=self.op_name,
267 + op_func=configured_kernel._decorated_kernel, # Register decorated kernel
268 + # TODO(gmagogsfm): Implement automatic mutation/aliasing detection
269 + # for Helion kernels.
270 + mutates_args=None,
271 + fake_impl=self._fake_impl,
272 + target_lib=vllm_helion_lib,
273 + )
274 + return getattr(torch.ops.vllm_helion, self.op_name)
vllm/kernels/helion/utils.py new
+33
@@ -0,0 +1,33 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Utility functions for Helion kernel management."""
4 +
5 +import torch
6 +
7 +
8 +def get_gpu_name(device_id: int | None = None) -> str:
9 + if device_id is None:
10 + device_id = torch.cuda.current_device()
11 + props = torch.cuda.get_device_properties(device_id)
12 + return props.name
13 +
14 +
15 +def canonicalize_gpu_name(name: str) -> str:
16 + """
17 + Canonicalize GPU name for use as a platform identifier.
18 +
19 + Converts to lowercase and replaces spaces and hyphens with underscores.
20 + e.g., "NVIDIA A100-SXM4-80GB" -> "nvidia_a100_sxm4_80gb"
21 +
22 + Raises ValueError if name is empty.
23 + """
24 + if not name or not name.strip():
25 + raise ValueError("GPU name cannot be empty")
26 + name = name.lower()
27 + name = name.replace(" ", "_")
28 + name = name.replace("-", "_")
29 + return name
30 +
31 +
32 +def get_canonical_gpu_name(device_id: int | None = None) -> str:
33 + return canonicalize_gpu_name(get_gpu_name(device_id))