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 6, 13:15 UTC
[Kernel][Helion] [16/N] Refactor register_kernel API to be more Dynamo-friendly (#36705)
Signed-off-by: Yanan Cao <gmagogsfm@gmail.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Yanan Cao committed
Mar 17, 2026 at 18:23 UTC
ff9fbc9aff345108c2f7baa335b8d26be8f25b09
6 files changed
+627
-272
tests/kernels/helion/helpers.py
new
+67
@@ -0,0 +1,67 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import json
5
+import tempfile
6
+from collections.abc import Callable
7
+from contextlib import contextmanager
8
+from pathlib import Path
9
+from unittest.mock import patch
10
+
11
+import helion
12
+
13
+from vllm.kernels.helion.config_manager import ConfigManager
14
+from vllm.kernels.helion.register import register_kernel
15
+from vllm.kernels.helion.utils import get_canonical_gpu_name
16
+
17
+GPU_PLATFORM = get_canonical_gpu_name()
18
+
19
+DEFAULT_CONFIGS: dict[str, helion.Config] = {
20
+ "default": helion.Config(block_sizes=[32]),
21
+}
22
+
23
+
24
+@contextmanager
25
+def dummy_kernel_registry(
26
+ configs: dict[str, helion.Config] | None = None,
27
+):
28
+ """Context manager providing a register function with automatic config setup.
29
+
30
+ Yields a ``register`` callable with the same signature as
31
+ ``register_kernel``. Before applying the real decorator it writes a
32
+ config JSON for the kernel name (from ``op_name`` or ``fn.__name__``)
33
+ into a temporary directory backed by a fresh ``ConfigManager``.
34
+ """
35
+ if configs is None:
36
+ configs = DEFAULT_CONFIGS
37
+ config_data = {k: v.__dict__["config"] for k, v in configs.items()}
38
+
39
+ with tempfile.TemporaryDirectory() as tmpdir:
40
+ config_dir = Path(tmpdir)
41
+ ConfigManager.reset_instance()
42
+ cm = ConfigManager(base_dir=config_dir)
43
+
44
+ with patch(
45
+ "vllm.kernels.helion.config_manager.ConfigManager",
46
+ return_value=cm,
47
+ ):
48
+
49
+ def register(
50
+ op_name: str | None = None,
51
+ **kwargs,
52
+ ) -> Callable:
53
+ def decorator(fn: Callable) -> Callable:
54
+ name = op_name or fn.__name__
55
+ kernel_dir = config_dir / name
56
+ kernel_dir.mkdir(parents=True, exist_ok=True)
57
+ (kernel_dir / f"{GPU_PLATFORM}.json").write_text(
58
+ json.dumps(config_data)
59
+ )
60
+ return register_kernel(op_name, **kwargs)(fn)
61
+
62
+ return decorator
63
+
64
+ try:
65
+ yield register
66
+ finally:
67
+ ConfigManager.reset_instance()
tests/kernels/helion/test_autotune.py
new
+91
@@ -0,0 +1,91 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for autotuning Helion kernels, including disabled kernels with no configs."""
4
+
5
+import pytest
6
+import torch
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
+import helion
17
+import helion.language as hl
18
+from helion.autotuner.base_search import BaseSearch
19
+
20
+from tests.kernels.helion.helpers import dummy_kernel_registry
21
+from vllm.kernels.helion.register import create_helion_decorated_kernel
22
+
23
+
24
+def _add_kernel(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
25
+ out = torch.empty_like(x)
26
+ for tile in hl.tile(x.size()):
27
+ out[tile] = x[tile] + y[tile]
28
+ return out
29
+
30
+
31
+class NoCompileSearch(BaseSearch):
32
+ """Autotuner that returns the default config without GPU compilation.
33
+
34
+ Modeled after helion's test BasicSearch (pytorch/helion#1649).
35
+ """
36
+
37
+ def autotune(self, *, skip_cache: bool = False):
38
+ return self.config_spec.default_config()
39
+
40
+
41
+def _no_compile_autotuner_fn(bound_kernel, args, **kwargs):
42
+ return NoCompileSearch(bound_kernel, args, **kwargs)
43
+
44
+
45
+class TestAutotuneDisabledKernel:
46
+ """Test autotuning flow on disabled kernels (no platform configs)."""
47
+
48
+ def setup_method(self):
49
+ from vllm.kernels.helion.register import _REGISTERED_KERNELS
50
+
51
+ self._saved_registry = dict(_REGISTERED_KERNELS)
52
+ _REGISTERED_KERNELS.clear()
53
+
54
+ def teardown_method(self):
55
+ from vllm.kernels.helion.register import _REGISTERED_KERNELS
56
+
57
+ _REGISTERED_KERNELS.clear()
58
+ _REGISTERED_KERNELS.update(self._saved_registry)
59
+
60
+ def test_autotune_disabled_kernel_produces_valid_config(self):
61
+ """Register a kernel with no configs (disabled), run autotune,
62
+ verify it produces a valid helion.Config."""
63
+ with dummy_kernel_registry(configs={}) as register:
64
+ wrapper = register(
65
+ "autotune_test_kernel",
66
+ config_picker=lambda args, keys: "default",
67
+ fake_impl=lambda *a, **kw: None,
68
+ input_generator=lambda: {
69
+ "small": (
70
+ torch.randn(4, 4, device="cuda"),
71
+ torch.randn(4, 4, device="cuda"),
72
+ ),
73
+ },
74
+ )(_add_kernel)
75
+
76
+ assert wrapper._disabled is True
77
+
78
+ inputs = wrapper.get_inputs()
79
+ assert "small" in inputs
80
+
81
+ settings = helion.Settings()
82
+ settings.autotuner_fn = _no_compile_autotuner_fn
83
+ wrapper.helion_settings = settings
84
+
85
+ config = wrapper.run_autotune(inputs["small"])
86
+ expected_default = (
87
+ create_helion_decorated_kernel(_add_kernel, helion_settings=settings)
88
+ .bind(inputs["small"])
89
+ .config_spec.default_config()
90
+ )
91
+ assert config == expected_default
tests/kernels/helion/test_pattern_matching.py
+3
-3
@@ -52,7 +52,7 @@ def _helion_mock_context():
52
53
with (
54
patch(
55
- "vllm.kernels.helion.config_manager.ConfigManager.get_instance",
55
+ "vllm.kernels.helion.config_manager.ConfigManager",
56
return_value=mock_config_manager,
57
),
58
patch(
@@ -87,8 +87,8 @@ class TestMakeFxHop:
87
raw_kernel_func=raw_add_scale,
88
op_name="test_make_fx",
89
fake_impl=lambda *a, **kw: None,
90
+ config_picker=lambda args, keys: "default",
91
)
91
- wrapper.register_config_picker(lambda args, keys: "default")
92
93
def fn(x, y):
94
return wrapper(x, y, scale)
@@ -143,8 +143,8 @@ class TestMakeFxHop:
143
raw_kernel_func=raw_silu_mul,
144
op_name="test_pm_silu_mul",
145
fake_impl=lambda *a, **kw: None,
146
+ config_picker=lambda args, keys: "default",
147
)
147
- wrapper.register_config_picker(lambda args, keys: "default")
148
149
def pattern(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
150
return torch.nn.functional.silu(x) * y
tests/kernels/helion/test_register.py
+348
-142
@@ -21,7 +21,9 @@ if not has_helion():
21
)
22
23
import helion
24
+import helion.language as hl
25
26
+from tests.kernels.helion.helpers import dummy_kernel_registry
27
from vllm.kernels.helion.config_manager import ConfigManager
28
from vllm.kernels.helion.register import (
29
_HOP_AVAILABLE,
@@ -34,6 +36,13 @@ from vllm.kernels.helion.register import (
36
)
37
38
39
+def _add_kernel(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
40
+ out = torch.empty_like(x)
41
+ for tile in hl.tile(x.size()):
42
+ out[tile] = x[tile] + y[tile]
43
+ return out
44
+
45
+
46
@pytest.fixture
47
def sample_configs():
48
"""Create real Helion config objects for testing."""
@@ -90,7 +99,7 @@ def configured_kernel(sample_kernel, sample_configs, config_manager_with_test_co
99
100
with (
101
patch(
93
- "vllm.kernels.helion.config_manager.ConfigManager.get_instance",
102
+ "vllm.kernels.helion.config_manager.ConfigManager",
103
return_value=config_manager_with_test_configs,
104
),
105
patch(
@@ -158,7 +167,7 @@ def create_configured_kernel_with_configs(
167
168
with (
169
patch(
161
- "vllm.kernels.helion.config_manager.ConfigManager.get_instance",
170
+ "vllm.kernels.helion.config_manager.ConfigManager",
171
return_value=mock_config_manager,
172
),
173
patch(
@@ -189,7 +198,7 @@ class TestConfiguredHelionKernel:
198
199
with (
200
patch(
192
- "vllm.kernels.helion.config_manager.ConfigManager.get_instance",
201
+ "vllm.kernels.helion.config_manager.ConfigManager",
202
return_value=mock_config_manager,
203
),
204
patch(
@@ -266,7 +275,7 @@ class TestConfiguredHelionKernel:
275
with (
276
patch("vllm.kernels.helion.register.helion.kernel") as mock_kernel,
277
patch(
269
- "vllm.kernels.helion.config_manager.ConfigManager.get_instance",
278
+ "vllm.kernels.helion.config_manager.ConfigManager",
279
return_value=mock_config_manager,
280
),
281
patch(
@@ -310,7 +319,7 @@ class TestConfiguredHelionKernel:
319
with (
320
patch("vllm.kernels.helion.register.helion.kernel") as mock_helion_kernel,
321
patch(
313
- "vllm.kernels.helion.config_manager.ConfigManager.get_instance",
322
+ "vllm.kernels.helion.config_manager.ConfigManager",
323
return_value=mock_config_manager,
324
),
325
patch(
@@ -346,23 +355,15 @@ class TestConfiguredHelionKernel:
355
class TestHelionKernelWrapper:
356
"""Test suite for HelionKernelWrapper."""
357
349
- def test_get_configured_op_validates_configs_available(self, sample_kernel):
350
- """Test get_configured_op validates configs are available."""
358
+ def test_init_disables_on_missing_configs(self, sample_kernel):
359
+ """Test __init__ marks wrapper as disabled when configs are missing."""
360
361
def fake_impl(*args, **kwargs):
362
return torch.zeros_like(args[0])
363
355
- wrapper = HelionKernelWrapper(
356
- raw_kernel_func=sample_kernel,
357
- op_name="test_kernel",
358
- fake_impl=fake_impl,
359
- )
360
-
364
def default_picker(args, config_keys):
365
return "default"
366
364
- wrapper._config_picker = default_picker
365
-
367
mock_config_manager = Mock(spec=ConfigManager)
368
mock_config_manager.get_platform_configs = Mock(
369
return_value={}
@@ -370,52 +371,99 @@ class TestHelionKernelWrapper:
371
372
with (
373
patch(
373
- "vllm.kernels.helion.config_manager.ConfigManager.get_instance",
374
+ "vllm.kernels.helion.config_manager.ConfigManager",
375
return_value=mock_config_manager,
376
),
377
patch(
378
"vllm.kernels.helion.utils.get_canonical_gpu_name",
379
return_value="nvidia_h200",
380
),
380
- pytest.raises(ValueError, match="No configs available"),
381
+ patch("vllm.kernels.helion.register.helion.kernel") as mock_kernel,
382
):
382
- wrapper.get_configured_op()
383
+ mock_kernel.return_value = Mock(return_value=sample_kernel)
384
384
- def test_get_configured_op_validates_config_picker(
385
- self, sample_kernel, sample_configs
386
- ):
387
- """Test get_configured_op validates config picker."""
385
+ wrapper = HelionKernelWrapper(
386
+ raw_kernel_func=sample_kernel,
387
+ op_name="test_kernel",
388
+ fake_impl=fake_impl,
389
+ config_picker=default_picker,
390
+ )
391
+
392
+ assert wrapper._disabled is True
393
+ assert "No configs available" in wrapper._disabled_reason
394
+
395
+ def test_disabled_wrapper_raises_on_call(self, sample_kernel):
396
+ """Test __call__ raises RuntimeError on a disabled wrapper."""
397
398
def fake_impl(*args, **kwargs):
399
return torch.zeros_like(args[0])
400
392
- wrapper = HelionKernelWrapper(
393
- raw_kernel_func=sample_kernel,
394
- op_name="test_kernel",
395
- fake_impl=fake_impl,
396
- )
397
- # Don't set config picker - should raise assertion error
401
+ def default_picker(args, config_keys):
402
+ return "default"
403
404
mock_config_manager = Mock(spec=ConfigManager)
400
- mock_config_manager.get_platform_configs = Mock(return_value=sample_configs)
405
+ mock_config_manager.get_platform_configs = Mock(return_value={})
406
+
407
+ with (
408
+ patch(
409
+ "vllm.kernels.helion.config_manager.ConfigManager",
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
+ patch("vllm.kernels.helion.register.helion.kernel") as mock_kernel,
417
+ ):
418
+ mock_kernel.return_value = Mock(return_value=sample_kernel)
419
+
420
+ wrapper = HelionKernelWrapper(
421
+ raw_kernel_func=sample_kernel,
422
+ op_name="test_kernel",
423
+ fake_impl=fake_impl,
424
+ config_picker=default_picker,
425
+ )
426
+
427
+ with pytest.raises(RuntimeError, match="is disabled"):
428
+ wrapper(torch.randn(4, 4), torch.randn(4, 4))
429
+
430
+ def test_disabled_wrapper_get_configured_op_raises(self, sample_kernel):
431
+ """Test get_configured_op raises RuntimeError on a disabled wrapper."""
432
+
433
+ def fake_impl(*args, **kwargs):
434
+ return torch.zeros_like(args[0])
435
+
436
+ def default_picker(args, config_keys):
437
+ return "default"
438
+
439
+ mock_config_manager = Mock(spec=ConfigManager)
440
+ mock_config_manager.get_platform_configs = Mock(return_value={})
441
442
with (
443
patch(
404
- "vllm.kernels.helion.config_manager.ConfigManager.get_instance",
444
+ "vllm.kernels.helion.config_manager.ConfigManager",
445
return_value=mock_config_manager,
446
),
447
patch(
448
"vllm.kernels.helion.utils.get_canonical_gpu_name",
449
return_value="nvidia_h200",
450
),
411
- pytest.raises(AssertionError, match="No config picker registered"),
451
+ patch("vllm.kernels.helion.register.helion.kernel") as mock_kernel,
452
):
453
+ mock_kernel.return_value = Mock(return_value=sample_kernel)
454
+
455
+ wrapper = HelionKernelWrapper(
456
+ raw_kernel_func=sample_kernel,
457
+ op_name="test_kernel",
458
+ fake_impl=fake_impl,
459
+ config_picker=default_picker,
460
+ )
461
+
462
+ with pytest.raises(RuntimeError, match="is disabled"):
463
wrapper.get_configured_op()
464
415
- def test_get_configured_op_returns_cached_kernel(
416
- self, sample_kernel, sample_configs
417
- ):
418
- """Test get_configured_op returns cached ConfiguredHelionKernel."""
465
+ def test_disabled_wrapper_supports_get_inputs(self, sample_kernel):
466
+ """Test get_inputs works on a disabled wrapper."""
467
468
def fake_impl(*args, **kwargs):
469
return torch.zeros_like(args[0])
@@ -423,19 +471,99 @@ class TestHelionKernelWrapper:
471
def default_picker(args, config_keys):
472
return "default"
473
426
- wrapper = HelionKernelWrapper(
427
- raw_kernel_func=sample_kernel,
428
- op_name="test_kernel",
429
- fake_impl=fake_impl,
430
- )
431
- wrapper._config_picker = default_picker
474
+ expected_inputs = {"key1": (torch.randn(4),)}
475
+ input_gen = Mock(return_value=expected_inputs)
476
+
477
+ mock_config_manager = Mock(spec=ConfigManager)
478
+ mock_config_manager.get_platform_configs = Mock(return_value={})
479
+
480
+ with (
481
+ patch(
482
+ "vllm.kernels.helion.config_manager.ConfigManager",
483
+ return_value=mock_config_manager,
484
+ ),
485
+ patch(
486
+ "vllm.kernels.helion.utils.get_canonical_gpu_name",
487
+ return_value="nvidia_h200",
488
+ ),
489
+ patch("vllm.kernels.helion.register.helion.kernel") as mock_kernel,
490
+ ):
491
+ mock_kernel.return_value = Mock(return_value=sample_kernel)
492
+
493
+ wrapper = HelionKernelWrapper(
494
+ raw_kernel_func=sample_kernel,
495
+ op_name="test_kernel",
496
+ fake_impl=fake_impl,
497
+ config_picker=default_picker,
498
+ input_generator=input_gen,
499
+ )
500
+
501
+ assert wrapper._disabled is True
502
+ result = wrapper.get_inputs()
503
+ assert result is expected_inputs
504
+
505
+ def test_disabled_wrapper_supports_run_autotune(self, sample_kernel):
506
+ """Test run_autotune works on a disabled wrapper."""
507
+
508
+ def fake_impl(*args, **kwargs):
509
+ return torch.zeros_like(args[0])
510
+
511
+ def default_picker(args, config_keys):
512
+ return "default"
513
+
514
+ mock_config_manager = Mock(spec=ConfigManager)
515
+ mock_config_manager.get_platform_configs = Mock(return_value={})
516
+
517
+ mock_config = Mock()
518
+
519
+ with (
520
+ patch(
521
+ "vllm.kernels.helion.config_manager.ConfigManager",
522
+ return_value=mock_config_manager,
523
+ ),
524
+ patch(
525
+ "vllm.kernels.helion.utils.get_canonical_gpu_name",
526
+ return_value="nvidia_h200",
527
+ ),
528
+ patch("vllm.kernels.helion.register.helion.kernel") as mock_kernel,
529
+ ):
530
+ mock_kernel.return_value = Mock(return_value=sample_kernel)
531
+
532
+ wrapper = HelionKernelWrapper(
533
+ raw_kernel_func=sample_kernel,
534
+ op_name="test_kernel",
535
+ fake_impl=fake_impl,
536
+ config_picker=default_picker,
537
+ )
538
+
539
+ assert wrapper._disabled is True
540
+
541
+ with patch(
542
+ "vllm.kernels.helion.register.create_helion_decorated_kernel"
543
+ ) as mock_create:
544
+ mock_autotune_kernel = Mock()
545
+ mock_autotune_kernel.autotune.return_value = mock_config
546
+ mock_create.return_value = mock_autotune_kernel
547
+
548
+ inputs = (torch.randn(4, 4),)
549
+ result = wrapper.run_autotune(inputs)
550
+ assert result is mock_config
551
+
552
+ def test_init_caches_configured_kernel(self, sample_kernel, sample_configs):
553
+ """Test __init__ eagerly builds and caches ConfiguredHelionKernel."""
554
+
555
+ def fake_impl(*args, **kwargs):
556
+ return torch.zeros_like(args[0])
557
+
558
+ def default_picker(args, config_keys):
559
+ return "default"
560
561
mock_config_manager = Mock(spec=ConfigManager)
562
mock_config_manager.get_platform_configs = Mock(return_value=sample_configs)
563
564
with (
565
patch(
438
- "vllm.kernels.helion.config_manager.ConfigManager.get_instance",
566
+ "vllm.kernels.helion.config_manager.ConfigManager",
567
return_value=mock_config_manager,
568
),
569
patch(
@@ -444,13 +572,77 @@ class TestHelionKernelWrapper:
572
),
573
patch("vllm.kernels.helion.register.helion.kernel") as mock_kernel,
574
):
447
- mock_decorated = Mock()
448
- mock_kernel.return_value = Mock(return_value=mock_decorated)
575
+ mock_kernel.return_value = Mock(return_value=sample_kernel)
576
577
+ wrapper = HelionKernelWrapper(
578
+ raw_kernel_func=sample_kernel,
579
+ op_name="test_kernel",
580
+ fake_impl=fake_impl,
581
+ config_picker=default_picker,
582
+ )
583
+
584
+ assert wrapper._configured_kernel is not None
585
result1 = wrapper.get_configured_op()
586
result2 = wrapper.get_configured_op()
587
assert result1 is result2
588
589
+ @pytest.mark.skipif(
590
+ not _HOP_AVAILABLE, reason="HOP path only used when HOP available"
591
+ )
592
+ def test_init_eagerly_initializes_hop_path(self):
593
+ """Test that register_kernel eagerly builds the configured kernel
594
+ on the HOP path (no custom op registration needed)."""
595
+ from vllm.kernels.helion.utils import get_canonical_gpu_name
596
+
597
+ configs = {"default": helion.Config(block_sizes=[4, 4])}
598
+ with (
599
+ dummy_kernel_registry(configs=configs) as register,
600
+ patch(
601
+ "vllm.kernels.helion.utils.get_canonical_gpu_name",
602
+ wraps=get_canonical_gpu_name,
603
+ ) as mock_gpu,
604
+ ):
605
+ wrapper = register(
606
+ config_picker=lambda args, keys: "default",
607
+ )(_add_kernel)
608
+
609
+ mock_gpu.assert_called_once()
610
+ assert wrapper._configured_kernel is not None
611
+
612
+ with patch(
613
+ "vllm.kernels.helion.utils.get_canonical_gpu_name",
614
+ side_effect=AssertionError("get_canonical_gpu_name called during __call__"),
615
+ ):
616
+ x = torch.randn(4, 4, device="cuda")
617
+ y = torch.randn(4, 4, device="cuda")
618
+ result = wrapper(x, y)
619
+ expected = x + y
620
+ assert torch.allclose(result, expected)
621
+
622
+ @pytest.mark.skipif(
623
+ _HOP_AVAILABLE, reason="CustomOp path not used when HOP available"
624
+ )
625
+ def test_init_eagerly_initializes(self):
626
+ """Test that register_kernel eagerly loads configs and detects GPU
627
+ during construction so __call__ needs no further initialization."""
628
+ from vllm.kernels.helion.utils import get_canonical_gpu_name
629
+
630
+ with (
631
+ dummy_kernel_registry() as register,
632
+ patch(
633
+ "vllm.kernels.helion.utils.get_canonical_gpu_name",
634
+ wraps=get_canonical_gpu_name,
635
+ ) as mock_gpu,
636
+ ):
637
+ wrapper = register(
638
+ config_picker=lambda args, keys: "default",
639
+ )(_add_kernel)
640
+
641
+ # Init must have detected GPU and built the kernel
642
+ mock_gpu.assert_called_once()
643
+ assert wrapper._configured_kernel is not None
644
+ assert hasattr(torch.ops.vllm_helion, wrapper.op_name)
645
+
646
@pytest.mark.skipif(
647
_HOP_AVAILABLE, reason="CustomOp path not used when HOP available"
648
)
@@ -463,13 +655,6 @@ class TestHelionKernelWrapper:
655
def default_picker(args, config_keys):
656
return "default"
657
466
- wrapper = HelionKernelWrapper(
467
- raw_kernel_func=sample_kernel,
468
- op_name="test_kernel",
469
- fake_impl=fake_impl,
470
- )
471
- wrapper._config_picker = default_picker
472
-
658
mock_config_manager = Mock(spec=ConfigManager)
659
mock_config_manager.get_platform_configs = Mock(return_value=sample_configs)
660
@@ -479,7 +664,7 @@ class TestHelionKernelWrapper:
664
665
with (
666
patch(
482
- "vllm.kernels.helion.config_manager.ConfigManager.get_instance",
667
+ "vllm.kernels.helion.config_manager.ConfigManager",
668
return_value=mock_config_manager,
669
),
670
patch(
@@ -491,6 +676,13 @@ class TestHelionKernelWrapper:
676
):
677
mock_decorated = Mock()
678
mock_kernel.return_value = Mock(return_value=mock_decorated)
679
+
680
+ wrapper = HelionKernelWrapper(
681
+ raw_kernel_func=sample_kernel,
682
+ op_name="test_kernel",
683
+ fake_impl=fake_impl,
684
+ config_picker=default_picker,
685
+ )
686
result = wrapper._get_or_register_custom_op()
687
assert result is existing_op
688
@@ -506,13 +698,6 @@ class TestHelionKernelWrapper:
698
def default_picker(args, config_keys):
699
return "default"
700
509
- wrapper = HelionKernelWrapper(
510
- raw_kernel_func=sample_kernel,
511
- op_name="test_kernel",
512
- fake_impl=fake_impl,
513
- )
514
- wrapper._config_picker = default_picker
515
-
701
mock_config_manager = Mock(spec=ConfigManager)
702
mock_config_manager.get_platform_configs = Mock(return_value=sample_configs)
703
@@ -532,7 +717,7 @@ class TestHelionKernelWrapper:
717
718
with (
719
patch(
535
- "vllm.kernels.helion.config_manager.ConfigManager.get_instance",
720
+ "vllm.kernels.helion.config_manager.ConfigManager",
721
return_value=mock_config_manager,
722
),
723
patch(
@@ -548,6 +733,13 @@ class TestHelionKernelWrapper:
733
):
734
mock_decorated = Mock()
735
mock_kernel.return_value = Mock(return_value=mock_decorated)
736
+
737
+ wrapper = HelionKernelWrapper(
738
+ raw_kernel_func=sample_kernel,
739
+ op_name="test_kernel",
740
+ fake_impl=fake_impl,
741
+ config_picker=default_picker,
742
+ )
743
result = wrapper._get_or_register_custom_op()
744
745
mock_register.assert_called_once()
@@ -584,11 +776,10 @@ class TestKernelRegistry:
776
777
def test_get_kernel_by_name_returns_kernel(self):
778
"""Test get_kernel_by_name returns registered kernel."""
587
- wrapper = HelionKernelWrapper(
588
- raw_kernel_func=Mock(),
589
- op_name="test_kernel",
590
- fake_impl=Mock(),
591
- )
779
+ with dummy_kernel_registry() as register:
780
+ wrapper = register(
781
+ "test_kernel", config_picker=lambda args, keys: "default"
782
+ )(_add_kernel)
783
784
from vllm.kernels.helion.register import _REGISTERED_KERNELS
785
@@ -604,112 +795,87 @@ class TestKernelRegistry:
795
796
def test_register_kernel_auto_generates_fake_impl(self):
797
"""Test register_kernel auto-generates fake_impl when not provided."""
607
- with patch("vllm.kernels.helion.register.infer_fake_impl") as mock_infer:
798
+ with (
799
+ dummy_kernel_registry() as register,
800
+ patch("vllm.kernels.helion.register.infer_fake_impl") as mock_infer,
801
+ ):
802
mock_fake = Mock()
803
mock_infer.return_value = mock_fake
804
+ wrapper = register(
805
+ config_picker=lambda args, keys: "default",
806
+ )(_add_kernel)
807
611
- def original_kernel(x):
612
- return x
613
-
614
- wrapper = register_kernel(original_kernel)
615
-
616
- mock_infer.assert_called_once_with(original_kernel, None)
617
- assert wrapper._fake_impl is mock_fake
808
+ mock_infer.assert_called_once_with(_add_kernel, None)
809
+ assert wrapper._fake_impl is mock_fake
810
811
def test_register_kernel_creates_wrapper(self):
812
"""Test register_kernel creates HelionKernelWrapper."""
621
-
622
- def test_kernel(x):
623
- return x
624
-
625
- result = register_kernel("test_name")(test_kernel)
813
+ with dummy_kernel_registry() as register:
814
+ result = register("test_name", config_picker=lambda args, keys: "default")(
815
+ _add_kernel
816
+ )
817
818
assert isinstance(result, HelionKernelWrapper)
819
assert result.op_name == "test_name"
629
- assert result.raw_kernel_func is test_kernel
820
+ assert result.raw_kernel_func is _add_kernel
821
822
def test_register_kernel_auto_detects_name(self):
823
"""Test register_kernel uses function name when no name provided."""
824
+ with dummy_kernel_registry() as register:
825
+ wrapper = register(config_picker=lambda args, keys: "default")(_add_kernel)
826
634
- @register_kernel
635
- def my_test_kernel(x):
636
- return x
637
-
638
- assert my_test_kernel.op_name == "my_test_kernel"
827
+ assert wrapper.op_name == "_add_kernel"
828
829
def test_register_kernel_registers_in_global_registry(self):
830
"""Test register_kernel adds wrapper to global registry."""
642
-
643
- @register_kernel
644
- def test_kernel(x):
645
- return x
831
+ with dummy_kernel_registry() as register:
832
+ wrapper = register(
833
+ "test_kernel", config_picker=lambda args, keys: "default"
834
+ )(_add_kernel)
835
836
registered_kernels = get_registered_kernels()
837
assert "test_kernel" in registered_kernels
649
- assert registered_kernels["test_kernel"] is test_kernel
838
+ assert registered_kernels["test_kernel"] is wrapper
839
840
def test_register_kernel_passes_helion_settings(self):
841
"""Test register_kernel passes helion_settings to wrapper."""
653
- mock_settings = Mock()
654
- mock_settings.to_dict.return_value = {"debug": True}
842
+ settings = helion.Settings()
843
+ settings.print_output_code = True
844
656
- @register_kernel("test_name", helion_settings=mock_settings)
657
- def test_kernel(x):
658
- return x
845
+ with dummy_kernel_registry() as register:
846
+ result = register(
847
+ "test_name",
848
+ config_picker=lambda args, keys: "default",
849
+ helion_settings=settings,
850
+ )(_add_kernel)
851
660
- assert test_kernel.helion_settings is mock_settings
852
+ assert result.helion_settings is settings
853
854
def test_register_kernel_supports_decorator_syntax(self):
855
"""Test register_kernel works with decorator arguments."""
856
mock_fake = Mock()
857
666
- wrapper = register_kernel("custom_name", fake_impl=mock_fake)
667
-
668
- def test_kernel(x):
669
- return x
670
-
671
- result = wrapper(test_kernel)
858
+ with dummy_kernel_registry() as register:
859
+ result = register(
860
+ "custom_name",
861
+ config_picker=lambda args, keys: "default",
862
+ fake_impl=mock_fake,
863
+ )(_add_kernel)
864
865
assert result.op_name == "custom_name"
866
assert result._fake_impl is mock_fake
867
676
- def test_register_kernel_bare_decorator(self):
677
- """Test register_kernel works as bare decorator."""
678
-
679
- @register_kernel
680
- def test_kernel(x):
681
- return x
682
-
683
- assert isinstance(test_kernel, HelionKernelWrapper)
684
- assert test_kernel.op_name == "test_kernel"
685
-
686
- def test_registered_wrapper_can_register_config_picker(self):
687
- """Test that registered wrapper can register config picker."""
688
-
689
- @register_kernel
690
- def test_kernel(x):
691
- return x
692
-
693
- def my_picker(args, config_keys):
694
- return "default"
695
-
696
- result = test_kernel.register_config_picker(my_picker)
697
-
698
- assert result is my_picker
699
- assert test_kernel._config_picker is my_picker
700
-
868
def test_register_kernel_raises_on_duplicate_registration(self):
869
"""Test register_kernel raises error on duplicate names."""
870
+ with dummy_kernel_registry() as register:
871
+ register("duplicate_name", config_picker=lambda args, keys: "default")(
872
+ _add_kernel
873
+ )
874
704
- @register_kernel("duplicate_name")
705
- def kernel1(x):
706
- return x
707
-
708
- with pytest.raises(ValueError, match="already registered"):
709
-
710
- @register_kernel("duplicate_name")
711
- def kernel2(x):
712
- return x
875
+ with pytest.raises(ValueError, match="already registered"):
876
+ register("duplicate_name", config_picker=lambda args, keys: "default")(
877
+ _add_kernel
878
+ )
879
880
def test_register_kernel_rejects_autotuner_fn_in_settings(self):
881
"""Test register_kernel rejects conflicting autotuner_fn."""
@@ -718,7 +884,11 @@ class TestKernelRegistry:
884
885
with pytest.raises(ValueError, match="uses a custom autotuner"):
886
721
- @register_kernel("test", helion_settings=mock_settings)
887
+ @register_kernel(
888
+ "test",
889
+ config_picker=lambda args, keys: "default",
890
+ helion_settings=mock_settings,
891
+ )
892
def test_kernel(x):
893
return x
894
@@ -727,11 +897,47 @@ class TestKernelRegistry:
897
mock_settings = Mock()
898
mock_settings.to_dict.return_value = {"static_shapes": False}
899
730
- with patch("vllm.kernels.helion.register.logger") as mock_logger:
900
+ with (
901
+ dummy_kernel_registry() as register,
902
+ patch("vllm.kernels.helion.register.logger") as mock_logger,
903
+ ):
904
+ register(
905
+ "test",
906
+ config_picker=lambda args, keys: "default",
907
+ helion_settings=mock_settings,
908
+ )(_add_kernel)
909
732
- @register_kernel("test", helion_settings=mock_settings)
733
- def test_kernel(x):
734
- return x
910
+ mock_logger.warning.assert_not_called()
911
736
- # Should not call warning
737
- mock_logger.warning.assert_not_called()
912
+ def test_disabled_kernel_appears_in_registry(self):
913
+ """Test that a disabled wrapper is still in the global registry."""
914
+
915
+ def fake_impl(*args, **kwargs):
916
+ return torch.zeros_like(args[0])
917
+
918
+ mock_config_manager = Mock(spec=ConfigManager)
919
+ mock_config_manager.get_platform_configs = Mock(return_value={})
920
+
921
+ with (
922
+ patch(
923
+ "vllm.kernels.helion.config_manager.ConfigManager",
924
+ return_value=mock_config_manager,
925
+ ),
926
+ patch(
927
+ "vllm.kernels.helion.utils.get_canonical_gpu_name",
928
+ return_value="nvidia_h200",
929
+ ),
930
+ patch("vllm.kernels.helion.register.helion.kernel") as mock_kernel,
931
+ ):
932
+ mock_kernel.return_value = Mock(return_value=_add_kernel)
933
+
934
+ wrapper = register_kernel(
935
+ "disabled_kernel",
936
+ config_picker=lambda args, keys: "default",
937
+ fake_impl=fake_impl,
938
+ )(_add_kernel)
939
+
940
+ assert wrapper._disabled is True
941
+ registered = get_registered_kernels()
942
+ assert "disabled_kernel" in registered
943
+ assert registered["disabled_kernel"] is wrapper
vllm/kernels/helion/ops/silu_mul_fp8.py
+35
-36
@@ -22,39 +22,6 @@ from vllm.kernels.helion.register import register_kernel
22
logger = init_logger(__name__)
23
24
25
-@register_kernel # type: ignore[misc]
26
-def silu_mul_fp8(input: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
27
- original_shape = input.shape
28
- two_d = hl.specialize(original_shape[-1])
29
- d = two_d // 2
30
- output_shape = original_shape[:-1] + (d,)
31
-
32
- input_2d = input.view(-1, original_shape[-1])
33
- m = input_2d.shape[0]
34
-
35
- # TODO(gmagogsfm): Support for more float8 subtypes (e4m3fnuz, e5m2) coming
36
- out = torch.empty((m, d), device=input.device, dtype=torch.float8_e4m3fn)
37
-
38
- input_part_a = input_2d[:, :d]
39
- input_part_b = input_2d[:, d:]
40
-
41
- assert scale.numel() == 1, "Scale must be a scalar Tensor"
42
-
43
- for tile_m, tile_n in hl.tile([m, d]):
44
- a_vals = input_part_a[tile_m, tile_n]
45
- silu_result = torch.nn.functional.silu(a_vals)
46
- b_vals = input_part_b[tile_m, tile_n]
47
- result = silu_result * b_vals
48
- result_f32 = result.to(torch.float32)
49
- scale_val = hl.load(scale, [0])
50
- inv_scale = 1.0 / scale_val
51
- result_scaled = result_f32 * inv_scale
52
- out[tile_m, tile_n] = result_scaled.to(out.dtype)
53
-
54
- return out.view(output_shape)
55
-
56
-
57
-@silu_mul_fp8.register_input_generator # type: ignore[misc]
25
def generate_silu_mul_fp8_inputs() -> dict[str, tuple[Any, ...]]:
26
intermediate_sizes = [2048, 2880, 4096, 8192, 11008, 14336]
27
@@ -65,8 +32,6 @@ def generate_silu_mul_fp8_inputs() -> dict[str, tuple[Any, ...]]:
32
inputs = {}
33
for num_tokens in num_tokens_list:
34
for intermediate_size in intermediate_sizes:
68
- # Input tensor has shape (num_tokens, 2 * intermediate_size)
69
- # because silu_mul splits it into two halves
35
input_tensor = torch.randn(
36
num_tokens,
37
2 * intermediate_size,
@@ -81,7 +46,6 @@ def generate_silu_mul_fp8_inputs() -> dict[str, tuple[Any, ...]]:
46
return inputs
47
48
84
-@silu_mul_fp8.register_config_picker # type: ignore[misc]
49
def pick_silu_mul_fp8_config(
50
args: tuple[Any, ...], config_keys: list[str]
51
) -> str | None:
@@ -128,6 +92,41 @@ def pick_silu_mul_fp8_config(
92
return f"intermediate_{best_isize}_numtokens_{best_ntokens}"
93
94
95
+@register_kernel(
96
+ config_picker=pick_silu_mul_fp8_config,
97
+ input_generator=generate_silu_mul_fp8_inputs,
98
+)
99
+def silu_mul_fp8(input: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
100
+ original_shape = input.shape
101
+ two_d = hl.specialize(original_shape[-1])
102
+ d = two_d // 2
103
+ output_shape = original_shape[:-1] + (d,)
104
+
105
+ input_2d = input.view(-1, original_shape[-1])
106
+ m = input_2d.shape[0]
107
+
108
+ # TODO(gmagogsfm): Support for more float8 subtypes (e4m3fnuz, e5m2) coming
109
+ out = torch.empty((m, d), device=input.device, dtype=torch.float8_e4m3fn)
110
+
111
+ input_part_a = input_2d[:, :d]
112
+ input_part_b = input_2d[:, d:]
113
+
114
+ assert scale.numel() == 1, "Scale must be a scalar Tensor"
115
+
116
+ for tile_m, tile_n in hl.tile([m, d]):
117
+ a_vals = input_part_a[tile_m, tile_n]
118
+ silu_result = torch.nn.functional.silu(a_vals)
119
+ b_vals = input_part_b[tile_m, tile_n]
120
+ result = silu_result * b_vals
121
+ result_f32 = result.to(torch.float32)
122
+ scale_val = hl.load(scale, [0])
123
+ inv_scale = 1.0 / scale_val
124
+ result_scaled = result_f32 * inv_scale
125
+ out[tile_m, tile_n] = result_scaled.to(out.dtype)
126
+
127
+ return out.view(output_shape)
128
+
129
+
130
def silu_mul_fp8_baseline(input: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
131
output_shape = input.shape[:-1] + (input.shape[-1] // 2,)
132
out = torch.empty(output_shape, dtype=torch.float8_e4m3fn, device=input.device)
vllm/kernels/helion/register.py
+83
-91
@@ -37,7 +37,7 @@ Key Classes
37
"""
38
39
from collections.abc import Callable
40
-from typing import Any, cast, overload
40
+from typing import Any, cast
41
42
import torch
43
from torch.library import Library
@@ -95,7 +95,7 @@ def validate_helion_settings(
95
raise ValueError(
96
f"HelionKernelWrapper for '{op_name}' uses a custom autotuner via "
97
f"config picker. Remove 'autotuner_fn' from helion_settings and use "
98
- f"@{op_name}.register_config_picker instead."
98
+ f"register_kernel(..., config_picker=...) instead."
99
)
100
101
if settings_dict.get("static_shapes") is True:
@@ -169,7 +169,7 @@ class ConfiguredHelionKernel:
169
if self.config_picker is None:
170
raise RuntimeError(
171
f"No config picker registered for kernel '{self.op_name}'. "
172
- f"Use @{self.op_name}.register_config_picker to register one."
172
+ f"A config_picker must be provided to register_kernel()."
173
)
174
175
# After None check, config_picker is guaranteed to be non-None
@@ -215,7 +215,7 @@ class ConfiguredHelionKernel:
215
from vllm.kernels.helion.utils import get_canonical_gpu_name
216
217
self.platform = get_canonical_gpu_name()
218
- config_manager = ConfigManager.get_instance()
218
+ config_manager = ConfigManager()
219
self.configs = config_manager.get_platform_configs(self.op_name, self.platform)
220
221
if not self.configs:
@@ -253,7 +253,9 @@ class HelionKernelWrapper:
253
raw_kernel_func: Callable,
254
op_name: str,
255
fake_impl: Callable,
256
+ config_picker: Callable[[tuple[Any, ...], list[str]], str | None],
257
helion_settings: "helion.Settings | None" = None,
258
+ input_generator: Callable[[], dict[str, tuple[Any, ...]]] | None = None,
259
):
260
# Validate helion_settings doesn't conflict with our custom autotuner
261
validate_helion_settings(helion_settings, op_name)
@@ -262,23 +264,43 @@ class HelionKernelWrapper:
264
self.op_name = op_name
265
self._fake_impl = fake_impl
266
self.helion_settings = helion_settings
265
- self._config_picker: (
266
- Callable[[tuple[Any, ...], list[str]], str | None] | None
267
- ) = None
267
+ self._config_picker = config_picker
268
+ self._input_generator = input_generator
269
self._configured_kernel: ConfiguredHelionKernel | None = None
269
- self._input_generator: Callable[[], dict[str, tuple[Any, ...]]] | None = None
270
+ # TODO(@gmagogsfm): Remove this disable flag once integrated with vLLM IR,
271
+ # which handles op enablement/disablement.
272
+ self._disabled = False
273
+ self._disabled_reason: str | None = None
274
+
275
+ try:
276
+ if not _HOP_AVAILABLE:
277
+ self._get_or_register_custom_op()
278
+ else:
279
+ self.get_configured_op()
280
+ except ValueError as e:
281
+ self._disabled = True
282
+ self._disabled_reason = str(e)
283
+ logger.warning(
284
+ "Helion kernel '%s' is disabled: %s",
285
+ op_name,
286
+ self._disabled_reason,
287
+ )
288
289
def __call__(self, *args, **kwargs):
272
- # CustomOp fallback: register as torch custom op for torch.compile
273
- # compatibility on older PyTorch lacking HOP/EffectType support
290
+ if self._disabled:
291
+ raise RuntimeError(
292
+ f"Helion kernel '{self.op_name}' is disabled: {self._disabled_reason}"
293
+ )
294
if not _HOP_AVAILABLE:
275
- custom_op = self._get_or_register_custom_op()
276
- return custom_op(*args, **kwargs)
277
- # HOP tracing: record HigherOrderOp in the FX graph
295
+ op = getattr(torch.ops.vllm_helion, self.op_name)
296
+ return op(*args, **kwargs)
297
+ assert self._configured_kernel is not None, (
298
+ f"Kernel '{self.op_name}' was not initialized. "
299
+ "Please open an issue on GitHub."
300
+ )
301
if get_proxy_mode() is not None:
302
return self._call_via_hop(args, kwargs)
280
- # Eager: run the configured kernel directly
281
- return self.get_configured_op()(*args, **kwargs)
303
+ return self._configured_kernel(*args, **kwargs)
304
305
def _call_via_hop(
306
self,
@@ -346,42 +368,11 @@ class HelionKernelWrapper:
368
constant_args[name] = val
369
return constant_args, tensor_args
370
349
- def register_config_picker(
350
- self, picker_func: Callable[[tuple[Any, ...], list[str]], str | None]
351
- ) -> Callable[[tuple[Any, ...], list[str]], str | None]:
352
- self._config_picker = picker_func
353
- return picker_func
354
-
355
- def register_input_generator(
356
- self, generator_func: Callable[[], dict[str, tuple[Any, ...]]]
357
- ) -> Callable[[], dict[str, tuple[Any, ...]]]:
358
- """
359
- Register a function to generate inputs for autotuning and benchmarking.
360
-
361
- Args:
362
- generator_func: Function that returns dict[str, tuple] where:
363
- - key: Configuration identifier (e.g., "4096", "hidden_4096")
364
- - value: Tuple of arguments to pass to the kernel
365
-
366
- Returns:
367
- The registered function (for decorator usage)
368
-
369
- Example:
370
- @kernel_wrapper.register_input_generator
371
- def generate_inputs():
372
- return {
373
- "4096": (torch.randn(4096, device="cuda"), 0.5),
374
- "8192": (torch.randn(8192, device="cuda"), 0.5),
375
- }
376
- """
377
- self._input_generator = generator_func
378
- return generator_func
379
-
371
def get_inputs(self) -> dict[str, tuple[Any, ...]]:
372
if self._input_generator is None:
373
raise NotImplementedError(
374
f"No input generator registered for kernel '{self.op_name}'. "
384
- f"Use @{self.op_name}.register_input_generator to register one."
375
+ f"Use register_kernel(..., input_generator=...) to register one."
376
)
377
return self._input_generator()
378
@@ -401,11 +392,10 @@ class HelionKernelWrapper:
392
return autotune_kernel.autotune(inputs)
393
394
def get_configured_op(self) -> ConfiguredHelionKernel:
404
- assert self._config_picker is not None, (
405
- f"No config picker registered for kernel '{self.op_name}'. "
406
- f"Use @{self.op_name}.register_config_picker to register one."
407
- )
408
-
395
+ if self._disabled:
396
+ raise RuntimeError(
397
+ f"Helion kernel '{self.op_name}' is disabled: {self._disabled_reason}"
398
+ )
399
if self._configured_kernel is None:
400
self._configured_kernel = ConfiguredHelionKernel(
401
op_name=self.op_name,
@@ -413,7 +403,6 @@ class HelionKernelWrapper:
403
raw_kernel_func=self.raw_kernel_func,
404
helion_settings=self.helion_settings,
405
)
416
-
406
return self._configured_kernel
407
408
def _get_or_register_custom_op(self) -> Any:
@@ -466,45 +455,51 @@ def infer_fake_impl(
455
return helion_fake_kernel
456
457
469
-# Overloads are necessary for proper mypy type inference.
470
-# Without overloads, the union return type HelionKernelWrapper | Callable[...]
471
-# causes mypy to complain about missing attributes when tests do:
472
-# wrapper = register_kernel(func) # Should return HelionKernelWrapper
473
-# wrapper._fake_impl # mypy error: "Callable has no attribute _fake_impl"
474
-# The overloads tell mypy the exact return type based on the argument pattern.
475
-@overload
458
def register_kernel(
477
- op_name_or_func: Callable,
459
+ op_name: str | None = None,
460
*,
461
+ config_picker: Callable[[tuple[Any, ...], list[str]], str | None],
462
fake_impl: Callable | None = None,
463
helion_settings: "helion.Settings | None" = None,
481
-) -> HelionKernelWrapper: ...
482
-
483
-
484
-@overload
485
-def register_kernel(
486
- op_name_or_func: str | None = None,
487
- *,
488
- fake_impl: Callable | None = None,
489
- helion_settings: "helion.Settings | None" = None,
490
-) -> Callable[[Callable], HelionKernelWrapper]: ...
491
-
492
-
493
-def register_kernel(
494
- op_name_or_func: str | Callable | None = None,
495
- *,
496
- fake_impl: Callable | None = None,
497
- helion_settings: "helion.Settings | None" = None,
498
-) -> HelionKernelWrapper | Callable[[Callable], HelionKernelWrapper]:
499
- """
500
- Decorator to register a Helion kernel function as a HelionKernelWrapper.
501
-
502
- Wraps the raw kernel function in a HelionKernelWrapper and registers it
503
- in the global kernel registry. Auto-generates fake_impl if not provided.
464
+ input_generator: Callable[[], dict[str, tuple[Any, ...]]] | None = None,
465
+) -> Callable[[Callable], HelionKernelWrapper]:
466
+ """Register a Helion kernel with pre-tuned config selection.
467
+
468
+ Wraps the kernel function in a HelionKernelWrapper that eagerly builds
469
+ the configured kernel and (on older PyTorch) registers a custom op.
470
+
471
+ Args:
472
+ config_picker: Required. Function with signature
473
+ ``(args: tuple, config_keys: list[str]) -> str | None``
474
+ that picks the best config key from available options.
475
+ Return ``None`` to fall back to ``"default"``.
476
+
477
+ Example::
478
+
479
+ def pick_config(args, config_keys):
480
+ x = args[0]
481
+ hidden_size = x.shape[-1]
482
+ batch_size = x.shape[0]
483
+ for key in config_keys:
484
+ if key == f"hiddensize_{hidden_size}_batchsize_{batch_size}":
485
+ return key
486
+ return "default" if "default" in config_keys else None
487
+
488
+ input_generator: Optional. Function that returns
489
+ ``dict[str, tuple]`` where each key is a configuration
490
+ identifier (e.g. ``"4096"``, ``"hidden_4096"``) and each
491
+ value is a tuple of arguments to pass to the kernel.
492
+
493
+ Example::
494
+
495
+ def generate_inputs():
496
+ return {
497
+ "4096": (torch.randn(4096, device="cuda"), 0.5),
498
+ "8192": (torch.randn(8192, device="cuda"), 0.5),
499
+ }
500
"""
501
502
def decorator(kernel_func: Callable) -> HelionKernelWrapper:
507
- op_name = op_name_or_func if isinstance(op_name_or_func, str) else None
503
final_op_name = op_name if op_name else kernel_func.__name__
504
505
if final_op_name in _REGISTERED_KERNELS:
@@ -525,7 +520,9 @@ def register_kernel(
520
raw_kernel_func=kernel_func,
521
op_name=final_op_name,
522
fake_impl=final_fake_impl,
523
+ config_picker=config_picker,
524
helion_settings=helion_settings,
525
+ input_generator=input_generator,
526
)
527
528
_REGISTERED_KERNELS[final_op_name] = kernel_wrapper
@@ -537,9 +534,4 @@ def register_kernel(
534
535
return kernel_wrapper
536
540
- if callable(op_name_or_func) and not isinstance(op_name_or_func, str):
541
- # Bare decorator usage: @register_kernel
542
- return decorator(op_name_or_func)
543
- else:
544
- # Decorator with arguments: @register_kernel(...)
545
- return decorator
537
+ return decorator