feat: lazy TP weight loading via SafetensorsSlice

When tensor parallelism is active, each rank previously read the full
weight tensor from disk and then called .narrow() to extract its shard,
discarding the rest. This change introduces SafetensorsSlice, a lazy
wrapper around safetensors PySafeSlice that defers disk I/O until
materialization. The .narrow() calls from weight loaders are recorded
without touching the disk, and the actual read at materialize() time
fetches only the needed sub-region.

The __torch_function__ protocol auto-materializes when the slice is
used in any torch operation (e.g. param_data.copy_()), so no downstream
weight loader code needs changes. Unsupported operations (.reshape(),
.view(), .t(), etc.) also auto-materialize as a safety fallback.

Only affects the default "lazy" safetensors loading strategy.

Co-authored-by: Claude

https://claude.ai/code/session_01Ngt6Nm9BtPKAEied3djkuE

.buildkite/scripts/hardware_ci/run-xpu-test.sh

@@ -40,16 +40,16 @@ docker run \
     python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend mp
     python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --attention-backend=TRITON_ATTN
     python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --quantization fp8
-    python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4  --block-size 64 --enforce-eager --max-model-len 8192
+    python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4  --block-size 64 --enforce-eager
     python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b  --block-size 64 --enforce-eager -tp 2
     python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b  --block-size 64 --enforce-eager -tp 2 --enable-expert-parallel
     cd tests
     pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py
     pytest -v -s v1/engine
     pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py
-    pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py
+    pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py
     pytest -v -s v1/structured_output
     pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_tree_attention.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py
-    pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_nixl_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py -k "not (test_register_kv_caches and FLASH_ATTN and True)"
+    pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_nixl_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py
     pytest -v -s v1/test_serial_utils.py
 '

tests/entrypoints/openai/chat_completion/test_audio_in_video.py

@@ -9,7 +9,7 @@ import pytest
 import pytest_asyncio
 
 from tests.conftest import VideoTestAssets
-from tests.utils import ROCM_EXTRA_ARGS, RemoteOpenAIServer
+from tests.utils import RemoteOpenAIServer
 
 MODEL_NAME = "Qwen/Qwen2.5-Omni-3B"
 
@@ -22,7 +22,6 @@ def server():
         "--enforce-eager",
         "--limit-mm-per-prompt",
         json.dumps({"audio": 3, "video": 3}),
-        *ROCM_EXTRA_ARGS,
     ]
 
     with RemoteOpenAIServer(

tests/kernels/attention/test_trtllm_kvfp8_dequant.py

@@ -12,12 +12,6 @@ import torch
 
 from vllm.platforms import current_platform
 
-if current_platform.is_rocm():
-    pytest.skip(
-        "trtllm kvfp8 dequant is not supported on ROCm.",
-        allow_module_level=True,
-    )
-
 FP8_DTYPE = current_platform.fp8_dtype()
 
 NUM_BLOCKS = 128

tests/kernels/helion/helpers.py

@@ -1,67 +0,0 @@
-# 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()

tests/kernels/helion/test_autotune.py

@@ -1,91 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for autotuning Helion kernels, including disabled kernels with no configs."""
-
-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.autotuner.base_search import BaseSearch
-
-from tests.kernels.helion.helpers import dummy_kernel_registry
-from vllm.kernels.helion.register import create_helion_decorated_kernel
-
-
-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
-
-
-class NoCompileSearch(BaseSearch):
-    """Autotuner that returns the default config without GPU compilation.
-
-    Modeled after helion's test BasicSearch (pytorch/helion#1649).
-    """
-
-    def autotune(self, *, skip_cache: bool = False):
-        return self.config_spec.default_config()
-
-
-def _no_compile_autotuner_fn(bound_kernel, args, **kwargs):
-    return NoCompileSearch(bound_kernel, args, **kwargs)
-
-
-class TestAutotuneDisabledKernel:
-    """Test autotuning flow on disabled kernels (no platform configs)."""
-
-    def setup_method(self):
-        from vllm.kernels.helion.register import _REGISTERED_KERNELS
-
-        self._saved_registry = dict(_REGISTERED_KERNELS)
-        _REGISTERED_KERNELS.clear()
-
-    def teardown_method(self):
-        from vllm.kernels.helion.register import _REGISTERED_KERNELS
-
-        _REGISTERED_KERNELS.clear()
-        _REGISTERED_KERNELS.update(self._saved_registry)
-
-    def test_autotune_disabled_kernel_produces_valid_config(self):
-        """Register a kernel with no configs (disabled), run autotune,
-        verify it produces a valid helion.Config."""
-        with dummy_kernel_registry(configs={}) as register:
-            wrapper = register(
-                "autotune_test_kernel",
-                config_picker=lambda args, keys: "default",
-                fake_impl=lambda *a, **kw: None,
-                input_generator=lambda: {
-                    "small": (
-                        torch.randn(4, 4, device="cuda"),
-                        torch.randn(4, 4, device="cuda"),
-                    ),
-                },
-            )(_add_kernel)
-
-        assert wrapper._disabled is True
-
-        inputs = wrapper.get_inputs()
-        assert "small" in inputs
-
-        settings = helion.Settings()
-        settings.autotuner_fn = _no_compile_autotuner_fn
-        wrapper.helion_settings = settings
-
-        config = wrapper.run_autotune(inputs["small"])
-        expected_default = (
-            create_helion_decorated_kernel(_add_kernel, helion_settings=settings)
-            .bind(inputs["small"])
-            .config_spec.default_config()
-        )
-        assert config == expected_default

tests/kernels/helion/test_pattern_matching.py

@@ -52,7 +52,7 @@ def _helion_mock_context():
 
     with (
         patch(
-            "vllm.kernels.helion.config_manager.ConfigManager",
+            "vllm.kernels.helion.config_manager.ConfigManager.get_instance",
             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

tests/kernels/helion/test_register.py

@@ -21,9 +21,7 @@ if not has_helion():
     )
 
 import helion
-import helion.language as hl
 
-from tests.kernels.helion.helpers import dummy_kernel_registry
 from vllm.kernels.helion.config_manager import ConfigManager
 from vllm.kernels.helion.register import (
     _HOP_AVAILABLE,
@@ -36,13 +34,6 @@ from vllm.kernels.helion.register import (
 )
 
 
-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
-
-
 @pytest.fixture
 def sample_configs():
     """Create real Helion config objects for testing."""
@@ -99,7 +90,7 @@ def configured_kernel(sample_kernel, sample_configs, config_manager_with_test_co
 
     with (
         patch(
-            "vllm.kernels.helion.config_manager.ConfigManager",
+            "vllm.kernels.helion.config_manager.ConfigManager.get_instance",
             return_value=config_manager_with_test_configs,
         ),
         patch(
@@ -167,7 +158,7 @@ def create_configured_kernel_with_configs(
 
     with (
         patch(
-            "vllm.kernels.helion.config_manager.ConfigManager",
+            "vllm.kernels.helion.config_manager.ConfigManager.get_instance",
             return_value=mock_config_manager,
         ),
         patch(
@@ -198,7 +189,7 @@ class TestConfiguredHelionKernel:
 
         with (
             patch(
-                "vllm.kernels.helion.config_manager.ConfigManager",
+                "vllm.kernels.helion.config_manager.ConfigManager.get_instance",
                 return_value=mock_config_manager,
             ),
             patch(
@@ -275,7 +266,7 @@ class TestConfiguredHelionKernel:
         with (
             patch("vllm.kernels.helion.register.helion.kernel") as mock_kernel,
             patch(
-                "vllm.kernels.helion.config_manager.ConfigManager",
+                "vllm.kernels.helion.config_manager.ConfigManager.get_instance",
                 return_value=mock_config_manager,
             ),
             patch(
@@ -319,7 +310,7 @@ class TestConfiguredHelionKernel:
         with (
             patch("vllm.kernels.helion.register.helion.kernel") as mock_helion_kernel,
             patch(
-                "vllm.kernels.helion.config_manager.ConfigManager",
+                "vllm.kernels.helion.config_manager.ConfigManager.get_instance",
                 return_value=mock_config_manager,
             ),
             patch(
@@ -355,15 +346,23 @@ class TestConfiguredHelionKernel:
 class TestHelionKernelWrapper:
     """Test suite for HelionKernelWrapper."""
 
-    def test_init_disables_on_missing_configs(self, sample_kernel):
-        """Test __init__ marks wrapper as disabled when configs are missing."""
+    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={}
@@ -371,199 +370,72 @@ class TestHelionKernelWrapper:
 
         with (
             patch(
-                "vllm.kernels.helion.config_manager.ConfigManager",
+                "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,
+            pytest.raises(ValueError, match="No configs available"),
         ):
-            mock_kernel.return_value = Mock(return_value=sample_kernel)
-
-            wrapper = HelionKernelWrapper(
-                raw_kernel_func=sample_kernel,
-                op_name="test_kernel",
-                fake_impl=fake_impl,
-                config_picker=default_picker,
-            )
-
-            assert wrapper._disabled is True
-            assert "No configs available" in wrapper._disabled_reason
-
-    def test_disabled_wrapper_raises_on_call(self, sample_kernel):
-        """Test __call__ raises RuntimeError on a disabled wrapper."""
-
-        def fake_impl(*args, **kwargs):
-            return torch.zeros_like(args[0])
-
-        def default_picker(args, config_keys):
-            return "default"
-
-        mock_config_manager = Mock(spec=ConfigManager)
-        mock_config_manager.get_platform_configs = Mock(return_value={})
-
-        with (
-            patch(
-                "vllm.kernels.helion.config_manager.ConfigManager",
-                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_kernel.return_value = Mock(return_value=sample_kernel)
-
-            wrapper = HelionKernelWrapper(
-                raw_kernel_func=sample_kernel,
-                op_name="test_kernel",
-                fake_impl=fake_impl,
-                config_picker=default_picker,
-            )
-
-        with pytest.raises(RuntimeError, match="is disabled"):
-            wrapper(torch.randn(4, 4), torch.randn(4, 4))
-
-    def test_disabled_wrapper_get_configured_op_raises(self, sample_kernel):
-        """Test get_configured_op raises RuntimeError on a disabled wrapper."""
-
-        def fake_impl(*args, **kwargs):
-            return torch.zeros_like(args[0])
-
-        def default_picker(args, config_keys):
-            return "default"
-
-        mock_config_manager = Mock(spec=ConfigManager)
-        mock_config_manager.get_platform_configs = Mock(return_value={})
-
-        with (
-            patch(
-                "vllm.kernels.helion.config_manager.ConfigManager",
-                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_kernel.return_value = Mock(return_value=sample_kernel)
-
-            wrapper = HelionKernelWrapper(
-                raw_kernel_func=sample_kernel,
-                op_name="test_kernel",
-                fake_impl=fake_impl,
-                config_picker=default_picker,
-            )
-
-        with pytest.raises(RuntimeError, match="is disabled"):
             wrapper.get_configured_op()
 
-    def test_disabled_wrapper_supports_get_inputs(self, sample_kernel):
-        """Test get_inputs works on a disabled wrapper."""
+    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])
 
-        def default_picker(args, config_keys):
-            return "default"
-
-        expected_inputs = {"key1": (torch.randn(4),)}
-        input_gen = Mock(return_value=expected_inputs)
-
-        mock_config_manager = Mock(spec=ConfigManager)
-        mock_config_manager.get_platform_configs = Mock(return_value={})
-
-        with (
-            patch(
-                "vllm.kernels.helion.config_manager.ConfigManager",
-                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_kernel.return_value = Mock(return_value=sample_kernel)
-
-            wrapper = HelionKernelWrapper(
-                raw_kernel_func=sample_kernel,
-                op_name="test_kernel",
-                fake_impl=fake_impl,
-                config_picker=default_picker,
-                input_generator=input_gen,
-            )
-
-        assert wrapper._disabled is True
-        result = wrapper.get_inputs()
-        assert result is expected_inputs
-
-    def test_disabled_wrapper_supports_run_autotune(self, sample_kernel):
-        """Test run_autotune works on a disabled wrapper."""
-
-        def fake_impl(*args, **kwargs):
-            return torch.zeros_like(args[0])
-
-        def default_picker(args, config_keys):
-            return "default"
-
-        mock_config_manager = Mock(spec=ConfigManager)
-        mock_config_manager.get_platform_configs = Mock(return_value={})
-
-        mock_config = Mock()
-
-        with (
-            patch(
-                "vllm.kernels.helion.config_manager.ConfigManager",
-                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_kernel.return_value = Mock(return_value=sample_kernel)
-
-            wrapper = HelionKernelWrapper(
-                raw_kernel_func=sample_kernel,
-                op_name="test_kernel",
-                fake_impl=fake_impl,
-                config_picker=default_picker,
-            )
-
-        assert wrapper._disabled is True
-
-        with patch(
-            "vllm.kernels.helion.register.create_helion_decorated_kernel"
-        ) as mock_create:
-            mock_autotune_kernel = Mock()
-            mock_autotune_kernel.autotune.return_value = mock_config
-            mock_create.return_value = mock_autotune_kernel
-
-            inputs = (torch.randn(4, 4),)
-            result = wrapper.run_autotune(inputs)
-            assert result is mock_config
-
-    def test_init_caches_configured_kernel(self, sample_kernel, sample_configs):
-        """Test __init__ eagerly builds and caches 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,
+        )
+        # 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",
+                "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(
@@ -572,77 +444,13 @@ class TestHelionKernelWrapper:
             ),
             patch("vllm.kernels.helion.register.helion.kernel") as mock_kernel,
         ):
