[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>
This commit is contained in:
@@ -52,7 +52,7 @@ def _helion_mock_context():
|
||||
|
||||
with (
|
||||
patch(
|
||||
"vllm.kernels.helion.config_manager.ConfigManager.get_instance",
|
||||
"vllm.kernels.helion.config_manager.ConfigManager",
|
||||
return_value=mock_config_manager,
|
||||
),
|
||||
patch(
|
||||
@@ -87,8 +87,8 @@ class TestMakeFxHop:
|
||||
raw_kernel_func=raw_add_scale,
|
||||
op_name="test_make_fx",
|
||||
fake_impl=lambda *a, **kw: None,
|
||||
config_picker=lambda args, keys: "default",
|
||||
)
|
||||
wrapper.register_config_picker(lambda args, keys: "default")
|
||||
|
||||
def fn(x, y):
|
||||
return wrapper(x, y, scale)
|
||||
@@ -143,8 +143,8 @@ class TestMakeFxHop:
|
||||
raw_kernel_func=raw_silu_mul,
|
||||
op_name="test_pm_silu_mul",
|
||||
fake_impl=lambda *a, **kw: None,
|
||||
config_picker=lambda args, keys: "default",
|
||||
)
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user