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 5, 02:15 UTC
[Kernel][Helion] Fix inductor fusion of Helion HOP (#39944)
Signed-off-by: Yanan Cao <gmagogsfm@gmail.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Yanan Cao committed
Apr 16, 2026 at 04:41 UTC
edc3648966e30e8bf5f34edeee50027ddd79dc16
5 files changed
+59
-6
.buildkite/test-amd.yaml
+1
-1
@@ -769,7 +769,7 @@ steps:
769
- tests/kernels/helion/
770
- vllm/platforms/rocm.py
771
commands:
772
- - pip install helion==0.3.3
772
+ - pip install helion==1.0.0
773
- pytest -v -s kernels/helion/
774
775
.buildkite/test_areas/kernels.yaml
+1
-1
@@ -155,7 +155,7 @@ steps:
155
- vllm/utils/import_utils.py
156
- tests/kernels/helion/
157
commands:
158
- - pip install helion==0.3.3
158
+ - pip install helion==1.0.0
159
- pytest -v -s kernels/helion/
160
161
setup.py
+1
-1
@@ -1104,7 +1104,7 @@ setup(
1104
# NOTE: When updating helion version, also update CI files:
1105
# - .buildkite/test_areas/kernels.yaml
1106
# - .buildkite/test-amd.yaml
1107
- "helion": ["helion==0.3.3"],
1107
+ "helion": ["helion==1.0.0"],
1108
# Optional deps for gRPC server (vllm serve --grpc)
1109
"grpc": ["smg-grpc-servicer[vllm] >= 0.5.0"],
1110
# Optional deps for OpenTelemetry tracing
tests/kernels/helion/test_register.py
+48
@@ -36,9 +36,11 @@ from vllm.kernels.helion.register import (
36
)
37
38
if _HOP_AVAILABLE:
39
+ from helion._compat import supports_torch_compile_fusion
40
from helion._compiler._dynamo.higher_order_ops import (
41
helion_kernel_wrapper_mutation,
42
)
43
+ from torch._inductor.utils import run_and_get_code
44
45
46
def _add_kernel(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
@@ -1003,3 +1005,49 @@ class TestTorchCompileHOP:
1005
"Compiled execution result doesn't match eager execution. "
1006
f"Max difference: {torch.max(torch.abs(compiled_result - eager_result))}"
1007
)
1008
+
1009
+ @pytest.mark.skipif(
1010
+ not (_HOP_AVAILABLE and supports_torch_compile_fusion()),
1011
+ reason="Requires PyTorch with Helion inductor fusion support",
1012
+ )
1013
+ def test_inductor_backend_compiles_helion_hop(self):
1014
+ """Test torch.compile with inductor backend and Helion fusion enabled."""
1015
+
1016
+ configs = {"default": helion.Config(block_sizes=[4, 4])}
1017
+
1018
+ with dummy_kernel_registry(configs=configs) as register:
1019
+ add_helion_kernel = register(
1020
+ op_name="test_inductor_add_kernel",
1021
+ config_picker=lambda args, keys: "default",
1022
+ helion_settings=helion.Settings(
1023
+ torch_compile_fusion=True, static_shapes=False
1024
+ ),
1025
+ )(_add_kernel)
1026
+
1027
+ def f(x, y):
1028
+ x = x * 2.0
1029
+ y = y + 1.0
1030
+ out = add_helion_kernel(x, y)
1031
+ return out.relu()
1032
+
1033
+ torch._dynamo.reset()
1034
+ compiled_f = torch.compile(f, backend="inductor", fullgraph=True)
1035
+
1036
+ x = torch.randn(4, 4, device="cuda")
1037
+ y = torch.randn(4, 4, device="cuda")
1038
+
1039
+ compiled_result, source_codes = run_and_get_code(compiled_f, x, y)
1040
+ eager_result = f(x, y)
1041
+
1042
+ assert torch.allclose(compiled_result, eager_result, atol=1e-5, rtol=1e-5), (
1043
+ "Inductor-compiled result doesn't match eager execution. "
1044
+ f"Max difference: {torch.max(torch.abs(compiled_result - eager_result))}"
1045
+ )
1046
+
1047
+ # With fusion enabled, prologue/epilogue ops should be fused into
1048
+ # a single triton kernel rather than generating separate kernels.
1049
+ kernel_count = sum(code.count("@triton.jit") for code in source_codes)
1050
+ assert kernel_count == 1, (
1051
+ f"Expected 1 fused triton kernel, got {kernel_count}. "
1052
+ "Prologue/epilogue ops were not fused into the Helion kernel."
1053
+ )
vllm/kernels/helion/register.py
+8
-3
@@ -63,8 +63,10 @@ from helion.runtime.settings import default_autotuner_fn
63
_HOP_AVAILABLE = requires_torch_version("2.11")
64
65
if _HOP_AVAILABLE:
66
+ from helion._compat import supports_torch_compile_fusion
67
from helion._compiler._dynamo.higher_order_ops import helion_kernel_side_table
68
from helion._compiler._dynamo.variables import HelionKernelVariable
69
+ from helion.runtime.kernel import Kernel
70
from torch._dynamo.guards import GuardBuilder
71
from torch._dynamo.variables.builder import VariableBuilder
72
@@ -475,19 +477,22 @@ if _HOP_AVAILABLE:
477
"""Register HelionKernelWrapper with Dynamo's VariableBuilder.
478
479
When Dynamo encounters a HelionKernelWrapper during tracing, this
478
- extracts the underlying Helion Kernel, registers it in the side table,
479
- and returns Helion's own HelionKernelVariable to handle HOP emission.
480
+ extracts the underlying Helion Kernel and delegates to Helion's own
481
+ registered Kernel handler, which handles HOP emission, side table
482
+ registration, and inductor lowering setup.
483
"""
484
485
def wrap_helion_kernel_wrapper(
486
builder: VariableBuilder, value: HelionKernelWrapper
487
):
488
kernel = value.get_configured_op()._decorated_kernel
489
+ if supports_torch_compile_fusion():
490
+ helion_handler = VariableBuilder._type_dispatch()[Kernel]
491
+ return helion_handler(builder, kernel)
492
kernel_idx = helion_kernel_side_table.add_kernel(kernel)
493
builder.install_guards(GuardBuilder.ID_MATCH)
494
return HelionKernelVariable(kernel, kernel_idx, source=builder.source)
495
490
- # Register with Dynamo's type dispatch system
496
dispatch = VariableBuilder._type_dispatch()
497
dispatch[HelionKernelWrapper] = wrap_helion_kernel_wrapper
498