-            mock_kernel.return_value = Mock(return_value=sample_kernel)
+            mock_decorated = Mock()
+            mock_kernel.return_value = Mock(return_value=mock_decorated)
 
-            wrapper = HelionKernelWrapper(
-                raw_kernel_func=sample_kernel,
-                op_name="test_kernel",
-                fake_impl=fake_impl,
-                config_picker=default_picker,
-            )
-
-            assert wrapper._configured_kernel is not None
             result1 = wrapper.get_configured_op()
             result2 = wrapper.get_configured_op()
             assert result1 is result2
 
-    @pytest.mark.skipif(
-        not _HOP_AVAILABLE, reason="HOP path only used when HOP available"
-    )
-    def test_init_eagerly_initializes_hop_path(self):
-        """Test that register_kernel eagerly builds the configured kernel
-        on the HOP path (no custom op registration needed)."""
-        from vllm.kernels.helion.utils import get_canonical_gpu_name
-
-        configs = {"default": helion.Config(block_sizes=[4, 4])}
-        with (
-            dummy_kernel_registry(configs=configs) as register,
-            patch(
-                "vllm.kernels.helion.utils.get_canonical_gpu_name",
-                wraps=get_canonical_gpu_name,
-            ) as mock_gpu,
-        ):
-            wrapper = register(
-                config_picker=lambda args, keys: "default",
-            )(_add_kernel)
-
-            mock_gpu.assert_called_once()
-            assert wrapper._configured_kernel is not None
-
-        with patch(
-            "vllm.kernels.helion.utils.get_canonical_gpu_name",
-            side_effect=AssertionError("get_canonical_gpu_name called during __call__"),
-        ):
-            x = torch.randn(4, 4, device="cuda")
-            y = torch.randn(4, 4, device="cuda")
-            result = wrapper(x, y)
-            expected = x + y
-            assert torch.allclose(result, expected)
-
-    @pytest.mark.skipif(
-        _HOP_AVAILABLE, reason="CustomOp path not used when HOP available"
-    )
-    def test_init_eagerly_initializes(self):
-        """Test that register_kernel eagerly loads configs and detects GPU
-        during construction so __call__ needs no further initialization."""
-        from vllm.kernels.helion.utils import get_canonical_gpu_name
-
-        with (
-            dummy_kernel_registry() as register,
-            patch(
-                "vllm.kernels.helion.utils.get_canonical_gpu_name",
-                wraps=get_canonical_gpu_name,
-            ) as mock_gpu,
-        ):
-            wrapper = register(
-                config_picker=lambda args, keys: "default",
-            )(_add_kernel)
-
-            # Init must have detected GPU and built the kernel
-            mock_gpu.assert_called_once()
-            assert wrapper._configured_kernel is not None
-            assert hasattr(torch.ops.vllm_helion, wrapper.op_name)
-
     @pytest.mark.skipif(
         _HOP_AVAILABLE, reason="CustomOp path not used when HOP available"
     )
@@ -655,6 +463,13 @@ class TestHelionKernelWrapper:
         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)
 
@@ -664,7 +479,7 @@ class TestHelionKernelWrapper:
 
         with (
             patch(
-                "vllm.kernels.helion.config_manager.ConfigManager",
+                "vllm.kernels.helion.config_manager.ConfigManager.get_instance",
                 return_value=mock_config_manager,
             ),
             patch(
@@ -676,13 +491,6 @@ class TestHelionKernelWrapper:
         ):
             mock_decorated = Mock()
             mock_kernel.return_value = Mock(return_value=mock_decorated)
-
-            wrapper = HelionKernelWrapper(
-                raw_kernel_func=sample_kernel,
-                op_name="test_kernel",
-                fake_impl=fake_impl,
-                config_picker=default_picker,
-            )
             result = wrapper._get_or_register_custom_op()
             assert result is existing_op
 
@@ -698,6 +506,13 @@ class TestHelionKernelWrapper:
         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)
 
