From 18a7db87b0693be14e7138e723df25bd9138754d Mon Sep 17 00:00:00 2001 From: Iyed Sebai Date: Tue, 22 Sep 2026 20:55:23 +0100 Subject: [PATCH 1/3] [Bugfix][CPU] Install spec decode and sampling Triton fallbacks on dispatchers `CPUModelRunner._postprocess_triton` swaps Triton kernels for their C++ CPU implementations when Triton is unavailable, which is the production path on arm64 and macOS (`docker/Dockerfile.cpu` builds triton-cpu only for amd64). #56323 moved the sampling and DFlash kernels behind `triton_kernel_dispatcher_with_warmup`, which captures the kernel on the dispatcher instance at import time and launches `self.kernel`. Three of the patches were updated to assign `.kernel`; six were left rebinding the module-level kernel name, which nothing reads at call time any more, so those fallbacks were never installed. Point the remaining six at the dispatchers, matching the three that were already correct. The `dflash` module `setattr` block is dropped because `dflash.py` imports and calls the dispatcher rather than the kernel name; it was the only user of `import sys` in this file. Add a regression test covering all nine dispatchers, which fails on the unpatched module-level assignments. Signed-off-by: Iyed Sebai --- tests/v1/worker/test_cpu_triton_fallbacks.py | 85 ++++++++++++++++++++ vllm/v1/worker/cpu_model_runner.py | 21 ++--- 2 files changed, 91 insertions(+), 15 deletions(-) create mode 100644 tests/v1/worker/test_cpu_triton_fallbacks.py diff --git a/tests/v1/worker/test_cpu_triton_fallbacks.py b/tests/v1/worker/test_cpu_triton_fallbacks.py new file mode 100644 index 000000000000..610ae557fbbc --- /dev/null +++ b/tests/v1/worker/test_cpu_triton_fallbacks.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU Triton fallbacks must land on the dispatchers that launch the kernels. + +`CPUModelRunner._postprocess_triton` swaps Triton kernels for C++ CPU +implementations when Triton is unavailable (arm64, macOS). Kernels wrapped by +`triton_kernel_dispatcher_with_warmup` are captured on the dispatcher instance +at import time, so rebinding the module-level kernel name never reaches the +launch path and the fallback is silently not installed. +""" + +import pytest + +import vllm.utils.cpu_triton_utils as cpu_tl +import vllm.v1.sample.rejection_sampler as rejection_sampler +import vllm.v1.spec_decode.utils as spec_decode_utils +from vllm.v1.worker.cpu_model_runner import CPUModelRunner + +# (module, dispatcher attribute, expected CPU fallback in cpu_triton_utils) +DISPATCHER_FALLBACKS = [ + ( + spec_decode_utils, + "_eagle_prepare_inputs_padded", + "eagle_prepare_inputs_padded_kernel", + ), + ( + spec_decode_utils, + "_eagle_prepare_next_token_padded", + "eagle_prepare_next_token_padded_kernel", + ), + ( + spec_decode_utils, + "_copy_and_expand_eagle_inputs", + "copy_and_expand_eagle_inputs_kernel", + ), + ( + spec_decode_utils, + "_copy_and_expand_dflash_inputs", + "copy_and_expand_dflash_inputs_kernel", + ), + ( + spec_decode_utils, + "_eagle_step_slot_mapping_metadata", + "eagle_step_slot_mapping_metadata_kernel", + ), + (rejection_sampler, "_rejection_greedy_sample", "rejection_greedy_sample_kernel"), + (rejection_sampler, "_rejection_random_sample", "rejection_random_sample_kernel"), + (rejection_sampler, "_expand", "expand_kernel"), + (rejection_sampler, "_sample_recovered_tokens", "sample_recovered_tokens_kernel"), +] + + +@pytest.fixture +def restore_patched_kernels(monkeypatch): + """Undo the global kernel swaps `_postprocess_triton` performs.""" + import vllm.v1.worker.block_table as block_table + import vllm.v1.worker.mamba_utils as mamba_utils + + for module, dispatcher_name, _ in DISPATCHER_FALLBACKS: + dispatcher = getattr(module, dispatcher_name) + monkeypatch.setattr(dispatcher, "kernel", dispatcher.kernel) + monkeypatch.setattr( + block_table._COMPUTE_SLOT_MAPPING_KERNEL, + "kernel", + block_table._COMPUTE_SLOT_MAPPING_KERNEL.kernel, + ) + monkeypatch.setattr( + mamba_utils, "batch_memcpy_kernel", mamba_utils.batch_memcpy_kernel + ) + + +@pytest.mark.parametrize( + "module, dispatcher_name, fallback_name", DISPATCHER_FALLBACKS +) +def test_postprocess_triton_installs_cpu_fallback( + monkeypatch, restore_patched_kernels, module, dispatcher_name, fallback_name +): + """Each dispatcher launches the CPU fallback once Triton is unavailable.""" + monkeypatch.setattr("vllm.triton_utils.HAS_TRITON", False) + + # `_postprocess_triton` reads no instance state, so it can run unbound. + CPUModelRunner._postprocess_triton(None) + + dispatcher = getattr(module, dispatcher_name) + assert dispatcher.kernel is getattr(cpu_tl, fallback_name) diff --git a/vllm/v1/worker/cpu_model_runner.py b/vllm/v1/worker/cpu_model_runner.py index 04ca2b9194e4..5bcf768fa5b0 100644 --- a/vllm/v1/worker/cpu_model_runner.py +++ b/vllm/v1/worker/cpu_model_runner.py @@ -1,6 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import sys from contextlib import AbstractContextManager, contextmanager from typing import Any @@ -91,28 +90,20 @@ def _postprocess_triton(self) -> None: spec_decode_utils._copy_and_expand_eagle_inputs.kernel = ( cpu_tl.copy_and_expand_eagle_inputs_kernel ) - spec_decode_utils.copy_and_expand_dflash_inputs_kernel = ( + spec_decode_utils._copy_and_expand_dflash_inputs.kernel = ( cpu_tl.copy_and_expand_dflash_inputs_kernel ) - dflash_module = sys.modules.get("vllm.v1.spec_decode.dflash") - if dflash_module is not None: - dflash_kernel_name = "copy_and_expand_dflash_inputs_kernel" - setattr( - dflash_module, - dflash_kernel_name, - cpu_tl.copy_and_expand_dflash_inputs_kernel, - ) - spec_decode_utils.eagle_step_slot_mapping_metadata_kernel = ( + spec_decode_utils._eagle_step_slot_mapping_metadata.kernel = ( cpu_tl.eagle_step_slot_mapping_metadata_kernel ) - vllm.v1.sample.rejection_sampler.rejection_greedy_sample_kernel = ( + vllm.v1.sample.rejection_sampler._rejection_greedy_sample.kernel = ( cpu_tl.rejection_greedy_sample_kernel ) - vllm.v1.sample.rejection_sampler.rejection_random_sample_kernel = ( + vllm.v1.sample.rejection_sampler._rejection_random_sample.kernel = ( cpu_tl.rejection_random_sample_kernel ) - vllm.v1.sample.rejection_sampler.expand_kernel = cpu_tl.expand_kernel - vllm.v1.sample.rejection_sampler.sample_recovered_tokens_kernel = ( + vllm.v1.sample.rejection_sampler._expand.kernel = cpu_tl.expand_kernel + vllm.v1.sample.rejection_sampler._sample_recovered_tokens.kernel = ( cpu_tl.sample_recovered_tokens_kernel ) From 25ce3ec26e3483df8278f1a6ad82d8473b8dbfc0 Mon Sep 17 00:00:00 2001 From: Iyed Sebai Date: Tue, 22 Sep 2026 23:11:43 +0100 Subject: [PATCH 2/3] [Bugfix][CPU] Match CPU fallback shim signatures to dispatch argument names The dispatcher binds each dispatch input to a kernel argument named `name`, `name_ptr`, `NAME` or `NAME_PTR` (`_launch_binding_plan`). Two CPU shims use names that match none of those, so the fallback raises `TypeError` for a missing positional argument once it is installed: - `_expand_kernel_impl` took `input_val`; the dispatch input is `input`. - `_eagle_step_slot_mapping_metadata_kernel_impl` took `stride`; the dispatch input is `block_table_stride`, and that dispatch passes no launch kwargs. Rename both to the dispatch argument names so the launch binds. Signed-off-by: Iyed Sebai --- vllm/utils/cpu_triton_utils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/vllm/utils/cpu_triton_utils.py b/vllm/utils/cpu_triton_utils.py index d1696317aab8..6f9b62c8a04c 100644 --- a/vllm/utils/cpu_triton_utils.py +++ b/vllm/utils/cpu_triton_utils.py @@ -119,7 +119,7 @@ def _eagle_prepare_next_token_padded_kernel_impl( def _eagle_step_slot_mapping_metadata_kernel_impl( positions, block_table, - stride, + block_table_stride, seq_lens, out_clamped_positions, out_slot_mapping, @@ -405,7 +405,7 @@ def _rejection_random_sample_kernel_impl( def _expand_kernel_impl( output, - input_val, + input, cu_num_tokens, replace_from, replace_to, @@ -413,7 +413,7 @@ def _expand_kernel_impl( ): torch.ops._C.expand_kernel_impl( _ensure_int64(output), - _ensure_int64(input_val), + _ensure_int64(input), _ensure_int64(cu_num_tokens), replace_from, replace_to, From f4f9520d9e1162725087c71b921f9f6203161db9 Mon Sep 17 00:00:00 2001 From: The-iyed Date: Wed, 23 Sep 2026 22:48:15 +0100 Subject: [PATCH 3/3] Format the CPU fallback regression test with ruff Signed-off-by: Iyed Sebai --- tests/v1/worker/test_cpu_triton_fallbacks.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/v1/worker/test_cpu_triton_fallbacks.py b/tests/v1/worker/test_cpu_triton_fallbacks.py index 610ae557fbbc..c93420bfc809 100644 --- a/tests/v1/worker/test_cpu_triton_fallbacks.py +++ b/tests/v1/worker/test_cpu_triton_fallbacks.py @@ -69,9 +69,7 @@ def restore_patched_kernels(monkeypatch): ) -@pytest.mark.parametrize( - "module, dispatcher_name, fallback_name", DISPATCHER_FALLBACKS -) +@pytest.mark.parametrize("module, dispatcher_name, fallback_name", DISPATCHER_FALLBACKS) def test_postprocess_triton_installs_cpu_fallback( monkeypatch, restore_patched_kernels, module, dispatcher_name, fallback_name ):