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] [3/N] Helion kernel registry (#33203)
Signed-off-by: Yanan Cao <gmagogsfm@gmail.com>
Yanan Cao committed
Jan 30, 2026 at 23:38 UTC
d5c41db35b33d348bc8b913eab6d11a20a64f168
3 files changed
+319
-3
tests/kernels/helion/test_register.py
+191
@@ -27,6 +27,9 @@ from vllm.kernels.helion.config_manager import ConfigManager
27
from vllm.kernels.helion.register import (
28
ConfiguredHelionKernel,
29
HelionKernelWrapper,
30
+ get_kernel_by_name,
31
+ get_registered_kernels,
32
+ register_kernel,
33
validate_helion_settings,
34
)
35
@@ -545,3 +548,191 @@ class TestHelionKernelWrapper:
548
assert result is new_op
549
# Check that op_func is the decorated kernel, not ConfiguredHelionKernel
550
assert mock_register.call_args[1]["op_func"] is mock_decorated
551
+
552
+
553
+class TestKernelRegistry:
554
+ """Test suite for kernel registry functionality."""
555
+
556
+ def setup_method(self):
557
+ """Clear the registry before each test."""
558
+ from vllm.kernels.helion.register import _REGISTERED_KERNELS
559
+
560
+ _REGISTERED_KERNELS.clear()
561
+
562
+ def test_get_registered_kernels_returns_copy(self):
563
+ """Test get_registered_kernels returns copy of registry."""
564
+ result1 = get_registered_kernels()
565
+ result2 = get_registered_kernels()
566
+
567
+ # Should be separate objects
568
+ assert result1 is not result2
569
+ # Should have same content
570
+ assert result1 == result2
571
+
572
+ def test_get_kernel_by_name_returns_kernel(self):
573
+ """Test get_kernel_by_name returns registered kernel."""
574
+ wrapper = HelionKernelWrapper(
575
+ raw_kernel_func=Mock(),
576
+ op_name="test_kernel",
577
+ fake_impl=Mock(),
578
+ )
579
+
580
+ from vllm.kernels.helion.register import _REGISTERED_KERNELS
581
+
582
+ _REGISTERED_KERNELS["test_kernel"] = wrapper
583
+
584
+ result = get_kernel_by_name("test_kernel")
585
+ assert result is wrapper
586
+
587
+ def test_get_kernel_by_name_returns_none_for_missing(self):
588
+ """Test get_kernel_by_name returns None for missing kernel."""
589
+ result = get_kernel_by_name("nonexistent")
590
+ assert result is None
591
+
592
+ def test_register_kernel_auto_generates_fake_impl(self):
593
+ """Test register_kernel auto-generates fake_impl when not provided."""
594
+ with patch("vllm.kernels.helion.register.infer_fake_impl") as mock_infer:
595
+ mock_fake = Mock()
596
+ mock_infer.return_value = mock_fake
597
+
598
+ def original_kernel(x):
599
+ return x
600
+
601
+ wrapper = register_kernel(original_kernel)
602
+
603
+ mock_infer.assert_called_once_with(original_kernel, None)
604
+ assert wrapper._fake_impl is mock_fake
605
+
606
+ def test_register_kernel_creates_wrapper(self):
607
+ """Test register_kernel creates HelionKernelWrapper."""
608
+
609
+ def test_kernel(x):
610
+ return x
611
+
612
+ result = register_kernel("test_name")(test_kernel)
613
+
614
+ assert isinstance(result, HelionKernelWrapper)
615
+ assert result.op_name == "test_name"
616
+ assert result.raw_kernel_func is test_kernel
617
+
618
+ def test_register_kernel_auto_detects_name(self):
619
+ """Test register_kernel uses function name when no name provided."""
620
+
621
+ @register_kernel
622
+ def my_test_kernel(x):
623
+ return x
624
+
625
+ assert my_test_kernel.op_name == "my_test_kernel"
626
+
627
+ def test_register_kernel_registers_in_global_registry(self):
628
+ """Test register_kernel adds wrapper to global registry."""
629
+
630
+ @register_kernel
631
+ def test_kernel(x):
632
+ return x
633
+
634
+ registered_kernels = get_registered_kernels()
635
+ assert "test_kernel" in registered_kernels
636
+ assert registered_kernels["test_kernel"] is test_kernel
637
+
638
+ def test_register_kernel_passes_helion_settings(self):
639
+ """Test register_kernel passes helion_settings to wrapper."""
640
+ mock_settings = Mock()
641
+ mock_settings.to_dict.return_value = {"debug": True}
642
+
643
+ @register_kernel("test_name", helion_settings=mock_settings)
644
+ def test_kernel(x):
645
+ return x
646
+
647
+ assert test_kernel.helion_settings is mock_settings
648
+
649
+ def test_register_kernel_supports_decorator_syntax(self):
650
+ """Test register_kernel works with decorator arguments."""
651
+ mock_fake = Mock()
652
+
653
+ wrapper = register_kernel("custom_name", fake_impl=mock_fake)
654
+
655
+ def test_kernel(x):
656
+ return x
657
+
658
+ result = wrapper(test_kernel)
659
+
660
+ assert result.op_name == "custom_name"
661
+ assert result._fake_impl is mock_fake
662
+
663
+ def test_register_kernel_bare_decorator(self):
664
+ """Test register_kernel works as bare decorator."""
665
+
666
+ @register_kernel
667
+ def test_kernel(x):
668
+ return x
669
+
670
+ assert isinstance(test_kernel, HelionKernelWrapper)
671
+ assert test_kernel.op_name == "test_kernel"
672
+
673
+ def test_registered_wrapper_can_register_config_picker(self):
674
+ """Test that registered wrapper can register config picker."""
675
+
676
+ @register_kernel
677
+ def test_kernel(x):
678
+ return x
679
+
680
+ def my_picker(args, config_keys):
681
+ return "default"
682
+
683
+ result = test_kernel.register_config_picker(my_picker)
684
+
685
+ assert result is my_picker
686
+ assert test_kernel._config_picker is my_picker
687
+
688
+ def test_register_kernel_raises_on_duplicate_registration(self):
689
+ """Test register_kernel raises error on duplicate names."""
690
+
691
+ @register_kernel("duplicate_name")
692
+ def kernel1(x):
693
+ return x
694
+
695
+ with pytest.raises(ValueError, match="already registered"):
696
+
697
+ @register_kernel("duplicate_name")
698
+ def kernel2(x):
699
+ return x
700
+
701
+ def test_register_kernel_rejects_autotuner_fn_in_settings(self):
702
+ """Test register_kernel rejects conflicting autotuner_fn."""
703
+ mock_settings = Mock()
704
+ mock_settings.to_dict.return_value = {"autotuner_fn": Mock()}
705
+
706
+ with pytest.raises(ValueError, match="uses a custom autotuner"):
707
+
708
+ @register_kernel("test", helion_settings=mock_settings)
709
+ def test_kernel(x):
710
+ return x
711
+
712
+ def test_register_kernel_warns_with_static_shapes_true(self):
713
+ """Test register_kernel warns when static_shapes=True."""
714
+ mock_settings = Mock()
715
+ mock_settings.to_dict.return_value = {"static_shapes": True}
716
+
717
+ with patch("vllm.kernels.helion.register.logger") as mock_logger:
718
+
719
+ @register_kernel("test", helion_settings=mock_settings)
720
+ def test_kernel(x):
721
+ return x
722
+
723
+ mock_logger.warning.assert_called_once()
724
+ assert "static_shapes=True" in mock_logger.warning.call_args[0][0]
725
+
726
+ def test_register_kernel_no_warning_with_static_shapes_false(self):
727
+ """Test register_kernel doesn't warn with static_shapes=False."""
728
+ mock_settings = Mock()
729
+ mock_settings.to_dict.return_value = {"static_shapes": False}
730
+
731
+ with patch("vllm.kernels.helion.register.logger") as mock_logger:
732
+
733
+ @register_kernel("test", helion_settings=mock_settings)
734
+ def test_kernel(x):
735
+ return x
736
+
737
+ # Should not call warning
738
+ mock_logger.warning.assert_not_called()
vllm/kernels/helion/__init__.py
+6
@@ -9,6 +9,9 @@ from vllm.kernels.helion.config_manager import (
9
from vllm.kernels.helion.register import (
10
ConfiguredHelionKernel,
11
HelionKernelWrapper,
12
+ get_kernel_by_name,
13
+ get_registered_kernels,
14
+ register_kernel,
15
vllm_helion_lib,
16
)
17
from vllm.kernels.helion.utils import canonicalize_gpu_name, get_canonical_gpu_name
@@ -20,6 +23,9 @@ __all__ = [
23
# Kernel registration
24
"ConfiguredHelionKernel",
25
"HelionKernelWrapper",
26
+ "get_kernel_by_name",
27
+ "get_registered_kernels",
28
+ "register_kernel",
29
"vllm_helion_lib",
30
# Utilities
31
"canonicalize_gpu_name",
vllm/kernels/helion/register.py
+122
-3
@@ -37,7 +37,7 @@ Key Classes
37
"""
38
39
from collections.abc import Callable
40
-from typing import Any
40
+from typing import Any, cast, overload
41
42
import torch
43
from torch.library import Library
@@ -114,7 +114,7 @@ class ConfiguredHelionKernel:
114
def __init__(
115
self,
116
op_name: str,
117
- config_picker: Callable[[tuple[Any, ...], list[str]], str | None],
117
+ config_picker: Callable[[tuple[Any, ...], list[str]], str | None] | None,
118
raw_kernel_func: Callable,
119
helion_settings: "helion.Settings | None" = None,
120
):
@@ -140,9 +140,16 @@ class ConfiguredHelionKernel:
140
f"Use @{self.op_name}.register_config_picker to register one."
141
)
142
143
+ # After None check, config_picker is guaranteed to be non-None
144
+ assert self.config_picker is not None
145
+
146
def key_computer(*args):
147
config_keys = list(self.configs.keys())
145
- selected_key = self.config_picker(args, config_keys)
148
+ # Cast is safe because we checked for None above
149
+ config_picker = cast(
150
+ Callable[[tuple[Any, ...], list[str]], str | None], self.config_picker
151
+ )
152
+ selected_key = config_picker(args, config_keys)
153
if selected_key:
154
return selected_key
155
return "default" if "default" in self.configs else None
@@ -272,3 +279,115 @@ class HelionKernelWrapper:
279
target_lib=vllm_helion_lib,
280
)
281
return getattr(torch.ops.vllm_helion, self.op_name)
282
+
283
+
284
+# Global registry for tracking all registered HelionKernelWrapper instances
285
+_REGISTERED_KERNELS: dict[str, HelionKernelWrapper] = {}
286
+
287
+
288
+def get_registered_kernels() -> dict[str, HelionKernelWrapper]:
289
+ return _REGISTERED_KERNELS.copy()
290
+
291
+
292
+def get_kernel_by_name(kernel_name: str) -> HelionKernelWrapper | None:
293
+ return _REGISTERED_KERNELS.get(kernel_name)
294
+
295
+
296
+def infer_fake_impl(
297
+ kernel_func: Callable,
298
+ helion_settings: "helion.Settings | None" = None,
299
+) -> Callable:
300
+ def helion_fake_kernel(*args, **kwargs):
301
+ kernel_kwargs = {}
302
+ if helion_settings:
303
+ kernel_kwargs.update(helion_settings.to_dict())
304
+
305
+ temp_decorated_kernel = helion.kernel(**kernel_kwargs)(kernel_func)
306
+
307
+ # Bind with args to get config_spec, then get a valid default config
308
+ bound = temp_decorated_kernel.bind(args)
309
+ default_config = bound.config_spec.default_config()
310
+ compiled_runner = bound.compile_config(default_config)
311
+
312
+ return compiled_runner(*args, **kwargs, _launcher=lambda *a, **kw: None)
313
+
314
+ return helion_fake_kernel
315
+
316
+
317
+# Overloads are necessary for proper mypy type inference.
318
+# Without overloads, the union return type HelionKernelWrapper | Callable[...]
319
+# causes mypy to complain about missing attributes when tests do:
320
+# wrapper = register_kernel(func) # Should return HelionKernelWrapper
321
+# wrapper._fake_impl # mypy error: "Callable has no attribute _fake_impl"
322
+# The overloads tell mypy the exact return type based on the argument pattern.
323
+@overload
324
+def register_kernel(
325
+ op_name_or_func: Callable,
326
+ *,
327
+ fake_impl: Callable | None = None,
328
+ helion_settings: "helion.Settings | None" = None,
329
+) -> HelionKernelWrapper: ...
330
+
331
+
332
+@overload
333
+def register_kernel(
334
+ op_name_or_func: str | None = None,
335
+ *,
336
+ fake_impl: Callable | None = None,
337
+ helion_settings: "helion.Settings | None" = None,
338
+) -> Callable[[Callable], HelionKernelWrapper]: ...
339
+
340
+
341
+def register_kernel(
342
+ op_name_or_func: str | Callable | None = None,
343
+ *,
344
+ fake_impl: Callable | None = None,
345
+ helion_settings: "helion.Settings | None" = None,
346
+) -> HelionKernelWrapper | Callable[[Callable], HelionKernelWrapper]:
347
+ """
348
+ Decorator to register a Helion kernel function as a HelionKernelWrapper.
349
+
350
+ Wraps the raw kernel function in a HelionKernelWrapper and registers it
351
+ in the global kernel registry. Auto-generates fake_impl if not provided.
352
+ """
353
+
354
+ def decorator(kernel_func: Callable) -> HelionKernelWrapper:
355
+ op_name = op_name_or_func if isinstance(op_name_or_func, str) else None
356
+ final_op_name = op_name if op_name else kernel_func.__name__
357
+
358
+ if final_op_name in _REGISTERED_KERNELS:
359
+ raise ValueError(
360
+ f"Helion kernel '{final_op_name}' is already registered. "
361
+ f"Use a different op_name or check for duplicate registrations."
362
+ )
363
+
364
+ final_fake_impl = fake_impl
365
+ if final_fake_impl is None:
366
+ final_fake_impl = infer_fake_impl(kernel_func, helion_settings)
367
+ logger.debug(
368
+ "Auto-generated fake_impl for Helion kernel '%s'",
369
+ kernel_func.__name__,
370
+ )
371
+
372
+ kernel_wrapper = HelionKernelWrapper(
373
+ raw_kernel_func=kernel_func,
374
+ op_name=final_op_name,
375
+ fake_impl=final_fake_impl,
376
+ helion_settings=helion_settings,
377
+ )
378
+
379
+ _REGISTERED_KERNELS[final_op_name] = kernel_wrapper
380
+
381
+ logger.info(
382
+ "Registered Helion kernel '%s' as HelionKernelWrapper",
383
+ kernel_func.__name__,
384
+ )
385
+
386
+ return kernel_wrapper
387
+
388
+ if callable(op_name_or_func) and not isinstance(op_name_or_func, str):
389
+ # Bare decorator usage: @register_kernel
390
+ return decorator(op_name_or_func)
391
+ else:
392
+ # Decorator with arguments: @register_kernel(...)
393
+ return decorator