@@ -717,7 +532,7 @@ class TestHelionKernelWrapper:
 
         with (
             patch(
-                "vllm.kernels.helion.config_manager.ConfigManager",
+                "vllm.kernels.helion.config_manager.ConfigManager.get_instance",
                 return_value=mock_config_manager,
             ),
             patch(
@@ -733,13 +548,6 @@ class TestHelionKernelWrapper:
         ):
             mock_decorated = Mock()
             mock_kernel.return_value = Mock(return_value=mock_decorated)
-
-            wrapper = HelionKernelWrapper(
-                raw_kernel_func=sample_kernel,
-                op_name="test_kernel",
-                fake_impl=fake_impl,
-                config_picker=default_picker,
-            )
             result = wrapper._get_or_register_custom_op()
 
             mock_register.assert_called_once()
@@ -776,10 +584,11 @@ class TestKernelRegistry:
 
     def test_get_kernel_by_name_returns_kernel(self):
         """Test get_kernel_by_name returns registered kernel."""
-        with dummy_kernel_registry() as register:
-            wrapper = register(
-                "test_kernel", config_picker=lambda args, keys: "default"
-            )(_add_kernel)
+        wrapper = HelionKernelWrapper(
+            raw_kernel_func=Mock(),
+            op_name="test_kernel",
+            fake_impl=Mock(),
+        )
 
         from vllm.kernels.helion.register import _REGISTERED_KERNELS
 
@@ -795,87 +604,112 @@ class TestKernelRegistry:
 
     def test_register_kernel_auto_generates_fake_impl(self):
         """Test register_kernel auto-generates fake_impl when not provided."""
-        with (
-            dummy_kernel_registry() as register,
-            patch("vllm.kernels.helion.register.infer_fake_impl") as mock_infer,
-        ):
+        with patch("vllm.kernels.helion.register.infer_fake_impl") as mock_infer:
             mock_fake = Mock()
             mock_infer.return_value = mock_fake
-            wrapper = register(
-                config_picker=lambda args, keys: "default",
-            )(_add_kernel)
 
-        mock_infer.assert_called_once_with(_add_kernel, None)
-        assert wrapper._fake_impl is 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."""
-        with dummy_kernel_registry() as register:
-            result = register("test_name", config_picker=lambda args, keys: "default")(
-                _add_kernel
-            )
+
+        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 _add_kernel
+        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."""
-        with dummy_kernel_registry() as register:
-            wrapper = register(config_picker=lambda args, keys: "default")(_add_kernel)
 
-        assert wrapper.op_name == "_add_kernel"
+        @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."""
-        with dummy_kernel_registry() as register:
-            wrapper = register(
-                "test_kernel", config_picker=lambda args, keys: "default"
-            )(_add_kernel)
+
+        @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 wrapper
+        assert registered_kernels["test_kernel"] is test_kernel
 
     def test_register_kernel_passes_helion_settings(self):
         """Test register_kernel passes helion_settings to wrapper."""
-        settings = helion.Settings()
-        settings.print_output_code = True
+        mock_settings = Mock()
+        mock_settings.to_dict.return_value = {"debug": True}
 
-        with dummy_kernel_registry() as register:
-            result = register(
-                "test_name",
-                config_picker=lambda args, keys: "default",
-                helion_settings=settings,
-            )(_add_kernel)
+        @register_kernel("test_name", helion_settings=mock_settings)
+        def test_kernel(x):
+            return x
 
-        assert result.helion_settings is settings
+        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()
 
-        with dummy_kernel_registry() as register:
-            result = register(
-                "custom_name",
-                config_picker=lambda args, keys: "default",
-                fake_impl=mock_fake,
-            )(_add_kernel)
+        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."""
-        with dummy_kernel_registry() as register:
-            register("duplicate_name", config_picker=lambda args, keys: "default")(
-                _add_kernel
-            )
 
-            with pytest.raises(ValueError, match="already registered"):
-                register("duplicate_name", config_picker=lambda args, keys: "default")(
-                    _add_kernel
-                )
+        @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."""
@@ -884,11 +718,7 @@ class TestKernelRegistry:
 
         with pytest.raises(ValueError, match="uses a custom autotuner"):
 
-            @register_kernel(
-                "test",
-                config_picker=lambda args, keys: "default",
-                helion_settings=mock_settings,
-            )
+            @register_kernel("test", helion_settings=mock_settings)
             def test_kernel(x):
                 return x
 
@@ -897,47 +727,11 @@ class TestKernelRegistry:
         mock_settings = Mock()
         mock_settings.to_dict.return_value = {"static_shapes": False}
 
-        with (
-            dummy_kernel_registry() as register,
-            patch("vllm.kernels.helion.register.logger") as mock_logger,
-        ):
-            register(
-                "test",
-                config_picker=lambda args, keys: "default",
-                helion_settings=mock_settings,
-            )(_add_kernel)
+        with patch("vllm.kernels.helion.register.logger") as mock_logger:
 
-        mock_logger.warning.assert_not_called()
+            @register_kernel("test", helion_settings=mock_settings)
+            def test_kernel(x):
+                return x
 
-    def test_disabled_kernel_appears_in_registry(self):
-        """Test that a disabled wrapper is still in the global registry."""
-
-        def fake_impl(*args, **kwargs):
-            return torch.zeros_like(args[0])
-
-        mock_config_manager = Mock(spec=ConfigManager)
-        mock_config_manager.get_platform_configs = Mock(return_value={})
-
-        with (
-            patch(
-                "vllm.kernels.helion.config_manager.ConfigManager",
-                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_kernel.return_value = Mock(return_value=_add_kernel)
-
-            wrapper = register_kernel(
-                "disabled_kernel",
-                config_picker=lambda args, keys: "default",
-                fake_impl=fake_impl,
-            )(_add_kernel)
-
-        assert wrapper._disabled is True
-        registered = get_registered_kernels()
-        assert "disabled_kernel" in registered
-        assert registered["disabled_kernel"] is wrapper
+            # Should not call warning
+            mock_logger.warning.assert_not_called()

tests/tool_parsers/test_glm47_moe_tool_parser.py

@@ -1,168 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-# ruff: noqa: E501
-"""Tests for the GLM-4.7 tool call parser."""
-
-import json
-from unittest.mock import Mock
-
-import pytest
-
-from vllm.entrypoints.openai.chat_completion.protocol import (
-    ChatCompletionRequest,
-    ChatCompletionToolsParam,
-    FunctionDefinition,
-)
-from vllm.tokenizers import get_tokenizer
-from vllm.tool_parsers.glm47_moe_tool_parser import Glm47MoeModelToolParser
-
-MODEL = "zai-org/GLM-4.5"
-
-
-@pytest.fixture(scope="module")
-def glm47_tokenizer():
-    return get_tokenizer(tokenizer_name=MODEL)
-
-
-@pytest.fixture
-def glm47_tool_parser(glm47_tokenizer):
-    return Glm47MoeModelToolParser(glm47_tokenizer)
-
-
-@pytest.fixture
-def mock_request() -> ChatCompletionRequest:
-    request = Mock(spec=ChatCompletionRequest)
-    request.tools = [
-        ChatCompletionToolsParam(
-            function=FunctionDefinition(name="get_current_date", parameters={}),
-        ),
-        ChatCompletionToolsParam(
-            function=FunctionDefinition(
-                name="get_weather",
-                parameters={
-                    "type": "object",
-                    "properties": {
-                        "city": {"type": "string"},
-                        "date": {"type": "string"},
-                    },
-                },
-            ),
-        ),
-    ]
-    request.tool_choice = "auto"
-    return request
-
-
-class TestGlm47ExtractToolCalls:
-    def test_no_tool_call(self, glm47_tool_parser, mock_request):
-        out = "This is a plain response."
-        r = glm47_tool_parser.extract_tool_calls(out, request=mock_request)
-        assert not r.tools_called
-        assert r.content == out
-
-    def test_zero_arg_inline(self, glm47_tool_parser, mock_request):
-        out = "<tool_call>get_current_date</tool_call>"
-        r = glm47_tool_parser.extract_tool_calls(out, request=mock_request)
-        assert r.tools_called
-        assert r.tool_calls[0].function.name == "get_current_date"
-        assert json.loads(r.tool_calls[0].function.arguments) == {}
-        assert r.content is None
-
-    def test_zero_arg_newline(self, glm47_tool_parser, mock_request):
-        out = "<tool_call>get_current_date\n</tool_call>"
-        r = glm47_tool_parser.extract_tool_calls(out, request=mock_request)
-        assert r.tools_called
-        assert r.tool_calls[0].function.name == "get_current_date"
-
-    def test_args_same_line(self, glm47_tool_parser, mock_request):
-        out = "<tool_call>get_weather<arg_key>city</arg_key><arg_value>Beijing</arg_value></tool_call>"
-        r = glm47_tool_parser.extract_tool_calls(out, request=mock_request)
-        assert r.tools_called
-        assert json.loads(r.tool_calls[0].function.arguments) == {"city": "Beijing"}
-
-    def test_args_with_newlines(self, glm47_tool_parser, mock_request):
-        out = "<tool_call>get_weather\n<arg_key>city</arg_key>\n<arg_value>Beijing</arg_value>\n</tool_call>"
-        r = glm47_tool_parser.extract_tool_calls(out, request=mock_request)
-        assert r.tools_called
-        assert json.loads(r.tool_calls[0].function.arguments) == {"city": "Beijing"}
-
-    def test_content_before(self, glm47_tool_parser, mock_request):
-        out = "Checking.<tool_call>get_current_date</tool_call>"
-        r = glm47_tool_parser.extract_tool_calls(out, request=mock_request)
-        assert r.tools_called
-        assert r.content == "Checking."
-
-    def test_multiple(self, glm47_tool_parser, mock_request):
-        out = (
-            "<tool_call>get_weather<arg_key>city</arg_key><arg_value>Beijing</arg_value></tool_call>"
-            "<tool_call>get_weather<arg_key>city</arg_key><arg_value>Shanghai</arg_value></tool_call>"
-        )
-        r = glm47_tool_parser.extract_tool_calls(out, request=mock_request)
-        assert len(r.tool_calls) == 2
-
-    def test_empty_content_none(self, glm47_tool_parser, mock_request):
-        out = "<tool_call>get_current_date</tool_call>"
-        r = glm47_tool_parser.extract_tool_calls(out, request=mock_request)
-        assert r.content is None
-
-    def test_whitespace_content_none(self, glm47_tool_parser, mock_request):
-        out = "  \n  <tool_call>get_current_date</tool_call>"
-        r = glm47_tool_parser.extract_tool_calls(out, request=mock_request)
-        assert r.content is None
-
-
-def _reset(parser):
-    parser._buffer = ""
-    parser._in_tool_call = False
-    parser.current_tool_name_sent = False
-    parser._current_tool_name = None
-    parser._pending_key = None
-    parser._streaming_string_value = False
-    parser.prev_tool_call_arr = []
-    parser.current_tool_id = -1
-    parser.streamed_args_for_tool = []
-    parser._tool_call_ids = []
-    parser._args_started = []
-    parser._args_closed = []
-    parser._seen_keys = []
-
-
-class TestGlm47Streaming:
-    def test_no_args(self, glm47_tool_parser, mock_request):
-        _reset(glm47_tool_parser)
-        for chunk in ["<tool_call>", "get_current_date", "</tool_call>"]:
-            glm47_tool_parser.extract_tool_calls_streaming(
-                previous_text="",
-                current_text="",
-                delta_text=chunk,
-                previous_token_ids=[],
-                current_token_ids=[],
-                delta_token_ids=[],
-                request=mock_request,
-            )
-        assert len(glm47_tool_parser.prev_tool_call_arr) >= 1
-
-    def test_with_args(self, glm47_tool_parser, mock_request):
-        _reset(glm47_tool_parser)
-        # Split chunks so that the incremental string streaming path
-        # processes the value, its closing tag, and the tool-call closing
-        # tag in separate calls.
-        for chunk in [
-            "<tool_call>",
-            "get_weather\n",
-            "<arg_key>city</arg_key>",
-            "<arg_value>",
-            "Beijing",
-            "</arg_value>",
-            "</tool_call>",
-        ]:
-            glm47_tool_parser.extract_tool_calls_streaming(
-                previous_text="",
-                current_text="",
-                delta_text=chunk,
-                previous_token_ids=[],
-                current_token_ids=[],
-                delta_token_ids=[],
-                request=mock_request,
-            )
-        assert glm47_tool_parser.prev_tool_call_arr[0]["arguments"]["city"] == "Beijing"

tests/tool_parsers/test_glm4_moe_tool_parser.py

@@ -107,7 +107,7 @@ def test_extract_tool_calls_no_tools(glm4_moe_tool_parser, mock_request):
                     )
                 )
             ],
-            None,
+            "",
         ),
         (
             """<tool_call>get_current_weather
@@ -152,7 +152,7 @@ def test_extract_tool_calls_no_tools(glm4_moe_tool_parser, mock_request):
                     )
                 ),
             ],
