[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>

tests/kernels/helion/helpers.py

@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import json
+import tempfile
+from collections.abc import Callable
+from contextlib import contextmanager
+from pathlib import Path
+from unittest.mock import patch
+
+import helion
+
+from vllm.kernels.helion.config_manager import ConfigManager
+from vllm.kernels.helion.register import register_kernel
+from vllm.kernels.helion.utils import get_canonical_gpu_name
+
+GPU_PLATFORM = get_canonical_gpu_name()
+
+DEFAULT_CONFIGS: dict[str, helion.Config] = {
+    "default": helion.Config(block_sizes=[32]),
+}
+
+
+@contextmanager
+def dummy_kernel_registry(
+    configs: dict[str, helion.Config] | None = None,
+):
+    """Context manager providing a register function with automatic config setup.
+
+    Yields a ``register`` callable with the same signature as
+    ``register_kernel``.  Before applying the real decorator it writes a
+    config JSON for the kernel name (from ``op_name`` or ``fn.__name__``)
+    into a temporary directory backed by a fresh ``ConfigManager``.
+    """
+    if configs is None:
+        configs = DEFAULT_CONFIGS
+    config_data = {k: v.__dict__["config"] for k, v in configs.items()}
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        config_dir = Path(tmpdir)
+        ConfigManager.reset_instance()
+        cm = ConfigManager(base_dir=config_dir)
+
+        with patch(
+            "vllm.kernels.helion.config_manager.ConfigManager",
+            return_value=cm,
+        ):
+
+            def register(
+                op_name: str | None = None,
+                **kwargs,
+            ) -> Callable:
+                def decorator(fn: Callable) -> Callable:
+                    name = op_name or fn.__name__
+                    kernel_dir = config_dir / name
+                    kernel_dir.mkdir(parents=True, exist_ok=True)
+                    (kernel_dir / f"{GPU_PLATFORM}.json").write_text(
+                        json.dumps(config_data)
+                    )
+                    return register_kernel(op_name, **kwargs)(fn)
+
+                return decorator
+
+            try:
+                yield register
+            finally:
+                ConfigManager.reset_instance()