gahow/agentic-kvc
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add vLLM v0.18.1 source tree with KV transfer abort fix

Browse Source 
third_party/vllm/ now tracked in git for direct patch management.
Based on vLLM v0.18.1 release with one patch applied:

  vllm/v1/core/sched/scheduler.py:
    Replace fatal assert with graceful skip when KV transfer callback
    arrives for an already-aborted request during PD disaggregated serving.

Future vLLM modifications should be made directly in third_party/vllm/
and committed normally. The patches/ directory is kept as documentation
of what changed from upstream.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
2026-05-22 00:30:38 +08:00
parent b6591950bc
commit 445e491123
4285 changed files with 1111303 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

366
third_party/vllm/tests/kernels/helion/test_config_manager.py vendored Normal file
View File

@@ -0,0 +1,366 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Unit tests for Helion ConfigManager and ConfigSet.
Tests the simplified configuration management system for Helion custom kernels.
"""
import json
import tempfile
from pathlib import Path
import pytest
from vllm.utils.import_utils import has_helion
# Skip entire module if helion is not available
if not has_helion():
pytest.skip(
"Helion is not installed. Install with: pip install vllm[helion]",
allow_module_level=True,
)
import helion
from vllm.kernels.helion.config_manager import (
ConfigManager,
ConfigSet,
)
@pytest.fixture(autouse=True)
def reset_config_manager_singleton():
"""Reset ConfigManager singleton before each test."""
ConfigManager.reset_instance()
yield
ConfigManager.reset_instance()
class TestConfigSet:
"""Test suite for ConfigSet class."""
def test_config_set_creation(self):
"""Test creating an empty ConfigSet."""
config_set = ConfigSet("test_kernel")
assert config_set.kernel_name == "test_kernel"
assert config_set.get_platforms() == []
def test_config_set_from_dict(self):
"""Test creating ConfigSet from dictionary data."""
# Use realistic config data that helion.Config can handle
config_data = {
"block_sizes": [32, 16],
"num_warps": 4,
"num_stages": 3,
"pid_type": "persistent_interleaved",
}
data = {"h100": {"batch_32_hidden_4096": config_data}}
config_set = ConfigSet.from_dict("test_kernel", data)
assert config_set.kernel_name == "test_kernel"
assert config_set.get_platforms() == ["h100"]
# Verify the config was created correctly
config = config_set.get_config("h100", "batch_32_hidden_4096")
assert isinstance(config, helion.Config)
assert config.block_sizes == [32, 16]
assert config.num_warps == 4
assert config.num_stages == 3
assert config.pid_type == "persistent_interleaved"
def test_config_set_get_config_keyerror(self):
"""Test that accessing non-existent configs raises informative KeyErrors."""
config_set = ConfigSet("test_kernel")
with pytest.raises(KeyError, match="platform 'h100' not found"):
config_set.get_config("h100", "batch_32_hidden_4096")
# Use realistic config data
config_data = {"num_warps": 8, "num_stages": 4}
data = {"h100": {"batch_64_hidden_2048": config_data}}
config_set = ConfigSet.from_dict("test_kernel", data)
with pytest.raises(
KeyError, match="config_key 'batch_32_hidden_4096' not found"
):
config_set.get_config("h100", "batch_32_hidden_4096")
def test_config_set_get_platforms(self):
"""Test get_platforms method."""
# Use realistic config data
config1 = {"num_warps": 4, "num_stages": 3}
config2 = {"num_warps": 8, "num_stages": 5}
data = {
"h100": {"batch_32_hidden_4096": config1},
"a100": {"batch_16_hidden_2048": config2},
}
config_set = ConfigSet.from_dict("test_kernel", data)
platforms = config_set.get_platforms()
assert platforms == ["a100", "h100"] # Should be sorted
def test_config_set_get_config_keys(self):
"""Test get_config_keys method."""
# Use realistic config data
config1 = {"num_warps": 4, "num_stages": 3}
config2 = {"num_warps": 8, "num_stages": 5}
data = {
"h100": {
"batch_32_hidden_4096": config1,
"batch_64_hidden_2048": config2,
}
}
config_set = ConfigSet.from_dict("test_kernel", data)
config_keys = config_set.get_config_keys("h100")
assert config_keys == ["batch_32_hidden_4096", "batch_64_hidden_2048"]
assert config_set.get_config_keys("v100") == []
def test_config_set_to_dict(self):
"""Test converting ConfigSet to dictionary."""
# Use realistic config data
original_config = {
"block_sizes": [64, 32],
"num_warps": 16,
"num_stages": 4,
"pid_type": "persistent_blocked",
}
original_data = {"h100": {"batch_32_hidden_4096": original_config}}
config_set = ConfigSet.from_dict("test_kernel", original_data)
result_data = config_set.to_dict()
# The result should match the original (Config roundtrip should work)
assert result_data == original_data
class TestConfigManager:
"""Test suite for ConfigManager class."""
def test_config_manager_creation_default_base_dir(self):
"""Test creating ConfigManager with default base directory."""
manager = ConfigManager()
assert manager._base_dir.name == "configs"
def test_config_manager_creation_custom_base_dir(self):
"""Test creating ConfigManager with custom base directory."""
custom_dir = "/tmp/custom_configs"
manager = ConfigManager(base_dir=custom_dir)
# Paths are resolved, so compare with resolved path
assert manager._base_dir == Path(custom_dir).resolve()
def test_get_config_file_path(self):
"""Test getting config file path for a kernel."""
manager = ConfigManager(base_dir="/tmp")
dir_path = manager.get_config_file_path("silu_mul_fp8")
assert dir_path == Path("/tmp/silu_mul_fp8")
file_path = manager.get_config_file_path("silu_mul_fp8", "nvidia_h100")
assert file_path == Path("/tmp/silu_mul_fp8/nvidia_h100.json")
def test_ensure_base_dir_exists(self):
"""Test ensuring base directory exists."""
with tempfile.TemporaryDirectory() as temp_dir:
base_dir = Path(temp_dir) / "non_existent" / "configs"
manager = ConfigManager(base_dir=base_dir)
assert not base_dir.exists()
returned_path = manager.ensure_base_dir_exists()
assert base_dir.exists()
assert base_dir.is_dir()
assert returned_path == base_dir
def test_load_config_set_file_not_exists(self):
"""Test loading config set when file doesn't exist."""
with tempfile.TemporaryDirectory() as temp_dir:
manager = ConfigManager(base_dir=temp_dir)
config_set = manager.load_config_set("non_existent_kernel")
assert isinstance(config_set, ConfigSet)
assert config_set.kernel_name == "non_existent_kernel"
assert config_set.get_platforms() == []
def test_load_config_set_valid_file(self):
"""Test loading config set from per-platform files."""
with tempfile.TemporaryDirectory() as temp_dir:
kernel_config = {
"block_sizes": [128, 64],
"num_warps": 8,
"num_stages": 6,
"pid_type": "persistent_interleaved",
}
kernel_dir = Path(temp_dir) / "test_kernel"
kernel_dir.mkdir()
platform_file = kernel_dir / "h100.json"
with open(platform_file, "w") as f:
json.dump({"batch_32_hidden_4096": kernel_config}, f)
manager = ConfigManager(base_dir=temp_dir)
config_set = manager.load_config_set("test_kernel")
assert isinstance(config_set, ConfigSet)
assert config_set.kernel_name == "test_kernel"
assert config_set.get_platforms() == ["h100"]
config = config_set.get_config("h100", "batch_32_hidden_4096")
assert isinstance(config, helion.Config)
assert config.block_sizes == [128, 64]
assert config.num_warps == 8
def test_load_config_set_invalid_json(self):
"""Test loading config set from file with invalid JSON."""
with tempfile.TemporaryDirectory() as temp_dir:
kernel_dir = Path(temp_dir) / "test_kernel"
kernel_dir.mkdir()
config_file = kernel_dir / "h100.json"
with open(config_file, "w") as f:
f.write("invalid json content {")
manager = ConfigManager(base_dir=temp_dir)
config_set = manager.load_config_set("test_kernel")
assert isinstance(config_set, ConfigSet)
assert config_set.kernel_name == "test_kernel"
assert config_set.get_platforms() == []
def test_save_config_set(self):
"""Test saving ConfigSet to per-platform files."""
with tempfile.TemporaryDirectory() as temp_dir:
kernel_config = {
"block_sizes": [256, 128],
"num_warps": 16,
"num_stages": 8,
"pid_type": "persistent_blocked",
}
data = {"h100": {"batch_32_hidden_4096": kernel_config}}
config_set = ConfigSet.from_dict("test_kernel", data)
manager = ConfigManager(base_dir=temp_dir)
saved_path = manager.save_config_set(config_set)
expected_dir = Path(temp_dir) / "test_kernel"
assert saved_path == expected_dir
assert saved_path.is_dir()
platform_file = expected_dir / "h100.json"
assert platform_file.exists()
with open(platform_file) as f:
loaded_data = json.load(f)
assert loaded_data == data["h100"]
def test_save_config_set_creates_directory(self):
"""Test that save_config_set creates parent directories if needed."""
with tempfile.TemporaryDirectory() as temp_dir:
nested_dir = Path(temp_dir) / "nested" / "configs"
data = {"h100": {"default": {"num_warps": 4}}}
config_set = ConfigSet.from_dict("test_kernel", data)
manager = ConfigManager(base_dir=nested_dir)
saved_path = manager.save_config_set(config_set)
assert nested_dir.exists()
assert nested_dir.is_dir()
assert saved_path.is_dir()
assert (saved_path / "h100.json").exists()
def test_get_platform_configs(self):
"""Test getting all configs for a specific platform."""
with tempfile.TemporaryDirectory() as temp_dir:
config_1 = {"num_warps": 4, "num_stages": 3, "block_sizes": [64, 32]}
config_2 = {"num_warps": 8, "num_stages": 5, "block_sizes": [128, 64]}
default_config = {
"num_warps": 16,
"num_stages": 7,
"block_sizes": [256, 128],
}
config_3 = {"num_warps": 2, "num_stages": 2, "block_sizes": [32, 16]}
kernel_dir = Path(temp_dir) / "test_kernel"
kernel_dir.mkdir()
with open(kernel_dir / "h100.json", "w") as f:
json.dump(
{
"batch_32_hidden_4096": config_1,
"batch_64_hidden_2048": config_2,
"default": default_config,
},
f,
)
with open(kernel_dir / "a100.json", "w") as f:
json.dump({"batch_16_hidden_1024": config_3}, f)
manager = ConfigManager(base_dir=temp_dir)
h100_configs = manager.get_platform_configs("test_kernel", "h100")
assert len(h100_configs) == 3
assert "batch_32_hidden_4096" in h100_configs
assert "batch_64_hidden_2048" in h100_configs
assert "default" in h100_configs
for config in h100_configs.values():
assert isinstance(config, helion.Config)
assert h100_configs["batch_32_hidden_4096"].num_warps == 4
assert h100_configs["default"].num_stages == 7
a100_configs = manager.get_platform_configs("test_kernel", "a100")
assert len(a100_configs) == 1
assert "batch_16_hidden_1024" in a100_configs
assert isinstance(a100_configs["batch_16_hidden_1024"], helion.Config)
assert a100_configs["batch_16_hidden_1024"].num_warps == 2
nonexistent_configs = manager.get_platform_configs("test_kernel", "v100")
assert len(nonexistent_configs) == 0
def test_singleton_returns_same_instance(self):
"""Test that ConfigManager returns the same instance on repeated calls."""
manager1 = ConfigManager(base_dir="/tmp/test_singleton")
manager2 = ConfigManager(base_dir="/tmp/test_singleton")
assert manager1 is manager2
def test_singleton_with_default_base_dir(self):
"""Test singleton behavior with default base directory."""
manager1 = ConfigManager()
manager2 = ConfigManager()
assert manager1 is manager2
assert manager1._base_dir == manager2._base_dir
def test_singleton_error_on_different_base_dir(self):
"""Test that ConfigManager raises error when created with different base_dir."""
ConfigManager(base_dir="/tmp/first_dir")
with pytest.raises(ValueError, match="singleton already exists"):
ConfigManager(base_dir="/tmp/different_dir")
def test_reset_instance_allows_new_base_dir(self):
"""Test that reset_instance allows creating with a new base_dir."""
manager1 = ConfigManager(base_dir="/tmp/first_dir")
assert manager1._base_dir == Path("/tmp/first_dir").resolve()
ConfigManager.reset_instance()
manager2 = ConfigManager(base_dir="/tmp/second_dir")
assert manager2._base_dir == Path("/tmp/second_dir").resolve()
assert manager1 is not manager2
def test_get_instance_returns_existing(self):
"""Test that get_instance returns the existing singleton."""
manager1 = ConfigManager(base_dir="/tmp/test_get_instance")
manager2 = ConfigManager.get_instance()
assert manager1 is manager2
def test_get_instance_raises_if_not_initialized(self):
"""Test that get_instance raises RuntimeError if no instance exists."""
with pytest.raises(RuntimeError, match="has not been created"):
ConfigManager.get_instance()