-            None,
+            "",
         ),
         (
             """I'll help you check the weather. <tool_call>get_current_weather
@@ -202,7 +202,7 @@ def test_extract_tool_calls_no_tools(glm4_moe_tool_parser, mock_request):
                     )
                 )
             ],
-            None,
+            "",
         ),
         (
             """I will help you get the weather.<tool_call>get_weather

tests/v1/executor/test_executor.py

@@ -14,35 +14,12 @@ from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs
 from vllm.sampling_params import SamplingParams
 from vllm.v1.engine.async_llm import AsyncLLM
 from vllm.v1.engine.llm_engine import LLMEngine
-from vllm.v1.executor.abstract import Executor
 from vllm.v1.executor.multiproc_executor import MultiprocExecutor
-from vllm.v1.executor.uniproc_executor import (
-    ExecutorWithExternalLauncher,
-    UniProcExecutor,
-)
 
 
 class Mock: ...
 
 
-def test_supports_async_scheduling_base_executor():
-    assert Executor.supports_async_scheduling() is False
-
-
-def test_supports_async_scheduling_uniproc_executor():
-    assert UniProcExecutor.supports_async_scheduling() is True
-
-
-def test_supports_async_scheduling_executor_with_external_launcher():
-    # ExecutorWithExternalLauncher inherits from UniProcExecutor and does not
-    # override supports_async_scheduling, so it should return True.
-    assert ExecutorWithExternalLauncher.supports_async_scheduling() is True
-
-
-def test_supports_async_scheduling_multiproc_executor():
-    assert MultiprocExecutor.supports_async_scheduling() is True
-
-
 class CustomMultiprocExecutor(MultiprocExecutor):
     def collective_rpc(
         self,

tests/v1/kv_connector/unit/test_multi_connector.py

@@ -231,11 +231,10 @@ def test_multi_example_connector_consistency():
     ]
     # First three events are from initialization (register_kv_caches,
     # set_host_xfer_buffer_ops, get_handshake_metadata), then generate() events.
-    assert events["storage1-WORKER"][:8] == [
+    assert events["storage1-WORKER"][:7] == [
         "register_kv_caches",
         "set_host_xfer_buffer_ops",
         "get_handshake_metadata",
-        "handle_preemptions",
         "bind_connector_metadata",
         "start_load_kv",
         "wait_for_layer_load",
@@ -247,11 +246,10 @@ def test_multi_example_connector_consistency():
         "update_state_after_alloc num_blocks=[0] 0",
         "build_connector_meta",
     ]
-    assert events["storage2-WORKER"][:8] == [
+    assert events["storage2-WORKER"][:7] == [
         "register_kv_caches",
         "set_host_xfer_buffer_ops",
         "get_handshake_metadata",
-        "handle_preemptions",
         "bind_connector_metadata",
         "start_load_kv",
         "wait_for_layer_load",
@@ -401,8 +399,8 @@ def test_multi_connector_handle_preemptions_integration():
         # testing the delegation behavior of MultiConnector here.
         # The connector attribute contains the KV connector.
         assert scheduler.connector is not None, "Scheduler should have a connector"
-        connector_md = scheduler.connector.build_connector_meta(scheduler.schedule())
-        scheduler.connector.handle_preemptions(connector_md)
+        preempted_req_ids = {"req-1", "req-2", "req-3"}
+        scheduler.connector.handle_preemptions(preempted_req_ids)
 
         # Verify both connectors received the handle_preemptions call
         events = get_connector_events()

tests/v1/kv_connector/unit/test_nixl_connector.py

@@ -1369,13 +1369,7 @@ def test_abort_timeout_on_prefiller(monkeypatch, distributed_executor_backend):
                     "NIXL_TELEMETRY_ENABLE": "1",
                 },
             }
-            # On XPU/ROCm, vLLM expects Ray's device key to be "GPU".
-            # Explicitly reserving GPU resources here prevents false negatives
-            # when Ray cannot auto-detect accelerator resources in test envs.
-            ray_init_kwargs: dict[str, Any] = {"runtime_env": runtime_env}
-            if not current_platform.is_cuda():
-                ray_init_kwargs["num_gpus"] = 1
-            ray.init(**ray_init_kwargs)
+            ray.init(runtime_env=runtime_env)
             try:
                 run_test_and_cleanup()
             finally:

tests/v1/kv_connector/unit/test_offloading_connector.py

@@ -363,7 +363,10 @@ class RequestRunner:
             assert kv_connector_metadata is not None
             assert isinstance(kv_connector_metadata, OffloadingConnectorMetadata)
 
-            self.worker_connector.handle_preemptions(kv_connector_metadata)
+            if scheduler_output.preempted_req_ids:
+                self.worker_connector.handle_preemptions(
+                    scheduler_output.preempted_req_ids
+                )
 
             self.worker_connector.bind_connector_metadata(kv_connector_metadata)
             self.worker_connector.start_load_kv(self._dummy_ctx)

tests/v1/kv_offload/test_cpu_offloading.py

@@ -22,17 +22,6 @@ if current_platform.is_cuda():
 elif current_platform.is_rocm():
     ATTN_BACKENDS = ["TRITON_ATTN"]
 
-# Maximum time (seconds) to wait for the async CPU offload transfer
-# to complete before giving up.
-_RESET_CACHE_TIMEOUT = 30 if current_platform.is_rocm() else 10
-
-# ZMQ poll timeout (ms) for the first event.
-_FIRST_EVENT_POLL_MS = 10_000 if current_platform.is_rocm() else 1000
-
-# Hard ceiling (seconds) on how long get_new_cpu_stored_events may loop,
-# to prevent hangs if non-CPU events keep arriving indefinitely.
-_EVENT_DRAIN_TIMEOUT = 60
-
 
 class MockSubscriber:
     """Helper class to receive and verify published events"""
@@ -58,10 +47,9 @@ class MockSubscriber:
         poller = zmq.Poller()
         poller.register(self.sub, zmq.POLLIN)
 
-        poll_ms = _FIRST_EVENT_POLL_MS
-        deadline = time.monotonic() + _EVENT_DRAIN_TIMEOUT
-        while time.monotonic() < deadline:
-            events = dict(poller.poll(poll_ms))
+        timeout = 1000  # 1 second
+        while True:
+            events = dict(poller.poll(timeout))
 
             if events.get(self.sub) != zmq.POLLIN:
                 return cpu_stored_events
@@ -75,32 +63,13 @@ class MockSubscriber:
             for event in event_batch.events:
                 if isinstance(event, BlockStored) and event.medium == "CPU":
                     cpu_stored_events.append(event)
-                    poll_ms = 100
-
-        return cpu_stored_events
+                    timeout = 100
 
     def close(self):
         """Clean up resources"""
         self.sub.close()
 
 
-def _wait_for_prefix_cache_reset(llm: LLM) -> None:
-    """Wait for async offload transfers to finish so prefix cache can reset.
-
-    The GPU-to-CPU offload runs on a CUDA stream asynchronously.  While blocks
-    are still held by the offload worker, ``reset_prefix_cache`` returns
-    ``False``.  Retry with a short sleep until it succeeds or we time out.
-    """
-    deadline = time.monotonic() + _RESET_CACHE_TIMEOUT
-    while not llm.reset_prefix_cache():
-        if time.monotonic() > deadline:
-            raise TimeoutError(
-                "reset_prefix_cache did not succeed within "
-                f"{_RESET_CACHE_TIMEOUT}s - async offload may be stuck"
-            )
-        time.sleep(0.1)
-
-
 def _latency_test(llm: LLM, subscriber: MockSubscriber):
     sampling_params = SamplingParams(max_tokens=1)
 
@@ -126,16 +95,10 @@ def _latency_test(llm: LLM, subscriber: MockSubscriber):
         gpu_hit_time = time.time() - start_time
         total_gpu_hit_time += gpu_hit_time
 
-        # Wait for the async CPU offload to finish, then reset prefix cache
-        # so the next generate() must reload from CPU rather than GPU.
-        _wait_for_prefix_cache_reset(llm)
+        # reset prefix cache to avoid GPU hit.
+        llm.reset_prefix_cache()
 
-        # Verify CPU stored events arrived (offload is done before we
-        # attempt to load from CPU).
-        assert subscriber.get_new_cpu_stored_events(), (
-            f"No CPU stored events received on iteration {i}; "
-            "async offload may not have completed in time"
-        )
+        assert subscriber.get_new_cpu_stored_events()
 
         # run generation again - this should trigger loading from CPU
         start_time = time.time()
@@ -222,8 +185,6 @@ def test_cpu_offloading(cpu_block_size: int, attn_backend: str) -> None:
         kv_events_config=kv_events_config,
         kv_transfer_config=kv_transfer_config,
         attention_config={"backend": attn_backend},
-        # ROCm: batch size 1 to reduce variability
-        **({"max_num_seqs": 1} if current_platform.is_rocm() else {}),
     )
 
     events_endpoint = events_endpoint.replace("*", "127.0.0.1")

vllm/config/parallel.py

@@ -138,13 +138,6 @@ class ParallelConfig:
     """Whether the deployed model is MoE (if known)."""
     enable_expert_parallel: bool = False
     """Use expert parallelism instead of tensor parallelism for MoE layers."""
-    enable_ep_weight_filter: bool = False
-    """Skip non-local expert weights during model loading when expert
-    parallelism is active.  Each rank only reads its own expert shard from
-    disk, which can drastically reduce storage I/O for MoE models with
-    per-expert weight tensors (e.g. DeepSeek, Mixtral, Kimi-K2.5).  Has no
-    effect on 3D fused-expert checkpoints (e.g. GPT-OSS) or non-MoE
-    models."""
     enable_eplb: bool = False
     """Enable expert parallelism load balancing for MoE layers."""
     eplb_config: EPLBConfig = Field(default_factory=EPLBConfig)

vllm/config/vllm.py

@@ -682,11 +682,12 @@ class VllmConfig:
                 self.model_config, self.load_config
             )
 
-        from vllm.v1.executor.abstract import Executor
-
         executor_backend = self.parallel_config.distributed_executor_backend
-        executor_class = Executor.get_class(self)
-        executor_supports_async_sched = executor_class.supports_async_scheduling()
+        executor_supports_async_sched = executor_backend in (
+            "mp",
+            "uni",
+            "external_launcher",
+        )
 
         if self.scheduler_config.async_scheduling:
             # Async scheduling explicitly enabled, hard fail any incompatibilities.
@@ -710,7 +711,9 @@ class VllmConfig:
                     )
             if not executor_supports_async_sched:
                 raise ValueError(
-                    f"`{executor_backend}` does not support async scheduling yet."
+                    "Currently, async scheduling only supports `mp`, `uni`, or "
+                    "`external_launcher` distributed executor backend, but you chose "
+                    f"`{executor_backend}`."
                 )
         elif self.scheduler_config.async_scheduling is None:
             # Enable async scheduling unless there is an incompatible option.
@@ -739,7 +742,8 @@ class VllmConfig:
             elif not executor_supports_async_sched:
                 logger.warning_once(
                     "Async scheduling will be disabled because it is not supported "
-                    "with the `%s` distributed executor backend. ",
+                    "with the `%s` distributed executor backend (only `mp`, `uni`, and "
+                    "`external_launcher` are supported).",
                     executor_backend,
                     scope="local",
                 )

vllm/distributed/kv_transfer/kv_connector/v1/base.py

@@ -25,8 +25,8 @@ The class provides the following primitives:
 
     Worker-side: runs in each worker, loads/saves KV cache to/from
     the Connector based on the metadata.
-        handle_preemptions() - called for handling preempted requests
-            or request evicted blocks before they are overwritten
+        handle_preemptions() - called if there are preempted requests,
+            before their blocks are overwritten
 
         start_load_kv() - starts loading all KVs (maybe async)
         wait_for_layer_load() - blocks until layer i load is done
@@ -288,9 +288,9 @@ class KVConnectorBase_V1(ABC):
         """
         return
 
-    def handle_preemptions(self, kv_connector_metadata: KVConnectorMetadata):
+    def handle_preemptions(self, preempted_req_ids: set[str]):
         """
