@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 9098ce690c
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 7, 14:00 UTC

[Kernel] [Helion] [7/N] Use HOP to represent Helion Kernel call to enable fx tracing and pattern matching (#34390)

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

Yanan Cao committed Feb 27, 2026 at 09:21 UTC 9098ce690c802887fb36a8f3ee95bb18aacd2f76
3 files changed +368 -26
tests/kernels/helion/test_pattern_matching.py new
+203
@@ -0,0 +1,203 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Test make_fx tracing and inductor pattern matching with HelionKernelWrapper."""
4 +
5 +import contextlib
6 +from unittest.mock import Mock, patch
7 +
8 +import pytest
9 +import torch
10 +
11 +from vllm.utils.import_utils import has_helion
12 +
13 +if not has_helion():
14 + pytest.skip(
15 + "Helion is not installed. Install with: pip install vllm[helion]",
16 + allow_module_level=True,
17 + )
18 +
19 +import helion
20 +import helion.language as hl
21 +from helion._compat import requires_torch_version
22 +
23 +if not requires_torch_version("2.11"):
24 + pytest.skip(
25 + "HigherOrderOp requires PyTorch >= 2.11",
26 + allow_module_level=True,
27 + )
28 +
29 +from helion._compiler._dynamo.higher_order_ops import (
30 + helion_kernel_side_table,
31 + helion_kernel_wrapper_mutation,
32 +)
33 +from torch._inductor.pattern_matcher import (
34 + PatternMatcherPass,
35 + fwd_only,
36 + register_replacement,
37 + select_decomp_table,
38 +)
39 +from torch.fx.experimental.proxy_tensor import make_fx
40 +
41 +from vllm.kernels.helion.config_manager import ConfigManager
42 +from vllm.kernels.helion.register import HelionKernelWrapper
43 +
44 +
45 +@contextlib.contextmanager
46 +def _helion_mock_context():
47 + configs = {
48 + "default": helion.Config(block_sizes=[64], num_warps=2, num_stages=2),
49 + }
50 + mock_config_manager = Mock(spec=ConfigManager)
51 + mock_config_manager.get_platform_configs = Mock(return_value=configs)
52 +
53 + with (
54 + patch(
55 + "vllm.kernels.helion.config_manager.ConfigManager.get_instance",
56 + return_value=mock_config_manager,
57 + ),
58 + patch(
59 + "vllm.kernels.helion.utils.get_canonical_gpu_name",
60 + return_value="nvidia_h200",
61 + ),
62 + ):
63 + yield
64 +
65 +
66 +class TestMakeFxHop:
67 + def setup_method(self):
68 + helion_kernel_side_table.reset_table()
69 +
70 + def test_make_fx_symbolic(self):
71 + def raw_add_scale(
72 + x: torch.Tensor, y: torch.Tensor, scale: float
73 + ) -> tuple[torch.Tensor, int, torch.Tensor]:
74 + out_x = torch.empty_like(x)
75 + out_y = torch.empty_like(x)
76 + for tile in hl.tile(x.size()):
77 + out_x[tile] = x[tile] + y[tile] * scale
78 + out_y[tile] = out_x[tile] * 2.0
79 + return out_x, 42, out_y
80 +
81 + input_x = torch.randn(7, 13)
82 + input_y = torch.randn(7, 13)
83 + scale = 0.5
84 +
85 + with _helion_mock_context():
86 + wrapper = HelionKernelWrapper(
87 + raw_kernel_func=raw_add_scale,
88 + op_name="test_make_fx",
89 + fake_impl=lambda *a, **kw: None,
90 + )
91 + wrapper.register_config_picker(lambda args, keys: "default")
92 +
93 + def fn(x, y):
94 + return wrapper(x, y, scale)
95 +
96 + gm = make_fx(fn, tracing_mode="symbolic")(input_x, input_y)
97 +
98 + hop_nodes = [
99 + n
100 + for n in gm.graph.nodes
101 + if n.op == "call_function" and n.target is helion_kernel_wrapper_mutation
102 + ]
103 + assert len(hop_nodes) == 1
104 + node = hop_nodes[0]
105 +
106 + assert node.kwargs["constant_args"]["scale"] == scale
107 + assert set(node.kwargs["tensor_args"]) == {"x", "y"}
108 +
109 + specs = node.kwargs["output_spec"]["leaf_specs"]
110 + tensor_specs = [s for s in specs if s["type"] == "tensor"]
111 + scalar_specs = [s for s in specs if s["type"] == "scalar"]
112 + assert len(tensor_specs) == 2
113 + assert len(scalar_specs) == 1
114 +
115 + for spec in tensor_specs:
116 + assert spec["dtype"] == input_x.dtype
117 +
118 + assert scalar_specs[0]["scalar_value"] == 42
119 +
120 + for val in node.meta["val"]:
121 + assert all(isinstance(s, torch.SymInt) for s in val.shape)
122 +
123 + # Both out_x and out_y are empty_like(x), so output shapes == input shape
124 + input_node = next(n for n in gm.graph.nodes if n.op == "placeholder")
125 + input_shape = input_node.meta["val"].shape
126 + for val in node.meta["val"]:
127 + assert len(val.shape) == len(input_shape)
128 + for out_s, in_s in zip(val.shape, input_shape):
129 + assert out_s == in_s
130 +
131 + def test_pattern_matcher_replaces_with_helion_hop(self):
132 + def raw_silu_mul(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
133 + M, N = x.size()
134 + out = torch.empty_like(x)
135 + for tile_m, tile_n in hl.tile([M, N]):
136 + out[tile_m, tile_n] = (
137 + torch.nn.functional.silu(x[tile_m, tile_n]) * y[tile_m, tile_n]
138 + )
139 + return out
140 +
141 + with _helion_mock_context():
142 + wrapper = HelionKernelWrapper(
143 + raw_kernel_func=raw_silu_mul,
144 + op_name="test_pm_silu_mul",
145 + fake_impl=lambda *a, **kw: None,
146 + )
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
151 +
152 + def replacement(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
153 + return wrapper(x, y)
154 +
155 + inputs = [torch.randn(8, 16), torch.randn(8, 16)]
156 +
157 + pm_pass = PatternMatcherPass(pass_name="test_helion_replacement")
158 + register_replacement(pattern, replacement, inputs, fwd_only, pm_pass)
159 +
160 + def model(x, y):
161 + return torch.nn.functional.silu(x) * y
162 +
163 + decompositions = select_decomp_table()
164 + input_x = torch.randn(8, 16)
165 + input_y = torch.randn(8, 16)
166 + gm = make_fx(model, decompositions, tracing_mode="symbolic")(
167 + input_x, input_y
168 + )
169 +
170 + def count_hop_nodes(graph):
171 + return sum(
172 + 1
173 + for n in graph.nodes
174 + if n.op == "call_function"
175 + and n.target is helion_kernel_wrapper_mutation
176 + )
177 +
178 + assert count_hop_nodes(gm.graph) == 0
179 +
180 + match_count = pm_pass.apply(gm.graph)
181 + gm.graph.lint()
182 + gm.recompile()
183 +
184 + assert match_count == 1
185 + assert count_hop_nodes(gm.graph) == 1
186 +
187 + hop_node = next(
188 + n
189 + for n in gm.graph.nodes
190 + if n.op == "call_function"
191 + and n.target is helion_kernel_wrapper_mutation
192 + )
193 +
194 + # raw_silu_mul returns empty_like(x), so output shape == input shape
195 + for val in hop_node.meta["val"]:
196 + assert all(isinstance(s, torch.SymInt) for s in val.shape)
197 +
198 + input_node = next(n for n in gm.graph.nodes if n.op == "placeholder")
199 + input_shape = input_node.meta["val"].shape
200 + output_shape = hop_node.meta["val"][0].shape
201 + assert len(output_shape) == len(input_shape)
202 + for out_s, in_s in zip(output_shape, input_shape):
203 + assert out_s == in_s
tests/kernels/helion/test_register.py
+54 -10
@@ -4,8 +4,7 @@
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.
7 +including config picker registration and custom autotuner integration.
8 """
9
10 from unittest.mock import Mock, patch
@@ -25,6 +24,7 @@ import helion
24
25 from vllm.kernels.helion.config_manager import ConfigManager
26 from vllm.kernels.helion.register import (
27 + _HOP_AVAILABLE,
28 ConfiguredHelionKernel,
29 HelionKernelWrapper,
30 get_kernel_by_name,
@@ -451,9 +451,51 @@ class TestHelionKernelWrapper:
451 ):
452 wrapper.get_configured_op()
453
454 - def test_get_configured_op_returns_cached_op(self, sample_kernel, sample_configs):
455 - """Test get_configured_op returns cached op when already registered."""
454 + def test_get_configured_op_returns_cached_kernel(
455 + self, sample_kernel, sample_configs
456 + ):
457 + """Test get_configured_op returns cached ConfiguredHelionKernel."""
458 +
459 + def fake_impl(*args, **kwargs):
460 + return torch.zeros_like(args[0])
461 +
462 + def default_picker(args, config_keys):
463 + return "default"
464 +
465 + wrapper = HelionKernelWrapper(
466 + raw_kernel_func=sample_kernel,
467 + op_name="test_kernel",
468 + fake_impl=fake_impl,
469 + )
470 + wrapper._config_picker = default_picker
471 +
472 + mock_config_manager = Mock(spec=ConfigManager)
473 + mock_config_manager.get_platform_configs = Mock(return_value=sample_configs)
474 +
475 + with (
476 + patch(
477 + "vllm.kernels.helion.config_manager.ConfigManager.get_instance",
478 + return_value=mock_config_manager,
479 + ),
480 + patch(
481 + "vllm.kernels.helion.utils.get_canonical_gpu_name",
482 + return_value="nvidia_h200",
483 + ),
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 +
489 + result1 = wrapper.get_configured_op()
490 + result2 = wrapper.get_configured_op()
491 + assert result1 is result2
492
493 + @pytest.mark.skipif(
494 + _HOP_AVAILABLE, reason="CustomOp path not used when HOP available"
495 + )
496 + def test_get_or_register_custom_op_returns_cached_op(
497 + self, sample_kernel, sample_configs
498 + ):
499 def fake_impl(*args, **kwargs):
500 return torch.zeros_like(args[0])
501
@@ -488,12 +530,15 @@ class TestHelionKernelWrapper:
530 ):
531 mock_decorated = Mock()
532 mock_kernel.return_value = Mock(return_value=mock_decorated)
491 - result = wrapper.get_configured_op()
533 + result = wrapper._get_or_register_custom_op()
534 assert result is existing_op
535
494 - def test_get_configured_op_registers_new_op(self, sample_kernel, sample_configs):
495 - """Test get_configured_op creates and registers new op."""
496 -
536 + @pytest.mark.skipif(
537 + _HOP_AVAILABLE, reason="CustomOp path not used when HOP available"
538 + )
539 + def test_get_or_register_custom_op_registers_new_op(
540 + self, sample_kernel, sample_configs
541 + ):
542 def fake_impl(*args, **kwargs):
543 return torch.zeros_like(args[0])
544
@@ -542,11 +587,10 @@ class TestHelionKernelWrapper:
587 ):
588 mock_decorated = Mock()
589 mock_kernel.return_value = Mock(return_value=mock_decorated)
545 - result = wrapper.get_configured_op()
590 + result = wrapper._get_or_register_custom_op()
591
592 mock_register.assert_called_once()
593 assert result is new_op
549 - # Check that op_func is the decorated kernel, not ConfiguredHelionKernel
594 assert mock_register.call_args[1]["op_func"] is mock_decorated
595
596
vllm/kernels/helion/register.py
+111 -16
@@ -31,8 +31,8 @@ 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
34 +- HelionKernelWrapper: Wraps raw kernel + config_picker, creates configured kernels
35 +- ConfiguredHelionKernel: Platform-specific kernel with pre-tuned configs
36 - PresetConfigSearch: Custom autotuner that returns pre-tuned configs
37 """
38
@@ -53,10 +53,27 @@ if not has_helion():
53 )
54
55 import helion
56 +from helion._compat import requires_torch_version
57 from helion.autotuner.base_search import BaseAutotuner
58 from helion.runtime.config import Config
59 from helion.runtime.settings import default_autotuner_fn
60
61 +# TODO(gmagogsfm): Remove CustomOp fallback path (_get_or_register_custom_op,
62 +# vllm_helion_lib, direct_register_custom_op) once vLLM requires PyTorch >= 2.11.
63 +_HOP_AVAILABLE = requires_torch_version("2.11")
64 +
65 +if _HOP_AVAILABLE:
66 + import torch.utils._pytree as pytree
67 + from helion._compiler._dynamo.higher_order_ops import (
68 + helion_kernel_side_table,
69 + helion_kernel_wrapper_mutation,
70 + )
71 + from helion._compiler._dynamo.variables import infer_output_spec
72 + from torch.fx.experimental.proxy_tensor import (
73 + disable_proxy_modes_tracing,
74 + get_proxy_mode,
75 + )
76 +
77 logger = init_logger(__name__)
78
79 vllm_helion_lib = Library("vllm_helion", "FRAGMENT") # noqa
@@ -233,7 +250,7 @@ class ConfiguredHelionKernel:
250
251
252 class HelionKernelWrapper:
236 - """Wrapper for Helion kernels that creates config-specific PyTorch custom ops."""
253 + """Wrapper for Helion kernels with pre-tuned config selection and HOP support."""
254
255 def __init__(
256 self,
@@ -252,11 +269,86 @@ class HelionKernelWrapper:
269 self._config_picker: (
270 Callable[[tuple[Any, ...], list[str]], str | None] | None
271 ) = None
272 + self._configured_kernel: ConfiguredHelionKernel | None = None
273 self._input_generator: Callable[[], dict[str, tuple[Any, ...]]] | None = None
274
275 def __call__(self, *args, **kwargs):
258 - configured_op = self.get_configured_op()
259 - return configured_op(*args, **kwargs)
276 + # CustomOp fallback: register as torch custom op for torch.compile
277 + # compatibility on older PyTorch lacking HOP/EffectType support
278 + if not _HOP_AVAILABLE:
279 + custom_op = self._get_or_register_custom_op()
280 + return custom_op(*args, **kwargs)
281 + # HOP tracing: record HigherOrderOp in the FX graph
282 + if get_proxy_mode() is not None:
283 + return self._call_via_hop(args, kwargs)
284 + # Eager: run the configured kernel directly
285 + return self.get_configured_op()(*args, **kwargs)
286 +
287 + def _call_via_hop(
288 + self,
289 + args: tuple[Any, ...],
290 + kwargs: dict[str, Any],
291 + ) -> Any:
292 + kernel = self.get_configured_op()._decorated_kernel
293 + kernel_idx = helion_kernel_side_table.add_kernel(kernel)
294 +
295 + constant_args, tensor_args = self._partition_args(kernel, args, kwargs)
296 +
297 + all_named = {**constant_args, **tensor_args}
298 + full_args = tuple(
299 + all_named.get(n, p.default)
300 + for n, p in kernel.signature.parameters.items() # type: ignore[attr-defined]
301 + if n in all_named or p.default is not p.empty
302 + )
303 +
304 + with disable_proxy_modes_tracing():
305 + output_spec = infer_output_spec(kernel, full_args)
306 +
307 + hop_result = helion_kernel_wrapper_mutation(
308 + kernel_idx=kernel_idx,
309 + constant_args=constant_args,
310 + tensor_args=tensor_args,
311 + output_spec=output_spec,
312 + )
313 +
314 + tree_spec_str = output_spec.get("tree_spec_str")
315 + if tree_spec_str is None:
316 + return None
317 + tree_spec = pytree.treespec_loads(tree_spec_str)
318 +
319 + hop_iter = iter(hop_result)
320 + reconstructed = []
321 + for spec in output_spec["leaf_specs"]:
322 + is_constant_scalar = spec["type"] == "scalar" and not isinstance(
323 + spec.get("scalar_value"), torch.SymInt
324 + )
325 + if is_constant_scalar:
326 + reconstructed.append(spec["scalar_value"])
327 + else:
328 + reconstructed.append(next(hop_iter))
329 + return pytree.tree_unflatten(reconstructed, tree_spec)
330 +
331 + @staticmethod
332 + def _partition_args(
333 + kernel: Any,
334 + args: tuple[Any, ...],
335 + kwargs: dict[str, Any],
336 + ) -> tuple[dict[str, Any], dict[str, Any]]:
337 + constant_args: dict[str, Any] = {}
338 + tensor_args: dict[str, Any] = {}
339 + params = list(kernel.signature.parameters.keys())
340 + for i, val in enumerate(args):
341 + name = params[i]
342 + if isinstance(val, torch.Tensor):
343 + tensor_args[name] = val
344 + else:
345 + constant_args[name] = val
346 + for name, val in kwargs.items():
347 + if isinstance(val, torch.Tensor):
348 + tensor_args[name] = val
349 + else:
350 + constant_args[name] = val
351 + return constant_args, tensor_args
352
353 def register_config_picker(
354 self, picker_func: Callable[[tuple[Any, ...], list[str]], str | None]
@@ -309,29 +401,32 @@ class HelionKernelWrapper:
401 )
402 return autotune_kernel.autotune(inputs)
403
312 - def get_configured_op(self) -> Any:
404 + def get_configured_op(self) -> ConfiguredHelionKernel:
405 assert self._config_picker is not None, (
406 f"No config picker registered for kernel '{self.op_name}'. "
407 f"Use @{self.op_name}.register_config_picker to register one."
408 )
409
410 + if self._configured_kernel is None:
411 + self._configured_kernel = ConfiguredHelionKernel(
412 + op_name=self.op_name,
413 + config_picker=self._config_picker,
414 + raw_kernel_func=self.raw_kernel_func,
415 + helion_settings=self.helion_settings,
416 + )
417 +
418 + return self._configured_kernel
419 +
420 + def _get_or_register_custom_op(self) -> Any:
421 if hasattr(torch.ops.vllm_helion, self.op_name):
319 - logger.debug("Op vllm_helion::%s already registered", self.op_name)
422 return getattr(torch.ops.vllm_helion, self.op_name)
423
322 - configured_kernel = ConfiguredHelionKernel(
323 - op_name=self.op_name,
324 - config_picker=self._config_picker,
325 - raw_kernel_func=self.raw_kernel_func,
326 - helion_settings=self.helion_settings,
327 - )
424 + configured_kernel = self.get_configured_op()
425
426 logger.info("Registering op: vllm_helion::%s", self.op_name)
427 direct_register_custom_op(
428 op_name=self.op_name,
332 - op_func=configured_kernel._decorated_kernel, # Register decorated kernel
333 - # TODO(gmagogsfm): Implement automatic mutation/aliasing detection
334 - # for Helion kernels.
429 + op_func=configured_kernel._decorated_kernel,
430 mutates_args=None,
431 fake_impl=self._fake_impl,
432 target_lib=vllm_helion_lib,