45
third_party/vllm/tests/kernels/helion/test_helion_available.py vendored Normal file
View File

@@ -0,0 +1,45 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Tests for Helion kernel availability and basic functionality.
This module demonstrates the pattern for testing optional Helion kernels.
Tests in this directory will be skipped if Helion is not installed.
"""
import pytest
from vllm.utils.import_utils import has_helion
# Skip entire module if helion is not available
if not has_helion():
pytest.skip(
"Helion is not installed. Install with: pip install vllm[helion]",
allow_module_level=True,
)
import helion
import helion.language as hl
import torch
def test_helion_kernel_compilation_smoke():
"""Smoke test: compile and run a simple Helion kernel."""
@helion.kernel(autotune_effort="none")
def add_kernel(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
out = torch.empty_like(x)
for tile in hl.tile(x.size()):
out[tile] = x[tile] + y[tile]
return out
# Create test tensors
x = torch.randn(1024, device="cuda", dtype=torch.float32)
y = torch.randn(1024, device="cuda", dtype=torch.float32)
# Run the helion kernel
result = add_kernel(x, y)
# Verify correctness
expected = x + y
assert torch.allclose(result, expected), "Helion kernel output mismatch"

203
third_party/vllm/tests/kernels/helion/test_pattern_matching.py vendored Normal file
View File

@@ -0,0 +1,203 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Test make_fx tracing and inductor pattern matching with HelionKernelWrapper."""
import contextlib
from unittest.mock import Mock, patch
import pytest
import torch
from vllm.utils.import_utils import has_helion
if not has_helion():
pytest.skip(
"Helion is not installed. Install with: pip install vllm[helion]",
allow_module_level=True,
)
import helion
import helion.language as hl
from helion._compat import requires_torch_version
if not requires_torch_version("2.11"):
pytest.skip(
"HigherOrderOp requires PyTorch >= 2.11",
allow_module_level=True,
)
from helion._compiler._dynamo.higher_order_ops import (
helion_kernel_side_table,
helion_kernel_wrapper_mutation,
)
from torch._inductor.pattern_matcher import (
PatternMatcherPass,
fwd_only,
register_replacement,
select_decomp_table,
)
from torch.fx.experimental.proxy_tensor import make_fx
from vllm.kernels.helion.config_manager import ConfigManager
from vllm.kernels.helion.register import HelionKernelWrapper
@contextlib.contextmanager
def _helion_mock_context():
configs = {
"default": helion.Config(block_sizes=[64], num_warps=2, num_stages=2),
}
mock_config_manager = Mock(spec=ConfigManager)
mock_config_manager.get_platform_configs = Mock(return_value=configs)
with (
patch(
"vllm.kernels.helion.config_manager.ConfigManager.get_instance",
return_value=mock_config_manager,
),
patch(
"vllm.kernels.helion.utils.get_canonical_gpu_name",
return_value="nvidia_h200",
),
):
yield
class TestMakeFxHop:
def setup_method(self):
helion_kernel_side_table.reset_table()
def test_make_fx_symbolic(self):
def raw_add_scale(
x: torch.Tensor, y: torch.Tensor, scale: float
) -> tuple[torch.Tensor, int, torch.Tensor]:
out_x = torch.empty_like(x)
out_y = torch.empty_like(x)
for tile in hl.tile(x.size()):
out_x[tile] = x[tile] + y[tile] * scale
out_y[tile] = out_x[tile] * 2.0
return out_x, 42, out_y
input_x = torch.randn(7, 13)
input_y = torch.randn(7, 13)
scale = 0.5
with _helion_mock_context():
wrapper = HelionKernelWrapper(
raw_kernel_func=raw_add_scale,
op_name="test_make_fx",
fake_impl=lambda *a, **kw: None,
)
wrapper.register_config_picker(lambda args, keys: "default")
def fn(x, y):
return wrapper(x, y, scale)
gm = make_fx(fn, tracing_mode="symbolic")(input_x, input_y)
hop_nodes = [
n
for n in gm.graph.nodes
if n.op == "call_function" and n.target is helion_kernel_wrapper_mutation
]
assert len(hop_nodes) == 1
node = hop_nodes[0]
assert node.kwargs["constant_args"]["scale"] == scale
assert set(node.kwargs["tensor_args"]) == {"x", "y"}
specs = node.kwargs["output_spec"]["leaf_specs"]
tensor_specs = [s for s in specs if s["type"] == "tensor"]
scalar_specs = [s for s in specs if s["type"] == "scalar"]
assert len(tensor_specs) == 2
assert len(scalar_specs) == 1
for spec in tensor_specs:
assert spec["dtype"] == input_x.dtype
assert scalar_specs[0]["scalar_value"] == 42
for val in node.meta["val"]:
assert all(isinstance(s, torch.SymInt) for s in val.shape)
# Both out_x and out_y are empty_like(x), so output shapes == input shape
input_node = next(n for n in gm.graph.nodes if n.op == "placeholder")
input_shape = input_node.meta["val"].shape
for val in node.meta["val"]:
assert len(val.shape) == len(input_shape)
for out_s, in_s in zip(val.shape, input_shape):
assert out_s == in_s
def test_pattern_matcher_replaces_with_helion_hop(self):
def raw_silu_mul(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
M, N = x.size()
out = torch.empty_like(x)
for tile_m, tile_n in hl.tile([M, N]):
out[tile_m, tile_n] = (
torch.nn.functional.silu(x[tile_m, tile_n]) * y[tile_m, tile_n]
)
return out
with _helion_mock_context():
wrapper = HelionKernelWrapper(
raw_kernel_func=raw_silu_mul,
op_name="test_pm_silu_mul",
fake_impl=lambda *a, **kw: None,
)
wrapper.register_config_picker(lambda args, keys: "default")
def pattern(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
return torch.nn.functional.silu(x) * y
def replacement(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
return wrapper(x, y)
inputs = [torch.randn(8, 16), torch.randn(8, 16)]
pm_pass = PatternMatcherPass(pass_name="test_helion_replacement")
register_replacement(pattern, replacement, inputs, fwd_only, pm_pass)
def model(x, y):
return torch.nn.functional.silu(x) * y
decompositions = select_decomp_table()
input_x = torch.randn(8, 16)
input_y = torch.randn(8, 16)
gm = make_fx(model, decompositions, tracing_mode="symbolic")(
input_x, input_y
)
def count_hop_nodes(graph):
return sum(
1
for n in graph.nodes
if n.op == "call_function"
and n.target is helion_kernel_wrapper_mutation
)
assert count_hop_nodes(gm.graph) == 0
match_count = pm_pass.apply(gm.graph)
gm.graph.lint()
gm.recompile()
assert match_count == 1
assert count_hop_nodes(gm.graph) == 1
hop_node = next(
n
for n in gm.graph.nodes
if n.op == "call_function"
and n.target is helion_kernel_wrapper_mutation
)
# raw_silu_mul returns empty_like(x), so output shape == input shape
for val in hop_node.meta["val"]:
assert all(isinstance(s, torch.SymInt) for s in val.shape)
input_node = next(n for n in gm.graph.nodes if n.op == "placeholder")
input_shape = input_node.meta["val"].shape
output_shape = hop_node.meta["val"][0].shape
assert len(output_shape) == len(input_shape)
for out_s, in_s in zip(output_shape, input_shape):
assert out_s == in_s

737
third_party/vllm/tests/kernels/helion/test_register.py vendored Normal file
View File

@@ -0,0 +1,737 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Unit tests for Helion kernel registration.
Tests ConfiguredHelionKernel, HelionKernelWrapper, and PresetConfigSearch
including config picker registration and custom autotuner integration.
"""
from unittest.mock import Mock, patch
import pytest
import torch
from vllm.utils.import_utils import has_helion
if not has_helion():
pytest.skip(
"Helion is not installed. Install with: pip install vllm[helion]",
allow_module_level=True,
)
import helion
from vllm.kernels.helion.config_manager import ConfigManager
from vllm.kernels.helion.register import (
_HOP_AVAILABLE,
ConfiguredHelionKernel,
HelionKernelWrapper,
get_kernel_by_name,
get_registered_kernels,
register_kernel,
validate_helion_settings,
)
@pytest.fixture
def sample_configs():
"""Create real Helion config objects for testing."""
return {
"hiddensize_4096_batchsize_32": helion.Config(
block_sizes=[128],
num_warps=4,
num_stages=3,
),
"hiddensize_4096_batchsize_64": helion.Config(
block_sizes=[256],
num_warps=8,
num_stages=4,
),
"hiddensize_4096_batchsize_128": helion.Config(
block_sizes=[512],
num_warps=16,
num_stages=2,
),
"default": helion.Config(
block_sizes=[64],
num_warps=2,
num_stages=2,
),
}
@pytest.fixture
def sample_kernel():
"""Create a simple test kernel function."""
def test_kernel(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
"""Simple test kernel that adds two tensors."""
return x + y
return test_kernel
@pytest.fixture
def config_manager_with_test_configs(sample_configs):
"""Set up ConfigManager with test configs for nvidia_h200 platform."""
mock_config_manager = Mock(spec=ConfigManager)
mock_config_manager.get_platform_configs = Mock(return_value=sample_configs)
return mock_config_manager
@pytest.fixture
def configured_kernel(sample_kernel, sample_configs, config_manager_with_test_configs):
"""Create a ConfiguredHelionKernel for testing."""
def test_config_picker(args, config_keys):
"""Simple config picker that returns default."""
return "default"
with (
patch(
"vllm.kernels.helion.config_manager.ConfigManager.get_instance",
return_value=config_manager_with_test_configs,
),
patch(
"vllm.kernels.helion.utils.get_canonical_gpu_name",
return_value="nvidia_h200",
),
patch("vllm.kernels.helion.register.helion.kernel") as mock_kernel,
):
# Mock just the helion.kernel decorator to avoid actual kernel compilation
mock_decorated = Mock()
mock_kernel.return_value = Mock(return_value=mock_decorated)
return ConfiguredHelionKernel(
op_name="test_kernel",
config_picker=test_config_picker,
raw_kernel_func=sample_kernel,
helion_settings=None,
)
class TestValidateHelionSettings:
"""Test suite for validate_helion_settings utility function."""
def test_accepts_none_settings(self):
"""Test that None settings are accepted without error."""
validate_helion_settings(None, "test_kernel") # Should not raise
def test_accepts_valid_settings(self):
"""Test that valid settings without conflicts are accepted."""
settings = helion.Settings()
settings.static_shapes = False
settings.print_output_code = True
validate_helion_settings(settings, "test_kernel") # Should not raise
def test_rejects_autotuner_fn(self):
"""Test that settings with custom autotuner_fn raise ValueError."""
settings = helion.Settings()
settings.autotuner_fn = lambda *args: None # Set custom autotuner function
with pytest.raises(ValueError, match="uses a custom autotuner"):
validate_helion_settings(settings, "test_kernel")
def test_warns_on_static_shapes_true(self):
"""Test that static_shapes=True emits a warning about being overridden."""
settings = helion.Settings()
settings.static_shapes = True
with patch("vllm.kernels.helion.register.logger") as mock_logger:
validate_helion_settings(settings, "test_kernel")
mock_logger.warning.assert_called_once()
assert "overridden to False" in mock_logger.warning.call_args[0][0]
def create_configured_kernel_with_configs(
op_name,
config_picker,
kernel_func,
configs,
platform="nvidia_h200",
helion_settings=None,
):
"""Helper to create ConfiguredHelionKernel with real config objects."""
mock_config_manager = Mock(spec=ConfigManager)
mock_config_manager.get_platform_configs = Mock(return_value=configs)
with (
patch(
"vllm.kernels.helion.config_manager.ConfigManager.get_instance",
return_value=mock_config_manager,
),
patch(
"vllm.kernels.helion.utils.get_canonical_gpu_name",
return_value=platform,
),
patch("vllm.kernels.helion.register.helion.kernel") as mock_kernel,
):
mock_decorated = Mock()
mock_kernel.return_value = Mock(return_value=mock_decorated)
return ConfiguredHelionKernel(
op_name=op_name,
config_picker=config_picker,
raw_kernel_func=kernel_func,
helion_settings=helion_settings,
)
class TestConfiguredHelionKernel:
"""Test suite for ConfiguredHelionKernel."""
def test_init_raises_without_picker(self, sample_kernel, sample_configs):
"""Test that __init__ raises when no picker registered."""
configs = {"default": sample_configs["default"]}
mock_config_manager = Mock(spec=ConfigManager)
mock_config_manager.get_platform_configs = Mock(return_value=configs)
with (
patch(
"vllm.kernels.helion.config_manager.ConfigManager.get_instance",
return_value=mock_config_manager,
),
patch(
"vllm.kernels.helion.utils.get_canonical_gpu_name",
return_value="nvidia_h200",
),
pytest.raises(RuntimeError, match="No config picker registered"),
):
ConfiguredHelionKernel(
op_name="test_kernel",
config_picker=None, # No picker registered
raw_kernel_func=sample_kernel,
helion_settings=None,
)
def test_config_selector_validates_picker_result(
self, sample_kernel, sample_configs
):
"""Test that config selector validates picker returns valid key."""
def invalid_picker(args, config_keys):
return "invalid_key"
kernel = create_configured_kernel_with_configs(
op_name="test_kernel",
config_picker=invalid_picker,
kernel_func=sample_kernel,
configs=sample_configs,
)
key_computer = kernel._create_key_computer()
selector = kernel._create_config_selector(key_computer)
with pytest.raises(
ValueError, match="Config picker returned invalid config key"
):
selector((torch.randn(32, 4096),))
def test_config_selector_handles_none_from_picker(
self, sample_kernel, sample_configs
):
"""Test that config selector falls back to 'default' on None."""
def none_picker(args, config_keys):
return None
kernel = create_configured_kernel_with_configs(
op_name="test_kernel",
config_picker=none_picker,
kernel_func=sample_kernel,
configs=sample_configs,
)
key_computer = kernel._create_key_computer()
selector = kernel._create_config_selector(key_computer)
result = selector((torch.randn(32, 4096),))
assert result is kernel.configs["default"]
def test_create_decorated_kernel_passes_helion_settings(
self, sample_kernel, sample_configs
):
"""Test that _create_decorated_kernel passes helion_settings."""
def default_picker(args, config_keys):
return "default"
settings = helion.Settings()
settings.print_output_code = True
mock_config_manager = Mock(spec=ConfigManager)
mock_config_manager.get_platform_configs = Mock(return_value=sample_configs)
with (
patch("vllm.kernels.helion.register.helion.kernel") as mock_kernel,
patch(
"vllm.kernels.helion.config_manager.ConfigManager.get_instance",
return_value=mock_config_manager,
),
patch(
"vllm.kernels.helion.utils.get_canonical_gpu_name",
return_value="nvidia_h200",
),
):
mock_decorated = Mock()
mock_kernel.return_value = Mock(return_value=mock_decorated)
ConfiguredHelionKernel(
op_name="test_kernel",
config_picker=default_picker,
raw_kernel_func=sample_kernel,
helion_settings=settings,
)
call_kwargs = mock_kernel.call_args[1]
assert "print_output_code" in call_kwargs
assert call_kwargs["print_output_code"] is True
# static_shapes is always forced to False by vLLM
assert call_kwargs["static_shapes"] is False
def test_key_and_config_selector_use_same_logic(
self, sample_kernel, sample_configs
):
"""Test that key and config_selector produce identical results."""
def tracking_picker(args, config_keys):
x = args[0]
batch_size = x.shape[0]
if batch_size <= 32:
return "hiddensize_4096_batchsize_32"
elif batch_size <= 64:
return "hiddensize_4096_batchsize_64"
return "hiddensize_4096_batchsize_128"
mock_config_manager = Mock(spec=ConfigManager)
mock_config_manager.get_platform_configs = Mock(return_value=sample_configs)
with (
patch("vllm.kernels.helion.register.helion.kernel") as mock_helion_kernel,
patch(
"vllm.kernels.helion.config_manager.ConfigManager.get_instance",
return_value=mock_config_manager,
),
patch(
"vllm.kernels.helion.utils.get_canonical_gpu_name",
return_value="nvidia_h200",
),
):
mock_decorated = Mock()
mock_helion_kernel.return_value = Mock(return_value=mock_decorated)
kernel = ConfiguredHelionKernel(
op_name="test_kernel",
config_picker=tracking_picker,
raw_kernel_func=sample_kernel,
helion_settings=None,
)
call_kwargs = mock_helion_kernel.call_args[1]
key_fn = call_kwargs["key"]
autotuner_fn = call_kwargs["autotuner_fn"]
tensor = torch.randn(50, 4096) # batch=50, should select batchsize_64
# key receives unpacked args, autotuner receives args as tuple
key_result = key_fn(tensor)
autotuner = autotuner_fn(None, (tensor,))
config = autotuner.autotune()
assert key_result == "hiddensize_4096_batchsize_64"
assert config is kernel.configs["hiddensize_4096_batchsize_64"]
class TestHelionKernelWrapper:
"""Test suite for HelionKernelWrapper."""
def test_get_configured_op_validates_configs_available(self, sample_kernel):
"""Test get_configured_op validates configs are available."""
def fake_impl(*args, **kwargs):
return torch.zeros_like(args[0])
wrapper = HelionKernelWrapper(
raw_kernel_func=sample_kernel,
op_name="test_kernel",
fake_impl=fake_impl,
)
def default_picker(args, config_keys):
return "default"
wrapper._config_picker = default_picker
mock_config_manager = Mock(spec=ConfigManager)
mock_config_manager.get_platform_configs = Mock(
return_value={}
) # Empty configs
with (
patch(
"vllm.kernels.helion.config_manager.ConfigManager.get_instance",
return_value=mock_config_manager,
),
patch(
"vllm.kernels.helion.utils.get_canonical_gpu_name",
return_value="nvidia_h200",
),
pytest.raises(ValueError, match="No configs available"),
):
wrapper.get_configured_op()
def test_get_configured_op_validates_config_picker(
self, sample_kernel, sample_configs
):
"""Test get_configured_op validates config picker."""
def fake_impl(*args, **kwargs):
return torch.zeros_like(args[0])
wrapper = HelionKernelWrapper(
raw_kernel_func=sample_kernel,
op_name="test_kernel",
fake_impl=fake_impl,
)
# Don't set config picker - should raise assertion error
mock_config_manager = Mock(spec=ConfigManager)
mock_config_manager.get_platform_configs = Mock(return_value=sample_configs)
with (
patch(
"vllm.kernels.helion.config_manager.ConfigManager.get_instance",
return_value=mock_config_manager,
),
patch(
"vllm.kernels.helion.utils.get_canonical_gpu_name",
return_value="nvidia_h200",
),
pytest.raises(AssertionError, match="No config picker registered"),
):
wrapper.get_configured_op()
def test_get_configured_op_returns_cached_kernel(
self, sample_kernel, sample_configs
):
"""Test get_configured_op returns cached ConfiguredHelionKernel."""
def fake_impl(*args, **kwargs):
return torch.zeros_like(args[0])
def default_picker(args, config_keys):
return "default"
wrapper = HelionKernelWrapper(
raw_kernel_func=sample_kernel,
op_name="test_kernel",
fake_impl=fake_impl,
)
wrapper._config_picker = default_picker
mock_config_manager = Mock(spec=ConfigManager)
mock_config_manager.get_platform_configs = Mock(return_value=sample_configs)
with (
patch(
"vllm.kernels.helion.config_manager.ConfigManager.get_instance",
return_value=mock_config_manager,
),
patch(
"vllm.kernels.helion.utils.get_canonical_gpu_name",
return_value="nvidia_h200",
),
patch("vllm.kernels.helion.register.helion.kernel") as mock_kernel,
):
mock_decorated = Mock()
mock_kernel.return_value = Mock(return_value=mock_decorated)
result1 = wrapper.get_configured_op()
result2 = wrapper.get_configured_op()
assert result1 is result2
@pytest.mark.skipif(
_HOP_AVAILABLE, reason="CustomOp path not used when HOP available"
)
def test_get_or_register_custom_op_returns_cached_op(
self, sample_kernel, sample_configs
):
def fake_impl(*args, **kwargs):
return torch.zeros_like(args[0])
def default_picker(args, config_keys):
return "default"
wrapper = HelionKernelWrapper(
raw_kernel_func=sample_kernel,
op_name="test_kernel",
fake_impl=fake_impl,
)
wrapper._config_picker = default_picker
mock_config_manager = Mock(spec=ConfigManager)
mock_config_manager.get_platform_configs = Mock(return_value=sample_configs)
existing_op = Mock()
mock_namespace = Mock()
mock_namespace.test_kernel = existing_op
with (
patch(
"vllm.kernels.helion.config_manager.ConfigManager.get_instance",
return_value=mock_config_manager,
),
patch(
"vllm.kernels.helion.utils.get_canonical_gpu_name",
return_value="nvidia_h200",
),
patch.object(torch.ops, "vllm_helion", mock_namespace),
patch("vllm.kernels.helion.register.helion.kernel") as mock_kernel,
):
mock_decorated = Mock()
mock_kernel.return_value = Mock(return_value=mock_decorated)
result = wrapper._get_or_register_custom_op()
assert result is existing_op
@pytest.mark.skipif(
_HOP_AVAILABLE, reason="CustomOp path not used when HOP available"
)
def test_get_or_register_custom_op_registers_new_op(
self, sample_kernel, sample_configs
):
def fake_impl(*args, **kwargs):
return torch.zeros_like(args[0])
def default_picker(args, config_keys):
return "default"
wrapper = HelionKernelWrapper(
raw_kernel_func=sample_kernel,
op_name="test_kernel",
fake_impl=fake_impl,
)
wrapper._config_picker = default_picker
mock_config_manager = Mock(spec=ConfigManager)
mock_config_manager.get_platform_configs = Mock(return_value=sample_configs)
new_op = Mock()
registered_ops: dict[str, Mock] = {}
class MockNamespace:
def __getattr__(self, name):
if name in registered_ops:
return registered_ops[name]
raise AttributeError(name)
mock_namespace = MockNamespace()
def register_side_effect(op_name, op_func, **kwargs):
registered_ops[op_name] = new_op
with (
patch(
"vllm.kernels.helion.config_manager.ConfigManager.get_instance",
return_value=mock_config_manager,
),
patch(
"vllm.kernels.helion.utils.get_canonical_gpu_name",
return_value="nvidia_h200",
),
patch.object(torch.ops, "vllm_helion", mock_namespace),
patch(
"vllm.kernels.helion.register.direct_register_custom_op",
side_effect=register_side_effect,
) as mock_register,
patch("vllm.kernels.helion.register.helion.kernel") as mock_kernel,
):
mock_decorated = Mock()
mock_kernel.return_value = Mock(return_value=mock_decorated)
result = wrapper._get_or_register_custom_op()
mock_register.assert_called_once()
assert result is new_op
assert mock_register.call_args[1]["op_func"] is mock_decorated
class TestKernelRegistry:
"""Test suite for kernel registry functionality."""
def setup_method(self):
"""Save and clear the registry before each test."""
from vllm.kernels.helion.register import _REGISTERED_KERNELS
self._saved_registry = dict(_REGISTERED_KERNELS)
_REGISTERED_KERNELS.clear()
def teardown_method(self):
"""Restore the registry after each test."""
from vllm.kernels.helion.register import _REGISTERED_KERNELS
_REGISTERED_KERNELS.clear()
_REGISTERED_KERNELS.update(self._saved_registry)
def test_get_registered_kernels_returns_copy(self):
"""Test get_registered_kernels returns copy of registry."""
result1 = get_registered_kernels()
result2 = get_registered_kernels()
# Should be separate objects
assert result1 is not result2
# Should have same content
assert result1 == result2
def test_get_kernel_by_name_returns_kernel(self):
"""Test get_kernel_by_name returns registered kernel."""
wrapper = HelionKernelWrapper(
raw_kernel_func=Mock(),
op_name="test_kernel",
fake_impl=Mock(),
)
from vllm.kernels.helion.register import _REGISTERED_KERNELS
_REGISTERED_KERNELS["test_kernel"] = wrapper
result = get_kernel_by_name("test_kernel")
assert result is wrapper
def test_get_kernel_by_name_returns_none_for_missing(self):
"""Test get_kernel_by_name returns None for missing kernel."""
result = get_kernel_by_name("nonexistent")
assert result is None
def test_register_kernel_auto_generates_fake_impl(self):
"""Test register_kernel auto-generates fake_impl when not provided."""
with patch("vllm.kernels.helion.register.infer_fake_impl") as mock_infer:
mock_fake = Mock()
mock_infer.return_value = mock_fake
def original_kernel(x):
return x
wrapper = register_kernel(original_kernel)
mock_infer.assert_called_once_with(original_kernel, None)
assert wrapper._fake_impl is mock_fake
def test_register_kernel_creates_wrapper(self):
"""Test register_kernel creates HelionKernelWrapper."""
def test_kernel(x):
return x
result = register_kernel("test_name")(test_kernel)
assert isinstance(result, HelionKernelWrapper)
assert result.op_name == "test_name"
assert result.raw_kernel_func is test_kernel
def test_register_kernel_auto_detects_name(self):
"""Test register_kernel uses function name when no name provided."""
@register_kernel
def my_test_kernel(x):
return x
assert my_test_kernel.op_name == "my_test_kernel"
def test_register_kernel_registers_in_global_registry(self):
"""Test register_kernel adds wrapper to global registry."""
@register_kernel
def test_kernel(x):
return x
registered_kernels = get_registered_kernels()
assert "test_kernel" in registered_kernels
assert registered_kernels["test_kernel"] is test_kernel
def test_register_kernel_passes_helion_settings(self):
"""Test register_kernel passes helion_settings to wrapper."""
mock_settings = Mock()
mock_settings.to_dict.return_value = {"debug": True}
@register_kernel("test_name", helion_settings=mock_settings)
def test_kernel(x):
return x
assert test_kernel.helion_settings is mock_settings
def test_register_kernel_supports_decorator_syntax(self):
"""Test register_kernel works with decorator arguments."""
mock_fake = Mock()
wrapper = register_kernel("custom_name", fake_impl=mock_fake)
def test_kernel(x):
return x
result = wrapper(test_kernel)
assert result.op_name == "custom_name"
assert result._fake_impl is mock_fake
def test_register_kernel_bare_decorator(self):
"""Test register_kernel works as bare decorator."""
@register_kernel
def test_kernel(x):
return x
assert isinstance(test_kernel, HelionKernelWrapper)
assert test_kernel.op_name == "test_kernel"
def test_registered_wrapper_can_register_config_picker(self):
"""Test that registered wrapper can register config picker."""
@register_kernel
def test_kernel(x):
return x
def my_picker(args, config_keys):
return "default"
result = test_kernel.register_config_picker(my_picker)
assert result is my_picker
assert test_kernel._config_picker is my_picker
def test_register_kernel_raises_on_duplicate_registration(self):
"""Test register_kernel raises error on duplicate names."""
@register_kernel("duplicate_name")
def kernel1(x):
return x
with pytest.raises(ValueError, match="already registered"):
@register_kernel("duplicate_name")
def kernel2(x):
return x
def test_register_kernel_rejects_autotuner_fn_in_settings(self):
"""Test register_kernel rejects conflicting autotuner_fn."""
mock_settings = Mock()
mock_settings.to_dict.return_value = {"autotuner_fn": Mock()}
with pytest.raises(ValueError, match="uses a custom autotuner"):
@register_kernel("test", helion_settings=mock_settings)
def test_kernel(x):
return x
def test_register_kernel_no_warning_with_static_shapes_false(self):
"""Test register_kernel doesn't warn with static_shapes=False."""
mock_settings = Mock()
mock_settings.to_dict.return_value = {"static_shapes": False}
with patch("vllm.kernels.helion.register.logger") as mock_logger:
@register_kernel("test", helion_settings=mock_settings)
def test_kernel(x):
return x
# Should not call warning
mock_logger.warning.assert_not_called()

395
third_party/vllm/tests/kernels/helion/test_silu_mul_fp8.py vendored Normal file
View File

@@ -0,0 +1,395 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
import torch.nn.functional as F
from vllm.utils.import_utils import has_helion
if not has_helion():
pytest.skip(
"Helion is not installed. Install with: pip install vllm[helion]",
allow_module_level=True,
)
from vllm.kernels.helion.config_manager import ConfigManager
from vllm.kernels.helion.ops.silu_mul_fp8 import (
pick_silu_mul_fp8_config,
silu_mul_fp8,
silu_mul_fp8_baseline,
)
def skip_if_platform_unsupported():
try:
from vllm.kernels.helion.utils import get_canonical_gpu_name
if not torch.cuda.is_available():
pytest.skip("CUDA not available")
platform = get_canonical_gpu_name()
try:
config_manager = ConfigManager.get_instance()
except RuntimeError:
config_manager = ConfigManager()
configs = config_manager.get_platform_configs("silu_mul_fp8", platform)
if len(configs) == 0:
pytest.skip("Current GPU platform not supported for silu_mul_fp8 kernel")
except (ImportError, RuntimeError, KeyError):
pytest.skip("Error detecting platform support for silu_mul_fp8 kernel")
@pytest.fixture(autouse=True)
def reset_config_manager_singleton():
ConfigManager.reset_instance()
ConfigManager()
yield
ConfigManager.reset_instance()
class TestSiluMulFp8ConfigPicker:
def test_config_picker_exact_match(self):
config_keys = [
"intermediate_2048_numtokens_256",
"intermediate_4096_numtokens_256",
]
input_tensor = torch.randn(32, 4096, dtype=torch.bfloat16, device="cuda")
scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
args = (input_tensor, scale)
selected_key = pick_silu_mul_fp8_config(args, config_keys)
assert selected_key == "intermediate_2048_numtokens_256"
def test_config_picker_closest_match(self):
config_keys = [
"intermediate_2048_numtokens_256",
"intermediate_4096_numtokens_256",
]
# Use 7000 (intermediate_size=3500) which is closer to 4096 than 2048
input_tensor = torch.randn(32, 7000, dtype=torch.bfloat16, device="cuda")
scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
args = (input_tensor, scale)
selected_key = pick_silu_mul_fp8_config(args, config_keys)
assert selected_key == "intermediate_4096_numtokens_256"
def test_config_picker_fallback_to_default(self):
config_keys = ["default"]
input_tensor = torch.randn(32, 4096, dtype=torch.bfloat16, device="cuda")
scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
args = (input_tensor, scale)
selected_key = pick_silu_mul_fp8_config(args, config_keys)
assert selected_key == "default"
def test_config_picker_no_configs(self):
config_keys: list[str] = []
input_tensor = torch.randn(32, 4096, dtype=torch.bfloat16, device="cuda")
scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
args = (input_tensor, scale)
selected_key = pick_silu_mul_fp8_config(args, config_keys)
assert selected_key is None
@pytest.mark.parametrize("intermediate_size", [2048, 4096, 5120])
def test_config_picker_different_sizes(self, intermediate_size):
config_keys = [
"intermediate_2048_numtokens_256",
"intermediate_4096_numtokens_256",
"intermediate_5120_numtokens_256",
]
input_tensor = torch.randn(
32, 2 * intermediate_size, dtype=torch.bfloat16, device="cuda"
)
scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
args = (input_tensor, scale)
selected_key = pick_silu_mul_fp8_config(args, config_keys)
expected_key = f"intermediate_{intermediate_size}_numtokens_256"
assert selected_key == expected_key
def test_config_picker_numtokens_ceiling(self):
"""Pick the smallest numtokens >= input num_tokens."""
config_keys = [
"intermediate_4096_numtokens_8",
"intermediate_4096_numtokens_32",
"intermediate_4096_numtokens_128",
"intermediate_4096_numtokens_256",
]
# 20 tokens -> should pick numtokens_32 (smallest >= 20)
input_tensor = torch.randn(20, 8192, dtype=torch.bfloat16, device="cuda")
scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
selected_key = pick_silu_mul_fp8_config((input_tensor, scale), config_keys)
assert selected_key == "intermediate_4096_numtokens_32"
def test_config_picker_numtokens_exact(self):
"""Exact num_tokens match is preferred over ceiling."""
config_keys = [
"intermediate_4096_numtokens_8",
"intermediate_4096_numtokens_32",
"intermediate_4096_numtokens_128",
]
input_tensor = torch.randn(32, 8192, dtype=torch.bfloat16, device="cuda")
scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
selected_key = pick_silu_mul_fp8_config((input_tensor, scale), config_keys)
assert selected_key == "intermediate_4096_numtokens_32"
def test_config_picker_numtokens_fallback_to_largest(self):
"""Fall back to the largest numtokens when input exceeds all."""
config_keys = [
"intermediate_4096_numtokens_8",
"intermediate_4096_numtokens_32",
"intermediate_4096_numtokens_128",
]
# 512 tokens -> exceeds all available, should pick largest (128)
input_tensor = torch.randn(512, 8192, dtype=torch.bfloat16, device="cuda")
scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
selected_key = pick_silu_mul_fp8_config((input_tensor, scale), config_keys)
assert selected_key == "intermediate_4096_numtokens_128"
def test_config_picker_malformed_key_raises(self):
"""Malformed config keys should raise ValueError."""
config_keys = ["intermediate_4096_badformat_256"]
input_tensor = torch.randn(32, 8192, dtype=torch.bfloat16, device="cuda")
scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
with pytest.raises(ValueError, match="Malformed config key"):
pick_silu_mul_fp8_config((input_tensor, scale), config_keys)
def test_config_picker_default_ignored_when_valid_keys_exist(self):
"""'default' is skipped in favor of a real match."""
config_keys = [
"default",
"intermediate_4096_numtokens_32",
"intermediate_4096_numtokens_128",
]
input_tensor = torch.randn(64, 8192, dtype=torch.bfloat16, device="cuda")
scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
selected_key = pick_silu_mul_fp8_config((input_tensor, scale), config_keys)
assert selected_key == "intermediate_4096_numtokens_128"
class TestSiluMulFp8Correctness:
@pytest.mark.parametrize("batch_size", [1, 8, 32, 128])
@pytest.mark.parametrize("intermediate_size", [2048, 3000, 3500, 4096, 5000])
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
def test_silu_mul_fp8_correctness(self, batch_size, intermediate_size, dtype):
skip_if_platform_unsupported()
input_size = 2 * intermediate_size
input_tensor = torch.randn(batch_size, input_size, dtype=dtype, device="cuda")
scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
reference_output = silu_mul_fp8_baseline(input_tensor, scale)
helion_output = silu_mul_fp8(input_tensor, scale)
assert helion_output.shape == reference_output.shape
assert helion_output.dtype == torch.float8_e4m3fn
assert reference_output.dtype == torch.float8_e4m3fn
ref_f32 = reference_output.to(torch.float32)
helion_f32 = helion_output.to(torch.float32)
# FP8 E4M3 has limited precision. Values near quantization boundaries
# can round differently due to intermediate precision differences.
torch.testing.assert_close(
helion_f32,
ref_f32,
atol=0.05,
rtol=0.05,
msg=f"Mismatch at batch={batch_size}, size={intermediate_size}",
)
def test_silu_mul_fp8_shape_inference(self):
skip_if_platform_unsupported()
batch_size, input_size = 32, 8192
intermediate_size = input_size // 2
input_tensor = torch.randn(
batch_size, input_size, dtype=torch.bfloat16, device="cuda"
)
scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
output = silu_mul_fp8(input_tensor, scale)
expected_shape = (batch_size, intermediate_size)
assert output.shape == expected_shape
assert output.dtype == torch.float8_e4m3fn
def test_silu_mul_fp8_scale_variations(self):
skip_if_platform_unsupported()
batch_size, input_size = 16, 4096
input_tensor = torch.randn(
batch_size, input_size, dtype=torch.bfloat16, device="cuda"
)
scales = [0.1, 0.5, 1.0, 2.0, 10.0]
for scale_val in scales:
scale = torch.tensor([scale_val], dtype=torch.float32, device="cuda")
reference_output = silu_mul_fp8_baseline(input_tensor, scale)
helion_output = silu_mul_fp8(input_tensor, scale)
ref_f32 = reference_output.to(torch.float32)
helion_f32 = helion_output.to(torch.float32)
torch.testing.assert_close(
helion_f32,
ref_f32,
atol=0.05,
rtol=0.05,
msg=f"Mismatch for scale={scale_val}",
)
@pytest.mark.parametrize(
"shape",
[
(1, 4096),
(16, 4096),
(128, 4096),
(1024, 4096),
(1, 8192),
(16, 8192),
(128, 8192),
],
)
def test_silu_mul_fp8_various_shapes(self, shape):
skip_if_platform_unsupported()
input_tensor = torch.randn(*shape, dtype=torch.bfloat16, device="cuda")
scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
reference_output = silu_mul_fp8_baseline(input_tensor, scale)
helion_output = silu_mul_fp8(input_tensor, scale)
assert helion_output.shape == reference_output.shape
ref_f32 = reference_output.to(torch.float32)
helion_f32 = helion_output.to(torch.float32)
torch.testing.assert_close(
helion_f32, ref_f32, atol=0.05, rtol=0.05, msg=f"Mismatch for shape={shape}"
)
def silu_mul_fp8_pytorch(input: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
"""Pure PyTorch reference using F.silu.
This matches vLLM's SiluAndMul.forward_native exactly:
F.silu(x[..., :d]) * x[..., d:]
"""
d = input.shape[-1] // 2
result = F.silu(input[..., :d]) * input[..., d:]
return (result.to(torch.float32) / scale).to(torch.float8_e4m3fn)
class TestSiluMulFp8PytorchReference:
"""Tests comparing Helion kernel against pure PyTorch implementation.
Uses tighter tolerance since both use PyTorch's FP8 conversion
(same rounding mode), unlike the vLLM C++ baseline which uses
NVIDIA's hardware FP8 conversion with different rounding.
"""
@pytest.mark.parametrize("batch_size", [1, 8, 32, 128, 256])
@pytest.mark.parametrize("intermediate_size", [1024, 2048, 4096])
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
def test_silu_mul_fp8_vs_pytorch(self, batch_size, intermediate_size, dtype):
skip_if_platform_unsupported()
input_tensor = torch.randn(
batch_size, 2 * intermediate_size, dtype=dtype, device="cuda"
)
scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
pytorch_output = silu_mul_fp8_pytorch(input_tensor, scale)
helion_output = silu_mul_fp8(input_tensor, scale)
assert helion_output.shape == pytorch_output.shape
assert helion_output.dtype == torch.float8_e4m3fn
pytorch_f32 = pytorch_output.to(torch.float32)
helion_f32 = helion_output.to(torch.float32)
# Tolerance accounts for FP8 quantization boundary effects
torch.testing.assert_close(
helion_f32,
pytorch_f32,
atol=0.05,
rtol=0.05,
msg=(
f"Mismatch at batch={batch_size}, size={intermediate_size}, "
f"dtype={dtype}"
),
)
@pytest.mark.parametrize(
"shape",
[
(1, 2, 4096), # 3D input
(2, 4, 2048), # 3D input
(1, 1, 1, 8192), # 4D input
],
)
def test_silu_mul_fp8_multidim_vs_pytorch(self, shape):
skip_if_platform_unsupported()
input_tensor = torch.randn(*shape, dtype=torch.bfloat16, device="cuda")
scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
pytorch_output = silu_mul_fp8_pytorch(input_tensor, scale)
helion_output = silu_mul_fp8(input_tensor, scale)
assert helion_output.shape == pytorch_output.shape
pytorch_f32 = pytorch_output.to(torch.float32)
helion_f32 = helion_output.to(torch.float32)
torch.testing.assert_close(
helion_f32,
pytorch_f32,
atol=0.05,
rtol=0.05,
msg=f"Mismatch for shape={shape}",
)
class TestSiluMulFp8Integration:
def test_kernel_registration_integration(self):
from vllm.kernels.helion.register import get_registered_kernels
registered_kernels = get_registered_kernels()
assert "silu_mul_fp8" in registered_kernels
kernel_wrapper = registered_kernels["silu_mul_fp8"]
assert kernel_wrapper.op_name == "silu_mul_fp8"
assert kernel_wrapper._config_picker is not None
def test_fake_impl_functionality(self):
skip_if_platform_unsupported()
from vllm.kernels.helion.register import get_registered_kernels
input_tensor = torch.randn(32, 4096, dtype=torch.bfloat16, device="cuda")
scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
registered_kernels = get_registered_kernels()
kernel_wrapper = registered_kernels["silu_mul_fp8"]
fake_impl = kernel_wrapper._fake_impl
fake_output = fake_impl(input_tensor, scale)
expected_shape = (32, 2048)
assert fake_output.shape == expected_shape
assert fake_output.dtype == torch.float8_e4m3fn
assert fake_output.device == input_tensor.device

32
third_party/vllm/tests/kernels/helion/test_utils.py vendored Normal file
View File

@@ -0,0 +1,32 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for Helion utility functions."""
import pytest
from vllm.kernels.helion.utils import canonicalize_gpu_name
@pytest.mark.parametrize(
"driver_reported_name,expected",
[
("NVIDIA H200", "nvidia_h200"),
("NVIDIA A100-SXM4-80GB", "nvidia_a100"),
("NVIDIA H100 80GB HBM3", "nvidia_h100"),
("NVIDIA H100 PCIe", "nvidia_h100"),
("NVIDIA H100 SXM5", "nvidia_h100"),
("NVIDIA GeForce RTX 4090", "nvidia_geforce_rtx_4090"),
("AMD Instinct MI300X", "amd_instinct_mi300x"),
("Tesla V100-SXM2-32GB", "tesla_v100"),
],
)
def test_canonicalize_gpu_name(driver_reported_name, expected):
"""Test GPU name canonicalization."""
assert canonicalize_gpu_name(driver_reported_name) == expected
@pytest.mark.parametrize("invalid_name", ["", " ", "\t", "\n"])
def test_canonicalize_gpu_name_rejects_empty(invalid_name):
"""Test that empty or whitespace-only names are rejected."""
with pytest.raises(ValueError, match="cannot be empty"):
canonicalize_gpu_name(invalid_name)