-        Handle preempted requests or evicted blocks BEFORE they are overwritten.
+        Handle preempted requests BEFORE their blocks are overwritten.
         Needed for connectors which use async saves (e.g., OffloadingConnector)
         """
         return

vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py

@@ -315,11 +315,10 @@ class MultiConnector(KVConnectorBase_V1):
         for c in self._connectors:
             c.set_host_xfer_buffer_ops(copy_operation)
 
-    def handle_preemptions(self, kv_connector_metadata: KVConnectorMetadata):
+    def handle_preemptions(self, preempted_req_ids: set[str]):
         """Handle preempted requests for all sub-connectors."""
-        assert isinstance(kv_connector_metadata, MultiKVConnectorMetadata)
-        for c, cm in zip(self._connectors, kv_connector_metadata.metadata):
-            c.handle_preemptions(cm)
+        for c in self._connectors:
+            c.handle_preemptions(preempted_req_ids)
 
     def get_finished_count(self) -> int | None:
         # TODO(https://github.com/vllm-project/vllm/issues/33400)

vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py

@@ -111,7 +111,6 @@ class OffloadingConnectorStats(KVConnectorStats):
 class OffloadingConnectorMetadata(KVConnectorMetadata):
     reqs_to_load: dict[ReqId, TransferSpec]
     reqs_to_store: dict[ReqId, TransferSpec]
-    reqs_to_flush: set[str] | None = None
 
 
 class OffloadingConnector(KVConnectorBase_V1):
@@ -147,10 +146,9 @@ class OffloadingConnector(KVConnectorBase_V1):
         assert self.connector_worker is not None
         self.connector_worker.register_cross_layers_kv_cache(kv_cache, attn_backend)
 
-    def handle_preemptions(self, kv_connector_metadata: KVConnectorMetadata):
+    def handle_preemptions(self, preempted_req_ids: set[str]):
         assert self.connector_worker is not None
-        assert isinstance(kv_connector_metadata, OffloadingConnectorMetadata)
-        self.connector_worker.handle_preemptions(kv_connector_metadata)
+        self.connector_worker.handle_preemptions(preempted_req_ids)
 
     def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
         assert self.connector_worker is not None
@@ -484,7 +482,6 @@ class OffloadingConnectorScheduler:
         meta = OffloadingConnectorMetadata(
             reqs_to_load=self._reqs_to_load,
             reqs_to_store=self._get_reqs_to_store(scheduler_output),
-            reqs_to_flush=scheduler_output.preempted_req_ids,
         )
         self._reqs_to_load = {}
 
@@ -622,13 +619,13 @@ class OffloadingConnectorWorker:
         attn_backends = {cross_layer_name: attn_backend}
         self._register_handlers(kv_caches, attn_backends)
 
-    def handle_preemptions(self, kv_connector_metadata: OffloadingConnectorMetadata):
+    def handle_preemptions(self, preempted_req_ids: set[str]):
         for job_id, transfer_spec in self._unsubmitted_store_jobs:
             success = self.worker.transfer_async(job_id, transfer_spec)
             assert success
         self._unsubmitted_store_jobs.clear()
 
-        for req_id in kv_connector_metadata.reqs_to_flush or ():
+        for req_id in preempted_req_ids:
             job_ids = self._store_jobs.get(req_id)
             if job_ids:
                 self.worker.wait(job_ids)

vllm/engine/arg_utils.py

@@ -419,7 +419,6 @@ class EngineArgs:
     data_parallel_external_lb: bool = False
     data_parallel_backend: DataParallelBackend = ParallelConfig.data_parallel_backend
     enable_expert_parallel: bool = ParallelConfig.enable_expert_parallel
-    enable_ep_weight_filter: bool = ParallelConfig.enable_ep_weight_filter
     moe_backend: MoEBackend = KernelConfig.moe_backend
     all2all_backend: All2AllBackend = ParallelConfig.all2all_backend
     enable_elastic_ep: bool = ParallelConfig.enable_elastic_ep
@@ -903,10 +902,6 @@ class EngineArgs:
             "-ep",
             **parallel_kwargs["enable_expert_parallel"],
         )
-        parallel_group.add_argument(
-            "--enable-ep-weight-filter",
-            **parallel_kwargs["enable_ep_weight_filter"],
-        )
         parallel_group.add_argument(
             "--all2all-backend", **parallel_kwargs["all2all_backend"]
         )
@@ -1736,7 +1731,6 @@ class EngineArgs:
             data_parallel_hybrid_lb=self.data_parallel_hybrid_lb,
             is_moe_model=model_config.is_moe,
             enable_expert_parallel=self.enable_expert_parallel,
-            enable_ep_weight_filter=self.enable_ep_weight_filter,
             all2all_backend=self.all2all_backend,
             enable_elastic_ep=self.enable_elastic_ep,
             enable_dbo=self.enable_dbo,

vllm/entrypoints/openai/responses/serving.py

@@ -1012,7 +1012,6 @@ class OpenAIServingResponses(OpenAIServing):
             parser = self.parser(tokenizer)
             return parser.extract_response_outputs(
                 model_output=final_output.text,
-                model_output_token_ids=final_output.token_ids,
                 request=request,
                 enable_auto_tools=self.enable_auto_tools,
                 tool_call_id_type=self.tool_call_id_type,

vllm/kernels/helion/ops/silu_mul_fp8.py

@@ -22,6 +22,39 @@ from vllm.kernels.helion.register import register_kernel
 logger = init_logger(__name__)
 
 
+@register_kernel  # type: ignore[misc]
+def silu_mul_fp8(input: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    original_shape = input.shape
+    two_d = hl.specialize(original_shape[-1])
+    d = two_d // 2
+    output_shape = original_shape[:-1] + (d,)
+
+    input_2d = input.view(-1, original_shape[-1])
+    m = input_2d.shape[0]
+
+    # TODO(gmagogsfm): Support for more float8 subtypes (e4m3fnuz, e5m2) coming
+    out = torch.empty((m, d), device=input.device, dtype=torch.float8_e4m3fn)
+
+    input_part_a = input_2d[:, :d]
+    input_part_b = input_2d[:, d:]
+
+    assert scale.numel() == 1, "Scale must be a scalar Tensor"
+
+    for tile_m, tile_n in hl.tile([m, d]):
+        a_vals = input_part_a[tile_m, tile_n]
+        silu_result = torch.nn.functional.silu(a_vals)
+        b_vals = input_part_b[tile_m, tile_n]
+        result = silu_result * b_vals
+        result_f32 = result.to(torch.float32)
+        scale_val = hl.load(scale, [0])
+        inv_scale = 1.0 / scale_val
+        result_scaled = result_f32 * inv_scale
+        out[tile_m, tile_n] = result_scaled.to(out.dtype)
+
+    return out.view(output_shape)
+
+
+@silu_mul_fp8.register_input_generator  # type: ignore[misc]
 def generate_silu_mul_fp8_inputs() -> dict[str, tuple[Any, ...]]:
     intermediate_sizes = [2048, 2880, 4096, 8192, 11008, 14336]
 
@@ -32,6 +65,8 @@ def generate_silu_mul_fp8_inputs() -> dict[str, tuple[Any, ...]]:
     inputs = {}
     for num_tokens in num_tokens_list:
         for intermediate_size in intermediate_sizes:
+            # Input tensor has shape (num_tokens, 2 * intermediate_size)
+            # because silu_mul splits it into two halves
             input_tensor = torch.randn(
                 num_tokens,
                 2 * intermediate_size,
@@ -46,6 +81,7 @@ def generate_silu_mul_fp8_inputs() -> dict[str, tuple[Any, ...]]:
     return inputs
 
 
+@silu_mul_fp8.register_config_picker  # type: ignore[misc]
 def pick_silu_mul_fp8_config(
     args: tuple[Any, ...], config_keys: list[str]
 ) -> str | None:
@@ -92,41 +128,6 @@ def pick_silu_mul_fp8_config(
     return f"intermediate_{best_isize}_numtokens_{best_ntokens}"
 
 
-@register_kernel(
-    config_picker=pick_silu_mul_fp8_config,
-    input_generator=generate_silu_mul_fp8_inputs,
-)
-def silu_mul_fp8(input: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
-    original_shape = input.shape
-    two_d = hl.specialize(original_shape[-1])
-    d = two_d // 2
-    output_shape = original_shape[:-1] + (d,)
-
-    input_2d = input.view(-1, original_shape[-1])
-    m = input_2d.shape[0]
-
-    # TODO(gmagogsfm): Support for more float8 subtypes (e4m3fnuz, e5m2) coming
-    out = torch.empty((m, d), device=input.device, dtype=torch.float8_e4m3fn)
-
-    input_part_a = input_2d[:, :d]
-    input_part_b = input_2d[:, d:]
-
-    assert scale.numel() == 1, "Scale must be a scalar Tensor"
-
-    for tile_m, tile_n in hl.tile([m, d]):
-        a_vals = input_part_a[tile_m, tile_n]
-        silu_result = torch.nn.functional.silu(a_vals)
-        b_vals = input_part_b[tile_m, tile_n]
-        result = silu_result * b_vals
-        result_f32 = result.to(torch.float32)
-        scale_val = hl.load(scale, [0])
-        inv_scale = 1.0 / scale_val
-        result_scaled = result_f32 * inv_scale
-        out[tile_m, tile_n] = result_scaled.to(out.dtype)
-
-    return out.view(output_shape)
-
-
 def silu_mul_fp8_baseline(input: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
     output_shape = input.shape[:-1] + (input.shape[-1] // 2,)
     out = torch.empty(output_shape, dtype=torch.float8_e4m3fn, device=input.device)

vllm/kernels/helion/register.py

@@ -37,7 +37,7 @@ Key Classes
 """
 
 from collections.abc import Callable
-from typing import Any, cast
+from typing import Any, cast, overload
 
 import torch
 from torch.library import Library
@@ -95,7 +95,7 @@ def validate_helion_settings(
         raise ValueError(
             f"HelionKernelWrapper for '{op_name}' uses a custom autotuner via "
             f"config picker. Remove 'autotuner_fn' from helion_settings and use "
-            f"register_kernel(..., config_picker=...) instead."
+            f"@{op_name}.register_config_picker instead."
         )
 
     if settings_dict.get("static_shapes") is True:
@@ -169,7 +169,7 @@ class ConfiguredHelionKernel:
         if self.config_picker is None:
             raise RuntimeError(
                 f"No config picker registered for kernel '{self.op_name}'. "
-                f"A config_picker must be provided to register_kernel()."
+                f"Use @{self.op_name}.register_config_picker to register one."
             )
 
         # After None check, config_picker is guaranteed to be non-None
@@ -215,7 +215,7 @@ class ConfiguredHelionKernel:
         from vllm.kernels.helion.utils import get_canonical_gpu_name
 
         self.platform = get_canonical_gpu_name()
-        config_manager = ConfigManager()
+        config_manager = ConfigManager.get_instance()
         self.configs = config_manager.get_platform_configs(self.op_name, self.platform)
 
         if not self.configs:
@@ -253,9 +253,7 @@ class HelionKernelWrapper:
         raw_kernel_func: Callable,
         op_name: str,
         fake_impl: Callable,
-        config_picker: Callable[[tuple[Any, ...], list[str]], str | None],
         helion_settings: "helion.Settings | None" = None,
-        input_generator: Callable[[], dict[str, tuple[Any, ...]]] | None = None,
     ):
         # Validate helion_settings doesn't conflict with our custom autotuner
         validate_helion_settings(helion_settings, op_name)
@@ -264,43 +262,23 @@ class HelionKernelWrapper:
         self.op_name = op_name
         self._fake_impl = fake_impl
         self.helion_settings = helion_settings
-        self._config_picker = config_picker
-        self._input_generator = input_generator
+        self._config_picker: (
+            Callable[[tuple[Any, ...], list[str]], str | None] | None
+        ) = None
         self._configured_kernel: ConfiguredHelionKernel | None = None
-        # TODO(@gmagogsfm): Remove this disable flag once integrated with vLLM IR,
-        # which handles op enablement/disablement.
-        self._disabled = False
-        self._disabled_reason: str | None = None
-
-        try:
-            if not _HOP_AVAILABLE:
-                self._get_or_register_custom_op()
-            else:
-                self.get_configured_op()
-        except ValueError as e:
-            self._disabled = True
-            self._disabled_reason = str(e)
-            logger.warning(
-                "Helion kernel '%s' is disabled: %s",
-                op_name,
-                self._disabled_reason,
-            )
+        self._input_generator: Callable[[], dict[str, tuple[Any, ...]]] | None = None
 
     def __call__(self, *args, **kwargs):
-        if self._disabled:
-            raise RuntimeError(
-                f"Helion kernel '{self.op_name}' is disabled: {self._disabled_reason}"
-            )
+        # CustomOp fallback: register as torch custom op for torch.compile
+        # compatibility on older PyTorch lacking HOP/EffectType support
         if not _HOP_AVAILABLE:
-            op = getattr(torch.ops.vllm_helion, self.op_name)
-            return op(*args, **kwargs)
-        assert self._configured_kernel is not None, (
-            f"Kernel '{self.op_name}' was not initialized. "
-            "Please open an issue on GitHub."
-        )
+            custom_op = self._get_or_register_custom_op()
+            return custom_op(*args, **kwargs)
+        # HOP tracing: record HigherOrderOp in the FX graph
         if get_proxy_mode() is not None:
             return self._call_via_hop(args, kwargs)
-        return self._configured_kernel(*args, **kwargs)
+        # Eager: run the configured kernel directly
+        return self.get_configured_op()(*args, **kwargs)
 
     def _call_via_hop(
         self,
@@ -368,11 +346,42 @@ class HelionKernelWrapper:
                 constant_args[name] = val
         return constant_args, tensor_args
 
+    def register_config_picker(
+        self, picker_func: Callable[[tuple[Any, ...], list[str]], str | None]
+    ) -> Callable[[tuple[Any, ...], list[str]], str | None]:
+        self._config_picker = picker_func
+        return picker_func
+
+    def register_input_generator(
+        self, generator_func: Callable[[], dict[str, tuple[Any, ...]]]
+    ) -> Callable[[], dict[str, tuple[Any, ...]]]:
+        """
+        Register a function to generate inputs for autotuning and benchmarking.
+
+        Args:
+            generator_func: Function that returns dict[str, tuple] where:
+                - key: Configuration identifier (e.g., "4096", "hidden_4096")
+                - value: Tuple of arguments to pass to the kernel
+
+        Returns:
+            The registered function (for decorator usage)
+
+        Example:
+            @kernel_wrapper.register_input_generator
+            def generate_inputs():
+                return {
+                    "4096": (torch.randn(4096, device="cuda"), 0.5),
+                    "8192": (torch.randn(8192, device="cuda"), 0.5),
+                }
+        """
+        self._input_generator = generator_func
+        return generator_func
+
     def get_inputs(self) -> dict[str, tuple[Any, ...]]:
         if self._input_generator is None:
             raise NotImplementedError(
                 f"No input generator registered for kernel '{self.op_name}'. "
-                f"Use register_kernel(..., input_generator=...) to register one."
+                f"Use @{self.op_name}.register_input_generator to register one."
             )
         return self._input_generator()
 
@@ -392,10 +401,11 @@ class HelionKernelWrapper:
         return autotune_kernel.autotune(inputs)
 
     def get_configured_op(self) -> ConfiguredHelionKernel:
-        if self._disabled:
-            raise RuntimeError(
-                f"Helion kernel '{self.op_name}' is disabled: {self._disabled_reason}"
-            )
+        assert self._config_picker is not None, (
+            f"No config picker registered for kernel '{self.op_name}'. "
+            f"Use @{self.op_name}.register_config_picker to register one."
+        )
+
         if self._configured_kernel is None:
             self._configured_kernel = ConfiguredHelionKernel(
                 op_name=self.op_name,
@@ -403,6 +413,7 @@ class HelionKernelWrapper:
                 raw_kernel_func=self.raw_kernel_func,
                 helion_settings=self.helion_settings,
             )
+
         return self._configured_kernel
 
     def _get_or_register_custom_op(self) -> Any:
@@ -455,51 +466,45 @@ def infer_fake_impl(
     return helion_fake_kernel
 
 
+# Overloads are necessary for proper mypy type inference.
+# Without overloads, the union return type HelionKernelWrapper | Callable[...]
+# causes mypy to complain about missing attributes when tests do:
+#   wrapper = register_kernel(func)  # Should return HelionKernelWrapper
+#   wrapper._fake_impl  # mypy error: "Callable has no attribute _fake_impl"
+# The overloads tell mypy the exact return type based on the argument pattern.
+@overload
 def register_kernel(
-    op_name: str | None = None,
+    op_name_or_func: Callable,
     *,
-    config_picker: Callable[[tuple[Any, ...], list[str]], str | None],
     fake_impl: Callable | None = None,
     helion_settings: "helion.Settings | None" = None,
-    input_generator: Callable[[], dict[str, tuple[Any, ...]]] | None = None,
-) -> Callable[[Callable], HelionKernelWrapper]:
-    """Register a Helion kernel with pre-tuned config selection.
+) -> HelionKernelWrapper: ...
 
-    Wraps the kernel function in a HelionKernelWrapper that eagerly builds
-    the configured kernel and (on older PyTorch) registers a custom op.
 
-    Args:
-        config_picker: Required. Function with signature
-            ``(args: tuple, config_keys: list[str]) -> str | None``
-            that picks the best config key from available options.
-            Return ``None`` to fall back to ``"default"``.
+@overload
+def register_kernel(
+    op_name_or_func: str | None = None,
+    *,
+    fake_impl: Callable | None = None,
+    helion_settings: "helion.Settings | None" = None,
+) -> Callable[[Callable], HelionKernelWrapper]: ...
 
-            Example::
 
-                def pick_config(args, config_keys):
-                    x = args[0]
-                    hidden_size = x.shape[-1]
-                    batch_size = x.shape[0]
-                    for key in config_keys:
-                        if key == f"hiddensize_{hidden_size}_batchsize_{batch_size}":
-                            return key
-                    return "default" if "default" in config_keys else None
+def register_kernel(
+    op_name_or_func: str | Callable | None = None,
+    *,
+    fake_impl: Callable | None = None,
+    helion_settings: "helion.Settings | None" = None,
+) -> HelionKernelWrapper | Callable[[Callable], HelionKernelWrapper]:
+    """
+    Decorator to register a Helion kernel function as a HelionKernelWrapper.
 
-        input_generator: Optional. Function that returns
-            ``dict[str, tuple]`` where each key is a configuration
-            identifier (e.g. ``"4096"``, ``"hidden_4096"``) and each
-            value is a tuple of arguments to pass to the kernel.
-
-            Example::
-
-                def generate_inputs():
-                    return {
-                        "4096": (torch.randn(4096, device="cuda"), 0.5),
-                        "8192": (torch.randn(8192, device="cuda"), 0.5),
-                    }
+    Wraps the raw kernel function in a HelionKernelWrapper and registers it
+    in the global kernel registry. Auto-generates fake_impl if not provided.
     """
 
     def decorator(kernel_func: Callable) -> HelionKernelWrapper:
+        op_name = op_name_or_func if isinstance(op_name_or_func, str) else None
         final_op_name = op_name if op_name else kernel_func.__name__
 
         if final_op_name in _REGISTERED_KERNELS:
@@ -520,9 +525,7 @@ def register_kernel(
             raw_kernel_func=kernel_func,
             op_name=final_op_name,
             fake_impl=final_fake_impl,
-            config_picker=config_picker,
             helion_settings=helion_settings,
-            input_generator=input_generator,
         )
 
         _REGISTERED_KERNELS[final_op_name] = kernel_wrapper
@@ -534,4 +537,9 @@ def register_kernel(
 
         return kernel_wrapper
 
-    return decorator
+    if callable(op_name_or_func) and not isinstance(op_name_or_func, str):
+        # Bare decorator usage: @register_kernel
+        return decorator(op_name_or_func)
+    else:
+        # Decorator with arguments: @register_kernel(...)
+        return decorator

vllm/lora/model_manager.py

@@ -161,9 +161,9 @@ class LoRAModelManager:
             device=self.device,
             lora_config=self.lora_config,
         )
-
         lm_prefix = self.mm_mapping.language_model[0]
         self.punica_wrapper_mapping[lm_prefix] = llm_punica_wrapper
+
         if self.lora_config.enable_tower_connector_lora:
             self.supports_tower_connector_lora = self.supports_mm and hasattr(
                 self.model, "get_num_mm_encoder_tokens"
@@ -171,18 +171,6 @@ class LoRAModelManager:
         if not self.supports_tower_connector_lora:
             return
 
-        if (
-            vllm_config.model_config.multimodal_config
-            and vllm_config.model_config.multimodal_config.language_model_only
-        ):
-            if self.supports_tower_connector_lora:
-                logger.warning(
-                    "Disabling `enable_tower_connector_lora` because the multimodal "
-                    "model is configured to initialize the language model only."
-                )
-                self.supports_tower_connector_lora = False
-            return
-
         logger.warning(
             "LoRA for the tower and connector of multimodal models is "
             "experimental and may contain bugs. Please report any related issues on "

vllm/lora/ops/triton_ops/fused_moe_lora_fp8_op.py

@@ -10,10 +10,11 @@ from vllm.distributed import (
     tensor_model_parallel_all_gather,
     tensor_model_parallel_all_reduce,
 )
-from vllm.lora.ops.triton_ops.utils import supports_pdl
 from vllm.triton_utils import tl, triton
 from vllm.utils.torch_utils import direct_register_custom_op
 
+from .utils import supports_pdl
+
 
 @triton.jit
 def _get_lora_id(

vllm/model_executor/model_loader/default_loader.py

@@ -313,18 +313,7 @@ class DefaultModelLoader(BaseModelLoader):
         vllm_config = get_current_vllm_config()
         parallel_config = vllm_config.parallel_config
 
-        if not (
-            model_config.is_moe
-            and parallel_config.enable_expert_parallel
-            and parallel_config.enable_ep_weight_filter
-        ):
-            return
-
-        # When EPLB is enabled, redundant physical expert slots may map to
-        # logical experts that belong to other ranks in the default partition.
-        # The weight loader needs to see ALL logical expert weights so it can
-        # populate these redundant slots.  Skip the filter entirely.
-        if parallel_config.enable_eplb:
+        if not (model_config.is_moe and parallel_config.enable_expert_parallel):
             return
 
         num_experts = model_config.get_num_experts()

vllm/model_executor/model_loader/ep_weight_filter.py

@@ -73,9 +73,4 @@ def should_skip_weight(
     if eid is None:
         # Not an expert weight (dense / shared-expert / embedding) → keep.
         return False
-    # Only skip heavy weight tensors, never scale/metadata tensors.
-    # Scale tensors are tiny and some backends need them from ALL experts
-    # (e.g. FlashInfer NVFP4 computes a global max of activation scales).
-    if not weight_name.endswith(".weight"):
-        return False
     return eid not in local_expert_ids

vllm/model_executor/model_loader/weight_utils.py

@@ -856,8 +856,12 @@ def safetensors_weights_iterator(
                 for name in f.keys():  # noqa: SIM118
                     if should_skip_weight(name, local_expert_ids):
                         continue
-                    param = f.get_tensor(name)
-                    yield name, param
+                    # Yield a SafetensorsSlice that defers the disk
+                    # read.  Downstream weight loaders call .narrow()
+                    # to select only their TP shard; the actual I/O
+                    # happens at .materialize() time so we never read
+                    # data that will be discarded.
+                    yield name, SafetensorsSlice(f.get_slice(name))
 
 
 def multi_thread_safetensors_weights_iterator(
@@ -1155,6 +1159,141 @@ def gguf_quant_weights_iterator(
             yield name, param
 
 
+class SafetensorsSlice:
+    """Lazy wrapper around a safetensors ``PySafeSlice`` for deferred,
+    partial disk reads.
+
+    When tensor-parallel weight loading is active, each rank only needs a
+    *shard* of most weight tensors.  Without this wrapper every rank reads
+    the **full** tensor from disk and then calls ``torch.Tensor.narrow()``
+    to extract its shard, throwing the rest away.
+
+    ``SafetensorsSlice`` intercepts ``.narrow()`` calls and records them
+    without touching the disk.  The actual (partial) read happens only when
+    ``.materialize()`` is called – or when the object is used in a context
+    that requires a real tensor (e.g. ``.reshape()``, ``.item()``,
+    ``param_data.copy_(lazy_weight)``).
+
+    Because the safetensors format stores tensors contiguously, slicing on
+    the outer-most dimension (dim 0) translates to a contiguous sub-read
+    and avoids pulling the full tensor into host memory.
+    """
+
+    __slots__ = ("_slice", "_ndim", "_dims")
+
+    def __init__(self, safe_slice: Any) -> None:
+        original_shape = safe_slice.get_shape()
+        self._slice = safe_slice
+        self._ndim = len(original_shape)
+        # Per-dimension bookkeeping: (absolute_offset, current_size)
+        self._dims: list[tuple[int, int]] = [(0, s) for s in original_shape]
+
+    # -- Tensor-like API used by weight loaders --------------------------
+
+    @property
+    def shape(self) -> torch.Size:
+        return torch.Size(s for _, s in self._dims)
+
+    def size(self, dim: int | None = None):  # type: ignore[override]
+        if dim is not None:
+            return self._dims[dim][1]
+        return self.shape
+
+    def numel(self) -> int:
+        result = 1
+        for _, s in self._dims:
+            result *= s
+        return result
+
+    def __len__(self) -> int:
+        return self._dims[0][1]
+
+    def dim(self) -> int:
+        return self._ndim
+
+    # -- Lazy operations (no disk I/O) -----------------------------------
+
+    def narrow(self, dim: int, start: int, length: int) -> "SafetensorsSlice":
+        """Return a new slice narrowed on *dim* – zero-cost, no I/O."""
+        offset, size = self._dims[dim]
+        assert start + length <= size, (
+            f"narrow(dim={dim}, start={start}, length={length}) "
+            f"out of bounds for size {size}"
+        )
+        new = SafetensorsSlice.__new__(SafetensorsSlice)
+        new._slice = self._slice
+        new._ndim = self._ndim
+        new._dims = list(self._dims)
+        new._dims[dim] = (offset + start, length)
+        return new
+
+    # -- Materialisation (disk I/O happens here) -------------------------
+
+    def materialize(self) -> torch.Tensor:
+        """Read only the recorded sub-region from disk."""
+        indexing = tuple(slice(offset, offset + size) for offset, size in self._dims)
+        return self._slice[indexing]
+
+    # -- Fallback: auto-materialise for ops we cannot defer ---------------
+
+    def reshape(self, *args: Any) -> torch.Tensor:
+        return self.materialize().reshape(*args)
+
+    def view(self, *args: Any) -> torch.Tensor:
+        return self.materialize().view(*args)
+
+    def t(self) -> torch.Tensor:
+        return self.materialize().t()
+
+    def item(self) -> Any:
+        return self.materialize().item()
+
+    def to(self, *args: Any, **kwargs: Any) -> torch.Tensor:
+        return self.materialize().to(*args, **kwargs)
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return self.materialize().dtype
+
+    def __getitem__(self, key: Any) -> torch.Tensor:
+        return self.materialize().__getitem__(key)
+
+    @classmethod
+    def __torch_function__(cls, func, types, args=(), kwargs=None):
+        """Auto-materialise when used as an argument to any torch operation.
+
+        This ensures that ``param_data.copy_(safetensors_slice)`` works
+        transparently in all weight loaders without requiring any
+        downstream code changes.
+
+        The ``torch.narrow`` / ``torch.Tensor.narrow`` case is kept lazy
+        so that TP shard extraction still avoids unnecessary I/O.
+        """
+        kwargs = kwargs or {}
+        # Keep narrow lazy – delegate to the instance method.
+        if func is torch.narrow or func is torch.Tensor.narrow:
+            self_arg = args[0]
+            if isinstance(self_arg, SafetensorsSlice):
+                return self_arg.narrow(*args[1:], **(kwargs or {}))
+        new_args = tuple(
+            a.materialize() if isinstance(a, SafetensorsSlice) else a for a in args
+        )
+        return func(*new_args, **kwargs)
+
+    def __repr__(self) -> str:
+        return f"SafetensorsSlice(shape={list(self.shape)}, dims={self._dims})"
+
+
+def materialize_weight(
+    weight: "torch.Tensor | SafetensorsSlice",
+) -> torch.Tensor:
+    """Materialise a ``SafetensorsSlice`` into a real tensor (no-op on
+    tensors that are already materialised)."""
+    if isinstance(weight, SafetensorsSlice):
+        return weight.materialize()
+    return weight
+
+
 def convert_pyslice_to_tensor(x: Any) -> torch.Tensor:
     """convert PySafeSlice object from safetensors to torch.Tensor
 

vllm/model_executor/models/qwen3_5.py

@@ -180,16 +180,12 @@ class Qwen3_5GatedDeltaNet(Qwen3NextGatedDeltaNet):
         # ============================================================
         # Part 1: Input Projection
         # ============================================================
-        mixed_qkvz, ba = torch.ops.vllm.gdn_in_proj(
-            hidden_states,
-            self.in_proj_qkvz.weight.shape[0],
-            self.in_proj_ba.weight.shape[0],
-            self.prefix,
-        )
+        mixed_qkvz, _ = self.in_proj_qkvz(hidden_states)
         qkv_size = (self.key_dim * 2 + self.value_dim) // self.tp_size
         z_size = self.value_dim // self.tp_size
         mixed_qkv, z = mixed_qkvz.split([qkv_size, z_size], dim=-1)
         z = z.reshape(z.size(0), -1, self.head_v_dim)
+        ba, _ = self.in_proj_ba(hidden_states)
         b, a = ba.chunk(2, dim=-1)
 
         b = b.contiguous()

vllm/model_executor/models/qwen3_next.py

@@ -82,11 +82,7 @@ from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
 from vllm.transformers_utils.configs import Qwen3NextConfig
 from vllm.triton_utils import tl, triton
-from vllm.utils.multi_stream_utils import maybe_execute_in_parallel
-from vllm.utils.torch_utils import (
-    aux_stream,
-    direct_register_custom_op,
-)
+from vllm.utils.torch_utils import direct_register_custom_op
 from vllm.v1.attention.backend import AttentionMetadata
 from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
 
@@ -423,12 +419,6 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
         self.act = ACT2FN[config.hidden_act]
         self.layer_norm_epsilon = config.rms_norm_eps
         self.prefix = prefix
-        self.aux_stream = aux_stream()
-        self.events = (
-            [torch.cuda.Event(), torch.cuda.Event()]
-            if current_platform.is_cuda()
-            else [None, None]
-        )
 
         self.config = config
         self.model_config = model_config
@@ -657,12 +647,8 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
         # ============================================================
         # Part 1: Input Projection
         # ============================================================
-        projected_states_qkvz, projected_states_ba = torch.ops.vllm.gdn_in_proj(
-            hidden_states,
-            self.in_proj_qkvz.weight.shape[0],
-            self.in_proj_ba.weight.shape[0],
-            self.prefix,
-        )
+        projected_states_qkvz, _ = self.in_proj_qkvz(hidden_states)
+        projected_states_ba, _ = self.in_proj_ba(hidden_states)
         query, key, value, z, b, a = self.fix_query_key_value_ordering(
             projected_states_qkvz, projected_states_ba
         )
@@ -797,18 +783,6 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
 
         torch.accelerator.empty_cache()
 
-    def _forward_in_proj(
-        self, hidden_states: torch.Tensor
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        projected_states_qkvz, projected_states_ba = maybe_execute_in_parallel(
-            lambda: self.in_proj_qkvz(hidden_states)[0],
-            lambda: self.in_proj_ba(hidden_states)[0],
-            self.events[0],
-            self.events[1],
-            self.aux_stream,
-        )
-        return projected_states_qkvz, projected_states_ba
-
     def _forward_core(
         self,
         mixed_qkv: torch.Tensor,
@@ -1696,32 +1670,6 @@ class Qwen3NextForCausalLM(
         return self.model.get_expert_mapping()
 
 
-def gdn_in_proj(
-    hidden_states: torch.Tensor,
-    qkvz_output_size: int,
-    ba_output_size: int,
-    layer_name: str,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """
-    Custom op for the input projection.
-    """
-    forward_context: ForwardContext = get_forward_context()
-    self = forward_context.no_compile_layers[layer_name]
-    return self._forward_in_proj(hidden_states)
-
-
-def gdn_in_proj_fake(
-    hidden_states: torch.Tensor,
-    qkvz_output_size: int,
-    ba_output_size: int,
-    layer_name: str,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """Fake implementation for torch.compile."""
-    return hidden_states.new_empty(
-        hidden_states.shape[0], qkvz_output_size
-    ), hidden_states.new_empty(hidden_states.shape[0], ba_output_size)
-
-
 def gdn_attention_core(
     mixed_qkv: torch.Tensor,
     b: torch.Tensor,
@@ -1755,12 +1703,6 @@ def gdn_attention_core_fake(
     return
 
 
-direct_register_custom_op(
-    op_name="gdn_in_proj",
-    op_func=gdn_in_proj,
-    fake_impl=gdn_in_proj_fake,
-)
-
 direct_register_custom_op(
     op_name="gdn_attention_core",
     op_func=gdn_attention_core,

vllm/parser/abstract_parser.py

@@ -155,9 +155,7 @@ class Parser:
     @abstractmethod
     def extract_response_outputs(
         self,
-        *,
         model_output: str,
-        model_output_token_ids: Sequence[int],
         request: ResponsesRequest,
         enable_auto_tools: bool = False,
         tool_call_id_type: str = "random",
@@ -172,7 +170,6 @@ class Parser:
 
         Args:
             model_output: The complete model-generated string.
-            model_output_token_ids: The token IDs of the model output.
             request: The request object used to generate the output.
             enable_auto_tools: Whether to enable automatic tool call parsing.
             tool_call_id_type: Type of tool call ID generation ("random", etc).
@@ -316,9 +313,7 @@ class DelegatingParser(Parser):
 
     def extract_response_outputs(
         self,
-        *,
         model_output: str,
-        model_output_token_ids: Sequence[int],
         request: ResponsesRequest,
         enable_auto_tools: bool = False,
         tool_call_id_type: str = "random",

vllm/platforms/cpu.py

@@ -246,7 +246,6 @@ class CpuPlatform(Platform):
                     "size_asserts": False,
                     "nan_asserts": False,
                     "epilogue_fusion": True,
-                    "cpp.dynamic_threads": True,
                 }
             )
 

vllm/tool_parsers/glm47_moe_tool_parser.py

@@ -1,16 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-GLM-4.7 Tool Call Parser.
 
-GLM-4.7 uses a slightly different tool call format compared to GLM-4.5:
-  - The function name may appear on the same line as ``<tool_call>`` without
-    a newline separator before the first ``<arg_key>``.
-  - Tool calls may have zero arguments
-    (e.g. ``<tool_call>func</tool_call>``).
-
-This parser overrides the parent regex patterns to handle both formats.
-"""
 
 import regex as re
 
@@ -24,14 +14,10 @@ logger = init_logger(__name__)
 class Glm47MoeModelToolParser(Glm4MoeModelToolParser):
     def __init__(self, tokenizer: TokenizerLike):
         super().__init__(tokenizer)
-        # GLM-4.7 format: <tool_call>func_name[<arg_key>...]*</tool_call>
-        # The function name can be followed by a newline, whitespace, or
-        # directly by <arg_key> tags (no separator).  The arg section is
-        # optional so that zero-argument calls are supported.
         self.func_detail_regex = re.compile(
-            r"<tool_call>\s*(\S+?)\s*(<arg_key>.*)?</tool_call>", re.DOTALL
+            r"<tool_call>(.*?)(<arg_key>.*?)?</tool_call>", re.DOTALL
         )
         self.func_arg_regex = re.compile(
-            r"<arg_key>(.*?)</arg_key>\s*<arg_value>(.*?)</arg_value>",
+            r"<arg_key>(.*?)</arg_key>(?:\\n|\s)*<arg_value>(.*?)</arg_value>",
             re.DOTALL,
         )

vllm/tool_parsers/glm4_moe_tool_parser.py

@@ -206,12 +206,7 @@ class Glm4MoeModelToolParser(ToolParser):
             )
         else:
             if len(tool_calls) > 0:
-                content: str | None = model_output[
-                    : model_output.find(self.tool_calls_start_token)
-                ]
-                # Normalize empty/whitespace-only content to None
-                if not content or not content.strip():
-                    content = None
+                content = model_output[: model_output.find(self.tool_calls_start_token)]
                 return ExtractedToolCallInformation(
                     tools_called=True, tool_calls=tool_calls, content=content
                 )

vllm/utils/multi_stream_utils.py

@@ -1,48 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from collections.abc import Callable
-from typing import Any
-
-import torch
-
-
-def maybe_execute_in_parallel(
-    fn0: Callable[[], Any],
-    fn1: Callable[[], Any],
-    event0: torch.cuda.Event,
-    event1: torch.cuda.Event,
-    aux_stream: torch.cuda.Stream | None = None,
-) -> tuple[Any, Any]:
-    """Run two functions potentially in parallel on separate CUDA streams.
-
-    When aux_stream is provided, fn0 runs on the current (default) stream and
-    fn1 runs on aux_stream, synchronized via CUDA events.  When aux_stream is
-    None, both functions execute sequentially on the current stream.
-
-    This design follows TensorRT-LLM's maybe_execute_in_parallel pattern
-    (tensorrt_llm/_torch/modules/multi_stream_utils.py).
-
-    Args:
-        fn0: Callable for the default stream.
-        fn1: Callable for the auxiliary stream.
-        event0: CUDA event recorded before fn0 so aux_stream can wait.
-        event1: CUDA event recorded after fn1 so default stream can wait.
-        aux_stream: The second CUDA stream for fn1.
-            Multi-stream is disabled when aux_stream is None.
-
-    Returns:
-        Tuple of (fn0_result, fn1_result).
-    """
-    if aux_stream is not None:
-        event0.record()
-        result0 = fn0()
-        with torch.cuda.stream(aux_stream):
-            event0.wait()
-            result1 = fn1()
-            event1.record()
-        event1.wait()
-    else:
-        result0 = fn0()
-        result1 = fn1()
-    return (result0, result1)

vllm/v1/attention/backends/rocm_aiter_unified_attn.py

@@ -125,7 +125,6 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
         from aiter.ops.triton.unified_attention import unified_attention
 
         self.unified_attention = unified_attention
-        self.supports_quant_query_input = True
 
     def forward(
         self,
@@ -191,20 +190,12 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
 
         key_cache, value_cache = kv_cache.unbind(0)
 
-        softmax_scale = self.scale
-        fp8_post_attn_v_rescale = False
         if self.kv_cache_dtype.startswith("fp8"):
             key_cache = key_cache.view(self.fp8_dtype)
             value_cache = value_cache.view(self.fp8_dtype)
-            # When Q is FP8, triton kernel skips K/V dequant (for fp8xfp8 matmul).
-            # Compensate by absorbing q_scale and k_scale into softmax_scale, and
-            # v_scale into output_scale (or post-multiplying if no fusion).
-            if query.dtype == self.fp8_dtype:
-                softmax_scale = self.scale * layer._q_scale_float * layer._k_scale_float
-                if output_scale is not None:
-                    output_scale = output_scale / layer._v_scale_float
-                else:
-                    fp8_post_attn_v_rescale = True
+            assert layer._q_scale_float == 1.0, (
+                "A non 1.0 q_scale is not currently supported."
+            )
 
         cu_seqlens_q = attn_metadata.query_start_loc
         seqused_k = attn_metadata.seq_lens
@@ -226,22 +217,19 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
             max_seqlen_q=max_seqlen_q,
             seqused_k=seqused_k,
             max_seqlen_k=max_seqlen_k,
-            softmax_scale=softmax_scale,
+            softmax_scale=self.scale,
             causal=True,
             alibi_slopes=self.alibi_slopes,
             window_size=self.sliding_window,
             block_table=block_table,
             softcap=self.logits_soft_cap,
-            q_descale=None,  # q_scale absorbed into softmax_scale
+            q_descale=None,  # Not supported
             k_descale=layer._k_scale.expand(descale_shape),
             v_descale=layer._v_scale.expand(descale_shape),
             sinks=self.sinks,
             output_scale=output_scale,
         )
 
-        if fp8_post_attn_v_rescale:
-            output[:num_actual_tokens].mul_(layer._v_scale_float)
-
         return output
 
     def do_kv_cache_update(

vllm/v1/executor/abstract.py

@@ -353,13 +353,6 @@ class Executor(ABC):
     ) -> None:
         raise NotImplementedError
 
-    @classmethod
-    def supports_async_scheduling(cls) -> bool:
-        """
-        Whether the executor supports async scheduling.
-        """
-        return False
-
 
 from vllm.v1.executor.uniproc_executor import (  # noqa: E402
     ExecutorWithExternalLauncher as _ExecutorWithExternalLauncher,

vllm/v1/executor/multiproc_executor.py

@@ -487,10 +487,6 @@ class MultiprocExecutor(Executor):
             * self.parallel_config.prefill_context_parallel_size
         )
 
-    @classmethod
-    def supports_async_scheduling(cls) -> bool:
-        return True
-
 
 @dataclass
 class UnreadyWorkerProcHandle:

vllm/v1/executor/uniproc_executor.py

@@ -134,10 +134,6 @@ class UniProcExecutor(Executor):
         if worker := self.driver_worker:
             worker.shutdown()
 
-    @classmethod
-    def supports_async_scheduling(cls) -> bool:
-        return True
-
 
 class ExecutorWithExternalLauncher(UniProcExecutor):
     """An executor that uses external launchers to launch engines,

vllm/v1/worker/gpu/kv_connector.py

@@ -63,10 +63,11 @@ class ActiveKVConnector(KVConnector):
         if self._disabled:
             return
 
+        if scheduler_output.preempted_req_ids:
+            self.kv_connector.handle_preemptions(scheduler_output.preempted_req_ids)
         kv_connector_metadata = scheduler_output.kv_connector_metadata
         assert kv_connector_metadata is not None
         self.kv_connector.bind_connector_metadata(kv_connector_metadata)
-        self.kv_connector.handle_preemptions(kv_connector_metadata)
 
         # TODO: sort out KV Connectors' use of forward_context
         if is_forward_context_available():

vllm/v1/worker/gpu_model_runner.py

@@ -3594,10 +3594,10 @@ class GPUModelRunner(
                 scheduled_spec_decode_tokens=spec_decode_tokens_copy,
             )
 
-        if has_kv_transfer_group():
-            kv_connector_metadata = scheduler_output.kv_connector_metadata
-            assert kv_connector_metadata is not None
-            get_kv_transfer_group().handle_preemptions(kv_connector_metadata)
+        if scheduler_output.preempted_req_ids and has_kv_transfer_group():
+            get_kv_transfer_group().handle_preemptions(
+                scheduler_output.preempted_req_ids
+            )
 
         num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
         with (