From ed0f4750f8766c2cc3984be4f1ba073a6915d4ed Mon Sep 17 00:00:00 2001 From: Uranus <109661872+UranusSeven@users.noreply.github.com> Date: Sun, 16 Aug 2026 02:59:22 +0800 Subject: [PATCH 001/839] [Bugfix][Model] Kimi-K3 MegaMoE: pass situ_beta/situ_linear_beta to fp8_fp4_mega_moe (#52445) Signed-off-by: UranusSeven <109661872+UranusSeven@users.noreply.github.com> Co-authored-by: Claude --- vllm/models/kimi_k3/nvidia/model.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index d7a08a6da2ae..f7e8c54f1bf3 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -486,8 +486,8 @@ def forward( symm_buffer, activation_clamp=activation_clamp, activation=self.activation, - activation_beta=self.activation_beta, - activation_linear_beta=self.activation_linear_beta, + situ_beta=self.activation_beta, + situ_linear_beta=self.activation_linear_beta, fast_math=fast_math, ) return y From c94cdd0ae03865c00b694ca6f13b2d2bc4360fcd Mon Sep 17 00:00:00 2001 From: Henry Su Date: Sat, 15 Aug 2026 16:13:02 -0500 Subject: [PATCH 002/839] [Bugfix][Sampling] Clear empty side on thinking-budget asymmetric SWAP (#49613) Signed-off-by: Henry Su --- tests/v1/sample/test_thinking_budget_state.py | 95 +++++++++++++++++++ vllm/v1/sample/thinking_budget_state.py | 4 +- 2 files changed, 97 insertions(+), 2 deletions(-) create mode 100644 tests/v1/sample/test_thinking_budget_state.py diff --git a/tests/v1/sample/test_thinking_budget_state.py b/tests/v1/sample/test_thinking_budget_state.py new file mode 100644 index 000000000000..943734e8eefd --- /dev/null +++ b/tests/v1/sample/test_thinking_budget_state.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for ThinkingBudgetStateHolder batch index moves.""" + +import torch + +from vllm.sampling_params import SamplingParams +from vllm.v1.sample.logits_processor.interface import ( + BatchUpdate, + MoveDirectionality, +) +from vllm.v1.sample.thinking_budget_state import ThinkingBudgetStateHolder + + +class _MockReasoningConfig: + reasoning_start_token_ids = [151667] + reasoning_end_token_ids = [151668] + + +def _make_holder() -> ThinkingBudgetStateHolder: + return ThinkingBudgetStateHolder( + _MockReasoningConfig(), + 8, + 0, + torch.device("cpu"), + False, + ) + + +def test_swap_budgeted_with_unbudgeted_clears_empty_side(): + """Asymmetric SWAP must not leave the empty index sharing state.""" + h = _make_holder() + h.sync_batch( + BatchUpdate( + batch_size=2, + removed=(), + added=[ + (0, SamplingParams(thinking_token_budget=5), None, []), + (1, SamplingParams(), None, []), + ], + moved=(), + ) + ) + assert list(h._state.keys()) == [0] + budget_state = h._state[0] + + h.sync_batch( + BatchUpdate( + batch_size=2, + removed=(), + added=(), + moved=[(0, 1, MoveDirectionality.SWAP)], + ) + ) + assert list(h._state.keys()) == [1] + assert h._state[1] is budget_state + assert h._state[1]["thinking_token_budget"] == 5 + + h.sync_batch( + BatchUpdate( + batch_size=2, + removed=(), + added=(), + moved=[(0, 1, MoveDirectionality.SWAP)], + ) + ) + assert list(h._state.keys()) == [0] + assert h._state[0] is budget_state + + +def test_swap_exchanges_two_budgeted_states(): + h = _make_holder() + h.sync_batch( + BatchUpdate( + batch_size=2, + removed=(), + added=[ + (0, SamplingParams(thinking_token_budget=3), None, []), + (1, SamplingParams(thinking_token_budget=7), None, []), + ], + moved=(), + ) + ) + b0 = h._state[0]["thinking_token_budget"] + b1 = h._state[1]["thinking_token_budget"] + h.sync_batch( + BatchUpdate( + batch_size=2, + removed=(), + added=(), + moved=[(0, 1, MoveDirectionality.SWAP)], + ) + ) + assert h._state[0]["thinking_token_budget"] == b1 + assert h._state[1]["thinking_token_budget"] == b0 diff --git a/vllm/v1/sample/thinking_budget_state.py b/vllm/v1/sample/thinking_budget_state.py index efac0111779e..43880bf656e5 100644 --- a/vllm/v1/sample/thinking_budget_state.py +++ b/vllm/v1/sample/thinking_budget_state.py @@ -100,8 +100,8 @@ def sync_batch(self, batch_update: BatchUpdate | None) -> None: for i1, i2, direction in batch_update.moved: if direction == MoveDirectionality.SWAP: - state1 = self._state.get(i1) - state2 = self._state.get(i2) + state1 = self._state.pop(i1, None) + state2 = self._state.pop(i2, None) if state1 is not None: self._state[i2] = state1 if state2 is not None: From fa9d67f7828e9bc105912ddf41dc384105732b1e Mon Sep 17 00:00:00 2001 From: omerpaz95 <73347585+omerpaz95@users.noreply.github.com> Date: Sun, 16 Aug 2026 01:40:38 +0300 Subject: [PATCH 003/839] [EC Connector] Added Build Connector Worker Meta for EC Connector (#49585) Signed-off-by: yewentao256 Signed-off-by: omerpaz95 Signed-off-by: omerpaz95 <73347585+omerpaz95@users.noreply.github.com> Co-authored-by: yewentao256 Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> --- .buildkite/test_areas/misc.yaml | 1 + .../unit/test_ec_output_aggregator.py | 115 ++++++++++++++++++ .../unit/test_worker_ec_connector.py | 77 ++++++++++++ tests/v1/executor/test_executor.py | 3 + .../unit/test_handshake_pp_aggregation.py | 1 + tests/v1/test_outputs.py | 19 ++- .../worker/test_gpu_model_runner_v2_eplb.py | 1 + .../ec_transfer/ec_connector/base.py | 33 +++++ .../ec_transfer/ec_connector/utils.py | 51 ++++++++ vllm/v1/engine/core.py | 2 + vllm/v1/executor/abstract.py | 5 + vllm/v1/executor/multiproc_executor.py | 26 +++- vllm/v1/executor/ray_executor.py | 13 ++ vllm/v1/outputs.py | 36 ++++++ vllm/v1/worker/gpu/ec_connector.py | 27 +++- vllm/v1/worker/gpu/model_runner.py | 36 ++++-- 16 files changed, 430 insertions(+), 16 deletions(-) create mode 100644 tests/v1/ec_connector/unit/test_ec_output_aggregator.py create mode 100644 tests/v1/ec_connector/unit/test_worker_ec_connector.py create mode 100644 vllm/distributed/ec_transfer/ec_connector/utils.py diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 8c80fead290f..126bc657048c 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -127,6 +127,7 @@ steps: - pytest -v -s v1/test_kv_cache_spec_registry.py - pytest -v -s v1/cudagraph/test_cudagraph_manager.py - pytest -v -s -m 'cpu_test' v1/kv_connector/unit + - pytest -v -s -m 'cpu_test' v1/ec_connector/unit - pytest -v -s -m 'cpu_test' v1/metrics - label: Extract Hidden States Integration diff --git a/tests/v1/ec_connector/unit/test_ec_output_aggregator.py b/tests/v1/ec_connector/unit/test_ec_output_aggregator.py new file mode 100644 index 000000000000..4d5b7dae1037 --- /dev/null +++ b/tests/v1/ec_connector/unit/test_ec_output_aggregator.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for ECOutputAggregator.""" + +import pytest + +from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorWorkerMetadata +from vllm.distributed.ec_transfer.ec_connector.utils import ECOutputAggregator +from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator +from vllm.v1.outputs import ( + EMPTY_MODEL_RUNNER_OUTPUT, + ECConnectorOutput, + KVConnectorOutput, + ModelRunnerOutput, +) + +pytestmark = pytest.mark.cpu_test + + +class FakeWorkerMeta(ECConnectorWorkerMetadata): + """Records merge order. `aggregate` returns a new object, as the base class + declares: an aggregator discarding the return value would lose the merge. + """ + + def __init__(self, saves: list[str]): + self.saves = saves + + def aggregate(self, other: "FakeWorkerMeta") -> "FakeWorkerMeta": + return FakeWorkerMeta(self.saves + other.saves) + + +def _worker_output(ec_output: ECConnectorOutput | None) -> ModelRunnerOutput: + return ModelRunnerOutput( + req_ids=[], req_id_to_index={}, ec_connector_output=ec_output + ) + + +def test_aggregate_folds_every_rank_onto_output_rank(): + """EC work done on any rank reaches the scheduler via output_rank's output. + + The middle rank reports no worker metadata: it must neither seed nor clobber + the accumulator. + """ + outputs = [ + _worker_output( + ECConnectorOutput( + finished_sending={"mm0"}, + ec_connector_worker_meta=FakeWorkerMeta(["mm0"]), + ) + ), + _worker_output(ECConnectorOutput(finished_recving={"mm1"})), + _worker_output( + ECConnectorOutput(ec_connector_worker_meta=FakeWorkerMeta(["mm2"])) + ), + ] + + result = ECOutputAggregator().aggregate(outputs, output_rank=2) + + assert result is outputs[2] + assert result.ec_connector_output.finished_sending == {"mm0"} + assert result.ec_connector_output.finished_recving == {"mm1"} + assert result.ec_connector_output.ec_connector_worker_meta.saves == ["mm0", "mm2"] + + +def test_aggregate_leaves_no_ec_output_when_no_worker_reported(): + """Empty per-worker reports must not reach the scheduler as an empty object.""" + outputs = [_worker_output(ECConnectorOutput()), _worker_output(ECConnectorOutput())] + + result = ECOutputAggregator().aggregate(outputs, output_rank=0) + + assert result is outputs[0] + assert result.ec_connector_output is None + assert ECOutputAggregator().aggregate([None], output_rank=0) is None + + +def test_aggregate_does_not_write_through_the_shared_empty_output(): + """A rank with nothing to report yields the shared empty output singleton. + + Folding another rank's metadata onto it must not write through to the + module-level object, which every later step would then carry. + """ + outputs = [ + _worker_output( + ECConnectorOutput(ec_connector_worker_meta=FakeWorkerMeta(["mm0"])) + ), + EMPTY_MODEL_RUNNER_OUTPUT, + ] + + result = ECOutputAggregator().aggregate(outputs, output_rank=1) + + assert EMPTY_MODEL_RUNNER_OUTPUT.ec_connector_output is None + assert result is not EMPTY_MODEL_RUNNER_OUTPUT + assert result.ec_connector_output.ec_connector_worker_meta.saves == ["mm0"] + + +def test_chaining_with_kv_aggregator_preserves_both_outputs(): + """MultiprocExecutor chains both aggregators and keeps only the last result, + so each must merge onto the same output_rank output rather than replace it. + """ + outputs = [ + _worker_output(ECConnectorOutput(finished_sending={"mm0"})), + _worker_output(None), + ] + outputs[1].kv_connector_output = KVConnectorOutput(invalid_block_ids={7}) + + result = None + for aggregator in ( + KVOutputAggregator(expected_finished_count=1), + ECOutputAggregator(), + ): + result = aggregator.aggregate(outputs, output_rank=1) + + assert result is outputs[1] + assert result.kv_connector_output.invalid_block_ids == {7} + assert result.ec_connector_output.finished_sending == {"mm0"} diff --git a/tests/v1/ec_connector/unit/test_worker_ec_connector.py b/tests/v1/ec_connector/unit/test_worker_ec_connector.py new file mode 100644 index 000000000000..3dcad1e50ae5 --- /dev/null +++ b/tests/v1/ec_connector/unit/test_worker_ec_connector.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the V2 GPU model runner's EC connector wrapper.""" + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +from vllm.distributed.ec_transfer.ec_connector.base import ( + ECConnectorBase, + ECConnectorMetadata, +) +from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT +from vllm.v1.worker.gpu.ec_connector import NO_OP_EC_CONNECTOR, ActiveECConnector + +pytestmark = pytest.mark.cpu_test + +WORKER_META = object() + + +def _scheduler_output() -> SimpleNamespace: + return SimpleNamespace( + ec_connector_metadata=ECConnectorMetadata(), finished_req_ids=frozenset() + ) + + +def _connector( + encoder_cache: dict | None = None, + is_producer: bool = True, + is_consumer: bool = False, +) -> tuple[ActiveECConnector, MagicMock]: + fake = MagicMock(spec=ECConnectorBase) + fake.is_producer = is_producer + fake.is_consumer = is_consumer + fake.get_finished.return_value = (None, None) + fake.build_connector_worker_meta.return_value = WORKER_META + with patch("vllm.v1.worker.gpu.ec_connector.get_ec_transfer", return_value=fake): + return ActiveECConnector(SimpleNamespace(), encoder_cache or {}), fake + + +@pytest.mark.parametrize( + ("is_producer", "is_consumer"), [(True, False), (True, True), (False, True)] +) +def test_saves_newly_added_caches_for_every_producer(is_producer, is_consumer): + """An ec_both node is also a producer: it must offload what it just computed.""" + encoder_cache = {"mm_old": None} + connector, fake = _connector(encoder_cache, is_producer, is_consumer) + + with connector.maybe_get_output(_scheduler_output()): + encoder_cache["mm_new"] = None + + saved = [call.kwargs["mm_hash"] for call in fake.save_caches.call_args_list] + assert saved == (["mm_new"] if is_producer else []) + assert fake.start_load_caches.called == is_consumer + + +def test_worker_meta_is_reported_on_context_exit(): + """Reported in the finally block, so is_empty() sees it only after the exit.""" + connector, fake = _connector() + + with connector.maybe_get_output(_scheduler_output()) as output: + assert output.ec_connector_worker_meta is None + + assert output.ec_connector_worker_meta is WORKER_META + assert fake.clear_connector_metadata.called + + +def test_no_forward_reports_without_running_the_model(): + connector, _ = _connector() + + output = connector.no_forward(_scheduler_output()) + + assert output.ec_connector_output.ec_connector_worker_meta is WORKER_META + + empty = NO_OP_EC_CONNECTOR.no_forward(_scheduler_output()) + assert empty is EMPTY_MODEL_RUNNER_OUTPUT diff --git a/tests/v1/executor/test_executor.py b/tests/v1/executor/test_executor.py index c529c3204d50..a21a8dafd328 100644 --- a/tests/v1/executor/test_executor.py +++ b/tests/v1/executor/test_executor.py @@ -9,6 +9,7 @@ import pytest +from vllm.distributed.ec_transfer.ec_connector.utils import ECOutputAggregator from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs from vllm.sampling_params import SamplingParams @@ -98,6 +99,7 @@ def collective_rpc( non_block: bool = False, unique_reply_rank: int | None = None, kv_output_aggregator: KVOutputAggregator = None, + ec_output_aggregator: ECOutputAggregator | None = None, ) -> Any | list[Any] | Future[Any | list[Any]]: # Drop marker to show that this was run with open(".marker", "w"): @@ -110,6 +112,7 @@ def collective_rpc( non_block, unique_reply_rank, kv_output_aggregator, + ec_output_aggregator, ) diff --git a/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py b/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py index 0c0f9f1f8998..e5cbd977b6d8 100644 --- a/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py +++ b/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py @@ -55,6 +55,7 @@ def _run_engine_core_handshake( class _FakeScheduler: def __init__(self, **kwargs: Any) -> None: self.connector = connector + self.ec_connector = None def get_kv_connector(self) -> KVConnectorBase_V1: return connector diff --git a/tests/v1/test_outputs.py b/tests/v1/test_outputs.py index 06fd9b409bdf..d93ae96e5438 100644 --- a/tests/v1/test_outputs.py +++ b/tests/v1/test_outputs.py @@ -6,7 +6,13 @@ import torch from vllm.platforms import current_platform -from vllm.v1.outputs import LogprobsLists, LogprobsTensors +from vllm.v1.outputs import ( + EMPTY_MODEL_RUNNER_OUTPUT, + ECConnectorOutput, + LogprobsLists, + LogprobsTensors, + ModelRunnerOutput, +) from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p from vllm.v1.worker.gpu.sample.output import SamplingMaskTensors @@ -200,3 +206,14 @@ def test_slice_all_requests(self): assert len(sliced.logprob_token_ids) == 9 # All tokens assert sliced.logprob_token_ids == self.logprobsLists.logprob_token_ids assert sliced.cu_num_generated_tokens is None + + +def test_with_ec_conn_output_copies_shared_empty_output(): + """The shared empty output is copied, never written to.""" + ec_output = ECConnectorOutput(finished_sending={"mm_hash"}) + + result = ModelRunnerOutput.with_ec_conn_output(EMPTY_MODEL_RUNNER_OUTPUT, ec_output) + + assert result is not EMPTY_MODEL_RUNNER_OUTPUT + assert result.ec_connector_output is ec_output + assert EMPTY_MODEL_RUNNER_OUTPUT.ec_connector_output is None diff --git a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py index fbeaa198d0d6..ebb4beb2a5a9 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py +++ b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py @@ -180,6 +180,7 @@ def test_v2_sample_tokens_runs_eplb_on_non_last_pp_rank(monkeypatch): hidden_states=None, aux_hidden_states=None, finished_req_ids=set(), + ec_connector_output=None, routed_experts=None, num_tokens_across_dp=None, ) diff --git a/vllm/distributed/ec_transfer/ec_connector/base.py b/vllm/distributed/ec_transfer/ec_connector/base.py index 3c20a4a1f749..9138203bf4dd 100644 --- a/vllm/distributed/ec_transfer/ec_connector/base.py +++ b/vllm/distributed/ec_transfer/ec_connector/base.py @@ -20,6 +20,8 @@ get_finished() - called with ids of finished requests, returns ids of requests that have completed async sending/recving. + build_connector_worker_meta() - builds metadata to be sent + back to the scheduler-side connector """ import enum @@ -56,6 +58,27 @@ class ECConnectorMetadata(ABC): # noqa: B024 pass +class ECConnectorWorkerMetadata(ABC): + """ + Abstract Metadata used to communicate back + Worker ECConnector -> Scheduler ECConnector. + + Each worker can output its own metadata. + For a single engine step, all metadata objects returned by workers + will be aggregated using the `aggregate` method below, before + being passed to the Scheduler ECConnector. + """ + + @abstractmethod + def aggregate( + self, other: "ECConnectorWorkerMetadata" + ) -> "ECConnectorWorkerMetadata": + """ + Aggregate metadata with another `ECConnectorWorkerMetadata` object. + """ + pass + + class ECConnectorBase(ABC): def __init__(self, vllm_config: "VllmConfig", role: ECConnectorRole): self._connector_metadata: ECConnectorMetadata | None = None @@ -190,6 +213,16 @@ def get_finished( """ return None, None + def build_connector_worker_meta(self) -> ECConnectorWorkerMetadata | None: + """ + Build the ECConnector worker metadata for this engine step. + + Returns: + ECConnectorWorkerMetadata: the worker metadata. + None if no worker metadata is available. + """ + return None + # ============================== # Scheduler-side methods # ============================== diff --git a/vllm/distributed/ec_transfer/ec_connector/utils.py b/vllm/distributed/ec_transfer/ec_connector/utils.py new file mode 100644 index 000000000000..f5f78e6c3c33 --- /dev/null +++ b/vllm/distributed/ec_transfer/ec_connector/utils.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""EC connector helper utilities.""" + +from vllm.v1.outputs import ECConnectorOutput, ModelRunnerOutput + + +class ECOutputAggregator: + """Merge every worker's EC connector output onto the single + ModelRunnerOutput that reaches the scheduler. + + Mirrors KVOutputAggregator: only `output_rank`'s output is returned to the + scheduler, but the EC connector may have run on any rank. + """ + + def aggregate( + self, outputs: list[ModelRunnerOutput | None], output_rank: int = 0 + ) -> ModelRunnerOutput | None: + output = outputs[output_rank] + if not output: + return None + + finished_sending = set[str]() + finished_recving = set[str]() + worker_meta = None + for model_runner_output in outputs: + assert model_runner_output is not None + ec_output = model_runner_output.ec_connector_output + if not ec_output: + continue + + finished_sending |= ec_output.finished_sending or set() + finished_recving |= ec_output.finished_recving or set() + + if meta := ec_output.ec_connector_worker_meta: + worker_meta = ( + meta if worker_meta is None else worker_meta.aggregate(meta) + ) + + aggregated = ECConnectorOutput( + finished_sending=finished_sending or None, + finished_recving=finished_recving or None, + ec_connector_worker_meta=worker_meta, + ) + if aggregated.is_empty(): + output.ec_connector_output = None + return output + + # `output` is the shared empty output whenever `output_rank` had no work, + # so attach through the copy-on-write helper. + return ModelRunnerOutput.with_ec_conn_output(output, aggregated) diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 55fe5b90c587..295c66b360ed 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -173,6 +173,8 @@ def __init__( ) if self.scheduler.connector is not None: # type: ignore self.model_executor.init_kv_output_aggregator(self.scheduler.connector) # type: ignore + if self.scheduler.ec_connector is not None: # type: ignore + self.model_executor.init_ec_output_aggregator() mm_registry = MULTIMODAL_REGISTRY self.mm_receiver_cache = mm_registry.engine_receiver_cache_from_config( diff --git a/vllm/v1/executor/abstract.py b/vllm/v1/executor/abstract.py index 864989e6ac17..6fff5f032d10 100644 --- a/vllm/v1/executor/abstract.py +++ b/vllm/v1/executor/abstract.py @@ -9,6 +9,7 @@ import vllm.envs as envs from vllm.config import VllmConfig +from vllm.distributed.ec_transfer.ec_connector.utils import ECOutputAggregator from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorHandshakeMetadata, @@ -110,6 +111,7 @@ def __init__( self.is_sleeping = False self.sleeping_tags: set[str] = set() self.kv_output_aggregator: KVOutputAggregator | None = None + self.ec_output_aggregator: ECOutputAggregator | None = None @abstractmethod def _init_executor(self) -> None: @@ -283,6 +285,9 @@ def init_kv_output_aggregator(self, connector: "KVConnectorBase") -> None: connector, self.parallel_config.world_size ) + def init_ec_output_aggregator(self) -> None: + self.ec_output_aggregator = ECOutputAggregator() + @cached_property # Avoid unnecessary RPC calls def supported_tasks(self) -> tuple[SupportedTask, ...]: output: list[tuple[SupportedTask, ...]] diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py index afc333723d0f..df20a90d51d3 100644 --- a/vllm/v1/executor/multiproc_executor.py +++ b/vllm/v1/executor/multiproc_executor.py @@ -29,6 +29,7 @@ from vllm.config import VllmConfig from vllm.distributed import destroy_distributed_environment, destroy_model_parallel from vllm.distributed.device_communicators.shm_broadcast import Handle, MessageQueue +from vllm.distributed.ec_transfer.ec_connector.utils import ECOutputAggregator from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator from vllm.distributed.parallel_state import ( get_dcp_group, @@ -343,6 +344,7 @@ def execute_model( # type: ignore[override] non_block=non_block, timeout=envs.VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS, kv_output_aggregator=self.kv_output_aggregator, + ec_output_aggregator=self.ec_output_aggregator, ) def sample_tokens( # type: ignore[override] @@ -355,6 +357,7 @@ def sample_tokens( # type: ignore[override] non_block=non_block, timeout=envs.VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS, kv_output_aggregator=self.kv_output_aggregator, + ec_output_aggregator=self.ec_output_aggregator, ) def execute_dummy_batch(self) -> None: @@ -375,9 +378,10 @@ def collective_rpc( # type: ignore[override] non_block: bool = False, unique_reply_rank: int | None = None, kv_output_aggregator: KVOutputAggregator | None = None, + ec_output_aggregator: ECOutputAggregator | None = None, ) -> Any: - """Returns single result if unique_reply_rank and/or kv_output_aggregator - is provided, otherwise list.""" + """Returns single result if unique_reply_rank and/or an output + aggregator is provided, otherwise list.""" assert self.rpc_broadcast_mq is not None, ( "collective_rpc should not be called on follower node" ) @@ -387,11 +391,21 @@ def collective_rpc( # type: ignore[override] deadline = None if timeout is None else time.monotonic() + timeout kwargs = kwargs or {} - if kv_output_aggregator is not None: + aggregators = [a for a in (kv_output_aggregator, ec_output_aggregator) if a] + aggregate: Callable[[Any], Any] + if aggregators: output_rank = None - aggregate: Callable[[Any], Any] = partial( - kv_output_aggregator.aggregate, output_rank=unique_reply_rank or 0 - ) + + def _aggregate(outputs: Any) -> Any: + # Each aggregator merges its own connector's output onto + # outputs[output_rank] in place and returns it, so chaining is safe. + rank = unique_reply_rank or 0 + result = outputs[rank] + for a in aggregators: + result = a.aggregate(outputs, output_rank=rank) + return result + + aggregate = _aggregate else: output_rank = unique_reply_rank aggregate = lambda x: x diff --git a/vllm/v1/executor/ray_executor.py b/vllm/v1/executor/ray_executor.py index 39749ffc257e..986e4d9bbb03 100644 --- a/vllm/v1/executor/ray_executor.py +++ b/vllm/v1/executor/ray_executor.py @@ -89,6 +89,19 @@ def _init_executor(self) -> None: # KV connector setup self.has_connector = self.vllm_config.kv_transfer_config is not None + if ( + self.vllm_config.ec_transfer_config is not None + and self.parallel_config.world_size > 1 + ): + raise NotImplementedError( + "EC connector worker metadata is not supported with the " + "legacy Ray executor when world_size > 1: only the output " + "of a single worker is fetched, silently dropping the " + "other workers' EC connector state. Set " + "VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1 to use RayExecutorV2, " + "or use the multiprocessing executor instead." + ) + self.uses_sampler = self.vllm_config.model_config.runner_type != "pooling" and ( self.vllm_config.ec_transfer_config is None or self.vllm_config.ec_transfer_config.is_ec_consumer diff --git a/vllm/v1/outputs.py b/vllm/v1/outputs.py index 0bbee7667527..80b909afbc8d 100644 --- a/vllm/v1/outputs.py +++ b/vllm/v1/outputs.py @@ -14,6 +14,7 @@ from vllm.v1.core.sched.output import SchedulerOutput if TYPE_CHECKING: + from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorWorkerMetadata from vllm.distributed.kv_events import KVConnectorKVEvents from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorWorkerMetadata, @@ -23,6 +24,7 @@ KVConnectorStats = object KVConnectorWorkerMetadata = object KVConnectorKVEvents = object + ECConnectorWorkerMetadata = object class LogprobsLists(NamedTuple): @@ -292,6 +294,14 @@ class ECConnectorOutput: # [mm_hash] finished_sending: set[str] | None = None finished_recving: set[str] | None = None + ec_connector_worker_meta: ECConnectorWorkerMetadata | None = None + + def is_empty(self): + return ( + not self.finished_sending + and not self.finished_recving + and not self.ec_connector_worker_meta + ) # ModelRunnerOutput is serialized and sent to the scheduler process. @@ -362,6 +372,32 @@ def with_kv_conn_output_only( output.kv_connector_output = kv_connector_output return output + @staticmethod + def with_ec_conn_output_only( + ec_connector_output: ECConnectorOutput | None, + ) -> "ModelRunnerOutput": + """Return an otherwise-empty output carrying `ec_connector_output`.""" + return ModelRunnerOutput.with_ec_conn_output( + EMPTY_MODEL_RUNNER_OUTPUT, ec_connector_output + ) + + @staticmethod + def with_ec_conn_output( + output: "ModelRunnerOutput", + ec_connector_output: ECConnectorOutput | None, + ) -> "ModelRunnerOutput": + """Return `output` carrying `ec_connector_output`. + + The shared empty output is copied rather than written to, so callers + must use the return value. + """ + if ec_connector_output is None or ec_connector_output.is_empty(): + return output + if output is EMPTY_MODEL_RUNNER_OUTPUT: + output = copy(EMPTY_MODEL_RUNNER_OUTPUT) + output.ec_connector_output = ec_connector_output + return output + # ModelRunnerOutput wrapper for async scheduling. class AsyncModelRunnerOutput(ABC): diff --git a/vllm/v1/worker/gpu/ec_connector.py b/vllm/v1/worker/gpu/ec_connector.py index 825763b82001..5dc8d92359a9 100644 --- a/vllm/v1/worker/gpu/ec_connector.py +++ b/vllm/v1/worker/gpu/ec_connector.py @@ -9,7 +9,11 @@ from vllm.config import VllmConfig from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorBase -from vllm.v1.outputs import ECConnectorOutput +from vllm.v1.outputs import ( + EMPTY_MODEL_RUNNER_OUTPUT, + ECConnectorOutput, + ModelRunnerOutput, +) if TYPE_CHECKING: from vllm.v1.core.sched.output import SchedulerOutput @@ -25,6 +29,12 @@ def maybe_get_output( ) -> Generator[ECConnectorOutput | None, None, None]: yield None + def no_forward( + self, + scheduler_output: "SchedulerOutput", + ) -> ModelRunnerOutput: + return EMPTY_MODEL_RUNNER_OUTPUT + class ActiveECConnector(ECConnector): def __init__( @@ -33,9 +43,11 @@ def __init__( encoder_cache: dict[str, torch.Tensor], ) -> None: self.encoder_cache = encoder_cache - self.save_new_caches = vllm_config.is_ec_producer_only self.ec_connector = get_ec_transfer() assert isinstance(self.ec_connector, ECConnectorBase) + # Every producer offloads freshly computed encoder outputs, including + # an ec_both node that also reloads them. + self.save_new_caches = self.ec_connector.is_producer @contextmanager def maybe_get_output( @@ -65,8 +77,19 @@ def maybe_get_output( output.finished_sending, output.finished_recving = ( ec_connector.get_finished(scheduler_output.finished_req_ids) ) + output.ec_connector_worker_meta = ec_connector.build_connector_worker_meta() ec_connector.clear_connector_metadata() + def no_forward( + self, + scheduler_output: "SchedulerOutput", + ) -> ModelRunnerOutput: + # EC send/recv even if no work to do. + with self.maybe_get_output(scheduler_output) as ec_connector_output: + pass + + return ModelRunnerOutput.with_ec_conn_output_only(ec_connector_output) + NO_OP_EC_CONNECTOR = ECConnector() diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index c92ffedf5961..abbe6d5d3de6 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -65,6 +65,7 @@ from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec from vllm.v1.outputs import ( DraftTokenIds, + ECConnectorOutput, ModelRunnerOutput, RoutedExpertsTensors, make_empty_encoder_model_runner_output, @@ -1393,6 +1394,15 @@ def postprocess_sampled( idx_mapping, num_sampled, self.req_states.num_computed_tokens.gpu ) + def _merge_ec_connector_no_forward( + self, scheduler_output: SchedulerOutput, output: ModelRunnerOutput + ) -> ModelRunnerOutput: + """Let the EC connector send/recv on a step with no work to run.""" + return ModelRunnerOutput.with_ec_conn_output( + output, + self.ec_connector.no_forward(scheduler_output).ec_connector_output, + ) + @torch.inference_mode() def execute_model( self, @@ -1414,7 +1424,9 @@ def execute_model( if scheduler_output.total_num_scheduled_tokens == 0: # No need to run the model. empty_output = self.kv_connector.no_forward(scheduler_output) - return empty_output + return self._merge_ec_connector_no_forward( + scheduler_output, empty_output + ) # Get batch descriptor and sync across DP ranks. num_reqs = len(scheduler_output.num_scheduled_tokens) @@ -1455,7 +1467,7 @@ def execute_model( if batch_desc.num_tokens == 0: # All DP ranks have zero tokens to run. empty_output = self.kv_connector.no_forward(scheduler_output) - return empty_output + return self._merge_ec_connector_no_forward(scheduler_output, empty_output) if not dummy_run: # Common case. @@ -1572,9 +1584,10 @@ def execute_model( input_ids = None if self.is_encoder_only: - output = make_empty_encoder_model_runner_output(scheduler_output) - output.ec_connector_output = ec_connector_output - return output + return ModelRunnerOutput.with_ec_conn_output( + make_empty_encoder_model_runner_output(scheduler_output), + ec_connector_output, + ) model_inputs = { "input_ids": input_ids, @@ -1679,6 +1692,7 @@ def execute_model( hidden_states=hidden_states, aux_hidden_states=aux_hidden_states, finished_req_ids=finished_req_ids, + ec_connector_output=ec_connector_output, routed_experts=routed_experts, ) @@ -1702,6 +1716,7 @@ def sample_tokens( hidden_states = self.execute_model_state.hidden_states aux_hidden_states = self.execute_model_state.aux_hidden_states finished_req_ids = self.execute_model_state.finished_req_ids + ec_connector_output = self.execute_model_state.ec_connector_output routed_experts = self.execute_model_state.routed_experts self.execute_model_state = None @@ -1721,7 +1736,9 @@ def sample_tokens( # Post-step KV connector related operations. kv_connector_output = self.kv_connector.post_forward(finished_req_ids) - return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) + # The first PP rank holds the encoder cache, so pass its EC output on. + output = ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) + return ModelRunnerOutput.with_ec_conn_output(output, ec_connector_output) # Last rank: sample tokens hidden_states, input_batch = pcp.maybe_restore_pcp_for_sampling( @@ -1838,6 +1855,7 @@ def sample_tokens( # Post-step KV connector related operations. kv_connector_output = self.kv_connector.post_forward(finished_req_ids) model_runner_output.kv_connector_output = kv_connector_output + model_runner_output.ec_connector_output = ec_connector_output return async_output @@ -1854,6 +1872,7 @@ def pool(self) -> AsyncPoolingOutput | ModelRunnerOutput | None: input_batch = self.execute_model_state.input_batch hidden_states = self.execute_model_state.hidden_states finished_req_ids = self.execute_model_state.finished_req_ids + ec_connector_output = self.execute_model_state.ec_connector_output self.execute_model_state = None # Post-step KV connector related operations. @@ -1861,7 +1880,8 @@ def pool(self) -> AsyncPoolingOutput | ModelRunnerOutput | None: if not self.is_last_pp_rank: self.postprocess_num_computed_tokens(input_batch) - return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) + output = ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) + return ModelRunnerOutput.with_ec_conn_output(output, ec_connector_output) assert self.pooling_runner is not None pooler_output, finished_mask = self.pooling_runner.pool( @@ -1873,6 +1893,7 @@ def pool(self) -> AsyncPoolingOutput | ModelRunnerOutput | None: req_ids=input_batch.req_ids, req_id_to_index={req_id: i for i, req_id in enumerate(input_batch.req_ids)}, kv_connector_output=kv_connector_output, + ec_connector_output=ec_connector_output, ) async_output = AsyncPoolingOutput( model_runner_output=model_runner_output, @@ -1962,6 +1983,7 @@ class ExecuteModelState(NamedTuple): hidden_states: torch.Tensor | None aux_hidden_states: list[torch.Tensor] | None finished_req_ids: set[str] + ec_connector_output: ECConnectorOutput | None routed_experts: RoutedExpertsTensors | None From edd4c8176cfd98ece8a29beda574378c42971967 Mon Sep 17 00:00:00 2001 From: Toby Mao Date: Sat, 15 Aug 2026 18:39:23 -0700 Subject: [PATCH 004/839] [Bugfix][DSv4] Revert adaptive C128A metadata packing (#51318) Signed-off-by: tobymao Signed-off-by: Yongye Zhu Co-authored-by: Claude Co-authored-by: Roger Wang Co-authored-by: Yongye Zhu Co-authored-by: OpenAI Codex --- .../kernels/attention/test_flashmla_sparse.py | 37 ------------------- vllm/models/deepseek_v4/sparse_mla.py | 26 ++++--------- 2 files changed, 7 insertions(+), 56 deletions(-) diff --git a/tests/kernels/attention/test_flashmla_sparse.py b/tests/kernels/attention/test_flashmla_sparse.py index 8082a359cc21..eb953faa4469 100644 --- a/tests/kernels/attention/test_flashmla_sparse.py +++ b/tests/kernels/attention/test_flashmla_sparse.py @@ -4,43 +4,6 @@ import torch -def test_deepseek_v4_c128a_dynamic_topk_packed_buffers(): - from vllm.models.deepseek_v4.sparse_mla import build_c128a_topk_metadata - - device = torch.device("cuda") - capacity_width = 256 - active_width = 128 - global_decode_buffer = torch.empty( - (2, capacity_width), dtype=torch.int32, device=device - ) - decode_lens_buffer = torch.empty(2, dtype=torch.int32, device=device) - prefill_buffer = torch.empty((2, capacity_width), dtype=torch.int32, device=device) - - global_decode, decode_lens, prefill_local = build_c128a_topk_metadata( - positions=torch.tensor([255, 511], dtype=torch.int64, device=device), - compress_ratio=128, - num_decode_tokens=1, - token_to_req_indices=torch.tensor([0, 0], dtype=torch.int32, device=device), - block_table=torch.tensor([[3]], dtype=torch.int32, device=device), - block_size=capacity_width, - slot_mapping=torch.tensor([0, 1], dtype=torch.int64, device=device), - global_decode_buffer=global_decode_buffer, - decode_lens_buffer=decode_lens_buffer, - prefill_buffer=prefill_buffer, - max_compressed_tokens=active_width, - ) - - assert global_decode.shape == (1, active_width) - assert prefill_local.shape == (1, active_width) - assert global_decode.stride() == (active_width, 1) - assert prefill_local.stride() == (active_width, 1) - assert global_decode[0, :2].cpu().tolist() == [768, 769] - assert decode_lens.cpu().tolist() == [2] - assert prefill_local[0, :4].cpu().tolist() == list(range(4)) - assert torch.all(global_decode[0, 2:] == -1) - assert torch.all(prefill_local[0, 4:] == -1) - - def test_sparse_flashmla_metadata_smoke(): import vllm.v1.attention.ops.flashmla as fm diff --git a/vllm/models/deepseek_v4/sparse_mla.py b/vllm/models/deepseek_v4/sparse_mla.py index db8ab96e90ff..0475e4c96df7 100644 --- a/vllm/models/deepseek_v4/sparse_mla.py +++ b/vllm/models/deepseek_v4/sparse_mla.py @@ -257,13 +257,6 @@ def _build_c128a_metadata( assert cm.positions is not None, ( "positions is required for C128A metadata build" ) - active_topk_width = min( - max( - triton.next_power_of_2(max(cm.max_seq_len // self.compress_ratio, 1)), - _C128A_TOPK_ALIGNMENT, - ), - self.c128a_max_compressed, - ) block_size = self.kv_cache_spec.block_size // self.compress_ratio global_decode, decode_lens, prefill_local = build_c128a_topk_metadata( cm.positions[:num_total], @@ -276,7 +269,7 @@ def _build_c128a_metadata( self.c128a_global_decode_buffer, self.c128a_decode_lens_buffer, self.c128a_prefill_buffer, - max_compressed_tokens=active_topk_width, + max_compressed_tokens=self.c128a_max_compressed, ) result: dict[str, torch.Tensor | None] = {} @@ -322,30 +315,25 @@ def build_c128a_topk_metadata( Decode tokens: position → block_table lookup → global slot ids + topk_lens. Prefill tokens: position → local indices [0, ..., n-1, -1, ...]. - Writes into packed views of pre-allocated buffers for CUDA graph stability. + Writes into pre-allocated buffers for CUDA graph address stability. + Returns slices of the buffers. """ num_tokens = positions.shape[0] num_prefill_tokens = num_tokens - num_decode_tokens - # view(-1) as 1-d array and then expanded to - # [num_decode_tokens, max_compressed_tokens] - global_decode = global_decode_buffer.view(-1)[ - : num_decode_tokens * max_compressed_tokens - ].view(num_decode_tokens, max_compressed_tokens) + global_decode = global_decode_buffer[:num_decode_tokens] decode_lens = decode_lens_buffer[:num_decode_tokens] - prefill_local = prefill_buffer.view(-1)[ - : num_prefill_tokens * max_compressed_tokens - ].view(num_prefill_tokens, max_compressed_tokens) + prefill_local = prefill_buffer[:num_prefill_tokens] if num_tokens == 0: return global_decode, decode_lens, prefill_local _build_c128a_topk_metadata_kernel[(num_tokens,)]( global_decode_buffer, - max_compressed_tokens, + global_decode_buffer.stride(0), decode_lens_buffer, prefill_buffer, - max_compressed_tokens, + prefill_buffer.stride(0), positions, compress_ratio, max_compressed_tokens, From 6593754e61a45b27a93e40fa0479568945d38661 Mon Sep 17 00:00:00 2001 From: mispa-ms <81828223+mispa-ms@users.noreply.github.com> Date: Sat, 15 Aug 2026 21:03:53 -0700 Subject: [PATCH 005/839] [Bugfix][Spec Decode] Keep EAGLE cache registration on the partial-hash-hit path (#52419) --- .../test_partial_prefix_cache_hits.py | 74 +++++++++++++++++++ vllm/v1/core/kv_cache_coordinator.py | 38 +++++----- 2 files changed, 92 insertions(+), 20 deletions(-) diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py index 97199a15ab49..d237ce9f345b 100644 --- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py @@ -261,6 +261,80 @@ def test_hybrid_mamba_align_partial_hash_hit(): assert manager.get_blocks("1").blocks[1][1].block_hash_num_tokens == 8 +def test_eagle_group_registers_unaligned_tail_under_partial_hash_hits(): + """An EAGLE group must not re-floor what partial hash hits leaves un-floored. + + ``cache_blocks`` decides once how far a request may be registered, and with + fine-grained partial hash hits that bound is the raw token count. The EAGLE + branch then re-derives its own bound for the lookahead block; if it rounds + down to ``scheduler_block_size`` again, everything between the last aligned + boundary and the tail stops being registered -- ``(n % scheduler_block_size) + - manager.block_size`` tokens per call, which is most of a segment whenever + the group's own block is much smaller than the scheduler block. + """ + hash_block_size = 2 + mamba_block_size = 4 * hash_block_size + kv_cache_config = KVCacheConfig( + num_blocks=40, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=hash_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=mamba_block_size, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ], + ) + manager = make_kv_cache_manager( + kv_cache_config=kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=hash_block_size, + ) + coordinator = manager.coordinator + assert coordinator.enable_partial_hash_hits + # The full-attention group is the EAGLE one, and its block is smaller than + # the scheduler block -- the geometry where re-flooring loses tokens. + eagle_manager = coordinator.single_type_managers[0] + eagle_manager.use_eagle = True + assert eagle_manager.block_size < coordinator.scheduler_block_size + + # Deliberately not a multiple of the scheduler block, so the two bounds + # differ: floor(22/8)*8 + 2 = 18 against 22. + num_tokens = coordinator.scheduler_block_size * 2 + hash_block_size * 3 + req = make_request("0", list(range(num_tokens)), hash_block_size, sha256) + + recorded: list[int] = [] + for single_type_manager in coordinator.single_type_managers: + original = single_type_manager.cache_blocks + + def spy(request, num_tokens_to_cache, *args, _orig=original, **kwargs): + recorded.append(num_tokens_to_cache) + return _orig(request, num_tokens_to_cache, *args, **kwargs) + + single_type_manager.cache_blocks = spy + + # allocate_slots caches on the way out, so this exercises the real path. + computed_blocks, num_computed, _ = manager.get_computed_blocks(req) + assert manager.allocate_slots(req, num_tokens, num_computed, computed_blocks) + + # Every group, EAGLE or not, may register the whole unaligned tail. + assert recorded == [num_tokens] * len(coordinator.single_type_managers) + + def test_hybrid_mamba_partial_tail_owner_uses_cow_on_continue(): hash_block_size = 2 block_size = 2 * hash_block_size diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index f5cd79b285f6..8efaf9252e8a 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -6,7 +6,7 @@ from vllm import envs from vllm.logger import init_logger -from vllm.utils.math_utils import cdiv +from vllm.utils.math_utils import cdiv, round_down from vllm.v1.core.block_pool import BlockPool from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector from vllm.v1.core.kv_cache_utils import ( @@ -704,39 +704,37 @@ def verify_and_split_kv_cache_groups(self) -> None: for gid in group.group_ids: self.single_type_managers[gid].use_eagle = True - def cache_blocks(self, request: Request, num_computed_tokens: int) -> None: + def _align_cacheable(self, num_tokens: int) -> int: + """Largest prefix of ``num_tokens`` a future cache hit could match. + + Hits are ``scheduler_block_size``-aligned (see + ``find_longest_cache_hit``) unless fine-grained partial hash hits are + enabled, in which case no rounding applies -- rounding even to + ``hash_block_size`` would re-register a privatized Mamba tail. + """ if self.enable_partial_hash_hits: - aligned_num_computed_tokens = num_computed_tokens - else: - # Cache hits in this coordinator are always a multiple of - # ``scheduler_block_size`` tokens (see ``find_longest_cache_hit``). - # Within an aligned region, SWA groups may only consult a subset of - # blocks per ``scheduler_block_size``-segment so the unused blocks - # also stay out of the prefix-cache hash map. - aligned_num_computed_tokens = ( - num_computed_tokens - // self.scheduler_block_size - * self.scheduler_block_size - ) + return num_tokens + return round_down(num_tokens, self.scheduler_block_size) + + def cache_blocks(self, request: Request, num_computed_tokens: int) -> None: + cached_num_computed_tokens = self._align_cacheable(num_computed_tokens) for manager in self.single_type_managers: - num_tokens_to_cache = aligned_num_computed_tokens + num_tokens_to_cache = cached_num_computed_tokens # EAGLE groups match one block past each aligned boundary and drop # it, so make that lookahead block eligible to be cached. - if manager.use_eagle and aligned_num_computed_tokens > 0: + if manager.use_eagle and cached_num_computed_tokens > 0: # Only cache tokens with finalized KV. The last # num_reprefillable_tokens tokens can be re-prefilled during # multi-module MTP. num_finalized_computed_tokens = max( 0, num_computed_tokens - self.num_reprefillable_tokens ) - aligned_num_finalized_computed_tokens = ( + cached_num_finalized_computed_tokens = self._align_cacheable( num_finalized_computed_tokens - // self.scheduler_block_size - * self.scheduler_block_size ) num_tokens_to_cache = min( num_finalized_computed_tokens, - aligned_num_finalized_computed_tokens + manager.block_size, + cached_num_finalized_computed_tokens + manager.block_size, ) # The manager already knows the fine hit granularity # (``scheduler_block_size``); retention is passed separately so it From 8efa13b700f1836657699cae2503dc2feab27fa0 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Sat, 15 Aug 2026 21:04:08 -0700 Subject: [PATCH 006/839] [Bugfix] Pick the DeepSeek V4 eager cudagraph region per model runner (#52401) --- tests/test_config.py | 61 ++++++++++-------------- vllm/config/vllm.py | 32 +++---------- vllm/models/deepseek_v4/attention.py | 69 ++++++++++++++++++++++++++-- 3 files changed, 96 insertions(+), 66 deletions(-) diff --git a/tests/test_config.py b/tests/test_config.py index 93bcdbc258d6..70c8728d25d4 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -66,41 +66,18 @@ def test_v2_model_runner_env_tri_state(monkeypatch, env_value, expected): assert envs.VLLM_USE_V2_MODEL_RUNNER is expected -@pytest.mark.parametrize( - "cudagraph_mode", - [CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL_AND_PIECEWISE], -) -def test_deepseek_v4_rejects_mrv1_piecewise_cudagraph(cudagraph_mode): - config = SimpleNamespace( - use_v2_model_runner=False, - model_config=SimpleNamespace(architectures=["DeepseekV4ForCausalLM"]), - compilation_config=SimpleNamespace(cudagraph_mode=cudagraph_mode), - ) - - with pytest.raises(ValueError, match="DeepSeek V4 does not support PIECEWISE"): - VllmConfig._validate_mrv1_piecewise_cudagraph(config) - - -@pytest.mark.parametrize( - ("use_v2_model_runner", "architecture", "cudagraph_mode"), - [ - (True, "DeepseekV4ForCausalLM", CUDAGraphMode.PIECEWISE), - (False, "DeepseekV4ForCausalLM", CUDAGraphMode.NONE), - (False, "DeepseekV4ForCausalLM", CUDAGraphMode.FULL), - (False, "DeepseekV4ForCausalLM", CUDAGraphMode.FULL_DECODE_ONLY), - (False, "LlamaForCausalLM", CUDAGraphMode.PIECEWISE), - ], -) -def test_mrv1_piecewise_cudagraph_allowed( - use_v2_model_runner, architecture, cudagraph_mode -): - config = SimpleNamespace( - use_v2_model_runner=use_v2_model_runner, - model_config=SimpleNamespace(architectures=[architecture]), - compilation_config=SimpleNamespace(cudagraph_mode=cudagraph_mode), - ) - - VllmConfig._validate_mrv1_piecewise_cudagraph(config) +def test_rocm_defaults_deepseek_v4_to_mrv1(monkeypatch): + """ROCm keeps DeepSeek V4 on MRV1, which is still faster there.""" + from vllm.config.vllm import default_v2_model_runner_architectures + from vllm.platforms import current_platform + + monkeypatch.setattr(current_platform, "is_rocm", lambda: True) + # The lookup is lru_cached against a fixed platform. + default_v2_model_runner_architectures.cache_clear() + try: + assert "DeepseekV4ForCausalLM" not in default_v2_model_runner_architectures() + finally: + default_v2_model_runner_architectures.cache_clear() @pytest.mark.parametrize( @@ -330,10 +307,20 @@ def test_resolve_cudagraph_mode_adjusts_spec_decode_sizes_only_for_v1( ), ], ) -def test_is_default_v2_model_runner_model(model_config, expected): +def test_is_default_v2_model_runner_model(model_config, expected, monkeypatch): + from vllm.config.vllm import default_v2_model_runner_architectures + from vllm.platforms import current_platform + + # The expectations below are the platform-independent defaults; ROCm's + # DeepSeek V4 carve-out is covered by test_rocm_defaults_deepseek_v4_to_mrv1. + monkeypatch.setattr(current_platform, "is_rocm", lambda: False) + default_v2_model_runner_architectures.cache_clear() config = SimpleNamespace(model_config=model_config) - assert VllmConfig._is_default_v2_model_runner_model(config) is expected + try: + assert VllmConfig._is_default_v2_model_runner_model(config) is expected + finally: + default_v2_model_runner_architectures.cache_clear() @pytest.mark.skip_global_cleanup diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 2b5e512ccacf..0a4c5cad8838 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -66,10 +66,6 @@ logger = init_logger(__name__) -MRV1_UNSUPPORTED_PIECEWISE_CUDAGRAPH_ARCHITECTURES = frozenset( - {"DeepseekV4ForCausalLM"} -) - DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset( { "DeepseekV2ForCausalLM", @@ -87,6 +83,13 @@ @lru_cache def default_v2_model_runner_architectures() -> frozenset[str]: """Architectures defaulting to the V2 model runner on this platform.""" + from vllm.platforms import current_platform + + if current_platform.is_rocm(): + # TODO(rocm): DeepSeek V4 is still faster on MRV1 on ROCm. The + # attention layer picks the eager cudagraph region MRV1 needs, so + # this is a perf default only; drop it once MRV2 catches up. + return DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES - {"DeepseekV4ForCausalLM"} return DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES @@ -691,25 +694,6 @@ def _is_default_v2_model_runner_model(self) -> bool: return False return is_default_v2_architecture or not model_config.is_moe - def _validate_mrv1_piecewise_cudagraph(self) -> None: - if self.use_v2_model_runner: - return - model_config = self.model_config - if model_config is None: - return - if not self.compilation_config.cudagraph_mode.has_piecewise_cudagraphs(): - return - architectures = getattr(model_config, "architectures", []) - if any( - arch in MRV1_UNSUPPORTED_PIECEWISE_CUDAGRAPH_ARCHITECTURES - for arch in architectures - ): - raise ValueError( - "DeepSeek V4 does not support PIECEWISE CUDA graphs with " - "Model Runner V1. Use Model Runner V2 or disable PIECEWISE " - "CUDA graphs." - ) - @property def needs_dp_coordinator(self) -> bool: """ @@ -1644,8 +1628,6 @@ def has_blocked_weights(): "pipeline parallelism", ) - self._validate_mrv1_piecewise_cudagraph() - # final check of cudagraph mode after all possible updates if current_platform.is_cuda_alike(): if ( diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index 74a481e4b5f6..e2dc1bb45a7a 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -298,6 +298,13 @@ def __init__( eager_scratch_pool=eager_scratch_pool, ) + self._prepare_and_attn_fn = self._prepare_and_attn + if not vllm_config.use_v2_model_runner: + # MRV1's piecewise capture only tolerates the wide eager region: with + # the narrow one the attention input preparation stays in the captured + # graph and MRV1 produces garbage (#51430). + self._prepare_and_attn_fn = self._prepare_and_attn_eager + # Will be None on ROCm for now. self.aux_stream_list = aux_stream_list # [0]: GEMM start / post-GEMM event0. [1..3]: GEMM done events; @@ -379,6 +386,64 @@ def forward( self.eps, ) + self._prepare_and_attn_fn( + hidden_states, + qr, + kv, + kv_score, + indexer_kv_score, + indexer_weights, + positions, + o_padded, + ) + o = o_padded[:, : self.n_local_heads, :] + + # Inverse-RoPE + wo_a + wo_b output projection (platform-specific). + return self._o_proj(o, positions) + + @eager_break_during_capture + def _prepare_and_attn_eager( + self, + hidden_states: torch.Tensor, + qr: torch.Tensor, + kv: torch.Tensor, + kv_score: torch.Tensor, + indexer_kv_score: torch.Tensor, + indexer_weights: torch.Tensor, + positions: torch.Tensor, + o_padded: torch.Tensor, + ) -> None: + """Wide eager region: the whole of ``_prepare_and_attn`` runs eagerly. + + The nested ``_sparse_indexer_and_attn`` break runs inline, since + ``add_eager`` clears ``_capturing`` before invoking this. + """ + self._prepare_and_attn( + hidden_states, + qr, + kv, + kv_score, + indexer_kv_score, + indexer_weights, + positions, + o_padded, + ) + + def _prepare_and_attn( + self, + hidden_states: torch.Tensor, + qr: torch.Tensor, + kv: torch.Tensor, + kv_score: torch.Tensor, + indexer_kv_score: torch.Tensor, + indexer_weights: torch.Tensor, + positions: torch.Tensor, + o_padded: torch.Tensor, + ) -> None: + """Attention input preparation followed by the sparse indexer and MLA. + + Only the latter runs in the eager break. + """ attn_metadata = get_forward_context().attn_metadata indexer = self.indexer compressor = self.compressor @@ -438,10 +503,6 @@ def project_query_and_cache_kv() -> torch.Tensor: positions, o_padded, ) - o = o_padded[:, : self.n_local_heads, :] - - # Inverse-RoPE + wo_a + wo_b output projection (platform-specific). - return self._o_proj(o, positions) def _fused_wqa_wkv_gemm(self, hidden_states: torch.Tensor) -> torch.Tensor: # Override point: the ROCm layer preshuffles this weight in place, so From 41f12a0daf193c80f2e748f6a65f0735623be3a5 Mon Sep 17 00:00:00 2001 From: Jeffrey Wang Date: Sun, 16 Aug 2026 01:04:51 -0700 Subject: [PATCH 007/839] [Bugfix] Raise `VLLMValidationError` from structured output validators (#52394) Signed-off-by: Jeffrey Wang Co-authored-by: Claude Opus 5 (1M context) --- .../llm/test_struct_output_generate.py | 7 ++- tests/v1/structured_output/test_validation.py | 59 ++++++++++++++++++- vllm/entrypoints/openai/engine/protocol.py | 2 +- vllm/sampling_params.py | 9 ++- vllm/v1/structured_output/backend_guidance.py | 14 +++-- .../backend_lm_format_enforcer.py | 9 +-- vllm/v1/structured_output/backend_outlines.py | 20 ++++--- vllm/v1/structured_output/backend_xgrammar.py | 19 +++--- 8 files changed, 109 insertions(+), 30 deletions(-) diff --git a/tests/entrypoints/llm/test_struct_output_generate.py b/tests/entrypoints/llm/test_struct_output_generate.py index 7b90aeacda12..0f55ad7adc3e 100644 --- a/tests/entrypoints/llm/test_struct_output_generate.py +++ b/tests/entrypoints/llm/test_struct_output_generate.py @@ -13,6 +13,7 @@ from tests.reasoning.utils import run_reasoning_extraction from vllm.config import StructuredOutputsConfig +from vllm.exceptions import VLLMValidationError from vllm.outputs import RequestOutput from vllm.platforms import current_platform from vllm.reasoning.abs_reasoning_parsers import ReasoningParserManager @@ -326,7 +327,7 @@ def test_structured_output( ) if backend.startswith("xgrammar"): with pytest.raises( - ValueError, + VLLMValidationError, match="The provided JSON schema contains features " "not supported by xgrammar.", ): @@ -453,7 +454,9 @@ def test_structured_output( max_tokens=1000, structured_outputs=StructuredOutputsParams(grammar="not a grammar"), ) - with pytest.raises(ValueError, match="Failed to convert the grammar "): + with pytest.raises( + VLLMValidationError, match="Failed to convert the grammar " + ): runner.llm.generate( ( "Generate a sql statement that selects col_1 from " diff --git a/tests/v1/structured_output/test_validation.py b/tests/v1/structured_output/test_validation.py index 08be56613ceb..af417bc03e6f 100644 --- a/tests/v1/structured_output/test_validation.py +++ b/tests/v1/structured_output/test_validation.py @@ -5,7 +5,7 @@ import pytest from vllm.config import StructuredOutputsConfig -from vllm.exceptions import VLLMValidationError +from vllm.exceptions import VLLMClientError, VLLMValidationError from vllm.sampling_params import SamplingParams, StructuredOutputsParams pytestmark = pytest.mark.cpu_test @@ -103,3 +103,60 @@ def test_regex_with_nul_byte_rejected(regex): with pytest.raises(ValueError, match="NUL"): validate_xgrammar_grammar(params) + + +INVALID_JSON_SCHEMA = {"type": "object", "properties": {"name": {"type": "str"}}} + + +@pytest.mark.parametrize( + "backend, structured_outputs", + [ + ("xgrammar", StructuredOutputsParams(json=INVALID_JSON_SCHEMA)), + ("outlines", StructuredOutputsParams(json=INVALID_JSON_SCHEMA)), + ("auto", StructuredOutputsParams(json=INVALID_JSON_SCHEMA)), + ("auto", StructuredOutputsParams(json='{"type": ')), + ("xgrammar", StructuredOutputsParams(grammar="not a grammar")), + ("guidance", StructuredOutputsParams(grammar="not a grammar")), + ("lm-format-enforcer", StructuredOutputsParams(grammar="not a grammar")), + ("outlines", StructuredOutputsParams(regex="(")), + ("guidance", StructuredOutputsParams(structural_tag='{"nope": 1}')), + ], +) +def test_unsupported_grammar_is_a_client_error(backend, structured_outputs): + """Only `VLLMClientError` survives `AsyncLLM.generate` untouched; anything else + is wrapped in `EngineGenerateError` and served as a 500 instead of a 400.""" + params = SamplingParams(structured_outputs=structured_outputs) + with pytest.raises(VLLMClientError): + params._validate_structured_outputs( + _StubModelConfig(is_diffusion=False), + StructuredOutputsConfig(backend=backend), + tokenizer=object(), + ) + + +@pytest.mark.parametrize( + "schema, expected_backend", + [ + # multipleOf is unsupported by xgrammar, patternProperties also by guidance. + ( + { + "type": "object", + "properties": {"n": {"type": "integer", "multipleOf": 2}}, + }, + "guidance", + ), + ( + {"type": "object", "patternProperties": {"^a": {"type": "string"}}}, + "outlines", + ), + ], +) +def test_auto_backend_falls_back_on_unsupported_schema(schema, expected_backend): + """`auto` falls back on rejection, so it must catch what the validators raise.""" + params = SamplingParams(structured_outputs=StructuredOutputsParams(json=schema)) + params._validate_structured_outputs( + _StubModelConfig(is_diffusion=False), + StructuredOutputsConfig(backend="auto"), + tokenizer=object(), + ) + assert params.structured_outputs._backend == expected_backend diff --git a/vllm/entrypoints/openai/engine/protocol.py b/vllm/entrypoints/openai/engine/protocol.py index dded1261dcbe..c32be4459c90 100644 --- a/vllm/entrypoints/openai/engine/protocol.py +++ b/vllm/entrypoints/openai/engine/protocol.py @@ -266,7 +266,7 @@ def validate_structural_tag_payload(payload: Any, *, parameter: str) -> None: structured_outputs=StructuredOutputsParams(structural_tag=payload) ) ) - except (TypeError, ValueError) as exc: + except (TypeError, ValueError, VLLMValidationError) as exc: raise VLLMValidationError( f"Invalid {parameter} structural_tag specification.", parameter=parameter, diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index 46c404608289..2f0af9e1b43b 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -1092,7 +1092,7 @@ def _validate_structured_outputs( try: validate_xgrammar_grammar(self) self.structured_outputs._backend = "xgrammar" - except ValueError: + except VLLMValidationError: # The request either failed validation # or includes some jsonschema feature(s) that # are not supported in xgrammar. @@ -1103,7 +1103,12 @@ def _validate_structured_outputs( so_params = self.structured_outputs if not skip_guidance and so_params.json: if isinstance(so_params.json, str): - schema = json_mod.loads(so_params.json) + try: + schema = json_mod.loads(so_params.json) + except json_mod.JSONDecodeError as e: + raise VLLMValidationError( + "Invalid JSON grammar specification." + ) from e else: schema = so_params.json skip_guidance = has_guidance_unsupported_json_features(schema) diff --git a/vllm/v1/structured_output/backend_guidance.py b/vllm/v1/structured_output/backend_guidance.py index 30ecbfa065a6..310576cd99a8 100644 --- a/vllm/v1/structured_output/backend_guidance.py +++ b/vllm/v1/structured_output/backend_guidance.py @@ -10,6 +10,7 @@ import torch from transformers import MistralCommonBackend +from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger from vllm.sampling_params import SamplingParams from vllm.utils.import_utils import LazyLoader @@ -269,7 +270,7 @@ def _process_schema( begin: str = s["begin"] trig = next((t for t in triggers if begin.startswith(t)), None) if trig is None: - raise ValueError( + raise VLLMValidationError( f"Trigger {begin} not found in triggers {triggers}" ) tags.append( @@ -281,7 +282,9 @@ def _process_schema( ) ) if not tags: - raise ValueError("No structural tags found in the grammar spec.") + raise VLLMValidationError( + "No structural tags found in the grammar spec." + ) return llguidance.StructTag.to_grammar(tags) else: logger.error( @@ -300,7 +303,10 @@ def validate_guidance_grammar( if sampling_params.structured_outputs is None: return tp, grm = get_structured_output_key(sampling_params.structured_outputs) - guidance_grm = serialize_guidance_grammar(tp, grm) + try: + guidance_grm = serialize_guidance_grammar(tp, grm) + except (ValueError, KeyError, TypeError) as e: + raise VLLMValidationError(f"Invalid grammar specification: {e}") from e err = llguidance.LLMatcher.validate_grammar(guidance_grm, tokenizer) if err: - raise ValueError(f"Grammar error: {err}") + raise VLLMValidationError(f"Grammar error: {err}") diff --git a/vllm/v1/structured_output/backend_lm_format_enforcer.py b/vllm/v1/structured_output/backend_lm_format_enforcer.py index 898aaa136e40..36eb5eb159e1 100644 --- a/vllm/v1/structured_output/backend_lm_format_enforcer.py +++ b/vllm/v1/structured_output/backend_lm_format_enforcer.py @@ -9,6 +9,7 @@ import torch from transformers import PreTrainedTokenizerBase +from vllm.exceptions import VLLMValidationError from vllm.sampling_params import SamplingParams from vllm.utils.import_utils import LazyLoader from vllm.utils.torch_utils import PIN_MEMORY @@ -166,7 +167,7 @@ def validate_structured_output_request_lm_format_enforcer(params: SamplingParams so_params.regex, ) except Exception as err: - raise ValueError( + raise VLLMValidationError( f"Failed to compile regex for lm-format-enforcer: {err}" ) from err return @@ -176,19 +177,19 @@ def validate_structured_output_request_lm_format_enforcer(params: SamplingParams # make sure schema is valid json json.loads(so_params.json) except json.JSONDecodeError as e: - raise ValueError("Invalid JSON grammar specification.") from e + raise VLLMValidationError("Invalid JSON grammar specification.") from e else: try: json.dumps(so_params.json) except Exception as e: - raise ValueError( + raise VLLMValidationError( f"Error serializing structured outputs jsonschema: {e}" ) from e return elif so_params.choice: return elif so_params.grammar: - raise ValueError( + raise VLLMValidationError( "LM Format Enforcer structured outputs backend " "does not support grammar specifications" ) diff --git a/vllm/v1/structured_output/backend_outlines.py b/vllm/v1/structured_output/backend_outlines.py index e66ef6361a70..dd460f96ec91 100644 --- a/vllm/v1/structured_output/backend_outlines.py +++ b/vllm/v1/structured_output/backend_outlines.py @@ -13,6 +13,7 @@ import torch from regex import escape as regex_escape +from vllm.exceptions import VLLMValidationError from vllm.sampling_params import SamplingParams from vllm.utils.import_utils import LazyLoader from vllm.utils.torch_utils import PIN_MEMORY @@ -186,22 +187,27 @@ def validate_structured_output_request_outlines(params: SamplingParams): json.loads(so_params.json) schema = so_params.json except json.JSONDecodeError as e: - raise ValueError("Invalid JSON grammar specification.") from e + raise VLLMValidationError("Invalid JSON grammar specification.") from e else: try: schema = json.dumps(so_params.json) except Exception as e: - raise ValueError( + raise VLLMValidationError( f"Error serializing structured outputs jsonschema: {e}" ) from e - pattern = json_schema.build_regex_from_schema(schema) + try: + pattern = json_schema.build_regex_from_schema(schema) + except Exception as e: + raise VLLMValidationError( + f"Failed to transform json schema into a regex: {e}" + ) from e validate_regex_is_buildable(pattern) elif so_params.choice: choices = [regex_escape(str(choice)) for choice in so_params.choice] regex = "(" + "|".join(choices) + ")" validate_regex_is_buildable(regex) elif so_params.grammar: - raise ValueError( + raise VLLMValidationError( "Outlines structured outputs backend " "does not support grammar specifications" ) @@ -315,19 +321,19 @@ def validate_regex_is_buildable(pattern: str) -> None: parsed = sre_parse.parse(pattern) except sre_constants.error as e: - raise ValueError(f"Error parsing regex: {e}") from e + raise VLLMValidationError(f"Error parsing regex: {e}") from e try: _check_unsupported(parsed) except ValueError as e: - raise ValueError( + raise VLLMValidationError( f"Regex uses unsupported feature for structured outputs: {e}. " "Only basic matching constructs are supported—lookarounds, " "backreferences, and unicode boundaries are not." ) from e if _prefix_needs_context(parsed): - raise ValueError( + raise VLLMValidationError( "Regex does not have a anchored universal start state" "This means that the Regex uses anchors (^) or look-arounds " "in a way which requires context before any token is matched." diff --git a/vllm/v1/structured_output/backend_xgrammar.py b/vllm/v1/structured_output/backend_xgrammar.py index e11b1601e118..258b1dff32f1 100644 --- a/vllm/v1/structured_output/backend_xgrammar.py +++ b/vllm/v1/structured_output/backend_xgrammar.py @@ -8,6 +8,7 @@ import torch import vllm.envs +from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger from vllm.sampling_params import SamplingParams from vllm.utils.import_utils import LazyLoader @@ -276,7 +277,7 @@ def check_object(obj: dict[str, Any]) -> bool: def validate_xgrammar_grammar(sampling_params: SamplingParams) -> None: """Validate that the request is supported by structured output. - Raises ValueError if the request is not supported. + Raises VLLMValidationError if the request is not supported. """ if sampling_params.structured_outputs is None: return @@ -298,7 +299,7 @@ def validate_xgrammar_grammar(sampling_params: SamplingParams) -> None: so_params.regex, ) except Exception as err: - raise ValueError( + raise VLLMValidationError( f"Failed to transform regex into a grammar: {err}" ) from err @@ -307,7 +308,7 @@ def validate_xgrammar_grammar(sampling_params: SamplingParams) -> None: try: xgr.Grammar.from_ebnf(choice_grammar) except Exception as err: - raise ValueError( + raise VLLMValidationError( f"Failed to transform choices into a grammar: {err}" ) from err so_params.choice = None @@ -319,19 +320,19 @@ def validate_xgrammar_grammar(sampling_params: SamplingParams) -> None: try: schema = json.loads(so_params.json) except json.JSONDecodeError as e: - raise ValueError("Invalid JSON grammar specification.") from e + raise VLLMValidationError("Invalid JSON grammar specification.") from e else: schema = so_params.json if has_xgrammar_unsupported_json_features(schema): - raise ValueError( + raise VLLMValidationError( "The provided JSON schema contains features not supported by xgrammar." ) try: xgr.Grammar.from_json_schema(schema) except Exception as err: - raise ValueError( + raise VLLMValidationError( f"Failed to transform json schema into a grammar: {err}" ) from err return @@ -342,7 +343,7 @@ def validate_xgrammar_grammar(sampling_params: SamplingParams) -> None: try: so_params.grammar = convert_lark_to_ebnf(so_params.grammar) except ValueError as e: - raise ValueError( + raise VLLMValidationError( "Failed to convert the grammar from Lark to EBNF. " ) from e @@ -351,7 +352,7 @@ def validate_xgrammar_grammar(sampling_params: SamplingParams) -> None: # parse the grammar, but we aren't compiling it. xgr.Grammar.from_ebnf(so_params.grammar) except Exception as e: - raise ValueError("Invalid grammar specification.") from e + raise VLLMValidationError("Invalid grammar specification.") from e return if so_params.structural_tag: @@ -372,4 +373,4 @@ def validate_xgrammar_grammar(sampling_params: SamplingParams) -> None: else: xgr.Grammar.from_structural_tag(so_params.structural_tag) except Exception as e: - raise ValueError("Invalid structural tag specification.") from e + raise VLLMValidationError("Invalid structural tag specification.") from e From 70aaec832bde70ada17cbad2061105bbd07541dc Mon Sep 17 00:00:00 2001 From: Do_it_now_! Date: Sun, 16 Aug 2026 16:05:17 +0800 Subject: [PATCH 008/839] [Bugfix][Anthropic] Return 4xx for client-caused errors in /v1/messages (#52246) Signed-off-by: Do_it_now_! --- .../test_anthropic_messages_conversion.py | 95 +++++++++++++++++++ .../test_error_sanitization.py | 1 - vllm/entrypoints/anthropic/api_router.py | 24 +---- 3 files changed, 100 insertions(+), 20 deletions(-) diff --git a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py index a54e13666f68..1cd7227bdc99 100644 --- a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py +++ b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py @@ -15,12 +15,14 @@ import json from argparse import Namespace from http import HTTPStatus +from typing import Annotated from unittest.mock import MagicMock import pytest from fastapi import FastAPI from fastapi.exceptions import RequestValidationError from fastapi.testclient import TestClient +from pydantic import BaseModel, Field, ValidationError from vllm.entrypoints.anthropic.api_router import attach_router from vllm.entrypoints.anthropic.protocol import ( @@ -47,6 +49,7 @@ from vllm.entrypoints.serve.exception_handling.handlers.validation import ( validation_exception_handler, ) +from vllm.exceptions import VLLMValidationError _convert = AnthropicServingMessages._convert_anthropic_to_openai_request _img_url = AnthropicServingMessages._convert_image_source_to_url @@ -1458,3 +1461,95 @@ def test_empty_cache_salt_returns_bad_request(self): assert response.status_code == HTTPStatus.BAD_REQUEST handler.create_messages.assert_not_awaited() + + +# ====================================================================== +# Client-caused errors are 4xx, not 500 (Issue #52088) +# ====================================================================== + + +class TestClientErrorResponses: + @staticmethod + def _make_api_app(handler: MagicMock): + app = FastAPI() + attach_router(app) + app.state.args = Namespace(log_error_stack=False) + app.exception_handler(RequestValidationError)(validation_exception_handler) + app.state.anthropic_serving_messages = handler + return app + + @staticmethod + def _request_body() -> dict: + return { + "model": "test-model", + "max_tokens": 1, + "messages": [{"role": "user", "content": "Hello"}], + } + + @staticmethod + def _conversion_error() -> ValidationError: + """A real pydantic ValidationError like the one ChatCompletionRequest + construction raises when Anthropic input violates the OpenAI schema.""" + + class _StubRequest(BaseModel): + stop: Annotated[list[str], Field(max_length=4)] | None = None + + with pytest.raises(ValidationError) as exc_info: + _StubRequest(stop=["a"] * 6) + return exc_info.value + + def test_validation_error_returns_bad_request(self): + """A pydantic ValidationError during Anthropic->OpenAI conversion is + surfaced as a 400 BadRequestError, not a 500.""" + handler = MagicMock(spec=AnthropicServingMessages) + handler.create_messages.side_effect = self._conversion_error() + + app = self._make_api_app(handler) + with TestClient(app, raise_server_exceptions=False) as client: + response = client.post("/v1/messages", json=self._request_body()) + + assert response.status_code == HTTPStatus.BAD_REQUEST + body = response.json() + assert body["type"] == "error" + assert body["error"]["type"] == "BadRequestError" + assert "at most 4 items" in body["error"]["message"] + + def test_vllm_client_error_returns_bad_request(self): + """VLLMClientError raised by the serving layer maps to 400.""" + handler = MagicMock(spec=AnthropicServingMessages) + handler.create_messages.side_effect = VLLMValidationError( + "Invalid value for stop", parameter="stop" + ) + + app = self._make_api_app(handler) + with TestClient(app, raise_server_exceptions=False) as client: + response = client.post("/v1/messages", json=self._request_body()) + + assert response.status_code == HTTPStatus.BAD_REQUEST + assert response.json()["error"]["type"] == "BadRequestError" + + def test_generic_error_still_returns_internal_server_error(self): + """Non-client errors keep the existing 500 behaviour.""" + handler = MagicMock(spec=AnthropicServingMessages) + handler.create_messages.side_effect = RuntimeError("boom") + + app = self._make_api_app(handler) + with TestClient(app, raise_server_exceptions=False) as client: + response = client.post("/v1/messages", json=self._request_body()) + + assert response.status_code == HTTPStatus.INTERNAL_SERVER_ERROR + assert response.json()["error"]["type"] == "InternalServerError" + + def test_count_tokens_validation_error_returns_bad_request(self): + """The count_tokens route maps conversion errors to 400 as well.""" + handler = MagicMock(spec=AnthropicServingMessages) + handler.count_tokens.side_effect = self._conversion_error() + + app = self._make_api_app(handler) + with TestClient(app, raise_server_exceptions=False) as client: + response = client.post( + "/v1/messages/count_tokens", json=self._request_body() + ) + + assert response.status_code == HTTPStatus.BAD_REQUEST + assert response.json()["error"]["type"] == "BadRequestError" diff --git a/tests/entrypoints/serve/exception_handling/test_error_sanitization.py b/tests/entrypoints/serve/exception_handling/test_error_sanitization.py index f98d870318db..426d1ede745d 100644 --- a/tests/entrypoints/serve/exception_handling/test_error_sanitization.py +++ b/tests/entrypoints/serve/exception_handling/test_error_sanitization.py @@ -116,7 +116,6 @@ class TestAffectedModulesUseSanitize: @pytest.mark.parametrize( "module", [ - "vllm.entrypoints.anthropic.api_router", "vllm.entrypoints.anthropic.serving", "vllm.entrypoints.speech_to_text.realtime.connection", ], diff --git a/vllm/entrypoints/anthropic/api_router.py b/vllm/entrypoints/anthropic/api_router.py index 7c6d23183d3b..584a2958018e 100644 --- a/vllm/entrypoints/anthropic/api_router.py +++ b/vllm/entrypoints/anthropic/api_router.py @@ -17,7 +17,9 @@ ) from vllm.entrypoints.anthropic.serving import AnthropicServingMessages from vllm.entrypoints.openai.engine.protocol import ErrorResponse -from vllm.entrypoints.serve.exception_handling.utils import sanitize_message +from vllm.entrypoints.serve.exception_handling.error_response import ( + create_error_response, +) from vllm.entrypoints.serve.utils.api_utils import ( load_aware_call, validate_json_request, @@ -71,15 +73,7 @@ async def create_messages(request: AnthropicMessagesRequest, raw_request: Reques generator = await handler.create_messages(request, raw_request) except Exception as e: logger.exception("Error in create_messages: %s", e) - return JSONResponse( - status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, - content=AnthropicErrorResponse( - error=AnthropicError( - type="internal_error", - message=sanitize_message(str(e)), - ) - ).model_dump(), - ) + return translate_error_response(create_error_response(e)) if isinstance(generator, ErrorResponse): return translate_error_response(generator) @@ -117,15 +111,7 @@ async def count_tokens(request: AnthropicCountTokensRequest, raw_request: Reques response = await handler.count_tokens(request, raw_request) except Exception as e: logger.exception("Error in count_tokens: %s", e) - return JSONResponse( - status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, - content=AnthropicErrorResponse( - error=AnthropicError( - type="internal_error", - message=sanitize_message(str(e)), - ) - ).model_dump(), - ) + return translate_error_response(create_error_response(e)) if isinstance(response, ErrorResponse): return translate_error_response(response) From 1b079c40ff9d2598d837f4ed1fc08342fca4fd6e Mon Sep 17 00:00:00 2001 From: Jyan-R <87905051+jyan-R@users.noreply.github.com> Date: Sun, 16 Aug 2026 17:02:20 +0800 Subject: [PATCH 009/839] [Bugfix][Model Runner V2][Spec Decode] Fix off-by-one in bad_words draft-prefix matching (#52311) Signed-off-by: jyan Co-authored-by: jyan Co-authored-by: Claude Fable 5 Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/v1/worker/test_gpu_bad_words.py | 105 +++++++++++++++++++++++++ vllm/v1/worker/gpu/sample/bad_words.py | 4 +- 2 files changed, 108 insertions(+), 1 deletion(-) create mode 100644 tests/v1/worker/test_gpu_bad_words.py diff --git a/tests/v1/worker/test_gpu_bad_words.py b/tests/v1/worker/test_gpu_bad_words.py new file mode 100644 index 000000000000..150c4ed00711 --- /dev/null +++ b/tests/v1/worker/test_gpu_bad_words.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import numpy as np +import pytest +import torch + +pytest.importorskip("triton") +if not torch.cuda.is_available(): + pytest.skip( + "CUDA required for Model Runner V2 bad words tests", + allow_module_level=True, + ) + +from vllm.sampling_params import SamplingParams +from vllm.v1.worker.gpu.sample.bad_words import BadWordsState +from vllm.v1.worker.gpu.states import RequestState + +DEVICE = torch.device("cuda") +VOCAB_SIZE = 128 + +# Committed tokens: prompt [5], output [10, 11]. Draft tokens: [12, 13]. +# The sampler passes input_ids gathered at logits_indices, so local position 0 +# holds the last committed token (11) and draft tokens start at position 1. +PROMPT_LEN = 1 +COMMITTED = [5, 10, 11] +INPUT_IDS = [11, 12, 13] +LOCAL_POS = [0, 1, 2] + + +def _make_state(bad_words_token_ids: list[list[int]]) -> tuple[BadWordsState, int]: + req_states = RequestState( + max_num_reqs=4, + max_model_len=64, + max_num_batched_tokens=16, + num_speculative_steps=4, + vocab_size=VOCAB_SIZE, + device=DEVICE, + ) + req_states.add_request( + req_id="req", + prompt_len=PROMPT_LEN, + all_token_ids=COMMITTED, + num_computed_tokens=len(COMMITTED), + max_tokens=32, + ) + req_states.apply_staged_writes() + + req_idx = req_states.req_id_to_index["req"] + state = BadWordsState(req_states) + state.add_request(req_idx, SamplingParams(_bad_words_token_ids=bad_words_token_ids)) + state.apply_staged_writes() + return state, req_idx + + +def _apply(bad_words_token_ids: list[list[int]]) -> torch.Tensor: + state, req_idx = _make_state(bad_words_token_ids) + num_logits = len(INPUT_IDS) + logits = torch.zeros((num_logits, VOCAB_SIZE), device=DEVICE) + idx_mapping_np = np.array([req_idx], dtype=np.intp) + expanded_idx_mapping = torch.tensor( + [req_idx] * num_logits, dtype=torch.int32, device=DEVICE + ) + state.apply_bad_words( + logits, + expanded_idx_mapping, + idx_mapping_np, + torch.tensor(INPUT_IDS, dtype=torch.int32, device=DEVICE), + torch.tensor(LOCAL_POS, dtype=torch.int32, device=DEVICE), + ) + return logits.cpu() + + +def test_v2_bad_words_prefix_inside_draft_tokens(): + """A prefix matching entirely within the draft tokens must mask the bad + word's last token at the draft position that completes the prefix.""" + out = _apply([[12, 13, 40]]) + expected = torch.zeros_like(out) + expected[2, 40] = -float("inf") + torch.testing.assert_close(out, expected) + + +def test_v2_bad_words_prefix_spanning_committed_and_draft_tokens(): + """A prefix spanning the committed/draft boundary must mask at the row + where the prefix completes, not one draft position later.""" + out = _apply([[11, 12, 30]]) + expected = torch.zeros_like(out) + expected[1, 30] = -float("inf") + torch.testing.assert_close(out, expected) + + +def test_v2_bad_words_no_spurious_match_from_last_committed_token(): + """The last committed token must not be double-counted as the first draft + token; [11, 11] never occurs in output [10, 11] + drafts [12, 13].""" + out = _apply([[11, 11, 50]]) + expected = torch.zeros_like(out) + torch.testing.assert_close(out, expected) + + +def test_v2_bad_words_committed_prefix(): + """Baseline: a fully committed prefix masks at the first row.""" + out = _apply([[10, 11, 60]]) + expected = torch.zeros_like(out) + expected[0, 60] = -float("inf") + torch.testing.assert_close(out, expected) diff --git a/vllm/v1/worker/gpu/sample/bad_words.py b/vllm/v1/worker/gpu/sample/bad_words.py index 768ff30a0f06..3fb9d148ecac 100644 --- a/vllm/v1/worker/gpu/sample/bad_words.py +++ b/vllm/v1/worker/gpu/sample/bad_words.py @@ -153,8 +153,10 @@ def _bad_words_kernel( from_spec_input = actual_pos >= output_len if from_spec_input: + # input_ids at local position 0 is the last committed token; + # draft tokens start at local position 1. spec_offset = actual_pos - output_len - actual = tl.load(input_ids_ptr + cur_req_first_pos + spec_offset) + actual = tl.load(input_ids_ptr + cur_req_first_pos + spec_offset + 1) else: actual = tl.load(output_base + actual_pos) From 84530eb235dc3d866fa7d4588217f2fb53f43e76 Mon Sep 17 00:00:00 2001 From: Chauncey Date: Sun, 16 Aug 2026 17:17:04 +0800 Subject: [PATCH 010/839] [Bugfix][Multimodal] Keep Gemma 4 video frame counts on CPU (#52441) Signed-off-by: chaunceyjiang --- tests/models/multimodal/generation/test_vit_cudagraph.py | 6 ++++++ vllm/model_executor/models/gemma4_mm.py | 2 +- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/models/multimodal/generation/test_vit_cudagraph.py b/tests/models/multimodal/generation/test_vit_cudagraph.py index d927c547f49f..1972064b1d93 100644 --- a/tests/models/multimodal/generation/test_vit_cudagraph.py +++ b/tests/models/multimodal/generation/test_vit_cudagraph.py @@ -286,6 +286,12 @@ def ernie45_vl_chat_template(content: str) -> str: "user\n<|video|>\nDescribe this video in one sentence." "\nmodel\n" ), + # The 16-frame test video produces 1056 vision tokens. Capture only + # the smallest supported bucket that covers it instead of all default + # buckets through max_model_len, which adds unrelated memory pressure. + compilation_config_overrides={ + "encoder_cudagraph_token_budgets": [1120], + }, needs_video_metadata=True, marks=[pytest.mark.core_model], ), diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py index 7263d778961a..aa9669aeb157 100644 --- a/vllm/model_executor/models/gemma4_mm.py +++ b/vllm/model_executor/models/gemma4_mm.py @@ -801,7 +801,7 @@ def _get_mm_fields_config( MultiModalFieldConfig.flat_from_sizes("video", vfc) ), video_frame_counts=MultiModalFieldConfig.batched( - "video", + "video", keep_on_cpu=True ), video_num_soft_tokens=MultiModalFieldConfig.batched( "video", keep_on_cpu=True From 4d2a68d64d9e05921ed5c4099146e768a92d71d5 Mon Sep 17 00:00:00 2001 From: oops-oom Date: Sun, 16 Aug 2026 19:09:23 +0800 Subject: [PATCH 011/839] [Bugfix][Spec Decode][Structured Output] DSpark: fix the grammar bitmask mapping when the draft budget is zero (#52436) Signed-off-by: oops-oom <73481342@qq.com> Co-authored-by: oops-oom <73481342@qq.com> Co-authored-by: Claude Opus 4.8 --- .../spec_decode/test_adaptive_verification.py | 49 +++++++++++++++++ vllm/v1/worker/gpu/model_runner.py | 1 + vllm/v1/worker/gpu/structured_outputs.py | 54 +++++++++++++------ 3 files changed, 89 insertions(+), 15 deletions(-) diff --git a/tests/v1/spec_decode/test_adaptive_verification.py b/tests/v1/spec_decode/test_adaptive_verification.py index 5fc88de9984a..99708f50f8f2 100644 --- a/tests/v1/spec_decode/test_adaptive_verification.py +++ b/tests/v1/spec_decode/test_adaptive_verification.py @@ -9,6 +9,7 @@ from vllm.v1.worker.gpu.spec_decode.adaptive_verification import ( AdaptiveVerificationManager, ) +from vllm.v1.worker.gpu.structured_outputs import _build_grammar_mapping def make_manager( @@ -168,3 +169,51 @@ def test_zero_budget_rebuilds_cpu_cu_num_logits(): assert cu_num_logits_np.dtype == scheduled_cu_num_logits.dtype # The prefill keeps its scheduled tokens; only drafts are dropped. assert np.array_equal(compacted, np.array([1, 1, 40], dtype=np.int32)) + + +def test_zero_budget_keeps_one_grammar_row_per_scheduled_draft(): + # The scheduler sizes the grammar bitmask from the *scheduled* drafts + # (len(drafts) + 1 rows per request), but a zero budget rewrites + # cu_num_logits_np to bonus-only. Deriving the bitmask -> logits mapping + # from those rewritten offsets drops rows and trips the + # `num_masks == len(mapping)` assert in apply_grammar_bitmask. + manager = make_manager( + np.array([[0.9, 0.9], [0.9, 0.9], [1.0, 1.0]], dtype=np.float32), + np.ones(64), + ) + manager.req_states.req_id_to_index["prefill"] = 2 + manager.req_states.num_computed_tokens_np = np.zeros(3, dtype=np.int32) + manager.req_states.prefill_len.np = np.array([0, 0, 60], dtype=np.int32) + manager._max_total_logits = 2 # < 3 requests * 1 bonus token + + scheduled_spec_decode_tokens = {"low": [1, 2], "high": [3, 4]} + manager.get_num_tokens( + {"low": 3, "high": 3, "prefill": 40}, scheduled_spec_decode_tokens + ) + assert manager._batch_budget[2] == 0 + + req_ids = ["low", "high", "prefill"] + num_draft_tokens_per_req = np.array([2, 2, 0], dtype=np.int32) + _, cu_num_logits_np = manager.compact_batch( + num_draft_tokens_per_req, + np.array([3, 3, 40], dtype=np.int32), + np.array([0, 3, 6, 7], dtype=np.int32), + ) + + mask_stride = manager.num_speculative_steps + manager.num_bonus_tokens + mapping = _build_grammar_mapping( + req_ids, + req_ids, + cu_num_logits_np, + num_draft_tokens_per_req, + manager.num_bonus_tokens, + mask_stride, + ) + + num_bitmask_rows = sum( + len(scheduled_spec_decode_tokens.get(req_id, ())) + 1 for req_id in req_ids + ) + assert len(mapping) == num_bitmask_rows + # (request, position) keys, so the kernel can mask rows the compacted + # device layout no longer has room for. + assert mapping == [0, 1, 2, 3, 4, 5, 6] diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index abbe6d5d3de6..d68666cf4a0c 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -425,6 +425,7 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None: vocab_size=self.vocab_size, device=self.device, mask_stride=self.decode_query_len, + num_bonus_tokens=self.model_state.num_new_sampled_tokens_per_step, ) if self.is_pooling_model and self.is_last_pp_rank: diff --git a/vllm/v1/worker/gpu/structured_outputs.py b/vllm/v1/worker/gpu/structured_outputs.py index 221163ae3652..8437b0aed9a4 100644 --- a/vllm/v1/worker/gpu/structured_outputs.py +++ b/vllm/v1/worker/gpu/structured_outputs.py @@ -10,6 +10,32 @@ from vllm.v1.worker.gpu.input_batch import InputBatch +def _build_grammar_mapping( + req_ids: list[str], + grammar_req_ids: list[str], + cu_num_logits_np: np.ndarray, + num_draft_tokens_per_req: np.ndarray | None, + num_bonus_tokens: int, + mask_stride: int, +) -> list[int]: + mapping: list[int] = [] + req_id_to_idx = {req_id: i for i, req_id in enumerate(req_ids)} + for grammar_req_id in grammar_req_ids: + req_idx = req_id_to_idx[grammar_req_id] + if num_draft_tokens_per_req is None: + num_positions = int( + cu_num_logits_np[req_idx + 1] - cu_num_logits_np[req_idx] + ) + else: + # Grammar masks follow the scheduled layout even when adaptive + # verification compacts the actual CPU logit offsets to bonus-only. + num_positions = int(num_draft_tokens_per_req[req_idx]) + num_bonus_tokens + mapping.extend( + req_idx * mask_stride + position for position in range(num_positions) + ) + return mapping + + class StructuredOutputsWorker: def __init__( self, @@ -17,6 +43,7 @@ def __init__( vocab_size: int, device: torch.device, mask_stride: int, + num_bonus_tokens: int, ): self.logits_indices = torch.zeros( max_num_logits, dtype=torch.int32, device=device @@ -27,6 +54,7 @@ def __init__( self.device = device self.copy_stream = torch.cuda.Stream() self.mask_stride = mask_stride + self.num_bonus_tokens = num_bonus_tokens def apply_grammar_bitmask( self, @@ -45,21 +73,17 @@ def apply_grammar_bitmask( ) # Construct bitmask -> logits mapping - mapping: list[int] = [] - req_ids = input_batch.req_ids - cu_num_logits = input_batch.cu_num_logits_np.tolist() - req_id_to_idx = {req_id: i for i, req_id in enumerate(req_ids)} - for grammar_req_id in grammar_req_ids: - req_idx = req_id_to_idx[grammar_req_id] - logits_start_idx = cu_num_logits[req_idx] - logits_end_idx = cu_num_logits[req_idx + 1] - # Key by (request, position) rather than absolute logit index: - # adaptive verification finalizes per-request logit offsets on - # device, so the kernel resolves them from the GPU cu_num_logits. - mapping.extend( - req_idx * self.mask_stride + position - for position in range(logits_end_idx - logits_start_idx) - ) + # Key by (request, position) rather than absolute logit index: + # adaptive verification finalizes per-request logit offsets on + # device, so the kernel resolves them from the GPU cu_num_logits. + mapping = _build_grammar_mapping( + input_batch.req_ids, + grammar_req_ids, + input_batch.cu_num_logits_np, + input_batch.num_draft_tokens_per_req, + self.num_bonus_tokens, + self.mask_stride, + ) # Asynchronously copy the mapping to GPU. with torch.cuda.stream(self.copy_stream): From 836aac92ffdaa337083934181cb6d00b64b2a1a6 Mon Sep 17 00:00:00 2001 From: Chauncey Date: Sun, 16 Aug 2026 21:15:47 +0800 Subject: [PATCH 012/839] [Perf][DSV4] Optimize sparse top-k metadata kernels for higher prefill throughput (#52084) Signed-off-by: chaunceyjiang --- vllm/models/deepseek_v4/common/ops/cache_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/models/deepseek_v4/common/ops/cache_utils.py b/vllm/models/deepseek_v4/common/ops/cache_utils.py index 811bf21d4975..2d085f88fa29 100644 --- a/vllm/models/deepseek_v4/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v4/common/ops/cache_utils.py @@ -580,7 +580,7 @@ def combine_topk_swa_indices( return combined_indices, combined_lens -_COMBINE_TOPK_SWA_NUM_WORKERS = 128 +_COMBINE_TOPK_SWA_NUM_WORKERS = 256 # Representative pointer alignment variants for Triton pointer specialization. From fe1c317157d4478fdc0e02096447e61305b871e9 Mon Sep 17 00:00:00 2001 From: Shantipriya Parida Date: Sun, 16 Aug 2026 16:21:31 +0300 Subject: [PATCH 013/839] [Bugfix][ROCm] Skip FP8 MLA prefill PS-metadata build for chunked-context batches (#52356) Signed-off-by: Shantipriya Parida --- vllm/v1/attention/backends/mla/rocm_aiter_mla.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py index ccbdf728903a..2f832a1357a9 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py @@ -777,7 +777,11 @@ def build( attn_metadata.reduce_indptr = self._mla_reduce_indptr attn_metadata.reduce_final_map = self._mla_reduce_final_map attn_metadata.reduce_partial_map = self._mla_reduce_partial_map - if self._fp8_prefill_enabled and attn_metadata.prefill is not None: + if ( + self._fp8_prefill_enabled + and attn_metadata.prefill is not None + and attn_metadata.prefill.chunked_context is None + ): self._build_fp8_prefill_ps_metadata(attn_metadata, common_attn_metadata) return attn_metadata From 9409f59e0963a033a87a6847d6cafef54a64a79c Mon Sep 17 00:00:00 2001 From: Dakai An <77474977+andakai@users.noreply.github.com> Date: Sun, 16 Aug 2026 23:15:14 +0800 Subject: [PATCH 014/839] [Core] Add CuMemAllocator.discard() for tag-selective GPU memory release (#52514) Signed-off-by: AlanFokCo Signed-off-by: Dakai An Co-authored-by: AlanFokCo Co-authored-by: OpenAI Codex --- tests/basic_correctness/test_mem.py | 35 +++++++++++++++++++++ vllm/device_allocator/__init__.py | 2 ++ vllm/device_allocator/cumem.py | 46 +++++++++++++++++++++++++++ vllm/device_allocator/xpumem.py | 49 +++++++++++++++++++++++++++++ 4 files changed, 132 insertions(+) diff --git a/tests/basic_correctness/test_mem.py b/tests/basic_correctness/test_mem.py index 0618562b3ec5..696d2f66bc5b 100644 --- a/tests/basic_correctness/test_mem.py +++ b/tests/basic_correctness/test_mem.py @@ -75,6 +75,41 @@ def test_basic_cumem(): assert torch.allclose(output, torch.ones_like(output) * 3) +@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn") +def test_discard_tags(): + """Test that discard(tags) selectively frees GPU memory for specific + tags while keeping other tags mapped and usable.""" + allocator = get_mem_allocator_instance() + + with allocator.use_memory_pool("weights"): + weights = torch.ones(1024, 1024, device=DEVICE_TYPE) + + with allocator.use_memory_pool("kv_cache"): + kv = torch.ones(512, 512, device=DEVICE_TYPE) + + free_bytes = torch.accelerator.get_memory_info()[0] + + # Discard kv_cache only — weights should remain valid + allocator.discard("kv_cache") + + free_bytes_after_discard = torch.accelerator.get_memory_info()[0] + assert free_bytes_after_discard > free_bytes + + # Weights are still usable + assert torch.allclose(weights, torch.ones_like(weights)) + + # Wake up and verify kv_cache is remapped (zeroed content) + allocator.wake_up() + # After wake_up the VA is remapped; content is not preserved + # but the allocation is valid + assert kv.shape == (512, 512) + + # Full sleep/wake cycle still works after discard + allocator.sleep(offload_tags="weights") + allocator.wake_up() + assert torch.allclose(weights, torch.ones_like(weights)) + + @create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn") @pytest.mark.skipif(current_platform.is_xpu(), reason="CUDA graph not supported on XPU") def test_cumem_with_cudagraph(): diff --git a/vllm/device_allocator/__init__.py b/vllm/device_allocator/__init__.py index 66e8b146d29b..02d51603163e 100644 --- a/vllm/device_allocator/__init__.py +++ b/vllm/device_allocator/__init__.py @@ -27,6 +27,8 @@ def use_memory_pool(self, tag: str | None = None) -> AbstractContextManager: ... def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None: ... + def discard(self, tags: tuple[str, ...] | str) -> None: ... + def wake_up(self, tags: list[str] | None = None) -> None: ... def get_current_usage(self) -> int: ... diff --git a/vllm/device_allocator/cumem.py b/vllm/device_allocator/cumem.py index 2cb9805bae39..a16f1b01486e 100644 --- a/vllm/device_allocator/cumem.py +++ b/vllm/device_allocator/cumem.py @@ -247,8 +247,15 @@ def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None: total_bytes = 0 backup_bytes = 0 + has_policy_conflict = False for ptr, data in self.pointer_to_data.items(): + if data.is_asleep: + requests_offload = data.tag in offload_tags + was_offloaded = data.cpu_backup_tensor is not None + if requests_offload != was_offloaded: + has_policy_conflict = True + continue handle = data.handle total_bytes += handle[1] if data.tag in offload_tags: @@ -277,9 +284,46 @@ def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None: (total_bytes - backup_bytes) / 1024**3, ) + if has_policy_conflict: + logger.warning( + "CuMemAllocator: sleep cannot change the policy of " + "already-asleep allocations; the existing policy was kept." + ) + gc.collect() torch.cuda.empty_cache() + def discard(self, tags: tuple[str, ...] | str) -> None: + """Discard mapped allocations with the given tags without CPU backup.""" + if isinstance(tags, str): + tags = (tags,) + + discarded_bytes = 0 + has_policy_conflict = False + for data in self.pointer_to_data.values(): + if data.tag not in tags: + continue + if data.is_asleep: + if data.cpu_backup_tensor is not None: + has_policy_conflict = True + continue + torch.accelerator.synchronize(data.handle[0]) + unmap_and_release(data.handle) + data.is_asleep = True + discarded_bytes += data.handle[1] + + logger.info( + "CuMemAllocator: discarded %.2f GiB for tags %s.", + discarded_bytes / 1024**3, + tags, + ) + + if has_policy_conflict: + logger.warning( + "CuMemAllocator: discard cannot change the policy of " + "already-asleep allocations; the existing policy was kept." + ) + def wake_up(self, tags: list[str] | None = None) -> None: """ Wake up the allocator from sleep mode. @@ -295,6 +339,8 @@ def wake_up(self, tags: list[str] | None = None) -> None: torch.accelerator.empty_cache() for ptr, data in self.pointer_to_data.items(): + if not data.is_asleep: + continue if tags is None or data.tag in tags: handle = data.handle create_and_map(handle) diff --git a/vllm/device_allocator/xpumem.py b/vllm/device_allocator/xpumem.py index e0f359b200d2..146ea965db07 100644 --- a/vllm/device_allocator/xpumem.py +++ b/vllm/device_allocator/xpumem.py @@ -174,12 +174,20 @@ def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None: total_bytes = 0 backup_bytes = 0 + has_policy_conflict = False for ptr, data in self.pointer_to_data.items(): + if data.is_asleep: + requests_offload = data.tag in offload_tags + was_offloaded = data.cpu_backup_tensor is not None + if requests_offload != was_offloaded: + has_policy_conflict = True + continue size_in_bytes = data.handle[1] total_bytes += size_in_bytes if data.tag not in offload_tags: unmap_and_release(data.handle) + data.is_asleep = True continue backup_bytes += size_in_bytes @@ -201,6 +209,7 @@ def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None: data.cpu_backup_tensor = cpu_backup_tensor unmap_and_release(data.handle) + data.is_asleep = True logger.info( "XpuMemAllocator: sleep freed %.2f GiB memory in total, of which " @@ -211,16 +220,56 @@ def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None: (total_bytes - backup_bytes) / 1024**3, ) + if has_policy_conflict: + logger.warning( + "XpuMemAllocator: sleep cannot change the policy of " + "already-asleep allocations; the existing policy was kept." + ) + gc.collect() xpu_empty_cache = getattr(torch.xpu, "empty_cache", None) if callable(xpu_empty_cache): xpu_empty_cache() + def discard(self, tags: tuple[str, ...] | str) -> None: + """Discard mapped allocations with the given tags without CPU backup.""" + if isinstance(tags, str): + tags = (tags,) + + discarded_bytes = 0 + has_policy_conflict = False + for data in self.pointer_to_data.values(): + if data.tag not in tags: + continue + if data.is_asleep: + if data.cpu_backup_tensor is not None: + has_policy_conflict = True + continue + torch.accelerator.synchronize(data.handle[0]) + unmap_and_release(data.handle) + data.is_asleep = True + discarded_bytes += data.handle[1] + + logger.info( + "XpuMemAllocator: discarded %.2f GiB for tags %s.", + discarded_bytes / 1024**3, + tags, + ) + + if has_policy_conflict: + logger.warning( + "XpuMemAllocator: discard cannot change the policy of " + "already-asleep allocations; the existing policy was kept." + ) + def wake_up(self, tags: list[str] | None = None) -> None: for ptr, data in self.pointer_to_data.items(): + if not data.is_asleep: + continue if tags is not None and data.tag not in tags: continue create_and_allocate(data.handle) + data.is_asleep = False cpu_backup_tensor = data.cpu_backup_tensor if cpu_backup_tensor is None: From 83f591d7f694a3ca3ae3bf22d646e818a1421872 Mon Sep 17 00:00:00 2001 From: Chauncey Date: Sun, 16 Aug 2026 23:24:14 +0800 Subject: [PATCH 015/839] [Perf][DSV4] Optimize global top-k index kernel with compile-time constants (#51967) Signed-off-by: chaunceyjiang Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- vllm/models/deepseek_v4/common/ops/cache_utils.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/vllm/models/deepseek_v4/common/ops/cache_utils.py b/vllm/models/deepseek_v4/common/ops/cache_utils.py index 2d085f88fa29..1f53c96e9367 100644 --- a/vllm/models/deepseek_v4/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v4/common/ops/cache_utils.py @@ -477,15 +477,15 @@ def compute_global_topk_indices_and_lens( @triton.jit def _compute_global_topk_indices_and_lens_kernel( global_topk_indices_ptr, - global_topk_indices_stride, + global_topk_indices_stride: tl.constexpr, topk_lens_ptr, topk_indices_ptr, - topk_indices_stride, - topk, + topk_indices_stride: tl.constexpr, + topk: tl.constexpr, token_to_req_indices_ptr, block_table_ptr, - block_table_stride, - block_size, + block_table_stride: tl.constexpr, + block_size: tl.constexpr, is_valid_token_ptr, TRITON_BLOCK_SIZE: tl.constexpr, ): From 6914d60b1e1a2594f0066fe81e3684574413fbe7 Mon Sep 17 00:00:00 2001 From: akii96 Date: Sun, 16 Aug 2026 19:50:31 +0300 Subject: [PATCH 016/839] [ROCm][Perf] gfx942: use FlyDSL fp8 MQA logits kernel (ROCm/aiter#3913) (#49544) Signed-off-by: Aakif Nawaz --- vllm/v1/attention/ops/rocm_aiter_mla_sparse.py | 11 ++--------- 1 file changed, 2 insertions(+), 9 deletions(-) diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index ed494ce2dae4..94eeaba6afaf 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -563,21 +563,14 @@ def rocm_fp8_mqa_logits( Logits tensor of shape [M, N], dtype `torch.float32`. """ - # TODO(ganyi): Temporarily workaround, will remove the module check and reference - # path after aiter merge this kernel into main from vllm._aiter_ops import rocm_aiter_ops k_fp8, scale = kv - # Temporarily route gfx942 to the vendored ROCm/aiter#3257 workaround. - # Remove this branch once vLLM bumps AITER to a version that includes - # ROCm/aiter#3257. if _ON_GFX942 and rocm_aiter_ops.is_enabled(): - from vllm.v1.attention.ops.triton_fp8_mqa_logits import ( - fp8_mqa_logits_gfx942, - ) + from aiter.ops.flydsl import flydsl_fp8_mqa_logits - return fp8_mqa_logits_gfx942( + return flydsl_fp8_mqa_logits( q, k_fp8, scale, weights, cu_seqlen_ks, cu_seqlen_ke ) From 7d7b6f26f4120b5db10d0a697a878ad605f09579 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Sun, 16 Aug 2026 13:10:16 -0400 Subject: [PATCH 017/839] [Refactor] Remove dead code for quantization (#52221) Signed-off-by: yewentao256 --- tests/kernels/quantization/test_block_int8.py | 75 ------ vllm/model_executor/kernels/linear/base.py | 24 +- .../layers/fused_moe/oracle/mxfp4.py | 8 - .../layers/quantization/auto_awq.py | 25 -- .../schemes/compressed_tensors_w4a8_fp8.py | 1 - .../schemes/compressed_tensors_w4a8_int.py | 1 - .../layers/quantization/humming.py | 22 -- .../layers/quantization/inc/inc.py | 1 - .../inc/schemes/inc_wna16_linear.py | 4 - .../layers/quantization/modelopt.py | 1 - .../layers/quantization/torchao.py | 3 - .../layers/quantization/utils/int8_utils.py | 223 ------------------ .../layers/quantization/utils/mxfp4_utils.py | 7 - .../layers/quantization/utils/ocp_mx_utils.py | 1 - 14 files changed, 1 insertion(+), 395 deletions(-) delete mode 100644 tests/kernels/quantization/test_block_int8.py diff --git a/tests/kernels/quantization/test_block_int8.py b/tests/kernels/quantization/test_block_int8.py deleted file mode 100644 index 4a12ecd50dad..000000000000 --- a/tests/kernels/quantization/test_block_int8.py +++ /dev/null @@ -1,75 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# Adapted from https://github.com/sgl-project/sglang/blob/main/test/srt/test_block_int8.py -import itertools - -import pytest -import torch - -from tests.kernels.quant_utils import native_w8a8_block_matmul -from vllm.config import VllmConfig -from vllm.model_executor.layers.quantization.utils.int8_utils import ( - w8a8_block_int8_matmul, -) -from vllm.platforms import current_platform - -if not (current_platform.is_cuda_alike() or current_platform.is_xpu()): - pytest.skip( - "INT8 Triton kernels require a CUDA-alike or XPU device", - allow_module_level=True, - ) - -if current_platform.is_cuda_alike() and not current_platform.has_device_capability( - (7, 0) -): - pytest.skip("INT8 Triton requires CUDA 7.0 or higher", allow_module_level=True) - -vllm_config = VllmConfig() - -DTYPES = [torch.half, torch.bfloat16] -M = [1, 33, 64, 222] -N = [128, 1024] -K = [256, 4096] -# BLOCK_SIZE = [[64, 64], [64, 128], [128, 64], [128, 128]] -BLOCK_SIZE = [[128, 128]] -SEEDS = [0] - - -@pytest.mark.parametrize( - "M,N,K,block_size,out_dtype,seed", - itertools.product(M, N, K, BLOCK_SIZE, DTYPES, SEEDS), -) -@torch.inference_mode() -def test_w8a8_block_int8_matmul(M, N, K, block_size, out_dtype, seed): - torch.manual_seed(seed) - device = current_platform.device_type - factor_for_scale = 1e-2 - int8_info = torch.iinfo(torch.int8) - int8_max, int8_min = int8_info.max, int8_info.min - - A_fp32 = torch.rand(M, K, dtype=torch.float32, device=device) - A_fp32 = (A_fp32 - 0.5) * 2 * int8_max - A_fp8 = A_fp32.clamp(min=int8_min, max=int8_max).to(torch.float8_e4m3fn) - - B_fp32 = torch.rand(N, K, dtype=torch.float32, device=device) - B_fp32 = (B_fp32 - 0.5) * 2 * int8_max - B_fp8 = B_fp32.clamp(min=int8_min, max=int8_max).to(torch.float8_e4m3fn) - - block_n, block_k = block_size[0], block_size[1] - n_tiles = (N + block_n - 1) // block_n - k_tiles = (K + block_k - 1) // block_k - - As = torch.rand(M, k_tiles, dtype=torch.float32, device=device) * factor_for_scale - Bs = ( - torch.rand(n_tiles, k_tiles, dtype=torch.float32, device=device) - * factor_for_scale - ) - - ref_out = native_w8a8_block_matmul(A_fp8, B_fp8, As, Bs, block_size, out_dtype) - out = w8a8_block_int8_matmul(A_fp8, B_fp8, As, Bs, block_size, out_dtype) - - rel_diff = torch.mean( - torch.abs(out.to(torch.float32) - ref_out.to(torch.float32)) - ) / torch.mean(torch.abs(ref_out.to(torch.float32))) - assert rel_diff < 0.001 diff --git a/vllm/model_executor/kernels/linear/base.py b/vllm/model_executor/kernels/linear/base.py index 416b6ea1c1b6..036b30b7c46d 100644 --- a/vllm/model_executor/kernels/linear/base.py +++ b/vllm/model_executor/kernels/linear/base.py @@ -89,28 +89,6 @@ def from_layer(cls, layer: torch.nn.Module) -> "FP8Params": ) -@dataclass -class Int8Params(Params): - """Int8 layer parameters with typed fields""" - - input_zero_point: torch.Tensor | None - azp_adj: torch.Tensor | None - - INPUT_ZERO_POINT: ClassVar[str] = "input_zero_point" - AZP_ADJ: ClassVar[str] = "azp_adj" - - @classmethod - def from_layer(cls, layer: torch.nn.Module) -> "Int8Params": - """Extract parameters from layer""" - return cls( - weight=getattr(layer, cls.WEIGHT), - weight_scale=getattr(layer, cls.WEIGHT_SCALE), - input_scale=getattr(layer, cls.INPUT_SCALE, None), - input_zero_point=getattr(layer, cls.INPUT_ZERO_POINT, None), - azp_adj=getattr(layer, cls.AZP_ADJ, None), - ) - - _ParamsT = TypeVar("_ParamsT", bound=Params) _ConfigT = TypeVar("_ConfigT", bound=MMLinearLayerConfig) @@ -130,7 +108,7 @@ class MMLinearKernel(ABC, Generic[_ConfigT, _ParamsT]): Typical Usage: 1. Define a config dataclass inheriting from MMLinearLayerConfig - 2. Define a params dataclass inheriting from Params (or FP8Params/Int8Params) + 2. Define a params dataclass inheriting from Params 3. Subclass MMLinearKernel with your config and params types 4. Implement all abstract methods 5. Register the kernel with the quantization method diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index d09587b6d5f6..e11833197032 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -130,14 +130,6 @@ class Mxfp4MoeBackend(Enum): HUMMING = "HUMMING" -# AITER backends group -AITER_BACKENDS = ( - Mxfp4MoeBackend.AITER_MXFP4_BF16, - Mxfp4MoeBackend.AITER_MXFP4_FP8, - Mxfp4MoeBackend.AITER_MXFP4_MXFP4, -) - - # Backends that share the same TRTLLM weight format TRTLLM_BACKENDS = ( Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16, diff --git a/vllm/model_executor/layers/quantization/auto_awq.py b/vllm/model_executor/layers/quantization/auto_awq.py index ef81b0b876df..713bd7bd5e8d 100644 --- a/vllm/model_executor/layers/quantization/auto_awq.py +++ b/vllm/model_executor/layers/quantization/auto_awq.py @@ -358,31 +358,6 @@ def get_quant_method( return None - @classmethod - def is_awq_marlin_compatible(cls, quant_config: dict[str, Any]): - # Extract data from quant config. - quant_method = quant_config.get("quant_method", "").lower() - num_bits = quant_config.get("bits") - group_size = quant_config.get("group_size") - zero_point = quant_config.get("zero_point") - - if not (current_platform.is_cuda_alike() or current_platform.is_cpu()): - return False - - if quant_method != "awq": - return False - - # If we cannot find the info needed in the config, cannot convert. - if num_bits is None or group_size is None or zero_point is None: - return False - - if num_bits not in cls.TYPE_MAP: - return False - - return check_marlin_supported( - quant_type=cls.TYPE_MAP[num_bits], group_size=group_size, has_zp=zero_point - ) - def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"): if self.modules_to_not_convert: self.modules_to_not_convert = hf_to_vllm_mapper.apply_list( diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_fp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_fp8.py index 22c3539e9aec..2c121c6edfe9 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_fp8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_fp8.py @@ -32,7 +32,6 @@ W4A8_SUPPORTED_TYPES_MAP = { 4: scalar_types.int4, } -W4A8_SUPPORTED_BITS = list(W4A8_SUPPORTED_TYPES_MAP.keys()) class CompressedTensorsW4A8Fp8(CompressedTensorsScheme): diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_int.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_int.py index 77933ea2c736..a26a35f4e313 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_int.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_int.py @@ -27,7 +27,6 @@ W4A8_SUPPORTED_TYPES_MAP = { 4: scalar_types.int4, } -W4A8_SUPPORTED_BITS = list(W4A8_SUPPORTED_TYPES_MAP.keys()) class CompressedTensorsW4A8Int(CompressedTensorsScheme): diff --git a/vllm/model_executor/layers/quantization/humming.py b/vllm/model_executor/layers/quantization/humming.py index b814205922c6..ea00c0f3c408 100644 --- a/vllm/model_executor/layers/quantization/humming.py +++ b/vllm/model_executor/layers/quantization/humming.py @@ -1,6 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import math from collections.abc import Callable from typing import TYPE_CHECKING, Any @@ -61,11 +60,6 @@ ) -def prepare_padded_shape(shape, x): - padded_shape = math.ceil(shape / x) * x - return padded_shape, padded_shape - shape - - def prepare_param(tensor, name, extra_attrs): extra_attrs = extra_attrs.copy() scale_type = extra_attrs.pop("scale_type", None) @@ -129,22 +123,6 @@ def prepare_moe_param(tensor: torch.Tensor, name: str, extra_attrs: dict[str, An return param -def may_pad_loaded_weight(param, loaded_weight): - pad_shape = getattr(param, "pad_shape", None) - if pad_shape is None: - return loaded_weight - value = 1 if loaded_weight.dtype == torch.float8_e8m0fnu else 0 - padding = [] - for x in pad_shape[::-1][: loaded_weight.ndim]: - padding += [0, x] - loaded_weight = torch.nn.functional.pad( - input=loaded_weight, - pad=padding, - value=value, - ) - return loaded_weight - - def compressed_tensors_get_config(config: dict[str, Any], key: str): assert key in ["weights", "input_activations"] target_group_config = None diff --git a/vllm/model_executor/layers/quantization/inc/inc.py b/vllm/model_executor/layers/quantization/inc/inc.py index 219fc2b04692..ff8eae3a4649 100644 --- a/vllm/model_executor/layers/quantization/inc/inc.py +++ b/vllm/model_executor/layers/quantization/inc/inc.py @@ -53,7 +53,6 @@ class INCConfig(QuantizationConfig): MXFP8_GROUP_SIZE = 32 MXFP8_DATA_TYPE = "mx_fp" MXFP8_PACKING_FORMAT = "auto_round:llm_compressor" - MXFP8_SUPPORTED_ACT_DTYPES = {"mx_fp", "mx_fp_rceil"} def __init__( self, diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py index 5c99fd98b54d..42f053b471a9 100644 --- a/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py +++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py @@ -458,7 +458,3 @@ def apply_weights( layer.ark_scale_type, not self.sym, ) - - -class INCXPUW4A16LinearScheme(INCXPULinearMethod): - pass diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index db98b76b3f5d..e24c5656add3 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -112,7 +112,6 @@ # MIXED_PRECISION, "MIXED_PRECISION", ] -KV_CACHE_QUANT_ALGOS = ["FP8", "NVFP4"] class ModelOptKVCacheMethod(BaseKVCacheMethod): diff --git a/vllm/model_executor/layers/quantization/torchao.py b/vllm/model_executor/layers/quantization/torchao.py index 15399cfd39b4..8862578d6f7e 100644 --- a/vllm/model_executor/layers/quantization/torchao.py +++ b/vllm/model_executor/layers/quantization/torchao.py @@ -283,9 +283,6 @@ def get_quant_method( return TorchAOLinearMethod(self) - def get_scaled_act_names(self) -> list[str]: - return [] - def torchao_quantize_param_data( param: torch.Tensor, torchao_config: Any diff --git a/vllm/model_executor/layers/quantization/utils/int8_utils.py b/vllm/model_executor/layers/quantization/utils/int8_utils.py index 4f624cf49630..db95d3588a2a 100644 --- a/vllm/model_executor/layers/quantization/utils/int8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/int8_utils.py @@ -2,17 +2,12 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # Adapted from https://github.com/sgl-project/sglang/blob/4cb53ecd0cffceb6dee5c011a58f65997a86f151/python/sglang/srt/layers/quantization/int8_kernel.py -import functools -import json import logging -import os -from typing import Any import torch from vllm.platforms import current_platform from vllm.triton_utils import tl, triton -from vllm.utils.platform_utils import get_device_name_as_file_name logger = logging.getLogger(__name__) @@ -228,221 +223,3 @@ def per_token_group_quant_int8( ) return x_q, x_s - - -@triton.jit -def _w8a8_block_int8_matmul( - # Pointers to inputs and output - A, - B, - C, - As, - Bs, - # Shape for matmul - M, - N, - K, - # Block size for block-wise quantization - group_n, - group_k, - # Stride for inputs and output - stride_am, - stride_ak, - stride_bk, - stride_bn, - stride_cm, - stride_cn, - stride_As_m, - stride_As_k, - stride_Bs_k, - stride_Bs_n, - # Meta-parameters - BLOCK_SIZE_M: tl.constexpr, - BLOCK_SIZE_N: tl.constexpr, - BLOCK_SIZE_K: tl.constexpr, - GROUP_SIZE_M: tl.constexpr, -): - """Triton-accelerated function used to perform linear operations (dot - product) on input tensors `A` and `B` with block-wise quantization, and - store the result in output tensor `C`. - """ - - pid = tl.program_id(axis=0) - num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) - num_pid_in_group = GROUP_SIZE_M * num_pid_n - group_id = pid // num_pid_in_group - first_pid_m = group_id * GROUP_SIZE_M - group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) - pid_m = first_pid_m + (pid % group_size_m) - pid_n = (pid % num_pid_in_group) // group_size_m - - offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N - offs_k = tl.arange(0, BLOCK_SIZE_K) - a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) - b_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) - - As_ptrs = As + offs_am * stride_As_m - offs_bsn = offs_bn // group_n - Bs_ptrs = Bs + offs_bsn * stride_Bs_n - - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) - for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): - a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0) - b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0) - - k_start = k * BLOCK_SIZE_K - offs_ks = k_start // group_k - a_s = tl.load(As_ptrs + offs_ks * stride_As_k) - b_s = tl.load(Bs_ptrs + offs_ks * stride_Bs_k) - - accumulator += tl.dot(a, b).to(tl.float32) * a_s[:, None] * b_s[None, :] - a_ptrs += BLOCK_SIZE_K * stride_ak - b_ptrs += BLOCK_SIZE_K * stride_bk - - if C.dtype.element_ty == tl.bfloat16: - c = accumulator.to(tl.bfloat16) - elif C.dtype.element_ty == tl.float16: - c = accumulator.to(tl.float16) - else: - c = accumulator.to(tl.float32) - - offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) - c_ptrs = C + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] - c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) - tl.store(c_ptrs, c, mask=c_mask) - - -@functools.lru_cache -def get_w8a8_block_int8_configs( - N: int, K: int, block_n: int, block_k: int -) -> dict[int, Any] | None: - """ - Return optimized configurations for the w8a8 block fp8 kernel. - - The return value will be a dictionary that maps an irregular grid of - batch sizes to configurations of the w8a8 block fp8 kernel. To evaluate the - kernel on a given batch size bs, the closest batch size in the grid should - be picked and the associated configuration chosen to invoke the kernel. - """ - - # First look up if an optimized configuration is available in the configs - # directory - device_name = get_device_name_as_file_name() - json_file_name = f"N={N},K={K},device_name={device_name},dtype=int8_w8a8,block_shape=[{block_n}, {block_k}].json" # noqa: E501 - - config_file_path = os.path.join( - os.path.dirname(os.path.realpath(__file__)), "configs", json_file_name - ) - if os.path.exists(config_file_path): - with open(config_file_path) as f: - logger.info( - "Using configuration from %s for W8A8 Block INT8 kernel.", - config_file_path, - ) - # If a configuration has been found, return it - return {int(key): val for key, val in json.load(f).items()} - - # If no optimized configuration is available, we will use the default - # configuration - logger.warning( - ( - "Using default W8A8 Block INT8 kernel config. Performance might " - "be sub-optimal! Config file not found at %s" - ), - config_file_path, - ) - return None - - -def w8a8_block_int8_matmul( - A: torch.Tensor, - B: torch.Tensor, - As: torch.Tensor, - Bs: torch.Tensor, - block_size: list[int], - output_dtype: torch.dtype = torch.float16, -) -> torch.Tensor: - """This function performs matrix multiplication with block-wise - quantization. - - It takes two input tensors `A` and `B` with scales `As` and `Bs`. - The output is returned in the specified `output_dtype`. - - Args: - A: The input tensor, e.g., activation. - B: The input tensor, e.g., weight. - As: The per-token-group quantization scale for `A`. - Bs: The per-block quantization scale for `B`. - block_size: The block size for per-block quantization. It should be - 2-dim, e.g., [128, 128]. - output_dtype: The dtype of the returned tensor. - - Returns: - torch.Tensor: The result of matmul. - """ - assert len(block_size) == 2 - block_n, block_k = block_size[0], block_size[1] - - assert A.shape[-1] == B.shape[-1] - assert A.shape[:-1] == As.shape[:-1] and A.is_contiguous() - assert triton.cdiv(A.shape[-1], block_k) == As.shape[-1] - M = A.numel() // A.shape[-1] - - assert B.ndim == 2 and B.is_contiguous() and Bs.ndim == 2 - N, K = B.shape - assert triton.cdiv(N, block_n) == Bs.shape[0] - assert triton.cdiv(K, block_k) == Bs.shape[1] - - C_shape = A.shape[:-1] + (N,) - C = A.new_empty(C_shape, dtype=output_dtype) - - configs = get_w8a8_block_int8_configs(N, K, block_size[0], block_size[1]) - if configs: - # If an optimal configuration map has been found, look up the - # optimal config - config = configs[min(configs.keys(), key=lambda x: abs(x - M))] - else: - # Default config - # Block-wise quant: BLOCK_SIZE_K must be divisible by block_size[1] - config = { - "BLOCK_SIZE_M": 64, - "BLOCK_SIZE_N": block_size[0], - "BLOCK_SIZE_K": block_size[1], - "GROUP_SIZE_M": 32, - "num_warps": 4, - "num_stages": 3, - } - - def grid(META): - return ( - triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]), - ) - - _w8a8_block_int8_matmul[grid]( - A, - B, - C, - As, - Bs, - M, - N, - K, - block_n, - block_k, - A.stride(-2), - A.stride(-1), - B.stride(1), - B.stride(0), - C.stride(-2), - C.stride(-1), - As.stride(-2), - As.stride(-1), - Bs.stride(1), - Bs.stride(0), - **config, - ) - - return C diff --git a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py index 65348a822214..2d0a4fa48c0d 100644 --- a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py +++ b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py @@ -12,13 +12,6 @@ logger = init_logger(__name__) -# CK's pre-compiled MXFP4 MoE GEMM kernel instances require the -# intermediate_size (after TP split) to be a multiple of this value. -# This arises from FP4 packing (2 values per byte) combined with CK -# tile size constraints. When violated, AITER raises: -# "device_gemm ... does not support this GEMM problem". -CK_MXFP4_MOE_DIM_ALIGNMENT = 256 - def should_use_cdna4_mx_scale_swizzle() -> bool: """Whether to use the CDNA4 swizzled scale layout for mxfp4 on gfx950. diff --git a/vllm/model_executor/layers/quantization/utils/ocp_mx_utils.py b/vllm/model_executor/layers/quantization/utils/ocp_mx_utils.py index a9157cbfb08b..0d56f179b93f 100644 --- a/vllm/model_executor/layers/quantization/utils/ocp_mx_utils.py +++ b/vllm/model_executor/layers/quantization/utils/ocp_mx_utils.py @@ -16,7 +16,6 @@ "mxfp8_e5m2", "mxint8", } -SUPPORTED_OCP_MX_DTYPES = {"mxfp4", "mxfp6_e3m2", "mxfp6_e2m3"} class OCP_MX_Scheme(str, Enum): From 1f0e0bf61210346a6bef4ad75172e62554d1b86c Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Sun, 16 Aug 2026 12:13:02 -0500 Subject: [PATCH 018/839] [Bugfix][Attention] Temporarily disable FA4 head-dim 256 (#52050) Signed-off-by: Taneem Ibrahim Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> --- .../models/multimodal/pooling/test_colpali.py | 5 ++++ vllm/v1/attention/backends/fa_utils.py | 23 ++++++++----------- vllm/v1/attention/backends/flash_attn.py | 1 - 3 files changed, 15 insertions(+), 14 deletions(-) diff --git a/tests/models/multimodal/pooling/test_colpali.py b/tests/models/multimodal/pooling/test_colpali.py index 7c91731065bb..41e6393dc793 100644 --- a/tests/models/multimodal/pooling/test_colpali.py +++ b/tests/models/multimodal/pooling/test_colpali.py @@ -223,7 +223,10 @@ def _run_multimodal_text_query_image_docs_test( max_model_len=4096, enforce_eager=True, gpu_memory_utilization=GPU_MEMORY_UTILIZATION, + attention_backend="FLASH_ATTN", + kernel_config={"enable_flashinfer_autotune": False}, ) as vllm_model: + assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner scores = vllm_model.llm.score(query, image_docs) assert len(scores) == 2 @@ -297,9 +300,11 @@ def _run_multimodal_image_query_text_docs_test( @pytest.mark.parametrize("dtype", [DTYPE]) def test_colpali_multimodal_text_query_image_docs( vllm_runner, + monkeypatch: pytest.MonkeyPatch, model: str, dtype: str, ) -> None: + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") _run_multimodal_text_query_image_docs_test(vllm_runner, model, dtype=dtype) diff --git a/vllm/v1/attention/backends/fa_utils.py b/vllm/v1/attention/backends/fa_utils.py index ae02c3080569..907881e18773 100644 --- a/vllm/v1/attention/backends/fa_utils.py +++ b/vllm/v1/attention/backends/fa_utils.py @@ -72,7 +72,6 @@ def get_flash_attn_version( head_size: int | None = None, head_size_v: int | None = None, has_sinks: bool = False, - requires_local_attention: bool = False, ) -> int | None: if current_platform.is_xpu(): return 2 @@ -169,28 +168,26 @@ def get_flash_attn_version( ) fa_version = 2 - if ( - fa_version == 4 - and device_capability.major >= 10 - and head_size == 256 - and requires_local_attention - ): + # TODO: Restore the `requires_local_attention` restriction when FA4 + # head-dim 256 is re-enabled. + if fa_version == 4 and device_capability.major >= 10 and head_size == 256: logger.warning_once( - "FA4 on Blackwell does not support local attention with " - "head_size=256, defaulting to FA version 2." + "FA4 on Blackwell is temporarily disabled for head_size=256, " + "defaulting to FA version 2." ) fa_version = 2 # FA4 on SM100 (Blackwell) has TMEM capacity limits that restrict - # supported head dimensions to ≤128, with exceptions for 256 and 192/128 (MLA - # prefill). Development of symmetric 192, 384, and 512 support is being tracked - # in https://github.com/Dao-AILab/flash-attention/issues/2456 + # supported head dimensions to ≤128. The 192/128 MLA prefill case is + # supported; 256 is temporarily disabled until upstream supports the + # required features. Development of symmetric 192, 384, and 512 support + # is tracked in https://github.com/Dao-AILab/flash-attention/issues/2456 if ( fa_version == 4 and device_capability.major >= 10 and head_size is not None and head_size > 128 - and not (head_size == 256 or (head_size == 192 and head_size_v == 128)) + and not (head_size == 192 and head_size_v == 128) ): logger.warning_once( "FA4 on Blackwell does not support head_size=%d due to TMEM " diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index 92eab051b4bf..d176bd5e928a 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -905,7 +905,6 @@ def __init__( self.attn_type = attn_type self.vllm_flash_attn_version = get_flash_attn_version( requires_alibi=alibi_slopes is not None, - requires_local_attention=sliding_window is not None, head_size=head_size, has_sinks=sinks is not None, ) From ef43e3101b8fda8cd2b52de150c76b4fc177fad2 Mon Sep 17 00:00:00 2001 From: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> Date: Sun, 16 Aug 2026 15:16:46 -0400 Subject: [PATCH 019/839] [ROCm][DSV4][Perf] Optimize Triton sparse-MLA decode on gfx950 (#52212) Signed-off-by: Fangzhou Ai Signed-off-by: fai Co-authored-by: Nick Hill Co-authored-by: Cursor Grok 4.6 --- .../attention/test_rocm_triton_attn_dsv4.py | 481 +++++++++- tests/kernels/test_compressor_kv_cache.py | 243 ++++- vllm/models/deepseek_v4/amd/rocm.py | 51 +- .../common/ops/fused_compress_quant_cache.py | 23 +- .../v1/attention/ops/rocm_aiter_mla_sparse.py | 846 ++++++++++++++++-- 5 files changed, 1550 insertions(+), 94 deletions(-) diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index 6fe2a3e77587..7a98f3ec7005 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -24,12 +24,27 @@ def _on_split_decode_arch() -> bool: return False +def _on_gfx950() -> bool: + if not current_platform.is_rocm(): + return False + try: + from vllm.platforms.rocm import _ON_GFX950 + + return _ON_GFX950 + except ImportError: + return False + + # The flash-decode split-K decode path is only tuned for AMD gfx942/gfx950; other # architectures take the fallback decode kernel, so its tests are skipped there. requires_split_decode_arch = pytest.mark.skipif( not _on_split_decode_arch(), reason="split-K decode kernel is only tuned for AMD gfx942/gfx950", ) +requires_gfx950 = pytest.mark.skipif( + not _on_gfx950(), + reason="optimized sparse decode partial is gfx950-only", +) NOPE_HEAD_DIM = 448 ROPE_HEAD_DIM = 64 @@ -118,28 +133,48 @@ def _pack_fp8_ds_mla_cache( return cache -def _read_fp8_ds_mla_cache( - cache: torch.Tensor, slot: int, block_size: int, use_fnuz: bool -) -> torch.Tensor: - cache_flat = cache.view(torch.uint8).flatten() +def _poison_fp8_ds_mla_cache_row( + cache: torch.Tensor, block_size: int, slot: int = 0 +) -> None: + flat = cache.flatten() block_idx = slot // block_size pos = slot % block_size block_base = block_idx * cache.stride(0) token_base = block_base + pos * 576 scale_base = block_base + block_size * 576 + pos * 8 + flat[token_base] = 0x7F + flat[scale_base : scale_base + 7] = 255 + flat[token_base + NOPE_HEAD_DIM : token_base + 576].view(torch.bfloat16)[0] = float( + "nan" + ) + + +def _read_fp8_ds_mla_cache_rows( + cache: torch.Tensor, + slots: torch.Tensor, + block_size: int, + use_fnuz: bool, +) -> torch.Tensor: + cache_flat = cache.view(torch.uint8).flatten() + block_idx = slots // block_size + pos = slots % block_size + block_base = block_idx * cache.stride(0) + token_base = block_base + pos * 576 + scale_base = block_base + block_size * 576 + pos * 8 fp8_dtype = torch.float8_e4m3fnuz if use_fnuz else torch.float8_e4m3fn - nope_u8 = cache_flat[token_base : token_base + NOPE_HEAD_DIM] + nope_offsets = torch.arange(NOPE_HEAD_DIM, device=cache.device) + nope_u8 = cache_flat[token_base[:, None] + nope_offsets] nope = nope_u8.view(fp8_dtype).to(torch.float32) + scale_offsets = torch.arange(7, device=cache.device) scales = torch.exp2( - cache_flat[scale_base : scale_base + 7].to(torch.float32) - 127.0 + cache_flat[scale_base[:, None] + scale_offsets].to(torch.float32) - 127.0 ) - nope = nope * scales.repeat_interleave(64) - rope_u8 = cache_flat[ - token_base + NOPE_HEAD_DIM : token_base + NOPE_HEAD_DIM + ROPE_HEAD_DIM * 2 - ] - rope = rope_u8.view(torch.bfloat16).to(torch.float32) - return torch.cat([nope, rope]) + nope = nope * scales.repeat_interleave(64, dim=1) + rope_offsets = torch.arange(ROPE_HEAD_DIM * 2, device=cache.device) + rope_u8 = cache_flat[token_base[:, None] + NOPE_HEAD_DIM + rope_offsets] + rope = rope_u8.contiguous().view(torch.bfloat16).to(torch.float32) + return torch.cat([nope, rope], dim=1) def _ref_sparse_decode_ragged( @@ -158,19 +193,30 @@ def _ref_sparse_decode_ragged( out = torch.empty_like(q_f32) for query_idx in range(q.shape[0]): - row_kv = [ - _read_fp8_ds_mla_cache(main_cache, int(slot), block_size, main_use_fnuz) - for slot in main_rows[query_idx] - ] - if extra_cache is not None and extra_rows is not None: - row_kv.extend( - _read_fp8_ds_mla_cache( - extra_cache, int(slot), block_size, extra_use_fnuz + row_kv = [] + if main_rows[query_idx]: + main_slots = torch.tensor( + main_rows[query_idx], dtype=torch.int64, device=q.device + ) + row_kv.append( + _read_fp8_ds_mla_cache_rows( + main_cache, main_slots, block_size, main_use_fnuz + ) + ) + if extra_cache is not None and extra_rows is not None and extra_rows[query_idx]: + extra_slots = torch.tensor( + extra_rows[query_idx], dtype=torch.int64, device=q.device + ) + row_kv.append( + _read_fp8_ds_mla_cache_rows( + extra_cache, extra_slots, block_size, extra_use_fnuz ) - for slot in extra_rows[query_idx] ) - kv = torch.stack(row_kv).to(q.device) + if not row_kv: + out[query_idx] = 0 + continue + kv = torch.cat(row_kv) for head_idx in range(q.shape[1]): scores = torch.mv(kv, q_f32[query_idx, head_idx]) * scale if attn_sink is not None: @@ -198,6 +244,46 @@ def _ragged_from_rows( ) +def _launch_sparse_decode_reduce( + part_m: torch.Tensor, + part_l: torch.Tensor, + part_acc: torch.Tensor, + adaptive_splits: bool, +) -> torch.Tensor: + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod + + num_queries, num_splits, num_heads = part_m.shape + out = torch.empty( + (num_queries, num_heads, HEAD_DIM), + dtype=torch.bfloat16, + device=part_m.device, + ) + attn_sink = torch.empty(1, dtype=torch.float32, device=part_m.device) + mod._sparse_attn_decode_reduce_kernel[(num_queries, num_heads)]( + part_m, + part_l, + part_acc, + attn_sink, + out, + out.stride(0), + out.stride(1), + part_m.stride(0), + part_m.stride(1), + part_acc.stride(0), + part_acc.stride(1), + part_acc.stride(2), + num_heads, + HAS_ATTN_SINK=False, + ADAPTIVE_SPLITS=adaptive_splits, + COMB_DIM=HEAD_DIM, + BLOCK_H=1, + NUM_SPLITS=num_splits, + SPLITS_PAD=1 << (num_splits - 1).bit_length(), + num_warps=4, + ) + return out + + @torch.inference_mode() def test_paged_mqa_logits_do_not_contain_nan(monkeypatch) -> None: from vllm._aiter_ops import rocm_aiter_ops @@ -312,6 +398,19 @@ def test_compute_global_topk_ragged_indices_and_indptr() -> None: torch.testing.assert_close(actual_lens, expected_lens) +def test_extra_cache_nan_free_provenance_gate(monkeypatch) -> None: + from vllm.models.deepseek_v4.amd import rocm as mod + + monkeypatch.setattr(mod, "_ON_GFX950", True) + assert mod._trust_dsv4_extra_cache_nan_free("fp8_ds_mla", False, True) + assert not mod._trust_dsv4_extra_cache_nan_free("fp8_ds_mla", True, True) + assert not mod._trust_dsv4_extra_cache_nan_free("bfloat16", False, True) + assert not mod._trust_dsv4_extra_cache_nan_free("fp8_ds_mla", False, False) + + monkeypatch.setattr(mod, "_ON_GFX950", False) + assert not mod._trust_dsv4_extra_cache_nan_free("fp8_ds_mla", False, True) + + @torch.inference_mode() def test_sparse_attn_prefill_ragged_kernel() -> None: from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( @@ -366,6 +465,7 @@ def test_sparse_attn_decode_ragged_kernel() -> None: attn_sink = torch.tensor([-0.1, 0.0, 0.1], dtype=torch.float32, device=device) scale = HEAD_DIM**-0.5 + out = torch.empty_like(q) actual = _rocm_sparse_attn_decode_ragged_triton( q=q, main_cache=main_cache, @@ -378,6 +478,7 @@ def test_sparse_attn_decode_ragged_kernel() -> None: extra_cache=extra_cache, extra_indices=extra_indices, extra_indptr=extra_indptr, + out=out, ) expected = _ref_sparse_decode_ragged( q=q, @@ -391,9 +492,111 @@ def test_sparse_attn_decode_ragged_kernel() -> None: main_use_fnuz=main_use_fnuz, ) + assert actual.data_ptr() == out.data_ptr() torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) +@requires_gfx950 +@torch.inference_mode() +def test_sparse_attn_decode_scrubs_untrusted_cache_by_default() -> None: + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + _rocm_sparse_attn_decode_ragged_triton, + ) + + device = torch.device("cuda") + block_size = 4 + main_cache = torch.zeros(1, block_size, 584, dtype=torch.uint8, device=device) + extra_cache = torch.zeros_like(main_cache) + _poison_fp8_ds_mla_cache_row(main_cache, block_size) + _poison_fp8_ds_mla_cache_row(extra_cache, block_size) + indices = torch.zeros(1, dtype=torch.int32, device=device) + indptr = torch.tensor([0, 1], dtype=torch.int32, device=device) + + actual = _rocm_sparse_attn_decode_ragged_triton( + q=torch.ones(1, 1, HEAD_DIM, dtype=torch.bfloat16, device=device), + main_cache=main_cache, + main_indices=indices, + main_indptr=indptr, + scale=HEAD_DIM**-0.5, + attn_sink=None, + nope_head_dim=NOPE_HEAD_DIM, + rope_head_dim=ROPE_HEAD_DIM, + extra_cache=extra_cache, + extra_indices=indices, + extra_indptr=indptr, + ) + + assert not torch.isnan(actual).any() + assert torch.equal(actual, torch.zeros_like(actual)) + + +@pytest.mark.parametrize("on_gfx950", [False, True]) +@torch.inference_mode() +def test_rocm_ragged_graph_buffer_view_tracks_source_width( + monkeypatch, on_gfx950: bool +) -> None: + from vllm.models.deepseek_v4.amd import rocm as rocm_mod + + monkeypatch.setattr(rocm_mod, "_ON_GFX950", on_gfx950) + + indices_buffer = torch.full((16,), -1, dtype=torch.int32) + indptr_buffer = torch.full((3,), -1, dtype=torch.int32) + first_indices = torch.tensor([3, 5, 7], dtype=torch.int32) + first_indptr = torch.tensor([0, 1, 3], dtype=torch.int32) + first_view, first_indptr_view = rocm_mod._copy_ragged_to_graph_buffers( + first_indices, + first_indptr, + indices_buffer, + indptr_buffer, + num_rows=2, + max_entries_per_row=8, + ) + + second_indices = torch.tensor([1, 2, 3, 4, 5, 6], dtype=torch.int32) + second_indptr = torch.tensor([0, 2, 6], dtype=torch.int32) + second_view, second_indptr_view = rocm_mod._copy_ragged_to_graph_buffers( + second_indices, + second_indptr, + indices_buffer, + indptr_buffer, + num_rows=2, + max_entries_per_row=8, + ) + + expected_first_entries = ( + first_indices.numel() if on_gfx950 else indices_buffer.numel() + ) + expected_second_entries = ( + second_indices.numel() if on_gfx950 else indices_buffer.numel() + ) + assert first_view.numel() == expected_first_entries + assert second_view.numel() == expected_second_entries + assert first_view.data_ptr() == second_view.data_ptr() == indices_buffer.data_ptr() + assert first_indptr_view.data_ptr() == second_indptr_view.data_ptr() + assert torch.equal(second_view[: second_indices.numel()], second_indices) + assert torch.equal(second_indptr_view, second_indptr) + + +def test_rocm_capture_metadata_sets_adaptive_marker(monkeypatch) -> None: + from vllm.models.deepseek_v4.amd import rocm as rocm_mod + from vllm.models.deepseek_v4.sparse_mla import ( + DeepseekV4SparseMLAMetadataBuilder, + ) + + metadata = SimpleNamespace(for_cudagraph_capture=False) + monkeypatch.setattr( + DeepseekV4SparseMLAMetadataBuilder, + "build_for_cudagraph_capture", + lambda *_: metadata, + ) + builder = object.__new__(rocm_mod.DeepseekV4ROCMAiterMLASparseMetadataBuilder) + + actual = builder.build_for_cudagraph_capture(SimpleNamespace()) + + assert actual is metadata + assert actual.for_cudagraph_capture is _on_gfx950() + + @requires_split_decode_arch @torch.inference_mode() def test_decode_num_splits_heuristic(monkeypatch) -> None: @@ -408,6 +611,9 @@ def test_decode_num_splits_heuristic(monkeypatch) -> None: # A tiny batch on a large device should split to add parallelism. assert mod._decode_num_splits(2, 1, avg_main_len=256.0, avg_extra_len=0.0) > 1 + # The shared gfx942 selector retains its original 16-split ceiling. + assert mod._decode_num_splits(1, 1, 128.0, 8192.0) == 16 + # The chosen count always stays within the searched [1, 16] range, and a # zero-length workload never splits (no work to parallelize). for num_queries in (1, 4, 24, 224, 1024): @@ -418,6 +624,16 @@ def test_decode_num_splits_heuristic(monkeypatch) -> None: assert mod._decode_num_splits(2, 1, avg_main_len=0.0, avg_extra_len=0.0) >= 1 +@torch.inference_mode() +def test_decode_num_splits_gfx950(monkeypatch) -> None: + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod + + monkeypatch.setattr(mod, "_decode_cu_count", lambda: 256) + assert mod._decode_gfx950_num_splits(1, 1, 128, 8192) == 32 + assert mod._decode_gfx950_num_splits(17, 1, 128, 32) == 4 + assert mod._decode_gfx950_num_splits(512, 1, 128, 7812) == 1 + + @requires_split_decode_arch @pytest.mark.parametrize("num_splits", [1, 2, 3, 4, 8]) @pytest.mark.parametrize("with_extra", [True, False]) @@ -473,7 +689,14 @@ def test_sparse_attn_decode_split_k_kernel( scale = HEAD_DIM**-0.5 # Pin the split count so each parametrized value is exercised deterministically. - monkeypatch.setattr(mod, "_decode_num_splits", lambda *args, **kwargs: num_splits) + split_fn = "_decode_gfx950_num_splits" if _on_gfx950() else "_decode_num_splits" + other_split_fn = ( + "_decode_num_splits" if _on_gfx950() else "_decode_gfx950_num_splits" + ) + monkeypatch.setattr(mod, split_fn, lambda *args, **kwargs: num_splits) + monkeypatch.setattr( + mod, other_split_fn, lambda *args, **kwargs: pytest.fail("wrong selector") + ) actual = mod._rocm_sparse_attn_decode_ragged_triton( q=q, @@ -503,6 +726,218 @@ def test_sparse_attn_decode_split_k_kernel( torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) +@requires_gfx950 +@torch.inference_mode() +def test_sparse_attn_decode_gfx950_adaptive_reduce_ignores_stale_scratch() -> None: + device = torch.device("cuda") + part_m = torch.full((1, 8, 1), torch.finfo(torch.float32).min, device=device) + part_l = torch.zeros_like(part_m) + part_acc = torch.full((1, 8, 1, HEAD_DIM), float("nan"), device=device) + part_m[:, :2] = 0 + part_l[:, :2] = 1 + part_acc[:, 0] = 1 + part_acc[:, 1] = 3 + + actual = _launch_sparse_decode_reduce(part_m, part_l, part_acc, True) + + assert torch.isfinite(actual).all() + assert torch.equal(actual, torch.full_like(actual, 2)) + + +@requires_gfx950 +@pytest.mark.parametrize("extra_len", [0, 1, 31, 32, 33, 63, 64, 65]) +@torch.inference_mode() +def test_sparse_attn_decode_gfx950_outer64_boundaries( + monkeypatch, extra_len: int +) -> None: + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod + + device = torch.device("cuda") + torch.manual_seed(13) + block_size = 4 + num_heads = 16 + num_extra_rows = 80 + q = torch.randn(2, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=device) + q *= 0.125 + main_cache = torch.zeros(1, block_size, 584, dtype=torch.uint8, device=device) + main_indices = torch.empty(0, dtype=torch.int32, device=device) + main_indptr = torch.zeros(3, dtype=torch.int32, device=device) + extra_cache = _pack_fp8_ds_mla_cache( + torch.randn(num_extra_rows, HEAD_DIM, dtype=torch.bfloat16, device=device) + * 0.125, + block_size, + use_fnuz=False, + ) + _poison_fp8_ds_mla_cache_row(extra_cache, block_size) + + raw_row = list(range(1, extra_len + 1)) + if extra_len > 3: + raw_row[3] = -1 + if extra_len > 40: + raw_row[40] = num_extra_rows + if extra_len > 64: + raw_row[64] = num_extra_rows + 1024 + extra_indices, extra_indptr = _ragged_from_rows([raw_row, []], device) + valid_row = [slot for slot in raw_row if 0 <= slot < num_extra_rows] + attn_sink = torch.linspace(-0.1, 0.1, num_heads, dtype=torch.float32, device=device) + + monkeypatch.setattr(mod, "_decode_gfx950_num_splits", lambda *args: 1) + actual = mod._rocm_sparse_attn_decode_ragged_triton( + q=q, + main_cache=main_cache, + main_indices=main_indices, + main_indptr=main_indptr, + scale=HEAD_DIM**-0.5, + attn_sink=attn_sink, + nope_head_dim=NOPE_HEAD_DIM, + rope_head_dim=ROPE_HEAD_DIM, + extra_cache=extra_cache, + extra_indices=extra_indices, + extra_indptr=extra_indptr, + ) + expected = _ref_sparse_decode_ragged( + q=q, + main_cache=main_cache, + main_rows=[[], []], + scale=HEAD_DIM**-0.5, + attn_sink=attn_sink, + block_size=block_size, + extra_cache=extra_cache, + extra_rows=[valid_row, []], + ) + + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) + assert torch.equal(actual[1], torch.zeros_like(actual[1])) + + +@requires_gfx950 +@torch.inference_mode() +def test_sparse_attn_decode_gfx950_graph_replay(monkeypatch) -> None: + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod + + device = torch.device("cuda") + torch.manual_seed(17) + block_size = 64 + num_queries = 16 + num_heads = 16 + num_splits = 8 + extra_per_query = 65 * num_splits + max_extra_per_query = 8192 + q = ( + torch.randn( + num_queries, + num_heads, + HEAD_DIM, + dtype=torch.bfloat16, + device=device, + ) + * 0.125 + ) + main_cache = _pack_fp8_ds_mla_cache( + torch.randn(num_queries, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125, + block_size, + use_fnuz=False, + ) + extra_cache = _pack_fp8_ds_mla_cache( + torch.randn( + num_queries * extra_per_query, + HEAD_DIM, + dtype=torch.bfloat16, + device=device, + ) + * 0.125, + block_size, + use_fnuz=False, + ) + main_rows = [[query_idx] for query_idx in range(num_queries)] + extra_rows = [ + list(range(query_idx * extra_per_query, (query_idx + 1) * extra_per_query)) + for query_idx in range(num_queries) + ] + main_indices, main_indptr = _ragged_from_rows(main_rows, device) + short_extra_rows = [row[:64] for row in extra_rows] + long_indices, long_indptr = _ragged_from_rows(extra_rows, device) + short_indices, short_indptr = _ragged_from_rows(short_extra_rows, device) + extra_indices = torch.full( + (num_queries * max_extra_per_query,), + -1, + dtype=torch.int32, + device=device, + ) + extra_indices[: long_indices.numel()].copy_(long_indices) + extra_indptr = long_indptr.clone() + extra_indices_ptr = extra_indices.data_ptr() + attn_sink = torch.linspace(-0.1, 0.1, num_heads, dtype=torch.float32, device=device) + out = torch.empty_like(q) + + monkeypatch.setattr(mod, "_decode_gfx950_num_splits", lambda *args: num_splits) + + def run_decode() -> torch.Tensor: + return mod._rocm_sparse_attn_decode_ragged_triton( + q=q, + main_cache=main_cache, + main_indices=main_indices, + main_indptr=main_indptr, + scale=HEAD_DIM**-0.5, + attn_sink=attn_sink, + nope_head_dim=NOPE_HEAD_DIM, + rope_head_dim=ROPE_HEAD_DIM, + extra_cache=extra_cache, + extra_indices=extra_indices, + extra_indptr=extra_indptr, + out=out, + extra_cache_nan_free=True, + adaptive_splits=True, + ) + + run_decode() + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured_out = run_decode() + torch.accelerator.synchronize() + captured_long = out.clone() + expected_long = _ref_sparse_decode_ragged( + q=q, + main_cache=main_cache, + main_rows=main_rows, + scale=HEAD_DIM**-0.5, + attn_sink=attn_sink, + block_size=block_size, + extra_cache=extra_cache, + extra_rows=extra_rows, + ) + torch.testing.assert_close(captured_long, expected_long, atol=2e-2, rtol=2e-2) + + extra_indices[: short_indices.numel()].copy_(short_indices) + extra_indptr.copy_(short_indptr) + graph.replay() + torch.accelerator.synchronize() + short_out = out.clone() + expected_short = _ref_sparse_decode_ragged( + q=q, + main_cache=main_cache, + main_rows=main_rows, + scale=HEAD_DIM**-0.5, + attn_sink=attn_sink, + block_size=block_size, + extra_cache=extra_cache, + extra_rows=short_extra_rows, + ) + + assert captured_out.data_ptr() == out.data_ptr() + assert extra_indices.data_ptr() == extra_indices_ptr + assert extra_indices.numel() == num_queries * max_extra_per_query + assert not torch.equal(short_out, captured_long) + torch.testing.assert_close(short_out, expected_short, atol=2e-2, rtol=2e-2) + + extra_indices[: long_indices.numel()].copy_(long_indices) + extra_indptr.copy_(long_indptr) + graph.replay() + torch.accelerator.synchronize() + torch.testing.assert_close(out, expected_long, atol=2e-2, rtol=2e-2) + + # --------------------------------------------------------------------------- # o-projection: fused inverse-RoPE + cached bf16 wo_a (rocm_inv_rope_einsum) # --------------------------------------------------------------------------- diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index c1d9a2106be3..cfdec96e61d0 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -28,16 +28,81 @@ _fused_kv_compress_norm_rope_insert_indexer_attn, _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn, _launch_two_stage_sparse_attn_compressor, + compress_norm_rope_store_triton, ) from vllm.models.deepseek_v4.compressor import _get_c128_boundary from vllm.platforms import current_platform from vllm.v1.attention.backends.mla.compressor_utils import ( get_dspark_swa_index_width, ) +from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + cp_gather_indexer_k_quant_cache_triton, + indexer_k_quant_and_cache_triton, +) from .test_fused_indexer_q_rope_quant import quantize_to_mxfp4 +def _on_gfx950() -> bool: + if not current_platform.is_rocm(): + return False + try: + from vllm.platforms.rocm import _ON_GFX950 + + return _ON_GFX950 + except Exception: + return False + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-only dispatch") +def test_cp_gather_despecialized_kernel_is_gfx950_only(monkeypatch): + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod + + class FakeKernel: + def __init__(self): + self.calls = [] + + def __getitem__(self, grid): + def launch(*args): + self.calls.append((grid, args)) + + return launch + + legacy_kernel = FakeKernel() + gfx950_kernel = FakeKernel() + monkeypatch.setattr(mod, "_cp_gather_indexer_quant_cache_kernel", legacy_kernel) + monkeypatch.setattr( + mod, + "_cp_gather_indexer_quant_cache_gfx950_kernel", + gfx950_kernel, + ) + + k_cache = torch.zeros((4, 1, 132), dtype=torch.uint8) + k_fp8 = torch.empty((5, 128), dtype=current_platform.fp8_dtype()) + k_scale = torch.empty((5, 4), dtype=torch.uint8) + block_table = torch.zeros((2, 7), dtype=torch.int32) + cu_seqlen = torch.tensor([0, 2, 5], dtype=torch.int32) + token_to_seq = torch.tensor([0, 0, 1, 1, 1], dtype=torch.int32) + args = (k_cache, k_fp8, k_scale, block_table, cu_seqlen, token_to_seq) + + monkeypatch.setattr(mod, "_ON_GFX950", True) + mod.cp_gather_indexer_k_quant_cache_triton(*args) + assert len(gfx950_kernel.calls) == 1 + assert not legacy_kernel.calls + gfx950_grid, gfx950_args = gfx950_kernel.calls[0] + assert gfx950_grid == (5,) + assert len(gfx950_args) == 18 + assert gfx950_args[-3:] == (2, 7, 4) + + monkeypatch.setattr(mod, "_ON_GFX950", False) + mod.cp_gather_indexer_k_quant_cache_triton(*args) + assert len(legacy_kernel.calls) == 1 + legacy_grid, legacy_args = legacy_kernel.calls[0] + assert legacy_grid == (5,) + assert len(legacy_args) == 19 + assert legacy_args[-4:] == (5, 2, 7, 4) + + @pytest.mark.parametrize( ("window_size", "num_speculative_tokens", "expected"), [(128, 5, 192), (512, 5, 576), (1024, 0, 1024)], @@ -91,6 +156,129 @@ def _ue8m0_reference(x: torch.Tensor, block_size: int, fp8_max: float): return x_fp8, scales +def _decode_dsv4_cache_row( + cache: torch.Tensor, block_size: int, scrub_nan: bool +) -> torch.Tensor: + flat = cache.flatten() + nope = flat[:448].view(torch.float8_e4m3fn).to(torch.bfloat16) + encoded = flat[block_size * 576 : block_size * 576 + 7] + scales = torch.exp2(encoded.to(torch.float32) - 127.0).to(torch.bfloat16) + nope = nope * scales.repeat_interleave(64) + rope = flat[448:576].view(torch.bfloat16) + decoded = torch.cat((nope, rope)) + if scrub_nan: + decoded = torch.where(decoded == decoded, decoded, 0.0) + return decoded + + +def _assert_nan_free_cache_matches_legacy_scrub( + cache: torch.Tensor, block_size: int +) -> None: + flat = cache.flatten() + scale_base = block_size * 576 + scale_codes = flat[scale_base : scale_base + 8] + assert scale_codes[0].item() == 254 + assert scale_codes[1].item() == 247 + assert scale_codes[:7].max().item() <= 254 + nope_bytes = flat[:448] + assert not ((nope_bytes == 0x7F) | (nope_bytes == 0xFF)).any() + + rope = flat[448:576].view(torch.bfloat16) + assert not torch.isnan(rope).any() + assert torch.isposinf(rope[0]) + assert torch.equal(rope[1:4], torch.zeros_like(rope[1:4])) + + legacy_cache = cache.clone() + legacy_flat = legacy_cache.flatten() + legacy_flat[scale_base] = 255 + legacy_rope = legacy_flat[448:576].view(torch.bfloat16) + legacy_rope[1:4] = float("nan") + canonical = _decode_dsv4_cache_row(cache, block_size, scrub_nan=False) + legacy = _decode_dsv4_cache_row(legacy_cache, block_size, scrub_nan=True) + torch.testing.assert_close(canonical, legacy, rtol=0, atol=0) + assert torch.isinf(canonical[0]) + assert torch.isposinf(canonical[64]) + + +@pytest.mark.skipif( + not _on_gfx950(), + reason="NaN-free fp8_ds_mla compressed-cache contract is gfx950-only", +) +@pytest.mark.parametrize("writer", ["single_pass", "two_stage_finalizer"]) +def test_gfx950_compressed_cache_canonicalizes_nonfinite(writer: str) -> None: + head_dim = 512 + rope_dim = 64 + block_size = 4 + device = "cuda" + + positions = torch.zeros(1, dtype=torch.int64, device=device) + slot_mapping = torch.zeros(1, dtype=torch.int64, device=device) + rms_weight = torch.ones(head_dim, dtype=torch.bfloat16, device=device) + rms_weight[0] = float("inf") + rms_weight[64] = torch.finfo(torch.bfloat16).max + rms_weight[448] = float("inf") + rms_weight[450] = float("nan") + cos_sin_cache = torch.zeros(1, rope_dim, dtype=torch.float32, device=device) + cos_sin_cache[:, : rope_dim // 2] = 1.0 + cache = torch.zeros(1, block_size, 584, dtype=torch.uint8, device=device) + + state_cache = torch.zeros(1, 1, 2 * head_dim, dtype=torch.float32, device=device) + state_cache[..., :head_dim] = 1.0 + token_to_req = torch.zeros(1, dtype=torch.int32, device=device) + block_table = torch.zeros(1, 1, dtype=torch.int32, device=device) + + if writer == "single_pass": + compress_norm_rope_store_triton( + state_cache=state_cache, + num_actual=1, + token_to_req_indices=token_to_req, + positions=positions, + slot_mapping=slot_mapping, + block_table=block_table, + block_size=1, + state_width=head_dim, + cos_sin_cache=cos_sin_cache, + kv_cache=cache, + k_cache_metadata=SimpleNamespace(slot_mapping=slot_mapping), + pdl_kwargs={}, + head_dim=head_dim, + rope_head_dim=rope_dim, + compress_ratio=1, + overlap=False, + use_fp4_cache=False, + rms_norm_weight=rms_weight, + rms_norm_eps=1e-6, + quant_block=64, + token_stride=576, + scale_dim=8, + ) + else: + _launch_two_stage_sparse_attn_compressor( + state_cache, + token_to_req, + positions, + slot_mapping, + block_table, + 1, + head_dim, + 1, + cos_sin_cache, + cache, + slot_mapping, + rms_weight, + 1e-6, + 64, + 576, + 8, + head_dim, + rope_dim, + 1, + torch.empty(1, head_dim, dtype=torch.float32, device=device), + ) + + _assert_nan_free_cache_matches_legacy_scrub(cache, block_size) + + @pytest.mark.parametrize( ("starts", "query_start_loc", "expected"), [ @@ -402,7 +590,8 @@ def test_indexer_gather_accepts_upper_bound_output(): valid_tokens = 9 upper_bound_tokens = 13 block_size = 16 - num_blocks = 2 + num_seqs = 3 + num_blocks = num_seqs sentinel = 123 device = "cuda" @@ -410,13 +599,15 @@ def test_indexer_gather_accepts_upper_bound_output(): kv_cache = torch.zeros( num_blocks, block_size, cache_stride, dtype=torch.uint8, device=device ) - slot_mapping = torch.arange(valid_tokens, dtype=torch.int64, device=device) + slot_mapping = torch.tensor( + [0, 1, 2, 16, 17, 18, 32, 33, 34], dtype=torch.int64, device=device + ) ops.indexer_k_quant_and_cache(k, kv_cache, slot_mapping, quant_block_size, "ue8m0") block_table = torch.arange(num_blocks, dtype=torch.int32, device=device).unsqueeze( - 0 + 1 ) - cu_seq_lens = torch.tensor([0, valid_tokens], dtype=torch.int32, device=device) + cu_seq_lens = torch.tensor([0, 3, 6, 9], dtype=torch.int32, device=device) dst_k = torch.full( (upper_bound_tokens, head_dim), sentinel, dtype=torch.uint8, device=device ) @@ -431,8 +622,52 @@ def test_indexer_gather_accepts_upper_bound_output(): ops.cp_gather_indexer_k_quant_cache( kv_cache, dst_k, dst_scale, block_table, cu_seq_lens ) + + if current_platform.is_rocm(): + triton_kv_cache = torch.zeros_like(kv_cache) + indexer_k_quant_and_cache_triton( + k, + triton_kv_cache, + slot_mapping, + quant_block_size, + "ue8m0", + ) + triton_dst_k = torch.full_like(dst_k, sentinel) + triton_dst_scale = torch.full_like(dst_scale, sentinel) + token_to_seq = torch.cat( + ( + torch.repeat_interleave( + torch.arange(num_seqs, dtype=torch.int32, device=device), 3 + ), + torch.full( + (upper_bound_tokens - valid_tokens,), + -1, + dtype=torch.int32, + device=device, + ), + ) + ) + cp_gather_indexer_k_quant_cache_triton( + triton_kv_cache, + triton_dst_k.view(current_platform.fp8_dtype()), + triton_dst_scale, + block_table, + cu_seq_lens, + token_to_seq, + ) torch.accelerator.synchronize() + if current_platform.is_rocm(): + triton_recovered = triton_dst_k[:valid_tokens].view( + current_platform.fp8_dtype() + ).float() * triton_dst_scale[:valid_tokens].view(torch.float32) + triton_error = (triton_recovered - k.float()).abs().amax(dim=1) + max_triton_error = ( + 16.0 * triton_dst_scale[:valid_tokens].view(torch.float32).flatten() + ) + assert torch.all(triton_error <= max_triton_error) + assert torch.all(triton_dst_k[valid_tokens:] == sentinel) + assert torch.all(triton_dst_scale[valid_tokens:] == sentinel) k_recovered = dst_k[:valid_tokens].view(torch.float8_e4m3fn).float() * dst_scale[ :valid_tokens ].view(torch.float32) diff --git a/vllm/models/deepseek_v4/amd/rocm.py b/vllm/models/deepseek_v4/amd/rocm.py index 6a457298dc27..5e7609bd1191 100644 --- a/vllm/models/deepseek_v4/amd/rocm.py +++ b/vllm/models/deepseek_v4/amd/rocm.py @@ -19,6 +19,7 @@ DeepseekV4SparseMLAMetadataBuilder, ) from vllm.platforms import current_platform +from vllm.platforms.rocm import _ON_GFX950 from vllm.triton_utils import tl, triton from vllm.v1.attention.backend import ( CommonAttentionMetadata, @@ -36,6 +37,19 @@ from vllm.v1.worker.workspace import current_workspace_manager +def _trust_dsv4_extra_cache_nan_free( + kv_cache_dtype: str, + has_kv_transfer: bool, + has_extra_cache: bool, +) -> bool: + return ( + _ON_GFX950 + and kv_cache_dtype == "fp8_ds_mla" + and not has_kv_transfer + and has_extra_cache + ) + + def _build_indptr_from_lengths(lengths: torch.Tensor) -> torch.Tensor: lengths = lengths.to(dtype=torch.int32).contiguous() indptr = torch.zeros(lengths.shape[0] + 1, dtype=torch.int32, device=lengths.device) @@ -294,9 +308,13 @@ def _copy_ragged_to_graph_buffers( max_entries = max(num_rows * max_entries_per_row, 1) ragged_out = ragged_indices_buffer[:max_entries] - nnz = ragged_indices.numel() - if nnz > 0: - ragged_out[:nnz].copy_(ragged_indices, non_blocking=True) + source_entries = ragged_indices.numel() + if source_entries > 0: + ragged_out[:source_entries].copy_(ragged_indices, non_blocking=True) + if _ON_GFX950: + # Preserve the graph-stable base pointer while exposing source capacity + # to the sync-free split selector; indptr still carries the true NNZ. + ragged_out = ragged_out[: max(source_entries, 1)] return ragged_out, indptr_out @@ -306,6 +324,7 @@ class DeepseekV4ROCMAiterMLASparseMetadata(DeepseekV4FlashMLAMetadata): c128a_decode_topk_ragged_indices: torch.Tensor | None = None c128a_decode_topk_ragged_indptr: torch.Tensor | None = None + for_cudagraph_capture: bool = False @dataclass @@ -370,6 +389,16 @@ def build( c128a_decode_topk_ragged_indptr=ragged_indptr, ) + def build_for_cudagraph_capture( + self, common_attn_metadata: CommonAttentionMetadata + ) -> DeepseekV4ROCMAiterMLASparseMetadata: + metadata = cast( + DeepseekV4ROCMAiterMLASparseMetadata, + super().build_for_cudagraph_capture(common_attn_metadata), + ) + metadata.for_cudagraph_capture = _ON_GFX950 + return metadata + class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuilder): # Keep fused multi-step decode disabled until update_draft_decode_metadata() @@ -451,7 +480,9 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): backend_cls = DeepseekV4ROCMAiterMLASparseBackend def __init__(self, *args, **kwargs): + vllm_config = args[0] if args else kwargs["vllm_config"] super().__init__(*args, **kwargs) + self._has_kv_transfer = vllm_config.kv_transfer_config is not None # Block scale for the preshuffled weight; None = not preshuffled. self._wqa_wkv_scale: torch.Tensor | None = None self._wo_b_scale: torch.Tensor | None = None @@ -605,6 +636,13 @@ def forward_mqa( attn_metadata=rocm_metadata, swa_only=swa_only, output=output[:num_decode_tokens], + adaptive_splits=( + _ON_GFX950 + and not swa_only + and self.compress_ratio == 128 + and rocm_metadata is not None + and rocm_metadata.for_cudagraph_capture + ), ) def _forward_decode( @@ -615,6 +653,7 @@ def _forward_decode( attn_metadata: DeepseekV4ROCMAiterMLASparseMetadata | None, swa_only: bool, output: torch.Tensor, + adaptive_splits: bool, ) -> None: num_decodes = swa_metadata.num_decodes num_decode_tokens = swa_metadata.num_decode_tokens @@ -666,6 +705,12 @@ def _forward_decode( nope_head_dim=self.nope_head_dim, rope_head_dim=self.rope_head_dim, output=output, + adaptive_splits=adaptive_splits, + extra_cache_nan_free=_trust_dsv4_extra_cache_nan_free( + self.kv_cache_dtype, + self._has_kv_transfer, + not swa_only and kv_cache is not None, + ), ) def _forward_prefill( diff --git a/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py b/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py index a2085cd220f1..4c9f464ef067 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py +++ b/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py @@ -24,8 +24,14 @@ import torch +from vllm.platforms import current_platform from vllm.triton_utils import tl, triton +if current_platform.is_rocm(): + from vllm.platforms.rocm import _ON_GFX950 +else: + _ON_GFX950 = False + from .fused_indexer_q import _fp32x2_to_fp4x2 @@ -61,12 +67,15 @@ def compress_norm_rope_store_triton( if head_dim == 512: kernel = _fused_kv_compress_norm_rope_insert_sparse_attn num_warps = 4 + kernel_kwargs = {"SANITIZE_CACHE_NANS": _ON_GFX950} elif use_fp4_cache: kernel = _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn num_warps = 1 + kernel_kwargs = {} else: kernel = _fused_kv_compress_norm_rope_insert_indexer_attn num_warps = 1 + kernel_kwargs = {} kernel[(num_actual,)]( # state cache @@ -103,6 +112,7 @@ def compress_norm_rope_store_triton( SCALE_DIM=scale_dim, KV_BLOCK_STRIDE=kv_cache.stride(0), num_warps=num_warps, + **kernel_kwargs, **pdl_kwargs, ) @@ -145,6 +155,7 @@ def _fused_kv_compress_norm_rope_insert_sparse_attn( TOKEN_STRIDE: tl.constexpr, # 576 for DeepseekV4 SCALE_DIM: tl.constexpr, # 8 for DeepseekV4 (7 real + 1 pad) KV_BLOCK_STRIDE: tl.constexpr, + SANITIZE_CACHE_NANS: tl.constexpr, ): """Fused compress → RMSNorm → FP8 quant (nope) → RoPE → bf16 store (rope). @@ -261,7 +272,8 @@ def _fused_kv_compress_norm_rope_insert_sparse_attn( scale_idx = tl.arange(0, N_QUANT_BLOCKS) encoded = exponents + 127.0 - encoded = tl.maximum(tl.minimum(encoded, 255.0), 0.0) + max_encoded: tl.constexpr = 254.0 if SANITIZE_CACHE_NANS else 255.0 + encoded = tl.maximum(tl.minimum(encoded, max_encoded), 0.0) tl.store( scale_ptr + scale_idx, encoded.to(tl.uint8), @@ -289,6 +301,8 @@ def _fused_kv_compress_norm_rope_insert_sparse_attn( new_even = even * cos_v - odd * sin_v new_odd = odd * cos_v + even * sin_v result = tl.interleave(new_even, new_odd) # [TRITON_BLOCK_SIZE] fp32 + if SANITIZE_CACHE_NANS: + result = tl.where(result == result, result, 0.0) # Store rotated rope portion as bf16 into the cache's bf16 area. bf16_ptr = (fp8_ptr + NOPE_HEAD_DIM).to(tl.pointer_type(tl.bfloat16)) @@ -417,6 +431,7 @@ def _finalize_norm_rope_quant_store_sparse_attn( TOKEN_STRIDE: tl.constexpr, SCALE_DIM: tl.constexpr, KV_BLOCK_STRIDE: tl.constexpr, + SANITIZE_CACHE_NANS: tl.constexpr, ): """Stage 2: read compressed_kv[512] from scratch buffer, then RMSNorm + FP8 quant (nope) + RoPE + bf16 store @@ -474,7 +489,8 @@ def _finalize_norm_rope_quant_store_sparse_attn( tl.store(fp8_ptr + block, x_uint8, mask=block < NOPE_HEAD_DIM) scale_idx = tl.arange(0, N_QUANT_BLOCKS) - encoded = tl.maximum(tl.minimum(exponents + 127.0, 255.0), 0.0) + max_encoded: tl.constexpr = 254.0 if SANITIZE_CACHE_NANS else 255.0 + encoded = tl.maximum(tl.minimum(exponents + 127.0, max_encoded), 0.0) tl.store( scale_ptr + scale_idx, encoded.to(tl.uint8), mask=scale_idx < N_NOPE_BLOCKS ) @@ -494,6 +510,8 @@ def _finalize_norm_rope_quant_store_sparse_attn( new_even = even * cos_v - odd * sin_v new_odd = odd * cos_v + even * sin_v result = tl.interleave(new_even, new_odd) + if SANITIZE_CACHE_NANS: + result = tl.where(result == result, result, 0.0) bf16_ptr = (fp8_ptr + NOPE_HEAD_DIM).to(tl.pointer_type(tl.bfloat16)) rope_local = block - NOPE_HEAD_DIM is_rope = (block >= NOPE_HEAD_DIM) & mask @@ -564,6 +582,7 @@ def _launch_two_stage_sparse_attn_compressor( TOKEN_STRIDE=token_stride, SCALE_DIM=scale_dim, KV_BLOCK_STRIDE=kv_cache.stride(0), + SANITIZE_CACHE_NANS=_ON_GFX950, ) diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 94eeaba6afaf..5c9ab386d4e0 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -214,6 +214,87 @@ def _cp_gather_indexer_quant_cache_kernel( tl.store(dst_k_ptr + offset, val, mask=valid_block) +@triton.jit(do_not_specialize=["num_batches"]) +def _cp_gather_indexer_quant_cache_gfx950_kernel( + kv_cache_ptr, # [n_blks,blk_size//tile_blk,head_dim//16B,tile_blk,16B] + # [n_blks, blk_size, head_dim] + kv_cache_scale_ptr, # [n_blks, blk_size] + k_fp8_ptr, # [num_tokens, head_dim] + k_scale_ptr, # [num_tokens] + block_table_ptr, # [batch_size, block_table_stride] + cu_seqlen_ptr, # [batch_size + 1] + token_to_seq_ptr, # [num_tokens] + block_size, + block_table_stride, + kv_cache_stride, + kv_cache_scale_stride, + LAYOUT: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_TILE_SIZE: tl.constexpr, + HEAD_TILE_SIZE: tl.constexpr, + num_batches, + BLOCK_TABLE_WIDTH: tl.constexpr, + NUM_BLOCKS: tl.constexpr, +): + tid = tl.program_id(0) + offset = tl.arange(0, HEAD_DIM) + batch_id = tl.load(token_to_seq_ptr + tid) + valid_batch = (batch_id >= 0) & (batch_id < num_batches) + safe_batch_id = tl.where(valid_batch, batch_id, 0) + batch_start = tl.load(cu_seqlen_ptr + safe_batch_id, mask=valid_batch, other=0) + batch_end = tl.load(cu_seqlen_ptr + safe_batch_id + 1, mask=valid_batch, other=0) + batch_offset = tid - batch_start + valid_token = valid_batch & (tid >= batch_start) & (tid < batch_end) + if not valid_token: + return + block_table_id = batch_offset // block_size + block_offset = batch_offset % block_size + valid_block_table = ( + valid_token + & (block_table_id >= 0) + & (block_table_id < BLOCK_TABLE_WIDTH) + & (block_offset >= 0) + & (block_offset < block_size) + ) + safe_block_table_id = tl.where(valid_block_table, block_table_id, 0) + block_table_offset = safe_batch_id * block_table_stride + safe_block_table_id + block_id = tl.load( + block_table_ptr + block_table_offset, mask=valid_block_table, other=-1 + ) + valid_block = valid_block_table & (block_id >= 0) & (block_id < NUM_BLOCKS) + # The packed KV layout makes per-block strides large + # enough that block_id * stride can exceed 32-bit range. + safe_block_id = tl.where(valid_block, block_id, 0).to(tl.int64) + safe_block_offset = tl.where(valid_block, block_offset, 0) + tiled_block_offset = safe_block_offset % BLOCK_TILE_SIZE + if LAYOUT == "SHUFFLE": + src_cache_offset = ( + safe_block_id * kv_cache_stride + + (safe_block_offset // BLOCK_TILE_SIZE) * HEAD_DIM * BLOCK_TILE_SIZE + + tiled_block_offset * HEAD_TILE_SIZE + ) + else: + src_cache_offset = ( + safe_block_id * kv_cache_stride + safe_block_offset * HEAD_DIM + ) + src_scale_offset = safe_block_id * kv_cache_scale_stride + safe_block_offset + dst_offset = tid * HEAD_DIM + src_scale_ptr = kv_cache_scale_ptr + src_scale_offset + src_cache_ptr = kv_cache_ptr + src_cache_offset + dst_k_ptr = k_fp8_ptr + dst_offset + scale_val = tl.load(src_scale_ptr, mask=valid_block, other=0.0) + tl.store(k_scale_ptr + tid, scale_val) + if LAYOUT == "SHUFFLE": + tiled_src_offset = ( + offset // HEAD_TILE_SIZE * HEAD_TILE_SIZE * BLOCK_TILE_SIZE + + offset % HEAD_TILE_SIZE + ) + else: + tiled_src_offset = offset + val = tl.load(src_cache_ptr + tiled_src_offset) + tl.store(dst_k_ptr + offset, val, mask=valid_block) + + def cp_gather_indexer_k_quant_cache_triton( k_cache: torch.Tensor, # [num_blocks, block_size, head_dim + 4] k_fp8: torch.Tensor, @@ -237,7 +318,7 @@ def cp_gather_indexer_k_quant_cache_triton( grid = (num_tokens,) k_fp8_scale = k_fp8_scale.view(torch.float32) layout = "NORMAL" if block_size == 1 else "SHUFFLE" - _cp_gather_indexer_quant_cache_kernel[grid]( + kernel_args = ( k_cache_value, k_cache_scale, k_fp8, @@ -253,11 +334,22 @@ def cp_gather_indexer_k_quant_cache_triton( head_dim, block_tile_size, head_tile_size, - num_tokens, - cu_seqlen.shape[0] - 1, - block_table.shape[1], - num_blocks, ) + if _ON_GFX950: + _cp_gather_indexer_quant_cache_gfx950_kernel[grid]( + *kernel_args, + cu_seqlen.shape[0] - 1, + block_table.shape[1], + num_blocks, + ) + else: + _cp_gather_indexer_quant_cache_kernel[grid]( + *kernel_args, + num_tokens, + cu_seqlen.shape[0] - 1, + block_table.shape[1], + num_blocks, + ) # Taken from https://github.com/deepseek-ai/DeepGEMM/blob/main/tests/test_attention.py#L156 @@ -1211,6 +1303,90 @@ def _sparse_attn_prefill_ragged_kernel( ) +@triton.jit +def _decode_e8m0_scales_triton(encoded_scales): + scale_bits = encoded_scales.to(tl.int32) << 23 + scale_bits = tl.where(encoded_scales == 0, 1 << 22, scale_bits) + return scale_bits.to(tl.float32, bitcast=True) + + +@triton.jit +def _load_fp8_ds_mla_gfx950_nope_exact_chunk( + token_data_ptr, + token_scale_ptr, + valid, + CHUNK_START: tl.constexpr, + CHUNK_SIZE: tl.constexpr, + BLOCK_K: tl.constexpr, + IS_FNUZ: tl.constexpr, +): + offsets = CHUNK_START + tl.arange(0, CHUNK_SIZE) + x_uint8 = tl.load( + token_data_ptr[:, None] + offsets[None, :], + mask=valid[:, None], + other=0, + ) + scale_offsets = CHUNK_START // 64 + tl.arange(0, CHUNK_SIZE // 64) + encoded_scales = tl.load( + token_scale_ptr[:, None] + scale_offsets[None, :], + mask=valid[:, None], + other=127, + ) + scales = _decode_e8m0_scales_triton(encoded_scales) + scales = tl.broadcast_to(scales[:, :, None], (BLOCK_K, CHUNK_SIZE // 64, 64)) + scales = tl.reshape(scales, (BLOCK_K, CHUNK_SIZE)) + if IS_FNUZ: + x_f32 = x_uint8.to(tl.float8e4b8, bitcast=True).to(tl.bfloat16).to(tl.float32) + else: + x_f32 = x_uint8.to(tl.float8e4nv, bitcast=True).to(tl.float32) + value = (x_f32 * scales).to(tl.bfloat16) + zero = tl.zeros((BLOCK_K, CHUNK_SIZE), dtype=tl.bfloat16) + return tl.where(valid[:, None], value, zero) + + +@triton.jit +def _load_fp8_ds_mla_gfx950_tail128( + token_data_ptr, + token_scale_ptr, + valid, + NOPE_DIM: tl.constexpr, + BLOCK_K: tl.constexpr, + IS_FNUZ: tl.constexpr, +): + tail_offsets = 384 + tl.arange(0, 128) + nope_mask = tail_offsets < NOPE_DIM + x_uint8 = tl.load( + token_data_ptr[:, None] + tail_offsets[None, :], + mask=valid[:, None] & nope_mask[None, :], + other=0, + ) + scale_offsets = 6 + tl.arange(0, 2) + scale_mask = scale_offsets < NOPE_DIM // 64 + encoded_scales = tl.load( + token_scale_ptr[:, None] + scale_offsets[None, :], + mask=valid[:, None] & scale_mask[None, :], + other=127, + ) + scales = _decode_e8m0_scales_triton(encoded_scales) + scales = tl.broadcast_to(scales[:, :, None], (BLOCK_K, 2, 64)) + scales = tl.reshape(scales, (BLOCK_K, 128)) + if IS_FNUZ: + x_f32 = x_uint8.to(tl.float8e4b8, bitcast=True).to(tl.bfloat16).to(tl.float32) + else: + x_f32 = x_uint8.to(tl.float8e4nv, bitcast=True).to(tl.float32) + nope = (x_f32 * scales).to(tl.bfloat16) + + rope_ptr = (token_data_ptr + NOPE_DIM).to(tl.pointer_type(tl.bfloat16)) + rope = tl.load( + rope_ptr[:, None] + (tail_offsets[None, :] - NOPE_DIM), + mask=valid[:, None] & ~nope_mask[None, :], + other=0.0, + ) + value = tl.where(nope_mask[None, :], nope, rope) + zero = tl.zeros((BLOCK_K, 128), dtype=tl.bfloat16) + return tl.where(valid[:, None], value, zero) + + @triton.jit def _sparse_attn_decode_ragged_kernel( q_ptr, @@ -1685,6 +1861,407 @@ def _sparse_attn_decode_partial_kernel( ) +@triton.jit +def _sparse_attn_decode_gfx950_partial_loaded_tile( + q_combined, + cache_ptr, + slot, + valid, + cache_stride0, + scale: tl.constexpr, + head_mask, + m_i, + l_i, + acc_nope_0a, + acc_nope_0b, + acc_nope_1, + acc_tail, + BLOCK_SIZE: tl.constexpr, + NOPE_DIM: tl.constexpr, + BLOCK_K: tl.constexpr, + IS_FNUZ: tl.constexpr, + TRUST_EXTRA_CACHE_NAN_FREE: tl.constexpr, +): + safe_slot = tl.where(valid, slot, 0) + block_idx = safe_slot // BLOCK_SIZE + pos_in_block = safe_slot % BLOCK_SIZE + cache_block_ptr = cache_ptr + block_idx.to(tl.int64) * cache_stride0 + token_data_ptr = cache_block_ptr + pos_in_block * 576 + token_scale_ptr = cache_block_ptr + BLOCK_SIZE * 576 + pos_in_block * 8 + k_nope_0a = _load_fp8_ds_mla_gfx950_nope_exact_chunk( + token_data_ptr, + token_scale_ptr, + valid, + 0, + 128, + BLOCK_K, + IS_FNUZ, + ) + k_nope_0b = _load_fp8_ds_mla_gfx950_nope_exact_chunk( + token_data_ptr, + token_scale_ptr, + valid, + 128, + 128, + BLOCK_K, + IS_FNUZ, + ) + k_nope_1 = _load_fp8_ds_mla_gfx950_nope_exact_chunk( + token_data_ptr, + token_scale_ptr, + valid, + 256, + 128, + BLOCK_K, + IS_FNUZ, + ) + k_tail = _load_fp8_ds_mla_gfx950_tail128( + token_data_ptr, + token_scale_ptr, + valid, + NOPE_DIM, + BLOCK_K, + IS_FNUZ, + ) + if not TRUST_EXTRA_CACHE_NAN_FREE: + zero = tl.zeros((BLOCK_K, 128), dtype=tl.bfloat16) + k_nope_0a = tl.where(k_nope_0a == k_nope_0a, k_nope_0a, zero) + k_nope_0b = tl.where(k_nope_0b == k_nope_0b, k_nope_0b, zero) + k_nope_1 = tl.where(k_nope_1 == k_nope_1, k_nope_1, zero) + k_tail = tl.where(k_tail == k_tail, k_tail, zero) + k_nope_0 = tl.cat(k_nope_0a, k_nope_0b, dim=1) + k_tail_256 = tl.cat(k_nope_1, k_tail, dim=1) + k_combined = tl.cat(k_nope_0, k_tail_256, dim=1) + + scores = tl.dot(q_combined, tl.trans(k_combined)) + scores *= scale * 1.4426950408889634 + scores = tl.where( + head_mask[:, None] & valid[None, :], + scores, + -3.4028234663852886e38, + ) + m_block = tl.max(scores, axis=1) + m_new = tl.maximum(m_i, m_block) + alpha = tl.exp2(m_i - m_new) + p = tl.exp2(scores - m_new[:, None]) + p = tl.where(head_mask[:, None] & valid[None, :], p, 0.0) + l_new = l_i * alpha + tl.sum(p, axis=1) + p_bf16 = p.to(k_nope_0a.dtype) + acc_nope_0a = acc_nope_0a * alpha[:, None] + tl.dot(p_bf16, k_nope_0a) + acc_nope_0b = acc_nope_0b * alpha[:, None] + tl.dot(p_bf16, k_nope_0b) + acc_nope_1 = acc_nope_1 * alpha[:, None] + tl.dot(p_bf16, k_nope_1) + acc_tail = acc_tail * alpha[:, None] + tl.dot(p_bf16, k_tail) + return ( + m_new, + l_new, + acc_nope_0a, + acc_nope_0b, + acc_nope_1, + acc_tail, + ) + + +@triton.jit +def _sparse_attn_decode_gfx950_partial_kernel( + q_ptr, + main_cache_ptr, + main_indices_ptr, + main_indptr_ptr, + extra_cache_ptr, + extra_indices_ptr, + extra_indptr_ptr, + part_m_ptr, + part_l_ptr, + part_acc_ptr, + q_stride0: tl.constexpr, + q_stride1: tl.constexpr, + main_cache_stride0: tl.constexpr, + extra_cache_stride0: tl.constexpr, + main_num_rows, + extra_num_rows, + MAIN_BLOCK_SIZE: tl.constexpr, + EXTRA_BLOCK_SIZE: tl.constexpr, + scale: tl.constexpr, + num_heads: tl.constexpr, + HAS_EXTRA: tl.constexpr, + NOPE_DIM: tl.constexpr, + ROPE_DIM: tl.constexpr, + IS_FNUZ_MAIN: tl.constexpr, + IS_FNUZ_EXTRA: tl.constexpr, + TRUST_EXTRA_CACHE_NAN_FREE: tl.constexpr, + ADAPTIVE_SPLITS: tl.constexpr, + ONE_WAVE_SPLITS: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_K: tl.constexpr, + NUM_SPLITS: tl.constexpr, + NUM_STAGES: tl.constexpr, +): + query_idx = tl.program_id(0) + split_id = tl.program_id(1) + pid_h = tl.program_id(2) + + tl.static_assert(NOPE_DIM == 448) + tl.static_assert(ROPE_DIM == 64) + + head_offsets = pid_h * BLOCK_H + tl.arange(0, BLOCK_H) + if num_heads % BLOCK_H == 0: + head_mask = tl.full((BLOCK_H,), True, tl.int1) + else: + head_mask = head_offsets < num_heads + neg_large = -3.4028234663852886e38 + + if ADAPTIVE_SPLITS: + main_start = tl.load(main_indptr_ptr + query_idx) + main_end = tl.load(main_indptr_ptr + query_idx + 1) + main_len = main_end - main_start + if HAS_EXTRA: + extra_start = tl.load(extra_indptr_ptr + query_idx) + extra_end = tl.load(extra_indptr_ptr + query_idx + 1) + extra_len = extra_end - extra_start + else: + extra_start = 0 + extra_len = 0 + split4_span: tl.constexpr = 4 * BLOCK_K + split4_iters = (main_len + split4_span - 1) // split4_span + split4_iters += (extra_len + split4_span - 1) // split4_span + use_four_splits = split4_iters <= 3 + work_splits = NUM_SPLITS + if ONE_WAVE_SPLITS > 4 and ONE_WAVE_SPLITS < NUM_SPLITS: + one_wave_span: tl.constexpr = ONE_WAVE_SPLITS * BLOCK_K + one_wave_iters = (main_len + one_wave_span - 1) // one_wave_span + one_wave_iters += (extra_len + one_wave_span - 1) // one_wave_span + work_splits = tl.where(one_wave_iters <= 3, ONE_WAVE_SPLITS, work_splits) + work_splits = tl.where(use_four_splits, 4, work_splits) + if split_id >= work_splits: + pm_base = (query_idx * NUM_SPLITS + split_id) * num_heads + head_offsets + tl.store(part_m_ptr + pm_base, neg_large, mask=head_mask) + tl.store(part_l_ptr + pm_base, 0.0, mask=head_mask) + return + else: + work_splits = NUM_SPLITS + + nope_offsets_0a = tl.arange(0, 128) + nope_offsets_0b = 128 + tl.arange(0, 128) + nope_offsets_0 = tl.arange(0, 256) + tail_offsets = 256 + tl.arange(0, 256) + nope_offsets_1 = 256 + tl.arange(0, 128) + tail_offsets_128 = 384 + tl.arange(0, 128) + + q_row_ptr = q_ptr + query_idx * q_stride0 + head_offsets[:, None] * q_stride1 + q_nope_0 = tl.load( + q_row_ptr + nope_offsets_0[None, :], + mask=head_mask[:, None], + other=0.0, + ) + q_tail = tl.load( + q_row_ptr + tail_offsets[None, :], + mask=head_mask[:, None], + other=0.0, + ) + q_combined = tl.cat(q_nope_0, q_tail, dim=1) + + m_i = tl.full((BLOCK_H,), neg_large, dtype=tl.float32) + l_i = tl.zeros((BLOCK_H,), dtype=tl.float32) + acc_nope_0a = tl.zeros((BLOCK_H, 128), dtype=tl.float32) + acc_nope_0b = tl.zeros((BLOCK_H, 128), dtype=tl.float32) + acc_nope_1 = tl.zeros((BLOCK_H, 128), dtype=tl.float32) + acc_tail = tl.zeros((BLOCK_H, 128), dtype=tl.float32) + k_offsets = tl.arange(0, BLOCK_K) + + if not ADAPTIVE_SPLITS: + main_start = tl.load(main_indptr_ptr + query_idx) + main_end = tl.load(main_indptr_ptr + query_idx + 1) + main_len = main_end - main_start + main_chunk = (main_len + work_splits - 1) // work_splits + main_lo = split_id * main_chunk + main_hi = tl.minimum(main_lo + main_chunk, main_len) + + for k_start in tl.range( + main_lo, + main_hi, + BLOCK_K, + num_stages=NUM_STAGES, + ): + k_pos = k_start + k_offsets + in_range = k_pos < main_hi + slot = tl.load(main_indices_ptr + main_start + k_pos, mask=in_range, other=-1) + valid = in_range & (slot >= 0) & (slot < main_num_rows) + ( + m_i, + l_i, + acc_nope_0a, + acc_nope_0b, + acc_nope_1, + acc_tail, + ) = _sparse_attn_decode_gfx950_partial_loaded_tile( + q_combined, + main_cache_ptr, + slot, + valid, + main_cache_stride0, + scale, + head_mask, + m_i, + l_i, + acc_nope_0a, + acc_nope_0b, + acc_nope_1, + acc_tail, + MAIN_BLOCK_SIZE, + NOPE_DIM, + BLOCK_K, + IS_FNUZ_MAIN, + False, + ) + + if HAS_EXTRA: + if not ADAPTIVE_SPLITS: + extra_start = tl.load(extra_indptr_ptr + query_idx) + extra_end = tl.load(extra_indptr_ptr + query_idx + 1) + extra_len = extra_end - extra_start + extra_chunk = (extra_len + work_splits - 1) // work_splits + extra_lo = split_id * extra_chunk + extra_hi = tl.minimum(extra_lo + extra_chunk, extra_len) + + outer_block_k: tl.constexpr = 2 * BLOCK_K + outer_k_offsets = tl.arange(0, outer_block_k) + extra_hi_full = ( + extra_lo + ((extra_hi - extra_lo) // outer_block_k) * outer_block_k + ) + for k_start in tl.range( + extra_lo, + extra_hi_full, + outer_block_k, + num_stages=NUM_STAGES, + ): + slot = tl.load(extra_indices_ptr + extra_start + k_start + outer_k_offsets) + valid = (slot >= 0) & (slot < extra_num_rows) + slot_pairs = tl.trans(tl.reshape(slot, (2, BLOCK_K))) + valid_pairs = tl.trans(tl.reshape(valid, (2, BLOCK_K))) + slot_lo, slot_hi = tl.split(slot_pairs) + valid_lo, valid_hi = tl.split(valid_pairs) + ( + m_i, + l_i, + acc_nope_0a, + acc_nope_0b, + acc_nope_1, + acc_tail, + ) = _sparse_attn_decode_gfx950_partial_loaded_tile( + q_combined, + extra_cache_ptr, + slot_lo, + valid_lo, + extra_cache_stride0, + scale, + head_mask, + m_i, + l_i, + acc_nope_0a, + acc_nope_0b, + acc_nope_1, + acc_tail, + EXTRA_BLOCK_SIZE, + NOPE_DIM, + BLOCK_K, + IS_FNUZ_EXTRA, + TRUST_EXTRA_CACHE_NAN_FREE, + ) + ( + m_i, + l_i, + acc_nope_0a, + acc_nope_0b, + acc_nope_1, + acc_tail, + ) = _sparse_attn_decode_gfx950_partial_loaded_tile( + q_combined, + extra_cache_ptr, + slot_hi, + valid_hi, + extra_cache_stride0, + scale, + head_mask, + m_i, + l_i, + acc_nope_0a, + acc_nope_0b, + acc_nope_1, + acc_tail, + EXTRA_BLOCK_SIZE, + NOPE_DIM, + BLOCK_K, + IS_FNUZ_EXTRA, + TRUST_EXTRA_CACHE_NAN_FREE, + ) + for tail_idx in tl.static_range(2): + tail_start = extra_hi_full + tail_idx * BLOCK_K + if tail_start < extra_hi: + k_pos = tail_start + k_offsets + in_range = k_pos < extra_hi + slot = tl.load( + extra_indices_ptr + extra_start + k_pos, + mask=in_range, + other=-1, + ) + valid = in_range & (slot >= 0) & (slot < extra_num_rows) + ( + m_i, + l_i, + acc_nope_0a, + acc_nope_0b, + acc_nope_1, + acc_tail, + ) = _sparse_attn_decode_gfx950_partial_loaded_tile( + q_combined, + extra_cache_ptr, + slot, + valid, + extra_cache_stride0, + scale, + head_mask, + m_i, + l_i, + acc_nope_0a, + acc_nope_0b, + acc_nope_1, + acc_tail, + EXTRA_BLOCK_SIZE, + NOPE_DIM, + BLOCK_K, + IS_FNUZ_EXTRA, + TRUST_EXTRA_CACHE_NAN_FREE, + ) + + pm_base = (query_idx * NUM_SPLITS + split_id) * num_heads + head_offsets + m_store = tl.where(l_i > 0.0, m_i * 0.6931471805599453, neg_large) + tl.store(part_m_ptr + pm_base, m_store, mask=head_mask) + tl.store(part_l_ptr + pm_base, l_i, mask=head_mask) + acc_base = part_acc_ptr + ( + (query_idx * NUM_SPLITS + split_id) * num_heads + head_offsets[:, None] + ) * (NOPE_DIM + ROPE_DIM) + tl.store( + acc_base + nope_offsets_0a[None, :], + acc_nope_0a, + mask=head_mask[:, None], + ) + tl.store( + acc_base + nope_offsets_0b[None, :], + acc_nope_0b, + mask=head_mask[:, None], + ) + tl.store( + acc_base + nope_offsets_1[None, :], + acc_nope_1, + mask=head_mask[:, None], + ) + tl.store( + acc_base + tail_offsets_128[None, :], + acc_tail, + mask=head_mask[:, None], + ) + + @triton.jit def _sparse_attn_decode_reduce_kernel( part_m_ptr, @@ -1701,6 +2278,7 @@ def _sparse_attn_decode_reduce_kernel( pa_stride_h, num_heads, HAS_ATTN_SINK: tl.constexpr, + ADAPTIVE_SPLITS: tl.constexpr, COMB_DIM: tl.constexpr, BLOCK_H: tl.constexpr, NUM_SPLITS: tl.constexpr, @@ -1769,17 +2347,27 @@ def _sparse_attn_decode_reduce_kernel( other=neg_large, ) w_s = tl.exp(m_s - m_final) + if ADAPTIVE_SPLITS: + active_split = m_s > neg_large + w_s = tl.where(head_mask & active_split, w_s, 0.0) acc_base = ( part_acc_ptr + query_idx * pa_stride0 + s * pa_stride_s + head_offsets[:, None] * pa_stride_h ) - acc_s = tl.load( - acc_base + comb_offsets[None, :], - mask=head_mask[:, None], - other=0.0, - ) + if ADAPTIVE_SPLITS: + acc_s = tl.load( + acc_base + comb_offsets[None, :], + mask=head_mask[:, None] & active_split[:, None], + other=0.0, + ) + else: + acc_s = tl.load( + acc_base + comb_offsets[None, :], + mask=head_mask[:, None], + other=0.0, + ) acc += w_s[:, None] * acc_s out = tl.where(l_final[:, None] > 0.0, acc / denom[:, None], 0.0) @@ -1986,6 +2574,49 @@ def _decode_num_splits( return best_splits +def _decode_gfx950_num_splits( + num_queries: int, + heads_blocks: int, + avg_main_len: float = 0.0, + avg_extra_len: float = 0.0, + block_k: int = 32, +) -> int: + base = max(1, num_queries * heads_blocks) + cu = max(1, _decode_cu_count()) + target_workgroups = 2 * cu + num_splits = min( + 32, + max( + 1, + math.ceil(target_workgroups / base), + ), + ) + if ( + base >= 16 + and num_splits > 4 + and _decode_partial_iters(avg_main_len, avg_extra_len, 4, block_k) <= 3 + ): + return 4 + if 16 <= base < 64: + one_wave_splits = max(1, cu // base) + one_wave_iters = _decode_partial_iters( + avg_main_len, avg_extra_len, one_wave_splits, block_k + ) + target_waves = 1 if one_wave_iters <= 9 else 2 + num_splits = min(num_splits, max(1, target_waves * cu // base)) + if base >= 16 and num_splits > 1: + target_waves = (base * num_splits + cu - 1) // cu + target_iters = _decode_partial_iters( + avg_main_len, avg_extra_len, num_splits, block_k + ) + for splits in range(1, num_splits): + waves = (base * splits + cu - 1) // cu + iters = _decode_partial_iters(avg_main_len, avg_extra_len, splits, block_k) + if waves == target_waves and iters == target_iters: + return splits + return num_splits + + def _rocm_sparse_attn_decode_ragged_triton( q: torch.Tensor, main_cache: torch.Tensor, @@ -1998,6 +2629,9 @@ def _rocm_sparse_attn_decode_ragged_triton( extra_cache: torch.Tensor | None = None, extra_indices: torch.Tensor | None = None, extra_indptr: torch.Tensor | None = None, + out: torch.Tensor | None = None, + extra_cache_nan_free: bool = False, + adaptive_splits: bool = False, ) -> torch.Tensor: assert q.ndim == 3, f"expected q=[b,h,d], got {q.shape}" assert main_cache.ndim == 3, ( @@ -2038,6 +2672,10 @@ def _rocm_sparse_attn_decode_ragged_triton( and extra_indices is not None and extra_indptr is not None ) + assert not extra_cache_nan_free or (_ON_GFX950 and has_extra), ( + "extra_cache_nan_free requires a gfx950 compressed cache with trusted " + "canonical-writer provenance" + ) if has_extra: assert extra_cache is not None assert extra_indices is not None @@ -2059,7 +2697,16 @@ def _rocm_sparse_attn_decode_ragged_triton( extra_indptr = torch.zeros(num_queries + 1, device=q.device, dtype=torch.int32) block_h = 16 - out = torch.empty_like(q, dtype=torch.bfloat16) + if out is None: + out = torch.empty_like(q, dtype=torch.bfloat16) + else: + assert out.shape == q.shape, f"expected out shape {q.shape}, got {out.shape}" + assert out.device == q.device, ( + f"expected out on device {q.device}, got {out.device}" + ) + assert out.dtype == torch.bfloat16, ( + f"expected out dtype {torch.bfloat16}, got {out.dtype}" + ) heads_blocks = triton.cdiv(num_heads, block_h) nope_block = triton.next_power_of_2(nope_head_dim) comb_dim = nope_head_dim + rope_head_dim @@ -2103,14 +2750,35 @@ def _rocm_sparse_attn_decode_ragged_triton( return out block_k = 32 # KV tokens walked per split-K iteration. Tuned on gfx950. - # Average per-query segment lengths, read sync-free from the ragged index - # sizes, let the split heuristic avoid over-splitting - # main_indices/extra_indices are flat [nnz] int32. - inv_q = 1.0 / max(1, num_queries) - avg_main_len = main_indices.numel() * inv_q - avg_extra_len = (extra_indices.numel() * inv_q) if has_extra else 0.0 - num_splits = _decode_num_splits( - num_queries, heads_blocks, avg_main_len, avg_extra_len, block_k + if _ON_GFX950: + inv_q = 1.0 / max(1, num_queries) + avg_main_len = main_indices.numel() * inv_q + avg_extra_len = (extra_indices.numel() * inv_q) if has_extra else 0.0 + num_splits = _decode_gfx950_num_splits( + num_queries, + heads_blocks, + avg_main_len, + avg_extra_len, + block_k, + ) + else: + # Average per-query segment lengths, read sync-free from the ragged + # index sizes, let the split heuristic avoid over-splitting. + inv_q = 1.0 / max(1, num_queries) + avg_main_len = main_indices.numel() * inv_q + avg_extra_len = (extra_indices.numel() * inv_q) if has_extra else 0.0 + num_splits = _decode_num_splits( + num_queries, heads_blocks, avg_main_len, avg_extra_len, block_k + ) + + base_workgroups = num_queries * heads_blocks + adaptive_splits = ( + _ON_GFX950 and adaptive_splits and base_workgroups >= 16 and num_splits > 4 + ) + one_wave_splits = ( + max(1, _decode_cu_count() // base_workgroups) + if adaptive_splits and 16 <= base_workgroups < 64 + else num_splits ) part_m = torch.empty( @@ -2123,48 +2791,88 @@ def _rocm_sparse_attn_decode_ragged_triton( device=q.device, ) - _sparse_attn_decode_partial_kernel[(num_queries, num_splits, heads_blocks)]( - q, - main_cache, - main_indices, - main_indptr, - extra_cache, - extra_indices, - extra_indptr, - part_m, - part_l, - part_acc, - q.stride(0), - q.stride(1), - main_cache.stride(0), - extra_cache.stride(0), - part_m.stride(0), - part_m.stride(1), - part_acc.stride(0), - part_acc.stride(1), - part_acc.stride(2), - main_cache.shape[0] * main_cache.shape[1], - extra_cache.shape[0] * extra_cache.shape[1], - main_cache.shape[1], - extra_cache.shape[1], - scale, - num_heads, - HAS_EXTRA=has_extra, - NOPE_DIM=nope_head_dim, - NOPE_BLOCK=nope_block, - ROPE_DIM=rope_head_dim, - # main_cache = swa_k_cache (C++ encoder, FNUZ on gfx942 / OCP on gfx950). - # extra_cache = compressed kv_cache (Triton encoder, OCP everywhere). - # Reading both with a single IS_FNUZ would decode one of them with the - # wrong FNUZ/OCP scale ratio (~1.87×). - IS_FNUZ_MAIN=is_fnuz, - IS_FNUZ_EXTRA=False, - BLOCK_H=block_h, - BLOCK_K=block_k, - NUM_SPLITS=num_splits, - NUM_STAGES=1, - num_warps=4, - ) + if _ON_GFX950: + _sparse_attn_decode_gfx950_partial_kernel[ + (num_queries, num_splits, heads_blocks) + ]( + q, + main_cache, + main_indices, + main_indptr, + extra_cache, + extra_indices, + extra_indptr, + part_m, + part_l, + part_acc, + q.stride(0), + q.stride(1), + main_cache.stride(0), + extra_cache.stride(0), + main_cache.shape[0] * main_cache.shape[1], + extra_cache.shape[0] * extra_cache.shape[1], + main_cache.shape[1], + extra_cache.shape[1], + scale, + num_heads, + HAS_EXTRA=has_extra, + NOPE_DIM=nope_head_dim, + ROPE_DIM=rope_head_dim, + IS_FNUZ_MAIN=is_fnuz, + IS_FNUZ_EXTRA=False, + TRUST_EXTRA_CACHE_NAN_FREE=extra_cache_nan_free, + ADAPTIVE_SPLITS=adaptive_splits, + ONE_WAVE_SPLITS=one_wave_splits, + BLOCK_H=block_h, + BLOCK_K=block_k, + NUM_SPLITS=num_splits, + NUM_STAGES=1, + num_warps=4, + waves_per_eu=0, + ) + else: + _sparse_attn_decode_partial_kernel[(num_queries, num_splits, heads_blocks)]( + q, + main_cache, + main_indices, + main_indptr, + extra_cache, + extra_indices, + extra_indptr, + part_m, + part_l, + part_acc, + q.stride(0), + q.stride(1), + main_cache.stride(0), + extra_cache.stride(0), + part_m.stride(0), + part_m.stride(1), + part_acc.stride(0), + part_acc.stride(1), + part_acc.stride(2), + main_cache.shape[0] * main_cache.shape[1], + extra_cache.shape[0] * extra_cache.shape[1], + main_cache.shape[1], + extra_cache.shape[1], + scale, + num_heads, + HAS_EXTRA=has_extra, + NOPE_DIM=nope_head_dim, + NOPE_BLOCK=nope_block, + ROPE_DIM=rope_head_dim, + # main_cache = swa_k_cache (C++ encoder, FNUZ on gfx942 / OCP on gfx950). + # extra_cache = compressed kv_cache (Triton encoder, OCP everywhere). + # Reading both with a single IS_FNUZ would decode one of them with the + # wrong FNUZ/OCP scale ratio (~1.87×). + IS_FNUZ_MAIN=is_fnuz, + IS_FNUZ_EXTRA=False, + BLOCK_H=block_h, + BLOCK_K=block_k, + NUM_SPLITS=num_splits, + NUM_STAGES=1, + num_warps=4, + ) _sparse_attn_decode_reduce_kernel[(num_queries, num_heads)]( part_m, @@ -2181,6 +2889,7 @@ def _rocm_sparse_attn_decode_ragged_triton( part_acc.stride(2), num_heads, HAS_ATTN_SINK=has_attn_sink, + ADAPTIVE_SPLITS=adaptive_splits, COMB_DIM=comb_dim, BLOCK_H=1, NUM_SPLITS=num_splits, @@ -2206,6 +2915,9 @@ def _rocm_sparse_attn_decode_triton( main_ragged_indptr: torch.Tensor | None = None, extra_ragged_indices: torch.Tensor | None = None, extra_ragged_indptr: torch.Tensor | None = None, + out: torch.Tensor | None = None, + extra_cache_nan_free: bool = False, + adaptive_splits: bool = False, ) -> torch.Tensor: if main_ragged_indices is None or main_ragged_indptr is None: main_ragged_indices, main_ragged_indptr = build_ragged_indices_from_dense( @@ -2241,6 +2953,9 @@ def _rocm_sparse_attn_decode_triton( extra_cache=extra_cache, extra_indices=extra_ragged_indices, extra_indptr=extra_ragged_indptr, + out=out, + extra_cache_nan_free=extra_cache_nan_free, + adaptive_splits=adaptive_splits, ) @@ -2312,6 +3027,8 @@ def rocm_sparse_attn_decode( nope_head_dim: int, rope_head_dim: int, output: torch.Tensor, + extra_cache_nan_free: bool = False, + adaptive_splits: bool = False, ) -> None: assert swa_k_cache.dtype == torch.uint8, ( "ROCm Triton sparse decode expects uint8 fp8_ds_mla SWA cache, " @@ -2341,6 +3058,7 @@ def rocm_sparse_attn_decode( if topk_indices is not None: extra_indices = topk_indices.reshape(topk_indices.shape[0], -1) + direct_out = output if _ON_GFX950 and output.dtype == torch.bfloat16 else None attn_out = _rocm_sparse_attn_decode_triton( q=q, main_cache=swa_k_cache, @@ -2357,5 +3075,9 @@ def rocm_sparse_attn_decode( main_ragged_indptr=swa_ragged_indptr, extra_ragged_indices=topk_ragged_indices, extra_ragged_indptr=topk_ragged_indptr, + out=direct_out, + extra_cache_nan_free=extra_cache_nan_free, + adaptive_splits=adaptive_splits, ) - output.copy_(attn_out.to(output.dtype)) + if direct_out is None: + output.copy_(attn_out.to(output.dtype)) From fdab2b10bcac00a16c406f8b17a75a1c3f729e59 Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Sun, 16 Aug 2026 14:25:49 -0500 Subject: [PATCH 020/839] [ModelRunner v2] Support Transformers pooling model (#52425) Signed-off-by: Taneem Ibrahim --- tests/models/transformers/test_backend.py | 11 ++++++++++- vllm/v1/worker/gpu/model_states/__init__.py | 12 ++++++++---- 2 files changed, 18 insertions(+), 5 deletions(-) diff --git a/tests/models/transformers/test_backend.py b/tests/models/transformers/test_backend.py index 4b6b2796af64..5fd39a5b9d54 100644 --- a/tests/models/transformers/test_backend.py +++ b/tests/models/transformers/test_backend.py @@ -263,7 +263,16 @@ def test_embed_loading(vllm_runner, model): @pytest.mark.parametrize( "arch", ["TransformersEmbeddingModel", "TransformersForSequenceClassification"] ) -def test_pooling(hf_runner, vllm_runner, example_prompts, arch): +@pytest.mark.parametrize("use_v2_model_runner", [False, True], ids=["v1", "v2"]) +def test_pooling( + hf_runner, + vllm_runner, + example_prompts, + arch, + monkeypatch, + use_v2_model_runner, +): + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", str(int(use_v2_model_runner))) model = get_model(arch) vllm_kwargs = dict(max_model_len=None, model_impl="transformers") diff --git a/vllm/v1/worker/gpu/model_states/__init__.py b/vllm/v1/worker/gpu/model_states/__init__.py index 3717b804ee80..ed95d8285684 100644 --- a/vllm/v1/worker/gpu/model_states/__init__.py +++ b/vllm/v1/worker/gpu/model_states/__init__.py @@ -3,8 +3,9 @@ import torch import torch.nn as nn -from vllm.config import VllmConfig -from vllm.model_executor.layers.attention import CrossAttention, EncoderOnlyAttention +from vllm.config import VllmConfig, get_layers_from_vllm_config +from vllm.model_executor.layers.attention import Attention, CrossAttention +from vllm.v1.attention.backend import AttentionType from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache @@ -27,8 +28,11 @@ def init_model_state( return EncoderDecoderModelState(vllm_config, model, encoder_cache, device) - # Encoder-only models (BERT/RoBERTa): non-causal self-attention, no KV cache. - if any(isinstance(m, EncoderOnlyAttention) for m in model.modules()): + # Encoder-only attention is non-causal and needs no KV cache. + if any( + layer.attn_type == AttentionType.ENCODER_ONLY + for layer in get_layers_from_vllm_config(vllm_config, Attention).values() + ): from vllm.v1.worker.gpu.model_states.encoder_only import EncoderOnlyModelState return EncoderOnlyModelState(vllm_config, model, encoder_cache, device) From 6b0b850a8b1764a66d7ffbb023c0b0e0bbdb900b Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Sun, 16 Aug 2026 15:16:51 -0500 Subject: [PATCH 021/839] [CI] Fit small KV-offload evals within shared memory (#52496) Signed-off-by: Taneem Ibrahim --- tests/evals/gsm8k/test_gsm8k_offloading.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/evals/gsm8k/test_gsm8k_offloading.py b/tests/evals/gsm8k/test_gsm8k_offloading.py index 7ed97f7efcd9..c140a41e2644 100644 --- a/tests/evals/gsm8k/test_gsm8k_offloading.py +++ b/tests/evals/gsm8k/test_gsm8k_offloading.py @@ -133,6 +133,7 @@ class OffloadingModelConfig: connector="OffloadingConnector", # Baseline ~0.49 on 200 questions (measured on GB200). accuracy_threshold=0.39, + cpu_offload_gib=1, ), OffloadingModelConfig( id="offloading-gemma-4-e4b-it", @@ -140,6 +141,7 @@ class OffloadingModelConfig: connector="OffloadingConnector", # Baseline ~0.64 on 200 questions (measured on GB200). accuracy_threshold=0.55, + cpu_offload_gib=1, ), OffloadingModelConfig( id="offloading-qwen3.5-35b", From e3c1cb54fcafe796bc2a0354bb0962edbb4abbb8 Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Mon, 17 Aug 2026 06:17:25 +0800 Subject: [PATCH 022/839] [CI/Build] Avoid duplicate runner startup for multimodal test (#52417) Signed-off-by: Isotr0py --- .../multimodal/generation/test_phi4mm.py | 35 +++-------- .../multimodal/generation/test_qwen2_vl.py | 49 ++++----------- .../generation/vlm_utils/case_filtering.py | 15 ++--- .../generation/vlm_utils/runners.py | 59 ++++++++++++++----- .../multimodal/generation/vlm_utils/types.py | 12 ++-- 5 files changed, 77 insertions(+), 93 deletions(-) diff --git a/tests/models/multimodal/generation/test_phi4mm.py b/tests/models/multimodal/generation/test_phi4mm.py index 5ab75e145aee..71dbfcd5f357 100644 --- a/tests/models/multimodal/generation/test_phi4mm.py +++ b/tests/models/multimodal/generation/test_phi4mm.py @@ -66,6 +66,11 @@ def vllm_to_hf_output( target_dtype = "half" +IMAGE_SIZE_FACTOR_GROUPS = ( + (1.0,), + (1.0, 1.0, 1.0), + (0.25, 0.5, 1.0), +) def run_test( @@ -167,17 +172,6 @@ def patch_hf_processor( @pytest.mark.parametrize("model", models) -@pytest.mark.parametrize( - "size_factors", - [ - # Single-scale - [1.0], - # Single-scale, batched - [1.0, 1.0, 1.0], - # Multi-scale - [0.25, 0.5, 1.0], - ], -) @pytest.mark.parametrize("dtype", [target_dtype]) @pytest.mark.parametrize("max_model_len", [12800]) @pytest.mark.parametrize("max_tokens", [128]) @@ -187,7 +181,6 @@ def test_models( vllm_runner, image_assets, model, - size_factors, dtype: str, max_model_len: int, max_tokens: int, @@ -201,6 +194,7 @@ def test_models( [rescale_image_size(image, factor) for factor in size_factors], None, ) + for size_factors in IMAGE_SIZE_FACTOR_GROUPS for image, prompt in zip(images, HF_IMAGE_PROMPTS) ] @@ -220,19 +214,6 @@ def test_models( @large_gpu_test(min_gb=48) @pytest.mark.parametrize("model", models) -@pytest.mark.parametrize( - "size_factors", - [ - # No image - # [], - # Single-scale - [1.0], - # Single-scale, batched - [1.0, 1.0, 1.0], - # Multi-scale - [0.25, 0.5, 1.0], - ], -) @pytest.mark.parametrize("dtype", [target_dtype]) @pytest.mark.parametrize("max_model_len", [25600]) @pytest.mark.parametrize("max_tokens", [128]) @@ -242,7 +223,6 @@ def test_multi_images_models( vllm_runner, image_assets, model, - size_factors, dtype: str, max_model_len: int, max_tokens: int, @@ -258,7 +238,8 @@ def test_multi_images_models( for factor in size_factors ], None, - ), + ) + for size_factors in IMAGE_SIZE_FACTOR_GROUPS ] run_test( diff --git a/tests/models/multimodal/generation/test_qwen2_vl.py b/tests/models/multimodal/generation/test_qwen2_vl.py index 6148c0bcda7d..19328cad7a23 100644 --- a/tests/models/multimodal/generation/test_qwen2_vl.py +++ b/tests/models/multimodal/generation/test_qwen2_vl.py @@ -29,6 +29,16 @@ def enable_pickle(monkeypatch): models = ["Qwen/Qwen2-VL-2B-Instruct"] target_dtype = "half" +IMAGE_SIZE_FACTOR_GROUPS = ( + (0.5,), + (0.5, 0.5), + (0.25, 0.5, 0.5), +) +VIDEO_SIZE_FACTOR_GROUPS = ( + (0.5,), + (0.5, 0.5), + (0.25, 0.25, 0.5), +) IMAGE_PLACEHOLDER = "<|vision_start|><|image_pad|><|vision_end|>" VIDEO_PLACEHOLDER = "<|vision_start|><|video_pad|><|vision_end|>" @@ -323,17 +333,6 @@ def run_embedding_input_test( @pytest.mark.core_model @pytest.mark.parametrize("model", models) -@pytest.mark.parametrize( - "size_factors", - [ - # Single-scale - [0.5], - # Single-scale, batched - [0.5, 0.5], - # Multi-scale - [0.25, 0.5, 0.5], - ], -) @pytest.mark.parametrize("dtype", [target_dtype]) @pytest.mark.parametrize("max_tokens", [128]) @pytest.mark.parametrize("num_logprobs", [10]) @@ -341,7 +340,6 @@ def test_qwen2_vl_image_embeddings_input( vllm_runner, image_assets, model, - size_factors, dtype, max_tokens, num_logprobs, @@ -355,6 +353,7 @@ def test_qwen2_vl_image_embeddings_input( [rescale_image_size(image, factor) for factor in size_factors], [], ) + for size_factors in IMAGE_SIZE_FACTOR_GROUPS for image, prompt in zip(images, IMAGE_PROMPTS) ] @@ -372,17 +371,6 @@ def test_qwen2_vl_image_embeddings_input( @pytest.mark.core_model @pytest.mark.parametrize("model", models) -@pytest.mark.parametrize( - "size_factors", - [ - # Single-scale - [0.5], - # Single-scale, batched - [0.5, 0.5], - # Multi-scale - [0.25, 0.5, 0.5], - ], -) @pytest.mark.parametrize("dtype", [target_dtype]) @pytest.mark.parametrize("max_tokens", [128]) @pytest.mark.parametrize("num_logprobs", [10]) @@ -390,7 +378,6 @@ def test_qwen2_vl_multiple_image_embeddings_input( vllm_runner, image_assets, model, - size_factors, dtype: str, max_tokens: int, num_logprobs: int, @@ -406,6 +393,7 @@ def test_qwen2_vl_multiple_image_embeddings_input( ], [], ) + for size_factors in IMAGE_SIZE_FACTOR_GROUPS ] run_embedding_input_test( @@ -422,17 +410,6 @@ def test_qwen2_vl_multiple_image_embeddings_input( @pytest.mark.core_model @pytest.mark.parametrize("model", models) -@pytest.mark.parametrize( - "size_factors", - [ - # Single-scale - [0.5], - # Single-scale, batched - [0.5, 0.5], - # Multi-scale - [0.25, 0.25, 0.5], - ], -) @pytest.mark.parametrize("dtype", [target_dtype]) @pytest.mark.parametrize("max_tokens", [128]) @pytest.mark.parametrize("num_logprobs", [10]) @@ -440,7 +417,6 @@ def test_qwen2_vl_video_embeddings_input( vllm_runner, video_assets, model, - size_factors, dtype: str, max_tokens: int, num_logprobs: int, @@ -457,6 +433,7 @@ def test_qwen2_vl_video_embeddings_input( [], [rescale_video_size(video, factor) for factor in size_factors], ) + for size_factors in VIDEO_SIZE_FACTOR_GROUPS for video, prompt in zip(sampled_vids, VIDEO_PROMPTS) ] diff --git a/tests/models/multimodal/generation/vlm_utils/case_filtering.py b/tests/models/multimodal/generation/vlm_utils/case_filtering.py index 116eead7a70a..cbe660425751 100644 --- a/tests/models/multimodal/generation/vlm_utils/case_filtering.py +++ b/tests/models/multimodal/generation/vlm_utils/case_filtering.py @@ -94,7 +94,8 @@ def get_model_type_cases( test_info.needs_video_metadata ) - # No sizes passed for custom inputs, since inputs are directly provided + # Keep all size batches in one test case so they share the same model + # instances. No sizes are passed for preprocessed audio or custom inputs. if test_type not in ( VLMTestType.CUSTOM_INPUTS, VLMTestType.AUDIO, @@ -102,7 +103,9 @@ def get_model_type_cases( wrapped_sizes = get_wrapped_test_sizes(test_info, test_type) if wrapped_sizes is None: raise ValueError(f"Sizes must be set for test type {test_type}") - iter_kwargs["size_wrapper"] = wrapped_sizes + if not wrapped_sizes: + return [] + iter_kwargs["size_wrappers"] = (wrapped_sizes,) # Otherwise expand the custom test options instead elif test_type == VLMTestType.CUSTOM_INPUTS: @@ -127,9 +130,8 @@ def get_parametrized_options( create_new_process_for_each_test: bool, ): """Converts all of our VLMTestInfo into an expanded list of parameters. - This is similar to nesting pytest parametrize calls, but done directly - through an itertools product so that each test can set things like - size factors etc, while still running in isolated test cases. + Runner configuration values are expanded through an itertools product. + Input size batches stay grouped so they can share model instances. """ matching_tests = get_filtered_test_settings( test_settings, test_type, create_new_process_for_each_test @@ -149,8 +151,7 @@ def get_wrapped_test_sizes( test_info: VLMTestInfo, test_type: VLMTestType ) -> tuple[ImageSizeWrapper, ...]: """Given a test info which may have size factors or fixed sizes, wrap them - and combine them into an iterable, each of which will be used in parameter - expansion. + and combine them into an iterable of request batches. Args: test_info: Test configuration to be expanded. diff --git a/tests/models/multimodal/generation/vlm_utils/runners.py b/tests/models/multimodal/generation/vlm_utils/runners.py index 218339ef1dff..571ea54250b7 100644 --- a/tests/models/multimodal/generation/vlm_utils/runners.py +++ b/tests/models/multimodal/generation/vlm_utils/runners.py @@ -4,6 +4,7 @@ types / modalities. """ +import itertools from pathlib import PosixPath from .....conftest import ( @@ -27,9 +28,14 @@ def run_single_image_test( vllm_runner: type[VllmRunner], image_assets: ImageTestAssets, ): - assert test_case.size_wrapper is not None - inputs = builders.build_single_image_inputs_from_test_info( - model_test_info, image_assets, test_case.size_wrapper, tmp_path + assert test_case.size_wrappers + inputs = list( + itertools.chain.from_iterable( + builders.build_single_image_inputs_from_test_info( + model_test_info, image_assets, size_wrapper, tmp_path + ) + for size_wrapper in test_case.size_wrappers + ) ) core.run_test( @@ -55,9 +61,14 @@ def run_multi_image_test( vllm_runner: type[VllmRunner], image_assets: ImageTestAssets, ): - assert test_case.size_wrapper is not None - inputs = builders.build_multi_image_inputs_from_test_info( - model_test_info, image_assets, test_case.size_wrapper, tmp_path + assert test_case.size_wrappers + inputs = list( + itertools.chain.from_iterable( + builders.build_multi_image_inputs_from_test_info( + model_test_info, image_assets, size_wrapper, tmp_path + ) + for size_wrapper in test_case.size_wrappers + ) ) core.run_test( @@ -82,9 +93,20 @@ def run_embedding_test( vllm_runner: type[VllmRunner], image_assets: ImageTestAssets, ): - assert test_case.size_wrapper is not None - inputs, vllm_embeddings = builders.build_embedding_inputs_from_test_info( - model_test_info, image_assets, test_case.size_wrapper + assert test_case.size_wrappers + inputs_and_embeddings = [ + builders.build_embedding_inputs_from_test_info( + model_test_info, image_assets, size_wrapper + ) + for size_wrapper in test_case.size_wrappers + ] + inputs = list( + itertools.chain.from_iterable(inputs for inputs, _ in inputs_and_embeddings) + ) + vllm_embeddings = list( + itertools.chain.from_iterable( + embeddings for _, embeddings in inputs_and_embeddings + ) ) core.run_test( @@ -110,14 +132,19 @@ def run_video_test( vllm_runner: type[VllmRunner], video_assets: VideoTestAssets, ): - assert test_case.size_wrapper is not None + assert test_case.size_wrappers assert test_case.num_video_frames is not None - inputs = builders.build_video_inputs_from_test_info( - model_test_info, - video_assets, - test_case.size_wrapper, - test_case.num_video_frames, - test_case.needs_video_metadata, + inputs = list( + itertools.chain.from_iterable( + builders.build_video_inputs_from_test_info( + model_test_info, + video_assets, + size_wrapper, + test_case.num_video_frames, + test_case.needs_video_metadata, + ) + for size_wrapper in test_case.size_wrappers + ) ) core.run_test( diff --git a/tests/models/multimodal/generation/vlm_utils/types.py b/tests/models/multimodal/generation/vlm_utils/types.py index af48a1479bad..3722caed2cb3 100644 --- a/tests/models/multimodal/generation/vlm_utils/types.py +++ b/tests/models/multimodal/generation/vlm_utils/types.py @@ -158,11 +158,9 @@ class VLMTestInfo(NamedTuple): num_video_frames: int | tuple[int] = 16 needs_video_metadata: bool = False - # Fixed image sizes / image size factors; most tests use image_size_factors - # The values provided for these two fields will be stacked and expanded - # such that each model will consider each image size factor / image size - # once per tests (much like concatenating and wrapping in one parametrize - # call) + # Fixed image sizes / image size factors; most tests use image_size_factors. + # Each inner iterable defines one request batch. All batches are run against + # the same model instance. image_size_factors: Iterable[Iterable[float]] = IMAGE_SIZE_FACTORS image_sizes: Iterable[Iterable[tuple[int, int]]] | None = None @@ -211,8 +209,8 @@ class ExpandableVLMTestArgs(NamedTuple): num_logprobs: int dtype: str distributed_executor_backend: str | None - # Sizes are used for everything except for custom input tests - size_wrapper: ImageSizeWrapper | None = None + # Sizes are used for everything except audio and custom input tests. + size_wrappers: tuple[ImageSizeWrapper, ...] = () # Video only num_video_frames: int | None = None needs_video_metadata: bool = False From eee538d5daa0c8c969f20d0c48c972155a6f1859 Mon Sep 17 00:00:00 2001 From: Tianyu Guo Date: Mon, 17 Aug 2026 06:21:55 +0800 Subject: [PATCH 023/839] [Bugfix][V1][Multimodal] Ignore stale same-step encoder cache evictions (#52482) Signed-off-by: Tianyu Guo --- tests/v1/core/test_encoder_cache_manager.py | 19 +++++++++++++++++++ vllm/v1/core/encoder_cache_manager.py | 4 +++- 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/tests/v1/core/test_encoder_cache_manager.py b/tests/v1/core/test_encoder_cache_manager.py index e225666f8443..e56bcbf5c63a 100644 --- a/tests/v1/core/test_encoder_cache_manager.py +++ b/tests/v1/core/test_encoder_cache_manager.py @@ -167,6 +167,25 @@ def test_get_freed_mm_hashes_clears_freed_list(): assert manager.get_freed_mm_hashes() == [] +def test_reallocated_hash_is_not_reported_as_freed(): + manager = EncoderCacheManager(cache_size=8) + req_a = MockRequest("reqA", ["a"], [4]) + req_b = MockRequest("reqB", ["b"], [4]) + req_c = MockRequest("reqC", ["c"], [4]) + + manager.allocate(req_a, 0) + manager.allocate(req_b, 0) + manager.free(req_a) + manager.free(req_b) + + assert manager.can_allocate(req_c, 0, int(1e9), 0) + manager.allocate(req_c, 0) + assert manager.can_allocate(req_a, 0, int(1e9), 0) + manager.allocate(req_a, 0) + + assert manager.get_freed_mm_hashes() == ["b"] + + def test_schedule_request_multi_images_respect_space_limit(): manager = EncoderCacheManager(cache_size=10) req = MockRequest("reqA", ["a", "b"], [5, 6]) diff --git a/vllm/v1/core/encoder_cache_manager.py b/vllm/v1/core/encoder_cache_manager.py index 02133ff5b888..8d2a81a11b13 100644 --- a/vllm/v1/core/encoder_cache_manager.py +++ b/vllm/v1/core/encoder_cache_manager.py @@ -269,7 +269,9 @@ def get_freed_mm_hashes(self) -> list[str]: encoder outputs can be removed from their caches. The internal list is cleared after this call. """ - freed = self.freed + # An entry evicted early in the scheduling pass can be allocated again + # later in the same pass. Keep its worker-side tensor in that case. + freed = [mm_hash for mm_hash in self.freed if mm_hash not in self.cached] self.freed = [] return freed From dc9ae4b8ac2331991ad7091812ef82ece4f8fdc2 Mon Sep 17 00:00:00 2001 From: Guanyi Chen <939416532@qq.com> Date: Mon, 17 Aug 2026 07:09:46 +0800 Subject: [PATCH 024/839] [Bugfix][Mooncake] Reference GPU blocks for in-flight store jobs and key the store ledger by store_job_id (#52372) Signed-off-by: Guanyi Chen <939416532@qq.com> --- .../unit/test_mooncake_store_hma_e2e.py | 15 +- .../unit/test_mooncake_store_scheduler.py | 179 ++++++++++----- .../unit/test_mooncake_store_worker.py | 216 ++++++++++-------- .../v1/mooncake/store/connector.py | 22 +- .../kv_connector/v1/mooncake/store/data.py | 24 +- .../v1/mooncake/store/scheduler.py | 104 ++++++--- .../kv_connector/v1/mooncake/store/worker.py | 164 ++++++++----- 7 files changed, 474 insertions(+), 250 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py index 6c8af7fb73c1..32457f13d525 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py @@ -228,12 +228,11 @@ def _fake_thread_init(*args, **kwargs): block_ids=([0, 1, 2, 3], [0, 1, 2, 3]), block_hashes=hs, can_save=True, + store_job_id=1, ) - send_thread.add_stored_request("r0") - # Put the request in the queue so task_done() doesn't underflow. - send_thread.request_queue.put(save_req) - req = send_thread.request_queue.get() - send_thread._handle_request(req) + # add_request also queues the job, so task_done() doesn't underflow. + send_thread.add_request(save_req) + send_thread._handle_request(send_thread.request_queue.get()) # Point worker.store at the dict store (the worker constructor captured # the MagicMock; replace with the real dict store for lookup). @@ -499,15 +498,15 @@ def batch_put_from_multi_buffers(self, keys, addrs, sizes, *a, **k): can_save=True, num_prompt_tokens=12, partial_tail_offloads=[(1, 7, 12)], + store_job_id=1, ) req.current_event = event - send.add_stored_request("r1") - send.request_queue.put(req) + send.add_request(req) send._handle_request(send.request_queue.get()) assert send.request_queue.qsize() == 0 assert store._data - assert send.stored_requests["r1"] == 0 + assert send.stored_requests["r1"] == set() event.synchronize.assert_called_once() diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py b/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py index 13960c40340e..998a8408e651 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py @@ -5,12 +5,14 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( LoadSpec, + MooncakeStoreWorkerMetadata, ReqMeta, RequestTracker, ) from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.scheduler import ( MooncakeStoreScheduler, ) +from vllm.v1.core.block_pool import BlockPool def _make_bare_scheduler( @@ -27,6 +29,12 @@ def _make_bare_scheduler( scheduler._unfinished_request_ids = {"req-0"} scheduler._unfinished_requests = {} scheduler._request_trackers = {} + scheduler._gpu_block_pool = BlockPool( + num_gpu_blocks=64, enable_caching=True, hash_block_size=hash_block_size + ) + scheduler._num_workers = 1 + scheduler._next_store_job_id = 0 + scheduler._pinned_saves = {} return scheduler @@ -64,6 +72,14 @@ def _make_preemption_scheduler_output(): ) +def _make_worker_output(completed_saves: dict[int, int]) -> SimpleNamespace: + return SimpleNamespace( + kv_connector_worker_meta=MooncakeStoreWorkerMetadata( + completed_saves=completed_saves + ) + ) + + def _add_unfinished_request( scheduler: MooncakeStoreScheduler, *, @@ -133,7 +149,7 @@ def test_cached_request_without_spec_decode_keeps_current_step_save_overlap(): assert tracker.num_saved_tokens == 48 -def test_preemption_resets_tracker_before_request_finished(): +def test_preemption_resets_tracker(): scheduler = _make_bare_scheduler() _add_unfinished_request( scheduler, @@ -152,8 +168,6 @@ def test_preemption_resets_tracker_before_request_finished(): assert tracker.token_ids is None assert tracker.has_pending_offload is False assert tracker.prefill_end_tokens == 0 - request = SimpleNamespace(request_id="req-0") - assert scheduler.request_finished(request, ([0, 1],)) == (False, None) def test_preemption_clears_stale_load_state(): @@ -215,7 +229,7 @@ def test_pending_load_does_not_co_queue_save(): # enqueue a save in the same scheduling step. Co-queuing both produces a # recv+send pair for the same req_id, and the scheduler's # _update_from_kv_xfer_finished then trips `assert req_id in self.requests` - # when both completions land for the delay-freed request. + # when a completion lands for a request it has already dropped. scheduler = _make_bare_scheduler() _make_pending_load_unfinished_request( scheduler, @@ -238,9 +252,7 @@ def test_pending_load_does_not_co_queue_save(): # Load is still issued as planned. assert req_meta.load_spec is not None assert req_meta.load_spec.can_load is True - # And the tracker's saved-tokens watermark stays at 0 so request_finished - # later sees `num_saved_tokens <= 0` and frees immediately rather than - # waiting for a finished_sending that will never come. + # And the save watermark does not advance for a save that was never queued. tracker = scheduler._request_trackers["req-0"] assert tracker.num_saved_tokens == 0 @@ -661,30 +673,34 @@ def test_disabled_lookup_reports_no_hit_without_querying_client(): assert scheduler.load_specs == {} -def test_pending_partial_tail_emits_offload_only_reqmeta(): - # A sub-block prompt never produces a block-aligned save, so the partial- - # tail offload arriving this step is emitted as an offload-only ReqMeta - # (can_save=True so it takes the normal enqueue path, token_len_chunk=0 so - # the worker skips the normal save). Pending-offload state delays the free - # without advancing the normal-save watermark before the put succeeds. - scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True) +def _add_pending_partial_tail_request( + scheduler: MooncakeStoreScheduler, + *, + num_tokens: int, + block_hashes: list[bytes], + block_ids: tuple[list[int], ...], +) -> SimpleNamespace: + """Register a sub-block request and return the step that offloads its tail. + + The CoW block holding the tail is block 7, which the core deliberately keeps + out of the request's block table. + """ request = SimpleNamespace( - all_token_ids=list(range(12)), - block_hashes=[b"h0", b"h1", b"h2"], + all_token_ids=list(range(num_tokens)), + block_hashes=block_hashes, num_output_placeholders=0, num_prompt_tokens=12, ) - scheduler._unfinished_requests["req-0"] = (request, ([0],)) + scheduler._unfinished_requests["req-0"] = (request, block_ids) scheduler._request_trackers["req-0"] = RequestTracker( req_id="req-0", - token_len=12, - allocated_block_ids=([0],), + token_len=num_tokens, + allocated_block_ids=block_ids, num_saved_tokens=0, - token_ids=list(range(12)), - prefill_end_tokens=12, + token_ids=list(range(num_tokens)), + prefill_end_tokens=num_tokens, ) - - out = SimpleNamespace( + return SimpleNamespace( finished_req_ids=set(), preempted_req_ids=set(), scheduled_new_reqs=[], @@ -699,6 +715,21 @@ def test_pending_partial_tail_emits_offload_only_reqmeta(): partial_tail_offloads={"req-0": [(1, 7, 12)]}, ) + +def test_pending_partial_tail_emits_offload_only_reqmeta(): + # A sub-block prompt never produces a block-aligned save, so the partial- + # tail offload arriving this step is emitted as an offload-only ReqMeta + # (can_save=True so it takes the normal enqueue path, token_len_chunk=0 so + # the worker skips the normal save), without advancing the normal-save + # watermark before the put succeeds. + scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True) + out = _add_pending_partial_tail_request( + scheduler, + num_tokens=12, + block_hashes=[b"h0", b"h1", b"h2"], + block_ids=([0],), + ) + meta = scheduler.build_connector_meta(out) assert len(meta.requests) == 1 @@ -712,42 +743,16 @@ def test_pending_partial_tail_emits_offload_only_reqmeta(): tracker = scheduler._request_trackers["req-0"] assert tracker.num_saved_tokens == 0 assert tracker.has_pending_offload is True - request = SimpleNamespace(request_id="req-0") - assert scheduler.request_finished(request, ([0],)) == (True, None) def test_resumed_partial_tail_uses_handoff_boundary(): scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True) - request = SimpleNamespace( - all_token_ids=list(range(20)), + # Resumption replays prompt + previously generated tokens. + out = _add_pending_partial_tail_request( + scheduler, + num_tokens=20, block_hashes=[b"h0", b"h1", b"h2", b"h3", b"h4"], - num_output_placeholders=0, - num_prompt_tokens=12, - ) - scheduler._unfinished_requests["req-0"] = (request, ([0, 1],)) - scheduler._request_trackers["req-0"] = RequestTracker( - req_id="req-0", - token_len=20, - allocated_block_ids=([0, 1],), - num_saved_tokens=0, - token_ids=list(range(20)), - # Resumption replays prompt + previously generated tokens. - prefill_end_tokens=20, - ) - - out = SimpleNamespace( - finished_req_ids=set(), - preempted_req_ids=set(), - scheduled_new_reqs=[], - scheduled_cached_reqs=SimpleNamespace( - req_ids=[], - new_block_ids=[], - num_computed_tokens=[], - resumed_req_ids=set(), - ), - num_scheduled_tokens={}, - scheduled_spec_decode_tokens={}, - partial_tail_offloads={"req-0": [(1, 7, 12)]}, + block_ids=([0, 1],), ) meta = scheduler.build_connector_meta(out) @@ -791,3 +796,69 @@ def test_resumed_partial_tail_attached_to_save_keeps_handoff_boundary(): tracker = scheduler._request_trackers["req-0"] assert tracker.num_saved_tokens == 48 assert tracker.has_pending_offload is True + + +def test_partial_tail_cow_block_is_referenced_for_the_job(): + # The CoW block a partial-tail offload reads is deliberately kept out of the + # request block table, so it is absent from ReqMeta.block_ids. The worker + # DMAs out of it just as asynchronously, so it needs its own reference. + scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True) + out = _add_pending_partial_tail_request( + scheduler, + num_tokens=12, + block_hashes=[b"h0", b"h1", b"h2"], + block_ids=([0],), + ) + pool = scheduler._gpu_block_pool + + meta = scheduler.build_connector_meta(out) + + store_job_id = meta.requests[0].store_job_id + # It leads the list, as in `pop_blocks_for_free`, so that the reversed free + # puts it last in eviction priority. + assert scheduler._pinned_saves[store_job_id][0] == [7, 0] + assert pool.blocks[7].ref_cnt == 1 + + scheduler.update_connector_output(_make_worker_output({store_job_id: 1})) + assert pool.blocks[7].ref_cnt == 0 + + +def test_store_job_blocks_are_released_once_every_rank_reports(): + # Every rank DMAs the job's blocks on its own, so the reference can only be + # dropped once the last of them reports. Until then the engine has to keep + # stepping: a completion only reaches the scheduler as worker metadata + # attached to a step, and a finishing request no longer defers its own free. + scheduler = _make_bare_scheduler() + scheduler._num_workers = 2 + _add_unfinished_request( + scheduler, + token_ids=list(range(48)), + block_hashes=[b"h0", b"h1", b"h2"], + prefill_end_tokens=48, + ) + pool = scheduler._gpu_block_pool + assert scheduler.has_pending_push_work() is False + + meta = scheduler.build_connector_meta( + _make_scheduler_output(scheduled_spec_tokens=None) + ) + store_job_id = meta.requests[0].store_job_id + assert pool.blocks[2].ref_cnt == 1 + assert scheduler.has_pending_push_work() is True + + scheduler.update_connector_output(_make_worker_output({store_job_id: 1})) + assert pool.blocks[2].ref_cnt == 1 + assert scheduler.has_pending_push_work() is True + + scheduler.update_connector_output(_make_worker_output({store_job_id: 1})) + assert pool.blocks[2].ref_cnt == 0 + assert scheduler.has_pending_push_work() is False + + +def test_worker_metadata_aggregates_completions_across_ranks(): + # The engine merges each rank's metadata before the scheduler sees it, so a + # job that every rank finished in one step arrives as a single count. + merged = MooncakeStoreWorkerMetadata(completed_saves={1: 1}).aggregate( + MooncakeStoreWorkerMetadata(completed_saves={1: 1, 2: 1}) + ) + assert merged.completed_saves == {1: 2, 2: 1} diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index 034297100a7a..4655c5716ec8 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import contextlib +import itertools import json import logging import math @@ -156,6 +158,17 @@ def _make_load_req( ) +_TEST_SAVE_SEQ = itertools.count(1) + + +def _run_store_req(thread, req_meta: ReqMeta) -> None: + """Register, enqueue and run a store job the way the worker does.""" + if req_meta.store_job_id is None: + req_meta.store_job_id = next(_TEST_SAVE_SEQ) + thread.add_request(req_meta) + thread._handle_request(thread.request_queue.get()) + + def _make_store_req(req_id: str, block_hashes: list[bytes]) -> ReqMeta: return ReqMeta( req_id=req_id, @@ -385,27 +398,23 @@ def test_store_sending_thread_skips_request_during_cpu_pressure(): ] thread = _make_store_sending_thread(store) - thread.add_stored_request("req-a") - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"])) + _run_store_req(thread, _make_store_req("req-a", [b"a0", b"a1"])) assert thread._store_pressure_active is True assert "req-a" in thread._skip_store_requests assert store.batch_put_from_multi_buffers.call_count == 1 - thread.add_stored_request("req-a") - thread._handle_request(_make_store_req("req-a", [b"a2", b"a3"])) + _run_store_req(thread, _make_store_req("req-a", [b"a2", b"a3"])) assert store.batch_put_from_multi_buffers.call_count == 1 - thread.add_stored_request("req-b") - thread._handle_request(_make_store_req("req-b", [b"b0", b"b1"])) + _run_store_req(thread, _make_store_req("req-b", [b"b0", b"b1"])) assert thread._store_pressure_active is False assert "req-a" not in thread._skip_store_requests assert store.batch_put_from_multi_buffers.call_count == 2 - thread.add_stored_request("req-a") - thread._handle_request(_make_store_req("req-a", [b"a4", b"a5"])) + _run_store_req(thread, _make_store_req("req-a", [b"a4", b"a5"])) assert store.batch_put_from_multi_buffers.call_count == 3 @@ -418,8 +427,7 @@ def test_store_sending_thread_records_mooncake_metrics(): stats = MooncakeStoreConnectorStats() thread._record_operation_cb = stats.record_operation - thread.add_stored_request("req-a") - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"])) + _run_store_req(thread, _make_store_req("req-a", [b"a0", b"a1"])) assert len(stats.data["save_exists"]) == 1 assert stats.data["save_exists"][0]["num_keys"] == 2 @@ -497,16 +505,16 @@ def test_store_sending_thread_delta_saves_only_new_full_attention_chunks(): store.batch_put_from_multi_buffers.return_value = [256, 256] thread = _make_store_sending_thread(store) - thread.add_stored_request("req-a") thread._saved_offset["req-a"] = 32 - thread._handle_request( + _run_store_req( + thread, ReqMeta( req_id="req-a", token_len_chunk=64, block_ids=([0, 1, 2, 3],), block_hashes=[b"a0", b"a1", b"a2", b"a3"], can_save=True, - ) + ), ) keys = store.batch_is_exist.call_args.args[0] @@ -523,16 +531,16 @@ def test_store_sending_thread_delta_strides_with_local_phase(): store.batch_put_from_multi_buffers.return_value = [256] thread = _make_store_sending_thread(store, tp_rank=0, put_step=2) - thread.add_stored_request("req-a") thread._saved_offset["req-a"] = 16 - thread._handle_request( + _run_store_req( + thread, ReqMeta( req_id="req-a", token_len_chunk=64, block_ids=([0, 1, 2, 3],), block_hashes=[b"a0", b"a1", b"a2", b"a3"], can_save=True, - ) + ), ) keys = store.batch_is_exist.call_args.args[0] @@ -550,15 +558,15 @@ def test_tp_sharded_group_saves_every_block_on_every_rank(): thread = _make_store_sending_thread(store, tp_rank=0, put_step=2) thread.group_put_steps = [1] - thread.add_stored_request("req-a") - thread._handle_request( + _run_store_req( + thread, ReqMeta( req_id="req-a", token_len_chunk=64, block_ids=([0, 1, 2, 3],), block_hashes=[b"a0", b"a1", b"a2", b"a3"], can_save=True, - ) + ), ) keys = store.batch_is_exist.call_args.args[0] @@ -576,15 +584,15 @@ def test_store_sending_thread_retries_skipped_range_after_pressure(): thread._store_pressure_active = True thread._skip_store_requests.add("req-a") - thread.add_stored_request("req-a") - thread._handle_request( + _run_store_req( + thread, ReqMeta( req_id="req-a", token_len_chunk=16, block_ids=([0],), block_hashes=[b"a0"], can_save=True, - ) + ), ) store.batch_is_exist.assert_not_called() @@ -595,15 +603,15 @@ def test_store_sending_thread_retries_skipped_range_after_pressure(): # The next batch resumes from offset 0, re-covering the chunk skipped under # pressure (chunk 0) rather than losing it. - thread.add_stored_request("req-a") - thread._handle_request( + _run_store_req( + thread, ReqMeta( req_id="req-a", token_len_chunk=64, block_ids=([0, 1, 2, 3],), block_hashes=[b"a0", b"a1", b"a2", b"a3"], can_save=True, - ) + ), ) keys = store.batch_is_exist.call_args.args[0] @@ -703,13 +711,11 @@ def test_partial_tail_offload_honors_active_pressure_gate(): thread = _make_partial_tail_send_thread(store) thread._store_pressure_active = True thread._skip_store_requests.add("req-a") - thread.add_stored_request("req-a") - thread._handle_request(_make_partial_tail_req([1, 2, 3])) + _run_store_req(thread, _make_partial_tail_req([1, 2, 3])) store.batch_is_exist.assert_not_called() store.batch_put_from_multi_buffers.assert_not_called() - assert thread.stored_requests["req-a"] == 0 def test_partial_tail_put_failure_activates_pressure_gate(): @@ -717,17 +723,14 @@ def test_partial_tail_put_failure_activates_pressure_gate(): store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) store.batch_put_from_multi_buffers.return_value = [256, -200, 256] thread = _make_partial_tail_send_thread(store) - thread.add_stored_request("req-a") - thread._handle_request(_make_partial_tail_req([1, 2, 3])) + _run_store_req(thread, _make_partial_tail_req([1, 2, 3])) assert thread._store_pressure_active is True assert thread._skip_store_requests == {"req-a"} assert thread._saved_offset.get("req-a", 0) == 0 - assert thread.stored_requests["req-a"] == 0 - thread.add_stored_request("req-a") - thread._handle_request(_make_partial_tail_req([1, 2, 3])) + _run_store_req(thread, _make_partial_tail_req([1, 2, 3])) assert store.batch_put_from_multi_buffers.call_count == 1 @@ -737,16 +740,16 @@ def test_store_sending_thread_delta_start_rank_saves_second_local_chunk(): store.batch_put_from_multi_buffers.return_value = [256, 256] thread = _make_store_sending_thread(store, tp_rank=1, put_step=2) - thread.add_stored_request("req-a") thread._saved_offset["req-a"] = 16 - thread._handle_request( + _run_store_req( + thread, ReqMeta( req_id="req-a", token_len_chunk=64, block_ids=([0, 1, 2, 3],), block_hashes=[b"a0", b"a1", b"a2", b"a3"], can_save=True, - ) + ), ) keys = store.batch_is_exist.call_args.args[0] @@ -790,16 +793,16 @@ def test_store_sending_thread_delta_saves_only_new_masked_chunks(): token_databases=[db_full, db_masked], ) - thread.add_stored_request("req-a") thread._saved_offset["req-a"] = 32 - thread._handle_request( + _run_store_req( + thread, ReqMeta( req_id="req-a", token_len_chunk=64, block_ids=([0, 1, 2, 3], [0, 1, 2, 3]), block_hashes=[b"a0", b"a1", b"a2", b"a3"], can_save=True, - ) + ), ) keys = store.batch_is_exist.call_args.args[0] @@ -845,15 +848,15 @@ def test_store_sending_thread_prepares_missing_chunks_once_per_group(): coord=coord, token_databases=[db0, db1], ) - thread.add_stored_request("req-a") - thread._handle_request( + _run_store_req( + thread, ReqMeta( req_id="req-a", token_len_chunk=48, block_ids=([0, 1, 2], [2, 1, 0]), block_hashes=[b"a0", b"a1", b"a2"], can_save=True, - ) + ), ) db0.prepare_value.assert_not_called() @@ -881,46 +884,71 @@ def test_store_sending_thread_only_skips_on_no_available_handle(): ] thread = _make_store_sending_thread(store) - thread.add_stored_request("req-a") - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"])) + _run_store_req(thread, _make_store_req("req-a", [b"a0", b"a1"])) assert thread._store_pressure_active is False assert "req-a" not in thread._skip_store_requests assert store.batch_put_from_multi_buffers.call_count == 1 - thread.add_stored_request("req-a") - thread._handle_request(_make_store_req("req-a", [b"a2", b"a3"])) + _run_store_req(thread, _make_store_req("req-a", [b"a2", b"a3"])) assert store.batch_put_from_multi_buffers.call_count == 2 -def test_store_sending_thread_releases_pin_on_batch_is_exist_failure(): - # `batch_is_exist` raising must still decrement `stored_requests` so the - # scheduler can drop `delay_free_blocks` and release the pinned GPU blocks. +@pytest.mark.parametrize( + "failing_call", ["batch_is_exist", "batch_put_from_multi_buffers"] +) +def test_store_sending_thread_reports_job_when_store_raises(failing_call): + # A store that blows up must still report its job, or the scheduler keeps + # the job's GPU block references for the rest of the run. store = MagicMock() - store.batch_is_exist.side_effect = RuntimeError("mooncake down") + store.batch_is_exist.return_value = [0, 0] + getattr(store, failing_call).side_effect = RuntimeError("mooncake down") thread = _make_store_sending_thread(store) + req = _make_store_req("req-a", [b"a0", b"a1"]) - thread.add_stored_request("req-a") - with pytest.raises(RuntimeError): - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"])) + with contextlib.suppress(RuntimeError): + _run_store_req(thread, req) - assert thread.stored_requests["req-a"] == 0 - store.batch_put_from_multi_buffers.assert_not_called() + assert thread.take_completed_saves() == {req.store_job_id: 1} -def test_store_sending_thread_releases_pin_on_batch_put_failure(): - # `batch_put_from_multi_buffers` raising is logged (not re-raised), and the - # pin must still be released through the finally block. - store = MagicMock() - store.batch_is_exist.return_value = [0, 0] - store.batch_put_from_multi_buffers.side_effect = RuntimeError("rdma error") - thread = _make_store_sending_thread(store) +def test_store_sending_thread_reports_job_when_the_preamble_raises(): + # The report has to survive a failure before the first store call too, so + # every dequeue leaves through the same exit. + thread = _make_store_sending_thread(MagicMock()) + req = _make_store_req("req-a", [b"a0", b"a1"]) + req.token_len_chunk = None # type: ignore[assignment] + + with contextlib.suppress(TypeError): + _run_store_req(thread, req) - thread.add_stored_request("req-a") - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"])) + assert thread.take_completed_saves() == {req.store_job_id: 1} + thread.request_queue.task_done.assert_called_once() - assert thread.stored_requests["req-a"] == 0 + +def test_stale_store_job_cannot_touch_a_reused_request_id(): + # A preempted request resumes under its original id, so a job left over from + # the retired generation carries a req_id that now belongs to a live one. + thread = _make_store_sending_thread(MagicMock()) + stale = _make_store_req("req-a", [b"a0", b"a1"]) + stale.store_job_id = 1 + thread.add_request(stale) + thread._record_saved(stale, 32) + thread.delete_finished_stored_request("req-a") + + live = _make_store_req("req-a", [b"a0", b"a1"]) + live.store_job_id = 2 + thread.add_request(live) + + thread.finish_store_job(stale) + thread._record_saved(stale, 64) + thread._mark_request_skipped_for_pressure(stale) + + assert thread.is_live_store_job(live) + assert not thread.is_live_store_job(stale) + assert thread._saved_offset.get("req-a") is None + assert "req-a" not in thread._skip_store_requests def test_store_recving_thread_reports_failed_block_ids(): @@ -991,8 +1019,7 @@ def test_store_sending_thread_passes_replicate_config_when_preferred_segment_set replicate_config = SimpleNamespace(preferred_segment="10.0.0.7:50053") thread = _make_store_sending_thread(store, replicate_config=replicate_config) - thread.add_stored_request("req-a") - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"])) + _run_store_req(thread, _make_store_req("req-a", [b"a0", b"a1"])) assert store.batch_put_from_multi_buffers.call_count == 1 call_args = store.batch_put_from_multi_buffers.call_args.args @@ -1010,8 +1037,7 @@ def test_store_sending_thread_passes_default_replicate_config_when_no_preferred_ replicate_config = SimpleNamespace() thread = _make_store_sending_thread(store, replicate_config=replicate_config) - thread.add_stored_request("req-a") - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"])) + _run_store_req(thread, _make_store_req("req-a", [b"a0", b"a1"])) assert store.batch_put_from_multi_buffers.call_count == 1 call_args = store.batch_put_from_multi_buffers.call_args.args @@ -1067,8 +1093,7 @@ def test_store_sending_thread_sets_group_ids_when_enabled(): supports_group_ids=True, ) - thread.add_stored_request("req-a") - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"])) + _run_store_req(thread, _make_store_req("req-a", [b"a0", b"a1"])) assert store.batch_put_from_multi_buffers.call_count == 1 keys, _addrs, _sizes, config = store.batch_put_from_multi_buffers.call_args.args @@ -1092,8 +1117,7 @@ def test_store_sending_thread_leaves_group_ids_unchanged_when_flag_disabled(): supports_group_ids=True, ) - thread.add_stored_request("req-a") - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"])) + _run_store_req(thread, _make_store_req("req-a", [b"a0", b"a1"])) assert store.batch_put_from_multi_buffers.call_count == 1 assert store.batch_put_from_multi_buffers.call_args.args[3] is replicate_config @@ -1112,8 +1136,7 @@ def test_store_sending_thread_leaves_group_ids_unchanged_when_unsupported(): supports_group_ids=False, ) - thread.add_stored_request("req-a") - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"])) + _run_store_req(thread, _make_store_req("req-a", [b"a0", b"a1"])) assert store.batch_put_from_multi_buffers.call_count == 1 assert store.batch_put_from_multi_buffers.call_args.args[3] is replicate_config @@ -1180,8 +1203,7 @@ def test_store_sending_thread_group_id_excludes_physical_sharding(): supports_group_ids=True, ) - thread.add_stored_request("req-a") - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"])) + _run_store_req(thread, _make_store_req("req-a", [b"a0", b"a1"])) keys, _addrs, _sizes, config = store.batch_put_from_multi_buffers.call_args.args assert keys == [ @@ -1210,8 +1232,7 @@ def test_store_sending_thread_multiple_segments_share_logical_group_id(): supports_group_ids=True, ) - thread.add_stored_request("req-a") - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"])) + _run_store_req(thread, _make_store_req("req-a", [b"a0", b"a1"])) keys, addrs, sizes, config = store.batch_put_from_multi_buffers.call_args.args assert keys == [ @@ -1260,8 +1281,7 @@ def test_store_sending_thread_group_ids_share_across_kv_cache_groups(): supports_group_ids=True, ) - thread.add_stored_request("req-a") - thread._handle_request(_make_multi_group_store_req("req-a", [b"a0", b"a1"])) + _run_store_req(thread, _make_multi_group_store_req("req-a", [b"a0", b"a1"])) keys, addrs, sizes, config = store.batch_put_from_multi_buffers.call_args.args assert keys == [ @@ -1296,8 +1316,7 @@ def test_store_sending_thread_group_ids_follow_missing_key_filter(): supports_group_ids=True, ) - thread.add_stored_request("req-a") - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"])) + _run_store_req(thread, _make_store_req("req-a", [b"a0", b"a1"])) keys, _addrs, _sizes, config = store.batch_put_from_multi_buffers.call_args.args assert keys == ["test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6131"] @@ -1866,15 +1885,15 @@ def test_store_sending_thread_clamps_token_len_to_lcm(): # token_len_chunk=33 clamps to 32 → 2 chunks (not 3 with a partial 1-token chunk). thread = _make_store_sending_thread(store) - thread.add_stored_request("r0") - thread._handle_request( + _run_store_req( + thread, ReqMeta( req_id="r0", token_len_chunk=33, block_ids=([0, 1, 2],), block_hashes=[b"a0", b"a1", b"a2"], can_save=True, - ) + ), ) keys = store.batch_put_from_multi_buffers.call_args.args[0] @@ -1905,20 +1924,19 @@ def test_store_sending_thread_skips_when_token_len_below_lcm(): store, coord=coord, token_databases=[db], block_size=64 ) - thread.add_stored_request("r0") - thread._handle_request( + _run_store_req( + thread, ReqMeta( req_id="r0", token_len_chunk=32, block_ids=([0, 1],), block_hashes=[b"a0", b"a1"], can_save=True, - ) + ), ) store.batch_is_exist.assert_not_called() store.batch_put_from_multi_buffers.assert_not_called() - assert thread.stored_requests["r0"] == 0 def test_store_sending_thread_only_stores_swa_blocks_in_window(): @@ -1982,15 +2000,15 @@ def test_store_sending_thread_only_stores_swa_blocks_in_window(): ) hs = [bytes([i + 1]) * 4 for i in range(8)] - thread.add_stored_request("r0") - thread._handle_request( + _run_store_req( + thread, ReqMeta( req_id="r0", token_len_chunk=64, block_ids=([0, 1], list(range(8))), block_hashes=hs, can_save=True, - ) + ), ) keys = store.batch_put_from_multi_buffers.call_args.args[0] @@ -2057,16 +2075,16 @@ def test_store_sending_thread_delta_saves_only_new_swa_boundary_chunks(): ) hs = [bytes([i + 1]) * 4 for i in range(8)] - thread.add_stored_request("r0") thread._saved_offset["r0"] = 32 - thread._handle_request( + _run_store_req( + thread, ReqMeta( req_id="r0", token_len_chunk=64, block_ids=([0, 1], list(range(8))), block_hashes=hs, can_save=True, - ) + ), ) keys = store.batch_put_from_multi_buffers.call_args.args[0] @@ -2128,8 +2146,8 @@ def test_store_sending_thread_kv_events_use_group_chunk_metadata(): thread.enable_kv_event = True hs = [bytes([i + 1]) * 4 for i in range(4)] - thread.add_stored_request("r0") - thread._handle_request( + _run_store_req( + thread, ReqMeta( req_id="r0", token_len_chunk=32, @@ -2137,7 +2155,7 @@ def test_store_sending_thread_kv_events_use_group_chunk_metadata(): block_hashes=hs, can_save=True, token_ids=list(range(32)), - ) + ), ) full_event, swa_event = thread.get_kv_events() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py index ca32c3345c38..297fc22ad8a0 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py @@ -26,6 +26,7 @@ KVConnectorBase_V1, KVConnectorMetadata, KVConnectorRole, + KVConnectorWorkerMetadata, SupportsHMA, ) from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( @@ -37,6 +38,7 @@ from vllm.forward_context import ForwardContext from vllm.logger import init_logger from vllm.v1.attention.backend import AttentionMetadata +from vllm.v1.core.block_pool import BlockPool from vllm.v1.core.kv_cache_manager import KVCacheBlocks from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.kv_cache_interface import KVCacheConfig @@ -201,6 +203,14 @@ def update_state_after_alloc( request, blocks, num_external_tokens ) + def bind_gpu_block_pool(self, gpu_block_pool: BlockPool) -> None: + assert self.connector_scheduler is not None + self.connector_scheduler.bind_gpu_block_pool(gpu_block_pool) + + def has_pending_push_work(self) -> bool: + assert self.connector_scheduler is not None + return self.connector_scheduler.has_pending_push_work() + def build_connector_meta( self, scheduler_output: SchedulerOutput, @@ -208,6 +218,10 @@ def build_connector_meta( assert self.connector_scheduler is not None return self.connector_scheduler.build_connector_meta(scheduler_output) + def build_connector_worker_meta(self) -> KVConnectorWorkerMetadata | None: + assert self.connector_worker is not None + return self.connector_worker.build_connector_worker_meta() + def request_finished( self, request: Request, @@ -220,8 +234,9 @@ def request_finished_all_groups( request: Request, block_ids: tuple[list[int], ...], ) -> tuple[bool, dict[str, Any] | None]: - assert self.connector_scheduler is not None - return self.connector_scheduler.request_finished(request, block_ids) + # An in-flight store job holds its own reference on the blocks it reads, + # so a finishing request never has to defer freeing them. + return False, None def reset_cache(self) -> bool | None: """Reset the external Mooncake store on prefix-cache reset. @@ -241,6 +256,9 @@ def reset_cache(self) -> bool | None: return None def update_connector_output(self, connector_output: KVConnectorOutput): + assert self.connector_scheduler is not None + self.connector_scheduler.update_connector_output(connector_output) + kv_cache_events = connector_output.kv_cache_events if not kv_cache_events or not isinstance( kv_cache_events, MooncakeStoreKVEvents diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py index 6daa1e82ea6d..55ffc040cfd1 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py @@ -6,7 +6,7 @@ """Data classes for MooncakeStoreConnector.""" from collections.abc import Iterable, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import cast import numpy as np @@ -14,6 +14,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorMetadata, + KVConnectorWorkerMetadata, ) from vllm.logger import init_logger from vllm.utils.math_utils import cdiv @@ -366,6 +367,10 @@ class ReqMeta: token_ids: list[int] | None = None num_prompt_tokens: int | None = None + # Identifies this store job for the engine's lifetime. A request id cannot + # serve that purpose: it is reused once a preempted request resumes, so it + # would release the wrong job's blocks. + store_job_id: int | None = None # Core-provided per-mamba-group # (group_id, cow_block_id, boundary_tokens) for this request's partial tail. # Present only on the producer's CoW step; drives the connector's offload @@ -434,6 +439,23 @@ def from_request_tracker( ) +@dataclass +class MooncakeStoreWorkerMetadata(KVConnectorWorkerMetadata): + """Maps ``ReqMeta.store_job_id`` to the number of ranks done with that job.""" + + completed_saves: dict[int, int] = field(default_factory=dict) + + def aggregate( + self, other: "KVConnectorWorkerMetadata" + ) -> "MooncakeStoreWorkerMetadata": + assert isinstance(other, MooncakeStoreWorkerMetadata) + for store_job_id, count in other.completed_saves.items(): + self.completed_saves[store_job_id] = ( + self.completed_saves.get(store_job_id, 0) + count + ) + return self + + class MooncakeStoreConnectorMetadata(KVConnectorMetadata): """Metadata passed from scheduler to worker.""" diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py index 42c6f3fa99af..f583b3fcfdf1 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py @@ -5,8 +5,6 @@ # (vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/). """Scheduler-side logic for MooncakeStoreConnector.""" -from typing import Any - from vllm.config import VllmConfig from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorMetadata, @@ -17,6 +15,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( # noqa: E501 LoadSpec, MooncakeStoreConnectorMetadata, + MooncakeStoreWorkerMetadata, ReqMeta, RequestTracker, ) @@ -24,10 +23,12 @@ LookupKeyClient, ) from vllm.logger import init_logger +from vllm.v1.core.block_pool import BlockPool from vllm.v1.core.kv_cache_manager import KVCacheBlocks from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes from vllm.v1.core.sched.output import NewRequestData, SchedulerOutput from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.outputs import KVConnectorOutput from vllm.v1.request import Request logger = init_logger(__name__) @@ -78,6 +79,15 @@ def __init__( self._unfinished_requests: dict[str, tuple[Request, tuple[list[int], ...]]] = {} self._unfinished_request_ids: set[str] = set() + self._gpu_block_pool: BlockPool | None = None + self._num_workers = vllm_config.parallel_config.world_size + self._next_store_job_id = 0 + # store_job_id -> (referenced block ids, ranks yet to report completion) + self._pinned_saves: dict[int, tuple[list[int], int]] = {} + + def bind_gpu_block_pool(self, gpu_block_pool: BlockPool) -> None: + self._gpu_block_pool = gpu_block_pool + def get_num_new_matched_tokens( self, request: Request, @@ -392,33 +402,75 @@ def build_connector_meta( ) ) + self._reference_save_blocks(meta) return meta - def request_finished( - self, - request: Request, - block_ids: tuple[list[int], ...], - ) -> tuple[bool, dict[str, Any] | None]: - """Determine whether to delay freeing blocks for async save.""" - if self.kv_role == "kv_consumer": - return False, None - tracker = self._request_trackers.get(request.request_id) - # Missing tracker can happen when the request is aborted before the - # connector observes the normal finished lifecycle or is preempted - # before finishing. - if tracker is None or ( - tracker.num_saved_tokens <= 0 and not tracker.has_pending_offload - ): - return False, None - total_blocks = sum(len(g) for g in block_ids) - delay_free_blocks = total_blocks > 0 - if delay_free_blocks: - logger.debug( - "Delaying free of %d blocks for request %s", - total_blocks, - request.request_id, + def _reference_save_blocks(self, meta: MooncakeStoreConnectorMetadata) -> None: + """Take a GPU block reference for every store job this step emits. + + The worker DMAs out of these blocks after the step that scheduled them, + so a reference keeps them out of the free queue even once the request + itself is freed, until every rank reports the job done. + """ + pool = self._gpu_block_pool + for req_meta in meta.requests: + if not req_meta.can_save: + continue + assert pool is not None, ( + "GPU block pool must be bound before any store job is emitted" + ) + req_meta.store_job_id = store_job_id = self._next_store_job_id + self._next_store_job_id += 1 + block_ids: list[int] = [] + if req_meta.partial_tail_offloads: + # A partial-tail CoW block is deliberately kept out of the + # request's block table, so it is absent from `block_ids` even + # though the worker DMAs out of it just as asynchronously. + # It leads the list, as in `pop_blocks_for_free`. + block_ids += [bid for _, bid, _ in req_meta.partial_tail_offloads] + # Every allocated block is referenced, not just the ones covering + # this job's token range: a rank resumes from its own last + # successful offset, which lags the scheduler's whenever a save was + # skipped or failed, so it may read anywhere below the range. + block_ids += [bid for group in req_meta.block_ids for bid in group] + if not block_ids: + continue + self._pinned_saves[store_job_id] = (block_ids, self._num_workers) + pool.touch([pool.blocks[bid] for bid in block_ids]) + + def update_connector_output(self, connector_output: KVConnectorOutput) -> None: + """Drop the block references of store jobs every rank has finished.""" + meta = connector_output.kv_connector_worker_meta + if not isinstance(meta, MooncakeStoreWorkerMetadata): + return + pool = self._gpu_block_pool + assert pool is not None + for store_job_id, count in meta.completed_saves.items(): + pinned = self._pinned_saves.get(store_job_id) + if pinned is None: + # The job referenced no blocks, so nothing was recorded for it. + continue + block_ids, remaining = pinned + remaining -= count + if remaining > 0: + self._pinned_saves[store_job_id] = (block_ids, remaining) + continue + assert remaining == 0, ( + f"store job {store_job_id} reported by too many ranks" ) - return delay_free_blocks, None + del self._pinned_saves[store_job_id] + # Tail-first, as elsewhere, so the shared prefix is evicted last. + pool.free_blocks(pool.blocks[bid] for bid in reversed(block_ids)) + + def has_pending_push_work(self) -> bool: + """Keep the engine stepping while any store job still holds block refs. + + Completions only reach the scheduler as worker metadata on a step, so an + engine that quiesced with jobs in flight would leave those references + held indefinitely. Nothing else keeps it alive now that a finishing + request no longer defers its own free. + """ + return bool(self._pinned_saves) def reset_store(self) -> bool: """Trigger a global ``remove_all(force=True)`` on the Mooncake master. diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index 762e5eba9263..e25aec3fbde9 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -17,7 +17,6 @@ import socket import threading import time -from collections import defaultdict from collections.abc import Callable, Sequence from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import dataclass @@ -46,6 +45,7 @@ ChunkedTokenDatabase, KeyMetadata, MooncakeStoreConnectorMetadata, + MooncakeStoreWorkerMetadata, PoolKey, ReqMeta, ) @@ -506,7 +506,16 @@ def __init__( self.group_put_steps = group_put_steps self.coord = coord self.kv_role = kv_role - self.stored_requests: defaultdict[str, int] = defaultdict(int) + # req_id -> ids of its store jobs that are still queued or running. + # Keying by store_job_id, which never repeats for the engine's lifetime, + # rather than counting jobs per request id makes the ledger immune to id + # reuse across preemption: a job left over from a retired generation is + # missing from the set its resumed generation builds, so it can no longer + # retire that generation, rewind its resume offset, or mark it skipped. + self.stored_requests: dict[str, set[int]] = {} + # store_job_id -> times this rank finished with it, drained every step + # so the scheduler can release the blocks it referenced for those jobs. + self._completed_saves: dict[int, int] = {} self.enable_kv_event = enable_kv_event # Caller always passes a non-None ReplicateConfig — see # MooncakeStoreWorker.__init__ where store_replicate_config is built. @@ -522,14 +531,20 @@ def __init__( # batch resumes here, so pressure-skipped or failed ranges are retried. self._saved_offset: dict[str, int] = {} - def add_stored_request(self, req_id: str): + def add_request(self, request: ReqMeta) -> None: + # Register before enqueueing so a job is never picked up unledgered. + assert request.store_job_id is not None with self.done_task_lock: - self.stored_requests[req_id] += 1 + self.stored_requests.setdefault(request.req_id, set()).add( + request.store_job_id + ) + super().add_request(request) - def dec_stored_request(self, req_id: str): + def is_live_store_job(self, req_meta: ReqMeta) -> bool: with self.done_task_lock: - if req_id in self.stored_requests: - self.stored_requests[req_id] -= 1 + return req_meta.store_job_id in self.stored_requests.get( + req_meta.req_id, () + ) def delete_finished_stored_request(self, req_id: str): with self.done_task_lock: @@ -538,21 +553,51 @@ def delete_finished_stored_request(self, req_id: str): self._skip_store_requests.discard(req_id) self._saved_offset.pop(req_id, None) - def _record_saved(self, req_id: str, token_len: int) -> None: - # Guard on liveness so a concurrent finish/preempt pop isn't recreated. + def finish_store_job(self, req_meta: ReqMeta) -> None: + """Retire a job from the ledger and report its blocks as no longer read. + + Every path out of a job must reach this, skips and failures included: a + job that never reports leaves its blocks referenced for the rest of the + run. The discard is a no-op for a job whose generation already retired. + """ + store_job_id = req_meta.store_job_id + assert store_job_id is not None, ( + "a queued store job always carries a store_job_id" + ) with self.done_task_lock: - if req_id in self.stored_requests: - self._saved_offset[req_id] = token_len + live = self.stored_requests.get(req_meta.req_id) + if live is not None: + live.discard(store_job_id) + self._completed_saves[store_job_id] = ( + self._completed_saves.get(store_job_id, 0) + 1 + ) + + def take_completed_saves(self) -> dict[int, int]: + with self.done_task_lock: + completed = self._completed_saves + self._completed_saves = {} + return completed + + def _record_saved(self, req_meta: ReqMeta, token_len: int) -> None: + # Guard on job liveness so neither a concurrent finish/preempt pop nor a + # stale job's offset is written back over the live generation's. + with self.done_task_lock: + if req_meta.store_job_id in self.stored_requests.get(req_meta.req_id, ()): + self._saved_offset[req_meta.req_id] = token_len def _should_skip_request(self, req_id: str) -> bool: with self.done_task_lock: return self._store_pressure_active and req_id in self._skip_store_requests - def _mark_request_skipped_for_pressure(self, req_id: str) -> bool: + def _mark_request_skipped_for_pressure(self, req_meta: ReqMeta) -> bool: + req_id = req_meta.req_id with self.done_task_lock: already_skipped = req_id in self._skip_store_requests self._store_pressure_active = True - self._skip_store_requests.add(req_id) + # The pressure itself is global, but only a live job may sentence its + # own request to being skipped. + if req_meta.store_job_id in self.stored_requests.get(req_id, ()): + self._skip_store_requests.add(req_id) return already_skipped def _clear_store_pressure(self) -> bool: @@ -727,7 +772,7 @@ def _maybe_offload_partial_tail(self, req_meta: ReqMeta) -> bool: failed_codes, ) if MOONCAKE_NO_AVAILABLE_HANDLE in failed_codes: - self._mark_request_skipped_for_pressure(req_meta.req_id) + self._mark_request_skipped_for_pressure(req_meta) return False if self._clear_store_pressure(): @@ -738,22 +783,20 @@ def _maybe_offload_partial_tail(self, req_meta: ReqMeta) -> bool: return True def _handle_request(self, req_meta: ReqMeta): - # Cache hits are always a multiple of ``lcm_block_size`` tokens, which - # is also ``store_mask``'s precondition. - lcm_block_size = self.coord.lcm_block_size - token_len = req_meta.token_len_chunk // lcm_block_size * lcm_block_size - block_ids_per_group = req_meta.block_ids - req_id = req_meta.req_id - current_event = req_meta.current_event - - if req_id not in self.stored_requests: - self.request_queue.task_done() - return - - # Decrement the in-flight counter and signal task_done() in `finally` - # so the scheduler can release the GPU blocks it pinned for this - # request (via `delay_free_blocks`) even when the store path raises. + # The single `finally` is the only way out, so the scheduler releases + # this job's GPU block references however the job ends. try: + # Cache hits are always a multiple of ``lcm_block_size`` tokens, + # which is also ``store_mask``'s precondition. + lcm_block_size = self.coord.lcm_block_size + token_len = req_meta.token_len_chunk // lcm_block_size * lcm_block_size + block_ids_per_group = req_meta.block_ids + req_id = req_meta.req_id + current_event = req_meta.current_event + + if not self.is_live_store_job(req_meta): + return + if self._should_skip_request(req_id): logger.debug( "Skipping Mooncake store for request %s while CPU/disk " @@ -808,7 +851,7 @@ def _handle_request(self, req_meta: ReqMeta): group_indices.append(g_idx) if not keys: - self._record_saved(req_id, token_len) + self._record_saved(req_meta, token_len) return # Check which blocks already exist (dedup) @@ -834,7 +877,7 @@ def _handle_request(self, req_meta: ReqMeta): ] if not missing_indices: - self._record_saved(req_id, token_len) + self._record_saved(req_meta, token_len) return if len(missing_indices) != len(keys): @@ -954,7 +997,7 @@ def _handle_request(self, req_meta: ReqMeta): ) if ( MOONCAKE_NO_AVAILABLE_HANDLE in failed_codes - and not self._mark_request_skipped_for_pressure(req_id) + and not self._mark_request_skipped_for_pressure(req_meta) ): logger.warning( "Detected Mooncake CPU/disk offloading pressure " @@ -964,7 +1007,7 @@ def _handle_request(self, req_meta: ReqMeta): req_id, ) else: - self._record_saved(req_id, token_len) + self._record_saved(req_meta, token_len) if self._clear_store_pressure(): logger.info( "Mooncake CPU/disk offloading pressure cleared " @@ -984,7 +1027,7 @@ def _handle_request(self, req_meta: ReqMeta): if self.enable_kv_event and stored_events: self.update_kv_event(stored_events) finally: - self.dec_stored_request(req_id) + self.finish_store_job(req_meta) self.request_queue.task_done() @@ -1670,16 +1713,15 @@ def get_finished( continue request.current_event = current_event assert self.kv_send_thread is not None - self.kv_send_thread.add_stored_request(request.req_id) self.kv_send_thread.add_request(request) - # Check completion of previously queued transfers - done_sending = ( - self._get_and_clear_finished_sending(finished_req_ids, meta) - if self.kv_role in ["kv_producer", "kv_both"] - else set() - ) + if self.kv_role in ["kv_producer", "kv_both"]: + self._close_ended_store_requests(finished_req_ids, meta) + # Blocks read by a store job are released by the scheduler when the job + # reports back (see build_connector_worker_meta), so no request ever waits + # on a `finished_sending` signal to get its blocks back. + done_sending: set[str] = set() done_recving: set[str] = set() if self.load_async: for recv_thread in self.kv_recv_threads: @@ -1727,35 +1769,37 @@ def get_kv_connector_stats(self) -> MooncakeStoreConnectorStats | None: self.kv_connector_stats = MooncakeStoreConnectorStats() return kv_connector_stats - def _get_and_clear_finished_sending( + def _close_ended_store_requests( self, finished_req_ids: set[str], meta: MooncakeStoreConnectorMetadata, - ) -> set[str]: + ) -> None: + """Retire the ledger entries of requests that finished or were preempted. + + An entry may only go once its jobs have drained, because they still read + the resume offset it owns; a request that comes back after preemption + then saves from the start rather than from where the last attempt got to. + """ assert self.kv_send_thread is not None - finished_sending: set[str] = set() for req_id in meta.preempted_req_ids: self.kv_send_thread.delete_finished_stored_request(req_id) - for req_id in self.kv_send_thread.stored_requests.copy(): - if ( - self.kv_send_thread.stored_requests[req_id] == 0 - and req_id in self.finished_store_req - ): - self.finished_store_req.remove(req_id) - finished_sending.add(req_id) - self.kv_send_thread.delete_finished_stored_request(req_id) - - for req_id in finished_req_ids: - req_remain_jobs = self.kv_send_thread.stored_requests.get(req_id) - if req_remain_jobs == 0: - finished_sending.add(req_id) - self.kv_send_thread.delete_finished_stored_request(req_id) - elif req_remain_jobs is not None: + for req_id in finished_req_ids | self.finished_store_req: + if self.kv_send_thread.stored_requests.get(req_id): + # Queued jobs still need the resume offset; retire on a later step. self.finished_store_req.add(req_id) + else: + self.finished_store_req.discard(req_id) + self.kv_send_thread.delete_finished_stored_request(req_id) - return finished_sending + def build_connector_worker_meta(self) -> MooncakeStoreWorkerMetadata | None: + if self.kv_send_thread is None: + return None + completed_saves = self.kv_send_thread.take_completed_saves() + if not completed_saves: + return None + return MooncakeStoreWorkerMetadata(completed_saves=completed_saves) def lookup(self, num_tokens: int, block_hashes: Sequence[BlockHash]) -> int: """Check how many prefix tokens exist in the store. From a18c9b56ff9feb617d8f7c2eeca9263b328cc9c1 Mon Sep 17 00:00:00 2001 From: kiroxu <148877251+BabyDrangoner@users.noreply.github.com> Date: Mon, 17 Aug 2026 09:04:38 +0800 Subject: [PATCH 025/839] [Kimi-K3][Perf] Update FlashKDA for automatic K2 V-split (#52458) --- cmake/external_projects/flashkda.cmake | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cmake/external_projects/flashkda.cmake b/cmake/external_projects/flashkda.cmake index 1d3d163c61bf..988a27d7c755 100644 --- a/cmake/external_projects/flashkda.cmake +++ b/cmake/external_projects/flashkda.cmake @@ -13,7 +13,7 @@ else() FetchContent_Declare( flashkda GIT_REPOSITORY https://github.com/vllm-project/FlashKDA.git - GIT_TAG b5d11010ff01c1d4a683c0dde42e76cbeaa8107f + GIT_TAG 053de1b716ef3255873e02d2d28f4adf09951978 GIT_PROGRESS TRUE GIT_SUBMODULES cutlass ) From 0ad04cff1b10267a2642becbf31f62c91018cad5 Mon Sep 17 00:00:00 2001 From: Shanshan Shen <467638484@qq.com> Date: Mon, 17 Aug 2026 09:52:02 +0800 Subject: [PATCH 026/839] [ROCm][CI] Enable ViT CUDA graph tests on AMD gfx950 GPUs (#52256) Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com> Co-authored-by: Andreas Karatzas --- .buildkite/test-amd.yaml | 7 +- docs/design/cuda_graphs_multimodal.md | 78 +++++++++++++------ .../generation/test_vit_cudagraph.py | 8 +- tests/v1/cudagraph/test_encoder_cudagraph.py | 8 +- 4 files changed, 71 insertions(+), 30 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index df65ee762dba..c1658e7a8331 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -2092,7 +2092,7 @@ steps: - tests/models/multimodal/generation - tests/models/multimodal/test_mapping.py commands: - - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing + - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/generation/test_vit_cudagraph.py --ignore models/multimodal/processing - pytest -v -s models/multimodal/generation/test_memory_leak.py -m core_model - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model @@ -3671,7 +3671,8 @@ steps: - vllm/ - tests/models/multimodal/generation commands: - - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing + - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/generation/test_vit_cudagraph.py --ignore models/multimodal/processing + - pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model - pytest -v -s models/multimodal/generation/test_memory_leak.py -m core_model - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model @@ -3771,6 +3772,7 @@ steps: - tests/v1/executor - tests/v1/kv_offload - tests/v1/worker + - tests/v1/cudagraph - tests/v1/kv_connector/unit - tests/v1/metrics - tests/entrypoints/openai/correctness/test_lmeval.py @@ -3780,6 +3782,7 @@ steps: - pytest -v -s v1/executor - pytest -v -s v1/kv_offload - pytest -v -s v1/worker + - pytest -v -s v1/cudagraph/test_encoder_cudagraph.py - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit - pytest -v -s -m 'not cpu_test' v1/metrics - pip install -U git+https://github.com/vllm-project/lm-evaluation-harness.git@streaming-api diff --git a/docs/design/cuda_graphs_multimodal.md b/docs/design/cuda_graphs_multimodal.md index f171c42e808b..0fbb6801e44e 100644 --- a/docs/design/cuda_graphs_multimodal.md +++ b/docs/design/cuda_graphs_multimodal.md @@ -7,6 +7,59 @@ For two-tower vision encoders (e.g., DeepSeek-OCR's SAM + CLIP with dynamic tili !!! note Encoder CUDA Graphs are orthogonal to decoder CUDA Graphs — both can be enabled simultaneously. Encoder graphs capture the vision encoder execution (e.g., ViT in Qwen3-VL), while decoder graphs capture the language model execution as described in the [CUDA Graphs design document](cuda_graphs.md). +## Compatibility Matrix + +!!! note + The symbols used below have the following meanings: + + - ✅ = Full compatibility + - 🟠 = Partial compatibility + - ❌ = No compatibility + - ❔ = Unknown or TBD + +### Model x Feature + +| Architecture | Models | CG for Image | CG for Video | Multi-Path Graph | +| ------------ | ------ | ------------ | ------------ | --------------- | +| `DeepseekOCRForCausalLM` | `DeepSeek-OCR` | ✅︎ | ❌︎ | ✅︎ | +| `Ernie4_5_VLMoeForConditionalGeneration` | `ERNIE-4.5-VL` | ✅︎ | ❌︎ | ❌︎ | +| `Gemma3ForConditionalGeneration` | `Gemma3` | ✅︎ | ❌︎ | ❌︎ | +| `Glm4vForConditionalGeneration` | `GLM-4.1V, GLM-4.6V-Flash` | ✅︎ | ✅︎ | ❌︎ | +| `Gemma4ForConditionalGeneration` | `Gemma-4` | ✅︎ | ✅︎ | ❌︎ | +| `InternVLChatModel` | `InternVL3.5`, `InternVL3`, `InternVL2.5`, `InternVL2` | ✅︎ | ✅︎ | ❌︎ | +| `KimiVLForConditionalGeneration` | `Kimi-VL` | ✅︎ | ❌︎ | ❌︎ | +| `Llama4ForConditionalGeneration` | `Llama 4` | ✅︎ | ❌︎ | ❌︎ | +| `Qwen2VLForConditionalGeneration` | `Qwen2-VL` | ✅︎ | ✅︎ | ❌︎ | +| `Qwen2_5_VLForConditionalGeneration` | `Qwen2.5-VL` | ✅︎ | ✅︎ | ❌︎ | +| `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ | ❌︎ | +| `Qwen3_5ForConditionalGeneration` | `Qwen3.5`, `Qwen3.6` | ✅︎ | ✅︎ | ❌︎ | +| `Qwen3_5MoeForConditionalGeneration` | `Qwen3.5-MoE`, `Qwen3.6-MoE` | ✅︎ | ✅︎ | ❌︎ | +| `Step3VLForConditionalGeneration` | `Step3-VL` | ✅︎ | ❌︎ | ✅︎ | + +### Model x Hardware + +| Architecture | NV Blackwell | NV Ampere | AMD MI300X | AMD MI350X / MI355X | +| ------------ | ---------------- | ------------- | -------------- | --------------------- | +| `DeepseekOCRForCausalLM` | ✅︎ | ✅︎ | ❔ | ✅︎ | +| `Ernie4_5_VLMoeForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | +| `Gemma3ForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | +| `Glm4vForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | +| `Gemma4ForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | +| `InternVLChatModel` | ✅︎ | ✅︎ | ❔ | ✅︎ | +| `KimiVLForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | +| `Llama4ForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | +| `Qwen2VLForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | +| `Qwen2_5_VLForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | +| `Qwen3VLForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | +| `Qwen3_5ForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | +| `Qwen3_5MoeForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | +| `Step3VLForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | + +!!! note + Encoder CUDA Graph has currently been tested with `--mm-encoder-attn-backend=FLASH_ATTN` and `--mm-encoder-attn-backend=FLASHINFER` on Blackwell GPUs. + For Qwen2-VL and Qwen2.5-VL only FA2 and FA3 has been tested. + Encoder CUDA Graph has also been tested with AMD MI350X (gfx950) used `--mm-encoder-attn-backend=FLASH_ATTN` (the ROCm default). + ## Motivation Vision encoder inference incurs CUDA kernel launch overhead on the host side. The overhead is more significant when the batch size is small or image size is small. @@ -113,29 +166,6 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra !!! note The `SupportsEncoderCudaGraph` protocol is designed to be model-agnostic. New vision encoder models can opt-in by implementing the protocol methods without modifying the manager. -**Supported models:** - -| Architecture | Models | CG for Image | CG for Video | Dual-Path Graph | -| ------------ | ------ | ------------ | ------------ | --------------- | -| `DeepseekOCRForCausalLM` | `DeepSeek-OCR` | ✅︎ | ❌︎ | ✅︎ | -| `Ernie4_5_VLMoeForConditionalGeneration` | `ERNIE-4.5-VL` | ✅︎ | ❌︎ | ❌︎ | -| `Gemma3ForConditionalGeneration` | `Gemma3` | ✅︎ | ❌︎ | ❌︎ | -| `Glm4vForConditionalGeneration` | `GLM-4.1V, GLM-4.6V-Flash` | ✅︎ | ✅︎ | ❌︎ | -| `Gemma4ForConditionalGeneration` | `Gemma-4` | ✅︎ | ✅︎ | ❌︎ | -| `InternVLChatModel` | `InternVL3.5`, `InternVL3`, `InternVL2.5`, `InternVL2` | ✅︎ | ✅︎ | ❌︎ | -| `KimiVLForConditionalGeneration` | `Kimi-VL` | ✅︎ | ❌︎ | ❌︎ | -| `Llama4ForConditionalGeneration` | `Llama 4` | ✅︎ | ❌︎ | ❌︎ | -| `Qwen2VLForConditionalGeneration` | `Qwen2-VL` | ✅︎ | ✅︎ | ❌︎ | -| `Qwen2_5_VLForConditionalGeneration` | `Qwen2.5-VL` | ✅︎ | ✅︎ | ❌︎ | -| `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ | ❌︎ | -| `Qwen3_5ForConditionalGeneration` | `Qwen3.5`, `Qwen3.6` | ✅︎ | ✅︎ | ❌︎ | -| `Qwen3_5MoeForConditionalGeneration` | `Qwen3.5-MoE`, `Qwen3.6-MoE` | ✅︎ | ✅︎ | ❌︎ | -| `Step3VLForConditionalGeneration` | `Step3-VL` | ✅︎ | ❌︎ | ✅︎ | - -!!! note - Encoder CUDA Graphs have currently been tested with `--mm-encoder-attn-backend=FLASH_ATTN` and `--mm-encoder-attn-backend=FLASHINFER` on Blackwell GPUs. - For Qwen2-VL and Qwen2.5-VL only FA2 and FA3 has been tested. - ## Configuration Four fields in `CompilationConfig` control encoder CUDA Graphs: @@ -228,7 +258,7 @@ model = vllm.LLM( ) ``` -## About the Performance +## Benchmark Results The following benchmarks were run on Blackwell GPUs (GB200) using `vllm bench mm-processor`. See [#35963](https://github.com/vllm-project/vllm/pull/35963) for full details. diff --git a/tests/models/multimodal/generation/test_vit_cudagraph.py b/tests/models/multimodal/generation/test_vit_cudagraph.py index 1972064b1d93..8e12790925dd 100644 --- a/tests/models/multimodal/generation/test_vit_cudagraph.py +++ b/tests/models/multimodal/generation/test_vit_cudagraph.py @@ -313,7 +313,9 @@ def get_compilation_config(config: VitCudagraphTestConfig): @pytest.mark.parametrize("model_id", params_with_marks(MODEL_CONFIGS)) -@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="Skip if not cuda or rocm" +) def test_vit_cudagraph_image(model_id, vllm_runner, image_assets): config = MODEL_CONFIGS[model_id] @@ -357,7 +359,9 @@ def test_vit_cudagraph_image(model_id, vllm_runner, image_assets): @pytest.mark.parametrize("model_id", params_with_marks(MODEL_CONFIGS)) -@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="Skip if not cuda or rocm" +) def test_vit_cudagraph_video(model_id, vllm_runner, video_assets): config = MODEL_CONFIGS[model_id] diff --git a/tests/v1/cudagraph/test_encoder_cudagraph.py b/tests/v1/cudagraph/test_encoder_cudagraph.py index f2a54c55c809..a01214a40d29 100644 --- a/tests/v1/cudagraph/test_encoder_cudagraph.py +++ b/tests/v1/cudagraph/test_encoder_cudagraph.py @@ -475,7 +475,9 @@ def _make_video_mm_kwargs( # --------------------------------------------------------------------------- -@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda") +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="Skip if not cuda or rocm" +) class TestEncoderCudaGraphCaptureReplay: def setup_method(self): self.device = torch.device("cuda:0") @@ -758,7 +760,9 @@ def test_video_model_returns_video_for_video_kwargs(self): _VIDEO_MAX_FRAMES = 8 # 2 frames per item at max_batch_size=4 -@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda") +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="Skip if not cuda or rocm" +) class TestEncoderCudaGraphVideoReplay: def setup_method(self): self.device = torch.device("cuda:0") From 967e104fad12758e1fc4f16b7f9fdebfc3b6b1bb Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Sun, 16 Aug 2026 19:20:04 -0700 Subject: [PATCH 027/839] [Config] Unify indexer cache dtype under attention_config.indexer_kv_dtype (#52550) Signed-off-by: Yongye Zhu Co-authored-by: Claude Opus 5 (1M context) --- ...epSeek-V4-Flash-DSpark-confidence-TP4.yaml | 2 +- .../DeepSeek-V4-Flash-deepgemm-mega-moe.yaml | 2 +- vllm/config/attention.py | 41 +++++++++++++++---- vllm/models/deepseek_v4/attention.py | 3 +- vllm/models/minimax_m3/nvidia/model.py | 4 +- vllm/v1/attention/backends/mla/indexer.py | 37 +++++++++++------ 6 files changed, 65 insertions(+), 24 deletions(-) diff --git a/tests/evals/gsm8k/configs/DeepSeek-V4-Flash-DSpark-confidence-TP4.yaml b/tests/evals/gsm8k/configs/DeepSeek-V4-Flash-DSpark-confidence-TP4.yaml index 61fbaf33e29a..84c7fb30196a 100644 --- a/tests/evals/gsm8k/configs/DeepSeek-V4-Flash-DSpark-confidence-TP4.yaml +++ b/tests/evals/gsm8k/configs/DeepSeek-V4-Flash-DSpark-confidence-TP4.yaml @@ -15,7 +15,7 @@ server_args: >- --block-size 256 --gpu-memory-utilization 0.5 --kv-cache-dtype fp8 - --attention_config.use_fp4_indexer_cache=True + --attention_config.indexer_kv_dtype=mxfp4 --max-num-batched-tokens 16384 --max-num-seqs 128 --speculative-config '{"method":"dspark", diff --git a/tests/evals/gsm8k/configs/moe-refactor/DeepSeek-V4-Flash-deepgemm-mega-moe.yaml b/tests/evals/gsm8k/configs/moe-refactor/DeepSeek-V4-Flash-deepgemm-mega-moe.yaml index 742d9e40b8a0..955f9d3e90ec 100644 --- a/tests/evals/gsm8k/configs/moe-refactor/DeepSeek-V4-Flash-deepgemm-mega-moe.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor/DeepSeek-V4-Flash-deepgemm-mega-moe.yaml @@ -2,4 +2,4 @@ model_name: "deepseek-ai/DeepSeek-V4-Flash" accuracy_threshold: 0.95 num_questions: 1319 num_fewshot: 5 -server_args: "--trust-remote-code --kv-cache-dtype fp8 --block-size 256 --enable-expert-parallel --tensor-parallel-size 2 --attention_config.use_fp4_indexer_cache=True --moe-backend deep_gemm_mega_moe --tokenizer-mode deepseek_v4 --tool-call-parser deepseek_v4 --enable-auto-tool-choice --reasoning-parser deepseek_v4 --speculative_config.method=mtp --speculative_config.num_speculative_tokens=2" +server_args: "--trust-remote-code --kv-cache-dtype fp8 --block-size 256 --enable-expert-parallel --tensor-parallel-size 2 --attention_config.indexer_kv_dtype=mxfp4 --moe-backend deep_gemm_mega_moe --tokenizer-mode deepseek_v4 --tool-call-parser deepseek_v4 --enable-auto-tool-choice --reasoning-parser deepseek_v4 --speculative_config.method=mtp --speculative_config.num_speculative_tokens=2" diff --git a/vllm/config/attention.py b/vllm/config/attention.py index 994be05f54ea..907917d091cb 100644 --- a/vllm/config/attention.py +++ b/vllm/config/attention.py @@ -7,10 +7,13 @@ from pydantic import field_validator from vllm.config.utils import config +from vllm.logger import init_logger from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum from vllm.v1.attention.backends.registry import AttentionBackendEnum -IndexerKVDType = Literal["bf16", "fp8", "mxfp4", "nvfp4"] +logger = init_logger(__name__) + +IndexerKVDType = Literal["auto", "bf16", "fp8", "mxfp4", "nvfp4"] MiniMaxM3MSADecodeBackend = Literal["triton", "cutlass"] @@ -65,12 +68,15 @@ class AttentionConfig: use_prefill_query_quantization: bool = False """If set, quantize query for attention in prefill.""" - use_fp4_indexer_cache: bool = False - """If set, use fp4 indexer cache for dsv32 family model (not support yet)""" + use_fp4_indexer_cache: bool | None = None + """Deprecated alias for `indexer_kv_dtype`; use that instead. True maps to + `mxfp4`, False is a no-op (it selected the model default already).""" - indexer_kv_dtype: IndexerKVDType = "bf16" - """Data type for the sparse-attention indexer K cache. Quantized formats - (fp8, mxfp4, nvfp4) require indexer kernel support in the backend.""" + indexer_kv_dtype: IndexerKVDType = "auto" + """Data type for the sparse-attention indexer K cache. "auto" picks the + model's default (bf16 for MiniMax M3, fp8 for the DeepSeek sparse + indexer). Quantized formats (fp8, mxfp4, nvfp4) require indexer kernel + support in the backend.""" use_non_causal: bool = False """Whether to use non-causal (bidirectional) attention.""" @@ -114,6 +120,26 @@ def __post_init__(self) -> None: # layers still use the platform's normal automatic backend. self.backend = None + if self.use_fp4_indexer_cache is not None: + logger.warning( + "use_fp4_indexer_cache is deprecated and will be removed in " + "v0.19. Use indexer_kv_dtype instead (True -> 'mxfp4')." + ) + if self.use_fp4_indexer_cache: + if self.indexer_kv_dtype not in ("auto", "mxfp4"): + raise ValueError( + "use_fp4_indexer_cache=True conflicts with " + f"indexer_kv_dtype={self.indexer_kv_dtype!r}. Set only " + "indexer_kv_dtype." + ) + self.indexer_kv_dtype = "mxfp4" + + def resolve_indexer_kv_dtype(self, default: IndexerKVDType) -> IndexerKVDType: + """Resolve `indexer_kv_dtype`, substituting `default` for "auto".""" + if self.indexer_kv_dtype == "auto": + return default + return self.indexer_kv_dtype + def compute_hash(self) -> str: """ Provide a hash that uniquely identifies all the configs @@ -124,7 +150,8 @@ def compute_hash(self) -> str: """ from vllm.config.utils import get_hash_factors, hash_factors - ignored_factors: set[str] = set() + # Folded into indexer_kv_dtype by __post_init__. + ignored_factors: set[str] = {"use_fp4_indexer_cache"} factors = get_hash_factors(self, ignored_factors) return hash_factors(factors) diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index e2dc1bb45a7a..ff7b92c6b22b 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -56,6 +56,7 @@ from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata from vllm.v1.attention.backends.mla.indexer import ( DeepseekV4IndexerBackend, + dsa_indexer_uses_fp4, get_max_prefill_buffer_size, ) from vllm.v1.attention.backends.mla.sparse_swa import DeepseekV4SWACache @@ -805,7 +806,7 @@ def __init__( self.q_lora_rank = q_lora_rank # 1536 self.compress_ratio = compress_ratio self.eager_scratch_pool = eager_scratch_pool - self.use_fp4_kv = self.vllm_config.attention_config.use_fp4_indexer_cache + self.use_fp4_kv = dsa_indexer_uses_fp4(vllm_config) logger.info_once( "Using %s indexer cache for Lightning Indexer.", "MXFP4" if self.use_fp4_kv else "FP8", diff --git a/vllm/models/minimax_m3/nvidia/model.py b/vllm/models/minimax_m3/nvidia/model.py index 68f9805b1fdc..a75eca2b630e 100644 --- a/vllm/models/minimax_m3/nvidia/model.py +++ b/vllm/models/minimax_m3/nvidia/model.py @@ -498,7 +498,9 @@ def __init__( set_default_quant_scales(self, register_buffer=True) # Indexer side-cache dtype, mirroring --kv-cache-dtype for the main # cache (--attention-config '{"indexer_kv_dtype": ...}'). - self.indexer_kv_dtype = vllm_config.attention_config.indexer_kv_dtype + self.indexer_kv_dtype = vllm_config.attention_config.resolve_indexer_kv_dtype( + "bf16" + ) # Shared top-k buffer: the indexer writes the selected blocks into it and # the attend impl reads them back (so nothing crosses the eager break as a diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index baecb391fd6a..1177d8b27afe 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -40,6 +40,28 @@ logger = init_logger(__name__) +# The DSA indexer K cache is always quantized; "auto" means fp8 (V3.2 layout) +# and mxfp4 is the opt-in Blackwell path. +DSA_INDEXER_KV_DTYPES = ("fp8", "mxfp4") + + +def dsa_indexer_uses_fp4(vllm_config: VllmConfig) -> bool: + """Whether the DeepSeek sparse indexer should use the MXFP4 K cache.""" + kv_dtype = vllm_config.attention_config.resolve_indexer_kv_dtype("fp8") + if kv_dtype not in DSA_INDEXER_KV_DTYPES: + raise ValueError( + f"indexer_kv_dtype={kv_dtype!r} is not supported by the DeepSeek " + f"sparse indexer (expected one of {DSA_INDEXER_KV_DTYPES})." + ) + use_fp4 = kv_dtype == "mxfp4" + if use_fp4 and not current_platform.is_device_capability_family(100): + raise ValueError( + "indexer_kv_dtype='mxfp4' requires Blackwell datacenter GPUs " + "(sm_10x, e.g. B200/GB200); sm_120 (consumer Blackwell) and " + "earlier architectures are not supported." + ) + return use_fp4 + @triton.jit def _prepare_uniform_decode_kernel( @@ -526,18 +548,7 @@ def __init__(self, *args, block_table_width: int, **kwargs) -> None: if self.vllm_config.speculative_config else 0 ) - self.use_fp4_indexer_cache = ( - self.vllm_config.attention_config.use_fp4_indexer_cache - ) - - assert ( - current_platform.is_device_capability_family(100) - or not self.use_fp4_indexer_cache - ), ( - "use_fp4_indexer_cache requires Blackwell datacenter GPUs " - "(sm_10x, e.g. B200/GB200); sm_120 (consumer Blackwell) and " - "earlier architectures are not supported." - ) + self.use_fp4_indexer_cache = dsa_indexer_uses_fp4(self.vllm_config) next_n = self.num_speculative_tokens + 1 self.decode_threshold = next_n @@ -546,7 +557,7 @@ def __init__(self, *args, block_table_width: int, **kwargs) -> None: self.supports_varlen = _supports_varlen_paged_mqa_logits() logger.info_once( "DSA indexer decode path: use_flattening=%s supports_varlen=%s " - "(next_n=%d, use_fp4_indexer_cache=%s)", + "(next_n=%d, use_fp4_cache=%s)", self.use_flattening, self.supports_varlen, next_n, From a6a2a93f9b9f3b2c155b50c8c6f2badf41cbfd66 Mon Sep 17 00:00:00 2001 From: Kaif Kohari Date: Mon, 17 Aug 2026 03:27:59 +0100 Subject: [PATCH 028/839] [Bugfix][Frontend] Guard remaining before-validators against non-object JSON bodies (#52528) Signed-off-by: Kaif --- .../test_non_object_body_validation.py | 71 +++++++++++++++++++ .../entrypoints/openai/completion/protocol.py | 12 ++++ vllm/entrypoints/openai/responses/protocol.py | 6 ++ vllm/entrypoints/pooling/base/protocol.py | 2 + vllm/entrypoints/serve/tokenize/protocol.py | 2 + .../speech_to_text/transcription/protocol.py | 2 + .../speech_to_text/translation/protocol.py | 2 + 7 files changed, 97 insertions(+) create mode 100644 tests/entrypoints/unit_tests/test_non_object_body_validation.py diff --git a/tests/entrypoints/unit_tests/test_non_object_body_validation.py b/tests/entrypoints/unit_tests/test_non_object_body_validation.py new file mode 100644 index 000000000000..a40b73e5433c --- /dev/null +++ b/tests/entrypoints/unit_tests/test_non_object_body_validation.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Non-object JSON bodies must fail validation cleanly (4xx), not AttributeError (500). + +mode=before validators that call data.get(...) without an isinstance(data, dict) +guard raise AttributeError for string/list/scalar bodies and surface as HTTP 500. + +This extends the chat completion coverage added in #51654 to the remaining +request models whose before-validators were missing the same guard. +""" + +import pytest +from pydantic import ValidationError + +from vllm.entrypoints.openai.completion.protocol import CompletionRequest +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.entrypoints.pooling.classify.protocol import ClassificationChatRequest +from vllm.entrypoints.pooling.embed.protocol import EmbeddingChatRequest +from vllm.entrypoints.pooling.pooling.protocol import PoolingChatRequest +from vllm.entrypoints.serve.tokenize.protocol import TokenizeChatRequest +from vllm.entrypoints.speech_to_text.transcription.protocol import TranscriptionRequest +from vllm.entrypoints.speech_to_text.translation.protocol import TranslationRequest +from vllm.exceptions import VLLMValidationError + +pytestmark = pytest.mark.skip_global_cleanup + +REQUEST_MODELS = [ + CompletionRequest, + ResponsesRequest, + EmbeddingChatRequest, + ClassificationChatRequest, + PoolingChatRequest, + TokenizeChatRequest, + TranscriptionRequest, + TranslationRequest, +] + + +@pytest.mark.parametrize("request_model", REQUEST_MODELS, ids=lambda m: m.__name__) +@pytest.mark.parametrize( + "payload", + [ + "this is not valid json{{{", + ["not", "an", "object"], + 42, + None, + True, + ], +) +def test_request_models_reject_non_object_body(request_model, payload): + with pytest.raises(ValidationError): + request_model.model_validate(payload) + + +def test_completion_request_still_validates_dict_bodies(): + """The guard must not swallow real field-level errors on object bodies.""" + with pytest.raises(VLLMValidationError, match="prompt"): + CompletionRequest.model_validate({"model": "qwen", "prompt": ""}) + + +def test_tokenize_chat_request_still_validates_dict_bodies(): + with pytest.raises(VLLMValidationError, match="add_generation_prompt"): + TokenizeChatRequest.model_validate( + { + "model": "qwen", + "messages": [{"role": "user", "content": "hello"}], + "continue_final_message": True, + "add_generation_prompt": True, + } + ) diff --git a/vllm/entrypoints/openai/completion/protocol.py b/vllm/entrypoints/openai/completion/protocol.py index e4feece2e56f..0cb830294687 100644 --- a/vllm/entrypoints/openai/completion/protocol.py +++ b/vllm/entrypoints/openai/completion/protocol.py @@ -413,6 +413,8 @@ def normalize_null_max_tokens(cls, data): @model_validator(mode="before") @classmethod def validate_response_format(cls, data): + if not isinstance(data, dict): + return data response_format = data.get("response_format") if response_format is None: return data @@ -444,6 +446,8 @@ def validate_response_format(cls, data): @model_validator(mode="before") @classmethod def check_structured_outputs_count(cls, data): + if not isinstance(data, dict): + return data if data.get("structured_outputs", None) is None: return data @@ -472,6 +476,8 @@ def check_structured_outputs_count(cls, data): @model_validator(mode="before") @classmethod def check_logprobs(cls, data): + if not isinstance(data, dict): + return data if data.get("logprob_token_ids") and data.get("use_beam_search"): raise VLLMValidationError( "`logprob_token_ids` is not supported with beam search.", @@ -532,6 +538,8 @@ def check_logprobs(cls, data): @model_validator(mode="before") @classmethod def validate_stream_options(cls, data): + if not isinstance(data, dict): + return data if data.get("stream_options") and not data.get("stream"): raise VLLMValidationError( "Stream options can only be defined when `stream=True`.", @@ -543,6 +551,8 @@ def validate_stream_options(cls, data): @model_validator(mode="before") @classmethod def validate_prompt_and_prompt_embeds(cls, data): + if not isinstance(data, dict): + return data prompt = data.get("prompt") prompt_embeds = data.get("prompt_embeds") @@ -562,6 +572,8 @@ def validate_prompt_and_prompt_embeds(cls, data): @model_validator(mode="before") @classmethod def validate_prompt_list_length(cls, data): + if not isinstance(data, dict): + return data max_prompts = envs.VLLM_MAX_COMPLETION_PROMPTS prompt = data.get("prompt") diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index 7ad9f5711623..ff3331612f9f 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -467,6 +467,8 @@ def is_include_output_logprobs(self) -> bool: @model_validator(mode="before") @classmethod def validate_background(cls, data): + if not isinstance(data, dict): + return data if not data.get("background"): return data if not data.get("store", True): @@ -479,6 +481,8 @@ def validate_background(cls, data): @model_validator(mode="before") @classmethod def validate_prompt(cls, data): + if not isinstance(data, dict): + return data if data.get("prompt") is not None: raise VLLMValidationError( "prompt template is not supported", parameter="prompt" @@ -499,6 +503,8 @@ def input_item_parsing(cls, data): Invalid structures are left for Pydantic to reject. """ + if not isinstance(data, dict): + return data input_data = data.get("input") # Early return for None, strings, or bytes diff --git a/vllm/entrypoints/pooling/base/protocol.py b/vllm/entrypoints/pooling/base/protocol.py index fcc7c24b28c5..6e592e30f88c 100644 --- a/vllm/entrypoints/pooling/base/protocol.py +++ b/vllm/entrypoints/pooling/base/protocol.py @@ -239,6 +239,8 @@ class ChatRequestOptionsMixin(OpenAIBaseModel): @model_validator(mode="before") @classmethod def check_generation_prompt(cls, data): + if not isinstance(data, dict): + return data if data.get("continue_final_message") and data.get("add_generation_prompt"): raise VLLMValidationError( "Cannot set both `continue_final_message` and " diff --git a/vllm/entrypoints/serve/tokenize/protocol.py b/vllm/entrypoints/serve/tokenize/protocol.py index 66c122da87de..e68302d1f95b 100644 --- a/vllm/entrypoints/serve/tokenize/protocol.py +++ b/vllm/entrypoints/serve/tokenize/protocol.py @@ -120,6 +120,8 @@ class TokenizeChatRequest(OpenAIBaseModel): @model_validator(mode="before") @classmethod def check_generation_prompt(cls, data): + if not isinstance(data, dict): + return data if data.get("continue_final_message") and data.get("add_generation_prompt"): raise VLLMValidationError( "Cannot set both `continue_final_message` and " diff --git a/vllm/entrypoints/speech_to_text/transcription/protocol.py b/vllm/entrypoints/speech_to_text/transcription/protocol.py index 3220e1505099..bde69a49bd34 100644 --- a/vllm/entrypoints/speech_to_text/transcription/protocol.py +++ b/vllm/entrypoints/speech_to_text/transcription/protocol.py @@ -285,6 +285,8 @@ def to_sampling_params( @model_validator(mode="before") @classmethod def validate_transcription_request(cls, data): + if not isinstance(data, dict): + return data if isinstance(data.get("file"), str): raise HTTPException( status_code=HTTPStatus.UNPROCESSABLE_ENTITY, diff --git a/vllm/entrypoints/speech_to_text/translation/protocol.py b/vllm/entrypoints/speech_to_text/translation/protocol.py index d8836554c565..1a7319e2f105 100644 --- a/vllm/entrypoints/speech_to_text/translation/protocol.py +++ b/vllm/entrypoints/speech_to_text/translation/protocol.py @@ -271,6 +271,8 @@ def to_sampling_params( @model_validator(mode="before") @classmethod def validate_stream_options(cls, data): + if not isinstance(data, dict): + return data stream_opts = ["stream_include_usage", "stream_continuous_usage_stats"] stream = data.get("stream", False) if any(bool(data.get(so, False)) for so in stream_opts) and not stream: From 502af5ed007da93c2a7747b31d0dc0e6110e17a3 Mon Sep 17 00:00:00 2001 From: Yiting Jiang <59356937+yitingdc@users.noreply.github.com> Date: Mon, 17 Aug 2026 10:32:23 +0800 Subject: [PATCH 029/839] [Doc] Add MatrixHub as a model loading source (#50492) Signed-off-by: yiting.jiang --- docs/models/supported_models.md | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 584238af9bf1..f9d55cb97e7a 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -289,6 +289,21 @@ os.environ["http_proxy"] = "http://your.proxy.server:port" os.environ["https_proxy"] = "http://your.proxy.server:port" ``` +### MatrixHub + +[MatrixHub](https://github.com/matrixhub-ai/matrixhub) is a self-hosted model registry and distribution layer that caches models from upstream hubs and serves them over a Hugging Face-compatible API inside your own network. + +Since the API is Hugging Face-compatible, you only need to point `HF_ENDPOINT` at your MatrixHub instance: + +```shell +export HF_ENDPOINT="http://" +vllm serve Qwen/Qwen3-0.6B +``` + +vLLM then downloads model weights from MatrixHub over the internal network instead of the public Hugging Face Hub, which is useful for air-gapped clusters and for avoiding repeated downloads across nodes. + +See the [MatrixHub guide for vLLM](https://matrixhub.ai/docs/guides/use-with-vllm/) for an end-to-end walkthrough, including Docker and Kubernetes deployment examples. + ### ModelScope To use models from [ModelScope](https://www.modelscope.cn) instead of Hugging Face Hub, set an environment variable: From 6664d397bf091cb9371cba481d4efb8233436fe6 Mon Sep 17 00:00:00 2001 From: Cheng Rui <286040359@qq.com> Date: Sun, 16 Aug 2026 19:36:41 -0700 Subject: [PATCH 030/839] [Performance][MRV2] Cache logits-processing request state (#52329) Signed-off-by: Rui Rui4 Cheng Co-authored-by: Rui Rui4 Cheng --- tests/v1/worker/test_gpu_sampler_flags.py | 90 +++++++++++++++++++++ tests/v1/worker/test_gpu_thinking_budget.py | 39 +++++++++ vllm/v1/worker/gpu/sample/sampler.py | 37 +++++---- 3 files changed, 147 insertions(+), 19 deletions(-) create mode 100644 tests/v1/worker/test_gpu_sampler_flags.py diff --git a/tests/v1/worker/test_gpu_sampler_flags.py b/tests/v1/worker/test_gpu_sampler_flags.py new file mode 100644 index 000000000000..3fd1cb3b3c68 --- /dev/null +++ b/tests/v1/worker/test_gpu_sampler_flags.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import numpy as np +import pytest +import torch + +pytest.importorskip("triton") +if not torch.cuda.is_available(): + pytest.skip("CUDA required for sampler flag tests", allow_module_level=True) + +from vllm.sampling_params import SamplingParams +from vllm.v1.worker.gpu.sample.sampler import Sampler +from vllm.v1.worker.gpu.states import RequestState + +DEVICE = torch.device("cuda") +VOCAB_SIZE = 128 + + +class MockReasoningConfig: + reasoning_start_token_ids = [90] + reasoning_end_token_ids = [91] + natural_reasoning_end_token_ids = [91] + + +def _make_sampler() -> Sampler: + req_states = RequestState( + max_num_reqs=4, + max_model_len=64, + max_num_batched_tokens=16, + num_speculative_steps=1, + vocab_size=VOCAB_SIZE, + device=DEVICE, + ) + return Sampler( + max_num_reqs=4, + vocab_size=VOCAB_SIZE, + device=DEVICE, + req_states=req_states, + reasoning_config=MockReasoningConfig(), + ) + + +@pytest.mark.parametrize( + ("sampling_params", "expected"), + [ + pytest.param(SamplingParams(), False, id="defaults"), + pytest.param(SamplingParams(temperature=0.0), False, id="greedy"), + pytest.param( + SamplingParams(thinking_token_budget=3), True, id="thinking-budget" + ), + pytest.param(SamplingParams(logit_bias={1: 1.0}), True, id="logit-bias"), + pytest.param(SamplingParams(frequency_penalty=0.1), True, id="penalty"), + pytest.param(SamplingParams(_bad_words_token_ids=[[1]]), True, id="bad-words"), + pytest.param(SamplingParams(temperature=0.7), True, id="temperature"), + pytest.param(SamplingParams(min_p=0.1), True, id="min-p"), + pytest.param(SamplingParams(top_k=10), True, id="top-k"), + pytest.param(SamplingParams(top_p=0.9), True, id="top-p"), + pytest.param( + SamplingParams.for_sampler_warmup(), True, id="all-logits-processors" + ), + ], +) +def test_logits_processing_cache_matches_request_features( + sampling_params: SamplingParams, expected: bool +): + sampler = _make_sampler() + sampler.add_request(3, prompt_len=1, sampling_params=sampling_params) + + assert sampler.needs_logits_processing[3] == expected + + +def test_logits_processing_cache_is_overwritten_when_slot_is_reused(): + sampler = _make_sampler() + sampler.add_request(3, 1, SamplingParams.for_sampler_warmup()) + sampler.add_request(3, 1, SamplingParams()) + + assert not sampler.needs_logits_processing[3] + + +def test_logits_processing_cache_only_checks_active_requests(): + sampler = _make_sampler() + sampler.add_request(0, 1, SamplingParams(temperature=0.0)) + sampler.add_request(2, 1, SamplingParams.for_sampler_warmup()) + + sampling_only = np.array([0], dtype=np.int32) + with_processing = np.array([0, 2], dtype=np.int32) + + assert not np.any(sampler.needs_logits_processing[sampling_only]) + assert np.any(sampler.needs_logits_processing[with_processing]) diff --git a/tests/v1/worker/test_gpu_thinking_budget.py b/tests/v1/worker/test_gpu_thinking_budget.py index 0b0f19f7dbf0..2bc697b08438 100644 --- a/tests/v1/worker/test_gpu_thinking_budget.py +++ b/tests/v1/worker/test_gpu_thinking_budget.py @@ -12,6 +12,7 @@ ) from vllm.sampling_params import SamplingParams +from vllm.v1.worker.gpu.sample.sampler import Sampler from vllm.v1.worker.gpu.sample.thinking_budget import ThinkingBudgetState from vllm.v1.worker.gpu.states import RequestState @@ -185,6 +186,44 @@ def test_v2_thinking_budget_ignores_plain_request(): assert torch.all(out == 0) +def test_v2_greedy_sampling_applies_thinking_budget(): + """Greedy-only requests must not bypass thinking-budget processing.""" + req_states = _make_req_states([1, START, 10, 11, 12], prompt_len=1) + sampler = Sampler( + max_num_reqs=4, + vocab_size=VOCAB_SIZE, + device=DEVICE, + req_states=req_states, + reasoning_config=MockReasoningConfig(), + ) + sampler.add_request( + req_idx=3, + prompt_len=1, + sampling_params=SamplingParams( + temperature=0.0, + thinking_token_budget=3, + ), + ) + sampler.apply_staged_writes() + + idx_mapping = torch.tensor([3], dtype=torch.int32, device=DEVICE) + idx_mapping_np = idx_mapping.cpu().numpy() + expanded_idx_mapping = idx_mapping.clone() + input_ids = torch.tensor([12], dtype=torch.int32, device=DEVICE) + logits = torch.zeros((1, VOCAB_SIZE), device=DEVICE) + out = sampler.apply_sampling_params( + logits, + expanded_idx_mapping, + idx_mapping, + idx_mapping_np, + torch.tensor([4], dtype=torch.int32, device=DEVICE), + input_ids, + torch.tensor([0], dtype=torch.int32, device=DEVICE), + ) + + assert out[0, END].item() == pytest.approx(1.0e9) + + def test_v2_thinking_budget_latest_prefill_end_disables_forcing(): req_states = _make_req_states( [1, START, 10, 11, 12, END, 13], diff --git a/vllm/v1/worker/gpu/sample/sampler.py b/vllm/v1/worker/gpu/sample/sampler.py index bb7dc40b0211..26371f014e33 100644 --- a/vllm/v1/worker/gpu/sample/sampler.py +++ b/vllm/v1/worker/gpu/sample/sampler.py @@ -53,6 +53,7 @@ def __init__( self.bad_words_state = BadWordsState(req_states) self.logprob_token_ids_state = LogprobTokenIdsState(max_num_reqs, device) self.thinking_budget_state = ThinkingBudgetState(req_states, reasoning_config) + self.needs_logits_processing = np.zeros(max_num_reqs, dtype=bool) self.num_speculative_tokens = num_speculative_tokens self.return_sampling_mask = return_sampling_mask self.use_flashinfer = ( @@ -69,6 +70,22 @@ def add_request( self.logprob_token_ids_state.add_request(req_idx, sampling_params) self.thinking_budget_state.add_request(req_idx, sampling_params) + states = self.sampling_states + temperature = states.temperature.np[req_idx] + self.needs_logits_processing[req_idx] = ( + self.logit_bias_state.use_logit_bias[req_idx] + or self.penalties_state.use_penalty[req_idx] + or self.bad_words_state.num_bad_words.np[req_idx] > 0 + or ( + self.thinking_budget_state.enabled + and self.thinking_budget_state.use_thinking_budget[req_idx] + ) + or (temperature != 0.0 and temperature != 1.0) + or states.min_p.np[req_idx] != 0.0 + or states.top_k.np[req_idx] != states.vocab_size + or states.top_p.np[req_idx] != 1.0 + ) + def apply_staged_writes(self) -> None: self.sampling_states.apply_staged_writes() self.penalties_state.apply_staged_writes() @@ -172,7 +189,7 @@ def apply_sampling_params( expanded_local_pos: torch.Tensor, skip_top_k_top_p: bool = False, ) -> torch.Tensor: - if not self._requires_logits_processing(idx_mapping_np): + if not np.any(self.needs_logits_processing[idx_mapping_np]): return logits # Copy logits to a new FP32 tensor. @@ -228,24 +245,6 @@ def apply_sampling_params( logits, expanded_idx_mapping, idx_mapping_np ) - def _requires_logits_processing(self, idx_mapping_np: np.ndarray) -> bool: - if np.any(self.logit_bias_state.use_logit_bias[idx_mapping_np]): - return True - if np.any(self.penalties_state.use_penalty[idx_mapping_np]): - return True - if np.any(self.bad_words_state.num_bad_words.np[idx_mapping_np] > 0): - return True - - states = self.sampling_states - temperatures = states.temperature.np[idx_mapping_np] - if np.any((temperatures != 0.0) & (temperatures != 1.0)): - return True - if np.any(states.min_p.np[idx_mapping_np] != 0.0): - return True - if np.any(states.top_k.np[idx_mapping_np] != states.vocab_size): - return True - return bool(np.any(states.top_p.np[idx_mapping_np] != 1.0)) - def sample( self, logits: torch.Tensor, From 292187dd8ca1b1bfa195f25b2886262527269999 Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Sun, 16 Aug 2026 23:01:13 -0400 Subject: [PATCH 031/839] [Bugfix][DSv4] Keep indexer scoring in breakable graphs (#52492) Signed-off-by: Lucas Wilkinson Co-authored-by: OpenAI Codex Co-authored-by: Yongye Zhu --- vllm/models/deepseek_v4/attention.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index ff7b92c6b22b..128debf70cb1 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -906,7 +906,10 @@ def forward( attn_metadata = get_forward_context().attn_metadata if isinstance(attn_metadata, dict): indexer_metadata = cast(Any, attn_metadata[self.k_cache.prefix]) - if indexer_metadata.max_seq_len // self.compress_ratio <= self.topk_tokens: + if ( + indexer_metadata.max_seq_len // self.compress_ratio <= self.topk_tokens + and not torch.cuda.is_current_stream_capturing() + ): # candidates num smaller than topk, every candidate is selected # but we still need to build k cache compressor(compressed_kv_score, positions, rotary_emb) From 7ea4b40954084bbafa595faab6b322a5eb124ec0 Mon Sep 17 00:00:00 2001 From: pavelzak Date: Sun, 16 Aug 2026 20:04:15 -0700 Subject: [PATCH 032/839] [Hardware][NVIDIA] Add GB10 fused-MoE fp8 tuning configs (E=256, E=512) (#52502) --- ...,dtype=fp8_w8a8,block_shape=[128,128].json | 147 ++++++++++++++++++ ...,dtype=fp8_w8a8,block_shape=[128,128].json | 147 ++++++++++++++++++ 2 files changed, 294 insertions(+) create mode 100644 vllm/model_executor/layers/fused_moe/configs/E=256,N=512,device_name=NVIDIA_GB10,dtype=fp8_w8a8,block_shape=[128,128].json create mode 100644 vllm/model_executor/layers/fused_moe/configs/E=512,N=512,device_name=NVIDIA_GB10,dtype=fp8_w8a8,block_shape=[128,128].json diff --git a/vllm/model_executor/layers/fused_moe/configs/E=256,N=512,device_name=NVIDIA_GB10,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/fused_moe/configs/E=256,N=512,device_name=NVIDIA_GB10,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 000000000000..f47e4c67b456 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/configs/E=256,N=512,device_name=NVIDIA_GB10,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,147 @@ +{ + "triton_version": "3.5.0", + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + } +} diff --git a/vllm/model_executor/layers/fused_moe/configs/E=512,N=512,device_name=NVIDIA_GB10,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/fused_moe/configs/E=512,N=512,device_name=NVIDIA_GB10,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 000000000000..f47e4c67b456 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/configs/E=512,N=512,device_name=NVIDIA_GB10,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,147 @@ +{ + "triton_version": "3.5.0", + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + } +} From 71b578b9cc37c262aeeb23815e6c478112c97d8b Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Sun, 16 Aug 2026 22:52:24 -0500 Subject: [PATCH 033/839] [ROCm][CI] Use the same-build wheel in Python-only CI (#49514) Signed-off-by: Andreas Karatzas Signed-off-by: Andreas Karatzas --- .../scripts/hardware_ci/run-amd-test.sh | 59 ++++- .buildkite/test-amd.yaml | 5 +- .buildkite/test_areas/misc.yaml | 3 +- tests/standalone_tests/python_only_compile.sh | 203 ++++++++++-------- 4 files changed, 182 insertions(+), 88 deletions(-) diff --git a/.buildkite/scripts/hardware_ci/run-amd-test.sh b/.buildkite/scripts/hardware_ci/run-amd-test.sh index eb8b53c6ccbd..86e4b34869bb 100755 --- a/.buildkite/scripts/hardware_ci/run-amd-test.sh +++ b/.buildkite/scripts/hardware_ci/run-amd-test.sh @@ -261,6 +261,8 @@ validate_native_workspace() { } prepare_native_workspace() { + local test_commands="${1:-}" + if [[ "${VLLM_CI_USE_ARTIFACTS:-0}" != "1" ]]; then echo "Native CI requires VLLM_CI_USE_ARTIFACTS=1" return 1 @@ -281,6 +283,8 @@ prepare_native_workspace() { local recorded_base="" local recorded_commit="" local recorded_wheel="" + local checkout="" + local checkout_commit="" local workspace_dir="${VLLM_CI_WORKSPACE:-/vllm-workspace}" local wheel_dir="" local attempt=0 @@ -400,6 +404,53 @@ prepare_native_workspace() { return 1 fi + # The ROCm artifact intentionally contains only the installed wheel and the + # test workspace. The Python-only compilation job also needs setup.py and the + # vllm source tree, so overlay the matching Buildkite checkout for that job. + if [[ "${test_commands}" == *python_only_compile.sh* ]]; then + checkout="${BUILDKITE_BUILD_CHECKOUT_PATH:-}" + if [[ -z "${checkout}" || ! -d "${checkout}" ]]; then + echo "Python-only native CI requires BUILDKITE_BUILD_CHECKOUT_PATH" >&2 + return 1 + fi + if ! git -c "safe.directory=${checkout}" -C "${checkout}" \ + rev-parse --is-inside-work-tree >/dev/null 2>&1; then + echo "Buildkite checkout is not a Git worktree: ${checkout}" >&2 + return 1 + fi + checkout_commit=$( + git -c "safe.directory=${checkout}" -C "${checkout}" rev-parse HEAD + ) || return 1 + if [[ "${checkout_commit}" != "${recorded_commit}" ]]; then + echo "Buildkite checkout ${checkout_commit} does not match ROCm artifact ${recorded_commit}" >&2 + return 1 + fi + + # setup.py normally derives this from .git via setuptools-scm. The native + # source overlay deliberately excludes Git metadata, so preserve the exact + # version from the already installed, artifact-matched wheel. + VLLM_VERSION_OVERRIDE=$( + python3 -c 'import importlib.metadata as m; print(m.version("vllm"))' + ) || return 1 + export VLLM_VERSION_OVERRIDE + VLLM_PRECOMPILED_WHEEL_LOCATION="${wheels[0]}" + export VLLM_PRECOMPILED_WHEEL_LOCATION + echo "INFO: native Python-only wheel=${VLLM_PRECOMPILED_WHEEL_LOCATION}" + + echo "--- Overlaying full source checkout for Python-only compilation" + # Archive the verified commit instead of copying the worktree so dirty or + # untracked agent files cannot contaminate the artifact-matched workspace. + git -c "safe.directory=${checkout}" -C "${checkout}" \ + archive --format=tar "${recorded_commit}" \ + | tar --no-same-owner -C "${workspace_dir}" -xf - || return 1 + for required_source in setup.py pyproject.toml vllm; do + if [[ ! -e "${workspace_dir}/${required_source}" ]]; then + echo "Full source checkout is missing ${required_source}" >&2 + return 1 + fi + done + fi + return 0 } @@ -954,7 +1005,13 @@ if is_native_runtime; then echo "Failed to initialize the native test environment" exit 1 fi - if ! prepare_native_workspace; then + if [[ "${commands}" == *python_only_compile.sh* ]]; then + # This no-GPU job validates the ROCm precompiled/editable install path, + # rather than CPU runtime platform selection. + VLLM_TARGET_DEVICE=rocm + export VLLM_TARGET_DEVICE + fi + if ! prepare_native_workspace "${commands}"; then echo "Failed to prepare native test workspace" exit 1 fi diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index c1658e7a8331..b477b298bb10 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -411,8 +411,9 @@ steps: - label: Python-only Installation # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_1 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 126bc657048c..8f168cbd367f 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -275,7 +275,8 @@ steps: - bash standalone_tests/python_only_compile.sh mirror: amd: - device: mi250_1 + dind: false + device: mi300_1 timeout_in_minutes: 55 soft_fail: true depends_on: diff --git a/tests/standalone_tests/python_only_compile.sh b/tests/standalone_tests/python_only_compile.sh index 8e3d08d45665..7036dd4314f2 100644 --- a/tests/standalone_tests/python_only_compile.sh +++ b/tests/standalone_tests/python_only_compile.sh @@ -4,37 +4,70 @@ set -e -# ROCm CI runs this script inside `run-amd-test.sh` where /vllm-workspace often has no .git -# (wheel artifact layout). The wrapper passes CI_STANDALONE_MERGE_BASE from the agent checkout. merge_base_commit="" -if [[ -n "${CI_STANDALONE_MERGE_BASE:-}" ]]; then - merge_base_commit="${CI_STANDALONE_MERGE_BASE}" -elif merge_base_commit="$(git -C /vllm-workspace merge-base HEAD origin/main 2>/dev/null)"; then - : -elif merge_base_commit="$(git merge-base HEAD origin/main 2>/dev/null)"; then - : -else - echo "ERROR: need a git checkout or CI_STANDALONE_MERGE_BASE to resolve wheels.vllm.ai commit." >&2 - exit 1 +rocm_wheel="" +is_rocm=0 +_vllm_target_lower="$(printf '%s' "${VLLM_TARGET_DEVICE:-}" | tr '[:upper:]' '[:lower:]')" +if [[ "${_vllm_target_lower}" == "rocm" || -n "${ROCM_PATH:-}" || -d /opt/rocm ]] \ + || command -v rocminfo >/dev/null 2>&1; then + is_rocm=1 fi +unset -v _vllm_target_lower -echo "INFO: current merge base commit with main: $merge_base_commit" -if git show --oneline -s "$merge_base_commit" 2>/dev/null; then - : -else - echo "INFO: git show unavailable in this environment; using SHA above for precompiled metadata." +if [[ "${is_rocm}" == "1" ]]; then + # Native CI passes the verified wheel artifact explicitly. Legacy ROCm + # images carry the same-build wheel in /opt/vllm-wheels. + if [[ -n "${VLLM_PRECOMPILED_WHEEL_LOCATION:-}" ]]; then + rocm_wheel="${VLLM_PRECOMPILED_WHEEL_LOCATION}" + if [[ ! -f "${rocm_wheel}" || "$(basename "${rocm_wheel}")" != vllm-*.whl ]]; then + echo "ERROR: invalid ROCm wheel location: ${rocm_wheel}" >&2 + exit 1 + fi + rocm_wheel="$(realpath -- "${rocm_wheel}")" + elif [[ -d /opt/vllm-wheels ]]; then + shopt -s nullglob + rocm_wheels=(/opt/vllm-wheels/vllm-*.whl) + shopt -u nullglob + if [[ "${#rocm_wheels[@]}" -ne 1 ]]; then + echo "ERROR: expected exactly one vLLM wheel in /opt/vllm-wheels, found ${#rocm_wheels[@]}." >&2 + exit 1 + fi + rocm_wheel="${rocm_wheels[0]}" + fi fi -# test whether the metadata.json url is valid, retry each 3 minutes up to 5 times -# this avoids cumbersome error messages & manual retries in case the precompiled wheel -# for the given commit is still being built in the release pipeline -_vllm_target_lower="$(printf '%s' "${VLLM_TARGET_DEVICE:-}" | tr '[:upper:]' '[:lower:]')" -if [[ "${_vllm_target_lower}" == "rocm" ]] || [[ -d /opt/rocm ]] || command -v rocminfo >/dev/null 2>&1; then - _rocm_env_variant="$(python3 - <<'PY' +if [[ -n "${rocm_wheel}" ]]; then + echo "INFO: using same-build ROCm wheel: ${rocm_wheel}" +else + # Some CI images do not include .git under /vllm-workspace. Their wrapper + # passes CI_STANDALONE_MERGE_BASE from the agent checkout. + if [[ -n "${CI_STANDALONE_MERGE_BASE:-}" ]]; then + merge_base_commit="${CI_STANDALONE_MERGE_BASE}" + elif merge_base_commit="$(git -C /vllm-workspace merge-base HEAD origin/main 2>/dev/null)"; then + : + elif merge_base_commit="$(git merge-base HEAD origin/main 2>/dev/null)"; then + : + else + echo "ERROR: need a git checkout or CI_STANDALONE_MERGE_BASE to resolve wheels.vllm.ai commit." >&2 + exit 1 + fi + + echo "INFO: current merge base commit with main: $merge_base_commit" + if git show --oneline -s "$merge_base_commit" 2>/dev/null; then + : + else + echo "INFO: git show unavailable in this environment; using SHA above for precompiled metadata." + fi + + # Test whether the metadata.json URL is valid, retry each 5 minutes up to 5 times. + # This avoids manual retries while a new main-branch wheel is still publishing. + if [[ "${is_rocm}" == "1" ]]; then + _rocm_env_variant="$(python3 - <<'PY' import ctypes import os from pathlib import Path + def get_rocm_version() -> str | None: rocm_home = os.environ.get("ROCM_HOME") or os.environ.get("ROCM_PATH") or "/opt/rocm" try: @@ -54,76 +87,77 @@ def get_rocm_version() -> str | None: return None return None + version = get_rocm_version() if version: print(f"rocm{version.replace('.', '')}", end="") PY )" - _available_variants="$(curl -sf "https://wheels.vllm.ai/rocm/${merge_base_commit}/" \ - | grep -oP 'rocm\d+' | sort -u | tr '\n' ' ' || true)" - if [[ -n "${VLLM_PRECOMPILED_WHEEL_VARIANT:-}" ]]; then - _rocm_variant="${VLLM_PRECOMPILED_WHEEL_VARIANT}" - if [[ -n "${_rocm_env_variant}" && "${_rocm_variant}" != "${_rocm_env_variant}" ]]; then - echo "ERROR: VLLM_PRECOMPILED_WHEEL_VARIANT=${_rocm_variant} does not match detected environment ROCm variant ${_rocm_env_variant}" >&2 + _available_variants="$(curl -sf "https://wheels.vllm.ai/rocm/${merge_base_commit}/" \ + | grep -oP 'rocm\d+' | sort -u | tr '\n' ' ' || true)" + if [[ -n "${VLLM_PRECOMPILED_WHEEL_VARIANT:-}" ]]; then + _rocm_variant="${VLLM_PRECOMPILED_WHEEL_VARIANT}" + if [[ -n "${_rocm_env_variant}" && "${_rocm_variant}" != "${_rocm_env_variant}" ]]; then + echo "ERROR: VLLM_PRECOMPILED_WHEEL_VARIANT=${_rocm_variant} does not match detected environment ROCm variant ${_rocm_env_variant}" >&2 + exit 1 + fi + else + _rocm_variant="${_rocm_env_variant}" + fi + if [[ -z "${_rocm_variant}" ]]; then + echo "ERROR: Could not detect ROCm variant from the environment for commit ${merge_base_commit}" >&2 + exit 1 + fi + if [[ -z "${_available_variants}" ]] \ + || [[ " ${_available_variants} " != *" ${_rocm_variant} "* ]]; then + echo "ERROR: Environment ROCm variant '${_rocm_variant}' is not published for commit ${merge_base_commit} (available:${_available_variants:-none})" >&2 exit 1 fi + meta_json_url="https://wheels.vllm.ai/rocm/${merge_base_commit}/${_rocm_variant}/vllm/metadata.json" + unset -v _rocm_env_variant _available_variants _rocm_variant else - _rocm_variant="${_rocm_env_variant}" - fi - if [[ -z "${_rocm_variant}" ]]; then - echo "ERROR: Could not detect ROCm variant from the environment for commit ${merge_base_commit}" >&2 - exit 1 + meta_json_url="https://wheels.vllm.ai/${merge_base_commit}/vllm/metadata.json" fi - if [[ -z "${_available_variants}" ]] \ - || [[ " ${_available_variants} " != *" ${_rocm_variant} "* ]]; then - echo "ERROR: Environment ROCm variant '${_rocm_variant}' is not published for commit ${merge_base_commit} (available:${_available_variants:-none})" >&2 - exit 1 - fi - meta_json_url="https://wheels.vllm.ai/rocm/${merge_base_commit}/${_rocm_variant}/vllm/metadata.json" - unset -v _rocm_env_variant _available_variants _rocm_variant -else - meta_json_url="https://wheels.vllm.ai/${merge_base_commit}/vllm/metadata.json" -fi -unset -v _vllm_target_lower -echo "INFO: will use metadata.json from ${meta_json_url}" - -for i in {1..5}; do - echo "Checking metadata.json URL (attempt $i)..." - if curl --fail "$meta_json_url" > metadata.json; then - echo "INFO: metadata.json URL is valid." - # check whether it is valid json by python (printed to stdout) - if python3 -m json.tool metadata.json; then - echo "INFO: metadata.json is valid JSON. Proceeding with the check." - # check whether there is an object in the json matching: - # "package_name": "vllm", and "platform_tag" matches the current architecture - # see `determine_wheel_url` in setup.py for more details - if python3 -c "import platform as p,json as j,sys as s; d = j.load(open('metadata.json')); \ - s.exit(int(not any(o.get('package_name') == 'vllm' and p.machine() in o.get('platform_tag') \ - for o in d)))" 2>/dev/null; then - echo "INFO: metadata.json contains a pre-compiled wheel for the current architecture." - break + echo "INFO: will use metadata.json from ${meta_json_url}" + + for i in {1..5}; do + echo "Checking metadata.json URL (attempt $i)..." + if curl --fail "$meta_json_url" > metadata.json; then + echo "INFO: metadata.json URL is valid." + # check whether it is valid json by python (printed to stdout) + if python3 -m json.tool metadata.json; then + echo "INFO: metadata.json is valid JSON. Proceeding with the check." + # check whether there is an object in the json matching: + # "package_name": "vllm", and "platform_tag" matches the current architecture + # see `determine_wheel_url` in setup.py for more details + if python3 -c "import platform as p,json as j,sys as s; d = j.load(open('metadata.json')); \ + s.exit(int(not any(o.get('package_name') == 'vllm' and p.machine() in o.get('platform_tag') \ + for o in d)))" 2>/dev/null; then + echo "INFO: metadata.json contains a pre-compiled wheel for the current architecture." + break + else + echo "WARN: metadata.json does not have a pre-compiled wheel for the current architecture." + fi else - echo "WARN: metadata.json does not have a pre-compiled wheel for the current architecture." + echo "CRITICAL: metadata.json exists but is not valid JSON, please do report in #sig-ci channel!" + echo "INFO: metadata.json content:" + cat metadata.json + exit 1 fi - else - echo "CRITICAL: metadata.json exists but is not valid JSON, please do report in #sig-ci channel!" - echo "INFO: metadata.json content:" - cat metadata.json + fi + # failure handling & retry logic + if [ "$i" -eq 5 ]; then + echo "ERROR: metadata is still not available after 5 attempts." + echo "ERROR: Please check whether the precompiled wheel for commit $merge_base_commit is available." + echo " NOTE: If $merge_base_commit is a new commit on main, maybe try again after its release pipeline finishes." + echo " NOTE: If it fails, please report in #sig-ci channel." exit 1 + else + echo "WARNING: metadata is not available. Retrying after 5 minutes..." + sleep 300 fi - fi - # failure handling & retry logic - if [ "$i" -eq 5 ]; then - echo "ERROR: metadata is still not available after 5 attempts." - echo "ERROR: Please check whether the precompiled wheel for commit $merge_base_commit is available." - echo " NOTE: If $merge_base_commit is a new commit on main, maybe try again after its release pipeline finishes." - echo " NOTE: If it fails, please report in #sig-ci channel." - exit 1 - else - echo "WARNING: metadata is not available. Retrying after 5 minutes..." - sleep 300 - fi -done + done +fi set -x @@ -143,16 +177,17 @@ fi apt remove --purge build-essential -y apt autoremove -y +rm -f /tmp/changed.file echo 'import os; os.system("touch /tmp/changed.file")' >> vllm/__init__.py # ROCm CI uses setuptools develop for editable installs (see Dockerfile.rocm and run-amd-test.sh). -_vllm_target_lower="$(printf '%s' "${VLLM_TARGET_DEVICE:-}" | tr '[:upper:]' '[:lower:]')" -if [[ "${_vllm_target_lower}" == "rocm" ]]; then - VLLM_PRECOMPILED_WHEEL_COMMIT=$merge_base_commit VLLM_USE_PRECOMPILED=1 python3 setup.py develop +if [[ -n "${rocm_wheel}" ]]; then + VLLM_PRECOMPILED_WHEEL_LOCATION="${rocm_wheel}" VLLM_USE_PRECOMPILED=1 python3 setup.py develop --no-deps +elif [[ "${is_rocm}" == "1" ]]; then + VLLM_PRECOMPILED_WHEEL_COMMIT=$merge_base_commit VLLM_USE_PRECOMPILED=1 python3 setup.py develop --no-deps else - VLLM_PRECOMPILED_WHEEL_COMMIT=$merge_base_commit VLLM_USE_PRECOMPILED=1 pip3 install -vvv -e . + VLLM_PRECOMPILED_WHEEL_COMMIT=$merge_base_commit VLLM_USE_PRECOMPILED=1 pip3 install -vvv -e . fi -unset -v _vllm_target_lower # Run the script python3 -c 'import vllm' From 311b3513af33bc29b4acb2fde2e9313e5e9966a0 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Sun, 16 Aug 2026 23:08:35 -0500 Subject: [PATCH 034/839] [ROCm][CI] Avoid forcing FlashAttention in the ColPali pooling test (#52565) Signed-off-by: Andreas Karatzas --- tests/models/multimodal/pooling/test_colpali.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/models/multimodal/pooling/test_colpali.py b/tests/models/multimodal/pooling/test_colpali.py index 41e6393dc793..86e595824d9c 100644 --- a/tests/models/multimodal/pooling/test_colpali.py +++ b/tests/models/multimodal/pooling/test_colpali.py @@ -19,6 +19,7 @@ ChatCompletionContentPartTextParam, ) from vllm.entrypoints.pooling.scoring.typing import ScoreMultiModalParam +from vllm.platforms import current_platform from ....conftest import VllmRunner @@ -215,6 +216,7 @@ def _run_multimodal_text_query_image_docs_test( _make_image_mm_param(red_image), _make_image_mm_param(blue_image), ] + attention_backend = "FLASH_ATTN" if current_platform.is_cuda() else None with vllm_runner( model, @@ -223,7 +225,7 @@ def _run_multimodal_text_query_image_docs_test( max_model_len=4096, enforce_eager=True, gpu_memory_utilization=GPU_MEMORY_UTILIZATION, - attention_backend="FLASH_ATTN", + attention_backend=attention_backend, kernel_config={"enable_flashinfer_autotune": False}, ) as vllm_model: assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner From 53e211d2923192ea1f7442de93d59cfbb150e089 Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Mon, 17 Aug 2026 13:19:07 +0800 Subject: [PATCH 035/839] [CI/Build] Reduce more duplicate runner startup in tests (#52570) Signed-off-by: Isotr0py --- tests/models/language/pooling/test_colbert.py | 166 ++++----- .../pooling/test_truncation_control.py | 55 ++- .../multimodal/generation/test_whisper.py | 88 ++--- tests/models/multimodal/pooling/test_clip.py | 135 +++----- .../multimodal/pooling/test_colmodernvbert.py | 111 +++--- .../models/multimodal/pooling/test_colpali.py | 239 +++++-------- .../multimodal/pooling/test_colqwen3_5.py | 131 +++---- .../pooling/test_llama_nemotron_vl.py | 321 ++++++++---------- .../models/multimodal/pooling/test_siglip.py | 163 ++++----- tests/models/quantization/test_awq.py | 32 +- 10 files changed, 586 insertions(+), 855 deletions(-) diff --git a/tests/models/language/pooling/test_colbert.py b/tests/models/language/pooling/test_colbert.py index 3057e14060c6..975c4d6e8fb0 100644 --- a/tests/models/language/pooling/test_colbert.py +++ b/tests/models/language/pooling/test_colbert.py @@ -212,14 +212,13 @@ def colbert_extra_kwargs(colbert_spec): return colbert_spec["extra_kwargs"] -def test_colbert_token_embed( +@pytest.fixture(scope="module") +def colbert_model( vllm_runner, colbert_model_name, - colbert_dim, colbert_max_model_len, colbert_extra_kwargs, ): - """Test that ColBERT model produces token embeddings.""" with vllm_runner( colbert_model_name, runner="pooling", @@ -228,112 +227,84 @@ def test_colbert_token_embed( enforce_eager=True, **colbert_extra_kwargs, ) as vllm_model: - outputs = vllm_model.token_embed([TEXTS_1[0]]) + yield vllm_model + + +def test_colbert_token_embed( + colbert_model, + colbert_dim, +): + """Test that ColBERT model produces token embeddings.""" + outputs = colbert_model.token_embed([TEXTS_1[0]]) - assert len(outputs) == 1 - emb = torch.as_tensor(outputs[0]) - assert emb.dim() == 2 - assert emb.shape[1] == colbert_dim - assert emb.shape[0] > 1 + assert len(outputs) == 1 + emb = torch.as_tensor(outputs[0]) + assert emb.dim() == 2 + assert emb.shape[1] == colbert_dim + assert emb.shape[0] > 1 def test_colbert_late_interaction_1_to_1( - vllm_runner, - colbert_model_name, - colbert_max_model_len, - colbert_extra_kwargs, + colbert_model, ): """Test ColBERT late interaction scoring with 1:1 query-document pair.""" - with vllm_runner( - colbert_model_name, - runner="pooling", - dtype=DTYPE, - max_model_len=colbert_max_model_len, - enforce_eager=True, - **colbert_extra_kwargs, - ) as vllm_model: - q_outputs = vllm_model.token_embed([TEXTS_1[0]]) - d_outputs = vllm_model.token_embed([TEXTS_2[0]]) + q_outputs = colbert_model.token_embed([TEXTS_1[0]]) + d_outputs = colbert_model.token_embed([TEXTS_2[0]]) - q_emb = torch.as_tensor(q_outputs[0]) - d_emb = torch.as_tensor(d_outputs[0]) + q_emb = torch.as_tensor(q_outputs[0]) + d_emb = torch.as_tensor(d_outputs[0]) - manual_score = compute_maxsim_score(q_emb, d_emb).item() + manual_score = compute_maxsim_score(q_emb, d_emb).item() - vllm_scores = vllm_model.score(TEXTS_1[0], TEXTS_2[0]) + vllm_scores = colbert_model.score(TEXTS_1[0], TEXTS_2[0]) - assert len(vllm_scores) == 1 - assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01) + assert len(vllm_scores) == 1 + assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01) def test_colbert_late_interaction_1_to_N( - vllm_runner, - colbert_model_name, - colbert_max_model_len, - colbert_extra_kwargs, + colbert_model, ): """Test ColBERT late interaction scoring with 1:N query-documents.""" - with vllm_runner( - colbert_model_name, - runner="pooling", - dtype=DTYPE, - max_model_len=colbert_max_model_len, - enforce_eager=True, - **colbert_extra_kwargs, - ) as vllm_model: - q_outputs = vllm_model.token_embed([TEXTS_1[0]]) - d_outputs = vllm_model.token_embed(TEXTS_2) + q_outputs = colbert_model.token_embed([TEXTS_1[0]]) + d_outputs = colbert_model.token_embed(TEXTS_2) - q_emb = torch.as_tensor(q_outputs[0]) + q_emb = torch.as_tensor(q_outputs[0]) - manual_scores = [] - for d_out in d_outputs: - d_emb = torch.as_tensor(d_out) - manual_scores.append(compute_maxsim_score(q_emb, d_emb).item()) + manual_scores = [] + for d_out in d_outputs: + d_emb = torch.as_tensor(d_out) + manual_scores.append(compute_maxsim_score(q_emb, d_emb).item()) - vllm_scores = vllm_model.score(TEXTS_1[0], TEXTS_2) + vllm_scores = colbert_model.score(TEXTS_1[0], TEXTS_2) - assert len(vllm_scores) == 2 - for i in range(2): - assert vllm_scores[i] == pytest.approx(manual_scores[i], rel=0.01) + assert len(vllm_scores) == 2 + for i in range(2): + assert vllm_scores[i] == pytest.approx(manual_scores[i], rel=0.01) def test_colbert_late_interaction_N_to_N( - vllm_runner, - colbert_model_name, - colbert_max_model_len, - colbert_extra_kwargs, + colbert_model, ): """Test ColBERT late interaction scoring with N:N query-documents.""" - with vllm_runner( - colbert_model_name, - runner="pooling", - dtype=DTYPE, - max_model_len=colbert_max_model_len, - enforce_eager=True, - **colbert_extra_kwargs, - ) as vllm_model: - q_outputs = vllm_model.token_embed(TEXTS_1) - d_outputs = vllm_model.token_embed(TEXTS_2) + q_outputs = colbert_model.token_embed(TEXTS_1) + d_outputs = colbert_model.token_embed(TEXTS_2) - manual_scores = [] - for q_out, d_out in zip(q_outputs, d_outputs): - q_emb = torch.as_tensor(q_out) - d_emb = torch.as_tensor(d_out) - manual_scores.append(compute_maxsim_score(q_emb, d_emb).item()) + manual_scores = [] + for q_out, d_out in zip(q_outputs, d_outputs): + q_emb = torch.as_tensor(q_out) + d_emb = torch.as_tensor(d_out) + manual_scores.append(compute_maxsim_score(q_emb, d_emb).item()) - vllm_scores = vllm_model.score(TEXTS_1, TEXTS_2) + vllm_scores = colbert_model.score(TEXTS_1, TEXTS_2) - assert len(vllm_scores) == 2 - for i in range(2): - assert vllm_scores[i] == pytest.approx(manual_scores[i], rel=0.01) + assert len(vllm_scores) == 2 + for i in range(2): + assert vllm_scores[i] == pytest.approx(manual_scores[i], rel=0.01) def test_colbert_relevance_ordering( - vllm_runner, - colbert_model_name, - colbert_max_model_len, - colbert_extra_kwargs, + colbert_model, ): """Test that ColBERT scores relevant documents higher than irrelevant.""" query = "What is machine learning?" @@ -343,40 +314,19 @@ def test_colbert_relevance_ordering( "Deep learning uses neural networks.", ] - with vllm_runner( - colbert_model_name, - runner="pooling", - dtype=DTYPE, - max_model_len=colbert_max_model_len, - enforce_eager=True, - **colbert_extra_kwargs, - ) as vllm_model: - scores = vllm_model.score(query, documents) + scores = colbert_model.score(query, documents) - assert len(scores) == 3 - assert scores[0] > scores[1], "ML doc should score higher than Python doc" - assert scores[2] > scores[1], "DL doc should score higher than Python doc" + assert len(scores) == 3 + assert scores[0] > scores[1], "ML doc should score higher than Python doc" + assert scores[2] > scores[1], "DL doc should score higher than Python doc" def test_colbert_embed_not_supported( - vllm_runner, - colbert_model_name, - colbert_max_model_len, - colbert_extra_kwargs, + colbert_model, ): - """Test that ColBERT model does not support 'embed' task.""" - with ( - vllm_runner( - colbert_model_name, - runner="pooling", - dtype=DTYPE, - max_model_len=colbert_max_model_len, - enforce_eager=True, - **colbert_extra_kwargs, - ) as vllm_model, - pytest.raises(ValueError, match="Embedding API is not supported"), - ): - vllm_model.embed([TEXTS_1[0]]) + """Test that ColBERT model does not support the embed task.""" + with pytest.raises(ValueError, match="Embedding API is not supported"): + colbert_model.embed([TEXTS_1[0]]) @pytest.mark.parametrize( diff --git a/tests/models/language/pooling/test_truncation_control.py b/tests/models/language/pooling/test_truncation_control.py index 50e8cdbd064c..c4485195558a 100644 --- a/tests/models/language/pooling/test_truncation_control.py +++ b/tests/models/language/pooling/test_truncation_control.py @@ -22,60 +22,47 @@ field.""" +@pytest.fixture(scope="module") +def vllm_model(vllm_runner): + with vllm_runner( + MODEL_NAME, runner="pooling", max_model_len=max_model_len + ) as model: + yield model + + def test_smaller_truncation_size( - vllm_runner, model_name=MODEL_NAME, input_str=input_str + vllm_model, ): truncate_prompt_tokens = 10 - with vllm_runner( - model_name, runner="pooling", max_model_len=max_model_len - ) as vllm_model: - vllm_output = vllm_model.llm.embed( - input_str, - tokenization_kwargs=dict(truncate_prompt_tokens=truncate_prompt_tokens), - ) + vllm_output = vllm_model.llm.embed( + input_str, + tokenization_kwargs=dict(truncate_prompt_tokens=truncate_prompt_tokens), + ) prompt_tokens = vllm_output[0].prompt_token_ids assert len(prompt_tokens) == truncate_prompt_tokens -def test_max_truncation_size(vllm_runner, model_name=MODEL_NAME, input_str=input_str): +def test_max_truncation_size(vllm_model): truncate_prompt_tokens = -1 - with vllm_runner( - model_name, runner="pooling", max_model_len=max_model_len - ) as vllm_model: - vllm_output = vllm_model.llm.embed( - input_str, - tokenization_kwargs=dict(truncate_prompt_tokens=truncate_prompt_tokens), - ) + vllm_output = vllm_model.llm.embed( + input_str, + tokenization_kwargs=dict(truncate_prompt_tokens=truncate_prompt_tokens), + ) prompt_tokens = vllm_output[0].prompt_token_ids assert len(prompt_tokens) == max_model_len -def test_bigger_truncation_size( - vllm_runner, model_name=MODEL_NAME, input_str=input_str -): +def test_bigger_truncation_size(vllm_model): truncate_prompt_tokens = max_model_len + 1 - with ( - pytest.raises(VLLMValidationError), - vllm_runner( - model_name, runner="pooling", max_model_len=max_model_len - ) as vllm_model, - ): - llm_output = vllm_model.llm.embed( + with pytest.raises(VLLMValidationError): + vllm_model.llm.embed( input_str, tokenization_kwargs=dict(truncate_prompt_tokens=truncate_prompt_tokens), ) - - assert ( - llm_output - == f"""truncate_prompt_tokens value - ({truncate_prompt_tokens}) is greater than - max_model_len ({max_model_len}). Please, select - a smaller truncation size.""" - ) diff --git a/tests/models/multimodal/generation/test_whisper.py b/tests/models/multimodal/generation/test_whisper.py index 310c7fe0f563..fef6b69911b7 100644 --- a/tests/models/multimodal/generation/test_whisper.py +++ b/tests/models/multimodal/generation/test_whisper.py @@ -20,6 +20,7 @@ HF_PROMPT = "" # Whisper expects 16kHz audio WHISPER_SAMPLE_RATE = 16000 +BEAM_WIDTHS = (1, 2) @pytest.fixture(autouse=True) @@ -127,13 +128,11 @@ def check_model_available(model: str) -> None: @pytest.mark.parametrize("dtype", ["half"]) @pytest.mark.parametrize("max_tokens", [64]) -@pytest.mark.parametrize("beam_width", [1, 2]) def test_beam_search_encoder_decoder( hf_runner, vllm_runner, dtype: str, max_tokens: int, - beam_width: int, resampled_assets, ) -> None: """Test beam search with encoder-decoder models (Whisper).""" @@ -146,12 +145,15 @@ def test_beam_search_encoder_decoder( ] with hf_runner(model, dtype=dtype, auto_cls=AutoModelForSpeechSeq2Seq) as hf_model: - hf_outputs = hf_model.generate_beam_search( - hf_prompts, - beam_width=beam_width, - max_tokens=max_tokens, - audios=resampled_assets, - ) + hf_outputs_by_beam_width = [ + hf_model.generate_beam_search( + hf_prompts, + beam_width=beam_width, + max_tokens=max_tokens, + audios=resampled_assets, + ) + for beam_width in BEAM_WIDTHS + ] # Test both explicit encoder/decoder prompts vllm_prompts = [ @@ -179,38 +181,46 @@ def test_beam_search_encoder_decoder( limit_mm_per_prompt={"audio": 2}, enforce_eager=True, ) as vllm_model: - vllm_outputs = vllm_model.generate_beam_search( - vllm_prompts, - beam_width=beam_width, - max_tokens=max_tokens, - ) - - for i in range(len(vllm_prompts)): - hf_output_ids, hf_output_texts = hf_outputs[i] - vllm_output_ids, vllm_output_texts = vllm_outputs[i] - - for j, (hf_text, vllm_text) in enumerate( - zip(hf_output_texts, vllm_output_texts) - ): - print(f">>>{j}-th hf output [NOTE: special tokens are filtered]:") - print(hf_text) - print(f">>>{j}-th vllm output:") - print(vllm_text) - - # Check that we got the same number of beams - assert len(hf_output_ids) == len(vllm_output_ids) - - # For encoder-decoder models, we primarily want to verify that: - # 1. Beam search completes without errors - # 2. We get the expected number of beams - # 3. Outputs are reasonable (non-empty, diverse beams) - for j in range(len(vllm_output_ids)): - # Check that outputs are not empty - assert len(vllm_output_ids[j]) > 0, f"Prompt {i}, beam {j}: empty output" - # Check that decoded text is not empty - assert len(vllm_output_texts[j].strip()) > 0, ( - f"Prompt {i}, beam {j}: empty text output" + vllm_outputs_by_beam_width = [ + vllm_model.generate_beam_search( + vllm_prompts, + beam_width=beam_width, + max_tokens=max_tokens, ) + for beam_width in BEAM_WIDTHS + ] + + for beam_width, hf_outputs, vllm_outputs in zip( + BEAM_WIDTHS, hf_outputs_by_beam_width, vllm_outputs_by_beam_width + ): + for i in range(len(vllm_prompts)): + hf_output_ids, hf_output_texts = hf_outputs[i] + vllm_output_ids, vllm_output_texts = vllm_outputs[i] + + for j, (hf_text, vllm_text) in enumerate( + zip(hf_output_texts, vllm_output_texts) + ): + print(f">>>{j}-th hf output [NOTE: special tokens are filtered]:") + print(hf_text) + print(f">>>{j}-th vllm output:") + print(vllm_text) + + # Check that we got the same number of beams + assert len(hf_output_ids) == len(vllm_output_ids) == beam_width + + # For encoder-decoder models, we primarily want to verify that: + # 1. Beam search completes without errors + # 2. We get the expected number of beams + # 3. Outputs are reasonable (non-empty, diverse beams) + for j in range(len(vllm_output_ids)): + # Check that outputs are not empty + assert len(vllm_output_ids[j]) > 0, ( + f"Prompt {i}, beam {j}: empty output" + ) + # Check that decoded text is not empty + assert len(vllm_output_texts[j].strip()) > 0, ( + f"Prompt {i}, beam {j}: empty text output" + ) def test_parse_language_detection_output(): diff --git a/tests/models/multimodal/pooling/test_clip.py b/tests/models/multimodal/pooling/test_clip.py index 14ede6c1d328..54dd8556e903 100644 --- a/tests/models/multimodal/pooling/test_clip.py +++ b/tests/models/multimodal/pooling/test_clip.py @@ -26,8 +26,7 @@ def _run_test( hf_runner: type[HfRunner], vllm_runner: type[VllmRunner], - input_texts: list[str], - input_images: PromptImageInput, + input_cases: list[tuple[list[str], PromptImageInput]], model: str, *, dtype: str, @@ -39,105 +38,75 @@ def _run_test( with vllm_runner( model, runner="pooling", dtype=dtype, enforce_eager=True, max_model_len=77 ) as vllm_model: - vllm_outputs = vllm_model.embed(input_texts, images=input_images) - - with hf_runner(model, dtype=dtype, auto_cls=CLIPModel) as hf_model: - all_inputs = hf_model.get_inputs(input_texts, images=input_images) - - all_outputs = [] - for inputs in all_inputs: - inputs = hf_model.wrap_device(inputs) - - if "pixel_values" in inputs: - pooled_output = hf_model.model.get_image_features( - pixel_values=inputs.pixel_values, - ) - else: - pooled_output = hf_model.model.get_text_features( - input_ids=inputs.input_ids, - attention_mask=inputs.attention_mask, - ) - - if not isinstance(pooled_output, torch.Tensor): - pooled_output = pooled_output.pooler_output - pooled_output = pooled_output.squeeze(0) - all_outputs.append(pooled_output.tolist()) - - hf_outputs = all_outputs - - check_embeddings_close( - embeddings_0_lst=hf_outputs, - embeddings_1_lst=vllm_outputs, - name_0="hf", - name_1="vllm", - ) + vllm_outputs_per_case = [ + vllm_model.embed(input_texts, images=input_images) + for input_texts, input_images in input_cases + ] + texts = [HF_TEXT_PROMPTS[0]] + images = [input_cases[1][1][0]] + with pytest.raises(ValueError, match="not both"): + vllm_model.embed(texts, images=images) -@pytest.mark.parametrize("model", MODELS) -@pytest.mark.parametrize("dtype", ["float"]) -def test_models_text( - hf_runner, - vllm_runner, - image_assets, - model: str, - dtype: str, -) -> None: - input_texts_images = [(text, None) for text in HF_TEXT_PROMPTS] - input_texts = [text for text, _ in input_texts_images] - input_images = [image for _, image in input_texts_images] + # Should still be able to run subsequent requests + vllm_model.embed(texts) + vllm_model.embed([""], images=images) - _run_test( - hf_runner, - vllm_runner, - input_texts, - input_images, # type: ignore - model, - dtype=dtype, - ) + with hf_runner(model, dtype=dtype, auto_cls=CLIPModel) as hf_model: + hf_outputs_per_case = [] + for input_texts, input_images in input_cases: + all_inputs = hf_model.get_inputs(input_texts, images=input_images) + + hf_outputs = [] + for inputs in all_inputs: + inputs = hf_model.wrap_device(inputs) + + if "pixel_values" in inputs: + pooled_output = hf_model.model.get_image_features( + pixel_values=inputs.pixel_values, + ) + else: + pooled_output = hf_model.model.get_text_features( + input_ids=inputs.input_ids, + attention_mask=inputs.attention_mask, + ) + + if not isinstance(pooled_output, torch.Tensor): + pooled_output = pooled_output.pooler_output + pooled_output = pooled_output.squeeze(0) + hf_outputs.append(pooled_output.tolist()) + + hf_outputs_per_case.append(hf_outputs) + + for hf_outputs, vllm_outputs in zip(hf_outputs_per_case, vllm_outputs_per_case): + check_embeddings_close( + embeddings_0_lst=hf_outputs, + embeddings_1_lst=vllm_outputs, + name_0="hf", + name_1="vllm", + ) @pytest.mark.parametrize("model", MODELS) @pytest.mark.parametrize("dtype", ["float"]) -def test_models_image( +def test_models( hf_runner, vllm_runner, image_assets, model: str, dtype: str, ) -> None: - input_texts_images = [ - (text, asset.pil_image) for text, asset in zip(HF_IMAGE_PROMPTS, image_assets) + text_images = [None] * len(HF_TEXT_PROMPTS) + images = [asset.pil_image for asset in image_assets] + input_cases = [ + (HF_TEXT_PROMPTS, text_images), + (HF_IMAGE_PROMPTS, images), ] - input_texts = [text for text, _ in input_texts_images] - input_images = [image for _, image in input_texts_images] _run_test( hf_runner, vllm_runner, - input_texts, - input_images, + input_cases, # type: ignore[arg-type] model, dtype=dtype, ) - - -@pytest.mark.parametrize("model", MODELS) -@pytest.mark.parametrize("dtype", ["float"]) -def test_models_text_image_no_crash( - vllm_runner, - image_assets, - model: str, - dtype: str, -) -> None: - texts = [HF_TEXT_PROMPTS[0]] - images = [image_assets[0].pil_image] - - with vllm_runner( - model, runner="pooling", dtype=dtype, enforce_eager=True, max_model_len=77 - ) as vllm_model: - with pytest.raises(ValueError, match="not both"): - vllm_model.embed(texts, images=images) - - # Should still be able to run subsequent requests - vllm_model.embed(texts) - vllm_model.embed([""], images=images) diff --git a/tests/models/multimodal/pooling/test_colmodernvbert.py b/tests/models/multimodal/pooling/test_colmodernvbert.py index efeb3195b15b..3dffc66e7a8d 100644 --- a/tests/models/multimodal/pooling/test_colmodernvbert.py +++ b/tests/models/multimodal/pooling/test_colmodernvbert.py @@ -17,29 +17,34 @@ DTYPE = "half" -# ----------------------------------------------------------------------- -# Text-only tests -# ----------------------------------------------------------------------- - - -def test_colmodernvbert_text_token_embed(vllm_runner): - """Text query produces per-token embeddings with shape (seq_len, 128).""" +@pytest.fixture(scope="module") +def colmodernvbert_model(vllm_runner): with vllm_runner( MODEL_NAME, runner="pooling", dtype=DTYPE, enforce_eager=True, ) as vllm_model: - outputs = vllm_model.token_embed(["What is machine learning?"]) + yield vllm_model + + +# ----------------------------------------------------------------------- +# Text-only tests +# ----------------------------------------------------------------------- + - assert len(outputs) == 1 - emb = torch.tensor(outputs[0]) - assert emb.dim() == 2 - assert emb.shape[1] == COLBERT_DIM - assert emb.shape[0] > 1 +def test_colmodernvbert_text_token_embed(colmodernvbert_model): + """Text query produces per-token embeddings with shape (seq_len, 128).""" + outputs = colmodernvbert_model.token_embed(["What is machine learning?"]) + + assert len(outputs) == 1 + emb = torch.tensor(outputs[0]) + assert emb.dim() == 2 + assert emb.shape[1] == COLBERT_DIM + assert emb.shape[0] > 1 -def test_colmodernvbert_text_relevance_ordering(vllm_runner): +def test_colmodernvbert_text_relevance_ordering(colmodernvbert_model): """Relevant documents score higher than irrelevant ones.""" query = "What is machine learning?" documents = [ @@ -47,40 +52,28 @@ def test_colmodernvbert_text_relevance_ordering(vllm_runner): "The weather in Paris is mild in spring.", ] - with vllm_runner( - MODEL_NAME, - runner="pooling", - dtype=DTYPE, - enforce_eager=True, - ) as vllm_model: - scores = vllm_model.score(query, documents) + scores = colmodernvbert_model.score(query, documents) - assert len(scores) == 2 - assert scores[0] > scores[1], "ML doc should score higher than weather doc" + assert len(scores) == 2 + assert scores[0] > scores[1], "ML doc should score higher than weather doc" -def test_colmodernvbert_text_late_interaction(vllm_runner): +def test_colmodernvbert_text_late_interaction(colmodernvbert_model): """MaxSim scoring via vLLM matches manual computation.""" query = "What is the capital of France?" doc = "The capital of France is Paris." - with vllm_runner( - MODEL_NAME, - runner="pooling", - dtype=DTYPE, - enforce_eager=True, - ) as vllm_model: - q_out = vllm_model.token_embed([query]) - d_out = vllm_model.token_embed([doc]) + q_out = colmodernvbert_model.token_embed([query]) + d_out = colmodernvbert_model.token_embed([doc]) - q_emb = torch.tensor(q_out[0]) - d_emb = torch.tensor(d_out[0]) - manual_score = compute_maxsim_score(q_emb, d_emb).item() + q_emb = torch.tensor(q_out[0]) + d_emb = torch.tensor(d_out[0]) + manual_score = compute_maxsim_score(q_emb, d_emb).item() - vllm_scores = vllm_model.score(query, doc) + vllm_scores = colmodernvbert_model.score(query, doc) - assert len(vllm_scores) == 1 - assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01) + assert len(vllm_scores) == 1 + assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01) # ----------------------------------------------------------------------- @@ -88,28 +81,22 @@ def test_colmodernvbert_text_late_interaction(vllm_runner): # ----------------------------------------------------------------------- -def test_colmodernvbert_image_token_embed(vllm_runner, image_assets): +def test_colmodernvbert_image_token_embed(colmodernvbert_model, image_assets): """Image input produces per-token embeddings including vision tokens.""" - with vllm_runner( - MODEL_NAME, - runner="pooling", - dtype=DTYPE, - enforce_eager=True, - ) as vllm_model: - image = image_assets[0].pil_image - inputs = vllm_model.get_inputs( - [""], - images=[image], - ) - req_outputs = vllm_model.llm.encode( - inputs, - pooling_task="token_embed", - ) - outputs = [req_output.outputs.data for req_output in req_outputs] - - assert len(outputs) == 1 - emb = torch.tensor(outputs[0]) - assert emb.dim() == 2 - assert emb.shape[1] == COLBERT_DIM - # Should have at least the image tokens (64 after pixel shuffle) - assert emb.shape[0] >= 64 + image = image_assets[0].pil_image + inputs = colmodernvbert_model.get_inputs( + [""], + images=[image], + ) + req_outputs = colmodernvbert_model.llm.encode( + inputs, + pooling_task="token_embed", + ) + outputs = [req_output.outputs.data for req_output in req_outputs] + + assert len(outputs) == 1 + emb = torch.tensor(outputs[0]) + assert emb.dim() == 2 + assert emb.shape[1] == COLBERT_DIM + # Should have at least the image tokens (64 after pixel shuffle) + assert emb.shape[0] >= 64 diff --git a/tests/models/multimodal/pooling/test_colpali.py b/tests/models/multimodal/pooling/test_colpali.py index 86e595824d9c..4d4e0bd88b21 100644 --- a/tests/models/multimodal/pooling/test_colpali.py +++ b/tests/models/multimodal/pooling/test_colpali.py @@ -75,75 +75,51 @@ def _make_image_mm_param( def _run_token_embed_test( - vllm_runner: type[VllmRunner], + vllm_model: VllmRunner, model: str, - *, - dtype: str, ) -> None: """Verify per-token embedding shape and L2 normalization.""" - with vllm_runner( - model, - runner="pooling", - dtype=dtype, - max_model_len=4096, - enforce_eager=True, - gpu_memory_utilization=GPU_MEMORY_UTILIZATION, - ) as vllm_model: - outputs = vllm_model.token_embed([TEXT_QUERIES[0]]) - - assert len(outputs) == 1 - emb = torch.tensor(outputs[0]) - # Token embeddings should be 2D: [num_tokens, embed_dim] - assert emb.dim() == 2 - assert emb.shape[1] == EMBED_DIMS[model] - assert emb.shape[0] > 1 - - # Verify L2 normalization - norms = torch.norm(emb, p=2, dim=-1) - torch.testing.assert_close( - norms, - torch.ones_like(norms), - rtol=1e-2, - atol=1e-2, - ) + outputs = vllm_model.token_embed([TEXT_QUERIES[0]]) + + assert len(outputs) == 1 + emb = torch.tensor(outputs[0]) + # Token embeddings should be 2D: [num_tokens, embed_dim] + assert emb.dim() == 2 + assert emb.shape[1] == EMBED_DIMS[model] + assert emb.shape[0] > 1 + + # Verify L2 normalization + norms = torch.norm(emb, p=2, dim=-1) + torch.testing.assert_close( + norms, + torch.ones_like(norms), + rtol=1e-2, + atol=1e-2, + ) def _run_late_interaction_test( - vllm_runner: type[VllmRunner], - model: str, - *, - dtype: str, + vllm_model: VllmRunner, ) -> None: """Verify MaxSim scoring matches manual computation.""" from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score - with vllm_runner( - model, - runner="pooling", - dtype=dtype, - max_model_len=4096, - enforce_eager=True, - gpu_memory_utilization=GPU_MEMORY_UTILIZATION, - ) as vllm_model: - q_outputs = vllm_model.token_embed([TEXT_QUERIES[0]]) - d_outputs = vllm_model.token_embed([TEXT_DOCUMENTS[0]]) + q_outputs = vllm_model.token_embed([TEXT_QUERIES[0]]) + d_outputs = vllm_model.token_embed([TEXT_DOCUMENTS[0]]) - q_emb = torch.tensor(q_outputs[0]) - d_emb = torch.tensor(d_outputs[0]) + q_emb = torch.tensor(q_outputs[0]) + d_emb = torch.tensor(d_outputs[0]) - manual_score = compute_maxsim_score(q_emb, d_emb).item() + manual_score = compute_maxsim_score(q_emb, d_emb).item() - vllm_scores = vllm_model.score(TEXT_QUERIES[0], TEXT_DOCUMENTS[0]) + vllm_scores = vllm_model.score(TEXT_QUERIES[0], TEXT_DOCUMENTS[0]) - assert len(vllm_scores) == 1 - assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01) + assert len(vllm_scores) == 1 + assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01) def _run_relevance_test( - vllm_runner: type[VllmRunner], - model: str, - *, - dtype: str, + vllm_model: VllmRunner, ) -> None: """Verify that relevant documents score higher than irrelevant ones.""" query = "What is machine learning?" @@ -153,59 +129,18 @@ def _run_relevance_test( "Deep learning uses neural networks for complex tasks.", ] - with vllm_runner( - model, - runner="pooling", - dtype=dtype, - max_model_len=4096, - enforce_eager=True, - gpu_memory_utilization=GPU_MEMORY_UTILIZATION, - ) as vllm_model: - scores = vllm_model.score(query, documents) + scores = vllm_model.score(query, documents) - assert len(scores) == 3 - assert scores[0] > scores[1], "ML doc should score higher than weather doc" - assert scores[2] > scores[1], "DL doc should score higher than weather doc" - - -@pytest.mark.parametrize("model", MODELS) -@pytest.mark.parametrize("dtype", [DTYPE]) -def test_colpali_token_embed( - vllm_runner, - model: str, - dtype: str, -) -> None: - _run_token_embed_test(vllm_runner, model, dtype=dtype) - - -@pytest.mark.parametrize("model", MODELS) -@pytest.mark.parametrize("dtype", [DTYPE]) -def test_colpali_late_interaction_scoring( - vllm_runner, - model: str, - dtype: str, -) -> None: - _run_late_interaction_test(vllm_runner, model, dtype=dtype) - - -@pytest.mark.parametrize("model", MODELS) -@pytest.mark.parametrize("dtype", [DTYPE]) -def test_colpali_relevance_ordering( - vllm_runner, - model: str, - dtype: str, -) -> None: - _run_relevance_test(vllm_runner, model, dtype=dtype) + assert len(scores) == 3 + assert scores[0] > scores[1], "ML doc should score higher than weather doc" + assert scores[2] > scores[1], "DL doc should score higher than weather doc" # ── Multimodal scoring tests ──────────────────────────────── def _run_multimodal_text_query_image_docs_test( - vllm_runner: type[VllmRunner], - model: str, - *, - dtype: str, + vllm_model: VllmRunner, ) -> None: """Score a text query against image documents via the multimodal path.""" red_image = _make_base64_image(64, 64, color=(255, 0, 0)) @@ -218,29 +153,15 @@ def _run_multimodal_text_query_image_docs_test( ] attention_backend = "FLASH_ATTN" if current_platform.is_cuda() else None - with vllm_runner( - model, - runner="pooling", - dtype=dtype, - max_model_len=4096, - enforce_eager=True, - gpu_memory_utilization=GPU_MEMORY_UTILIZATION, - attention_backend=attention_backend, - kernel_config={"enable_flashinfer_autotune": False}, - ) as vllm_model: - assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner - scores = vllm_model.llm.score(query, image_docs) + scores = vllm_model.llm.score(query, image_docs) - assert len(scores) == 2 - for s in scores: - assert isinstance(s.outputs.score, float) + assert len(scores) == 2 + for s in scores: + assert isinstance(s.outputs.score, float) def _run_multimodal_mixed_docs_test( - vllm_runner: type[VllmRunner], - model: str, - *, - dtype: str, + vllm_model: VllmRunner, ) -> None: """Score a text query against a mix of text and image documents.""" red_image = _make_base64_image(64, 64, color=(255, 0, 0)) @@ -251,28 +172,17 @@ def _run_multimodal_mixed_docs_test( _make_image_mm_param(red_image), ] - with vllm_runner( - model, - runner="pooling", - dtype=dtype, - max_model_len=4096, - enforce_eager=True, - gpu_memory_utilization=GPU_MEMORY_UTILIZATION, - ) as vllm_model: - scores = vllm_model.llm.score(query, documents) + scores = vllm_model.llm.score(query, documents) - assert len(scores) == 2 - for s in scores: - assert isinstance(s.outputs.score, float) - # Text document about France should score higher than a random image - assert scores[0].outputs.score > scores[1].outputs.score + assert len(scores) == 2 + for s in scores: + assert isinstance(s.outputs.score, float) + # Text document about France should score higher than a random image + assert scores[0].outputs.score > scores[1].outputs.score def _run_multimodal_image_query_text_docs_test( - vllm_runner: type[VllmRunner], - model: str, - *, - dtype: str, + vllm_model: VllmRunner, ) -> None: """Score an image query against text documents.""" red_image = _make_base64_image(64, 64, color=(255, 0, 0)) @@ -283,6 +193,20 @@ def _run_multimodal_image_query_text_docs_test( "The weather forecast shows rain tomorrow.", ] + scores = vllm_model.llm.score(image_query, documents) + + assert len(scores) == 2 + for s in scores: + assert isinstance(s.outputs.score, float) + + +@pytest.mark.parametrize("model", MODELS) +@pytest.mark.parametrize("dtype", [DTYPE]) +def test_colpali_default_runner( + vllm_runner, + model: str, + dtype: str, +) -> None: with vllm_runner( model, runner="pooling", @@ -291,40 +215,31 @@ def _run_multimodal_image_query_text_docs_test( enforce_eager=True, gpu_memory_utilization=GPU_MEMORY_UTILIZATION, ) as vllm_model: - scores = vllm_model.llm.score(image_query, documents) - - assert len(scores) == 2 - for s in scores: - assert isinstance(s.outputs.score, float) + _run_token_embed_test(vllm_model, model) + _run_late_interaction_test(vllm_model) + _run_relevance_test(vllm_model) + _run_multimodal_mixed_docs_test(vllm_model) + _run_multimodal_image_query_text_docs_test(vllm_model) @pytest.mark.parametrize("model", MODELS) @pytest.mark.parametrize("dtype", [DTYPE]) -def test_colpali_multimodal_text_query_image_docs( +def test_colpali_v2_multimodal_text_query_image_docs( vllm_runner, monkeypatch: pytest.MonkeyPatch, model: str, dtype: str, ) -> None: monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") - _run_multimodal_text_query_image_docs_test(vllm_runner, model, dtype=dtype) - - -@pytest.mark.parametrize("model", MODELS) -@pytest.mark.parametrize("dtype", [DTYPE]) -def test_colpali_multimodal_mixed_docs( - vllm_runner, - model: str, - dtype: str, -) -> None: - _run_multimodal_mixed_docs_test(vllm_runner, model, dtype=dtype) - - -@pytest.mark.parametrize("model", MODELS) -@pytest.mark.parametrize("dtype", [DTYPE]) -def test_colpali_multimodal_image_query_text_docs( - vllm_runner, - model: str, - dtype: str, -) -> None: - _run_multimodal_image_query_text_docs_test(vllm_runner, model, dtype=dtype) + with vllm_runner( + model, + runner="pooling", + dtype=dtype, + max_model_len=4096, + enforce_eager=True, + gpu_memory_utilization=GPU_MEMORY_UTILIZATION, + attention_backend=attention_backend, + kernel_config={"enable_flashinfer_autotune": False}, + ) as vllm_model: + assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner + _run_multimodal_text_query_image_docs_test(vllm_model) diff --git a/tests/models/multimodal/pooling/test_colqwen3_5.py b/tests/models/multimodal/pooling/test_colqwen3_5.py index 43914d819b8a..2aac465d4d0a 100644 --- a/tests/models/multimodal/pooling/test_colqwen3_5.py +++ b/tests/models/multimodal/pooling/test_colqwen3_5.py @@ -35,74 +35,65 @@ DTYPE = "half" -def _run_token_embed_test( - vllm_runner: type[VllmRunner], - model: str, - *, - dtype: str, -) -> None: - """Verify per-token embedding shape and L2 normalization.""" +@pytest.fixture(scope="module", params=MODELS) +def colqwen3_5_model(request, vllm_runner): + model = request.param with vllm_runner( model, runner="pooling", - dtype=dtype, + dtype=DTYPE, max_model_len=4096, enforce_eager=True, ) as vllm_model: - outputs = vllm_model.token_embed([TEXT_QUERIES[0]]) + yield model, vllm_model - assert len(outputs) == 1 - emb = torch.tensor(outputs[0]) - # Token embeddings should be 2D: [num_tokens, embed_dim] - assert emb.dim() == 2 - assert emb.shape[1] == EMBED_DIMS[model] - assert emb.shape[0] > 1 - # Verify L2 normalization - norms = torch.norm(emb, p=2, dim=-1) - torch.testing.assert_close( - norms, - torch.ones_like(norms), - rtol=1e-2, - atol=1e-2, - ) +def _run_token_embed_test( + vllm_model: VllmRunner, + model: str, +) -> None: + """Verify per-token embedding shape and L2 normalization.""" + outputs = vllm_model.token_embed([TEXT_QUERIES[0]]) + + assert len(outputs) == 1 + emb = torch.tensor(outputs[0]) + # Token embeddings should be 2D: [num_tokens, embed_dim] + assert emb.dim() == 2 + assert emb.shape[1] == EMBED_DIMS[model] + assert emb.shape[0] > 1 + + # Verify L2 normalization + norms = torch.norm(emb, p=2, dim=-1) + torch.testing.assert_close( + norms, + torch.ones_like(norms), + rtol=1e-2, + atol=1e-2, + ) def _run_late_interaction_test( - vllm_runner: type[VllmRunner], - model: str, - *, - dtype: str, + vllm_model: VllmRunner, ) -> None: """Verify MaxSim scoring matches manual computation.""" from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score - with vllm_runner( - model, - runner="pooling", - dtype=dtype, - max_model_len=4096, - enforce_eager=True, - ) as vllm_model: - q_outputs = vllm_model.token_embed([TEXT_QUERIES[0]]) - d_outputs = vllm_model.token_embed([TEXT_DOCUMENTS[0]]) + q_outputs = vllm_model.token_embed([TEXT_QUERIES[0]]) + d_outputs = vllm_model.token_embed([TEXT_DOCUMENTS[0]]) - q_emb = torch.tensor(q_outputs[0]) - d_emb = torch.tensor(d_outputs[0]) + q_emb = torch.tensor(q_outputs[0]) + d_emb = torch.tensor(d_outputs[0]) - manual_score = compute_maxsim_score(q_emb, d_emb).item() + manual_score = compute_maxsim_score(q_emb, d_emb).item() - vllm_scores = vllm_model.score(TEXT_QUERIES[0], TEXT_DOCUMENTS[0]) + vllm_scores = vllm_model.score(TEXT_QUERIES[0], TEXT_DOCUMENTS[0]) - assert len(vllm_scores) == 1 - assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01) + assert len(vllm_scores) == 1 + assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01) def _run_relevance_test( - vllm_runner: type[VllmRunner], - model: str, - *, - dtype: str, + vllm_model: VllmRunner, ) -> None: """Verify that relevant documents score higher than irrelevant ones.""" query = "What is machine learning?" @@ -112,48 +103,26 @@ def _run_relevance_test( "Deep learning uses neural networks for complex tasks.", ] - with vllm_runner( - model, - runner="pooling", - dtype=dtype, - max_model_len=4096, - enforce_eager=True, - ) as vllm_model: - scores = vllm_model.score(query, documents) + scores = vllm_model.score(query, documents) - assert len(scores) == 3 - assert scores[0] > scores[1], "ML doc should score higher than weather doc" - assert scores[2] > scores[1], "DL doc should score higher than weather doc" + assert len(scores) == 3 + assert scores[0] > scores[1], "ML doc should score higher than weather doc" + assert scores[2] > scores[1], "DL doc should score higher than weather doc" -@pytest.mark.parametrize("model", MODELS) -@pytest.mark.parametrize("dtype", [DTYPE]) -def test_colqwen3_5_token_embed( - vllm_runner, - model: str, - dtype: str, -) -> None: - _run_token_embed_test(vllm_runner, model, dtype=dtype) +def test_colqwen3_5_token_embed(colqwen3_5_model) -> None: + model, vllm_model = colqwen3_5_model + _run_token_embed_test(vllm_model, model) -@pytest.mark.parametrize("model", MODELS) -@pytest.mark.parametrize("dtype", [DTYPE]) -def test_colqwen3_5_late_interaction_scoring( - vllm_runner, - model: str, - dtype: str, -) -> None: - _run_late_interaction_test(vllm_runner, model, dtype=dtype) +def test_colqwen3_5_late_interaction_scoring(colqwen3_5_model) -> None: + _, vllm_model = colqwen3_5_model + _run_late_interaction_test(vllm_model) -@pytest.mark.parametrize("model", MODELS) -@pytest.mark.parametrize("dtype", [DTYPE]) -def test_colqwen3_5_relevance_ordering( - vllm_runner, - model: str, - dtype: str, -) -> None: - _run_relevance_test(vllm_runner, model, dtype=dtype) +def test_colqwen3_5_relevance_ordering(colqwen3_5_model) -> None: + _, vllm_model = colqwen3_5_model + _run_relevance_test(vllm_model) @pytest.mark.parametrize( diff --git a/tests/models/multimodal/pooling/test_llama_nemotron_vl.py b/tests/models/multimodal/pooling/test_llama_nemotron_vl.py index a2f1d3424c34..9516ab95d52d 100644 --- a/tests/models/multimodal/pooling/test_llama_nemotron_vl.py +++ b/tests/models/multimodal/pooling/test_llama_nemotron_vl.py @@ -9,12 +9,14 @@ Both variants share a SigLIP vision encoder with a bidirectional LLaMA backbone. """ +from collections.abc import Sequence from io import BytesIO from pathlib import Path import pybase64 as base64 import pytest import torch +from PIL import Image from transformers import AutoModel, AutoModelForSequenceClassification, AutoProcessor from vllm.entrypoints.chat_utils import ( @@ -54,17 +56,15 @@ def _run_test( hf_runner: type[HfRunner], vllm_runner: type[VllmRunner], - input_texts: list[str], - input_images: PromptImageInput, + input_cases: list[tuple[list[str], PromptImageInput]], model: str, *, dtype: str, ) -> None: - """Run embedding comparison test between HF and vLLM. + """Compare HF and vLLM embeddings for all input cases. NOTE: Run vLLM first to avoid CUDA initialization issues with multiprocessing. """ - # Run vLLM inference first with vllm_runner( model, runner="pooling", @@ -74,91 +74,70 @@ def _run_test( trust_remote_code=True, **ROCM_ENGINE_KWARGS, ) as vllm_model: - vllm_outputs = vllm_model.embed(input_texts, images=input_images) + vllm_outputs_per_case = [ + vllm_model.embed(input_texts, images=input_images) + for input_texts, input_images in input_cases + ] - # Run HF inference using the model's encode_queries/encode_documents API with hf_runner(model, dtype=dtype, auto_cls=AutoModel) as hf_model: - hf_outputs = [] - for text, image in zip(input_texts, input_images): - with torch.inference_mode(): - if text.startswith(QUERY_PREFIX): - # Strip prefix and use encode_queries for query texts - query_text = text[len(QUERY_PREFIX) :] - embedding = hf_model.model.encode_queries([query_text]) - elif text.startswith(PASSAGE_PREFIX): - # Strip prefix and use encode_documents for passages/images - passage_text = text[len(PASSAGE_PREFIX) :] - if image is not None: - # Image document - pass image to encode_documents - embedding = hf_model.model.encode_documents( - images=[image], - texts=[passage_text], - ) + hf_outputs_per_case = [] + for input_texts, input_images in input_cases: + hf_outputs = [] + for text, image in zip(input_texts, input_images): + with torch.inference_mode(): + if text.startswith(QUERY_PREFIX): + query_text = text[len(QUERY_PREFIX) :] + embedding = hf_model.model.encode_queries([query_text]) + elif text.startswith(PASSAGE_PREFIX): + passage_text = text[len(PASSAGE_PREFIX) :] + if image is not None: + embedding = hf_model.model.encode_documents( + images=[image], + texts=[passage_text], + ) + else: + embedding = hf_model.model.encode_documents( + texts=[passage_text] + ) else: - # Text-only document - embedding = hf_model.model.encode_documents( - texts=[passage_text] + raise ValueError( + f"Text must start with {QUERY_PREFIX!r} " + f"or {PASSAGE_PREFIX!r}" ) - else: - raise ValueError( - f"Text must start with '{QUERY_PREFIX}' or '{PASSAGE_PREFIX}'" - ) - - hf_outputs.append(embedding[0].tolist()) - - check_embeddings_close( - embeddings_0_lst=hf_outputs, - embeddings_1_lst=vllm_outputs, - name_0="hf", - name_1="vllm", - ) - -@pytest.mark.parametrize("model", MODELS) -@pytest.mark.parametrize("dtype", ["half"]) -def test_models_text( - hf_runner, - vllm_runner, - image_assets, - model: str, - dtype: str, -) -> None: - """Test text-only embedding.""" - input_texts_images = [(text, None) for text in HF_TEXT_PROMPTS] - input_texts = [text for text, _ in input_texts_images] - input_images = [image for _, image in input_texts_images] + hf_outputs.append(embedding[0].tolist()) + hf_outputs_per_case.append(hf_outputs) - _run_test( - hf_runner, - vllm_runner, - input_texts, - input_images, # type: ignore - model, - dtype=dtype, - ) + for hf_outputs, vllm_outputs in zip(hf_outputs_per_case, vllm_outputs_per_case): + check_embeddings_close( + embeddings_0_lst=hf_outputs, + embeddings_1_lst=vllm_outputs, + name_0="hf", + name_1="vllm", + ) @pytest.mark.parametrize("model", MODELS) @pytest.mark.parametrize("dtype", ["half"]) -def test_models_image( +def test_models( hf_runner, vllm_runner, image_assets, model: str, dtype: str, ) -> None: - """Test image embedding.""" - input_texts_images = [ - (text, asset.pil_image) for text, asset in zip(HF_IMAGE_PROMPTS, image_assets) + """Test text and image embedding.""" + text_images = [None] * len(HF_TEXT_PROMPTS) + images = [asset.pil_image for asset in image_assets] + input_cases = [ + (HF_TEXT_PROMPTS, text_images), + (HF_IMAGE_PROMPTS, images), ] - input_texts = [text for text, _ in input_texts_images] - input_images = [image for _, image in input_texts_images] _run_test( hf_runner, vllm_runner, - input_texts, - input_images, + input_cases, # type: ignore[arg-type] model, dtype=dtype, ) @@ -188,8 +167,11 @@ def test_models_image( RERANKER_IMAGE_QUERY = "photo of a red stop sign on a street" +RerankerDocument = tuple[str | None, Image.Image | None] +RerankerCase = tuple[str, Sequence[RerankerDocument]] + -def _pil_to_data_uri(image) -> str: +def _pil_to_data_uri(image: Image.Image) -> str: buf = BytesIO() image.save(buf, format="PNG") b64 = base64.b64encode(buf.getvalue()).decode() @@ -200,10 +182,9 @@ def _run_hf_reranker( hf_runner: type[HfRunner], model: str, dtype: str, - query: str, - docs: list, -) -> list[float]: - """Run HF reranker inference; docs is a list of (doc_text, doc_image|None).""" + input_cases: Sequence[RerankerCase], +) -> list[list[float]]: + """Run all HF reranker cases in one model lifecycle.""" with hf_runner( model, dtype=dtype, @@ -217,35 +198,37 @@ def _run_hf_reranker( use_thumbnail=True, rerank_max_length=2048, ) - examples = [ - { - "question": query, - "doc_text": doc_text if doc_text is not None else "", - "doc_image": doc_image if doc_image is not None else "", + scores_per_case = [] + for query, docs in input_cases: + examples = [ + { + "question": query, + "doc_text": doc_text if doc_text is not None else "", + "doc_image": doc_image if doc_image is not None else "", + } + for doc_text, doc_image in docs + ] + batch_dict = processor.process_queries_documents_crossencoder(examples) + batch_dict = { + k: v.to(hf_model.model.device) if isinstance(v, torch.Tensor) else v + for k, v in batch_dict.items() } - for doc_text, doc_image in docs - ] - batch_dict = processor.process_queries_documents_crossencoder(examples) - batch_dict = { - k: v.to(hf_model.model.device) if isinstance(v, torch.Tensor) else v - for k, v in batch_dict.items() - } - with torch.inference_mode(): - logits = hf_model.model(**batch_dict, return_dict=True).logits - # vLLM applies sigmoid activation to the raw logits before returning - # scores; apply the same here so both sides are comparable. - scores = torch.sigmoid(logits.squeeze(-1).float()) - return scores.detach().cpu().tolist() + with torch.inference_mode(): + logits = hf_model.model(**batch_dict, return_dict=True).logits + # vLLM applies sigmoid activation to raw logits before returning scores. + scores = torch.sigmoid(logits.squeeze(-1).float()) + scores_per_case.append(scores.detach().cpu().tolist()) + + return scores_per_case def _run_vllm_reranker( vllm_runner: type[VllmRunner], model: str, dtype: str, - query: str, - docs: list, -) -> list[float]: - """Run vLLM reranker inference; docs is a list of (doc_text, doc_image|None).""" + input_cases: Sequence[RerankerCase], +) -> list[list[float]]: + """Run all vLLM reranker cases in one model lifecycle.""" with vllm_runner( model, runner="pooling", @@ -255,57 +238,59 @@ def _run_vllm_reranker( trust_remote_code=True, **ROCM_ENGINE_KWARGS, ) as vllm_model: - has_images = any(img is not None for _, img in docs) - - if not has_images: - # Text-only path: use the simple string score API. - queries = [query] * len(docs) - doc_texts = [doc_text for doc_text, _ in docs] - outputs = vllm_model.score( - queries, - doc_texts, - chat_template=_RERANKER_SCORE_TEMPLATE, - ) - else: - # Multimodal path: build ScoreMultiModalParam for each pair. - query_params = [ - ScoreMultiModalParam( - content=[ - ChatCompletionContentPartTextParam( - type="text", - text=query, - ) - ] + scores_per_case = [] + for query, docs in input_cases: + has_images = any(img is not None for _, img in docs) + + if not has_images: + queries = [query] * len(docs) + doc_texts = [doc_text for doc_text, _ in docs] + outputs = vllm_model.score( + queries, + doc_texts, + chat_template=_RERANKER_SCORE_TEMPLATE, ) - ] * len(docs) - - doc_params = [] - for doc_text, doc_image in docs: - content: list = [] - if doc_image is not None: - content.append( - ChatCompletionContentPartImageParam( - type="image_url", - image_url={"url": _pil_to_data_uri(doc_image)}, - ) + else: + query_params = [ + ScoreMultiModalParam( + content=[ + ChatCompletionContentPartTextParam( + type="text", + text=query, + ) + ] ) - if doc_text: - content.append( - ChatCompletionContentPartTextParam( - type="text", - text=doc_text, + ] * len(docs) + + doc_params = [] + for doc_text, doc_image in docs: + content: list = [] + if doc_image is not None: + content.append( + ChatCompletionContentPartImageParam( + type="image_url", + image_url={"url": _pil_to_data_uri(doc_image)}, + ) ) - ) - doc_params.append(ScoreMultiModalParam(content=content)) + if doc_text: + content.append( + ChatCompletionContentPartTextParam( + type="text", + text=doc_text, + ) + ) + doc_params.append(ScoreMultiModalParam(content=content)) - raw_outputs = vllm_model.llm.score( - query_params, - doc_params, - chat_template=_RERANKER_SCORE_TEMPLATE, - ) - outputs = [o.outputs.score for o in raw_outputs] + raw_outputs = vllm_model.llm.score( + query_params, + doc_params, + chat_template=_RERANKER_SCORE_TEMPLATE, + ) + outputs = [output.outputs.score for output in raw_outputs] + + scores_per_case.append(outputs) - return outputs + return scores_per_case def _run_reranker_test( @@ -313,50 +298,44 @@ def _run_reranker_test( vllm_runner: type[VllmRunner], model: str, dtype: str, - query: str, - docs: list, + input_cases: Sequence[RerankerCase], ) -> None: - """Compare HF and vLLM reranker scores. + """Compare HF and vLLM reranker scores for all input cases. NOTE: Run vLLM first to avoid CUDA initialization issues with multiprocessing. """ - vllm_scores = _run_vllm_reranker(vllm_runner, model, dtype, query, docs) - hf_scores = _run_hf_reranker(hf_runner, model, dtype, query, docs) + vllm_scores_per_case = _run_vllm_reranker(vllm_runner, model, dtype, input_cases) + hf_scores_per_case = _run_hf_reranker(hf_runner, model, dtype, input_cases) - assert len(hf_scores) == len(vllm_scores), ( - f"Output length mismatch: HF={len(hf_scores)}, vLLM={len(vllm_scores)}" - ) - # NOTE: ROCm shows slightly higher numerical variance dues to different attention - # backend between vLLM and HF; use a marginally looser tolerance + # ROCm has slightly higher variance because vLLM and HF use different + # attention backends. rel_tol = 0.022 if current_platform.is_rocm() else 0.02 - for i, (hf_score, vllm_score) in enumerate(zip(hf_scores, vllm_scores)): - assert hf_score == pytest.approx(vllm_score, rel=rel_tol), ( - f"Score mismatch at index {i}: HF={hf_score:.4f}, vLLM={vllm_score:.4f}" + for hf_scores, vllm_scores in zip(hf_scores_per_case, vllm_scores_per_case): + assert len(hf_scores) == len(vllm_scores), ( + f"Output length mismatch: HF={len(hf_scores)}, vLLM={len(vllm_scores)}" ) + for i, (hf_score, vllm_score) in enumerate(zip(hf_scores, vllm_scores)): + assert hf_score == pytest.approx(vllm_score, rel=rel_tol), ( + f"Score mismatch at index {i}: HF={hf_score:.4f}, vLLM={vllm_score:.4f}" + ) @pytest.mark.parametrize("model", RERANKER_MODELS) @pytest.mark.parametrize("dtype", ["half"]) -def test_reranker_text( - hf_runner, - vllm_runner, - model: str, - dtype: str, -) -> None: - """Test reranking with text-only query and text documents.""" - docs = [(text, None) for text in RERANKER_TEXT_DOCS] - _run_reranker_test(hf_runner, vllm_runner, model, dtype, RERANKER_TEXT_QUERY, docs) - - -@pytest.mark.parametrize("model", RERANKER_MODELS) -@pytest.mark.parametrize("dtype", ["half"]) -def test_reranker_image_doc( +def test_reranker( hf_runner, vllm_runner, image_assets, model: str, dtype: str, ) -> None: - """Test reranking with text query against image documents.""" - docs = [(None, asset.pil_image) for asset in image_assets] - _run_reranker_test(hf_runner, vllm_runner, model, dtype, RERANKER_IMAGE_QUERY, docs) + """Test reranking with text and image documents.""" + text_docs: list[RerankerDocument] = [(text, None) for text in RERANKER_TEXT_DOCS] + image_docs: list[RerankerDocument] = [ + (None, asset.pil_image) for asset in image_assets + ] + input_cases: list[RerankerCase] = [ + (RERANKER_TEXT_QUERY, text_docs), + (RERANKER_IMAGE_QUERY, image_docs), + ] + _run_reranker_test(hf_runner, vllm_runner, model, dtype, input_cases) diff --git a/tests/models/multimodal/pooling/test_siglip.py b/tests/models/multimodal/pooling/test_siglip.py index bca598b42c64..8eeb594db226 100644 --- a/tests/models/multimodal/pooling/test_siglip.py +++ b/tests/models/multimodal/pooling/test_siglip.py @@ -33,16 +33,11 @@ def _run_test( hf_runner: type[HfRunner], vllm_runner: type[VllmRunner], - input_texts: list[str], - input_images: PromptImageInput, + input_cases: list[tuple[list[str], PromptImageInput, dict[str, Any]]], model: str, *, dtype: str, - tokenization_kwargs: dict[str, Any] | None = None, ) -> None: - if tokenization_kwargs is None: - tokenization_kwargs = {} - with vllm_runner( model, runner="pooling", @@ -51,116 +46,88 @@ def _run_test( max_model_len=64, gpu_memory_utilization=0.7, ) as vllm_model: - vllm_outputs = vllm_model.embed( - input_texts, images=input_images, tokenization_kwargs=tokenization_kwargs - ) + vllm_outputs_per_case = [ + vllm_model.embed( + input_texts, + images=input_images, + tokenization_kwargs=tokenization_kwargs, + ) + for input_texts, input_images, tokenization_kwargs in input_cases + ] + + texts = [HF_TEXT_PROMPTS[0]] + images = [input_cases[1][1][0]] + with pytest.raises(ValueError, match="not both"): + vllm_model.embed(texts, images=images) + + vllm_model.embed(texts) + vllm_model.embed([""], images=images) with hf_runner(model, dtype=dtype, auto_cls=SiglipModel) as hf_model: - all_inputs = hf_model.get_inputs( - input_texts, images=input_images, tokenization_kwargs=tokenization_kwargs + hf_outputs_per_case = [] + for input_texts, input_images, tokenization_kwargs in input_cases: + all_inputs = hf_model.get_inputs( + input_texts, + images=input_images, + tokenization_kwargs=tokenization_kwargs, + ) + + hf_outputs = [] + for inputs in all_inputs: + inputs = hf_model.wrap_device(inputs) + + if "pixel_values" in inputs: + pooled_output = hf_model.model.get_image_features( + pixel_values=inputs.pixel_values, + ) + else: + pooled_output = hf_model.model.get_text_features( + input_ids=inputs.input_ids, + ) + + if not isinstance(pooled_output, torch.Tensor): + pooled_output = pooled_output.pooler_output + pooled_output = pooled_output.squeeze(0) + hf_outputs.append(pooled_output.tolist()) + + hf_outputs_per_case.append(hf_outputs) + + for hf_outputs, vllm_outputs in zip(hf_outputs_per_case, vllm_outputs_per_case): + check_embeddings_close( + embeddings_0_lst=hf_outputs, + embeddings_1_lst=vllm_outputs, + name_0="hf", + name_1="vllm", ) - all_outputs = [] - for inputs in all_inputs: - inputs = hf_model.wrap_device(inputs) - - if "pixel_values" in inputs: - pooled_output = hf_model.model.get_image_features( - pixel_values=inputs.pixel_values, - ) - else: - pooled_output = hf_model.model.get_text_features( - input_ids=inputs.input_ids, - ) - - if not isinstance(pooled_output, torch.Tensor): - pooled_output = pooled_output.pooler_output - pooled_output = pooled_output.squeeze(0) - all_outputs.append(pooled_output.tolist()) - - hf_outputs = all_outputs - - check_embeddings_close( - embeddings_0_lst=hf_outputs, - embeddings_1_lst=vllm_outputs, - name_0="hf", - name_1="vllm", - ) - @pytest.mark.parametrize("model", MODELS) @pytest.mark.parametrize("dtype", ["float"]) -def test_models_text( +def test_models( hf_runner, vllm_runner, image_assets, model: str, dtype: str, ) -> None: - input_texts_images = [(text, None) for text in HF_TEXT_PROMPTS] - input_texts = [text for text, _ in input_texts_images] - input_images = [image for _, image in input_texts_images] - - _run_test( - hf_runner, - vllm_runner, - input_texts, - input_images, # type: ignore - model, - dtype=dtype, - tokenization_kwargs={ - "padding": "max_length", - "max_length": 64, - }, # siglip2 was trained with this padding setting. - ) - - -@pytest.mark.parametrize("model", MODELS) -@pytest.mark.parametrize("dtype", ["float"]) -def test_models_image( - hf_runner, - vllm_runner, - image_assets, - model: str, - dtype: str, -) -> None: - input_texts_images = [ - (text, asset.pil_image) for text, asset in zip(HF_IMAGE_PROMPTS, image_assets) + text_images = [None] * len(HF_TEXT_PROMPTS) + images = [asset.pil_image for asset in image_assets] + input_cases = [ + ( + HF_TEXT_PROMPTS, + text_images, + { + "padding": "max_length", + "max_length": 64, + }, + ), + (HF_IMAGE_PROMPTS, images, {}), ] - input_texts = [text for text, _ in input_texts_images] - input_images = [image for _, image in input_texts_images] _run_test( hf_runner, vllm_runner, - input_texts, - input_images, + input_cases, # type: ignore[arg-type] model, dtype=dtype, ) - - -@pytest.mark.parametrize("model", MODELS) -@pytest.mark.parametrize("dtype", ["float"]) -def test_models_text_image_no_crash( - vllm_runner, - image_assets, - model: str, - dtype: str, -) -> None: - texts = [HF_TEXT_PROMPTS[0]] - images = [image_assets[0].pil_image] - - with vllm_runner( - model, - runner="pooling", - dtype=dtype, - enforce_eager=True, - max_model_len=64, - gpu_memory_utilization=0.7, - ) as vllm_model: - with pytest.raises(ValueError, match="not both"): - vllm_model.embed(texts, images=images) - - vllm_model.embed(texts) - vllm_model.embed([""], images=images) diff --git a/tests/models/quantization/test_awq.py b/tests/models/quantization/test_awq.py index 25a63f6bd907..38f60be6bfca 100644 --- a/tests/models/quantization/test_awq.py +++ b/tests/models/quantization/test_awq.py @@ -17,6 +17,15 @@ } ) +IMAGE_SIZE_FACTOR_GROUPS = ( + # Single-scale + (1.0,), + # Single-scale, batched + (1.0, 1.0, 1.0), + # Multi-scale + (0.25, 0.5, 1.0), +) + def run_awq_test( vllm_runner: type[VllmRunner], @@ -24,7 +33,7 @@ def run_awq_test( source_model: str, quant_model: str, *, - size_factors: list[float], + size_factor_groups: tuple[tuple[float, ...], ...], dtype: str, max_tokens: int, num_logprobs: int, @@ -33,11 +42,12 @@ def run_awq_test( ): images = [asset.pil_image for asset in image_assets] - inputs_per_image = [ + inputs_per_image_and_size_group = [ ( [prompt for _ in size_factors], [rescale_image_size(image, factor) for factor in size_factors], ) + for size_factors in size_factor_groups for image, prompt in zip(images, HF_IMAGE_PROMPTS) ] @@ -60,7 +70,7 @@ def run_awq_test( vllm_model.generate_greedy_logprobs( prompts, max_tokens, num_logprobs=num_logprobs, images=images ) - for prompts, images in inputs_per_image + for prompts, images in inputs_per_image_and_size_group ] with vllm_runner( @@ -77,7 +87,7 @@ def run_awq_test( vllm_model.generate_greedy_logprobs( prompts, max_tokens, num_logprobs=num_logprobs, images=images ) - for prompts, images in inputs_per_image + for prompts, images in inputs_per_image_and_size_group ] for source_outputs, quant_outputs in zip( @@ -128,17 +138,6 @@ def test_awq_load( ("source_model", "quant_model"), [("OpenGVLab/InternVL2-2B", "OpenGVLab/InternVL2-2B-AWQ")], ) -@pytest.mark.parametrize( - "size_factors", - [ - # Single-scale - [1.0], - # Single-scale, batched - [1.0, 1.0, 1.0], - # Multi-scale - [0.25, 0.5, 1.0], - ], -) @pytest.mark.parametrize("dtype", ["half"]) @pytest.mark.parametrize("max_tokens", [128]) @pytest.mark.parametrize("num_logprobs", [5]) @@ -148,7 +147,6 @@ def test_awq_models( image_assets, source_model, quant_model, - size_factors, dtype, max_tokens, num_logprobs, @@ -158,7 +156,7 @@ def test_awq_models( image_assets, source_model, quant_model, - size_factors=size_factors, + size_factor_groups=IMAGE_SIZE_FACTOR_GROUPS, dtype=dtype, max_tokens=max_tokens, num_logprobs=num_logprobs, From 93550cc4cd7e97bebe47b9b57333b60921c0d662 Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Mon, 17 Aug 2026 14:17:15 +0800 Subject: [PATCH 036/839] [Frontend] Consolidate entrypoint middleware (#52309) Signed-off-by: wang.yuqi --- .../entrypoints/serve/middleware/__init__.py | 0 .../test_authentication_middleware.py | 150 +++++++++++ .../test_optional_middleware.py | 0 vllm/entrypoints/launchers/__init__.py | 0 .../launchers/api_server/__init__.py | 0 .../launchers/api_server/routers.py | 74 ++++++ vllm/entrypoints/openai/api_server.py | 119 +-------- vllm/entrypoints/serve/middleware/__init__.py | 0 .../serve/middleware/authenticate.py | 62 +++++ .../serve/middleware/log_response.py | 162 +++++++++++ vllm/entrypoints/serve/middleware/register.py | 75 ++++++ .../serve/middleware/x_request_id.py | 38 +++ vllm/entrypoints/serve/utils/server_utils.py | 251 +----------------- vllm/entrypoints/speech_to_text/factories.py | 6 - 14 files changed, 567 insertions(+), 370 deletions(-) create mode 100644 tests/entrypoints/serve/middleware/__init__.py create mode 100644 tests/entrypoints/serve/middleware/test_authentication_middleware.py rename tests/entrypoints/serve/{instrumentator => middleware}/test_optional_middleware.py (100%) create mode 100644 vllm/entrypoints/launchers/__init__.py create mode 100644 vllm/entrypoints/launchers/api_server/__init__.py create mode 100644 vllm/entrypoints/launchers/api_server/routers.py create mode 100644 vllm/entrypoints/serve/middleware/__init__.py create mode 100644 vllm/entrypoints/serve/middleware/authenticate.py create mode 100644 vllm/entrypoints/serve/middleware/log_response.py create mode 100644 vllm/entrypoints/serve/middleware/register.py create mode 100644 vllm/entrypoints/serve/middleware/x_request_id.py diff --git a/tests/entrypoints/serve/middleware/__init__.py b/tests/entrypoints/serve/middleware/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/entrypoints/serve/middleware/test_authentication_middleware.py b/tests/entrypoints/serve/middleware/test_authentication_middleware.py new file mode 100644 index 000000000000..528261e644db --- /dev/null +++ b/tests/entrypoints/serve/middleware/test_authentication_middleware.py @@ -0,0 +1,150 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from argparse import Namespace +from typing import get_args + +import pytest +import regex as re +from fastapi import FastAPI +from starlette.responses import JSONResponse +from starlette.routing import Route +from starlette.testclient import TestClient + +from vllm.entrypoints.launchers.api_server.routers import register_api_routers +from vllm.entrypoints.serve.middleware.authenticate import ( + GUARDED_PREFIX, + AuthenticationMiddleware, +) +from vllm.tasks import POOLING_TASKS, SupportedTask + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def get_all_http_routes(app: FastAPI) -> list[tuple[str, list[str]]]: + """Extract all HTTP routes (path, methods) from the FastAPI app.""" + routes = [] + for route in app.routes: + if not isinstance(route, Route): + continue + path = route.path + methods = list(route.methods or {"GET"}) + routes.append((path, methods)) + return routes + + +def generate_test_path(path_template: str) -> str: + """Replace path parameters (e.g. {response_id}) with 'test'.""" + return re.sub(r"\{[^}]+\}", "test", path_template) + + +def _create_app_with_mock_routes(routes: list[tuple[str, list[str]]]) -> FastAPI: + """Create a FastAPI app with AuthenticationMiddleware and mock endpoints.""" + app = FastAPI() + app.add_middleware(AuthenticationMiddleware, tokens=["valid-token"]) + + async def mock_endpoint(): + return JSONResponse({"status": "ok"}) + + for path_template, methods in routes: + allowed_methods = list(set(methods + ["OPTIONS"])) + app.add_api_route( + path_template, + mock_endpoint, + methods=allowed_methods, + include_in_schema=False, + ) + return app + + +class MockModelConfig: + def __init__(self): + self.hf_config = Namespace() + self.hf_config.num_labels = 1 + + def get_pooling_task(self, supported_tasks: tuple["SupportedTask", ...]): + pooling_tasks = [s for s in supported_tasks if s in POOLING_TASKS] + return pooling_tasks[0] if len(pooling_tasks) > 0 else None + + +@pytest.fixture(params=get_args(SupportedTask)) +def task_routes(request, monkeypatch) -> tuple[str, list[tuple[str, list[str]]]]: + """For each supported task, build an app with only that task's routers, + extract all routes, and return the task name and routes.""" + task = request.param + # Enable development mode to register all routes (including dev-only routes). + monkeypatch.setenv("VLLM_SERVER_DEV_MODE", "1") + + app = FastAPI() + args = Namespace() + app.state = Namespace() + app.state.args = args + + # Register routers for this specific task (development mode already enabled). + register_api_routers( + args, app, supported_tasks=(task,), model_config=MockModelConfig() + ) + + routes = get_all_http_routes(app) + return task, routes + + +# --------------------------------------------------------------------------- +# Tests for auto-discovered routes +# --------------------------------------------------------------------------- + + +def test_auto_discovered_protected_routes_require_auth(task_routes): + """For every auto-discovered route that starts with a guarded prefix, + verify that authentication is enforced.""" + task, routes = task_routes + app = _create_app_with_mock_routes(routes) + client = TestClient(app) + + for path_template, methods in routes: + if not path_template.startswith(GUARDED_PREFIX): + continue + + test_path = generate_test_path(path_template) + test_method = methods[0] if methods else "GET" + + resp = client.request(test_method, test_path) + assert resp.status_code == 401, ( + f"[{task}] {test_method} {test_path} should reject missing token" + ) + + resp = client.request( + test_method, test_path, headers={"Authorization": "Bearer wrong"} + ) + assert resp.status_code == 401, ( + f"[{task}] {test_method} {test_path} should reject invalid token" + ) + + resp = client.request( + test_method, test_path, headers={"Authorization": "Bearer valid-token"} + ) + assert resp.status_code == 200, ( + f"[{task}] {test_method} {test_path} should accept valid token" + ) + + +def test_auto_discovered_unprotected_routes_no_auth(task_routes): + """For every auto-discovered route that does NOT start with a guarded + prefix, verify that no authentication is required.""" + task, routes = task_routes + app = _create_app_with_mock_routes(routes) + client = TestClient(app) + + for path_template, methods in routes: + if path_template.startswith(GUARDED_PREFIX): + continue + + test_path = generate_test_path(path_template) + test_method = methods[0] if methods else "GET" + + resp = client.request(test_method, test_path) + assert resp.status_code == 200, ( + f"[{task}] {test_method} {test_path} should be accessible without token" + ) diff --git a/tests/entrypoints/serve/instrumentator/test_optional_middleware.py b/tests/entrypoints/serve/middleware/test_optional_middleware.py similarity index 100% rename from tests/entrypoints/serve/instrumentator/test_optional_middleware.py rename to tests/entrypoints/serve/middleware/test_optional_middleware.py diff --git a/vllm/entrypoints/launchers/__init__.py b/vllm/entrypoints/launchers/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/vllm/entrypoints/launchers/api_server/__init__.py b/vllm/entrypoints/launchers/api_server/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/vllm/entrypoints/launchers/api_server/routers.py b/vllm/entrypoints/launchers/api_server/routers.py new file mode 100644 index 000000000000..39d670c019b7 --- /dev/null +++ b/vllm/entrypoints/launchers/api_server/routers.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from argparse import Namespace + +from fastapi import FastAPI + +from vllm import envs +from vllm.config import ModelConfig +from vllm.tasks import POOLING_TASKS, SupportedTask + + +def register_api_routers( + args: Namespace, + app: FastAPI, + supported_tasks: tuple["SupportedTask", ...], + model_config: ModelConfig | None = None, +): + from vllm.entrypoints.serve import register_vllm_serve_api_routers + + register_vllm_serve_api_routers(app) + + from vllm.entrypoints.openai.models.api_router import ( + attach_router as register_models_api_router, + ) + + register_models_api_router(app) + + from vllm.entrypoints.serve.sagemaker.api_router import ( + attach_router as register_sagemaker_api_router, + ) + + register_sagemaker_api_router(app, supported_tasks, model_config) + + if envs.VLLM_SERVER_DEV_MODE: + from vllm.entrypoints.serve import register_vllm_dev_api_routers + + register_vllm_dev_api_routers(app) + + if "generate" in supported_tasks: + from vllm.entrypoints.generate.api_router import ( + register_generate_api_routers, + ) + + register_generate_api_routers(app) + + from vllm.entrypoints.serve.elastic_ep.api_router import ( + attach_router as elastic_ep_attach_router, + ) + + elastic_ep_attach_router(app) + + if "generate" in supported_tasks or "render" in supported_tasks: + from vllm.entrypoints.scale_out.factories import register_scale_out_api_routers + + register_scale_out_api_routers(app, supported_tasks) + + if "transcription" in supported_tasks or "realtime" in supported_tasks: + from vllm.entrypoints.speech_to_text.factories import ( + register_speech_to_text_api_routers, + ) + + register_speech_to_text_api_routers(app, supported_tasks) + + if any(task in POOLING_TASKS for task in supported_tasks): + from vllm.entrypoints.pooling.factories import register_pooling_api_routers + + register_pooling_api_routers(app, supported_tasks, model_config) + + if getattr(args, "enable_fault_tolerance", False): + from vllm.entrypoints.serve.fault_tolerance.api_router import ( + register_fault_tolerance_api_router, + ) + + register_fault_tolerance_api_router(app) diff --git a/vllm/entrypoints/openai/api_server.py b/vllm/entrypoints/openai/api_server.py index 808420d191f4..c541f03f445e 100644 --- a/vllm/entrypoints/openai/api_server.py +++ b/vllm/entrypoints/openai/api_server.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio -import importlib -import inspect import multiprocessing import multiprocessing.forkserver as forkserver import os @@ -17,7 +15,6 @@ import uvloop from fastapi import FastAPI -from fastapi.middleware.cors import CORSMiddleware from starlette.datastructures import State import vllm.envs as envs @@ -26,11 +23,12 @@ from vllm.engine.protocol import EngineClient from vllm.entrypoints.chat_utils import load_chat_template from vllm.entrypoints.launcher import serve_http +from vllm.entrypoints.launchers.api_server.routers import register_api_routers from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args from vllm.entrypoints.openai.models.protocol import BaseModelPath from vllm.entrypoints.openai.models.serving import OpenAIServingModels -from vllm.entrypoints.serve.elastic_ep.middleware import ScalingMiddleware from vllm.entrypoints.serve.exception_handling.register import init_exception_handler +from vllm.entrypoints.serve.middleware.register import init_entrypoints_middleware from vllm.entrypoints.serve.sagemaker.api_router import sagemaker_standards_bootstrap from vllm.entrypoints.serve.tokenize.serving import ServingTokenization from vllm.entrypoints.serve.utils.api_utils import ( @@ -43,7 +41,6 @@ from vllm.entrypoints.serve.utils.server_utils import ( get_uvicorn_log_config, lifespan, - log_response, ) from vllm.logger import init_logger from vllm.reasoning import ReasoningParserManager @@ -205,64 +202,9 @@ def build_app( else: app = FastAPI(lifespan=lifespan) app.state.args = args + app.root_path = args.root_path - from vllm.entrypoints.serve import register_vllm_serve_api_routers - - register_vllm_serve_api_routers(app) - - from vllm.entrypoints.openai.models.api_router import ( - attach_router as register_models_api_router, - ) - - register_models_api_router(app) - - from vllm.entrypoints.serve.sagemaker.api_router import ( - attach_router as register_sagemaker_api_router, - ) - - register_sagemaker_api_router(app, supported_tasks, model_config) - - if envs.VLLM_SERVER_DEV_MODE: - from vllm.entrypoints.serve import register_vllm_dev_api_routers - - register_vllm_dev_api_routers(app) - - if "generate" in supported_tasks: - from vllm.entrypoints.generate.api_router import ( - register_generate_api_routers, - ) - - register_generate_api_routers(app) - - from vllm.entrypoints.serve.elastic_ep.api_router import ( - attach_router as elastic_ep_attach_router, - ) - - elastic_ep_attach_router(app) - - if "generate" in supported_tasks or "render" in supported_tasks: - from vllm.entrypoints.scale_out.factories import register_scale_out_api_routers - - register_scale_out_api_routers(app, supported_tasks) - - if "transcription" in supported_tasks or "realtime" in supported_tasks: - from vllm.entrypoints.speech_to_text.factories import ( - register_speech_to_text_api_routers, - ) - - register_speech_to_text_api_routers(app, supported_tasks) - - if any(task in POOLING_TASKS for task in supported_tasks): - from vllm.entrypoints.pooling.factories import register_pooling_api_routers - - register_pooling_api_routers(app, supported_tasks, model_config) - - if args.enable_fault_tolerance: - from vllm.entrypoints.serve.fault_tolerance.api_router import ( - register_fault_tolerance_api_router, - ) - - register_fault_tolerance_api_router(app) + register_api_routers(args, app, supported_tasks, model_config) # Endpoint plugins are attached last so their routes are registered after all core # routers. This runs even for the CPU only render server. A plugin eligible for @@ -271,58 +213,7 @@ def build_app( _attach_endpoint_plugins(app, supported_tasks) init_exception_handler(app) - - app.root_path = args.root_path - app.add_middleware( - CORSMiddleware, - allow_origins=args.allowed_origins, - allow_credentials=args.allow_credentials, - allow_methods=args.allowed_methods, - allow_headers=args.allowed_headers, - ) - - # Ensure --api-key option from CLI takes precedence over VLLM_API_KEY - if tokens := [key for key in (args.api_key or [envs.VLLM_API_KEY]) if key]: - from vllm.entrypoints.serve.utils.server_utils import AuthenticationMiddleware - - app.add_middleware(AuthenticationMiddleware, tokens=tokens) - - if args.enable_request_id_headers: - from vllm.entrypoints.serve.utils.server_utils import XRequestIdMiddleware - - app.add_middleware(XRequestIdMiddleware) - - # Add scaling middleware to check for scaling state - app.add_middleware(ScalingMiddleware) - - if "realtime" in supported_tasks: - # Add WebSocket metrics middleware - from vllm.entrypoints.speech_to_text.factories import ( - add_websocket_metrics_middleware, - ) - - add_websocket_metrics_middleware(app) - - if envs.VLLM_DEBUG_LOG_API_SERVER_RESPONSE: - logger.warning( - "CAUTION: Enabling log response in the API Server. " - "This can include sensitive information and should be " - "avoided in production." - ) - app.middleware("http")(log_response) - - for middleware in args.middleware: - module_path, object_name = middleware.rsplit(".", 1) - imported = getattr(importlib.import_module(module_path), object_name) - if inspect.isclass(imported): - app.add_middleware(imported) # type: ignore[arg-type] - elif inspect.iscoroutinefunction(imported): - app.middleware("http")(imported) - else: - raise ValueError( - f"Invalid middleware {middleware}. Must be a function or a class." - ) - + init_entrypoints_middleware(args, app, supported_tasks) app = sagemaker_standards_bootstrap(app) return app diff --git a/vllm/entrypoints/serve/middleware/__init__.py b/vllm/entrypoints/serve/middleware/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/vllm/entrypoints/serve/middleware/authenticate.py b/vllm/entrypoints/serve/middleware/authenticate.py new file mode 100644 index 000000000000..e8a61b21ee46 --- /dev/null +++ b/vllm/entrypoints/serve/middleware/authenticate.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import hashlib +import secrets +from collections.abc import Awaitable + +from starlette.datastructures import Headers +from starlette.responses import JSONResponse +from starlette.types import ASGIApp, Receive, Scope, Send + +GUARDED_PREFIX = ("/v1", "/v2", "/inference", "/cohere") + + +class AuthenticationMiddleware: + """ + Pure ASGI middleware that authenticates each request by checking + if the Authorization Bearer token exists and equals anyof "{api_key}". + + Notes + ----- + There are two cases in which authentication is skipped: + 1. The HTTP method is OPTIONS. + 2. The request path doesn't start with GUARDED_PREFIX (e.g. /health). + """ + + def __init__(self, app: ASGIApp, tokens: list[str]) -> None: + self.app = app + self.api_tokens = [hashlib.sha256(t.encode("utf-8")).digest() for t in tokens] + + def verify_token(self, headers: Headers) -> bool: + authorization_header_value = headers.get("Authorization") + if not authorization_header_value: + return False + + scheme, _, param = authorization_header_value.partition(" ") + if scheme.lower() != "bearer": + return False + + param_hash = hashlib.sha256(param.encode("utf-8")).digest() + + token_match = False + for token_hash in self.api_tokens: + token_match |= secrets.compare_digest(param_hash, token_hash) + + return token_match + + def __call__(self, scope: Scope, receive: Receive, send: Send) -> Awaitable[None]: + if ( + scope["type"] not in ("http", "websocket") + or scope.get("method") == "OPTIONS" + ): + # scope["type"] can be "lifespan" or "startup" for example, + # in which case we don't need to do anything + return self.app(scope, receive, send) + root_path = scope.get("root_path", "") + url_path = scope["path"].removeprefix(root_path) + headers = Headers(scope=scope) + # Type narrow to satisfy mypy. + if url_path.startswith(GUARDED_PREFIX) and not self.verify_token(headers): + response = JSONResponse(content={"error": "Unauthorized"}, status_code=401) + return response(scope, receive, send) + return self.app(scope, receive, send) diff --git a/vllm/entrypoints/serve/middleware/log_response.py b/vllm/entrypoints/serve/middleware/log_response.py new file mode 100644 index 000000000000..e2f5e5bde51e --- /dev/null +++ b/vllm/entrypoints/serve/middleware/log_response.py @@ -0,0 +1,162 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pydantic +from fastapi import Request +from starlette.concurrency import iterate_in_threadpool + +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +class SSEDecoder: + """Robust Server-Sent Events decoder for streaming responses.""" + + def __init__(self): + self.buffer = "" + self.content_buffer = [] + + def decode_chunk(self, chunk: bytes) -> list[dict]: + """Decode a chunk of SSE data and return parsed events.""" + import json + + try: + chunk_str = chunk.decode("utf-8") + except UnicodeDecodeError: + # Skip malformed chunks + return [] + + self.buffer += chunk_str + events = [] + + # Process complete lines + while "\n" in self.buffer: + line, self.buffer = self.buffer.split("\n", 1) + line = line.rstrip("\r") # Handle CRLF + + if line.startswith("data: "): + data_str = line[6:].strip() + if data_str == "[DONE]": + events.append({"type": "done"}) + elif data_str: + try: + event_data = json.loads(data_str) + events.append({"type": "data", "data": event_data}) + except json.JSONDecodeError: + # Skip malformed JSON + continue + + return events + + def extract_content(self, event_data: dict) -> str: + """Extract content from event data.""" + return _extract_content_from_chunk(event_data) + + def add_content(self, content: str) -> None: + """Add content to the buffer.""" + if content: + self.content_buffer.append(content) + + def get_complete_content(self) -> str: + """Get the complete buffered content.""" + return "".join(self.content_buffer) + + +def _extract_content_from_chunk(chunk_data: dict) -> str: + """Extract content from a streaming response chunk.""" + try: + from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionStreamResponse, + ) + from vllm.entrypoints.openai.completion.protocol import ( + CompletionStreamResponse, + ) + + # Try using Completion types for type-safe parsing + if chunk_data.get("object") == "chat.completion.chunk": + chat_response = ChatCompletionStreamResponse.model_validate(chunk_data) + if chat_response.choices and chat_response.choices[0].delta.content: + return chat_response.choices[0].delta.content + elif chunk_data.get("object") == "text_completion": + completion_response = CompletionStreamResponse.model_validate(chunk_data) + if completion_response.choices and completion_response.choices[0].text: + return completion_response.choices[0].text + except pydantic.ValidationError: + # Fallback to manual parsing + if "choices" in chunk_data and chunk_data["choices"]: + choice = chunk_data["choices"][0] + if "delta" in choice and choice["delta"].get("content"): + return choice["delta"]["content"] + elif choice.get("text"): + return choice["text"] + return "" + + +def _log_streaming_response(response, response_body: list) -> None: + """Log streaming response with robust SSE parsing.""" + + sse_decoder = SSEDecoder() + chunk_count = 0 + + def buffered_iterator(): + nonlocal chunk_count + + for chunk in response_body: + chunk_count += 1 + yield chunk + + # Parse SSE events from chunk + events = sse_decoder.decode_chunk(chunk) + + for event in events: + if event["type"] == "data": + content = sse_decoder.extract_content(event["data"]) + sse_decoder.add_content(content) + elif event["type"] == "done": + # Log complete content when done + full_content = sse_decoder.get_complete_content() + if full_content: + # Truncate if too long + if len(full_content) > 2048: + full_content = full_content[:2048] + "...[truncated]" + logger.info( + "response_body={streaming_complete: content=%r, chunks=%d}", + full_content, + chunk_count, + ) + else: + logger.info( + "response_body={streaming_complete: no_content, chunks=%d}", + chunk_count, + ) + return + + response.body_iterator = iterate_in_threadpool(buffered_iterator()) + logger.info("response_body={streaming_started: chunks=%d}", len(response_body)) + + +def _log_non_streaming_response(response_body: list) -> None: + """Log non-streaming response.""" + try: + decoded_body = response_body[0].decode() + logger.info("response_body={%s}", decoded_body) + except UnicodeDecodeError: + logger.info("response_body={}") + + +async def log_response(request: Request, call_next): + response = await call_next(request) + response_body = [section async for section in response.body_iterator] + response.body_iterator = iterate_in_threadpool(iter(response_body)) + # Check if this is a streaming response by looking at content-type + content_type = response.headers.get("content-type", "") + is_streaming = content_type == "text/event-stream; charset=utf-8" + + # Log response body based on type + if not response_body: + logger.info("response_body={}") + elif is_streaming: + _log_streaming_response(response, response_body) + else: + _log_non_streaming_response(response_body) + return response diff --git a/vllm/entrypoints/serve/middleware/register.py b/vllm/entrypoints/serve/middleware/register.py new file mode 100644 index 000000000000..3b132ce0886d --- /dev/null +++ b/vllm/entrypoints/serve/middleware/register.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import importlib +import inspect +from argparse import Namespace + +from fastapi import FastAPI +from starlette.middleware.cors import CORSMiddleware + +from vllm import envs +from vllm.logger import init_logger +from vllm.tasks import SupportedTask + +from .log_response import log_response + +logger = init_logger(__name__) + + +def init_entrypoints_middleware( + args: Namespace, + app: FastAPI, + supported_tasks: tuple["SupportedTask", ...], +): + app.add_middleware( + CORSMiddleware, + allow_origins=args.allowed_origins, + allow_credentials=args.allow_credentials, + allow_methods=args.allowed_methods, + allow_headers=args.allowed_headers, + ) + + # Ensure --api-key option from CLI takes precedence over VLLM_API_KEY + if tokens := [key for key in (args.api_key or [envs.VLLM_API_KEY]) if key]: + from .authenticate import AuthenticationMiddleware + + app.add_middleware(AuthenticationMiddleware, tokens=tokens) + + if args.enable_request_id_headers: + from .x_request_id import XRequestIdMiddleware + + app.add_middleware(XRequestIdMiddleware) + + if "generate" in supported_tasks: + # Add scaling middleware to check for scaling state + from vllm.entrypoints.serve.elastic_ep.middleware import ScalingMiddleware + + app.add_middleware(ScalingMiddleware) + + if "realtime" in supported_tasks: + # Add WebSocket metrics middleware + from vllm.entrypoints.speech_to_text.realtime.metrics import ( + WebSocketMetricsMiddleware, + ) + + app.add_middleware(WebSocketMetricsMiddleware) + + if envs.VLLM_DEBUG_LOG_API_SERVER_RESPONSE: + logger.warning( + "CAUTION: Enabling log response in the API Server. " + "This can include sensitive information and should be " + "avoided in production." + ) + app.middleware("http")(log_response) + + for middleware in args.middleware: + module_path, object_name = middleware.rsplit(".", 1) + imported = getattr(importlib.import_module(module_path), object_name) + if inspect.isclass(imported): + app.add_middleware(imported) # type: ignore[arg-type] + elif inspect.iscoroutinefunction(imported): + app.middleware("http")(imported) + else: + raise ValueError( + f"Invalid middleware {middleware}. Must be a function or a class." + ) diff --git a/vllm/entrypoints/serve/middleware/x_request_id.py b/vllm/entrypoints/serve/middleware/x_request_id.py new file mode 100644 index 000000000000..02e85306cd12 --- /dev/null +++ b/vllm/entrypoints/serve/middleware/x_request_id.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import uuid +from collections.abc import Awaitable + +from starlette.datastructures import Headers, MutableHeaders +from starlette.types import ASGIApp, Message, Receive, Scope, Send + + +class XRequestIdMiddleware: + """ + Middleware the set's the X-Request-Id header for each response + to a random uuid4 (hex) value if the header isn't already + present in the request, otherwise use the provided request id. + """ + + def __init__(self, app: ASGIApp) -> None: + self.app = app + + def __call__(self, scope: Scope, receive: Receive, send: Send) -> Awaitable[None]: + if scope["type"] not in ("http", "websocket"): + return self.app(scope, receive, send) + + # Extract the request headers. + request_headers = Headers(scope=scope) + + async def send_with_request_id(message: Message) -> None: + """ + Custom send function to mutate the response headers + and append X-Request-Id to it. + """ + if message["type"] == "http.response.start": + response_headers = MutableHeaders(raw=message["headers"]) + request_id = request_headers.get("X-Request-Id", uuid.uuid4().hex) + response_headers.append("X-Request-Id", request_id) + await send(message) + + return self.app(scope, receive, send_with_request_id) diff --git a/vllm/entrypoints/serve/utils/server_utils.py b/vllm/entrypoints/serve/utils/server_utils.py index 9eba5d3a674e..97910d019ec7 100644 --- a/vllm/entrypoints/serve/utils/server_utils.py +++ b/vllm/entrypoints/serve/utils/server_utils.py @@ -1,20 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio -import hashlib import json -import secrets -import uuid from argparse import Namespace -from collections.abc import Awaitable from contextlib import asynccontextmanager -import pydantic -from fastapi import FastAPI, Request -from fastapi.responses import JSONResponse -from starlette.concurrency import iterate_in_threadpool -from starlette.datastructures import Headers, MutableHeaders -from starlette.types import ASGIApp, Message, Receive, Scope, Send +from fastapi import FastAPI from vllm import envs from vllm.engine.protocol import EngineClient @@ -24,91 +15,6 @@ logger = init_logger("vllm.entrypoints.openai.server_utils") -GUARDED_PREFIX = ("/v1", "/v2", "/inference", "/cohere") - - -class AuthenticationMiddleware: - """ - Pure ASGI middleware that authenticates each request by checking - if the Authorization Bearer token exists and equals anyof "{api_key}". - - Notes - ----- - There are two cases in which authentication is skipped: - 1. The HTTP method is OPTIONS. - 2. The request path doesn't start with GUARDED_PREFIX (e.g. /health). - """ - - def __init__(self, app: ASGIApp, tokens: list[str]) -> None: - self.app = app - self.api_tokens = [hashlib.sha256(t.encode("utf-8")).digest() for t in tokens] - - def verify_token(self, headers: Headers) -> bool: - authorization_header_value = headers.get("Authorization") - if not authorization_header_value: - return False - - scheme, _, param = authorization_header_value.partition(" ") - if scheme.lower() != "bearer": - return False - - param_hash = hashlib.sha256(param.encode("utf-8")).digest() - - token_match = False - for token_hash in self.api_tokens: - token_match |= secrets.compare_digest(param_hash, token_hash) - - return token_match - - def __call__(self, scope: Scope, receive: Receive, send: Send) -> Awaitable[None]: - if ( - scope["type"] not in ("http", "websocket") - or scope.get("method") == "OPTIONS" - ): - # scope["type"] can be "lifespan" or "startup" for example, - # in which case we don't need to do anything - return self.app(scope, receive, send) - root_path = scope.get("root_path", "") - url_path = scope["path"].removeprefix(root_path) - headers = Headers(scope=scope) - # Type narrow to satisfy mypy. - if url_path.startswith(GUARDED_PREFIX) and not self.verify_token(headers): - response = JSONResponse(content={"error": "Unauthorized"}, status_code=401) - return response(scope, receive, send) - return self.app(scope, receive, send) - - -class XRequestIdMiddleware: - """ - Middleware the set's the X-Request-Id header for each response - to a random uuid4 (hex) value if the header isn't already - present in the request, otherwise use the provided request id. - """ - - def __init__(self, app: ASGIApp) -> None: - self.app = app - - def __call__(self, scope: Scope, receive: Receive, send: Send) -> Awaitable[None]: - if scope["type"] not in ("http", "websocket"): - return self.app(scope, receive, send) - - # Extract the request headers. - request_headers = Headers(scope=scope) - - async def send_with_request_id(message: Message) -> None: - """ - Custom send function to mutate the response headers - and append X-Request-Id to it. - """ - if message["type"] == "http.response.start": - response_headers = MutableHeaders(raw=message["headers"]) - request_id = request_headers.get("X-Request-Id", uuid.uuid4().hex) - response_headers.append("X-Request-Id", request_id) - await send(message) - - return self.app(scope, receive, send_with_request_id) - - def load_log_config(log_config_file: str | None) -> dict | None: if not log_config_file: return None @@ -155,161 +61,6 @@ def get_uvicorn_log_config(args: Namespace) -> dict | None: return None -def _extract_content_from_chunk(chunk_data: dict) -> str: - """Extract content from a streaming response chunk.""" - try: - from vllm.entrypoints.openai.chat_completion.protocol import ( - ChatCompletionStreamResponse, - ) - from vllm.entrypoints.openai.completion.protocol import ( - CompletionStreamResponse, - ) - - # Try using Completion types for type-safe parsing - if chunk_data.get("object") == "chat.completion.chunk": - chat_response = ChatCompletionStreamResponse.model_validate(chunk_data) - if chat_response.choices and chat_response.choices[0].delta.content: - return chat_response.choices[0].delta.content - elif chunk_data.get("object") == "text_completion": - completion_response = CompletionStreamResponse.model_validate(chunk_data) - if completion_response.choices and completion_response.choices[0].text: - return completion_response.choices[0].text - except pydantic.ValidationError: - # Fallback to manual parsing - if "choices" in chunk_data and chunk_data["choices"]: - choice = chunk_data["choices"][0] - if "delta" in choice and choice["delta"].get("content"): - return choice["delta"]["content"] - elif choice.get("text"): - return choice["text"] - return "" - - -class SSEDecoder: - """Robust Server-Sent Events decoder for streaming responses.""" - - def __init__(self): - self.buffer = "" - self.content_buffer = [] - - def decode_chunk(self, chunk: bytes) -> list[dict]: - """Decode a chunk of SSE data and return parsed events.""" - import json - - try: - chunk_str = chunk.decode("utf-8") - except UnicodeDecodeError: - # Skip malformed chunks - return [] - - self.buffer += chunk_str - events = [] - - # Process complete lines - while "\n" in self.buffer: - line, self.buffer = self.buffer.split("\n", 1) - line = line.rstrip("\r") # Handle CRLF - - if line.startswith("data: "): - data_str = line[6:].strip() - if data_str == "[DONE]": - events.append({"type": "done"}) - elif data_str: - try: - event_data = json.loads(data_str) - events.append({"type": "data", "data": event_data}) - except json.JSONDecodeError: - # Skip malformed JSON - continue - - return events - - def extract_content(self, event_data: dict) -> str: - """Extract content from event data.""" - return _extract_content_from_chunk(event_data) - - def add_content(self, content: str) -> None: - """Add content to the buffer.""" - if content: - self.content_buffer.append(content) - - def get_complete_content(self) -> str: - """Get the complete buffered content.""" - return "".join(self.content_buffer) - - -def _log_streaming_response(response, response_body: list) -> None: - """Log streaming response with robust SSE parsing.""" - from starlette.concurrency import iterate_in_threadpool - - sse_decoder = SSEDecoder() - chunk_count = 0 - - def buffered_iterator(): - nonlocal chunk_count - - for chunk in response_body: - chunk_count += 1 - yield chunk - - # Parse SSE events from chunk - events = sse_decoder.decode_chunk(chunk) - - for event in events: - if event["type"] == "data": - content = sse_decoder.extract_content(event["data"]) - sse_decoder.add_content(content) - elif event["type"] == "done": - # Log complete content when done - full_content = sse_decoder.get_complete_content() - if full_content: - # Truncate if too long - if len(full_content) > 2048: - full_content = full_content[:2048] + "" - "...[truncated]" - logger.info( - "response_body={streaming_complete: content=%r, chunks=%d}", - full_content, - chunk_count, - ) - else: - logger.info( - "response_body={streaming_complete: no_content, chunks=%d}", - chunk_count, - ) - return - - response.body_iterator = iterate_in_threadpool(buffered_iterator()) - logger.info("response_body={streaming_started: chunks=%d}", len(response_body)) - - -def _log_non_streaming_response(response_body: list) -> None: - """Log non-streaming response.""" - try: - decoded_body = response_body[0].decode() - logger.info("response_body={%s}", decoded_body) - except UnicodeDecodeError: - logger.info("response_body={}") - - -async def log_response(request: Request, call_next): - response = await call_next(request) - response_body = [section async for section in response.body_iterator] - response.body_iterator = iterate_in_threadpool(iter(response_body)) - # Check if this is a streaming response by looking at content-type - content_type = response.headers.get("content-type", "") - is_streaming = content_type == "text/event-stream; charset=utf-8" - - # Log response body based on type - if not response_body: - logger.info("response_body={}") - elif is_streaming: - _log_streaming_response(response, response_body) - else: - _log_non_streaming_response(response_body) - return response - - _running_tasks: set[asyncio.Task] = set() diff --git a/vllm/entrypoints/speech_to_text/factories.py b/vllm/entrypoints/speech_to_text/factories.py index 1971e32b989e..21633b22f323 100644 --- a/vllm/entrypoints/speech_to_text/factories.py +++ b/vllm/entrypoints/speech_to_text/factories.py @@ -37,12 +37,6 @@ def register_speech_to_text_api_routers( app.include_router(translation_router) -def add_websocket_metrics_middleware(app: FastAPI): - from .realtime.metrics import WebSocketMetricsMiddleware - - app.add_middleware(WebSocketMetricsMiddleware) - - def init_speech_to_text_state( engine_client: "EngineClient", state: "State", From a02cfccbc6187344325e364f09f6d8c33c4b253b Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Mon, 17 Aug 2026 01:29:41 -0500 Subject: [PATCH 037/839] [Bugfix][Mamba] Fix overlapping state copy race (#50729) Signed-off-by: Andreas Karatzas Co-authored-by: OpenAI Codex --- tests/v1/worker/test_mamba_utils.py | 189 ++++++++++++++++++++++------ vllm/v1/worker/mamba_utils.py | 102 +++++++++++---- 2 files changed, 228 insertions(+), 63 deletions(-) diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index a338b934b54f..2ba1ce1c5937 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -18,6 +18,7 @@ from vllm.v1.worker.mamba_utils import ( MambaCopyBuffers, MambaSpecDecodeGPUContext, + batch_memcpy, collect_mamba_copy_meta, do_mamba_copy_block, preprocess_mamba, @@ -536,6 +537,34 @@ def device(self): def test_config(self): return _TestConfig() + def test_batch_memcpy_left_overlap_has_memmove_semantics(self, device): + batch = 128 + row_bytes = 32 * 1024 + shift = 16 + copy_size = row_bytes - shift + + pattern = (torch.arange(row_bytes, dtype=torch.int32, device=device) % 251).to( + torch.uint8 + ) + state = pattern.expand(batch, -1).clone() + snapshot = state.clone() + + row_stride_bytes = state.stride(0) * state.element_size() + row_offsets = ( + torch.arange(batch, dtype=torch.int64, device=device) * row_stride_bytes + ) + dst_ptrs = (row_offsets + state.data_ptr()).to(torch.uint64) + src_ptrs = (row_offsets + state.data_ptr() + shift).to(torch.uint64) + sizes = torch.full((batch,), copy_size, dtype=torch.int32, device=device) + + expected = snapshot.clone() + expected[:, :copy_size].copy_(snapshot[:, shift:]) + for _ in range(10): + state.copy_(snapshot) + batch_memcpy(src_ptrs, dst_ptrs, sizes) + torch.accelerator.synchronize() + torch.testing.assert_close(state, expected, rtol=0, atol=0) + def test_matches_python_postprocess_mamba(self, device, test_config): """ Golden test: GPU kernel produces identical results to Python impl. @@ -1190,12 +1219,27 @@ def test_same_block_idx_with_offset_copies_then_sets_accepted_to_1( # --- Verify Python behavior (ground truth) --- dest_block_id = block_ids_per_req[0][1] # dest_block_idx = 1 - # Conv state should be modified (shifted copy within block) - conv_changed = not torch.allclose( - conv_state_py[dest_block_id], conv_state_orig[dest_block_id] + # This is an overlapping in-place left shift, so comparing only the + # Python and fused paths can hide the same memcpy race in both. Build + # the memmove result from the untouched snapshot and check each path + # independently. + expected_conv_state = conv_state_orig.clone() + expected_conv_state[dest_block_id, :-1].copy_( + conv_state_orig[dest_block_id, 1:] ) - assert conv_changed, ( - "Python: Conv state should be modified when accept_token_bias > 0" + torch.testing.assert_close( + conv_state_py, + expected_conv_state, + rtol=0, + atol=0, + msg="Python: overlapping conv copy should have memmove semantics", + ) + torch.testing.assert_close( + conv_state_gpu, + expected_conv_state, + rtol=0, + atol=0, + msg="GPU: overlapping conv copy should have memmove semantics", ) # Temporal state should be modified (copy from different block) @@ -2122,13 +2166,10 @@ def test_temporal_copy_with_bias_ge_2(self, device, test_config): actual_src_block_idx = src_block_idx + accept_token_bias actual_src_block_id = block_table[req, actual_src_block_idx] - All prior regression tests exercise only ``bias == 1``, i.e. they - only ever read one slot ahead of ``src_block_idx`` in the block - table. An off-by-one (or missing scale) in the address computation - on line 143 of ``mamba_utils.py`` would be invisible to every - existing test but would silently read the wrong physical block on - any speculative-decode cycle that accepts multiple tokens across a - block boundary, feeding a stale hidden state forward one step. + A ``bias == 1`` case only reads one slot ahead of ``src_block_idx`` + in the block table. This test isolates the larger-stride case, where + an off-by-one would read the wrong physical block after multiple + tokens are accepted across a block boundary. Setup (block_size=16): - running = 28 + 2 - 0 = 30 @@ -2141,8 +2182,7 @@ def test_temporal_copy_with_bias_ge_2(self, device, test_config): With identity block_ids = [0,1,2,3,...], an off-by-one that used bias=1 would copy from block_ids[2]=2 instead of block_ids[3]=3, - producing a clear state-value mismatch against the Python - reference. + producing a clear mismatch against the untouched snapshot. """ cfg = test_config torch.manual_seed(7002) @@ -2166,6 +2206,7 @@ def test_temporal_copy_with_bias_ge_2(self, device, test_config): fwd_py, fwd_gpu, ) = _make_dual_layer_state(cfg, device) + conv_state_orig = conv_state_py.clone() temporal_state_orig = temporal_state_py.clone() # --- Python reference --- @@ -2212,12 +2253,22 @@ def test_temporal_copy_with_bias_ge_2(self, device, test_config): device=device, ) - # --- Ground truth: Python must have sourced temporal from block 3 --- + # --- Ground truth from untouched snapshots --- actual_src_block_id = block_ids_per_req[0][3] # == 3 dest_block_id = block_ids_per_req[0][1] # == 1 + expected_conv_state = conv_state_orig.clone() + expected_conv_state[dest_block_id, :-2].copy_( + conv_state_orig[dest_block_id, 2:] + ) + torch.testing.assert_close(conv_state_py, expected_conv_state, rtol=0, atol=0) + torch.testing.assert_close(conv_state_gpu, expected_conv_state, rtol=0, atol=0) + + # Python must have sourced temporal from block 3. torch.testing.assert_close( temporal_state_py[dest_block_id], temporal_state_orig[actual_src_block_id], + rtol=0, + atol=0, msg=( "Python reference did not copy from block_ids[src+bias]=3; " "test preconditions are wrong" @@ -2251,21 +2302,44 @@ def test_temporal_copy_with_bias_ge_2(self, device, test_config): msg="num_accepted_tokens mismatch at accept_token_bias=2", ) - def test_ds_conv_layout_bias_gt_0_byte_equal_to_sd( - self, device, test_config, monkeypatch + @pytest.mark.parametrize( + "same_physical_block", [True, False], ids=["same", "distinct"] + ) + @pytest.mark.parametrize("accept_token_bias", [1, 2, 3]) + @pytest.mark.parametrize( + "dtype", + [torch.float16, torch.float32, torch.float64], + ids=["fp16", "fp32", "fp64"], + ) + def test_sd_and_ds_conv_layouts_match_snapshot( + self, + device, + test_config, + monkeypatch, + accept_token_bias, + same_physical_block, + dtype, ): - """DS conv postprocess should match SD when accept_token_bias > 0.""" + """SD and DS copies should independently match memmove semantics.""" from vllm.model_executor.layers.mamba import mamba_utils as model_mamba_utils cfg = test_config + cfg.dtype = dtype torch.manual_seed(38898) req_ids = ["req_0"] - num_computed_tokens = [30] - num_scheduled_tokens = {"req_0": 1} + # Keep new_num_computed on an aligned boundary while varying how far + # below it the running state starts. This makes the copy bias exactly + # ``accept_token_bias`` for each case. The 32 boundary keeps source and + # destination in logical block 1; the 64 boundary copies block 2 -> 3. + aligned_boundary = 32 if same_physical_block else 64 + num_computed_tokens = [aligned_boundary - 2 * accept_token_bias] + num_scheduled_tokens = {"req_0": accept_token_bias} num_draft_tokens: dict[str, int] = {} - num_accepted_tokens = [2] # Results in accept_token_bias = 1 - mamba_state_idx = [1] # src_block_idx = 1 = dest_block_idx + num_accepted_tokens = [accept_token_bias + 1] + dest_block_idx = aligned_boundary // cfg.block_size - 1 + src_block_idx = dest_block_idx if same_physical_block else dest_block_idx - 1 + mamba_state_idx = [src_block_idx] block_ids_per_req = [list(range(8))] layer_names = ["layer_0"] @@ -2288,7 +2362,8 @@ def test_ds_conv_layout_bias_gt_0_byte_equal_to_sd( cfg.num_blocks, cfg.temporal_state_dim, dtype=cfg.dtype, device=device ) - # SD GPU path. Default layout is SD. + # SD GPU path. + monkeypatch.delenv("VLLM_SSM_CONV_STATE_LAYOUT", raising=False) model_mamba_utils.get_conv_state_layout.cache_clear() sd_conv = sd_source_conv.clone() sd_temporal = sd_source_temporal.clone() @@ -2312,9 +2387,32 @@ def test_ds_conv_layout_bias_gt_0_byte_equal_to_sd( ) torch.accelerator.synchronize() - # Sanity: SD path actually modified the state (copy was performed). - assert not torch.equal(sd_conv, sd_source_conv), ( - "SD baseline did not modify conv state; test setup is wrong" + src_block_id = block_ids_per_req[0][src_block_idx] + dest_block_id = block_ids_per_req[0][dest_block_idx] + expected_conv = sd_source_conv.clone() + expected_conv[dest_block_id, :-accept_token_bias].copy_( + sd_source_conv[src_block_id, accept_token_bias:] + ) + torch.testing.assert_close( + sd_conv, + expected_conv, + rtol=0, + atol=0, + msg="SD conv copy did not match the untouched source snapshot", + ) + + actual_temporal_src_idx = src_block_idx + accept_token_bias + actual_temporal_src_id = block_ids_per_req[0][actual_temporal_src_idx] + expected_temporal = sd_source_temporal.clone() + expected_temporal[dest_block_id].copy_( + sd_source_temporal[actual_temporal_src_id] + ) + torch.testing.assert_close( + sd_temporal, + expected_temporal, + rtol=0, + atol=0, + msg="SD temporal copy did not match the untouched source snapshot", ) # DS GPU path on the DS twin. @@ -2346,22 +2444,39 @@ def test_ds_conv_layout_bias_gt_0_byte_equal_to_sd( # Reset the lru cache so other tests see the default layout again. model_mamba_utils.get_conv_state_layout.cache_clear() - # DS bytes, un-permuted, should match the SD result. + # Validate DS independently against the snapshot; otherwise a shared + # SD/DS bug would remain invisible. + ds_conv_sd_layout = ds_conv.permute(0, 2, 1).contiguous() torch.testing.assert_close( - ds_conv.permute(0, 2, 1).contiguous(), - sd_conv, - msg=( - "DS conv post-kernel does not match SD baseline; the DS " - "row-loop in postprocess_mamba_fused_kernel is wrong." - ), + ds_conv_sd_layout, + expected_conv, + rtol=0, + atol=0, + msg="DS conv copy did not match the untouched source snapshot", ) torch.testing.assert_close( ds_temporal, - sd_temporal, - msg="DS temporal state diverged from SD", + expected_temporal, + rtol=0, + atol=0, + msg="DS temporal copy did not match the untouched source snapshot", + ) + + expected_accepted = 1 if same_physical_block else accept_token_bias + 1 + expected_accepted_tensor = torch.tensor( + [expected_accepted], dtype=torch.int32, device=device ) torch.testing.assert_close( - gpu_ctx_ds.num_accepted_tokens_out[:num_reqs], gpu_ctx_sd.num_accepted_tokens_out[:num_reqs], - msg="DS num_accepted_tokens diverged from SD", + expected_accepted_tensor, + rtol=0, + atol=0, + msg="SD num_accepted_tokens result is wrong", + ) + torch.testing.assert_close( + gpu_ctx_ds.num_accepted_tokens_out[:num_reqs], + expected_accepted_tensor, + rtol=0, + atol=0, + msg="DS num_accepted_tokens result is wrong", ) diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index 9b89af081206..82247e3ee401 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -189,19 +189,46 @@ def _copy_mamba_state_block( src_block_id = tl.load(block_table_base + src_col).to(tl.int64) dim_rows = tl.load(state_dim_row_count_ptr + state_idx) row_stride = tl.load(state_dim_row_stride_ptr + state_idx) - per_row_bytes = (conv_width - token_bias).to(tl.int64) * state_elem_size - bias_bytes = token_bias.to(tl.int64) * state_elem_size src_block_addr = state_base_addr + src_block_id * state_block_stride offsets = tl.arange(0, COPY_BLOCK_SIZE) - for d in range(0, dim_rows): - row_src = src_block_addr + d * row_stride + bias_bytes - row_dst = dst_addr + d * row_stride - for i in range(0, per_row_bytes, COPY_BLOCK_SIZE): - mask = (i + offsets) < per_row_bytes - curr_src = (row_src + i + offsets).to(tl.pointer_type(tl.uint8)) - curr_dst = (row_dst + i + offsets).to(tl.pointer_type(tl.uint8)) - data = tl.load(curr_src, mask=mask) - tl.store(curr_dst, data, mask=mask) + + # Stable row-to-lane ownership makes left shifts memmove-safe while + # exposing the dimension rows in parallel. All addresses retain + # state_elem_size alignment: tensor strides and token offsets are + # measured in whole elements before conversion to bytes. + num_dst_tokens = conv_width - token_bias + for token_idx in range(0, num_dst_tokens): + for row_base in range(0, dim_rows, COPY_BLOCK_SIZE): + rows = row_base + offsets + mask = rows < dim_rows + src_byte_addr = ( + src_block_addr + + rows * row_stride + + (token_idx + token_bias) * state_elem_size + ) + dst_byte_addr = ( + dst_addr + rows * row_stride + token_idx * state_elem_size + ) + if state_elem_size == 2: + src_u16 = src_byte_addr.to(tl.pointer_type(tl.uint16)) + dst_u16 = dst_byte_addr.to(tl.pointer_type(tl.uint16)) + data_u16 = tl.load(src_u16, mask=mask) + tl.store(dst_u16, data_u16, mask=mask) + elif state_elem_size == 4: + src_u32 = src_byte_addr.to(tl.pointer_type(tl.uint32)) + dst_u32 = dst_byte_addr.to(tl.pointer_type(tl.uint32)) + data_u32 = tl.load(src_u32, mask=mask) + tl.store(dst_u32, data_u32, mask=mask) + else: + for byte_idx in range(0, state_elem_size): + src_u8 = (src_byte_addr + byte_idx).to( + tl.pointer_type(tl.uint8) + ) + dst_u8 = (dst_byte_addr + byte_idx).to( + tl.pointer_type(tl.uint8) + ) + data_u8 = tl.load(src_u8, mask=mask) + tl.store(dst_u8, data_u8, mask=mask) return if is_conv_state: @@ -210,22 +237,40 @@ def _copy_mamba_state_block( # SD conv: copy # state[bt[src_col], token_bias:] -> # state[bt[dst_col], :conv_width - token_bias] - # Small per-block bytes (~60-80 KiB) make tiling degenerate, so - # conv runs as a single-CTA memcpy (NUM_TILES=1). src_block_id = tl.load(block_table_base + src_col).to(tl.int64) - src_offset = token_bias.to(tl.int64) * state_inner_size * state_elem_size - src_addr = state_base_addr + src_block_id * state_block_stride + src_offset - copy_size = ( - (conv_width - token_bias).to(tl.int64) * state_inner_size * state_elem_size - ) - _memcpy_u64_tiled( - src_addr, - dst_addr, - copy_size, - tile_idx, - COPY_BLOCK_SIZE=COPY_BLOCK_SIZE, - NUM_TILES=1, - ) + src_block_addr = state_base_addr + src_block_id * state_block_stride + token_bytes = state_inner_size * state_elem_size + num_dst_tokens = conv_width - token_bias + + # Distinct blocks and exact self-copies cannot have a destructive + # overlap, so retain the u64-vectorized single-CTA copy. + if src_block_id != dest_block_id or token_bias == 0: + src_addr = src_block_addr + token_bias.to(tl.int64) * token_bytes + copy_size = num_dst_tokens.to(tl.int64) * token_bytes + _memcpy_u64_tiled( + src_addr, + dst_addr, + copy_size, + tile_idx, + COPY_BLOCK_SIZE=COPY_BLOCK_SIZE, + NUM_TILES=1, + ) + return + + # Copy tokens from low to high. Each token-sized source and destination + # region is disjoint, so same-block left shifts are memmove-safe + # without a barrier. + for token_idx in range(0, num_dst_tokens): + src_token = src_block_addr + (token_idx + token_bias) * token_bytes + dst_token = dst_addr + token_idx * token_bytes + _memcpy_u64_tiled( + src_token, + dst_token, + token_bytes, + tile_idx, + COPY_BLOCK_SIZE=COPY_BLOCK_SIZE, + NUM_TILES=1, + ) return # Temporal state: copy state[bt[src_col + token_bias]] -> state[bt[dst_col]] @@ -522,6 +567,7 @@ def batch_memcpy_kernel(src_ptrs, dst_ptrs, sizes, BLOCK_SIZE: tl.constexpr): src_ptr = tl.load(src_ptrs + pid) dst_ptr = tl.load(dst_ptrs + pid) size = tl.load(sizes + pid) + is_left_overlap = dst_ptr < src_ptr and dst_ptr + size > src_ptr offsets = tl.arange(0, BLOCK_SIZE) for i in range(0, size, BLOCK_SIZE): @@ -531,6 +577,10 @@ def batch_memcpy_kernel(src_ptrs, dst_ptrs, sizes, BLOCK_SIZE: tl.constexpr): curr_dst_ptr = (dst_ptr + i + offsets).to(tl.pointer_type(tl.uint8)) data = tl.load(curr_src_ptr, mask=mask) + if is_left_overlap: + # Preserve each lane's source before a lower-address lane stores + # over it. The condition is uniform within the program. + tl.debug_barrier() tl.store(curr_dst_ptr, data, mask=mask) From 5fd7a888386cff800f32de6b5a33d1dd3ca1e397 Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Mon, 17 Aug 2026 14:40:36 +0800 Subject: [PATCH 038/839] [CI/Build] Fix accident pre-commit breakage due to concurrent merge (#52578) Signed-off-by: Isotr0py --- tests/models/multimodal/pooling/test_colpali.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/models/multimodal/pooling/test_colpali.py b/tests/models/multimodal/pooling/test_colpali.py index 4d4e0bd88b21..9e040bf9a1e5 100644 --- a/tests/models/multimodal/pooling/test_colpali.py +++ b/tests/models/multimodal/pooling/test_colpali.py @@ -151,8 +151,6 @@ def _run_multimodal_text_query_image_docs_test( _make_image_mm_param(red_image), _make_image_mm_param(blue_image), ] - attention_backend = "FLASH_ATTN" if current_platform.is_cuda() else None - scores = vllm_model.llm.score(query, image_docs) assert len(scores) == 2 @@ -231,6 +229,7 @@ def test_colpali_v2_multimodal_text_query_image_docs( dtype: str, ) -> None: monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") + attention_backend = "FLASH_ATTN" if current_platform.is_cuda() else None with vllm_runner( model, runner="pooling", From 0ff370b51c58a3072b85e68fa5686b77b5034965 Mon Sep 17 00:00:00 2001 From: Amal Sebastian Date: Mon, 17 Aug 2026 14:18:11 +0530 Subject: [PATCH 039/839] docs: fix incorrect --custom-skip-chat-template flag reference (#52588) --- docs/benchmarking/cli.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/benchmarking/cli.md b/docs/benchmarking/cli.md index 476171f73226..2437392d9be6 100644 --- a/docs/benchmarking/cli.md +++ b/docs/benchmarking/cli.md @@ -205,7 +205,7 @@ vllm bench serve --port 9001 --save-result --save-detailed \ --endpoint /v1/completions \ --dataset-name custom \ --dataset-path \ - --custom-skip-chat-template \ + --skip-chat-template \ --num-prompts 80 \ --max-concurrency 1 \ --temperature=0.3 \ @@ -213,7 +213,7 @@ vllm bench serve --port 9001 --save-result --save-detailed \ --result-dir "./log/" ``` -You can skip applying chat template if your data already has it by using `--custom-skip-chat-template`. +You can skip applying chat template if your data already has it by using `--skip-chat-template`. #### Custom Audio Dataset From c05d923f186842d4ac35bfddafdd5aa01b6cbaf8 Mon Sep 17 00:00:00 2001 From: TJian Date: Mon, 17 Aug 2026 16:56:44 +0800 Subject: [PATCH 040/839] [Doc] [ROCm] Update installation documentation (#52303) Signed-off-by: tjtanaa --- .../installation/gpu.rocm.inc.md | 71 +++++++++++++++++++ 1 file changed, 71 insertions(+) diff --git a/docs/getting_started/installation/gpu.rocm.inc.md b/docs/getting_started/installation/gpu.rocm.inc.md index 59c9723e666b..ad5af7994054 100644 --- a/docs/getting_started/installation/gpu.rocm.inc.md +++ b/docs/getting_started/installation/gpu.rocm.inc.md @@ -154,6 +154,77 @@ uv pip install vllm==${VLLM_VERSION} \ --8<-- [end:pre-built-wheels] --8<-- [start:build-wheel-from-source] +#### Set up using Python-only build (without compilation) {#python-only-build} + +If you only need to change Python code, you can build and install vLLM without +compilation. Changes you make to the code will be reflected when you run vLLM: + +```bash +git clone https://github.com/vllm-project/vllm.git +cd vllm +VLLM_USE_PRECOMPILED=1 python3 setup.py develop +``` + +This command will do the following: + +1. Look for the current branch in your vLLM clone. +1. Identify the corresponding base commit in the main branch. +1. Detect the ROCm version in your environment and select the matching wheel + variant. +1. Download the pre-built wheel of the base commit. +1. Use its compiled libraries and `vllm-rs` binary in the installation. + +!!! note + 1. If you change C++, HIP, or kernel code, you cannot use Python-only build; + otherwise you may see an import error about a library not being found or + an undefined symbol. + 2. If you rebase your development branch, it is recommended to uninstall + vLLM and re-run the above command to make sure your libraries are up to + date. + +!!! tip "Rebuilding the Rust frontend" +If you need to recompile the `vllm-rs` Rust frontend binary, you can rebuild and +install it without re-running the full installation: + + ```bash + ./build_rust.sh # release build + ./build_rust.sh --debug # faster build for development + ``` + + This will install the required Rust toolchain if needed, build the binary, + and place it in `vllm/vllm-rs`. + +If you see an error about a wheel not being found, the wheel for your base +commit and ROCm patch version might not be available. Check the available +variants under `https://wheels.vllm.ai/rocm//`. For example, ROCm 7.2.1 +uses the `rocm721` variant. + +There are more environment variables to control the behavior of Python-only +build: + +- `VLLM_PRECOMPILED_WHEEL_LOCATION`: specify the exact wheel URL or local file + path of a pre-compiled wheel to use. All other logic to find the wheel will be + skipped. +- `VLLM_PRECOMPILED_WHEEL_COMMIT`: override the full commit hash used to + download the pre-compiled wheel. +- `VLLM_PRECOMPILED_WHEEL_VARIANT`: specify the ROCm variant subdirectory, e.g., + `rocm700` or `rocm721`. If not specified, the variant is auto-detected based + on your system's ROCm version. An explicitly specified variant must match the + detected environment. + +You can find more information about vLLM's wheels in +[Install the latest code](#install-the-latest-code). + +!!! note + There is a possibility that your source code may have a different commit ID + compared to the vLLM wheel, which could potentially lead to unknown errors. + It is recommended to use the same commit ID for the source code as the vLLM + wheel you have installed. Please refer to + [Install the latest code](#install-the-latest-code) for instructions on how + to install a specified wheel. + +#### Full build (with compilation) {#full-build} + !!! tip - If you found that the following installation step does not work for you, please refer to [docker/Dockerfile.rocm_base](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm_base). Dockerfile is a form of installation steps. From bb233626caa31602728f7ee4625f3d2a4d1a3ad5 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Mon, 17 Aug 2026 02:03:33 -0700 Subject: [PATCH 041/839] [CI] Shard Humming A100 eval (#52325) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> --- .buildkite/test_areas/lm_eval.yaml | 5 +++-- tests/evals/gsm8k/configs/humming/config-a100-shard-0.txt | 4 ++++ tests/evals/gsm8k/configs/humming/config-a100-shard-1.txt | 5 +++++ tests/evals/gsm8k/configs/humming/config-a100-shard-2.txt | 5 +++++ 4 files changed, 17 insertions(+), 2 deletions(-) create mode 100644 tests/evals/gsm8k/configs/humming/config-a100-shard-0.txt create mode 100644 tests/evals/gsm8k/configs/humming/config-a100-shard-1.txt create mode 100644 tests/evals/gsm8k/configs/humming/config-a100-shard-2.txt diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index 4f7264573093..f9c21a54217b 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -195,12 +195,13 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt -- label: LM Eval Humming f16 (A100 - TEMPORARY) +- label: LM Eval Humming f16 (A100 - TEMPORARY) %N key: lm-eval-humming-f16-a100 timeout_in_minutes: 75 device: a100 optional: true num_devices: 1 + parallelism: 3 source_file_dependencies: - vllm/model_executor/layers/quantization/humming.py - vllm/model_executor/layers/quantization/utils/humming_utils.py @@ -208,7 +209,7 @@ steps: - vllm/model_executor/layers/fused_moe/oracle/ - vllm/model_executor/kernels/linear/ commands: - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-a100-shard-$$BUILDKITE_PARALLEL_JOB.txt - label: LM Eval Humming Act int8 (A100 - TEMPORARY) key: lm-eval-humming-act-a100 diff --git a/tests/evals/gsm8k/configs/humming/config-a100-shard-0.txt b/tests/evals/gsm8k/configs/humming/config-a100-shard-0.txt new file mode 100644 index 000000000000..63b428ac0e30 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/config-a100-shard-0.txt @@ -0,0 +1,4 @@ +Qwen3.5-35B-A3B-experts-int8-humming.yaml +gpt-oss-20b-humming.yaml +Qwen3-30B-A3B-GPTQ-Int4-humming.yaml +Qwen3-30B-A3B-FP8-block-humming.yaml diff --git a/tests/evals/gsm8k/configs/humming/config-a100-shard-1.txt b/tests/evals/gsm8k/configs/humming/config-a100-shard-1.txt new file mode 100644 index 000000000000..e0260c504e75 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/config-a100-shard-1.txt @@ -0,0 +1,5 @@ +Qwen3.6-35B-A3B-NVFP4-humming.yaml +Qwen3-30B-A3B-Instruct-2507-quantized.w8a8-humming.yaml +Qwen3-30B-A3B-Fp8-v1-humming.yaml +Qwen3-30B-A3B-AWQ-humming.yaml +Qwen3-0.6B-MXFP8-humming.yaml diff --git a/tests/evals/gsm8k/configs/humming/config-a100-shard-2.txt b/tests/evals/gsm8k/configs/humming/config-a100-shard-2.txt new file mode 100644 index 000000000000..8453e5d73cab --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/config-a100-shard-2.txt @@ -0,0 +1,5 @@ +Qwen3.5-35B-A3B-FP8-humming.yaml +NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4-humming.yaml +Qwen3-30B-A3B-MXFP4A16-humming.yaml +Qwen3-30B-A3B-NVFP4-humming.yaml +Qwen2-1.5B-Instruct-FP8W8-humming.yaml From cc7cf71fc819df579664adc2e944438d464dbc30 Mon Sep 17 00:00:00 2001 From: pmanczak Date: Mon, 17 Aug 2026 11:21:34 +0200 Subject: [PATCH 042/839] [XPU] Enable Kimi K3 KDA kernel tests on XPU (#51809) Signed-off-by: pmanczak --- tests/models/kimi_k3/test_kda.py | 8 +++++++- .../layers/mamba/ops/gather_initial_states.py | 2 +- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/models/kimi_k3/test_kda.py b/tests/models/kimi_k3/test_kda.py index 9fae1c2bcd3d..8ea5744f3b8b 100644 --- a/tests/models/kimi_k3/test_kda.py +++ b/tests/models/kimi_k3/test_kda.py @@ -30,9 +30,15 @@ fused_recurrent_kda_fwd, fused_recurrent_kda_packed_decode, ) +from vllm.platforms import current_platform from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd -DEVICE = "cuda" +DEVICE = current_platform.device_type + +pytestmark = pytest.mark.skipif( + not (current_platform.is_cuda_alike() or current_platform.is_xpu()), + reason="The KDA kernels require a CUDA-alike or XPU device.", +) # The AMD and NVIDIA copies of the KDA kernels are vendored separately and are # free to diverge, so the shared-semantics tests below run against both. diff --git a/vllm/model_executor/layers/mamba/ops/gather_initial_states.py b/vllm/model_executor/layers/mamba/ops/gather_initial_states.py index b952e3ebbce0..02c89823ee06 100644 --- a/vllm/model_executor/layers/mamba/ops/gather_initial_states.py +++ b/vllm/model_executor/layers/mamba/ops/gather_initial_states.py @@ -49,7 +49,7 @@ def gather_initial_states( ) -> torch.Tensor: """Gather dense state rows, replacing uninitialized rows with zeros.""" assert state.ndim >= 2 - assert state.is_cuda + assert state.is_cuda or state.is_xpu assert indices.ndim == 1 and has_initial_state.ndim == 1 assert indices.shape == has_initial_state.shape assert indices.device == state.device From 95901ce70a625cd152e465d1b15441fd199bde24 Mon Sep 17 00:00:00 2001 From: Ola <114643959+030611@users.noreply.github.com> Date: Mon, 17 Aug 2026 18:30:35 +0800 Subject: [PATCH 043/839] fix(pooling): validate BGE-M3 combined task ownership (#51823) Signed-off-by: Zhe Li <2843409461@qq.com> Co-authored-by: OpenAI Codex --- docs/models/pooling_models/specific_models.md | 95 +++++++++++++++---- tests/entrypoints/pooling/test_factories.py | 94 ++++++++++++++++++ vllm/entrypoints/pooling/factories.py | 10 +- .../pooling/pooling/io_processor.py | 12 +++ 4 files changed, 193 insertions(+), 18 deletions(-) create mode 100644 tests/entrypoints/pooling/test_factories.py diff --git a/docs/models/pooling_models/specific_models.md b/docs/models/pooling_models/specific_models.md index 8753f1fd07c3..0c7b0a1bf3fc 100644 --- a/docs/models/pooling_models/specific_models.md +++ b/docs/models/pooling_models/specific_models.md @@ -365,36 +365,99 @@ curl -s http://localhost:8000/rerank -H "Content-Type: application/json" -d '{ ## BAAI/bge-m3 -The `BAAI/bge-m3` model comes with extra weights for sparse and colbert embeddings but unfortunately in its `config.json` -the architecture is declared as `XLMRobertaModel`, which makes `vLLM` load it as a vanilla ROBERTA model without the -extra weights. To load the full model weights, override its architecture like this: +`BAAI/bge-m3` supports dense retrieval, lexical matching, and ColBERT-style +multi-vector retrieval. Its `config.json` declares `XLMRobertaModel`, so vLLM +otherwise loads it as a vanilla RoBERTa model without the extra sparse and +ColBERT weights. The examples below therefore override the architecture with +`BgeM3EmbeddingModel`. + +The three retrieval modes map to concrete pooling tasks as follows: + +| Retrieval mode | Pooling task | Output | +| -------------- | ------------ | ------ | +| Dense | `embed` | One embedding vector per input | +| Lexical/sparse | `token_classify` | One scalar weight per non-special token | +| ColBERT multi-vector | `token_embed` | One embedding vector per non-special token | + +Serve one concrete mode by selecting its task at load time: + +```shell +vllm serve BAAI/bge-m3 \ + --runner pooling \ + --hf-overrides '{"architectures": ["BgeM3EmbeddingModel"]}' \ + --pooler-config.task +``` + +For dense embeddings, replace `` with `embed` and use the Embeddings API: ```shell -vllm serve BAAI/bge-m3 --hf-overrides '{"architectures": ["BgeM3EmbeddingModel"]}' +curl -s http://localhost:8000/v1/embeddings \ + -H "Content-Type: application/json" -d '{ + "model": "BAAI/bge-m3", + "input": ["What is BGE M3?", "Definition of BM25"] + }' ``` -Then you obtain the sparse embeddings like this: +For lexical weights, replace `` with `token_classify` and use the +Pooling API: ```shell curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{ - "model": "BAAI/bge-m3", - "task": "token_classify", - "input": ["What is BGE M3?", "Definition of BM25"] + "model": "BAAI/bge-m3", + "task": "token_classify", + "input": ["What is BGE M3?", "Definition of BM25"] }' ``` Due to limitations in the output schema, the output consists of a list of -token scores for each token for each input. This means that you'll have to call -`/tokenize` as well to be able to pair tokens with scores. -Refer to the tests in `tests/models/language/pooling/test_bge_m3.py` to see how -to do that. +token scores for each input. Call `/tokenize` as well to pair token IDs with +their scores. See +[`test_bge_m3.py`](../../../tests/models/language/pooling/test_bge_m3.py) for a +complete example that also combines repeated token IDs. + +For ColBERT vectors, replace `` with `token_embed` and use the Pooling API: + +```shell +curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{ + "model": "BAAI/bge-m3", + "task": "token_embed", + "input": ["What is BGE M3?", "Definition of BM25"] +}' +``` + +### Dense and sparse output through an IO processor plugin + +The source tree includes a reference +[BGE-M3 IO processor plugin](../../../tests/plugins/bge_m3_sparse_plugin) that +formats dense embeddings, sparse token weights, or both in one response. From +a source checkout, install it in the vLLM environment and load it as follows: -You can obtain the colbert embeddings like this: +```shell +uv pip install ./tests/plugins/bge_m3_sparse_plugin + +vllm serve BAAI/bge-m3 \ + --runner pooling \ + --hf-overrides '{"architectures": ["BgeM3EmbeddingModel"]}' \ + --io-processor-plugin bge_m3_sparse_plugin +``` + +The plugin selects the internal `embed&token_classify` task so the model +computes dense and lexical outputs together. Public requests must use task +`plugin` and put the plugin-specific fields under `data`: ```shell curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{ - "model": "BAAI/bge-m3", - "task": "token_embed", - "input": ["What is BGE M3?", "Definition of BM25"] + "model": "BAAI/bge-m3", + "task": "plugin", + "data": { + "input": ["What is BGE M3?", "Definition of BM25"], + "embed_task": "dense&sparse", + "return_tokens": true + } }' ``` + +`embed_task` accepts `dense`, `sparse`, or `dense&sparse`. The combined +`embed&token_classify` task is an internal execution contract for this plugin, +not a generic Pooling API response format. Without the plugin, select one of +the three concrete tasks above. diff --git a/tests/entrypoints/pooling/test_factories.py b/tests/entrypoints/pooling/test_factories.py new file mode 100644 index 000000000000..8972ac49c927 --- /dev/null +++ b/tests/entrypoints/pooling/test_factories.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from unittest.mock import MagicMock + +import pytest + +from vllm.entrypoints.pooling import factories +from vllm.entrypoints.pooling.factories import init_pooling_io_processors +from vllm.entrypoints.pooling.pooling.io_processor import ( + PluginWithIOProcessorPlugins, + UnsupportedCombinedTaskIOProcessor, +) + + +def _bge_m3_config(io_processor_plugin=None): + model_config = MagicMock() + model_config.get_pooling_task.return_value = "embed&token_classify" + model_config.io_processor_plugin = io_processor_plugin + model_config.hf_config.to_dict.return_value = {} + model_config.architecture = "BgeM3EmbeddingModel" + + vllm_config = MagicMock(model_config=model_config) + renderer = MagicMock() + renderer._executor = MagicMock() + chat_template_config = MagicMock( + chat_template=None, + chat_template_content_format="auto", + trust_request_chat_template=False, + ) + return vllm_config, renderer, chat_template_config + + +def test_combined_task_without_plugin_uses_rejection_processor(): + vllm_config, renderer, chat_template_config = _bge_m3_config() + + processors = init_pooling_io_processors( + supported_tasks=("embed", "embed&token_classify"), + vllm_config=vllm_config, + renderer=renderer, + chat_template_config=chat_template_config, + ) + + assert processors.keys() == {"embed&token_classify"} + assert isinstance( + processors["embed&token_classify"], UnsupportedCombinedTaskIOProcessor + ) + + +def test_combined_task_with_plugin_uses_plugin_processor(monkeypatch): + vllm_config, renderer, chat_template_config = _bge_m3_config("bge_m3_sparse_plugin") + monkeypatch.setattr(factories, "has_io_processor", lambda *_: True) + monkeypatch.setattr( + "vllm.entrypoints.pooling.pooling.io_processor.get_io_processor", + lambda *_: MagicMock(), + ) + + processors = init_pooling_io_processors( + supported_tasks=("embed", "embed&token_classify"), + vllm_config=vllm_config, + renderer=renderer, + chat_template_config=chat_template_config, + ) + + assert processors.keys() == {"embed&token_classify", "plugin"} + assert isinstance(processors["plugin"], PluginWithIOProcessorPlugins) + + +def test_combined_task_plain_pooling_request_has_actionable_error(monkeypatch): + from vllm.entrypoints.pooling.pooling.protocol import PoolingCompletionRequest + from vllm.entrypoints.pooling.pooling.serving import ServingPooling + + vllm_config, renderer, chat_template_config = _bge_m3_config("bge_m3_sparse_plugin") + monkeypatch.setattr(factories, "has_io_processor", lambda *_: True) + monkeypatch.setattr( + "vllm.entrypoints.pooling.pooling.io_processor.get_io_processor", + lambda *_: MagicMock(), + ) + + engine_client = MagicMock(renderer=renderer, vllm_config=vllm_config) + models = MagicMock(model_config=vllm_config.model_config) + serving = ServingPooling( + engine_client, + models, + supported_tasks=("embed", "embed&token_classify"), + request_logger=None, + chat_template_config=chat_template_config, + ) + request = PoolingCompletionRequest(model="BAAI/bge-m3", input=["hola"]) + + assert serving.io_processors.keys() == {"embed&token_classify", "plugin"} + io_processor = serving.get_io_processor(request) + with pytest.raises(ValueError, match="plugin request with a 'data' field"): + io_processor.create_pooling_params(request) diff --git a/vllm/entrypoints/pooling/factories.py b/vllm/entrypoints/pooling/factories.py index dd3d873b3116..c7615054dd26 100644 --- a/vllm/entrypoints/pooling/factories.py +++ b/vllm/entrypoints/pooling/factories.py @@ -65,10 +65,16 @@ def init_pooling_io_processors( processors["token_embed"] = TokenEmbedIOProcessor - if has_io_processor( + if pooling_task == "embed&token_classify": + from .pooling.io_processor import UnsupportedCombinedTaskIOProcessor + + processors[pooling_task] = UnsupportedCombinedTaskIOProcessor + + has_plugin = has_io_processor( vllm_config, model_config.io_processor_plugin, - ): + ) + if has_plugin: from .pooling.io_processor import PluginWithIOProcessorPlugins processors["plugin"] = PluginWithIOProcessorPlugins diff --git a/vllm/entrypoints/pooling/pooling/io_processor.py b/vllm/entrypoints/pooling/pooling/io_processor.py index ecc356c477a8..4b18cfc9e184 100644 --- a/vllm/entrypoints/pooling/pooling/io_processor.py +++ b/vllm/entrypoints/pooling/pooling/io_processor.py @@ -24,6 +24,18 @@ logger = init_logger(__name__) +class UnsupportedCombinedTaskIOProcessor(PoolingIOProcessor): + name = "embed&token_classify" + + def create_pooling_params(self, request): + raise ValueError( + "The 'embed&token_classify' pooling task is only available " + "through an IO processor plugin. Send a plugin request with " + "a 'data' field, " + "or select a concrete task with --pooler-config.task." + ) + + class PluginWithoutIOProcessorPlugins(PoolingIOProcessor): # Some models, such as Terratorch (tests/models/test_terratorch.py), # use plugin tasks in the pooler but do not use IO Processor plugins. From f27ae25473af7520dde3f8b9e0705041940be0c9 Mon Sep 17 00:00:00 2001 From: Ganesh R Date: Mon, 17 Aug 2026 16:15:59 +0530 Subject: [PATCH 044/839] [Bugfix][CPU] Take an attention group's query head count from its layers (#51852) Signed-off-by: Ganesh R Signed-off-by: R Co-authored-by: Cursor Co-authored-by: Li, Jiang --- .buildkite/hardware_tests/cpu.yaml | 2 + tests/v1/attention/test_group_head_counts.py | 82 ++++++++++++++++++++ vllm/v1/attention/backends/cpu_attn.py | 16 ++-- 3 files changed, 92 insertions(+), 8 deletions(-) create mode 100644 tests/v1/attention/test_group_head_counts.py diff --git a/.buildkite/hardware_tests/cpu.yaml b/.buildkite/hardware_tests/cpu.yaml index 9b9f70b58d13..efd3e30f6f89 100644 --- a/.buildkite/hardware_tests/cpu.yaml +++ b/.buildkite/hardware_tests/cpu.yaml @@ -11,6 +11,7 @@ steps: - CMakeLists.txt - vllm/_custom_ops.py - tests/kernels/attention/test_cpu_attn.py + - tests/v1/attention/test_group_head_counts.py - tests/kernels/moe/test_cpu_fused_moe.py - tests/kernels/moe/test_cpu_quant_fused_moe.py - tests/kernels/test_onednn.py @@ -27,6 +28,7 @@ steps: - | bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m " pytest -x -v -s tests/kernels/attention/test_cpu_attn.py + pytest -x -v -s tests/v1/attention/test_group_head_counts.py pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py pytest -x -v -s tests/kernels/moe/test_cpu_quant_fused_moe.py pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py diff --git a/tests/v1/attention/test_group_head_counts.py b/tests/v1/attention/test_group_head_counts.py new file mode 100644 index 000000000000..b799bccc29d6 --- /dev/null +++ b/tests/v1/attention/test_group_head_counts.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Scheduler metadata sizes a scratchpad from the query head count, so it must +come from the builder's own group: the model-wide ``get_num_attention_heads()`` +is wrong for models that vary it per layer (e.g. Laguna), and too small a +scratchpad is indexed past its end. +""" + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from vllm.platforms import current_platform +from vllm.v1.attention.backends.cpu_attn import ( + CPUAttentionBackendImpl, + CPUAttentionMetadataBuilder, +) + +pytestmark = pytest.mark.skipif( + not current_platform.is_cpu(), reason="CPU attention backend" +) + +# Laguna's shape: 48 query heads model-wide, 64 on its sliding layers, both +# against 8 KV heads. +MODEL_WIDE_NUM_HEADS = 48 +NUM_KV_HEADS = 8 + + +def _layers(layer_num_heads: list[int]): + """Stand-in attention layers, one per head count, as one attention group.""" + return { + f"layer_{i}": SimpleNamespace( + impl=MagicMock( + spec=CPUAttentionBackendImpl, + num_heads=num_heads, + sliding_window=None, + ) + ) + for i, num_heads in enumerate(layer_num_heads) + } + + +def _build(layer_num_heads: list[int]) -> CPUAttentionMetadataBuilder: + layers = _layers(layer_num_heads) + vllm_config = MagicMock() + vllm_config.model_config.dtype = torch.bfloat16 + vllm_config.model_config.get_num_attention_heads.return_value = MODEL_WIDE_NUM_HEADS + vllm_config.cache_config.block_size = 16 + vllm_config.cache_config.cache_dtype = "auto" + kv_cache_spec = SimpleNamespace(num_kv_heads=NUM_KV_HEADS, head_size=64) + + with ( + patch( + "vllm.v1.attention.backends.utils.get_layers_from_vllm_config", + return_value=layers, + ), + patch( + "vllm.v1.attention.backends.cpu_attn.get_layers_from_vllm_config", + return_value=layers, + ), + ): + return CPUAttentionMetadataBuilder( + kv_cache_spec=kv_cache_spec, + layer_names=list(layers), + vllm_config=vllm_config, + device=torch.device("cpu"), + ) + + +@pytest.mark.parametrize("group_num_heads", [MODEL_WIDE_NUM_HEADS, 64, 16]) +def test_num_heads_comes_from_the_group(group_num_heads): + """The group's own count wins, even when it is not the model-wide one.""" + builder = _build([group_num_heads, group_num_heads]) + assert builder.num_heads == group_num_heads + + +def test_mixed_head_counts_in_one_group_are_rejected(): + """Grouping guarantees uniformity; a mixed group means that broke.""" + with pytest.raises(AssertionError, match="share num_heads"): + _build([MODEL_WIDE_NUM_HEADS, 64]) diff --git a/vllm/v1/attention/backends/cpu_attn.py b/vllm/v1/attention/backends/cpu_attn.py index 3d09d789f3cc..a6cd254fea06 100644 --- a/vllm/v1/attention/backends/cpu_attn.py +++ b/vllm/v1/attention/backends/cpu_attn.py @@ -31,6 +31,7 @@ ) from vllm.v1.attention.backends.utils import ( KVCacheLayoutType, + get_num_attention_heads_from_layers, ) from vllm.v1.kv_cache_interface import ( AttentionSpec, @@ -151,13 +152,15 @@ def __init__( parallel_config = vllm_config.parallel_config self.num_kv_heads = kv_cache_spec.num_kv_heads - self.num_heads = vllm_config.model_config.get_num_attention_heads( - parallel_config - ) + # The scheduler metadata built here sizes a scratchpad from the query + # head count, so it must come from this group's layers: the model-wide + # count is wrong for models that vary it per layer (e.g. Laguna). + self.num_heads = get_num_attention_heads_from_layers( + vllm_config, layer_names + ) or vllm_config.model_config.get_num_attention_heads(parallel_config) self.head_dim = kv_cache_spec.head_size self.dtype = vllm_config.model_config.dtype - # Resolved from the layers on the first build(), once they exist. - self.window_size: int | None = None + self.window_size = self._group_sliding_window() self.block_size = vllm_config.cache_config.block_size self.kv_cache_dtype = vllm_config.cache_config.cache_dtype self.isa = _get_attn_isa( @@ -198,9 +201,6 @@ def build( common_attn_metadata: CommonAttentionMetadata, fast_build: bool = False, ) -> CPUAttentionMetadata: - if self.window_size is None: - self.window_size = self._group_sliding_window() - num_reqs = common_attn_metadata.num_reqs num_actual_tokens = common_attn_metadata.num_actual_tokens max_query_len = common_attn_metadata.max_query_len From 70afdedc1081d28c3eaae53bece8292298484c86 Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Mon, 17 Aug 2026 19:34:14 +0800 Subject: [PATCH 045/839] [K3] support recoverssm for K3 (#51855) Signed-off-by: zjy0516 Signed-off-by: Benjamin Chislett Co-authored-by: OpenAI Codex Co-authored-by: Benjamin Chislett Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/models/kimi_k3/test_kda.py | 342 ++++++ tests/models/kimi_k3/test_kda_metadata.py | 149 ++- tests/models/test_registry.py | 4 +- tests/test_config.py | 47 + .../worker/test_mamba_hybrid_model_state.py | 67 ++ tests/v1/worker/test_mamba_utils.py | 36 + vllm/config/cache.py | 7 +- vllm/config/vllm.py | 41 +- vllm/model_executor/layers/mamba/abstract.py | 8 +- .../layers/mamba/mamba_utils.py | 30 + vllm/model_executor/models/interfaces.py | 10 +- vllm/models/kimi_k3/nvidia/kda.py | 95 +- vllm/models/kimi_k3/nvidia/kda_metadata.py | 188 ++- vllm/models/kimi_k3/nvidia/model.py | 33 +- vllm/models/kimi_k3/nvidia/ops/recoverssm.py | 1067 +++++++++++++++++ .../attention/backends/recoverssm_metadata.py | 26 + .../worker/gpu/model_states/mamba_hybrid.py | 50 +- vllm/v1/worker/gpu/model_states/recoverssm.py | 101 ++ vllm/v1/worker/mamba_utils.py | 9 +- 19 files changed, 2235 insertions(+), 75 deletions(-) create mode 100644 vllm/models/kimi_k3/nvidia/ops/recoverssm.py create mode 100644 vllm/v1/attention/backends/recoverssm_metadata.py create mode 100644 vllm/v1/worker/gpu/model_states/recoverssm.py diff --git a/tests/models/kimi_k3/test_kda.py b/tests/models/kimi_k3/test_kda.py index 8ea5744f3b8b..7416cca02a32 100644 --- a/tests/models/kimi_k3/test_kda.py +++ b/tests/models/kimi_k3/test_kda.py @@ -6,6 +6,8 @@ Uses torch.rand for q/k/v to match FLA's test pattern. """ +from types import SimpleNamespace + import pytest import torch import torch.nn.functional as F @@ -22,6 +24,12 @@ is_flashkda_supported, is_fused_kda_decode_supported, ) +from vllm.models.kimi_k3.nvidia.model import KimiLinearForCausalLM +from vllm.models.kimi_k3.nvidia.ops import recoverssm as recoverssm_ops +from vllm.models.kimi_k3.nvidia.ops.recoverssm import ( + KDARecoverSSMCommitContext, + kda_recoverssm_verify, +) from vllm.models.kimi_k3.nvidia.ops.third_party.kda import ( chunk_kda, chunk_kda_with_fused_gate, @@ -48,6 +56,38 @@ } +def test_kda_recoverssm_config_state_layout(): + vllm_config = SimpleNamespace( + model_config=SimpleNamespace( + dtype=torch.bfloat16, + hf_config=SimpleNamespace( + linear_attn_config={ + "num_heads": 4, + "head_dim": 32, + "short_conv_kernel_size": 4, + } + ), + ), + cache_config=SimpleNamespace( + mamba_cache_dtype="auto", + use_kda_recoverssm=True, + ), + parallel_config=SimpleNamespace(tensor_parallel_size=1), + speculative_config=SimpleNamespace(num_speculative_tokens=2), + ) + + assert KimiLinearForCausalLM.get_mamba_state_dtype_from_config(vllm_config) == ( + torch.bfloat16, + torch.float32, + torch.float32, + torch.bfloat16, + ) + assert KimiLinearForCausalLM.get_mamba_state_shape_from_config(vllm_config)[2:] == ( + (4, 3, 32), + (4, 3, 64), + ) + + @torch.inference_mode() def test_gather_initial_states_correctness(): row_size = 8 * 128 * 128 @@ -535,6 +575,308 @@ def test_kda_spec_decode_correctness( assert torch.isnan(output_storage[..., H * D :]).all() +@pytest.mark.parametrize( + ( + "conv_state_dim_first", + "use_request_indices", + "lower_bound", + "align_mode", + ), + [ + pytest.param(False, False, None, False, id="baseline"), + pytest.param(True, True, -5.0, True, id="all-features"), + pytest.param(False, True, -5.0, False, id="request-indexed"), + pytest.param(True, False, None, True, id="aligned"), + ], +) +@torch.inference_mode() +def test_kda_recoverssm_verify_and_group_commit( + monkeypatch: pytest.MonkeyPatch, + lower_bound: float | None, + use_request_indices: bool, + conv_state_dim_first: bool, + align_mode: bool, +): + monkeypatch.setattr( + recoverssm_ops, + "is_conv_state_dim_first", + lambda: conv_state_dim_first, + ) + num_layers, num_seqs, query_len = 2, 2, 8 + num_blocks, num_heads, dim = (7 if align_mode else 3), 4, 128 + total_tokens = num_seqs * query_len + torch.manual_seed(20260808) + + q, k, v, raw_g = [ + torch.randn( + 1, + total_tokens, + num_heads, + dim, + dtype=torch.bfloat16, + device=DEVICE, + ) + for _ in range(4) + ] + raw_beta = torch.randn( + 1, + total_tokens, + num_heads, + dtype=torch.bfloat16, + device=DEVICE, + ) + query_start_loc = torch.arange( + 0, + total_tokens + 1, + query_len, + dtype=torch.int32, + device=DEVICE, + ) + state_indices = torch.tensor( + [5, 6] if align_mode else [1, 2], dtype=torch.int32, device=DEVICE + ) + accepted = [2, 8] + if use_request_indices: + global_num_accepted = torch.tensor( + [0, accepted[0], 0, accepted[1]], + dtype=torch.int32, + device=DEVICE, + ) + request_indices = torch.tensor([1, 3], dtype=torch.int32, device=DEVICE) + else: + global_num_accepted = torch.tensor(accepted, dtype=torch.int32, device=DEVICE) + request_indices = None + + block_table = None + num_computed_tokens = None + mamba_block_size = None + if align_mode: + batch_size = 4 if use_request_indices else num_seqs + block_table = torch.full((batch_size, 2), -1, dtype=torch.int32, device=DEVICE) + rows = ( + request_indices + if request_indices is not None + else torch.arange(num_seqs, device=DEVICE) + ) + block_table[rows] = torch.tensor( + [[1, 5], [2, 6]], + dtype=torch.int32, + device=DEVICE, + ) + num_computed_tokens = torch.zeros(batch_size, dtype=torch.int32, device=DEVICE) + num_computed_tokens[rows] = 4 + mamba_block_size = 8 + + layers = [] + expected_outputs = [] + expected_states = [] + initial_states = [] + initial_conv_states = [] + history_len, conv_dim = 3, 12 + for layer_idx in range(num_layers): + A_log = ( + 0.2 * torch.randn(num_heads, dtype=torch.float32, device=DEVICE) + + layer_idx * 0.03 + ).contiguous() + dt_bias = ( + 0.1 * torch.randn(num_heads, dim, dtype=torch.float32, device=DEVICE) + ).contiguous() + checkpoint = 0.01 * torch.randn( + num_blocks, + num_heads, + dim, + dim, + dtype=torch.float32, + device=DEVICE, + ) + conv_shape = ( + (num_blocks, conv_dim, history_len + query_len - 1) + if conv_state_dim_first + else (num_blocks, history_len + query_len - 1, conv_dim) + ) + conv_state = torch.randn(conv_shape, dtype=torch.bfloat16, device=DEVICE) + correction_cache = torch.empty( + num_blocks, + num_heads, + query_len, + dim, + dtype=torch.float32, + device=DEVICE, + ) + kg_cache = torch.empty( + num_blocks, + num_heads, + query_len, + 2 * dim, + dtype=torch.bfloat16, + device=DEVICE, + ) + layer = SimpleNamespace( + kv_cache=( + conv_state, + checkpoint, + correction_cache, + kg_cache, + ), + A_log=A_log, + dt_bias=dt_bias, + local_num_heads=num_heads, + head_dim=dim, + gate_lower_bound=lower_bound, + ) + layers.append(layer) + initial_states.append(checkpoint.clone()) + initial_conv_states.append(conv_state.clone()) + + actual_output = kda_recoverssm_verify( + q=q, + k=k, + v=v, + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + lower_bound=lower_bound, + checkpoint_state=checkpoint, + correction_cache=correction_cache, + kg_cache=kg_cache, + query_start_loc=query_start_loc, + state_indices=state_indices, + spec_query_len=query_len, + ) + + normalized_q = q.float() * torch.rsqrt( + q.float().square().sum(dim=-1, keepdim=True) + 1e-6 + ) + normalized_k = k.float() * torch.rsqrt( + k.float().square().sum(dim=-1, keepdim=True) + 1e-6 + ) + gate_input = raw_g.float() + dt_bias.view(1, 1, num_heads, dim) + if lower_bound is None: + gate = -A_log.exp().view(1, 1, num_heads, 1) * F.softplus(gate_input) + else: + gate = lower_bound * torch.sigmoid( + A_log.exp().view(1, 1, num_heads, 1) * gate_input + ) + beta = raw_beta.float().sigmoid() + + reference_output = [] + committed_states = checkpoint.clone() + for seq_idx, commit_len in enumerate(accepted): + start = seq_idx * query_len + end = start + query_len + output, _ = naive_recurrent_kda( + normalized_q[:, start:end], + normalized_k[:, start:end], + v[:, start:end], + gate[:, start:end], + beta[:, start:end], + initial_state=checkpoint[state_indices[seq_idx]].transpose(-1, -2), + ) + reference_output.append(output) + _, committed_state = naive_recurrent_kda( + normalized_q[:, start : start + commit_len], + normalized_k[:, start : start + commit_len], + v[:, start : start + commit_len], + gate[:, start : start + commit_len], + beta[:, start : start + commit_len], + initial_state=checkpoint[state_indices[seq_idx]].transpose(-1, -2), + output_final_state=True, + ) + assert committed_state is not None + final_block = state_indices[seq_idx] + if align_mode: + assert block_table is not None + row = request_indices[seq_idx] if use_request_indices else seq_idx + final_block = block_table[row, (4 + commit_len) // 8] + committed_states[final_block] = committed_state.transpose(-1, -2) + if align_mode and 4 + commit_len >= 8: + _, boundary_state = naive_recurrent_kda( + normalized_q[:, start : start + 4], + normalized_k[:, start : start + 4], + v[:, start : start + 4], + gate[:, start : start + 4], + beta[:, start : start + 4], + initial_state=checkpoint[state_indices[seq_idx]].transpose(-1, -2), + output_final_state=True, + ) + assert boundary_state is not None + assert block_table is not None + row = request_indices[seq_idx] if use_request_indices else seq_idx + committed_states[block_table[row, 0]] = boundary_state.transpose(-1, -2) + expected_outputs.append(torch.cat(reference_output, dim=1)) + expected_states.append(committed_states) + torch.testing.assert_close(checkpoint, initial_states[-1]) + torch.testing.assert_close( + actual_output, + expected_outputs[-1], + atol=3e-2, + rtol=3e-2, + ) + + context = KDARecoverSSMCommitContext.create( + layers, + spec_query_len=query_len, + max_num_reqs=global_num_accepted.shape[0], + ) + context.commit( + global_num_accepted, + state_indices, + query_start_loc, + request_indices=request_indices, + block_table=block_table, + num_computed_tokens=num_computed_tokens, + mamba_block_size=mamba_block_size, + ) + + for layer_idx, layer in enumerate(layers): + torch.testing.assert_close( + layer.kv_cache[1], + expected_states[layer_idx], + atol=3e-3, + rtol=3e-3, + ) + for seq_idx, commit_len in enumerate(accepted): + block = state_indices[seq_idx] + if align_mode: + assert block_table is not None + row = request_indices[seq_idx] if use_request_indices else seq_idx + block = block_table[row, (4 + commit_len) // 8] + source_block = state_indices[seq_idx] if align_mode else block + if conv_state_dim_first: + actual_conv = layer.kv_cache[0][block, :, :history_len] + expected_conv = initial_conv_states[layer_idx][ + source_block, + :, + commit_len - 1 : commit_len - 1 + history_len, + ] + else: + actual_conv = layer.kv_cache[0][block, :history_len] + expected_conv = initial_conv_states[layer_idx][ + source_block, + commit_len - 1 : commit_len - 1 + history_len, + ] + torch.testing.assert_close(actual_conv, expected_conv) + if align_mode and 4 + commit_len >= 8: + assert block_table is not None + boundary_block = block_table[row, 0] + if conv_state_dim_first: + actual_boundary_conv = layer.kv_cache[0][ + boundary_block, :, :history_len + ] + expected_boundary_conv = initial_conv_states[layer_idx][ + state_indices[seq_idx], :, 3 : 3 + history_len + ] + else: + actual_boundary_conv = layer.kv_cache[0][ + boundary_block, :history_len + ] + expected_boundary_conv = initial_conv_states[layer_idx][ + state_indices[seq_idx], 3 : 3 + history_len + ] + torch.testing.assert_close(actual_boundary_conv, expected_boundary_conv) + + @pytest.mark.parametrize( ("num_heads", "num_seqs", "lower_bound", "fuse_output_norm"), [ diff --git a/tests/models/kimi_k3/test_kda_metadata.py b/tests/models/kimi_k3/test_kda_metadata.py index 5352ef7a7a64..069c9a4d7a8e 100644 --- a/tests/models/kimi_k3/test_kda_metadata.py +++ b/tests/models/kimi_k3/test_kda_metadata.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from dataclasses import fields +from unittest.mock import Mock import pytest import torch @@ -26,6 +27,9 @@ GDNAttentionMetadata, GDNAttentionMetadataBuilder, ) +from vllm.v1.attention.backends.recoverssm_metadata import ( + RecoverSSMPostprocessMetadata, +) from vllm.v1.attention.backends.utils import ( NULL_BLOCK_ID, mamba_get_block_table_tensor, @@ -44,8 +48,12 @@ } -def _assert_matches_shared_gdn(reference, actual: KimiK3KDAMetadata): - for field in fields(KimiK3KDAMetadata): +def _assert_matches_shared_gdn( + reference: GDNAttentionMetadata, actual: KimiK3KDAMetadata +): + assert actual.recoverssm_commit is None + assert actual.recoverssm_context is None + for field in fields(GDNAttentionMetadata): actual_value = getattr(actual, field.name) expected_value = getattr(reference, field.name) if field.name in PRUNED_METADATA_FIELDS: @@ -78,6 +86,7 @@ def _make_builder( full_cuda_graph: bool, device: torch.device = DEVICE, mamba_cache_mode: str = "none", + use_recoverssm: bool = False, ) -> AttentionMetadataBuilder: vllm_config = create_vllm_config( model_name="Qwen/Qwen3.5-0.8B", @@ -92,17 +101,24 @@ def _make_builder( CUDAGraphMode.FULL_AND_PIECEWISE if full_cuda_graph else CUDAGraphMode.NONE ) vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode - return builder_cls( + vllm_config.cache_config.use_replayssm = use_recoverssm + vllm_config.cache_config.use_kda_recoverssm = use_recoverssm + builder = builder_cls( kv_cache_spec=MambaSpec( block_size=BLOCK_SIZE, shapes=((16, 64),), dtypes=(torch.float16,), - num_speculative_blocks=num_speculative_tokens, + mamba_cache_mode=mamba_cache_mode, + num_speculative_blocks=(0 if use_recoverssm else num_speculative_tokens), ), layer_names=["layer.0"], vllm_config=vllm_config, device=device, ) + if use_recoverssm: + assert isinstance(builder, KimiK3KDAMetadataBuilder) + builder.recoverssm_context = Mock() + return builder @pytest.mark.parametrize( @@ -244,6 +260,99 @@ def test_mixed_regular_and_spec_decode_excludes_request_padding(): torch.testing.assert_close(actual.spec_token_indx, torch.tensor([1, 2, 3])) +@pytest.mark.parametrize("mamba_cache_mode", ["none", "align"]) +def test_recoverssm_spec_uses_one_state_slot_and_current_window( + mamba_cache_mode: str, +): + if mamba_cache_mode == "align" and not torch.cuda.is_available(): + pytest.skip("align metadata construction requires CUDA") + device = torch.device("cuda") if mamba_cache_mode == "align" else DEVICE + batch = BatchSpec(seq_lens=[100, 65, 20], query_lens=[1, 1, 3]) + common_attn_metadata = create_common_attn_metadata( + batch, BLOCK_SIZE, device + ).replace(is_prefilling=torch.tensor([True, True, False])) + builder = _make_builder( + KimiK3KDAMetadataBuilder, + num_speculative_tokens=2, + full_cuda_graph=False, + device=device, + mamba_cache_mode=mamba_cache_mode, + use_recoverssm=True, + ) + assert isinstance(builder, KimiK3KDAMetadataBuilder) + context = builder.recoverssm_context + assert context is not None + actual = builder.build( + 0, + common_attn_metadata, + num_decode_draft_tokens_cpu=torch.tensor([-1, -1, 2], dtype=torch.int32), + num_accepted_tokens=torch.tensor([3, 2, 2], dtype=torch.int32, device=device), + ) + + assert actual.spec_state_indices_tensor is not None + assert actual.spec_state_indices_tensor.shape == (1, 1) + torch.testing.assert_close( + actual.num_accepted_tokens, + torch.ones(1, dtype=torch.int32, device=device), + ) + commit_metadata = actual.recoverssm_commit + assert commit_metadata is not None + torch.testing.assert_close( + commit_metadata.request_indices, + torch.tensor([2], dtype=torch.int32, device=device), + ) + assert actual.recoverssm_context is context + num_accepted_tokens = torch.tensor([3, 2, 1], dtype=torch.int32, device=device) + + postprocess = actual.commit_recoverssm_state(num_accepted_tokens) + + if mamba_cache_mode == "none": + assert commit_metadata.align is None + assert postprocess is None + else: + assert isinstance(postprocess, RecoverSSMPostprocessMetadata) + assert postprocess.num_spec_decodes == 1 + assert postprocess.request_indices is commit_metadata.request_indices + assert postprocess.block_table is common_attn_metadata.block_table_tensor + assert ( + postprocess.num_computed_tokens + is common_attn_metadata.compute_num_computed_tokens() + ) + assert postprocess.block_size == BLOCK_SIZE + args = context.commit.call_args.args + assert args[0] is num_accepted_tokens + torch.testing.assert_close(args[1], commit_metadata.state_indices[:, 0]) + torch.testing.assert_close(args[2], commit_metadata.query_start_loc) + + +def test_recoverssm_distinguishes_draftless_decode_from_one_token_prefill(): + batch = BatchSpec(seq_lens=[40, 30], query_lens=[1, 1]) + common_attn_metadata = create_common_attn_metadata( + batch, BLOCK_SIZE, DEVICE + ).replace(is_prefilling=torch.tensor([False, True])) + actual = _make_builder( + KimiK3KDAMetadataBuilder, + num_speculative_tokens=2, + full_cuda_graph=False, + use_recoverssm=True, + ).build( + 0, + common_attn_metadata, + num_decode_draft_tokens_cpu=torch.full((2,), -1, dtype=torch.int32), + num_accepted_tokens=torch.ones(2, dtype=torch.int32), + ) + + assert actual.num_spec_decodes == 1 + assert actual.num_decodes == 0 + assert actual.num_prefills == 1 + assert actual.spec_state_indices_tensor is not None + assert actual.spec_state_indices_tensor.shape == (1, 1) + torch.testing.assert_close( + actual.spec_query_start_loc, + torch.tensor([0, 1], dtype=torch.int32), + ) + + @pytest.mark.parametrize( ("seq_len", "expected_has_initial_state"), [ @@ -307,6 +416,38 @@ def test_kimi_k3_kda_cudagraph_capture_matches_shared_gdn(): _assert_matches_shared_gdn(reference, actual) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_recoverssm_spec_cudagraph_stages_one_checkpoint_per_request(): + device = torch.device("cuda") + batch = BatchSpec(seq_lens=[50, 30], query_lens=[3, 3]) + common_attn_metadata = create_common_attn_metadata( + batch, BLOCK_SIZE, device + ).replace(is_prefilling=torch.tensor([False, False])) + builder = _make_builder( + KimiK3KDAMetadataBuilder, + num_speculative_tokens=2, + full_cuda_graph=True, + device=device, + use_recoverssm=True, + ) + assert isinstance(builder, KimiK3KDAMetadataBuilder) + assert builder.spec_state_indices_tensor.shape == ( + builder.vllm_config.scheduler_config.max_num_seqs, + 1, + ) + actual = builder.build_for_cudagraph_capture(common_attn_metadata) + + assert actual.spec_state_indices_tensor is not None + assert actual.spec_state_indices_tensor.shape == (batch.batch_size, 1) + assert actual.num_accepted_tokens is not None + torch.testing.assert_close( + actual.num_accepted_tokens, + torch.ones(batch.batch_size, dtype=torch.int32, device=device), + ) + assert actual.recoverssm_commit is not None + assert actual.recoverssm_commit.request_indices is None + + def test_kimi_k3_kda_backend_uses_private_metadata_builder(): assert KimiK3KDAAttentionBackend.get_builder_cls() is KimiK3KDAMetadataBuilder assert KimiK3KDAAttentionBackend.is_ssm() diff --git a/tests/models/test_registry.py b/tests/models/test_registry.py index 6eda6075d2f8..70b8b18f76f6 100644 --- a/tests/models/test_registry.py +++ b/tests/models/test_registry.py @@ -143,8 +143,10 @@ def test_registry_is_pp(model_arch, is_pp, init_cuda): @pytest.mark.parametrize( "model_arch,supported", [ - # ReplaySSM is opt-in per model; only Nemotron-H sets the flag today. + # ReplaySSM is opt-in per model. ("NemotronHForCausalLM", True), + ("KimiLinearForCausalLM", not current_platform.is_rocm()), + ("KimiK3ForConditionalGeneration", not current_platform.is_rocm()), ("Mamba2ForCausalLM", False), ("Zamba2ForCausalLM", False), ], diff --git a/tests/test_config.py b/tests/test_config.py index 70c8728d25d4..a2797e52126a 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -29,6 +29,7 @@ from vllm.config.compilation import CompilationMode, CUDAGraphMode from vllm.config.kernel import IrOpPriorityConfig from vllm.config.load import LoadConfig +from vllm.config.mamba import MambaBackendEnum from vllm.config.utils import get_field from vllm.config.vllm import OPTIMIZATION_LEVEL_TO_CONFIG, OptimizationLevel from vllm.platforms import current_platform @@ -37,6 +38,52 @@ DEVICE_TYPE = current_platform.device_type +def test_kda_recoverssm_derivation_is_revalidated(): + config = SimpleNamespace( + cache_config=SimpleNamespace( + use_replayssm=True, + use_kda_recoverssm=False, + mamba_cache_mode="none", + ), + num_speculative_tokens=3, + model_config=SimpleNamespace( + supports_replayssm=True, + architecture="KimiLinearForCausalLM", + ), + mamba_config=SimpleNamespace( + backend=MambaBackendEnum.TRITON, + enable_stochastic_rounding=False, + ), + parallel_config=SimpleNamespace(pipeline_parallel_size=1), + kv_transfer_config=None, + use_v2_model_runner=True, + ) + + VllmConfig.validate_mamba_cached_kernel(config) + assert config.cache_config.use_replayssm + assert config.cache_config.use_kda_recoverssm + + config.cache_config.mamba_cache_mode = "align" + VllmConfig.validate_mamba_cached_kernel(config) + config.use_v2_model_runner = False + with pytest.raises(ValueError, match="VLLM_USE_V2_MODEL_RUNNER=1"): + VllmConfig.validate_mamba_cached_kernel(config) + config.use_v2_model_runner = True + config.cache_config.mamba_cache_mode = "all" + with pytest.raises(ValueError, match="only none and align"): + VllmConfig.validate_mamba_cached_kernel(config) + config.cache_config.mamba_cache_mode = "none" + + config.model_config.architecture = "NemotronHForCausalLM" + with pytest.raises(ValueError, match="only supported for Kimi-K3 KDA"): + VllmConfig.validate_mamba_cached_kernel(config) + + config.model_config.architecture = "KimiLinearForCausalLM" + config.parallel_config.pipeline_parallel_size = 2 + with pytest.raises(ValueError, match="pipeline_parallel_size=1"): + VllmConfig.validate_mamba_cached_kernel(config) + + def test_compile_config_repr_succeeds(): # setup: VllmBackend mutates the config object config = VllmConfig() diff --git a/tests/v1/worker/test_mamba_hybrid_model_state.py b/tests/v1/worker/test_mamba_hybrid_model_state.py index 545d23f90912..749821274318 100644 --- a/tests/v1/worker/test_mamba_hybrid_model_state.py +++ b/tests/v1/worker/test_mamba_hybrid_model_state.py @@ -1,11 +1,19 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace +from unittest.mock import Mock + import pytest import torch from vllm.platforms import current_platform +from vllm.v1.attention.backends.recoverssm_metadata import ( + RecoverSSMMetadata, + RecoverSSMPostprocessMetadata, +) from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState +from vllm.v1.worker.gpu.model_states.recoverssm import RecoverSSMState @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") @@ -18,6 +26,7 @@ def test_postprocess_state_scalar_with_int32_mapping( (4,), 9, dtype=torch.int32, device="cuda" ) state._align_mode = False + state.recoverssm = None state._mamba_ctx = None idx_mapping = torch.tensor([2, -1, 0], dtype=torch.int32, device="cuda") @@ -27,3 +36,61 @@ def test_postprocess_state_scalar_with_int32_mapping( [expected_value, 9, expected_value, 9], dtype=torch.int32, device="cuda" ) torch.testing.assert_close(state.num_accepted_tokens_gpu, expected) + + +def test_recoverssm_commits_accepted_window_after_v2_sampling() -> None: + state = RecoverSSMState() + metadata = Mock(spec=RecoverSSMMetadata) + metadata.commit_recoverssm_state.return_value = None + num_sampled = torch.tensor([3, 1], dtype=torch.int32) + idx_mapping = torch.tensor([0, 1], dtype=torch.int32) + num_accepted_tokens = torch.ones(2, dtype=torch.int32) + group = SimpleNamespace(layer_names=["layer"]) + + state.record_step({"layer": metadata}, [[group]], for_capture=False) + state.commit_step( + num_sampled, + idx_mapping, + state_indices=None, + num_accepted_tokens=num_accepted_tokens, + ) + state.commit_step( + num_sampled, + idx_mapping, + state_indices=None, + num_accepted_tokens=num_accepted_tokens, + ) + + metadata.commit_recoverssm_state.assert_called_once_with(num_sampled) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") +def test_recoverssm_align_tracks_mixed_batch_state_and_neutralizes_copy_bias() -> None: + state = object.__new__(MambaHybridModelState) + state._align_mode = True + state._mamba_ctx = None + state._mamba_state_idx_gpu = torch.full((5,), -1, dtype=torch.int32, device="cuda") + state.recoverssm = RecoverSSMState() + state.num_accepted_tokens_gpu = torch.full( + (5,), 9, dtype=torch.int32, device="cuda" + ) + metadata = Mock(spec=RecoverSSMMetadata) + metadata.commit_recoverssm_state.return_value = RecoverSSMPostprocessMetadata( + num_spec_decodes=1, + request_indices=torch.tensor([1], dtype=torch.int32, device="cuda"), + num_computed_tokens=torch.tensor([6, 7], dtype=torch.int32, device="cuda"), + block_size=8, + block_table=torch.zeros((2, 4), dtype=torch.int32, device="cuda"), + ) + num_sampled = torch.tensor([2, 3], dtype=torch.int32, device="cuda") + idx_mapping = torch.tensor([3, 1], dtype=torch.int32, device="cuda") + group = SimpleNamespace(layer_names=["layer"]) + + state.recoverssm.record_step({"layer": metadata}, [[group]], for_capture=False) + + state.postprocess_state(idx_mapping, num_sampled) + + expected_state_indices = [-1, 1, -1, -1, -1] + assert state._mamba_state_idx_gpu.tolist() == expected_state_indices + expected_accepted = [9, 1, 9, 2, 9] + assert state.num_accepted_tokens_gpu.tolist() == expected_accepted diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index 2ba1ce1c5937..0534795f9e84 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -490,6 +490,42 @@ def test_stage_postprocess_inputs_to_gpu_asserts_on_missing_state_idx(): ) +def test_gpu_context_ignores_auxiliary_cache_tensors() -> None: + device = torch.device("cpu") + config = _TestConfig(num_layers=1) + layer_names = ["layer_0"] + kv_cache_config = _make_kv_cache_config(config, layer_names) + conv_state = torch.empty( + config.num_blocks, + config.conv_width, + config.conv_inner_dim, + dtype=config.dtype, + ) + temporal_state = torch.empty( + config.num_blocks, config.temporal_state_dim, dtype=config.dtype + ) + attention = MagicMock() + attention.kv_cache = [ + conv_state, + temporal_state, + *(torch.empty(config.num_blocks, 1) for _ in range(4)), + ] + context = _make_gpu_ctx(config, kv_cache_config, device) + + context.initialize_from_forward_context( + kv_cache_config, + {"layer_0": attention}, + _COPY_FUNCS, + [torch.zeros(1, 1, dtype=torch.int32)], + ) + + assert context.is_initialized + assert context.state_base_addrs.tolist() == [ + conv_state.data_ptr(), + temporal_state.data_ptr(), + ] + + def _run_gpu_postprocess( gpu_ctx: MambaSpecDecodeGPUContext, *, diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 654e9a53589b..0981f710ea57 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -144,9 +144,8 @@ class CacheConfig: caching is enabled. """ replayssm_buffer_len: int = Field(default=16, gt=0) - """ReplaySSM history buffer length B: with use_replayssm, standard decode - caches recent SSM inputs in a size-B ring buffer and flushes the checkpoint - state to HBM every B steps. Default 16.""" + """ReplaySSM history buffer length B for standard Mamba2 decode. Kimi-K3 + speculative decoding does not use B. Default 16.""" use_replayssm: bool = False """Use the ReplaySSM Mamba2 decode kernel: cache recent SSM inputs and skip the per-step full-state store, writing the checkpoint back only on flush. @@ -154,6 +153,8 @@ class CacheConfig: mamba backend; standard (non-speculative) decode only. In align mode flushes are most efficient when mamba_block_size is a multiple of replayssm_buffer_len, but this is not required.""" + use_kda_recoverssm: bool = field(default=False, init=False) + """Whether Kimi-K3 KDA uses RecoverSSM speculative decode.""" # Will be set after profiling. num_gpu_blocks: int | None = field(default=None, init=False) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 0a4c5cad8838..a2d6b293bb12 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -2540,22 +2540,47 @@ def validate_mamba_block_size(self) -> "VllmConfig": @model_validator(mode="after") def validate_mamba_cached_kernel(self) -> "VllmConfig": if not self.cache_config.use_replayssm: + self.cache_config.use_kda_recoverssm = False return self - # ReplaySSM adds a 3-tensor ring to the mamba state; only models that - # opt in (supports_replayssm) build a consistent shape on both the layer - # and config paths. Reject others so the mamba page size cannot desync. + self.cache_config.use_kda_recoverssm = self.num_speculative_tokens > 0 + if self.model_config is not None and not self.model_config.supports_replayssm: raise ValueError( - "--use-replayssm is only supported for Nemotron-H models " - f"(got architecture {self.model_config.architecture!r})" + "--use-replayssm is not supported for architecture " + f"{self.model_config.architecture!r}" ) - if self.cache_config.mamba_cache_mode == "all": + if self.cache_config.use_kda_recoverssm: + if self.model_config is not None and self.model_config.architecture not in ( + "KimiLinearForCausalLM", + "KimiK3ForConditionalGeneration", + ): + raise ValueError("RecoverSSM is only supported for Kimi-K3 KDA") + if self.mamba_config.enable_stochastic_rounding: + raise ValueError( + "RecoverSSM supports bfloat16/float32 " + "SSM state caches, not --enable-mamba-cache-stochastic-" + "rounding, which requires an explicit float16 cache" + ) + if self.cache_config.mamba_cache_mode not in ("none", "align"): + raise ValueError( + "RecoverSSM supports only none and align Mamba cache modes" + ) + if ( + self.cache_config.mamba_cache_mode == "align" + and not self.use_v2_model_runner + ): + raise ValueError( + "RecoverSSM with align mode requires VLLM_USE_V2_MODEL_RUNNER=1" + ) + if self.parallel_config.pipeline_parallel_size > 1: + raise ValueError( + "RecoverSSM currently requires pipeline_parallel_size=1" + ) + elif self.cache_config.mamba_cache_mode == "all": raise ValueError( "--use-replayssm supports prefix caching only in align mode; " "pass --mamba-cache-mode align" ) - if self.num_speculative_tokens > 0: - raise ValueError("--use-replayssm does not support speculative decoding") if self.mamba_config.backend != MambaBackendEnum.TRITON: raise ValueError("--use-replayssm requires --mamba-backend triton") if ( diff --git a/vllm/model_executor/layers/mamba/abstract.py b/vllm/model_executor/layers/mamba/abstract.py index d06916f697c7..ef350237dc70 100644 --- a/vllm/model_executor/layers/mamba/abstract.py +++ b/vllm/model_executor/layers/mamba/abstract.py @@ -71,10 +71,12 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: page_size_padded=page_size_padded, mamba_type=self.mamba_type, mamba_cache_mode=vllm_config.cache_config.mamba_cache_mode, + # RecoverSSM verifies the whole window off one checkpoint, so it + # never writes the baseline's per-draft-token state slots. num_speculative_blocks=( - vllm_config.speculative_config.num_speculative_tokens - if vllm_config.speculative_config - else 0 + 0 + if vllm_config.cache_config.use_kda_recoverssm + else vllm_config.num_speculative_tokens ), ) diff --git a/vllm/model_executor/layers/mamba/mamba_utils.py b/vllm/model_executor/layers/mamba/mamba_utils.py index 50f6e059c0f4..18a73c6554e6 100644 --- a/vllm/model_executor/layers/mamba/mamba_utils.py +++ b/vllm/model_executor/layers/mamba/mamba_utils.py @@ -136,6 +136,15 @@ def kda_state_dtype( state_dtype = get_kv_cache_torch_dtype(mamba_cache_dtype, model_dtype) return (state_dtype, torch.float32) + @classmethod + def append_kda_recoverssm_record( + cls, + base_dtypes: tuple[torch.dtype, ...], + model_dtype: ModelDType | torch.dtype, + ) -> tuple[torch.dtype, ...]: + activation_dtype = get_kv_cache_torch_dtype("auto", model_dtype) + return (*base_dtypes, torch.float32, activation_dtype) + class MambaStateShapeCalculator: @classmethod @@ -293,6 +302,27 @@ def kda_state_shape( recurrent_state_shape = (divide(num_heads, tp_world_size), head_dim, head_dim) return (conv_state_shape, recurrent_state_shape) + @classmethod + def append_kda_recoverssm_record( + cls, + base_shapes: tuple[tuple[int, int], tuple[int, int, int]], + num_heads: int, + head_dim: int, + tp_world_size: int, + spec_query_len: int, + ) -> tuple[ + tuple[int, int], + tuple[int, int, int], + tuple[int, int, int], + tuple[int, int, int], + ]: + local_num_heads = divide(num_heads, tp_world_size) + return ( + *base_shapes, + (local_num_heads, spec_query_len, head_dim), + (local_num_heads, spec_query_len, 2 * head_dim), + ) + @dataclass class MambaCopySpec: diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 88f0a2dc6ffc..555051237336 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -75,6 +75,12 @@ | tuple[tuple[int, int, int]] | tuple[tuple[int, int], tuple[int, int]] | tuple[tuple[int, int], tuple[int, int, int]] + | tuple[ + tuple[int, int], + tuple[int, int, int], + tuple[int, int, int], + tuple[int, int, int], + ] ) @@ -1099,8 +1105,8 @@ def supports_mamba_prefix_caching( @runtime_checkable class SupportsReplaySSM(Protocol): - """The interface for models whose Mamba2 layers support ReplaySSM cached - standard decode. + """The interface for models whose recurrent layers support ReplaySSM + cached decode. This is currently experimental. """ diff --git a/vllm/models/kimi_k3/nvidia/kda.py b/vllm/models/kimi_k3/nvidia/kda.py index c17a8ec4e9a6..e680180f26fa 100644 --- a/vllm/models/kimi_k3/nvidia/kda.py +++ b/vllm/models/kimi_k3/nvidia/kda.py @@ -282,23 +282,37 @@ def get_attn_backend(self) -> type[AttentionBackend]: def get_state_dtype( self, - ) -> tuple[torch.dtype, torch.dtype]: + ) -> tuple[torch.dtype, ...]: if self.model_config is None or self.cache_config is None: raise ValueError("model_config and cache_config must be set") - return MambaStateDtypeCalculator.kda_state_dtype( + base_dtypes = MambaStateDtypeCalculator.kda_state_dtype( self.model_config.dtype, self.cache_config.mamba_cache_dtype ) + if self.cache_config.use_kda_recoverssm: + return MambaStateDtypeCalculator.append_kda_recoverssm_record( + base_dtypes, self.model_config.dtype + ) + return base_dtypes def get_state_shape( self, - ) -> tuple[tuple[int, ...], tuple[int, ...]]: - return MambaStateShapeCalculator.kda_state_shape( + ) -> tuple[tuple[int, ...], ...]: + base_shapes = MambaStateShapeCalculator.kda_state_shape( self.tp_size, self.num_heads, self.head_dim, conv_kernel_size=self.conv_size, num_spec=self.num_spec, ) + if self.cache_config.use_kda_recoverssm: + return MambaStateShapeCalculator.append_kda_recoverssm_record( + base_shapes, + self.num_heads, + self.head_dim, + tp_world_size=self.tp_size, + spec_query_len=1 + self.num_spec, + ) + return base_shapes def __init__( self, @@ -308,6 +322,12 @@ def __init__( run_gemm_rs: bool = False, ) -> None: super().__init__(config, vllm_config, prefix) + self.use_recoverssm = self.cache_config.use_kda_recoverssm + if self.cache_config.use_replayssm and not self.use_recoverssm: + raise ValueError( + "Kimi-K3 supports --use-replayssm only with speculative decoding" + ) + self.spec_query_len = 1 + self.num_spec kda_config = config.linear_attn_config # type: ignore[attr-defined] assert kda_config is not None, "linear_attn_config must be set" @@ -371,7 +391,7 @@ def __init__( self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1) # Keep a width-major copy for fused decode without changing the layout # consumed by the prefill and fallback decode kernels. - conv_state_dtype, _ = self.get_state_dtype() + conv_state_dtype = self.get_state_dtype()[0] decode_conv1d_weight = None if is_fused_kda_decode_supported( self.local_num_heads, @@ -567,7 +587,7 @@ def _forward( g1 = g1[:, :num_actual_tokens] beta = beta[:, :num_actual_tokens] - conv_state, recurrent_state = self.kv_cache + conv_state, recurrent_state, *recoverssm_records = self.kv_cache # The convolution kernels consume (..., dim, width - 1). if not is_conv_state_dim_first(): conv_state = conv_state.transpose(-1, -2) @@ -634,7 +654,11 @@ def _forward( assert spec_state_indices_tensor is not None assert spec_query_start_loc is not None spec_conv_indices = spec_state_indices_tensor[:, 0][: m.num_spec_decodes] - spec_max_query_len = spec_state_indices_tensor.size(-1) + spec_max_query_len = ( + self.spec_query_len + if self.use_recoverssm + else spec_state_indices_tensor.size(-1) + ) spec_conv_out = torch.empty_like(mixed_qkv_spec) mixed_qkv_spec = causal_conv1d_update( mixed_qkv_spec, @@ -659,21 +683,48 @@ def _forward( if m.num_prefills == 0 and m.num_decodes == 0 else None ) - core_attn_out_spec, _ = fused_recurrent_kda( - q=q_spec, - k=k_spec, - v=v_spec, - raw_g=g1_spec, - raw_beta=beta_spec, - A_log=self.A_log, - dt_bias=self.dt_bias, - lower_bound=self.gate_lower_bound, - initial_state=recurrent_state, - cu_seqlens=spec_cu_seqlens, - ssm_state_indices=spec_state_indices_tensor, - num_accepted_tokens=num_accepted_tokens, - out=spec_out, - ) + if self.use_recoverssm: + from vllm.models.kimi_k3.nvidia.ops.recoverssm import ( + kda_recoverssm_verify, + ) + + if len(recoverssm_records) != 2: + raise ValueError( + "KDA RecoverSSM requires correction and key/gate buffers" + ) + core_attn_out_spec = kda_recoverssm_verify( + q=q_spec, + k=k_spec, + v=v_spec, + raw_g=g1_spec, + raw_beta=beta_spec, + A_log=self.A_log, + dt_bias=self.dt_bias, + lower_bound=self.gate_lower_bound, + checkpoint_state=recurrent_state, + correction_cache=recoverssm_records[0], + kg_cache=recoverssm_records[1], + query_start_loc=spec_cu_seqlens, + state_indices=spec_state_indices_tensor[: m.num_spec_decodes, 0], + spec_query_len=self.spec_query_len, + out=spec_out, + ) + else: + core_attn_out_spec, _ = fused_recurrent_kda( + q=q_spec, + k=k_spec, + v=v_spec, + raw_g=g1_spec, + raw_beta=beta_spec, + A_log=self.A_log, + dt_bias=self.dt_bias, + lower_bound=self.gate_lower_bound, + initial_state=recurrent_state, + cu_seqlens=spec_cu_seqlens, + ssm_state_indices=spec_state_indices_tensor, + num_accepted_tokens=num_accepted_tokens, + out=spec_out, + ) # Prefill or plain-decode path. core_attn_out_non_spec = None diff --git a/vllm/models/kimi_k3/nvidia/kda_metadata.py b/vllm/models/kimi_k3/nvidia/kda_metadata.py index 0bb4864d2918..2154eaa3c51f 100644 --- a/vllm/models/kimi_k3/nvidia/kda_metadata.py +++ b/vllm/models/kimi_k3/nvidia/kda_metadata.py @@ -6,13 +6,18 @@ ``GDNAttentionMetadataBuilder``. Kimi-K3 builds the metadata required by its prefill KDA kernel internally, so this builder omits the shared FLA chunk metadata construction. + +For Kimi-K3 speculative decoding, ``--use-replayssm`` selects the simplified +RecoverSSM path implemented here instead of the Mamba2 ReplaySSM kernel. """ -from dataclasses import dataclass +from dataclasses import dataclass, field from functools import cache +from typing import TYPE_CHECKING import torch +from vllm.config import VllmConfig from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.utils.torch_utils import async_tensor_h2d @@ -22,6 +27,10 @@ GDNAttentionMetadata, GDNAttentionMetadataBuilder, ) +from vllm.v1.attention.backends.recoverssm_metadata import ( + RecoverSSMMetadata, + RecoverSSMPostprocessMetadata, +) from vllm.v1.attention.backends.utils import ( NULL_BLOCK_ID, compute_causal_conv1d_metadata, @@ -29,6 +38,11 @@ ) from vllm.v1.kv_cache_interface import MambaSpec +if TYPE_CHECKING: + from vllm.models.kimi_k3.nvidia.ops.recoverssm import ( + KDARecoverSSMCommitContext, + ) + @cache def _metadata_launch_pdl() -> bool: @@ -229,11 +243,103 @@ def stage_spec_decode_metadata( @dataclass -class KimiK3KDAMetadata(GDNAttentionMetadata): - pass +class KDARecoverSSMAlignMetadata: + block_table: torch.Tensor + num_computed_tokens: torch.Tensor + block_size: int + + +@dataclass +class KDARecoverSSMCommitMetadata: + state_indices: torch.Tensor + query_start_loc: torch.Tensor + request_indices: torch.Tensor | None + align: KDARecoverSSMAlignMetadata | None + + +@dataclass +class KimiK3KDAMetadata(GDNAttentionMetadata, RecoverSSMMetadata): + recoverssm_commit: KDARecoverSSMCommitMetadata | None = None + recoverssm_context: "KDARecoverSSMCommitContext | None" = field( + default=None, repr=False, compare=False + ) + + def commit_recoverssm_state( + self, num_accepted_tokens: torch.Tensor + ) -> RecoverSSMPostprocessMetadata | None: + commit = self.recoverssm_commit + if commit is None: + return None + context = self.recoverssm_context + assert context is not None + align = commit.align + context.commit( + num_accepted_tokens, + commit.state_indices[: self.num_spec_decodes, 0], + commit.query_start_loc[: self.num_spec_decodes + 1], + request_indices=commit.request_indices, + block_table=align.block_table if align is not None else None, + num_computed_tokens=( + align.num_computed_tokens if align is not None else None + ), + mamba_block_size=align.block_size if align is not None else None, + ) + if align is None: + return None + return RecoverSSMPostprocessMetadata( + num_spec_decodes=self.num_spec_decodes, + request_indices=commit.request_indices, + block_table=align.block_table, + num_computed_tokens=align.num_computed_tokens, + block_size=align.block_size, + ) class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder): + def __init__( + self, + kv_cache_spec: MambaSpec, + layer_names: list[str], + vllm_config: VllmConfig, + device: torch.device, + ) -> None: + super().__init__(kv_cache_spec, layer_names, vllm_config, device) + self.use_recoverssm = vllm_config.cache_config.use_kda_recoverssm + self.spec_state_slots = 1 if self.use_recoverssm else self.num_spec + 1 + self.recoverssm_num_accepted_tokens: torch.Tensor | None = None + self.recoverssm_context: KDARecoverSSMCommitContext | None = None + if self.use_recoverssm: + max_num_reqs = vllm_config.scheduler_config.max_num_seqs + self.spec_state_indices_tensor = torch.empty( + (max_num_reqs, 1), + dtype=torch.int32, + device=device, + ) + self.recoverssm_num_accepted_tokens = torch.ones( + max_num_reqs, + dtype=torch.int32, + device=device, + ) + + def _get_recoverssm_context(self) -> "KDARecoverSSMCommitContext": + context = self.recoverssm_context + if context is not None: + return context + + from vllm.models.kimi_k3.nvidia.ops.recoverssm import ( + KDARecoverSSMCommitContext, + ) + + forward_context = self.vllm_config.compilation_config.static_forward_context + layers = [forward_context[layer_name] for layer_name in self.layer_names] + context = KDARecoverSSMCommitContext.create( + layers, + spec_query_len=1 + self.vllm_config.num_speculative_tokens, + max_num_reqs=self.vllm_config.scheduler_config.max_num_seqs, + ) + self.recoverssm_context = context + return context + def build( # type: ignore[override] self, common_prefix_len: int, @@ -263,14 +369,26 @@ def build( # type: ignore[override] num_spec_decodes = 0 else: spec_sequence_masks_cpu = num_decode_draft_tokens_cpu >= 0 - # A nonnegative entry identifies a spec request. If no draft token - # was scheduled, process the whole batch as non-spec instead. - if num_decode_draft_tokens_cpu[spec_sequence_masks_cpu].sum().item() == 0: + if self.use_recoverssm: + assert m.is_prefilling is not None + assert m.is_prefilling.device.type == "cpu" + active_decode_mask_cpu = (~m.is_prefilling) & ( + query_start_loc_cpu.diff() > 0 + ) + spec_sequence_masks_cpu |= active_decode_mask_cpu + # Native KDA can use its regular decode path when no draft token + # was scheduled. RecoverSSM must preserve its extended conv window. + if ( + not self.use_recoverssm + and num_decode_draft_tokens_cpu[spec_sequence_masks_cpu].sum().item() + == 0 + ): spec_sequence_masks_cpu = None num_spec_decodes = 0 else: num_spec_decodes = spec_sequence_masks_cpu.sum().item() + spec_request_indices = None if num_spec_decodes == 0: # The runner orders ordinary decodes before prefills. num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( @@ -289,6 +407,16 @@ def build( # type: ignore[override] assert spec_sequence_masks_cpu is not None assert num_accepted_tokens is not None query_lens_cpu = query_start_loc_cpu.diff() + if ( + self.use_recoverssm + and torch.any( + query_lens_cpu[spec_sequence_masks_cpu] > self.num_spec + 1 + ).item() + ): + raise ValueError( + "KDA RecoverSSM speculative decode query length exceeds " + f"its activation capacity ({self.num_spec + 1})" + ) num_query_tokens = query_start_loc_cpu[-1].item() # Exclude zero-length cudagraph padding from request-indexed @@ -325,7 +453,7 @@ def build( # type: ignore[override] non_spec_token_indx = None # Real requests precede trailing cudagraph padding. spec_state_indices_tensor = block_table_tensor[ - :num_spec_decodes, : self.num_spec + 1 + :num_spec_decodes, : self.spec_state_slots ] non_spec_state_indices_tensor = None # Padding trails real requests, so this prefix already contains @@ -339,6 +467,12 @@ def build( # type: ignore[override] spec_sequence_masks_gpu = async_tensor_h2d( spec_sequence_masks_cpu, device=query_start_loc.device ) + if self.use_recoverssm: + spec_request_indices = async_tensor_h2d( + spec_sequence_masks_cpu.nonzero(as_tuple=True)[0], + dtype=torch.int32, + device=query_start_loc.device, + ) spec_token_masks = torch.repeat_interleave( spec_sequence_masks_gpu, query_lens, @@ -351,10 +485,10 @@ def build( # type: ignore[override] non_spec_token_indx = index[:num_non_spec_tokens] spec_token_indx = index[num_non_spec_tokens:] - # Spec requests carry one state slot per speculative step; - # non-spec requests use only their current state slot. + # Native spec uses one state slot per step. RecoverSSM keeps + # only the current checkpoint slot. spec_state_indices_tensor = block_table_tensor[ - spec_sequence_masks_cpu, : self.num_spec + 1 + spec_sequence_masks_cpu, : self.spec_state_slots ] non_spec_state_indices_tensor = block_table_tensor[ active_non_spec_mask_cpu, 0 @@ -400,12 +534,18 @@ def build( # type: ignore[override] num_accepted_tokens = num_accepted_tokens[spec_sequence_masks_cpu] + if self.use_recoverssm: + assert self.recoverssm_num_accepted_tokens is not None + num_accepted_tokens = self.recoverssm_num_accepted_tokens[ + :num_spec_decodes + ] + # Unlike the shared GDN layer, Kimi-K3's prefill KDA wrapper prepares # its own chunk indices. Only causal-convolution metadata is needed here. nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None if num_prefills > 0: has_initial_state = m.compute_num_computed_tokens() > 0 - if spec_sequence_masks_cpu is not None: + if num_spec_decodes > 0: has_initial_state = has_initial_state[active_non_spec_mask_cpu] assert non_spec_query_start_loc_cpu is not None nums_dict, batch_ptr, token_chunk_offset_ptr = ( @@ -426,7 +566,7 @@ def build( # type: ignore[override] and num_spec_decodes > 0 and num_prefills == 0 and num_decodes == 0 - and num_spec_decodes <= self.decode_cudagraph_max_bs + and batch_size <= self.spec_state_indices_tensor.shape[0] and num_spec_decode_tokens <= self.decode_cudagraph_max_bs ): # Equivalent PyTorch staging: @@ -463,6 +603,24 @@ def build( # type: ignore[override] :batch_size ] + recoverssm_commit = None + if self.use_recoverssm and num_spec_decodes > 0: + assert spec_state_indices_tensor is not None + assert spec_query_start_loc is not None + align = None + if self.kv_cache_spec.mamba_cache_mode == "align": + align = KDARecoverSSMAlignMetadata( + block_table=m.block_table_tensor, + num_computed_tokens=m.compute_num_computed_tokens(), + block_size=self.kv_cache_spec.block_size, + ) + recoverssm_commit = KDARecoverSSMCommitMetadata( + state_indices=spec_state_indices_tensor, + query_start_loc=spec_query_start_loc, + request_indices=spec_request_indices, + align=align, + ) + return KimiK3KDAMetadata( num_prefills=num_prefills, num_prefill_tokens=num_prefill_tokens, @@ -480,6 +638,12 @@ def build( # type: ignore[override] spec_token_indx=spec_token_indx, non_spec_token_indx=non_spec_token_indx, num_accepted_tokens=num_accepted_tokens, + recoverssm_commit=recoverssm_commit, + recoverssm_context=( + self._get_recoverssm_context() + if recoverssm_commit is not None + else None + ), nums_dict=nums_dict, batch_ptr=batch_ptr, token_chunk_offset_ptr=token_chunk_offset_ptr, diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index f7e8c54f1bf3..f1166695fdeb 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -62,12 +62,14 @@ EagleModelMixin, HasInnerState, IsHybrid, + MambaStateShapes, MixtureOfExperts, SupportsEagle3, SupportsEncoderCudaGraph, SupportsMultiModal, SupportsPP, SupportsQuant, + SupportsReplaySSM, ) from vllm.model_executor.models.kimi_k25 import KimiK25MediaPixelInputs from vllm.model_executor.models.kimi_k25_vit import ( @@ -1559,7 +1561,13 @@ def finalize_mega_moe_weights(self) -> None: class KimiLinearForCausalLM( - nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid, SupportsEagle3 + nn.Module, + HasInnerState, + SupportsPP, + MixtureOfExperts, + IsHybrid, + SupportsEagle3, + SupportsReplaySSM, ): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() @@ -1613,15 +1621,20 @@ def forward( # type: ignore[override] def get_mamba_state_dtype_from_config( cls, vllm_config: "VllmConfig", - ) -> tuple[torch.dtype, torch.dtype]: - return MambaStateDtypeCalculator.kda_state_dtype( + ) -> tuple[torch.dtype, ...]: + dtypes = MambaStateDtypeCalculator.kda_state_dtype( vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype ) + if vllm_config.cache_config.use_kda_recoverssm: + dtypes = MambaStateDtypeCalculator.append_kda_recoverssm_record( + dtypes, vllm_config.model_config.dtype + ) + return dtypes @classmethod def get_mamba_state_shape_from_config( cls, vllm_config: "VllmConfig" - ) -> tuple[tuple[int, int], tuple[int, int, int]]: + ) -> MambaStateShapes: parallel_config = vllm_config.parallel_config hf_config = vllm_config.model_config.hf_config tp_size = parallel_config.tensor_parallel_size @@ -1630,13 +1643,22 @@ def get_mamba_state_shape_from_config( if vllm_config.speculative_config else 0 ) - return MambaStateShapeCalculator.kda_state_shape( + shapes = MambaStateShapeCalculator.kda_state_shape( tp_size, hf_config.linear_attn_config["num_heads"], hf_config.linear_attn_config["head_dim"], conv_kernel_size=hf_config.linear_attn_config["short_conv_kernel_size"], num_spec=num_spec, ) + if vllm_config.cache_config.use_kda_recoverssm: + return MambaStateShapeCalculator.append_kda_recoverssm_record( + shapes, + hf_config.linear_attn_config["num_heads"], + hf_config.linear_attn_config["head_dim"], + tp_world_size=tp_size, + spec_query_len=1 + num_spec, + ) + return shapes @classmethod def get_mamba_state_copy_func( @@ -1697,6 +1719,7 @@ class KimiK3ForConditionalGeneration( SupportsEagle3, HasInnerState, IsHybrid, + SupportsReplaySSM, ): """Kimi-K3 model with Kimi-K2.5 vision and KimiLinear text.""" diff --git a/vllm/models/kimi_k3/nvidia/ops/recoverssm.py b/vllm/models/kimi_k3/nvidia/ops/recoverssm.py new file mode 100644 index 000000000000..c603ec7ed1fc --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/recoverssm.py @@ -0,0 +1,1067 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Kimi-K3 RecoverSSM speculative verify and accepted-state recovery.""" + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any + +import torch + +from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first +from vllm.triton_utils import tl, triton +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID + + +@triton.jit +def _kda_gate( + raw_g, + dt_bias, + A, + lower_bound, + USE_LOWER_BOUND: tl.constexpr, +): + gate_input = raw_g + dt_bias + if USE_LOWER_BOUND: + return lower_bound * tl.sigmoid(A * gate_input) + softplus_gate = tl.where( + gate_input > 20.0, + gate_input, + tl.log(1.0 + tl.exp(gate_input)), + ) + return -A * softplus_gate + + +@triton.jit +def _kda_recurrent_step( + state, + k, + v, + raw_g, + raw_beta, + dt_bias, + A, + lower_bound, + USE_LOWER_BOUND: tl.constexpr, +): + normalized_k = k * tl.rsqrt(tl.sum(k * k) + 1e-6) + gate = _kda_gate( + raw_g, + dt_bias, + A, + lower_bound, + USE_LOWER_BOUND, + ) + + state *= tl.exp(gate)[None, :] + correction = v - tl.sum(state * normalized_k[None, :], axis=1) + correction *= tl.sigmoid(raw_beta) + return state + correction[:, None] * normalized_k[None, :], correction + + +@triton.jit +def _kda_recoverssm_verify_kernel( + q_ptr, + k_ptr, + v_ptr, + raw_g_ptr, + raw_beta_ptr, + A_log_ptr, + dt_bias_ptr, + state_ptr, + correction_cache_ptr, + kg_cache_ptr, + out_ptr, + query_start_loc_ptr, + state_indices_ptr, + lower_bound, + null_block_id, + stride_q_token, + stride_k_token, + stride_v_token, + stride_g_token, + stride_beta_token, + stride_state_block, + stride_state_head, + stride_state_v, + stride_state_k, + stride_correction_block, + stride_correction_head, + stride_correction_pos, + stride_correction_dim, + stride_kg_block, + stride_kg_head, + stride_kg_pos, + stride_kg_dim, + stride_out_token, + stride_query_start_loc, + stride_state_indices, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + SPEC_QUERY_LEN: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, +): + pid_v = tl.program_id(0) + pid_b = tl.program_id(1) + pid_h = tl.program_id(2) + + bos = tl.load(query_start_loc_ptr + pid_b * stride_query_start_loc).to(tl.int64) + eos = tl.load(query_start_loc_ptr + (pid_b + 1) * stride_query_start_loc).to( + tl.int64 + ) + query_len = eos - bos + state_idx = tl.load(state_indices_ptr + pid_b * stride_state_indices).to(tl.int64) + + offs_k = tl.arange(0, BK) + offs_v = pid_v * BV + tl.arange(0, BV) + mask_k = offs_k < K + mask_v = offs_v < V + mask_state = mask_v[:, None] & mask_k[None, :] + + if state_idx <= null_block_id: + for token_offset in tl.static_range(SPEC_QUERY_LEN): + token_valid = token_offset < query_len + tl.store( + out_ptr + (bos + token_offset) * stride_out_token + pid_h * V + offs_v, + tl.zeros([BV], dtype=tl.float32), + mask=token_valid & mask_v, + ) + return + + state_ptrs = ( + state_ptr + + state_idx * stride_state_block + + pid_h * stride_state_head + + offs_v[:, None] * stride_state_v + + offs_k[None, :] * stride_state_k + ) + state = tl.load(state_ptrs, mask=mask_state, other=0.0).to(tl.float32) + A = tl.exp(tl.load(A_log_ptr + pid_h).to(tl.float32)) + + for token_offset in tl.static_range(SPEC_QUERY_LEN): + token_valid = token_offset < query_len + token = bos + token_offset + q = tl.load( + q_ptr + token * stride_q_token + pid_h * K + offs_k, + mask=token_valid & mask_k, + other=0.0, + ).to(tl.float32) + k = tl.load( + k_ptr + token * stride_k_token + pid_h * K + offs_k, + mask=token_valid & mask_k, + other=0.0, + ).to(tl.float32) + v = tl.load( + v_ptr + token * stride_v_token + pid_h * V + offs_v, + mask=token_valid & mask_v, + other=0.0, + ).to(tl.float32) + raw_g = tl.load( + raw_g_ptr + token * stride_g_token + pid_h * K + offs_k, + mask=token_valid & mask_k, + other=0.0, + ).to(tl.float32) + raw_beta = tl.load( + raw_beta_ptr + token * stride_beta_token + pid_h, + mask=token_valid, + other=0.0, + ).to(tl.float32) + + q *= tl.rsqrt(tl.sum(q * q) + 1e-6) * (K**-0.5) + dt_bias = tl.load(dt_bias_ptr + pid_h * K + offs_k, mask=mask_k, other=0.0).to( + tl.float32 + ) + updated_state, correction = _kda_recurrent_step( + state, + k, + v, + raw_g, + raw_beta, + dt_bias, + A, + lower_bound, + USE_LOWER_BOUND, + ) + state = tl.where(token_valid, updated_state, state) + + out = tl.sum(state * q[None, :], axis=1) + tl.store( + out_ptr + token * stride_out_token + pid_h * V + offs_v, + out, + mask=token_valid & mask_v, + ) + + correction_ptr = ( + correction_cache_ptr + + state_idx * stride_correction_block + + pid_h * stride_correction_head + + token_offset * stride_correction_pos + ) + tl.store( + correction_ptr + offs_v * stride_correction_dim, + correction, + mask=token_valid & mask_v, + ) + if pid_v == 0: + kg_ptr = ( + kg_cache_ptr + + state_idx * stride_kg_block + + pid_h * stride_kg_head + + token_offset * stride_kg_pos + ) + tl.store( + kg_ptr + offs_k * stride_kg_dim, + k, + mask=token_valid & mask_k, + ) + tl.store( + kg_ptr + (K + offs_k) * stride_kg_dim, + raw_g, + mask=token_valid & mask_k, + ) + + +@triton.heuristics( + { + "HAS_REQUEST_INDICES": lambda args: args["request_indices_ptr"] is not None, + "ALIGN_MODE": lambda args: args["block_table_ptr"] is not None, + } +) +@triton.jit +def _prepare_commit_plan_kernel( + num_accepted_ptr, + request_indices_ptr, + state_indices_ptr, + query_start_loc_ptr, + block_table_ptr, + num_computed_ptr, + commit_lens_ptr, + final_state_indices_ptr, + boundary_state_indices_ptr, + boundary_recovery_lens_ptr, + null_block_id, + mamba_block_size, + block_table_width, + stride_num_accepted, + stride_request_indices, + stride_state_indices, + stride_query_start_loc, + stride_block_table_row, + stride_block_table_col, + stride_num_computed, + SPEC_QUERY_LEN: tl.constexpr, + HAS_REQUEST_INDICES: tl.constexpr, + ALIGN_MODE: tl.constexpr, +): + spec_idx = tl.program_id(0) + source_state_idx = tl.load(state_indices_ptr + spec_idx * stride_state_indices).to( + tl.int64 + ) + request_idx = spec_idx + if HAS_REQUEST_INDICES: + request_idx = tl.load( + request_indices_ptr + spec_idx * stride_request_indices + ).to(tl.int64) + num_accepted = tl.load(num_accepted_ptr + request_idx * stride_num_accepted).to( + tl.int32 + ) + bos = tl.load(query_start_loc_ptr + spec_idx * stride_query_start_loc).to(tl.int64) + eos = tl.load(query_start_loc_ptr + (spec_idx + 1) * stride_query_start_loc).to( + tl.int64 + ) + query_len = (eos - bos).to(tl.int32) + commit_len = tl.minimum(tl.maximum(num_accepted, 0), query_len) + commit_len = tl.minimum(commit_len, SPEC_QUERY_LEN) + + final_state_idx = source_state_idx + boundary_state_idx = null_block_id + boundary_recovery_len = 0 + if ALIGN_MODE: + num_computed = tl.load(num_computed_ptr + request_idx * stride_num_computed).to( + tl.int32 + ) + final_num_computed = num_computed + commit_len + final_state_col = tl.minimum( + final_num_computed // mamba_block_size, block_table_width - 1 + ) + final_state_idx = tl.load( + block_table_ptr + + request_idx * stride_block_table_row + + final_state_col * stride_block_table_col + ).to(tl.int64) + next_boundary = (num_computed // mamba_block_size + 1) * mamba_block_size + crosses_boundary = final_num_computed >= next_boundary + boundary_recovery_len = next_boundary - num_computed + boundary_state_idx = tl.load( + block_table_ptr + + request_idx * stride_block_table_row + + (next_boundary // mamba_block_size - 1) * stride_block_table_col, + mask=crosses_boundary, + other=null_block_id, + ).to(tl.int64) + valid = (source_state_idx > null_block_id) & (commit_len > 0) + tl.store(commit_lens_ptr + spec_idx, tl.where(valid, commit_len, 0)) + tl.store( + final_state_indices_ptr + spec_idx, + tl.where(valid, final_state_idx, null_block_id), + ) + tl.store( + boundary_state_indices_ptr + spec_idx, + tl.where(valid, boundary_state_idx, null_block_id), + ) + tl.store( + boundary_recovery_lens_ptr + spec_idx, + tl.where(valid, boundary_recovery_len, 0), + ) + + +@triton.jit +def _compact_conv_state_kernel( + conv_state_ref_ptr, + conv_state_base_addrs_ptr, + conv_state_block_strides_ptr, + conv_state_dim_strides_ptr, + conv_state_token_strides_ptr, + state_indices_ptr, + commit_lens_ptr, + final_state_indices_ptr, + boundary_state_indices_ptr, + boundary_recovery_lens_ptr, + null_block_id, + conv_dim, + conv_history_len, + stride_state_indices, + BLOCK_D: tl.constexpr, + BLOCK_HISTORY: tl.constexpr, + ALIGN_MODE: tl.constexpr, +): + pid_d = tl.program_id(0) + pid_b = tl.program_id(1) + pid_l = tl.program_id(2) + source_state_idx = tl.load(state_indices_ptr + pid_b * stride_state_indices).to( + tl.int64 + ) + if source_state_idx <= null_block_id: + return + + commit_len = tl.load(commit_lens_ptr + pid_b) + if commit_len == 0: + return + final_state_idx = tl.load(final_state_indices_ptr + pid_b).to(tl.int64) + boundary_state_idx = tl.load(boundary_state_indices_ptr + pid_b).to(tl.int64) + boundary_recovery_len = tl.load(boundary_recovery_lens_ptr + pid_b) + + if final_state_idx <= null_block_id: + return + + base_addr = tl.load(conv_state_base_addrs_ptr + pid_l) + block_stride = tl.load(conv_state_block_strides_ptr + pid_l) + dim_stride = tl.load(conv_state_dim_strides_ptr + pid_l) + token_stride = tl.load(conv_state_token_strides_ptr + pid_l) + conv_state_ptr = base_addr.to(tl.pointer_type(conv_state_ref_ptr.dtype.element_ty)) + source_ptr = conv_state_ptr + source_state_idx * block_stride + final_ptr = conv_state_ptr + final_state_idx * block_stride + + offs_d = pid_d * BLOCK_D + tl.arange(0, BLOCK_D) + offs_h = tl.arange(0, BLOCK_HISTORY) + mask = (offs_d[:, None] < conv_dim) & (offs_h[None, :] < conv_history_len) + final_values = tl.load( + source_ptr + + offs_d[:, None] * dim_stride + + (commit_len - 1 + offs_h[None, :]) * token_stride, + mask=mask, + ) + if ALIGN_MODE: + boundary_values = tl.load( + source_ptr + + offs_d[:, None] * dim_stride + + (boundary_recovery_len - 1 + offs_h[None, :]) * token_stride, + mask=mask & (boundary_state_idx > null_block_id), + ) + boundary_ptr = conv_state_ptr + boundary_state_idx * block_stride + tl.store( + boundary_ptr + + offs_d[:, None] * dim_stride + + offs_h[None, :] * token_stride, + boundary_values, + mask=mask & (boundary_state_idx > null_block_id), + ) + tl.store( + final_ptr + offs_d[:, None] * dim_stride + offs_h[None, :] * token_stride, + final_values, + mask=mask, + ) + + +@triton.jit +def _commit_kda_state_kernel( + state_ref_ptr, + state_base_addrs_ptr, + state_block_strides_ptr, + correction_cache_ref_ptr, + correction_cache_base_addrs_ptr, + correction_cache_block_strides_ptr, + kg_cache_ref_ptr, + kg_cache_base_addrs_ptr, + kg_cache_block_strides_ptr, + A_log_ptr, + dt_bias_ptr, + state_indices_ptr, + commit_lens_ptr, + final_state_indices_ptr, + boundary_state_indices_ptr, + boundary_recovery_lens_ptr, + lower_bound, + null_block_id, + stride_state_head, + stride_state_v, + stride_state_k, + stride_correction_cache_head, + stride_correction_cache_pos, + stride_correction_cache_dim, + stride_kg_cache_head, + stride_kg_cache_pos, + stride_kg_cache_dim, + stride_A_layer, + stride_A_head, + stride_dt_bias_layer, + stride_dt_bias_head, + stride_dt_bias_dim, + stride_state_indices, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + NUM_HEADS: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, + ALIGN_MODE: tl.constexpr, +): + pid_v = tl.program_id(0) + pid_b = tl.program_id(1) + pid_lh = tl.program_id(2) + pid_l = pid_lh // NUM_HEADS + pid_h = pid_lh % NUM_HEADS + + source_state_idx = tl.load(state_indices_ptr + pid_b * stride_state_indices).to( + tl.int64 + ) + if source_state_idx <= null_block_id: + return + commit_len = tl.load(commit_lens_ptr + pid_b) + if commit_len == 0: + return + final_state_idx = tl.load(final_state_indices_ptr + pid_b).to(tl.int64) + boundary_state_idx = tl.load(boundary_state_indices_ptr + pid_b).to(tl.int64) + boundary_recovery_len = tl.load(boundary_recovery_lens_ptr + pid_b) + + if final_state_idx <= null_block_id: + return + + state_base_addr = tl.load(state_base_addrs_ptr + pid_l) + state_block_stride = tl.load(state_block_strides_ptr + pid_l) + state_ptr = state_base_addr.to(tl.pointer_type(state_ref_ptr.dtype.element_ty)) + source_state_ptr = ( + state_ptr + source_state_idx * state_block_stride + pid_h * stride_state_head + ) + + correction_cache_base_addr = tl.load(correction_cache_base_addrs_ptr + pid_l) + correction_cache_block_stride = tl.load(correction_cache_block_strides_ptr + pid_l) + correction_cache_ptr = correction_cache_base_addr.to( + tl.pointer_type(correction_cache_ref_ptr.dtype.element_ty) + ) + correction_cache_ptr += ( + source_state_idx * correction_cache_block_stride + + pid_h * stride_correction_cache_head + ) + kg_cache_base_addr = tl.load(kg_cache_base_addrs_ptr + pid_l) + kg_cache_block_stride = tl.load(kg_cache_block_strides_ptr + pid_l) + kg_cache_ptr = kg_cache_base_addr.to( + tl.pointer_type(kg_cache_ref_ptr.dtype.element_ty) + ) + kg_cache_ptr += ( + source_state_idx * kg_cache_block_stride + pid_h * stride_kg_cache_head + ) + + offs_k = tl.arange(0, BK) + offs_v = pid_v * BV + tl.arange(0, BV) + mask_k = offs_k < K + mask_v = offs_v < V + mask_state = mask_v[:, None] & mask_k[None, :] + state_ptrs = ( + source_state_ptr + + offs_v[:, None] * stride_state_v + + offs_k[None, :] * stride_state_k + ) + initial_state = tl.load(state_ptrs, mask=mask_state, other=0.0).to(tl.float32) + A = tl.exp( + tl.load(A_log_ptr + pid_l * stride_A_layer + pid_h * stride_A_head).to( + tl.float32 + ) + ) + + dt_bias = tl.load( + dt_bias_ptr + + pid_l * stride_dt_bias_layer + + pid_h * stride_dt_bias_head + + offs_k * stride_dt_bias_dim, + mask=mask_k, + other=0.0, + ).to(tl.float32) + final_decay = tl.full([BK], 1.0, tl.float32) + final_correction = tl.zeros([BV, BK], tl.float32) + boundary_decay = tl.full([BK], 1.0, tl.float32) + boundary_correction = tl.zeros([BV, BK], tl.float32) + + for reverse_offset in range(commit_len): + token_offset = commit_len - reverse_offset - 1 + correction_ptr = ( + correction_cache_ptr + token_offset * stride_correction_cache_pos + ) + kg_ptr = kg_cache_ptr + token_offset * stride_kg_cache_pos + k = tl.load( + kg_ptr + offs_k * stride_kg_cache_dim, + mask=mask_k, + other=0.0, + ).to(tl.float32) + correction = tl.load( + correction_ptr + offs_v * stride_correction_cache_dim, + mask=mask_v, + other=0.0, + ).to(tl.float32) + raw_g = tl.load( + kg_ptr + (K + offs_k) * stride_kg_cache_dim, + mask=mask_k, + other=0.0, + ).to(tl.float32) + normalized_k = k * tl.rsqrt(tl.sum(k * k) + 1e-6) + gate = _kda_gate( + raw_g, + dt_bias, + A, + lower_bound, + USE_LOWER_BOUND, + ) + update = correction[:, None] * normalized_k[None, :] + decay = tl.exp(gate) + final_correction += update * final_decay[None, :] + final_decay *= decay + if ALIGN_MODE: + before_boundary = token_offset < boundary_recovery_len + boundary_correction += tl.where( + before_boundary, + update * boundary_decay[None, :], + 0.0, + ) + boundary_decay *= tl.where(before_boundary, decay, 1.0) + + state = initial_state * final_decay[None, :] + final_correction + if ALIGN_MODE: + boundary_ptrs = ( + state_ptr + + boundary_state_idx * state_block_stride + + pid_h * stride_state_head + + offs_v[:, None] * stride_state_v + + offs_k[None, :] * stride_state_k + ) + boundary_state = initial_state * boundary_decay[None, :] + boundary_correction + tl.store( + boundary_ptrs, + boundary_state, + mask=mask_state & (boundary_state_idx > null_block_id), + ) + + final_ptrs = ( + state_ptr + + final_state_idx * state_block_stride + + pid_h * stride_state_head + + offs_v[:, None] * stride_state_v + + offs_k[None, :] * stride_state_k + ) + tl.store(final_ptrs, state, mask=mask_state) + + +def kda_recoverssm_verify( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + lower_bound: float | None, + checkpoint_state: torch.Tensor, + correction_cache: torch.Tensor, + kg_cache: torch.Tensor, + query_start_loc: torch.Tensor, + state_indices: torch.Tensor, + spec_query_len: int, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Verify a KDA speculative window without modifying its checkpoint.""" + if q.ndim != 4 or q.shape[0] != 1: + raise ValueError("KDA RecoverSSM q must have shape [1, tokens, heads, dim]") + _, total_tokens, num_heads, key_dim = q.shape + value_dim = v.shape[-1] + if k.shape != q.shape or v.shape != (1, total_tokens, num_heads, value_dim): + raise ValueError("KDA RecoverSSM q, k, and v shapes are incompatible") + if raw_g.shape != q.shape or raw_beta.shape != (1, total_tokens, num_heads): + raise ValueError("KDA RecoverSSM gate or beta shape is incompatible") + if any(tensor.stride()[2:] != (key_dim, 1) for tensor in (q, k, raw_g)): + raise ValueError("KDA RecoverSSM q, k, and gate heads must be contiguous") + if v.stride()[2:] != (value_dim, 1) or raw_beta.stride(2) != 1: + raise ValueError("KDA RecoverSSM v and beta heads must be contiguous") + num_blocks = checkpoint_state.shape[0] + if checkpoint_state.shape[1:] != ( + num_heads, + value_dim, + key_dim, + ): + raise ValueError("KDA RecoverSSM checkpoint shape is incompatible") + expected_correction_shape = ( + num_blocks, + num_heads, + spec_query_len, + value_dim, + ) + if correction_cache.shape != expected_correction_shape: + raise ValueError( + f"KDA RecoverSSM correction buffer needs shape {expected_correction_shape}" + ) + expected_kg_shape = (num_blocks, num_heads, spec_query_len, 2 * key_dim) + if kg_cache.shape != expected_kg_shape: + raise ValueError( + f"KDA RecoverSSM key/gate buffer needs shape {expected_kg_shape}" + ) + if correction_cache.dtype != torch.float32: + raise ValueError("KDA RecoverSSM correction buffer must use float32") + if kg_cache.dtype != k.dtype: + raise ValueError("KDA RecoverSSM key/gate buffer must match activation dtype") + if A_log.shape != (num_heads,) or dt_bias.numel() != num_heads * key_dim: + raise ValueError("KDA RecoverSSM gate parameters are incompatible") + if not A_log.is_contiguous() or not dt_bias.is_contiguous(): + raise ValueError("KDA RecoverSSM gate parameters must be contiguous") + batch = state_indices.shape[0] + if query_start_loc.shape[0] != batch + 1: + raise ValueError("KDA RecoverSSM query metadata is incompatible") + if total_tokens > batch * spec_query_len: + raise ValueError( + "KDA RecoverSSM speculative decode input exceeds its activation capacity" + ) + if out is None: + out = torch.empty_like(v) + if out.shape != v.shape: + raise ValueError("KDA RecoverSSM output shape is incompatible") + if out.stride()[2:] != (value_dim, 1): + raise ValueError("KDA RecoverSSM output heads must be contiguous") + device = q.device + if any( + tensor.device != device + for tensor in ( + k, + v, + raw_g, + raw_beta, + A_log, + dt_bias, + checkpoint_state, + correction_cache, + kg_cache, + query_start_loc, + state_indices, + out, + ) + ): + raise ValueError("KDA RecoverSSM inputs must be on the same device") + if total_tokens == 0: + return out + + block_k = triton.next_power_of_2(key_dim) + block_v = min(triton.next_power_of_2(value_dim), 32) + grid = (triton.cdiv(value_dim, block_v), batch, num_heads) + _kda_recoverssm_verify_kernel[grid]( + q, + k, + v, + raw_g, + raw_beta, + A_log, + dt_bias, + checkpoint_state, + correction_cache, + kg_cache, + out, + query_start_loc, + state_indices, + lower_bound or 0.0, + NULL_BLOCK_ID, + q.stride(1), + k.stride(1), + v.stride(1), + raw_g.stride(1), + raw_beta.stride(1), + checkpoint_state.stride(0), + checkpoint_state.stride(1), + checkpoint_state.stride(2), + checkpoint_state.stride(3), + correction_cache.stride(0), + correction_cache.stride(1), + correction_cache.stride(2), + correction_cache.stride(3), + kg_cache.stride(0), + kg_cache.stride(1), + kg_cache.stride(2), + kg_cache.stride(3), + out.stride(1), + query_start_loc.stride(0), + state_indices.stride(0), + K=key_dim, + V=value_dim, + BK=block_k, + BV=block_v, + SPEC_QUERY_LEN=spec_query_len, + USE_LOWER_BOUND=lower_bound is not None, + num_warps=4, + num_stages=2, + ) + return out + + +@dataclass +class KDARecoverSSMCommitContext: + conv_states: tuple[torch.Tensor, ...] + conv_state_base_addrs: torch.Tensor + conv_state_block_strides: torch.Tensor + conv_state_dim_strides: torch.Tensor + conv_state_token_strides: torch.Tensor + conv_history_len: int + checkpoints: tuple[torch.Tensor, ...] + state_base_addrs: torch.Tensor + state_block_strides: torch.Tensor + correction_caches: tuple[torch.Tensor, ...] + correction_cache_base_addrs: torch.Tensor + correction_cache_block_strides: torch.Tensor + kg_caches: tuple[torch.Tensor, ...] + kg_cache_base_addrs: torch.Tensor + kg_cache_block_strides: torch.Tensor + commit_lens: torch.Tensor + final_state_indices: torch.Tensor + boundary_state_indices: torch.Tensor + boundary_recovery_lens: torch.Tensor + A_log: torch.Tensor + dt_bias: torch.Tensor + lower_bound: float | None + spec_query_len: int + + @classmethod + def create( + cls, + layers: Sequence[Any], + *, + spec_query_len: int, + max_num_reqs: int, + ) -> "KDARecoverSSMCommitContext": + if not layers: + raise ValueError("KDA RecoverSSM commit requires at least one layer") + if any(len(layer.kv_cache) != 4 for layer in layers): + raise ValueError( + "KDA RecoverSSM pages must contain conv, state, correction, " + "and key/gate" + ) + + conv_states = [layer.kv_cache[0] for layer in layers] + if not is_conv_state_dim_first(): + conv_states = [state.transpose(-1, -2) for state in conv_states] + checkpoints = [layer.kv_cache[1] for layer in layers] + correction_caches = [layer.kv_cache[2] for layer in layers] + kg_caches = [layer.kv_cache[3] for layer in layers] + A_log = [layer.A_log for layer in layers] + dt_bias = [ + layer.dt_bias.view(layer.local_num_heads, layer.head_dim) + for layer in layers + ] + lower_bounds = {layer.gate_lower_bound for layer in layers} + if len(lower_bounds) != 1: + raise ValueError("KDA RecoverSSM layers need matching gate bounds") + + state_ref = checkpoints[0] + if state_ref.ndim != 4: + raise ValueError("KDA RecoverSSM checkpoint must be four-dimensional") + num_blocks, num_heads, value_dim, key_dim = state_ref.shape + for state in checkpoints: + if ( + state.shape != state_ref.shape + or state.dtype != state_ref.dtype + or state.device != state_ref.device + or state.stride()[1:] != state_ref.stride()[1:] + ): + raise ValueError( + "KDA RecoverSSM layers need matching checkpoint layout" + ) + expected_correction_shape = ( + num_blocks, + num_heads, + spec_query_len, + value_dim, + ) + correction_ref = correction_caches[0] + for correction_cache in correction_caches: + if ( + correction_cache.shape != expected_correction_shape + or correction_cache.dtype != torch.float32 + or correction_cache.device != state_ref.device + or correction_cache.stride()[1:] != correction_ref.stride()[1:] + ): + raise ValueError( + "KDA RecoverSSM correction buffers need float32 shape " + f"{expected_correction_shape}" + ) + expected_kg_shape = (num_blocks, num_heads, spec_query_len, 2 * key_dim) + kg_ref = kg_caches[0] + for kg_cache in kg_caches: + if ( + kg_cache.shape != expected_kg_shape + or kg_cache.dtype != kg_ref.dtype + or kg_cache.device != state_ref.device + or kg_cache.stride()[1:] != kg_ref.stride()[1:] + ): + raise ValueError( + f"KDA RecoverSSM key/gate buffers need shape {expected_kg_shape}" + ) + if any(param.shape != (num_heads,) for param in A_log): + raise ValueError("KDA RecoverSSM A_log shape is incompatible") + if any(param.shape != (num_heads, key_dim) for param in dt_bias): + raise ValueError("KDA RecoverSSM dt_bias shape is incompatible") + + conv_ref = conv_states[0] + if conv_ref.ndim != 3: + raise ValueError("KDA RecoverSSM conv state must be three-dimensional") + conv_dim, conv_state_len = conv_ref.shape[1:] + conv_history_len = conv_state_len - spec_query_len + 1 + if conv_history_len <= 0: + raise ValueError("KDA RecoverSSM conv state is shorter than its window") + for conv_state in conv_states: + if ( + conv_state.shape != conv_ref.shape + or conv_state.dtype != conv_ref.dtype + or conv_state.device != state_ref.device + or conv_state.shape[0] != num_blocks + ): + raise ValueError("KDA RecoverSSM layers need matching conv state") + + device = state_ref.device + + def _base_addrs(tensors: Sequence[torch.Tensor]) -> torch.Tensor: + return torch.tensor( + [tensor.data_ptr() for tensor in tensors], + dtype=torch.int64, + device=device, + ) + + def _block_strides(tensors: Sequence[torch.Tensor]) -> torch.Tensor: + return torch.tensor( + [tensor.stride(0) for tensor in tensors], + dtype=torch.int64, + device=device, + ) + + return cls( + conv_states=tuple(conv_states), + conv_state_base_addrs=_base_addrs(conv_states), + conv_state_block_strides=_block_strides(conv_states), + conv_state_dim_strides=torch.tensor( + [state.stride(1) for state in conv_states], + dtype=torch.int64, + device=device, + ), + conv_state_token_strides=torch.tensor( + [state.stride(2) for state in conv_states], + dtype=torch.int64, + device=device, + ), + conv_history_len=conv_history_len, + checkpoints=tuple(checkpoints), + state_base_addrs=_base_addrs(checkpoints), + state_block_strides=_block_strides(checkpoints), + correction_caches=tuple(correction_caches), + correction_cache_base_addrs=_base_addrs(correction_caches), + correction_cache_block_strides=_block_strides(correction_caches), + kg_caches=tuple(kg_caches), + kg_cache_base_addrs=_base_addrs(kg_caches), + kg_cache_block_strides=_block_strides(kg_caches), + commit_lens=torch.empty(max_num_reqs, dtype=torch.int32, device=device), + final_state_indices=torch.empty( + max_num_reqs, dtype=torch.int32, device=device + ), + boundary_state_indices=torch.empty( + max_num_reqs, dtype=torch.int32, device=device + ), + boundary_recovery_lens=torch.empty( + max_num_reqs, dtype=torch.int32, device=device + ), + A_log=torch.stack(tuple(A_log)).contiguous(), + dt_bias=torch.stack(tuple(dt_bias)).contiguous(), + lower_bound=lower_bounds.pop(), + spec_query_len=spec_query_len, + ) + + def commit( + self, + num_accepted_tokens: torch.Tensor, + state_indices: torch.Tensor, + query_start_loc: torch.Tensor, + request_indices: torch.Tensor | None = None, + block_table: torch.Tensor | None = None, + num_computed_tokens: torch.Tensor | None = None, + mamba_block_size: int | None = None, + ) -> None: + """Fold accepted KDA and convolution inputs into every layer.""" + batch = state_indices.shape[0] + if batch == 0: + return + if batch > self.commit_lens.shape[0]: + raise ValueError("KDA RecoverSSM commit batch exceeds its plan capacity") + if query_start_loc.shape[0] != batch + 1: + raise ValueError("KDA RecoverSSM commit metadata is incompatible") + if request_indices is not None and request_indices.shape[0] < batch: + raise ValueError("KDA RecoverSSM request mapping is too short") + align_args = (block_table, num_computed_tokens, mamba_block_size) + if any(arg is not None for arg in align_args) and any( + arg is None for arg in align_args + ): + raise ValueError("KDA RecoverSSM align metadata is incomplete") + if mamba_block_size is not None and mamba_block_size < self.spec_query_len: + raise ValueError( + "KDA RecoverSSM align block size must cover one speculative window" + ) + if block_table is not None and block_table.ndim != 2: + raise ValueError("KDA RecoverSSM block table must be two-dimensional") + device = self.checkpoints[0].device + if ( + any( + tensor.device != device + for tensor in ( + num_accepted_tokens, + state_indices, + query_start_loc, + ) + ) + or (request_indices is not None and request_indices.device != device) + or (block_table is not None and block_table.device != device) + or ( + num_computed_tokens is not None and num_computed_tokens.device != device + ) + ): + raise ValueError("KDA RecoverSSM commit inputs must be on the same device") + + block_table_stride = (0, 0) if block_table is None else block_table.stride() + num_computed_stride = ( + 0 if num_computed_tokens is None else num_computed_tokens.stride(0) + ) + + num_layers = len(self.checkpoints) + conv_ref = self.conv_states[0] + conv_dim = conv_ref.shape[1] + block_history = triton.next_power_of_2(self.conv_history_len) + _prepare_commit_plan_kernel[(batch,)]( + num_accepted_tokens, + request_indices, + state_indices, + query_start_loc, + block_table, + num_computed_tokens, + self.commit_lens, + self.final_state_indices, + self.boundary_state_indices, + self.boundary_recovery_lens, + NULL_BLOCK_ID, + mamba_block_size or 1, + block_table.shape[1] if block_table is not None else 1, + num_accepted_tokens.stride(0), + request_indices.stride(0) if request_indices is not None else 0, + state_indices.stride(0), + query_start_loc.stride(0), + block_table_stride[0], + block_table_stride[1], + num_computed_stride, + SPEC_QUERY_LEN=self.spec_query_len, + num_warps=1, + ) + _compact_conv_state_kernel[(triton.cdiv(conv_dim, 256), batch, num_layers)]( + conv_ref, + self.conv_state_base_addrs, + self.conv_state_block_strides, + self.conv_state_dim_strides, + self.conv_state_token_strides, + state_indices, + self.commit_lens, + self.final_state_indices, + self.boundary_state_indices, + self.boundary_recovery_lens, + NULL_BLOCK_ID, + conv_dim, + self.conv_history_len, + state_indices.stride(0), + BLOCK_D=256, + BLOCK_HISTORY=block_history, + ALIGN_MODE=block_table is not None, + num_warps=4, + ) + + state_ref = self.checkpoints[0] + _, num_heads, value_dim, key_dim = state_ref.shape + block_k = triton.next_power_of_2(key_dim) + block_v = min(triton.next_power_of_2(value_dim), 32) + grid = ( + triton.cdiv(value_dim, block_v), + batch, + num_layers * num_heads, + ) + _commit_kda_state_kernel[grid]( + state_ref, + self.state_base_addrs, + self.state_block_strides, + self.correction_caches[0], + self.correction_cache_base_addrs, + self.correction_cache_block_strides, + self.kg_caches[0], + self.kg_cache_base_addrs, + self.kg_cache_block_strides, + self.A_log, + self.dt_bias, + state_indices, + self.commit_lens, + self.final_state_indices, + self.boundary_state_indices, + self.boundary_recovery_lens, + self.lower_bound or 0.0, + NULL_BLOCK_ID, + state_ref.stride(1), + state_ref.stride(2), + state_ref.stride(3), + self.correction_caches[0].stride(1), + self.correction_caches[0].stride(2), + self.correction_caches[0].stride(3), + self.kg_caches[0].stride(1), + self.kg_caches[0].stride(2), + self.kg_caches[0].stride(3), + self.A_log.stride(0), + self.A_log.stride(1), + self.dt_bias.stride(0), + self.dt_bias.stride(1), + self.dt_bias.stride(2), + state_indices.stride(0), + K=key_dim, + V=value_dim, + BK=block_k, + BV=block_v, + NUM_HEADS=num_heads, + USE_LOWER_BOUND=self.lower_bound is not None, + ALIGN_MODE=block_table is not None, + num_warps=4, + num_stages=2, + ) + + +__all__ = ["KDARecoverSSMCommitContext", "kda_recoverssm_verify"] diff --git a/vllm/v1/attention/backends/recoverssm_metadata.py b/vllm/v1/attention/backends/recoverssm_metadata.py new file mode 100644 index 000000000000..7e4cca68f8b3 --- /dev/null +++ b/vllm/v1/attention/backends/recoverssm_metadata.py @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import abc +from dataclasses import dataclass + +import torch + + +@dataclass(frozen=True) +class RecoverSSMPostprocessMetadata: + """Metadata used during postprocessing for align-mode prefix caching.""" + + num_spec_decodes: int + request_indices: torch.Tensor | None + block_table: torch.Tensor + num_computed_tokens: torch.Tensor + block_size: int + + +class RecoverSSMMetadata(abc.ABC): + @abc.abstractmethod + def commit_recoverssm_state( + self, num_accepted_tokens: torch.Tensor + ) -> RecoverSSMPostprocessMetadata | None: + raise NotImplementedError diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index 2715f790dcce..c61a56f9aabf 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -20,6 +20,7 @@ from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache from vllm.v1.worker.gpu.model_states.default import DefaultModelState from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata +from vllm.v1.worker.gpu.model_states.recoverssm import RecoverSSMState from vllm.v1.worker.mamba_utils import ( MambaSpecDecodeGPUContext, preprocess_mamba_align_fused_kernel, @@ -80,6 +81,9 @@ def __init__( # kernel reusing the postprocess copy machinery, so the per-step src # columns and the running state_idx are kept GPU-resident. self._align_mode = self.cache_config.mamba_cache_mode == "align" + self.recoverssm = ( + RecoverSSMState() if self.cache_config.use_kda_recoverssm else None + ) if self._align_mode: self._mamba_state_idx_gpu = torch.zeros( self.max_num_reqs, dtype=torch.int32, device=self.device @@ -267,7 +271,7 @@ def prepare_attn( num_accepted_tokens=num_accepted_tokens, num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, ) - return build_attn_metadata( + attn_metadata = build_attn_metadata( attn_groups=attn_groups, num_reqs=num_reqs, num_tokens=num_tokens, @@ -285,6 +289,13 @@ def prepare_attn( for_cudagraph_capture=for_capture, rswa_prefix_lens=input_batch.prompt_lens, ) + if self.recoverssm is not None: + self.recoverssm.record_step( + attn_metadata, + attn_groups, + for_capture=for_capture, + ) + return attn_metadata def postprocess_state( self, @@ -295,21 +306,34 @@ def postprocess_state( # Chunked prefill does not sample a token, so num_sampled can be 0. # Mamba treats num_accepted_tokens=1 as the neutral non-spec value. num_reqs = idx_mapping.shape[0] - if not num_reqs: - return + if num_reqs: + if not isinstance(num_sampled, int): + # idx_mapping may contain -1 sentinels (filtered rows) under PP; the + # kernel skips them rather than scattering with a host-side gather. + _scatter_num_accepted_kernel[(num_reqs,)]( + idx_mapping, + num_sampled, + self.num_accepted_tokens_gpu, + ) + else: + # Fill with single value. + _fill_num_accepted_kernel[(num_reqs,)]( + idx_mapping, + self.num_accepted_tokens_gpu, + max(num_sampled, 1), + ) - if not isinstance(num_sampled, int): - # idx_mapping may contain -1 sentinels (filtered rows) under PP; the - # kernel skips them rather than scattering with a host-side gather. - _scatter_num_accepted_kernel[(num_reqs,)]( - idx_mapping, num_sampled, self.num_accepted_tokens_gpu - ) - else: - # Fill with single value. - _fill_num_accepted_kernel[(num_reqs,)]( - idx_mapping, self.num_accepted_tokens_gpu, max(num_sampled, 1) + if self.recoverssm is not None: + self.recoverssm.commit_step( + num_sampled, + idx_mapping, + state_indices=(self._mamba_state_idx_gpu if self._align_mode else None), + num_accepted_tokens=self.num_accepted_tokens_gpu, ) + if not num_reqs: + return + # Align: save the running state to the block-aligned position when # spec-decode acceptance leaves the sequence non-block-aligned (mirrors # the V1 align postprocess). num_computed_tokens already holds the diff --git a/vllm/v1/worker/gpu/model_states/recoverssm.py b/vllm/v1/worker/gpu/model_states/recoverssm.py new file mode 100644 index 000000000000..cfbe1a4d45d1 --- /dev/null +++ b/vllm/v1/worker/gpu/model_states/recoverssm.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Any + +import torch + +from vllm.triton_utils import tl, triton +from vllm.v1.attention.backends.recoverssm_metadata import RecoverSSMMetadata +from vllm.v1.worker.utils import AttentionGroup + + +class RecoverSSMState: + """Coordinates RecoverSSM metadata between attention and postprocessing.""" + + def __init__(self) -> None: + self._step: tuple[RecoverSSMMetadata, ...] | None = None + + def record_step( + self, + attn_metadata: dict[str, Any], + attn_groups: list[list[AttentionGroup]], + *, + for_capture: bool, + ) -> None: + if for_capture: + self._step = None + return + + step: list[RecoverSSMMetadata] = [] + for group_list in attn_groups: + for group in group_list: + metadata = attn_metadata[group.layer_names[0]] + if isinstance(metadata, RecoverSSMMetadata): + step.append(metadata) + self._step = tuple(step) + + def commit_step( + self, + num_sampled: torch.Tensor | int, + idx_mapping: torch.Tensor, + *, + state_indices: torch.Tensor | None, + num_accepted_tokens: torch.Tensor, + ) -> None: + step = self._step + self._step = None + if isinstance(num_sampled, int) or step is None: + return + + for metadata in step: + postprocess_meta = metadata.commit_recoverssm_state(num_sampled) + if postprocess_meta is None: + continue + assert state_indices is not None + # RecoverSSM already restored the accepted state. Update its running + # column and reset the next-step copy bias to the neutral value. + _postprocess_recoverssm_align_kernel[(postprocess_meta.num_spec_decodes,)]( + idx_mapping, + num_sampled, + postprocess_meta.request_indices, + postprocess_meta.num_computed_tokens, + state_indices, + num_accepted_tokens, + MAMBA_BLOCK_SIZE=postprocess_meta.block_size, + BLOCK_TABLE_WIDTH=postprocess_meta.block_table.shape[1], + ) + + +@triton.heuristics( + {"HAS_REQUEST_INDICES": lambda args: args["request_indices_ptr"] is not None} +) +@triton.jit +def _postprocess_recoverssm_align_kernel( + idx_mapping_ptr, + num_sampled_ptr, + request_indices_ptr, + num_computed_ptr, + state_idx_ptr, + num_accepted_ptr, + HAS_REQUEST_INDICES: tl.constexpr, + MAMBA_BLOCK_SIZE: tl.constexpr, + BLOCK_TABLE_WIDTH: tl.constexpr, +): + spec_idx = tl.program_id(0) + batch_idx = spec_idx + if HAS_REQUEST_INDICES: + batch_idx = tl.load(request_indices_ptr + spec_idx) + req_state_idx = tl.load(idx_mapping_ptr + batch_idx) + if req_state_idx < 0: + return + num_sampled = tl.load(num_sampled_ptr + batch_idx) + num_computed = tl.load(num_computed_ptr + batch_idx) + tl.store( + state_idx_ptr + req_state_idx, + tl.minimum( + (num_computed + num_sampled) // MAMBA_BLOCK_SIZE, + BLOCK_TABLE_WIDTH - 1, + ), + ) + tl.store(num_accepted_ptr + req_state_idx, 1) diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index 82247e3ee401..a0b89303dd86 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -828,7 +828,13 @@ def _populate_metadata( attention = forward_context[layer_name] kv_caches: list[torch.Tensor] = attention.kv_cache - for state_type_idx, state in enumerate(kv_caches): + if len(kv_caches) < self.num_state_types: + raise ValueError( + f"Expected at least {self.num_state_types} Mamba state " + f"tensors, got {len(kv_caches)}" + ) + for state_type_idx, copy_func in enumerate(mamba_state_copy_funcs): + state = kv_caches[state_type_idx] # Base address self.state_base_addrs[idx] = state.data_ptr() @@ -845,7 +851,6 @@ def _populate_metadata( # Element size self.state_elem_sizes[idx] = state.element_size() - copy_func = mamba_state_copy_funcs[state_type_idx] assert ( copy_func is get_conv_copy_spec or copy_func is get_temporal_copy_spec From 1d3a8b9e220f1edc77c190c3370cf0f76dfdd2fd Mon Sep 17 00:00:00 2001 From: qli88 Date: Mon, 17 Aug 2026 07:26:05 -0500 Subject: [PATCH 046/839] [ROCm][Bugfix] Fix Triton W4A16 bug in determining if transpose is required for GPTQ/AutoGPTQ (#48998) Signed-off-by: Qiang Li --- .../kernels/quantization/test_triton_w4a16.py | 2 + .../linear/mixed_precision/triton_w4a16.py | 62 +++++++++++-------- 2 files changed, 39 insertions(+), 25 deletions(-) diff --git a/tests/kernels/quantization/test_triton_w4a16.py b/tests/kernels/quantization/test_triton_w4a16.py index 42f163dea44a..3c4682911656 100644 --- a/tests/kernels/quantization/test_triton_w4a16.py +++ b/tests/kernels/quantization/test_triton_w4a16.py @@ -306,6 +306,7 @@ class DummyLayer(torch.nn.Module): @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only") def test_triton_w4a16_process_weights_after_loading_keeps_gptq_qzeros_layout(): + """AutoGPTQ qzeros are already [K//G, N//8] (output_dim=1): no transpose.""" if not torch.cuda.is_available(): pytest.skip("CUDA/HIP device not available") @@ -413,6 +414,7 @@ class DummyLayer(torch.nn.Module): @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only") def test_triton_w4a16_symmetric_apply_ignores_qzeros(monkeypatch): + """For symmetric (uint4b8) layers, apply_weights must pass qzeros=None.""" from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( MPLinearLayerConfig, ) diff --git a/vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py b/vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py index bc0a587b6763..8a4a493c6c55 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py @@ -400,29 +400,40 @@ def repack_w_s(x: BasevLLMParameter) -> BasevLLMParameter: if self.w_zp_name is not None: zp = getattr(layer, self.w_zp_name, None) if zp is not None: - c = self.config - K, N = c.partition_weight_shape - group_size = c.group_size if c.group_size != -1 else K - expected_shape = (K // group_size, N // 8) - transposed_shape = (N // 8, K // group_size) - - if tuple(zp.data.shape) == expected_shape: - # GPTQ/AutoGPTQ already stores qzeros in the kernel layout. - qzeros = zp.data.contiguous() - elif tuple(zp.data.shape) == transposed_shape: - # Compressed-tensors stores qzeros transposed from what the - # kernel needs. - qzeros = zp.data.t().contiguous() + # Kernel needs [K//G, N//8]: + # input(K) at dim 0, output(N) packed at dim 1. + # AutoGPTQ: + # output_dim=1 -> already [K//G, N//8], no transpose. + # compressed-tensors: + # output_dim=0 -> [N//8, K//G], needs transpose. + # None (unknown): + # infer from shape; if square (ambiguous), default to transpose. + zp_output_dim = getattr(zp, "output_dim", None) + if zp_output_dim is not None: + needs_transpose = zp_output_dim != 1 else: - raise AssertionError( - f"{self.w_zp_name} shape mismatch: {zp.data.shape}; " - f"expected {expected_shape} or {transposed_shape}" - ) - + # in case output_dim is None + c = self.config + K, N = c.partition_weight_shape + group_size = c.group_size if c.group_size != -1 else K + expected_shape = (K // group_size, N // 8) + transposed_shape = (N // 8, K // group_size) + if ( + tuple(zp.data.shape) == expected_shape + and expected_shape != transposed_shape + ): + needs_transpose = False + else: + needs_transpose = True + zp_data = ( + zp.data.t().contiguous() + if needs_transpose + else zp.data.contiguous() + ) replace_parameter( layer, self.w_zp_name, - torch.nn.Parameter(qzeros, requires_grad=False), + torch.nn.Parameter(zp_data, requires_grad=False), ) def apply_weights( @@ -437,17 +448,18 @@ def apply_weights( K = c.partition_weight_shape[0] group_size = c.group_size if c.group_size != -1 else K - # For symmetric types (uint4b8), use the scalar bias; no zeros tensor. - # Some checkpoint loaders still register qzeros parameters for GPTQ - # layers, but they are not part of the symmetric kernel contract. - zp_bias = c.weight_type.bias if c.weight_type.has_bias() else 0 - qzeros = None if c.weight_type.has_bias() else w_zp + # For symmetric types (uint4b8), use the scalar bias; no zeros tensor + if c.weight_type.has_bias(): + zp_bias = c.weight_type.bias + w_zp = None # symmetric: ignore qzeros, use scalar bias instead + else: + zp_bias = 0 output = triton_w4a16_gemm( a=x_2d, b_q=w_q, scales=w_s, - qzeros=qzeros, + qzeros=w_zp, group_size=group_size, zp_bias=zp_bias, ) From 017e9f4448b700e85ee16023287b025693c72b9e Mon Sep 17 00:00:00 2001 From: Tyler Michael Smith Date: Mon, 17 Aug 2026 09:20:50 -0400 Subject: [PATCH 047/839] Promote `prefix_cache_retention_interval` to an argument and change the default to 0 (#52216) Signed-off-by: Tyler Michael Smith Co-authored-by: OpenAI Codex --- tests/config/test_config_utils.py | 5 + tests/engine/test_arg_utils.py | 23 ++++ tests/v1/core/test_contiguous_kv_packing.py | 2 + tests/v1/core/test_kv_cache_utils.py | 4 + tests/v1/core/test_prefix_caching.py | 110 +++++++++++------- tests/v1/engine/test_engine_args.py | 5 + tests/v1/engine/test_engine_core_client.py | 1 + .../unit/test_mooncake_store_worker.py | 6 +- vllm/config/cache.py | 23 +++- .../kv_connector/v1/mooncake/store/worker.py | 2 +- vllm/engine/arg_utils.py | 8 ++ vllm/envs.py | 6 - vllm/v1/core/kv_cache_coordinator.py | 11 +- vllm/v1/core/kv_cache_manager.py | 2 +- vllm/v1/core/kv_cache_utils.py | 6 + vllm/v1/kv_cache_interface.py | 2 + vllm/v1/simple_kv_offload/manager.py | 7 +- 17 files changed, 160 insertions(+), 63 deletions(-) diff --git a/tests/config/test_config_utils.py b/tests/config/test_config_utils.py index 35bc1e167b52..24ef1b52a95d 100644 --- a/tests/config/test_config_utils.py +++ b/tests/config/test_config_utils.py @@ -216,6 +216,11 @@ def test_cache_config_hash_ignores_kv_cache_sizing_knobs(): assert CacheConfig(gpu_memory_utilization=0.5).compute_hash() == base_hash +def test_cache_config_hash_ignores_prefix_cache_retention_interval(): + base_hash = CacheConfig().compute_hash() + assert CacheConfig(prefix_cache_retention_interval=64).compute_hash() == base_hash + + def test_envs_compile_factors_relocation_invariant(tmp_path): """Relocating HOME or the XDG roots must not change the compile-cache env hash. diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 2feb9f7a039d..6a799ee912d4 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -476,6 +476,7 @@ def test_prefix_cache_default(): # should be None by default (depends on model). engine_args = EngineArgs.from_cli_args(args=args) assert engine_args.enable_prefix_caching is None + assert engine_args.prefix_cache_retention_interval == 0 # with flag to turn it on. args = parser.parse_args(["--enable-prefix-caching"]) @@ -487,6 +488,28 @@ def test_prefix_cache_default(): engine_args = EngineArgs.from_cli_args(args=args) assert not engine_args.enable_prefix_caching + args = parser.parse_args(["--prefix-cache-retention-interval", "64"]) + engine_args = EngineArgs.from_cli_args(args=args) + assert engine_args.prefix_cache_retention_interval == 64 + + +def test_prefix_cache_retention_interval_from_deprecated_env( + monkeypatch, caplog, disable_log_dedup +): + monkeypatch.setenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", "64") + + engine_args = EngineArgs() + + assert engine_args.prefix_cache_retention_interval == 64 + assert "VLLM_PREFIX_CACHE_RETENTION_INTERVAL" in caplog.text + assert "deprecated" in caplog.text + assert "prefix_cache_retention_interval" in caplog.text + + parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) + args = parser.parse_args(["--prefix-cache-retention-interval", "32"]) + engine_args = EngineArgs.from_cli_args(args) + assert engine_args.prefix_cache_retention_interval == 32 + @pytest.mark.parametrize( ("arg", "expected", "option"), diff --git a/tests/v1/core/test_contiguous_kv_packing.py b/tests/v1/core/test_contiguous_kv_packing.py index 88d17e9acdc7..40906f621369 100644 --- a/tests/v1/core/test_contiguous_kv_packing.py +++ b/tests/v1/core/test_contiguous_kv_packing.py @@ -88,6 +88,7 @@ def _make_groups(n_c4, n_c128, n_swa): def _mock_vllm_config(kv_connector_extra_config: dict[str, str] | None = None): config = MagicMock() config.cache_config.num_gpu_blocks_override = None + config.cache_config.prefix_cache_retention_interval = 0 config.kv_transfer_config = None if kv_connector_extra_config is not None: config.kv_transfer_config = MagicMock() @@ -322,6 +323,7 @@ def test_hma_attention_groups_keep_default_backing(self): ) assert config.num_blocks == 32 + assert config.prefix_cache_retention_interval == 0 assert sum(t.size for t in config.kv_cache_tensors) == page_size * 2 * 32 assert config.kv_cache_tensors == [ KVCacheTensor(size=page_size * 32, shared_by=["full.0", "sw.0", "sw.1"]), diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 9b2afeb00181..bfe4996b1aa8 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -816,6 +816,7 @@ def test_metrics_empty_stats(): def test_get_kv_cache_configs_multiple_workers(): model_config = ModelConfig(max_model_len=16) vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.prefix_cache_retention_interval = None ref_kv_cache_spec = new_kv_cache_spec() same_kv_cache_specs = [ @@ -1173,6 +1174,7 @@ def test_get_kv_cache_configs_multiple_workers(): def test_get_kv_cache_configs_pp_sharding(asymmetric_memory): model_config = ModelConfig(max_model_len=512) vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.prefix_cache_retention_interval = None ref_kv_cache_spec = new_kv_cache_spec() pp_kv_cache_specs = [ @@ -1702,6 +1704,7 @@ def test_get_kv_cache_config_one_worker(): # pass max_model_len to pass check_enough_kv_cache_memory model_config = ModelConfig(max_model_len=16) vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.prefix_cache_retention_interval = None mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 # all layers are full attention -> single group @@ -2015,6 +2018,7 @@ def test_get_kv_cache_config_one_worker(): def test_get_kv_cache_configs_attention_free(): kv_cache_specs: dict[str, KVCacheSpec] = {} vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=16)) + vllm_config.cache_config.prefix_cache_retention_interval = None kv_cache_configs = get_kv_cache_configs(vllm_config, [kv_cache_specs], [0]) assert kv_cache_configs == [ KVCacheConfig( diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index 13ed7c7b9d8b..6cd117b1fa4a 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -4,6 +4,7 @@ import copy from collections.abc import Callable +from dataclasses import replace from math import lcm from types import SimpleNamespace @@ -110,6 +111,11 @@ def make_kv_cache_manager(kv_cache_config: KVCacheConfig, **kwargs) -> KVCacheMa "scheduler_block_size", lcm(*(g.kv_cache_spec.block_size for g in kv_cache_config.kv_cache_groups)), ) + if "retention_interval" in kwargs: + kv_cache_config = replace( + kv_cache_config, + prefix_cache_retention_interval=kwargs.pop("retention_interval"), + ) return KVCacheManager(kv_cache_config, **kwargs) @@ -3119,9 +3125,8 @@ def test_hybrid_cache_blocks_clamped_to_lcm(): ) -def test_hybrid_local_kv_retention_interval_aligns_in_manager(monkeypatch): +def test_hybrid_local_kv_retention_interval_aligns_in_manager(): """Verify fixed intervals retain sparse tails plus the latest replay tail.""" - monkeypatch.setenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", "64") block_size = 8 kv_cache_config = KVCacheConfig( num_blocks=100, @@ -3153,6 +3158,7 @@ def test_hybrid_local_kv_retention_interval_aligns_in_manager(monkeypatch): max_model_len=8192, enable_caching=True, hash_block_size=block_size, + retention_interval=64, ) # The SWA manager uses the configured 64-token interval (a multiple of the @@ -3184,18 +3190,15 @@ def test_hybrid_local_kv_retention_interval_aligns_in_manager(monkeypatch): "interval, expected_match", [ # scheduler_block_size is 32 (= lcm(4*8, 8)); 33 is not a multiple of it. - ("33", "multiple of scheduler_block_size"), + (33, "multiple of scheduler_block_size"), # A negative multiple (-32 % 32 == 0) must still be rejected explicitly, # otherwise it would pass the modulo check and silently degrade to dense. - ("-32", "non-negative"), + (-32, "non-negative"), ], ) -def test_hybrid_local_kv_retention_interval_rejects_invalid( - monkeypatch, interval, expected_match -): +def test_hybrid_local_kv_retention_interval_rejects_invalid(interval, expected_match): """A retention interval that is negative or not a multiple of scheduler_block_size errors out at construction time.""" - monkeypatch.setenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", interval) block_size = 8 kv_cache_config = KVCacheConfig( num_blocks=100, @@ -3228,12 +3231,36 @@ def test_hybrid_local_kv_retention_interval_rejects_invalid( max_model_len=8192, enable_caching=True, hash_block_size=block_size, + retention_interval=interval, ) -def test_hybrid_local_kv_retention_interval_survives_recycling(monkeypatch): +def test_zero_retention_is_ignored_for_full_attention(): + kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=10) + manager = make_kv_cache_manager( + kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=16, + retention_interval=0, + ) + assert manager.coordinator.retention_interval == 0 + + +def test_positive_retention_rejects_full_attention(): + kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=10) + with pytest.raises(ValueError, match="no sliding-window or Mamba"): + make_kv_cache_manager( + kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=16, + retention_interval=16, + ) + + +def test_hybrid_local_kv_retention_interval_survives_recycling(): """Verify retained local checkpoints are reused after block recycling.""" - monkeypatch.setenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", "1024") hash_block_size = 4 kv_cache_config = KVCacheConfig( num_blocks=800, @@ -3286,6 +3313,7 @@ def test_hybrid_local_kv_retention_interval_survives_recycling(monkeypatch): max_model_len=4096, enable_caching=True, hash_block_size=hash_block_size, + retention_interval=1024, ) def fill_request(request_id: str, token_offset: int) -> list[int]: @@ -3314,9 +3342,8 @@ def fill_request(request_id: str, token_offset: int) -> list[int]: assert [len(blocks) for blocks in computed_blocks.blocks] == [4, 16, 128, 256] -def test_hybrid_local_kv_retention_latest_only_reuses_replay_boundary(monkeypatch): +def test_hybrid_local_kv_retention_latest_only_reuses_replay_boundary(): """Verify latest-only retention reuses only the replayable prompt boundary.""" - monkeypatch.setenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", "0") block_size = 8 kv_cache_config = KVCacheConfig( num_blocks=100, @@ -3348,6 +3375,7 @@ def test_hybrid_local_kv_retention_latest_only_reuses_replay_boundary(monkeypatc max_model_len=8192, enable_caching=True, hash_block_size=block_size, + retention_interval=0, ) token_ids = [i for i in range(16) for _ in range(block_size)] @@ -3388,14 +3416,13 @@ def test_hybrid_local_kv_retention_latest_only_reuses_replay_boundary(monkeypatc assert len(computed_blocks.blocks[1]) == 0 -def test_hybrid_local_kv_retention_mtp_reuses_latest_boundary(monkeypatch): +def test_hybrid_local_kv_retention_mtp_reuses_latest_boundary(): """Verify MTP/EAGLE SWA retention keeps the extra proof block. EAGLE/MTP lookup matches one additional local block after the returned prefix and then drops it. Sparse retention must therefore cache the normal local tail at the latest replay boundary plus one extra SWA block. """ - monkeypatch.setenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", "0") block_size = 8 kv_cache_config = KVCacheConfig( num_blocks=100, @@ -3428,6 +3455,7 @@ def test_hybrid_local_kv_retention_mtp_reuses_latest_boundary(monkeypatch): max_model_len=8192, enable_caching=True, hash_block_size=block_size, + retention_interval=0, use_eagle=True, ) @@ -3824,12 +3852,11 @@ def test_cache_hit_local_and_external_two_groups_preempt_and_reallocate(): assert manager.get_blocks("test").get_block_ids() != ([], []) -def test_swa_free_split_keeps_cached_tail_ahead_of_scratch(monkeypatch): - """Default path (no retention): freeing an SWA request must place its +def test_swa_free_split_keeps_cached_tail_ahead_of_scratch(): + """Dense retention: freeing an SWA request must place its uncached scratch blocks at the front of the free queue (recycled first) and keep its cached checkpoint blocks at the back (retained for prefix hits). This split is always-on, independent of the retention interval.""" - monkeypatch.delenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", raising=False) block_size = 8 kv_cache_config = KVCacheConfig( num_blocks=100, @@ -3937,13 +3964,14 @@ def _make_pure_swa_manager(block_size, sliding_window, num_blocks=100, **kwargs) ) -def test_pure_swa_retention_interval_caches_sparse_tails(monkeypatch): +def test_pure_swa_retention_interval_caches_sparse_tails(): """Sparse retention must work for a pure-SWA single-group model, not just hybrid models: only the per-interval tails plus the latest replay tail are cached, and a replay still hits the latest replayable boundary.""" - monkeypatch.setenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", "64") block_size = 16 - manager = _make_pure_swa_manager(block_size, sliding_window=block_size) + manager = _make_pure_swa_manager( + block_size, sliding_window=block_size, retention_interval=64 + ) assert type(manager.coordinator).__name__ == "UnitaryKVCacheCoordinator" token_ids = [i for i in range(16) for _ in range(block_size)] @@ -3976,11 +4004,12 @@ def test_pure_swa_retention_interval_caches_sparse_tails(monkeypatch): assert num_computed == 240 -def test_pure_swa_retention_latest_only(monkeypatch): +def test_pure_swa_retention_latest_only(): """`=0` on a pure-SWA model keeps only the latest replay tail.""" - monkeypatch.setenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", "0") block_size = 16 - manager = _make_pure_swa_manager(block_size, sliding_window=block_size) + manager = _make_pure_swa_manager( + block_size, sliding_window=block_size, retention_interval=0 + ) token_ids = [i for i in range(16) for _ in range(block_size)] req = make_request("0", token_ids, block_size, sha256) @@ -4008,10 +4037,9 @@ def test_pure_swa_retention_latest_only(monkeypatch): assert num_computed == 240 -def test_pure_swa_retention_dense_default_caches_all(monkeypatch): - """With retention unset, a pure-SWA model must keep the dense behavior: +def test_pure_swa_dense_retention_caches_all(): + """With retention set to ``None``, a pure-SWA model keeps dense behavior: every block boundary is a potential hit, so all blocks are cached.""" - monkeypatch.delenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", raising=False) block_size = 16 manager = _make_pure_swa_manager(block_size, sliding_window=block_size) @@ -4037,7 +4065,7 @@ def test_pure_swa_retention_dense_default_caches_all(monkeypatch): def test_mamba_reachable_block_mask_sparsifies_retention(): - """Mamba state-snapshot retention: with VLLM_PREFIX_CACHE_RETENTION_INTERVAL + """Mamba state-snapshot retention: with a configured retention interval, the manager keeps one cached state per interval-sized segment (plus the latest replay boundary) instead of a snapshot per block, which is what lets a small attention block_size avoid Mamba dominating the KV pool.""" @@ -4063,7 +4091,7 @@ def retained(retention_interval, num_prompt_tokens=256, end_block=16): ) return None if m is None else {i for i, v in enumerate(m) if v} - # Dense default (None) -> no mask, every block cached (unchanged behavior). + # Dense retention (None) -> no mask, every block cached. assert retained(None) is None # interval == block_size -> every block is a boundary -> stays dense. assert retained(block_size) is None @@ -4111,7 +4139,7 @@ def retained(retention_interval, shared_prefix_boundary, end_block=16): assert retained(0, 100) == {5, 14} # Coexists with segment tails (interval 64 -> {3,7,11,15} + replay 14). assert retained(64, 96) == {3, 5, 7, 11, 14, 15} - # Dense default ignores the hint (nothing to sparsify). + # Dense retention ignores the hint (nothing to sparsify). assert retained(None, 96) is None # Out-of-range boundary is a no-op (only replay 14 remains). assert retained(0, 16 * block_size * 2) == {14} @@ -4120,14 +4148,13 @@ def retained(retention_interval, shared_prefix_boundary, end_block=16): assert retained(0, None) == {14} -def test_mamba_shared_prefix_survives_zero_retention(monkeypatch): +def test_mamba_shared_prefix_survives_zero_retention(): """Manager-level check of the full wiring: a pinned shared-prefix boundary (``Request.shared_prefix_boundary``, set by the scheduler on Marconi-style detection) keeps its Mamba state block cached under - ``VLLM_PREFIX_CACHE_RETENTION_INTERVAL=0``, which otherwise retains only the + ``prefix_cache_retention_interval=0``, which otherwise retains only the end-of-prompt replay boundary. Without this, a shared prefix (junction before ``num_prompt``) would be recomputed by every sharing request.""" - monkeypatch.setenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", "0") block_size = 16 # 16-block (256-token) prompt; replay boundary is block 240 // 16 - 1 = 14. @@ -4140,6 +4167,7 @@ def cached_mamba_blocks(shared_prefix_boundary): max_model_len=8192, enable_caching=True, hash_block_size=block_size, + retention_interval=0, ) req = make_request("r", token_ids, block_size, sha256) req.shared_prefix_boundary = shared_prefix_boundary @@ -4163,24 +4191,21 @@ def cached_mamba_blocks(shared_prefix_boundary): assert cached_mamba_blocks(96) == {5, 14} -def test_mamba_shared_prefix_reuse_under_zero_retention(monkeypatch): +def test_mamba_shared_prefix_reuse_under_zero_retention(): """Full cross-request Marconi flow: a partial shared prefix cached by the detecting request must stay reusable by a later request under - ``VLLM_PREFIX_CACHE_RETENTION_INTERVAL=0``. Without the pin the junction is + ``prefix_cache_retention_interval=0``. Without the pin the junction is masked out and the later request misses; with it (and under dense) the reuse is preserved.""" block_size = 16 def last_req_hit(retention, pin): - if retention is None: - monkeypatch.delenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", raising=False) - else: - monkeypatch.setenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", str(retention)) manager = make_kv_cache_manager( _make_hybrid_kv_cache_config(block_size, 200, ["full", "mamba_align"]), max_model_len=8192, enable_caching=True, hash_block_size=block_size, + retention_interval=retention, ) shared = [7 for _ in range(2 * block_size)] # 2-block shared prefix @@ -4261,23 +4286,20 @@ def retained(retention, boundary, window, end_block=16): assert retained(0, 0, block_size) == {14} -def test_swa_shared_prefix_reuse_under_zero_retention(monkeypatch): +def test_swa_shared_prefix_reuse_under_zero_retention(): """SWA cross-request analog: a partial shared prefix's sliding-window tail - must stay reusable under ``VLLM_PREFIX_CACHE_RETENTION_INTERVAL=0``. Without + must stay reusable under ``prefix_cache_retention_interval=0``. Without the pin the junction window is masked out and a later request misses; with it (and under dense) reuse is preserved.""" block_size = 16 def last_req_hit(retention, pin): - if retention is None: - monkeypatch.delenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", raising=False) - else: - monkeypatch.setenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", str(retention)) manager = make_kv_cache_manager( _make_hybrid_kv_cache_config(block_size, 200, ["full", "sliding_window"]), max_model_len=8192, enable_caching=True, hash_block_size=block_size, + retention_interval=retention, ) shared = [7 for _ in range(4 * block_size)] # 4-block shared prefix diff --git a/tests/v1/engine/test_engine_args.py b/tests/v1/engine/test_engine_args.py index 5033f4768bf7..b1ba1adcfd0a 100644 --- a/tests/v1/engine/test_engine_args.py +++ b/tests/v1/engine/test_engine_args.py @@ -18,6 +18,7 @@ def test_prefix_caching_from_cli(): assert vllm_config.cache_config.enable_prefix_caching, ( "V1 turns on prefix caching by default." ) + assert vllm_config.cache_config.prefix_cache_retention_interval == 0 # Turn it off possible with flag. args = parser.parse_args(["--no-enable-prefix-caching"]) @@ -47,6 +48,10 @@ def test_prefix_caching_from_cli(): with pytest.raises(ArgumentError): args = parser.parse_args(["--prefix-caching-hash-algo", "invalid"]) + args = parser.parse_args(["--prefix-cache-retention-interval", "64"]) + vllm_config = EngineArgs.from_cli_args(args=args).create_engine_config() + assert vllm_config.cache_config.prefix_cache_retention_interval == 64 + @pytest.mark.skipif(_xxhash is None, reason="xxhash not installed") def test_prefix_caching_xxhash_from_cli(): diff --git a/tests/v1/engine/test_engine_core_client.py b/tests/v1/engine/test_engine_core_client.py index 5a943e8d6044..74f7b803b890 100644 --- a/tests/v1/engine/test_engine_core_client.py +++ b/tests/v1/engine/test_engine_core_client.py @@ -1066,6 +1066,7 @@ def test_kv_cache_events( model=model_name, enforce_eager=True, enable_prefix_caching=True, + prefix_cache_retention_interval=None, block_size=block_size, ) engine_args.kv_events_config = publisher_config diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index 4655c5716ec8..50d2ee16149e 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -247,7 +247,9 @@ def _make_vllm_config( ) -def _make_kv_cache_config(*, block_size: int = 16) -> object: +def _make_kv_cache_config( + *, block_size: int = 16, prefix_cache_retention_interval: int | None = 0 +) -> object: """Minimal single-group KVCacheConfig for topology tests.""" from vllm.v1.kv_cache_interface import ( FullAttentionSpec, @@ -262,6 +264,7 @@ def _make_kv_cache_config(*, block_size: int = 16) -> object: num_blocks=10, kv_cache_tensors=[], kv_cache_groups=[KVCacheGroupSpec(["layer0"], spec)], + prefix_cache_retention_interval=prefix_cache_retention_interval, ) @@ -1594,6 +1597,7 @@ def test_requester_worker_init_uses_positional_setup(tmp_path, monkeypatch): "mlx5_0", "10.0.0.7:50051", ) + assert w.coord.retention_interval == 0 def test_requester_worker_init_prefers_local_hostname_override( diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 0981f710ea57..27e46534dcd3 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -7,7 +7,7 @@ from pydantic import Field, field_validator, model_validator -from vllm.config.utils import config +from vllm.config.utils import config, get_from_deprecated_env_if_set from vllm.logger import init_logger from vllm.utils.torch_utils import ( is_quantized_kv_cache, @@ -35,6 +35,17 @@ "nvfp4", "nvfp4_4over6", ] + + +def _get_prefix_cache_retention_interval() -> int | None: + env_value = get_from_deprecated_env_if_set( + "VLLM_PREFIX_CACHE_RETENTION_INTERVAL", + "v0.29", + "prefix_cache_retention_interval", + ) + return 0 if env_value is None else int(env_value) + + MambaDType = Literal["auto", "float32", "float16", "bfloat16"] MambaCacheMode = Literal["all", "align", "none"] PrefixCachingHashAlgo = Literal["sha256", "sha256_cbor", "xxhash", "xxhash_cbor"] @@ -111,6 +122,15 @@ class CacheConfig: security risk tolerance against the performance benefits before turning this on. - "xxhash_cbor" combines canonical CBOR serialization with xxHash for reproducible hashing. Requires the optional ``xxhash`` package.""" + prefix_cache_retention_interval: int | None = Field( + default_factory=_get_prefix_cache_retention_interval, ge=0 + ) + """Token interval between retained sliding-window and Mamba prefix-cache + checkpoints. ``0`` retains only semantic checkpoints, including the latest + replay boundary and shared-prefix junctions. Positive values additionally + retain periodic checkpoints at the specified interval, which must be a + multiple of the scheduler block size. ``None`` retains checkpoints densely. + Applies only to sliding-window and Mamba cache groups.""" kv_cache_dtype_skip_layers: list[str] = field(default_factory=list) """Layer patterns to skip KV cache quantization. Accepts layer indices (e.g., '0', '2', '4') or attention type names (e.g., 'sliding_window').""" @@ -218,6 +238,7 @@ def compute_hash(self) -> str: "num_gpu_blocks_override", "enable_prefix_caching", "prefix_caching_hash_algo", + "prefix_cache_retention_interval", # Prefix-caching implementation detail (doesn't affect compiled graph). "prefix_match_unit", "mamba_page_size_padded", diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index e25aec3fbde9..83c6d16e1b9e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -1433,7 +1433,7 @@ def __init__( scheduler_block_size=self.block_size, hash_block_size=self.hash_block_size, use_eagle=use_eagle, - retention_interval=envs.VLLM_PREFIX_CACHE_RETENTION_INTERVAL, + retention_interval=kv_cache_config.prefix_cache_retention_interval, ) # One ChunkedTokenDatabase per group; addresses populated in # register_kv_caches once the kv-cache layout is known. Each group's diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 44ee837c6fef..1fb1388f100c 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -522,6 +522,9 @@ class EngineArgs: prefix_caching_hash_algo: PrefixCachingHashAlgo = ( CacheConfig.prefix_caching_hash_algo ) + prefix_cache_retention_interval: int | None = get_field( + CacheConfig, "prefix_cache_retention_interval" + ) disable_sliding_window: bool = ModelConfig.disable_sliding_window disable_cascade_attn: bool = ModelConfig.disable_cascade_attn offload_backend: str = OffloadConfig.offload_backend @@ -1224,6 +1227,10 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: cache_group.add_argument( "--prefix-caching-hash-algo", **cache_kwargs["prefix_caching_hash_algo"] ) + cache_group.add_argument( + "--prefix-cache-retention-interval", + **cache_kwargs["prefix_cache_retention_interval"], + ) cache_group.add_argument( "--kv-cache-dtype-skip-layers", **cache_kwargs["kv_cache_dtype_skip_layers"] ) @@ -2001,6 +2008,7 @@ def create_engine_config( sliding_window=sliding_window, enable_prefix_caching=self.enable_prefix_caching, prefix_caching_hash_algo=self.prefix_caching_hash_algo, + prefix_cache_retention_interval=self.prefix_cache_retention_interval, kv_cache_dtype_skip_layers=self.kv_cache_dtype_skip_layers, kv_sharing_fast_prefill=self.kv_sharing_fast_prefill, mamba_cache_dtype=self.mamba_cache_dtype, diff --git a/vllm/envs.py b/vllm/envs.py index ed62659cfad3..15705b7c1779 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -1148,12 +1148,6 @@ def _resolve_rust_cli_path() -> str | None: if "VLLM_PLUGINS" not in os.environ else os.environ["VLLM_PLUGINS"].split(",") ), - # Retain local sliding-window KV checkpoints for prefix caching. - # Unset (default) preserves the dense local checkpointing behavior. `0` - # retains only the latest completed prompt boundary. Positive values retain - # checkpoints at the specified interval boundaries (rounded up to the - # prefix-cache alignment). - # Applies to sliding-window attention for now but not yet Mamba/linear attention. "VLLM_PREFIX_CACHE_RETENTION_INTERVAL": lambda: ( int(os.environ["VLLM_PREFIX_CACHE_RETENTION_INTERVAL"]) if "VLLM_PREFIX_CACHE_RETENTION_INTERVAL" in os.environ diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index 8efaf9252e8a..0829ec0d2be8 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -4,7 +4,6 @@ from collections.abc import Sequence from typing import NamedTuple -from vllm import envs from vllm.logger import init_logger from vllm.utils.math_utils import cdiv, round_down from vllm.v1.core.block_pool import BlockPool @@ -45,16 +44,18 @@ def _validate_prefix_cache_retention_interval( isinstance(g.kv_cache_spec, (SlidingWindowSpec, MambaSpec)) for g in kv_cache_config.kv_cache_groups ): + if retention_interval == 0: + return raise ValueError( - "VLLM_PREFIX_CACHE_RETENTION_INTERVAL is set but this model has " + "prefix_cache_retention_interval is set but this model has " "no sliding-window or Mamba KV cache group, so retention has no " - "effect. Unset it (it only applies to sliding-window and Mamba " + "effect. Set it to 0 (it only applies to sliding-window and Mamba " "attention)." ) if retention_interval < 0 or retention_interval % scheduler_block_size != 0: raise ValueError( - f"VLLM_PREFIX_CACHE_RETENTION_INTERVAL ({retention_interval}) " + f"prefix_cache_retention_interval ({retention_interval}) " "must be non-negative and a multiple of scheduler_block_size " f"({scheduler_block_size})." ) @@ -151,7 +152,7 @@ def __init__( # A positive retention interval must be a multiple of the base hit granularity # (``scheduler_block_size``) to land on real cache-hit boundaries. # 0 = keep only the latest replay boundary; None = dense; - self.retention_interval = envs.VLLM_PREFIX_CACHE_RETENTION_INTERVAL + self.retention_interval = kv_cache_config.prefix_cache_retention_interval _validate_prefix_cache_retention_interval( self.retention_interval, self.scheduler_block_size, kv_cache_config ) diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 44097c3da276..d1af91b65993 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -243,7 +243,7 @@ def get_computed_blocks(self, request: Request) -> tuple[KVCacheBlocks, int, int - ``shared_prefix_boundary``: the block-aligned token position of a shared prefix that a sparse-retention group (Mamba / sliding window) has not cached yet (Marconi-style APC), or 0 if none. - Pinned so ``VLLM_PREFIX_CACHE_RETENTION_INTERVAL`` does not drop + Pinned so sparse prefix-cache retention does not drop the junction and defeat cross-request reuse. """ # We skip finding the prefix cache hit when prefix caching is diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index ad51e4b2b397..d6e401d1184e 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -1347,6 +1347,9 @@ def get_kv_cache_config_from_groups( num_blocks=1, kv_cache_tensors=[], kv_cache_groups=kv_cache_groups, + prefix_cache_retention_interval=( + vllm_config.cache_config.prefix_cache_retention_interval + ), ) # Determine how model runners should initialize the KV cache tensors. @@ -1406,6 +1409,9 @@ def get_kv_cache_config_from_groups( num_blocks=num_blocks, kv_cache_tensors=kv_cache_tensors, kv_cache_groups=kv_cache_groups, + prefix_cache_retention_interval=( + vllm_config.cache_config.prefix_cache_retention_interval + ), ) diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 90391f6ef0a6..cbead6d3c885 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -970,6 +970,8 @@ class KVCacheConfig: For models with multiple types of attention, there will be multiple groups, see `_get_kv_cache_config_uniform_page_size` for more details. """ + prefix_cache_retention_interval: int | None = None + """Resolved retention policy for local prefix-cache checkpoints.""" @property def has_mamba_layers(self) -> bool: diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index 2e6839fef557..1e4fcae68cf9 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/vllm/v1/simple_kv_offload/manager.py @@ -4,7 +4,7 @@ import contextlib from collections.abc import Iterable -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from typing import TYPE_CHECKING, Any from vllm.config import VllmConfig @@ -190,7 +190,6 @@ def _derive_cpu_config( """Derive a CPU KVCacheConfig from the GPU config. Same kv_cache_groups, num_blocks scaled by CPU/GPU memory ratio.""" # Import here to avoid potential circular imports - from vllm.v1.kv_cache_interface import KVCacheConfig as KVCacheConfigCls from vllm.v1.kv_cache_interface import KVCacheTensor assert len(gpu_config.kv_cache_tensors) > 0 @@ -215,10 +214,10 @@ def _derive_cpu_config( for t in gpu_config.kv_cache_tensors ] - return KVCacheConfigCls( + return replace( + gpu_config, num_blocks=num_cpu_blocks, kv_cache_tensors=cpu_tensors, - kv_cache_groups=gpu_config.kv_cache_groups, ) @staticmethod From 4ab5e5012a27e2b751c679873f231bafa0f6b098 Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Mon, 17 Aug 2026 12:45:45 -0400 Subject: [PATCH 048/839] [Refactor] Simplify B12X linear kernels and warmup (#52368) Signed-off-by: mgoin --- .buildkite/test_areas/kernels.yaml | 22 + tests/kernels/quantization/test_block_fp8.py | 2 +- ...2x_mxfp8_linear.py => test_b12x_linear.py} | 618 +++++++----------- .../kernels/test_b12x_mxfp4_linear.py | 199 ------ .../kernels/test_b12x_nvfp4_linear.py | 246 ------- tests/model_executor/test_b12x_warmup.py | 226 +++++-- .../model_executor/kernels/linear/__init__.py | 4 +- .../kernels/linear/mxfp4/b12x.py | 110 +--- .../kernels/linear/mxfp8/b12x.py | 126 ++-- .../kernels/linear/nvfp4/b12x.py | 119 ++-- .../kernels/linear/scaled_mm/b12x.py | 339 ++++++++++ .../kernels/linear/scaled_mm/b12x_block.py | 214 ------ .../kernels/linear/scaled_mm/b12x_tensor.py | 238 ------- vllm/model_executor/warmup/b12x_warmup.py | 92 +-- vllm/utils/b12x.py | 11 +- 15 files changed, 959 insertions(+), 1607 deletions(-) rename tests/model_executor/kernels/{test_b12x_mxfp8_linear.py => test_b12x_linear.py} (62%) delete mode 100644 tests/model_executor/kernels/test_b12x_mxfp4_linear.py delete mode 100644 tests/model_executor/kernels/test_b12x_nvfp4_linear.py create mode 100644 vllm/model_executor/kernels/linear/scaled_mm/b12x.py delete mode 100644 vllm/model_executor/kernels/linear/scaled_mm/b12x_block.py delete mode 100644 vllm/model_executor/kernels/linear/scaled_mm/b12x_tensor.py diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 4e73d7620e51..887ed03e372a 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -339,6 +339,28 @@ steps: # e2e - pytest -v -s tests/models/quantization/test_nvfp4.py +- label: B12X Linear Kernels (DGX Spark) Nightly + key: b12x-linear-kernels-dgx-spark-nightly + timeout_in_minutes: 30 + device: dgx-spark + optional: true + num_devices: 1 + depends_on: + - arm64-image-build + source_file_dependencies: + - setup.py + - vllm/model_executor/kernels/linear/ + - vllm/model_executor/warmup/b12x_warmup.py + - vllm/utils/b12x.py + - tests/model_executor/kernels/test_b12x_linear.py + - tests/model_executor/test_b12x_warmup.py + - tests/kernels/quantization/test_block_fp8.py + commands: + - uv pip install --system b12x==1.2.4 + - pytest -v -s model_executor/kernels/test_b12x_linear.py + model_executor/test_b12x_warmup.py + kernels/quantization/test_block_fp8.py -k b12x + - label: Kernels Helion Test key: kernels-helion-test timeout_in_minutes: 115 diff --git a/tests/kernels/quantization/test_block_fp8.py b/tests/kernels/quantization/test_block_fp8.py index 6f86aed59e68..52f87c403e75 100644 --- a/tests/kernels/quantization/test_block_fp8.py +++ b/tests/kernels/quantization/test_block_fp8.py @@ -14,7 +14,7 @@ ) from tests.kernels.utils import fp8_ulp_distance from vllm.config import VllmConfig -from vllm.model_executor.kernels.linear.scaled_mm.b12x_block import ( +from vllm.model_executor.kernels.linear.scaled_mm.b12x import ( B12xFp8BlockScaledMMKernel, _run_b12x_fp8_block_scaled_mm, ) diff --git a/tests/model_executor/kernels/test_b12x_mxfp8_linear.py b/tests/model_executor/kernels/test_b12x_linear.py similarity index 62% rename from tests/model_executor/kernels/test_b12x_mxfp8_linear.py rename to tests/model_executor/kernels/test_b12x_linear.py index 74f74e2f4f49..1efc521bd8b4 100644 --- a/tests/model_executor/kernels/test_b12x_mxfp8_linear.py +++ b/tests/model_executor/kernels/test_b12x_linear.py @@ -3,6 +3,7 @@ from __future__ import annotations +import importlib import types from dataclasses import dataclass @@ -13,141 +14,120 @@ _LINEAR_BACKEND_KERNEL_MAP, _POSSIBLE_FP8_BLOCK_KERNELS, _POSSIBLE_FP8_KERNELS, + _POSSIBLE_MXFP4_KERNELS, _POSSIBLE_MXFP8_KERNELS, - init_fp8_linear_kernel, - init_mxfp8_linear_kernel, -) -from vllm.model_executor.kernels.linear.mxfp8.b12x import ( - B12xMxfp8LinearKernel, - _b12x_mxfp8_expected_m, - warmup_b12x_mxfp8_linear, -) -from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import ( - Mxfp8LinearLayerConfig, -) -from vllm.model_executor.kernels.linear.scaled_mm.b12x_block import ( + _POSSIBLE_NVFP4_KERNELS, B12xFp8BlockScaledMMKernel, - _run_b12x_fp8_block_scaled_mm, - warmup_b12x_block_fp8_linear, -) -from vllm.model_executor.kernels.linear.scaled_mm.b12x_tensor import ( + B12xMxFp4LinearKernel, + B12xMxfp8LinearKernel, + B12xNvFp4LinearKernel, B12xTensorFP8ScaledMMLinearKernel, - warmup_b12x_tensor_fp8_linear, -) -from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import ( FP8ScaledMMLinearLayerConfig, + Mxfp8LinearLayerConfig, + init_fp8_linear_kernel, + init_mxfp4_linear_kernel, + init_mxfp8_linear_kernel, + init_nvfp4_linear_kernel, ) -from vllm.model_executor.layers.linear import ( - ReplicatedLinear, - UnquantizedLinearMethod, +from vllm.model_executor.kernels.linear.nvfp4.marlin import ( + MarlinNvFp4LinearKernel, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8Dynamic128Sym, kFp8Static128BlockSym, kFp8StaticTensorSym, + kMxfp4Dynamic, ) -from vllm.platforms import PlatformEnum, current_platform -from vllm.utils.b12x import b12x_warmup_token_counts - - -def test_b12x_backend_maps_mxfp8_kernel() -> None: - assert B12xMxfp8LinearKernel in _LINEAR_BACKEND_KERNEL_MAP["b12x"] - assert B12xMxfp8LinearKernel in _POSSIBLE_MXFP8_KERNELS[PlatformEnum.CUDA] - - -def test_b12x_backend_maps_tensor_fp8_kernel() -> None: - assert B12xTensorFP8ScaledMMLinearKernel in _LINEAR_BACKEND_KERNEL_MAP["b12x"] - assert B12xTensorFP8ScaledMMLinearKernel in _POSSIBLE_FP8_KERNELS[PlatformEnum.CUDA] +from vllm.platforms import PlatformEnum @pytest.mark.parametrize( - ("kernels", "before", "b12x", "after"), + ("kernel_cls", "kernels", "before", "after", "initializer", "kwargs"), [ ( + B12xMxFp4LinearKernel, + _POSSIBLE_MXFP4_KERNELS[PlatformEnum.CUDA], + "HummingMxFp4LinearKernel", + "EmulationMxfp4LinearKernel", + init_mxfp4_linear_kernel, + {"activation_quant_key": kMxfp4Dynamic}, + ), + ( + B12xNvFp4LinearKernel, + _POSSIBLE_NVFP4_KERNELS[PlatformEnum.CUDA], + "FbgemmNvFp4LinearKernel", + "EmulationNvFp4LinearKernel", + init_nvfp4_linear_kernel, + {}, + ), + ( + B12xMxfp8LinearKernel, + _POSSIBLE_MXFP8_KERNELS[PlatformEnum.CUDA], + "MarlinMxfp8LinearKernel", + "EmulationMxfp8LinearKernel", + init_mxfp8_linear_kernel, + {}, + ), + ( + B12xTensorFP8ScaledMMLinearKernel, _POSSIBLE_FP8_KERNELS[PlatformEnum.CUDA], "CutlassFP8ScaledMMLinearKernel", - "B12xTensorFP8ScaledMMLinearKernel", "PerTensorTorchFP8ScaledMMLinearKernel", + init_fp8_linear_kernel, + { + "activation_quant_key": kFp8StaticTensorSym, + "weight_quant_key": kFp8StaticTensorSym, + "input_dtype": torch.bfloat16, + "out_dtype": torch.bfloat16, + "weight_shape": (2048, 2048), + }, ), ( + B12xFp8BlockScaledMMKernel, _POSSIBLE_FP8_BLOCK_KERNELS[PlatformEnum.CUDA], "CutlassFp8BlockScaledMMKernel", - "B12xFp8BlockScaledMMKernel", "MarlinFP8ScaledMMLinearKernel", - ), - ( - _POSSIBLE_MXFP8_KERNELS[PlatformEnum.CUDA], - "MarlinMxfp8LinearKernel", - "B12xMxfp8LinearKernel", - "EmulationMxfp8LinearKernel", + init_fp8_linear_kernel, + { + "activation_quant_key": kFp8Dynamic128Sym, + "weight_quant_key": kFp8Static128BlockSym, + "input_dtype": torch.bfloat16, + "out_dtype": torch.bfloat16, + "weight_shape": (2048, 2048), + }, ), ], ) -def test_b12x_fp8_fallback_priority( - kernels: list[type], - before: str, - b12x: str, - after: str, -) -> None: - names = [kernel.__name__ for kernel in kernels] - - assert names.index(before) < names.index(b12x) < names.index(after) - - -@torch.inference_mode() -def test_b12x_backend_does_not_intercept_unquantized_bf16( - default_vllm_config, - dist_init, -) -> None: - default_vllm_config.kernel_config.linear_backend = "b12x" - device = current_platform.device_type - layer = ReplicatedLinear( - 128, - 64, - bias=False, - params_dtype=torch.bfloat16, - quant_config=None, - prefix="bf16_linear", - ).to(device) - layer.weight.data.normal_() - x = torch.randn(8, 128, device=device, dtype=torch.bfloat16) - - output, output_bias = layer(x) - expected = torch.nn.functional.linear(x, layer.weight) - - assert isinstance(layer.quant_method, UnquantizedLinearMethod) - assert output_bias is None - torch.testing.assert_close(output, expected) - - -def test_b12x_explicit_backend_selects_per_tensor_fp8( +def test_b12x_backend_registration_priority_and_selection( monkeypatch, default_vllm_config, + kernel_cls, + kernels, + before: str, + after: str, + initializer, + kwargs: dict, ) -> None: import vllm.model_executor.kernels.linear as linear_mod + assert kernel_cls in _LINEAR_BACKEND_KERNEL_MAP["b12x"] + names = [kernel.__name__ for kernel in kernels] + assert names.index(before) < names.index(kernel_cls.__name__) < names.index(after) + monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA) monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x") monkeypatch.setattr( - B12xTensorFP8ScaledMMLinearKernel, + kernel_cls, "is_supported", classmethod(lambda cls, compute_capability=None: (True, None)), ) monkeypatch.setattr( - B12xTensorFP8ScaledMMLinearKernel, + kernel_cls, "can_implement", classmethod(lambda cls, config: (True, None)), ) - kernel = init_fp8_linear_kernel( - activation_quant_key=kFp8StaticTensorSym, - weight_quant_key=kFp8StaticTensorSym, - input_dtype=torch.bfloat16, - out_dtype=torch.bfloat16, - weight_shape=(2048, 2048), - ) - - assert isinstance(kernel, B12xTensorFP8ScaledMMLinearKernel) + assert isinstance(initializer(**kwargs), kernel_cls) def test_b12x_tensor_fp8_can_implement_supported_config() -> None: @@ -165,33 +145,8 @@ def test_b12x_tensor_fp8_can_implement_supported_config() -> None: assert reason is None -def test_b12x_explicit_backend_selects_block_fp8( - monkeypatch, - default_vllm_config, -) -> None: - import vllm.model_executor.kernels.linear as linear_mod - - monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA) - monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x") - monkeypatch.setattr( - B12xFp8BlockScaledMMKernel, - "is_supported", - classmethod(lambda cls, compute_capability=None: (True, None)), - ) - - kernel = init_fp8_linear_kernel( - activation_quant_key=kFp8Dynamic128Sym, - weight_quant_key=kFp8Static128BlockSym, - input_dtype=torch.bfloat16, - out_dtype=torch.bfloat16, - weight_shape=(2048, 2048), - ) - - assert isinstance(kernel, B12xFp8BlockScaledMMKernel) - - def test_b12x_block_fp8_checks_runtime_support(monkeypatch) -> None: - import vllm.model_executor.kernels.linear.scaled_mm.b12x_block as b12x_mod + import vllm.model_executor.kernels.linear.scaled_mm.b12x as b12x_mod platform = types.SimpleNamespace( is_cuda=lambda: True, @@ -261,75 +216,10 @@ def can_implement(weight_shape: tuple[int, int]): ) -def test_b12x_warmup_token_counts_cover_serving_regimes() -> None: - assert b12x_warmup_token_counts( - max_tokens=2048, - cudagraph_capture_sizes=[1, 2, 8, 128], - ) == (1, 2, 8, 128, 2048) - - -def test_warmup_b12x_block_fp8_dedupes_weight_signatures(monkeypatch) -> None: - import vllm.model_executor.kernels.linear.scaled_mm.b12x_block as b12x_mod - - calls = [] - - def run(a, weight, a_scale, weight_scale, out_dtype): - calls.append((a.shape, weight, a_scale.shape, weight_scale, out_dtype)) - return torch.empty((a.shape[0], weight.shape[0]), dtype=out_dtype) - - platform = types.SimpleNamespace( - is_cuda=lambda: True, - is_device_capability_family=lambda family: family == 120, - ) - monkeypatch.setattr(b12x_mod, "current_platform", platform) - - monkeypatch.setattr( - b12x_mod, - "_import_b12x_blockscaled", - lambda: types.SimpleNamespace(), - ) - monkeypatch.setattr(b12x_mod, "_run_b12x_fp8_block_scaled_mm", run) - - def layer(in_features: int, out_features: int): - return types.SimpleNamespace( - b12x_block_fp8_linear=True, - weight=torch.empty((out_features, in_features), dtype=torch.float8_e4m3fn), - weight_scale_inv=torch.empty( - (out_features // 128, in_features // 128), dtype=torch.float32 - ), - ) - - layer_a = layer(128, 256) - layer_b = layer(256, 128) - modules = [ - layer_a, - layer_a, - layer_b, - types.SimpleNamespace(), - ] - model = types.SimpleNamespace(modules=lambda: iter(modules)) - - warmed = warmup_b12x_block_fp8_linear( - model, - max_tokens=32, - cudagraph_capture_sizes=[2, 8], - output_dtype=torch.bfloat16, - ) - - assert warmed == 8 - assert [call[0][0] for call in calls] == [1, 2, 8, 32] * 2 - assert [call[2][0] for call in calls] == [1, 2, 8, 32] * 2 - assert calls[0][1] is layer_a.weight - assert calls[4][1] is layer_b.weight - assert calls[0][3] is layer_a.weight_scale_inv - assert calls[4][3] is layer_b.weight_scale_inv - assert all(call[4] == torch.bfloat16 for call in calls) - - def test_b12x_tensor_fp8_process_weights_packs_modelopt_layout( monkeypatch, ) -> None: - import vllm.model_executor.kernels.linear.scaled_mm.b12x_tensor as b12x_mod + import vllm.model_executor.kernels.linear.scaled_mm.b12x as b12x_mod calls = [] packed = types.SimpleNamespace(out_features=64) @@ -367,6 +257,7 @@ def pack(weight: torch.Tensor, output_scale: torch.Tensor): kernel.process_weights_after_loading(layer) assert layer.b12x_tensor_fp8_packed_weight is packed + assert layer.b12x_warmup_provider is kernel assert len(calls) == 1 weight, output_scale = calls[0] torch.testing.assert_close(weight, original_weight.T.contiguous()) @@ -378,67 +269,10 @@ def pack(weight: torch.Tensor, output_scale: torch.Tensor): torch.testing.assert_close(layer.input_scale, torch.tensor(0.5)) -def test_warmup_b12x_tensor_fp8_dedupes_weight_signatures(monkeypatch) -> None: - import vllm.model_executor.kernels.linear.scaled_mm.b12x_tensor as b12x_mod - - calls = [] - - def prewarm(packed_weight, token_counts, *, out_dtype, stream): - del stream - calls.append((packed_weight, tuple(token_counts), out_dtype)) - return len(tuple(token_counts)) - - platform = types.SimpleNamespace( - is_cuda=lambda: True, - is_device_capability_family=lambda family: family == 120, - ) - monkeypatch.setattr(b12x_mod, "current_platform", platform) - monkeypatch.setattr( - b12x_mod, - "_import_b12x_tensor_fp8", - lambda: types.SimpleNamespace(prewarm=prewarm), - ) - monkeypatch.setattr( - b12x_mod, - "current_stream", - lambda: types.SimpleNamespace(cuda_stream=object()), - ) - - def packed(in_features: int, padded_in_features: int, out_features: int): - return types.SimpleNamespace( - in_features=in_features, - padded_in_features=padded_in_features, - out_features=out_features, - values=torch.empty(1), - ) - - packed_a = packed(128, 128, 256) - packed_b = packed(160, 256, 512) - modules = [ - types.SimpleNamespace(b12x_tensor_fp8_packed_weight=packed_a), - types.SimpleNamespace(b12x_tensor_fp8_packed_weight=packed_a), - types.SimpleNamespace(b12x_tensor_fp8_packed_weight=packed_b), - types.SimpleNamespace(), - ] - model = types.SimpleNamespace(modules=lambda: iter(modules)) - - warmed = warmup_b12x_tensor_fp8_linear( - model, - max_tokens=2048, - cudagraph_capture_sizes=[1, 2], - ) - - assert warmed == 6 - assert calls == [ - (packed_a, (1, 2, 2048), torch.bfloat16), - (packed_b, (1, 2, 2048), torch.bfloat16), - ] - - def test_b12x_tensor_fp8_apply_quantizes_and_uses_packed_weight( monkeypatch, ) -> None: - import vllm.model_executor.kernels.linear.scaled_mm.b12x_tensor as b12x_mod + import vllm.model_executor.kernels.linear.scaled_mm.b12x as b12x_mod calls = [] @@ -506,27 +340,6 @@ def mm( assert expected_m == 6 -def test_b12x_mxfp8_explicit_backend_selects_kernel(monkeypatch) -> None: - import vllm.model_executor.kernels.linear as linear_mod - - monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA) - monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x") - monkeypatch.setattr( - B12xMxfp8LinearKernel, - "is_supported", - classmethod(lambda cls, compute_capability=None: (True, None)), - ) - monkeypatch.setattr( - B12xMxfp8LinearKernel, - "can_implement", - classmethod(lambda cls, c: (True, None)), - ) - - kernel = init_mxfp8_linear_kernel() - - assert isinstance(kernel, B12xMxfp8LinearKernel) - - def test_b12x_mxfp8_can_implement_supported_config() -> None: can_implement, reason = B12xMxfp8LinearKernel.can_implement( Mxfp8LinearLayerConfig() @@ -536,90 +349,6 @@ def test_b12x_mxfp8_can_implement_supported_config() -> None: assert reason is None -def test_b12x_mxfp8_expected_m_uses_live_m() -> None: - assert _b12x_mxfp8_expected_m(0) == 1 - assert _b12x_mxfp8_expected_m(1) == 1 - assert _b12x_mxfp8_expected_m(2) == 2 - assert _b12x_mxfp8_expected_m(8) == 8 - assert _b12x_mxfp8_expected_m(9) == 9 - assert _b12x_mxfp8_expected_m(128) == 128 - assert _b12x_mxfp8_expected_m(129) == 129 - assert _b12x_mxfp8_expected_m(2048) == 2048 - - -def test_warmup_b12x_mxfp8_linear_dedupes_weight_signatures( - monkeypatch, -) -> None: - import vllm.model_executor.kernels.linear.mxfp8.b12x as b12x_mod - - calls = [] - - def mm( - source: torch.Tensor, - packed_weight, - *, - bias: torch.Tensor | None = None, - expected_m: int | None = None, - stream: object = None, - ) -> torch.Tensor: - del stream - calls.append((source.shape, packed_weight, bias, expected_m)) - return source.new_empty((source.shape[0], packed_weight.out_features)) - - platform = types.SimpleNamespace( - is_cuda=lambda: True, - is_device_capability_family=lambda family: family == 120, - ) - monkeypatch.setattr(b12x_mod, "current_platform", platform) - monkeypatch.setattr( - b12x_mod, - "_import_b12x_mxfp8", - lambda: types.SimpleNamespace(mm=mm), - ) - monkeypatch.setattr( - b12x_mod, - "current_stream", - lambda: types.SimpleNamespace(cuda_stream=object()), - ) - - def packed(in_features: int, padded_in_features: int, out_features: int): - return types.SimpleNamespace( - in_features=in_features, - padded_in_features=padded_in_features, - out_features=out_features, - weight=types.SimpleNamespace(values=torch.empty(1)), - ) - - packed_a = packed(128, 128, 256) - packed_b = packed(128, 128, 512) - modules = [ - types.SimpleNamespace(b12x_mxfp8_packed_weight=packed_a), - types.SimpleNamespace(b12x_mxfp8_packed_weight=packed_a), - types.SimpleNamespace(b12x_mxfp8_packed_weight=packed_b), - types.SimpleNamespace(), - ] - model = types.SimpleNamespace(modules=lambda: iter(modules)) - - warmed = warmup_b12x_mxfp8_linear( - model, - max_tokens=2048, - cudagraph_capture_sizes=[1, 2], - ) - - assert warmed == 6 - assert [call[0] for call in calls] == [ - torch.Size([1, 128]), - torch.Size([2, 128]), - torch.Size([2048, 128]), - torch.Size([1, 128]), - torch.Size([2, 128]), - torch.Size([2048, 128]), - ] - assert [call[3] for call in calls] == [1, 2, 2048, 1, 2, 2048] - assert calls[0][1] is packed_a - assert calls[3][1] is packed_b - - def test_b12x_mxfp8_support_check_reports_missing_import(monkeypatch) -> None: import vllm.model_executor.kernels.linear.mxfp8.b12x as b12x_mod @@ -693,6 +422,7 @@ def pack(weight: torch.Tensor, weight_scale: torch.Tensor): kernel.process_weights_after_loading(layer) assert layer.b12x_mxfp8_packed_weight is packed + assert layer.b12x_warmup_provider is kernel assert len(calls) == 1 weight, weight_scale = calls[0] assert weight.shape == (48, 128) @@ -805,7 +535,7 @@ def test_b12x_block_fp8_process_weights_keeps_native_block_layout() -> None: kernel.process_weights_after_loading(layer) - assert layer.b12x_block_fp8_linear + assert layer.b12x_warmup_provider is kernel assert layer.weight.shape == (128, 128) assert layer.weight.dtype == torch.float8_e4m3fn assert layer.weight_scale_inv.shape == (1, 1) @@ -882,7 +612,7 @@ def mxfp8_linear( def test_b12x_block_fp8_apply_uses_b12x_recipe_api(monkeypatch) -> None: - import vllm.model_executor.kernels.linear.scaled_mm.b12x_block as b12x_mod + import vllm.model_executor.kernels.linear.scaled_mm.b12x as b12x_mod calls = [] @@ -919,37 +649,183 @@ def mm_block_fp8(*args, **kwargs): torch.testing.assert_close(output, torch.full_like(output, 13.0)) -def test_b12x_block_fp8_helper_uses_regular_compact_scale_api(monkeypatch) -> None: - import vllm.model_executor.kernels.linear.scaled_mm.b12x_block as b12x_mod +def test_b12x_mxfp4_requires_dynamic_activations() -> None: + config = types.SimpleNamespace(activation_quant_key=kMxfp4Dynamic) + can_implement, reason = B12xMxFp4LinearKernel.can_implement(config) - calls = [] + assert can_implement + assert reason is None - def mm_block_fp8(*args, **kwargs): + config.activation_quant_key = None + can_implement, reason = B12xMxFp4LinearKernel.can_implement(config) + + assert not can_implement + assert reason == "B12X MXFP4 GEMM requires dynamic MXFP4 activations" + + +@pytest.mark.parametrize( + ("kernel_cls", "module_name", "scale_dtype"), + [ + ( + B12xMxFp4LinearKernel, + "vllm.model_executor.kernels.linear.mxfp4.b12x", + torch.uint8, + ), + ( + B12xNvFp4LinearKernel, + "vllm.model_executor.kernels.linear.nvfp4.b12x", + torch.float8_e4m3fn, + ), + ], +) +def test_b12x_fp4_processes_scale_and_preserves_loader( + monkeypatch, + kernel_cls, + module_name: str, + scale_dtype: torch.dtype, +) -> None: + scale = torch.empty((48, 8), dtype=scale_dtype) + swizzled_scale = torch.empty((128, 8), dtype=scale_dtype) + intrinsics = types.SimpleNamespace(swizzle_block_scale=lambda value: swizzled_scale) + monkeypatch.setattr( + importlib.import_module(module_name), + "_import_b12x_intrinsics", + lambda: intrinsics, + ) + layer = torch.nn.Module() + layer.prefix = "model.layers.0.mlp.shared_expert.down_proj" + layer.weight_scale = torch.nn.Parameter(scale, requires_grad=False) + weight_loader = object() + layer.weight_scale.weight_loader = weight_loader + kernel = object.__new__(kernel_cls) + + kernel.process_weights_after_loading(layer) + + assert layer.weight_scale.data_ptr() == swizzled_scale.data_ptr() + assert layer.weight_scale.weight_loader is weight_loader + assert layer.b12x_warmup_provider is kernel + + +def test_b12x_mxfp4_apply_calls_native_blockscaled_gemm(monkeypatch) -> None: + import vllm.model_executor.kernels.linear.mxfp4.b12x as b12x_mod + import vllm.utils.flashinfer as flashinfer_utils + + calls: list[tuple] = [] + x_packed = torch.empty((6, 64), dtype=torch.uint8) + x_scale_storage = torch.empty((128, 4), dtype=torch.uint8) + + def mm_mxfp4(*args, **kwargs): calls.append((args, kwargs)) - return torch.full((6, 256), 17.0, dtype=torch.bfloat16) + return torch.full((6, 48), 3.0, dtype=torch.bfloat16) + monkeypatch.setattr( + flashinfer_utils, + "flashinfer_mxfp4_quantize", + lambda *args, **kwargs: (x_packed, x_scale_storage), + ) monkeypatch.setattr( b12x_mod, "_import_b12x_blockscaled", - lambda: types.SimpleNamespace(mm_block_fp8=mm_block_fp8), + lambda: types.SimpleNamespace(mm_mxfp4=mm_mxfp4), ) - a = torch.empty((6, 128), dtype=torch.float8_e4m3fn) - weight = torch.empty((256, 128), dtype=torch.float8_e4m3fn) - a_scale = torch.empty((6, 1), dtype=torch.float32) - weight_scale = torch.empty((2, 1), dtype=torch.float32) - output = _run_b12x_fp8_block_scaled_mm( - a, - weight, - a_scale, - weight_scale, - torch.bfloat16, + layer = torch.nn.Module() + layer.output_size_per_partition = 48 + layer.weight = torch.empty((48, 64), dtype=torch.uint8) + layer.weight_scale = torch.empty((128, 4), dtype=torch.uint8) + x = torch.empty((2, 3, 128), dtype=torch.bfloat16) + bias = torch.ones(48, dtype=torch.bfloat16) + kernel = object.__new__(B12xMxFp4LinearKernel) + + output = kernel.apply_weights(layer, x, bias) + + assert output.shape == (2, 3, 48) + torch.testing.assert_close(output, torch.full_like(output, 4.0)) + assert len(calls) == 1 + args, kwargs = calls[0] + assert args == ( + x_packed, + x_scale_storage, + layer.weight, + layer.weight_scale, ) + assert kwargs == {"out_dtype": torch.bfloat16} - assert output.shape == (6, 256) + +def test_b12x_nvfp4_can_implement_supported_config() -> None: + can_implement, reason = B12xNvFp4LinearKernel.can_implement(None) + + assert can_implement + assert reason is None + + +def test_b12x_backend_preserves_w4a16_fallback(monkeypatch) -> None: + import vllm.model_executor.kernels.linear as linear_mod + + monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA) + monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x") + monkeypatch.setattr( + MarlinNvFp4LinearKernel, + "is_supported", + classmethod(lambda cls, compute_capability=None: (True, None)), + ) + + kernel = init_nvfp4_linear_kernel(use_a16=True) + + assert isinstance(kernel, MarlinNvFp4LinearKernel) + + +def test_b12x_nvfp4_apply_calls_native_blockscaled_gemm(monkeypatch) -> None: + import vllm.model_executor.kernels.linear.nvfp4.b12x as b12x_mod + + calls: list[tuple] = [] + quant_calls: list[tuple] = [] + x_packed = torch.empty((6, 64), dtype=torch.uint8) + x_scale_storage = torch.empty((128, 8), dtype=torch.float8_e4m3fn) + + def quant(*args, **kwargs): + quant_calls.append((args, kwargs)) + return x_packed, x_scale_storage + + def mm_nvfp4(*args, **kwargs): + calls.append((args, kwargs)) + return torch.full((6, 48), 3.0, dtype=torch.bfloat16) + + monkeypatch.setattr(b12x_mod, "scaled_fp4_quant", quant) + monkeypatch.setattr( + b12x_mod, + "_import_b12x_blockscaled", + lambda: types.SimpleNamespace(mm_nvfp4=mm_nvfp4), + ) + + layer = torch.nn.Module() + layer.output_size_per_partition = 48 + layer.weight = torch.empty((48, 64), dtype=torch.uint8) + layer.weight_scale = torch.empty((128, 8), dtype=torch.float8_e4m3fn) + layer.input_global_scale_inv = torch.tensor(2.0) + layer.alpha = torch.tensor(0.25) + x = torch.empty((2, 3, 256), dtype=torch.bfloat16)[..., ::2] + bias = torch.ones(48, dtype=torch.bfloat16) + kernel = object.__new__(B12xNvFp4LinearKernel) + + output = kernel.apply_weights(layer, x, bias) + + assert output.shape == (2, 3, 48) + torch.testing.assert_close(output, torch.full_like(output, 4.0)) + assert len(quant_calls) == 1 + quant_args, quant_kwargs = quant_calls[0] + assert quant_args[0].shape == (6, 128) + assert quant_args[0].data_ptr() == x.data_ptr() + assert quant_args[1] is layer.input_global_scale_inv + assert not quant_args[0].is_contiguous() + assert quant_kwargs == {"is_sf_swizzled_layout": True} + assert len(calls) == 1 args, kwargs = calls[0] - assert args == (a, a_scale, weight, weight_scale) - assert kwargs == { - "out_dtype": torch.bfloat16, - } - torch.testing.assert_close(output, torch.full_like(output, 17.0)) + assert args == ( + x_packed, + x_scale_storage, + layer.weight, + layer.weight_scale, + layer.alpha, + ) + assert kwargs == {"out_dtype": torch.bfloat16} diff --git a/tests/model_executor/kernels/test_b12x_mxfp4_linear.py b/tests/model_executor/kernels/test_b12x_mxfp4_linear.py deleted file mode 100644 index 7ba1aad81717..000000000000 --- a/tests/model_executor/kernels/test_b12x_mxfp4_linear.py +++ /dev/null @@ -1,199 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from __future__ import annotations - -import types - -import torch - -from vllm.model_executor.kernels.linear import ( - _LINEAR_BACKEND_KERNEL_MAP, - _POSSIBLE_MXFP4_KERNELS, - init_mxfp4_linear_kernel, -) -from vllm.model_executor.kernels.linear.mxfp4.b12x import ( - B12xMxFp4LinearKernel, - warmup_b12x_mxfp4_linear, -) -from vllm.model_executor.layers.quantization.utils.quant_utils import ( - kMxfp4Dynamic, -) -from vllm.platforms import PlatformEnum - - -def test_b12x_backend_maps_mxfp4_kernel() -> None: - assert B12xMxFp4LinearKernel in _LINEAR_BACKEND_KERNEL_MAP["b12x"] - assert B12xMxFp4LinearKernel in _POSSIBLE_MXFP4_KERNELS[PlatformEnum.CUDA] - - -def test_b12x_mxfp4_fallback_priority() -> None: - kernels = _POSSIBLE_MXFP4_KERNELS[PlatformEnum.CUDA] - names = [kernel.__name__ for kernel in kernels] - - assert ( - names.index("HummingMxFp4LinearKernel") - < names.index("B12xMxFp4LinearKernel") - < names.index("EmulationMxfp4LinearKernel") - ) - - -def test_b12x_mxfp4_explicit_backend_selects_native_kernel(monkeypatch) -> None: - import vllm.model_executor.kernels.linear as linear_mod - - monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA) - monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x") - monkeypatch.setattr( - B12xMxFp4LinearKernel, - "is_supported", - classmethod(lambda cls, compute_capability=None: (True, None)), - ) - monkeypatch.setattr( - B12xMxFp4LinearKernel, - "can_implement", - classmethod(lambda cls, config: (True, None)), - ) - - kernel = init_mxfp4_linear_kernel(activation_quant_key=kMxfp4Dynamic) - - assert isinstance(kernel, B12xMxFp4LinearKernel) - - -def test_b12x_mxfp4_requires_dynamic_activations() -> None: - config = types.SimpleNamespace(activation_quant_key=kMxfp4Dynamic) - can_implement, reason = B12xMxFp4LinearKernel.can_implement(config) - - assert can_implement - assert reason is None - - config.activation_quant_key = None - can_implement, reason = B12xMxFp4LinearKernel.can_implement(config) - - assert not can_implement - assert reason == "B12X MXFP4 GEMM requires dynamic MXFP4 activations" - - -def test_b12x_mxfp4_processes_scale_and_preserves_loader(monkeypatch) -> None: - import vllm.model_executor.kernels.linear.mxfp4.b12x as b12x_mod - - scale = torch.empty((48, 8), dtype=torch.uint8) - swizzled_scale = torch.empty((128, 8), dtype=torch.uint8) - monkeypatch.setattr( - b12x_mod, - "_import_b12x_intrinsics", - lambda: types.SimpleNamespace(swizzle_block_scale=lambda value: swizzled_scale), - ) - - layer = torch.nn.Module() - layer.prefix = "model.layers.0.mlp.shared_expert.down_proj" - layer.weight_scale = torch.nn.Parameter(scale, requires_grad=False) - weight_loader = object() - layer.weight_scale.weight_loader = weight_loader - kernel = object.__new__(B12xMxFp4LinearKernel) - - kernel.process_weights_after_loading(layer) - - assert layer.weight_scale.data_ptr() == swizzled_scale.data_ptr() - assert layer.weight_scale.weight_loader is weight_loader - assert layer.b12x_mxfp4_linear - - -def test_b12x_mxfp4_apply_calls_native_blockscaled_gemm(monkeypatch) -> None: - import vllm.model_executor.kernels.linear.mxfp4.b12x as b12x_mod - import vllm.utils.flashinfer as flashinfer_utils - - calls: list[tuple] = [] - x_packed = torch.empty((6, 64), dtype=torch.uint8) - x_scale_storage = torch.empty((128, 4), dtype=torch.uint8) - - def mm_mxfp4(*args, **kwargs): - calls.append((args, kwargs)) - return torch.full((6, 48), 3.0, dtype=torch.bfloat16) - - monkeypatch.setattr( - flashinfer_utils, - "flashinfer_mxfp4_quantize", - lambda *args, **kwargs: (x_packed, x_scale_storage), - ) - monkeypatch.setattr( - b12x_mod, - "_import_b12x_blockscaled", - lambda: types.SimpleNamespace(mm_mxfp4=mm_mxfp4), - ) - - layer = torch.nn.Module() - layer.output_size_per_partition = 48 - layer.weight = torch.empty((48, 64), dtype=torch.uint8) - layer.weight_scale = torch.empty((128, 4), dtype=torch.uint8) - x = torch.empty((2, 3, 128), dtype=torch.bfloat16) - bias = torch.ones(48, dtype=torch.bfloat16) - kernel = object.__new__(B12xMxFp4LinearKernel) - - output = kernel.apply_weights(layer, x, bias) - - assert output.shape == (2, 3, 48) - torch.testing.assert_close(output, torch.full_like(output, 4.0)) - assert len(calls) == 1 - args, kwargs = calls[0] - assert args == ( - x_packed, - x_scale_storage, - layer.weight, - layer.weight_scale, - ) - assert kwargs == {"out_dtype": torch.bfloat16} - - -def test_warmup_b12x_mxfp4_dedupes_weight_signatures(monkeypatch) -> None: - import vllm.model_executor.kernels.linear.mxfp4.b12x as b12x_mod - - calls = [] - - def apply(source, weight, weight_scale, bias): - calls.append((source.shape, weight, weight_scale, bias)) - return source.new_empty((source.shape[0], weight.shape[0])) - - platform = types.SimpleNamespace( - is_cuda=lambda: True, - is_device_capability_family=lambda family: family == 120, - ) - monkeypatch.setattr(b12x_mod, "current_platform", platform) - monkeypatch.setattr( - b12x_mod, - "_import_b12x_blockscaled", - lambda: types.SimpleNamespace(), - ) - monkeypatch.setattr(b12x_mod, "_apply_b12x_mxfp4_linear", apply) - - def layer(n: int): - return types.SimpleNamespace( - b12x_mxfp4_linear=True, - weight=torch.empty((n, 64), dtype=torch.uint8), - weight_scale=torch.empty((128, 4), dtype=torch.uint8), - ) - - layer_a = layer(48) - layer_b = layer(48) - layer_c = layer(96) - model = types.SimpleNamespace( - modules=lambda: iter([layer_a, layer_b, layer_c, types.SimpleNamespace()]) - ) - - warmed = warmup_b12x_mxfp4_linear( - model, - max_tokens=8, - cudagraph_capture_sizes=[1, 2], - ) - - assert warmed == 6 - assert [call[0] for call in calls] == [ - torch.Size([1, 128]), - torch.Size([2, 128]), - torch.Size([8, 128]), - torch.Size([1, 128]), - torch.Size([2, 128]), - torch.Size([8, 128]), - ] - assert calls[0][1] is layer_a.weight - assert calls[3][1] is layer_c.weight - assert all(call[3] is None for call in calls) diff --git a/tests/model_executor/kernels/test_b12x_nvfp4_linear.py b/tests/model_executor/kernels/test_b12x_nvfp4_linear.py deleted file mode 100644 index 1e2717c4e05d..000000000000 --- a/tests/model_executor/kernels/test_b12x_nvfp4_linear.py +++ /dev/null @@ -1,246 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from __future__ import annotations - -import types - -import torch - -from vllm.model_executor.kernels.linear import ( - _LINEAR_BACKEND_KERNEL_MAP, - _POSSIBLE_NVFP4_KERNELS, - _resolve_backend_kernels, - init_nvfp4_linear_kernel, -) -from vllm.model_executor.kernels.linear.nvfp4.b12x import ( - B12xNvFp4LinearKernel, - warmup_b12x_nvfp4_linear, -) -from vllm.model_executor.kernels.linear.nvfp4.marlin import ( - MarlinNvFp4LinearKernel, -) -from vllm.platforms import PlatformEnum - - -def test_b12x_backend_maps_nvfp4_kernel() -> None: - assert B12xNvFp4LinearKernel in _LINEAR_BACKEND_KERNEL_MAP["b12x"] - assert B12xNvFp4LinearKernel in _POSSIBLE_NVFP4_KERNELS[PlatformEnum.CUDA] - - -def test_b12x_nvfp4_fallback_priority() -> None: - kernels = _POSSIBLE_NVFP4_KERNELS[PlatformEnum.CUDA] - names = [kernel.__name__ for kernel in kernels] - - assert ( - names.index("FbgemmNvFp4LinearKernel") - < names.index("B12xNvFp4LinearKernel") - < names.index("EmulationNvFp4LinearKernel") - ) - - -def test_b12x_nvfp4_explicit_backend_selects_native_kernel(monkeypatch) -> None: - import vllm.model_executor.kernels.linear as linear_mod - - monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA) - monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x") - monkeypatch.setattr( - B12xNvFp4LinearKernel, - "is_supported", - classmethod(lambda cls, compute_capability=None: (True, None)), - ) - monkeypatch.setattr( - B12xNvFp4LinearKernel, - "can_implement", - classmethod(lambda cls, config: (True, None)), - ) - - kernel = init_nvfp4_linear_kernel() - - assert isinstance(kernel, B12xNvFp4LinearKernel) - - -def test_b12x_nvfp4_can_implement_supported_config() -> None: - can_implement, reason = B12xNvFp4LinearKernel.can_implement(None) - - assert can_implement - assert reason is None - - -def test_b12x_backend_preserves_w4a16_fallback(monkeypatch) -> None: - import vllm.model_executor.kernels.linear as linear_mod - - monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA) - monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x") - monkeypatch.setattr( - MarlinNvFp4LinearKernel, - "is_supported", - classmethod(lambda cls, compute_capability=None: (True, None)), - ) - - kernel = init_nvfp4_linear_kernel(use_a16=True) - - assert isinstance(kernel, MarlinNvFp4LinearKernel) - - -def test_backend_without_w4a16_kernel_preserves_fallback(monkeypatch) -> None: - import vllm.model_executor.kernels.linear as linear_mod - - kernels = [MarlinNvFp4LinearKernel] - monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x") - assert _resolve_backend_kernels(kernels, "NVFP4") == kernels - - -def test_b12x_nvfp4_processes_scale_and_preserves_loader(monkeypatch) -> None: - import vllm.model_executor.kernels.linear.nvfp4.b12x as b12x_mod - - scale = torch.empty((48, 8), dtype=torch.float8_e4m3fn) - swizzled_scale = torch.empty((128, 8), dtype=torch.float8_e4m3fn) - intrinsics = types.SimpleNamespace(swizzle_block_scale=lambda value: swizzled_scale) - monkeypatch.setattr(b12x_mod, "_import_b12x_intrinsics", lambda: intrinsics) - - layer = torch.nn.Module() - layer.prefix = "model.layers.0.mlp.shared_expert.down_proj" - layer.weight_scale = torch.nn.Parameter(scale, requires_grad=False) - weight_loader = object() - layer.weight_scale.weight_loader = weight_loader - kernel = object.__new__(B12xNvFp4LinearKernel) - - kernel.process_weights_after_loading(layer) - - assert layer.weight_scale.data_ptr() == swizzled_scale.data_ptr() - assert layer.weight_scale.weight_loader is weight_loader - assert layer.b12x_nvfp4_linear - - -def test_b12x_nvfp4_apply_calls_native_blockscaled_gemm(monkeypatch) -> None: - import vllm.model_executor.kernels.linear.nvfp4.b12x as b12x_mod - - calls: list[tuple] = [] - quant_calls: list[tuple] = [] - x_packed = torch.empty((6, 64), dtype=torch.uint8) - x_scale_storage = torch.empty((128, 8), dtype=torch.float8_e4m3fn) - - def quant(*args, **kwargs): - quant_calls.append((args, kwargs)) - return x_packed, x_scale_storage - - def mm_nvfp4(*args, **kwargs): - calls.append((args, kwargs)) - return torch.full((6, 48), 3.0, dtype=torch.bfloat16) - - monkeypatch.setattr( - b12x_mod, - "scaled_fp4_quant", - quant, - ) - monkeypatch.setattr( - b12x_mod, - "_import_b12x_blockscaled", - lambda: types.SimpleNamespace(mm_nvfp4=mm_nvfp4), - ) - - layer = torch.nn.Module() - layer.output_size_per_partition = 48 - layer.weight = torch.empty((48, 64), dtype=torch.uint8) - layer.weight_scale = torch.empty((128, 8), dtype=torch.float8_e4m3fn) - layer.input_global_scale_inv = torch.tensor(2.0) - layer.alpha = torch.tensor(0.25) - x = torch.empty((2, 3, 256), dtype=torch.bfloat16)[..., ::2] - bias = torch.ones(48, dtype=torch.bfloat16) - kernel = object.__new__(B12xNvFp4LinearKernel) - - output = kernel.apply_weights(layer, x, bias) - - assert output.shape == (2, 3, 48) - torch.testing.assert_close(output, torch.full_like(output, 4.0)) - assert len(quant_calls) == 1 - quant_args, quant_kwargs = quant_calls[0] - assert quant_args[0].shape == (6, 128) - assert quant_args[0].data_ptr() == x.data_ptr() - assert quant_args[1] is layer.input_global_scale_inv - assert not quant_args[0].is_contiguous() - assert quant_kwargs == {"is_sf_swizzled_layout": True} - assert len(calls) == 1 - args, kwargs = calls[0] - assert args == ( - x_packed, - x_scale_storage, - layer.weight, - layer.weight_scale, - layer.alpha, - ) - assert kwargs == {"out_dtype": torch.bfloat16} - - -def test_warmup_b12x_nvfp4_dedupes_weight_signatures(monkeypatch) -> None: - import vllm.model_executor.kernels.linear.nvfp4.b12x as b12x_mod - - calls = [] - - def apply( - source, - weight, - weight_scale, - input_global_scale_inv, - alpha, - bias, - ): - calls.append( - ( - source.shape, - weight, - weight_scale, - input_global_scale_inv, - alpha, - bias, - ) - ) - return source.new_empty((source.shape[0], weight.shape[0])) - - platform = types.SimpleNamespace( - is_cuda=lambda: True, - is_device_capability_family=lambda family: family == 120, - ) - monkeypatch.setattr(b12x_mod, "current_platform", platform) - monkeypatch.setattr( - b12x_mod, - "_import_b12x_blockscaled", - lambda: types.SimpleNamespace(), - ) - monkeypatch.setattr(b12x_mod, "_apply_b12x_nvfp4_linear", apply) - - def layer(n: int): - return types.SimpleNamespace( - b12x_nvfp4_linear=True, - weight=torch.empty((n, 64), dtype=torch.uint8), - weight_scale=torch.empty((128, 8), dtype=torch.float8_e4m3fn), - input_global_scale_inv=torch.tensor(2.0), - alpha=torch.tensor(0.25), - ) - - layer_a = layer(48) - layer_b = layer(48) - layer_c = layer(96) - model = types.SimpleNamespace( - modules=lambda: iter([layer_a, layer_b, layer_c, types.SimpleNamespace()]) - ) - - warmed = warmup_b12x_nvfp4_linear( - model, - max_tokens=8, - cudagraph_capture_sizes=[1, 2], - ) - - assert warmed == 6 - assert [call[0] for call in calls] == [ - torch.Size([1, 128]), - torch.Size([2, 128]), - torch.Size([8, 128]), - torch.Size([1, 128]), - torch.Size([2, 128]), - torch.Size([8, 128]), - ] - assert calls[0][1] is layer_a.weight - assert calls[3][1] is layer_c.weight - assert all(call[5] is None for call in calls) diff --git a/tests/model_executor/test_b12x_warmup.py b/tests/model_executor/test_b12x_warmup.py index bd71f9bcd207..26069e0a9ad1 100644 --- a/tests/model_executor/test_b12x_warmup.py +++ b/tests/model_executor/test_b12x_warmup.py @@ -7,94 +7,200 @@ import pytest import torch +from vllm.model_executor.kernels.linear import ( + B12xFp8BlockScaledMMKernel, + B12xMxFp4LinearKernel, + B12xMxfp8LinearKernel, + B12xNvFp4LinearKernel, + B12xTensorFP8ScaledMMLinearKernel, +) from vllm.model_executor.warmup.b12x_warmup import b12x_warmup +from vllm.utils.b12x import B12xWarmupUnit, b12x_warmup_token_counts + + +def test_b12x_warmup_token_counts_cover_serving_regimes() -> None: + assert b12x_warmup_token_counts( + max_tokens=2048, + cudagraph_capture_sizes=[1, 2, 8, 128], + ) == (1, 2, 8, 128, 2048) @pytest.mark.parametrize( - ("module_name", "warmup_name", "import_name"), + ("kernel_cls", "module_name", "call_name", "layer", "name"), [ ( - "vllm.model_executor.kernels.linear.scaled_mm.b12x_block", - "warmup_b12x_block_fp8_linear", - "_import_b12x_blockscaled", - ), - ( - "vllm.model_executor.kernels.linear.scaled_mm.b12x_tensor", - "warmup_b12x_tensor_fp8_linear", - "_import_b12x_tensor_fp8", - ), - ( - "vllm.model_executor.kernels.linear.mxfp8.b12x", - "warmup_b12x_mxfp8_linear", - "_import_b12x_mxfp8", - ), - ( + B12xMxFp4LinearKernel, "vllm.model_executor.kernels.linear.mxfp4.b12x", - "warmup_b12x_mxfp4_linear", - "_import_b12x_blockscaled", + "_apply_b12x_mxfp4_linear", + SimpleNamespace( + weight=torch.empty((48, 64), dtype=torch.uint8), + weight_scale=torch.empty((128, 4), dtype=torch.uint8), + ), + "MXFP4", ), ( + B12xNvFp4LinearKernel, "vllm.model_executor.kernels.linear.nvfp4.b12x", - "warmup_b12x_nvfp4_linear", - "_import_b12x_blockscaled", + "_apply_b12x_nvfp4_linear", + SimpleNamespace( + weight=torch.empty((48, 64), dtype=torch.uint8), + weight_scale=torch.empty((128, 8), dtype=torch.float8_e4m3fn), + input_global_scale_inv=torch.tensor(2.0), + alpha=torch.tensor(0.25), + ), + "NVFP4", + ), + ( + B12xFp8BlockScaledMMKernel, + "vllm.model_executor.kernels.linear.scaled_mm.b12x", + "_run_b12x_fp8_block_scaled_mm", + SimpleNamespace( + weight=torch.empty((256, 128), dtype=torch.float8_e4m3fn), + weight_scale_inv=torch.empty((2, 1), dtype=torch.float32), + ), + "block-FP8", ), ], ) -def test_b12x_linear_warmup_skips_unused_provider( +def test_b12x_warmup_units_cover_token_counts( monkeypatch, + kernel_cls, module_name: str, - warmup_name: str, - import_name: str, + call_name: str, + layer, + name: str, ) -> None: - module = importlib.import_module(module_name) - platform = SimpleNamespace( - is_cuda=lambda: True, - is_device_capability_family=lambda family: family == 120, + calls = [] + monkeypatch.setattr( + importlib.import_module(module_name), + call_name, + lambda *args: calls.append(args), ) - monkeypatch.setattr(module, "current_platform", platform) + kernel = object.__new__(kernel_cls) - def fail_import(): - pytest.fail("unused B12X provider was imported") + unit = kernel.get_b12x_warmup_unit(layer, (1, 8), torch.bfloat16) + unit.compile() - monkeypatch.setattr(module, import_name, fail_import) + assert unit.name == name + assert [args[0].shape[0] for args in calls] == [1, 8] + assert unit.key[-1] == torch.bfloat16 - warmup = getattr(module, warmup_name) - assert warmup(torch.nn.Module(), max_tokens=128) == 0 +def test_b12x_mxfp8_warmup_unit(monkeypatch) -> None: + import vllm.model_executor.kernels.linear.mxfp8.b12x as b12x_mod -def test_b12x_warmup_covers_linear_serving_shapes(monkeypatch) -> None: - import vllm.model_executor.warmup.b12x_warmup as warmup_mod + calls = [] + monkeypatch.setattr( + b12x_mod, + "_import_b12x_mxfp8", + lambda: SimpleNamespace( + mm=lambda *args, **kwargs: calls.append((args, kwargs)) + ), + ) + monkeypatch.setattr( + b12x_mod, + "current_stream", + lambda: SimpleNamespace(cuda_stream=object()), + ) + packed_weight = SimpleNamespace( + in_features=128, + padded_in_features=128, + out_features=256, + weight=SimpleNamespace(values=torch.empty(1)), + ) + layer = SimpleNamespace(b12x_mxfp8_packed_weight=packed_weight) + kernel = object.__new__(B12xMxfp8LinearKernel) - model = torch.nn.Module() - linear_calls: list[tuple[torch.nn.Module, dict]] = [] + unit = kernel.get_b12x_warmup_unit(layer, (1, 8), torch.float16) + unit.compile() - def linear_warmup(model, **kwargs): - linear_calls.append((model, kwargs)) - return 0 + assert [args[0].shape for args, _ in calls] == [(1, 128), (8, 128)] + assert [kwargs["expected_m"] for _, kwargs in calls] == [1, 8] - monkeypatch.setattr(warmup_mod, "warmup_b12x_block_fp8_linear", linear_warmup) - monkeypatch.setattr(warmup_mod, "warmup_b12x_mxfp4_linear", linear_warmup) - monkeypatch.setattr(warmup_mod, "warmup_b12x_mxfp8_linear", linear_warmup) - monkeypatch.setattr(warmup_mod, "warmup_b12x_nvfp4_linear", linear_warmup) - monkeypatch.setattr(warmup_mod, "warmup_b12x_tensor_fp8_linear", linear_warmup) - worker = SimpleNamespace( - get_model=lambda: model, - scheduler_config=SimpleNamespace( - max_num_batched_tokens=256, - max_num_scheduled_tokens=320, +def test_b12x_tensor_fp8_warmup_unit(monkeypatch) -> None: + import vllm.model_executor.kernels.linear.scaled_mm.b12x as b12x_mod + + calls = [] + monkeypatch.setattr( + b12x_mod, + "_import_b12x_tensor_fp8", + lambda: SimpleNamespace( + prewarm=lambda *args, **kwargs: calls.append((args, kwargs)) ), - model_config=SimpleNamespace(dtype=torch.float16), - vllm_config=SimpleNamespace( - compilation_config=SimpleNamespace(compile_sizes=[32, "dynamic", 64]) + ) + monkeypatch.setattr( + b12x_mod, + "current_stream", + lambda: SimpleNamespace(cuda_stream=object()), + ) + packed_weight = SimpleNamespace( + in_features=128, + padded_in_features=128, + out_features=256, + values=torch.empty(1), + ) + layer = SimpleNamespace(b12x_tensor_fp8_packed_weight=packed_weight) + kernel = object.__new__(B12xTensorFP8ScaledMMLinearKernel) + + unit = kernel.get_b12x_warmup_unit(layer, (1, 8), torch.bfloat16) + unit.compile() + + assert calls[0][0] == (packed_weight, (1, 8)) + assert calls[0][1]["out_dtype"] == torch.bfloat16 + + +def test_b12x_warmup_deduplicates_registered_signatures(monkeypatch) -> None: + import vllm.model_executor.warmup.b12x_warmup as warmup_mod + + calls: list[tuple[str, tuple[int, ...], torch.dtype]] = [] + + class Provider: + def get_b12x_warmup_unit(self, layer, token_counts, output_dtype): + return B12xWarmupUnit( + name="fake", + key=(type(self), layer.shape, output_dtype), + compile=lambda: calls.append((layer.name, token_counts, output_dtype)), + ) + + provider = Provider() + layers = [ + SimpleNamespace(name="first", shape=(128, 256), b12x_warmup_provider=provider), + SimpleNamespace( + name="duplicate", shape=(128, 256), b12x_warmup_provider=provider ), + SimpleNamespace(name="second", shape=(256, 256), b12x_warmup_provider=provider), + SimpleNamespace(), + ] + scans = 0 + + def modules(): + nonlocal scans + scans += 1 + return iter(layers) + + worker = SimpleNamespace( + get_model=lambda: SimpleNamespace(modules=modules), + scheduler_config=SimpleNamespace(max_num_batched_tokens=8), + model_config=SimpleNamespace(dtype=torch.float32), + ) + platform = SimpleNamespace( + is_cuda=lambda: True, + is_device_capability_family=lambda family: family == 120, + ) + synchronized = [] + monkeypatch.setattr(warmup_mod, "current_platform", platform) + monkeypatch.setattr( + warmup_mod.torch.accelerator, + "synchronize", + lambda: synchronized.append(True), ) - b12x_warmup(worker, [8, 16]) + b12x_warmup(worker, [1, 2]) - expected_linear_kwargs = { - "max_tokens": 256, - "cudagraph_capture_sizes": [8, 16], - "output_dtype": torch.float16, - } - assert linear_calls == [(model, expected_linear_kwargs)] * 5 + assert scans == 1 + assert calls == [ + ("first", (1, 2, 8), torch.bfloat16), + ("second", (1, 2, 8), torch.bfloat16), + ] + assert synchronized == [True] diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 3413d4a69fa9..3ba499155b6c 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -172,10 +172,8 @@ AiterPerTokenFp8ScaledMMLinearKernel, AiterPreshuffledPerTokenFp8ScaledMMLinearKernel, ) -from vllm.model_executor.kernels.linear.scaled_mm.b12x_block import ( +from vllm.model_executor.kernels.linear.scaled_mm.b12x import ( B12xFp8BlockScaledMMKernel, -) -from vllm.model_executor.kernels.linear.scaled_mm.b12x_tensor import ( B12xTensorFP8ScaledMMLinearKernel, ) from vllm.model_executor.kernels.linear.scaled_mm.cpu import ( diff --git a/vllm/model_executor/kernels/linear/mxfp4/b12x.py b/vllm/model_executor/kernels/linear/mxfp4/b12x.py index 93e872c0680c..a4649d3c4f8e 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/b12x.py +++ b/vllm/model_executor/kernels/linear/mxfp4/b12x.py @@ -3,9 +3,6 @@ from __future__ import annotations -from collections.abc import Iterable -from typing import Any - import torch from vllm.model_executor.layers.quantization.utils.quant_utils import ( @@ -13,9 +10,7 @@ ) from vllm.model_executor.utils import replace_parameter from vllm.platforms import current_platform -from vllm.utils.b12x import ( - b12x_warmup_token_counts, -) +from vllm.utils.b12x import B12xWarmupUnit from vllm.utils.b12x import ( get_b12x_blockscaled as _import_b12x_blockscaled, ) @@ -51,73 +46,6 @@ def _apply_b12x_mxfp4_linear( return output.view(*output_shape) -def warmup_b12x_mxfp4_linear( - model: torch.nn.Module, - *, - max_tokens: int, - cudagraph_capture_sizes: Iterable[int] = (), - output_dtype: torch.dtype = torch.bfloat16, -) -> int: - if not current_platform.is_cuda(): - return 0 - if not current_platform.is_device_capability_family(120): - return 0 - if output_dtype not in (torch.bfloat16, torch.float16): - output_dtype = torch.bfloat16 - layer_map: dict[tuple[Any, ...], torch.nn.Module] = {} - for layer in model.modules(): - if not getattr(layer, "b12x_mxfp4_linear", False): - continue - weight = layer.weight - weight_scale = layer.weight_scale - n, packed_k = map(int, weight.shape) - signature = ( - weight.device, - n, - packed_k * 2, - weight.dtype, - weight_scale.dtype, - output_dtype, - ) - layer_map.setdefault(signature, layer) - if not layer_map: - return 0 - if _import_b12x_blockscaled() is None: - return 0 - - token_counts = b12x_warmup_token_counts( - max_tokens=max_tokens, - cudagraph_capture_sizes=cudagraph_capture_sizes, - ) - warmed = 0 - last_device: torch.device | None = None - - with torch.inference_mode(): - for signature, layer in layer_map.items(): - weight = layer.weight - weight_scale = layer.weight_scale - k = signature[2] - last_device = weight.device - for tokens in token_counts: - source = torch.zeros( - (tokens, k), - dtype=output_dtype, - device=weight.device, - ) - _apply_b12x_mxfp4_linear( - source, - weight, - weight_scale, - None, - ) - warmed += 1 - - if warmed > 0 and last_device is not None and last_device.type == "cuda": - torch.accelerator.synchronize(last_device) - - return warmed - - class B12xMxFp4LinearKernel(MxFp4LinearKernel): """MXFP4 linear through the native B12X SM120 dense GEMM.""" @@ -151,7 +79,39 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: "weight_scale", intrinsics.swizzle_block_scale(layer.weight_scale.data), ) - layer.b12x_mxfp4_linear = True + layer.b12x_warmup_provider = self + + def get_b12x_warmup_unit( + self, + layer: torch.nn.Module, + token_counts: tuple[int, ...], + output_dtype: torch.dtype, + ) -> B12xWarmupUnit: + weight = layer.weight + weight_scale = layer.weight_scale + n, packed_k = map(int, weight.shape) + k = packed_k * 2 + + def compile() -> None: + for tokens in token_counts: + source = torch.zeros( + (tokens, k), dtype=output_dtype, device=weight.device + ) + _apply_b12x_mxfp4_linear(source, weight, weight_scale, None) + + return B12xWarmupUnit( + name="MXFP4", + key=( + type(self), + weight.device, + n, + k, + weight.dtype, + weight_scale.dtype, + output_dtype, + ), + compile=compile, + ) def apply_weights( self, @@ -167,4 +127,4 @@ def apply_weights( ) -__all__ = ["B12xMxFp4LinearKernel", "warmup_b12x_mxfp4_linear"] +__all__ = ["B12xMxFp4LinearKernel"] diff --git a/vllm/model_executor/kernels/linear/mxfp8/b12x.py b/vllm/model_executor/kernels/linear/mxfp8/b12x.py index 4dd77dd6c00e..48343f3d694b 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/b12x.py +++ b/vllm/model_executor/kernels/linear/mxfp8/b12x.py @@ -3,9 +3,6 @@ from __future__ import annotations -from collections.abc import Iterable -from typing import Any - import torch from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( @@ -15,10 +12,7 @@ ) from vllm.model_executor.utils import replace_parameter from vllm.platforms import current_platform -from vllm.utils.b12x import ( - b12x_warmup_token_counts, - reuse_packed_weight_storage, -) +from vllm.utils.b12x import B12xWarmupUnit, reuse_packed_weight_storage from vllm.utils.b12x import ( get_b12x_mxfp8_linear as _import_b12x_mxfp8, ) @@ -27,21 +21,12 @@ from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig -def _b12x_mxfp8_expected_m(tokens: int) -> int: - return max(1, int(tokens)) - - def _apply_b12x_mxfp8_packed_linear( layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None, ) -> torch.Tensor: - packed_weight = getattr(layer, "b12x_mxfp8_packed_weight", None) - if packed_weight is None: - raise RuntimeError( - "b12x MXFP8 packed weights are missing; " - "process_weights_after_loading did not run for this layer" - ) + packed_weight = layer.b12x_mxfp8_packed_weight input_2d = x.reshape(-1, x.shape[-1]).contiguous() output_shape = [*x.shape[:-1], int(packed_weight.out_features)] @@ -52,77 +37,11 @@ def _apply_b12x_mxfp8_packed_linear( input_2d, packed_weight, bias=bias, - expected_m=_b12x_mxfp8_expected_m(int(input_2d.shape[0])), + expected_m=max(1, int(input_2d.shape[0])), ) return output.view(*output_shape) -def warmup_b12x_mxfp8_linear( - model: torch.nn.Module, - *, - max_tokens: int, - cudagraph_capture_sizes: Iterable[int] = (), - output_dtype: torch.dtype = torch.bfloat16, -) -> int: - if not current_platform.is_cuda(): - return 0 - if not current_platform.is_device_capability_family(120): - return 0 - if output_dtype not in (torch.bfloat16, torch.float16): - output_dtype = torch.bfloat16 - - layer_map: dict[tuple[Any, ...], Any] = {} - for layer in model.modules(): - packed_weight = getattr(layer, "b12x_mxfp8_packed_weight", None) - if packed_weight is None: - continue - device = torch.device(packed_weight.weight.values.device) - signature = ( - device, - int(packed_weight.in_features), - int(packed_weight.padded_in_features), - int(packed_weight.out_features), - output_dtype, - ) - layer_map.setdefault(signature, packed_weight) - if not layer_map: - return 0 - - mxfp8 = _import_b12x_mxfp8() - if mxfp8 is None: - return 0 - - token_counts = b12x_warmup_token_counts( - max_tokens=max_tokens, - cudagraph_capture_sizes=cudagraph_capture_sizes, - ) - warmed = 0 - last_device: torch.device | None = None - - with torch.inference_mode(): - for signature, packed_weight in layer_map.items(): - device = signature[0] - last_device = device - for tokens in token_counts: - source = torch.zeros( - (tokens, int(packed_weight.in_features)), - dtype=output_dtype, - device=device, - ) - mxfp8.mm( - source, - packed_weight, - expected_m=_b12x_mxfp8_expected_m(tokens), - stream=current_stream().cuda_stream, - ) - warmed += 1 - - if warmed > 0 and last_device is not None and last_device.type == "cuda": - torch.accelerator.synchronize(last_device) - - return warmed - - class B12xMxfp8LinearKernel(Mxfp8LinearKernel): """ModelOpt MXFP8 linear through the native b12x SM120 dense GEMM path.""" @@ -183,6 +102,45 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: ) replace_parameter(layer, "weight", weight.new_empty((0,))) replace_parameter(layer, "weight_scale", weight_scale.new_empty((0,))) + layer.b12x_warmup_provider = self + + def get_b12x_warmup_unit( + self, + layer: torch.nn.Module, + token_counts: tuple[int, ...], + output_dtype: torch.dtype, + ) -> B12xWarmupUnit: + packed_weight = layer.b12x_mxfp8_packed_weight + device = torch.device(packed_weight.weight.values.device) + + def compile() -> None: + mxfp8 = _import_b12x_mxfp8() + assert mxfp8 is not None + for tokens in token_counts: + source = torch.zeros( + (tokens, int(packed_weight.in_features)), + dtype=output_dtype, + device=device, + ) + mxfp8.mm( + source, + packed_weight, + expected_m=max(1, int(tokens)), + stream=current_stream().cuda_stream, + ) + + return B12xWarmupUnit( + name="MXFP8", + key=( + type(self), + device, + int(packed_weight.in_features), + int(packed_weight.padded_in_features), + int(packed_weight.out_features), + output_dtype, + ), + compile=compile, + ) def apply_weights( self, diff --git a/vllm/model_executor/kernels/linear/nvfp4/b12x.py b/vllm/model_executor/kernels/linear/nvfp4/b12x.py index a9e5a69283fc..4ff531dd48be 100644 --- a/vllm/model_executor/kernels/linear/nvfp4/b12x.py +++ b/vllm/model_executor/kernels/linear/nvfp4/b12x.py @@ -3,17 +3,12 @@ from __future__ import annotations -from collections.abc import Iterable -from typing import Any - import torch from vllm._custom_ops import scaled_fp4_quant from vllm.model_executor.utils import replace_parameter from vllm.platforms import current_platform -from vllm.utils.b12x import ( - b12x_warmup_token_counts, -) +from vllm.utils.b12x import B12xWarmupUnit from vllm.utils.b12x import ( get_b12x_blockscaled as _import_b12x_blockscaled, ) @@ -54,75 +49,6 @@ def _apply_b12x_nvfp4_linear( return output.view(*output_shape) -def warmup_b12x_nvfp4_linear( - model: torch.nn.Module, - *, - max_tokens: int, - cudagraph_capture_sizes: Iterable[int] = (), - output_dtype: torch.dtype = torch.bfloat16, -) -> int: - if not current_platform.is_cuda(): - return 0 - if not current_platform.is_device_capability_family(120): - return 0 - if output_dtype not in (torch.bfloat16, torch.float16): - output_dtype = torch.bfloat16 - layer_map: dict[tuple[Any, ...], torch.nn.Module] = {} - for layer in model.modules(): - if not getattr(layer, "b12x_nvfp4_linear", False): - continue - weight = layer.weight - weight_scale = layer.weight_scale - n, packed_k = map(int, weight.shape) - signature = ( - weight.device, - n, - packed_k * 2, - weight.dtype, - weight_scale.dtype, - output_dtype, - ) - layer_map.setdefault(signature, layer) - if not layer_map: - return 0 - if _import_b12x_blockscaled() is None: - return 0 - - token_counts = b12x_warmup_token_counts( - max_tokens=max_tokens, - cudagraph_capture_sizes=cudagraph_capture_sizes, - ) - warmed = 0 - last_device: torch.device | None = None - - with torch.inference_mode(): - for signature, layer in layer_map.items(): - weight = layer.weight - weight_scale = layer.weight_scale - k = signature[2] - last_device = weight.device - for tokens in token_counts: - source = torch.zeros( - (tokens, k), - dtype=output_dtype, - device=weight.device, - ) - _apply_b12x_nvfp4_linear( - source, - weight, - weight_scale, - layer.input_global_scale_inv, - layer.alpha, - None, - ) - warmed += 1 - - if warmed > 0 and last_device is not None and last_device.type == "cuda": - torch.accelerator.synchronize(last_device) - - return warmed - - class B12xNvFp4LinearKernel(NvFp4LinearKernel): """ModelOpt NVFP4 linear through the native B12X SM120 dense GEMM.""" @@ -155,7 +81,46 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: "weight_scale", intrinsics.swizzle_block_scale(layer.weight_scale.data), ) - layer.b12x_nvfp4_linear = True + layer.b12x_warmup_provider = self + + def get_b12x_warmup_unit( + self, + layer: torch.nn.Module, + token_counts: tuple[int, ...], + output_dtype: torch.dtype, + ) -> B12xWarmupUnit: + weight = layer.weight + weight_scale = layer.weight_scale + n, packed_k = map(int, weight.shape) + k = packed_k * 2 + + def compile() -> None: + for tokens in token_counts: + source = torch.zeros( + (tokens, k), dtype=output_dtype, device=weight.device + ) + _apply_b12x_nvfp4_linear( + source, + weight, + weight_scale, + layer.input_global_scale_inv, + layer.alpha, + None, + ) + + return B12xWarmupUnit( + name="NVFP4", + key=( + type(self), + weight.device, + n, + k, + weight.dtype, + weight_scale.dtype, + output_dtype, + ), + compile=compile, + ) def apply_weights( self, @@ -173,4 +138,4 @@ def apply_weights( ) -__all__ = ["B12xNvFp4LinearKernel", "warmup_b12x_nvfp4_linear"] +__all__ = ["B12xNvFp4LinearKernel"] diff --git a/vllm/model_executor/kernels/linear/scaled_mm/b12x.py b/vllm/model_executor/kernels/linear/scaled_mm/b12x.py new file mode 100644 index 000000000000..62e839c7291c --- /dev/null +++ b/vllm/model_executor/kernels/linear/scaled_mm/b12x.py @@ -0,0 +1,339 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import torch + +from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + _upcast_e8m0_to_fp32, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape +from vllm.model_executor.utils import replace_parameter +from vllm.platforms import current_platform +from vllm.utils.b12x import B12xWarmupUnit, reuse_packed_weight_storage +from vllm.utils.b12x import ( + get_b12x_blockscaled as _import_b12x_blockscaled, +) +from vllm.utils.b12x import ( + get_b12x_tensor_fp8_linear as _import_b12x_tensor_fp8, +) +from vllm.utils.torch_utils import current_stream + +from .BlockScaledMMLinearKernel import ( + Fp8BlockScaledMMLinearKernel, + FP8ScaledMMLinearLayerConfig, +) +from .ScaledMMLinearKernel import FP8ScaledMMLinearKernel + + +def _run_b12x_fp8_block_scaled_mm( + a: torch.Tensor, + weight: torch.Tensor, + a_scale: torch.Tensor, + weight_scale: torch.Tensor, + out_dtype: torch.dtype, +) -> torch.Tensor: + blockscaled = _import_b12x_blockscaled() + assert blockscaled is not None + + return blockscaled.mm_block_fp8( + a, + a_scale, + weight, + weight_scale, + out_dtype=out_dtype, + ) + + +class B12xFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): + """K128 block-FP8 linear through the native B12X SM120 dense GEMM.""" + + @classmethod + def is_supported( + cls, + compute_capability: int | None = None, + ) -> tuple[bool, str | None]: + del compute_capability + if not current_platform.is_cuda(): + return False, "B12X FP8 kernels are only available on CUDA" + if not current_platform.is_device_capability_family(120): + return False, "B12X FP8 kernels require a Blackwell 12x device" + blockscaled = _import_b12x_blockscaled() + if blockscaled is None: + return False, "Install the B12X backend with `pip install vllm[b12x]`" + if not blockscaled.is_supported(): + return False, "B12X regular block-FP8 GEMM is not supported" + return True, None + + @classmethod + def can_implement( + cls, + config: FP8ScaledMMLinearLayerConfig, + ) -> tuple[bool, str | None]: + can_implement_base, reason = super().can_implement(config) + if not can_implement_base: + return can_implement_base, reason + + if config.input_dtype not in (torch.bfloat16, torch.float16): + return False, "Supports only bf16/fp16 input dtype" + if config.input_dtype != config.out_dtype: + return False, "Input and output dtype must match" + + act_group_shape = config.activation_quant_key.scale.group_shape + if act_group_shape != GroupShape(1, 128): + return ( + False, + "Supports only dynamic per-token group activation quantization " + "with group_shape=(1,128)", + ) + weight_group_shape = config.weight_quant_key.scale.group_shape + if weight_group_shape != GroupShape(128, 128): + return False, "Supports only 128x128 block-scaled FP8 weights" + + out_features, in_features = config.weight_shape + if in_features <= 0 or in_features % 128 != 0: + return False, "Input features must be a positive multiple of 128" + if out_features <= 0 or out_features % 128 != 0: + return False, "Output features must be a positive multiple of 128" + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + super().process_weights_after_loading(layer) + params = self._get_layer_params(layer) + if params.weight_scale_inv is not None: + weight_scale = params.weight_scale_inv + scale_attr = params.WEIGHT_SCALE_INV + else: + weight_scale = params.weight_scale + scale_attr = params.WEIGHT_SCALE + if weight_scale is not None and weight_scale.dtype in ( + torch.float8_e8m0fnu, + torch.uint8, + ): + # TODO: Remove once B12X supports 128x128 UE8M0 block scales. + replace_parameter( + layer, + scale_attr, + _upcast_e8m0_to_fp32(weight_scale).contiguous(), + ) + layer.b12x_warmup_provider = self + + def get_b12x_warmup_unit( + self, + layer: torch.nn.Module, + token_counts: tuple[int, ...], + output_dtype: torch.dtype, + ) -> B12xWarmupUnit: + weight = layer.weight + weight_scale = getattr(layer, "weight_scale_inv", None) + if weight_scale is None: + weight_scale = layer.weight_scale + n, k = map(int, weight.shape) + + def compile() -> None: + for tokens in token_counts: + a = torch.empty((tokens, k), dtype=weight.dtype, device=weight.device) + a_scale = torch.empty( + (tokens, k // 128), + dtype=torch.float32, + device=weight.device, + ) + _run_b12x_fp8_block_scaled_mm( + a, weight, a_scale, weight_scale, output_dtype + ) + + return B12xWarmupUnit( + name="block-FP8", + key=( + type(self), + weight.device, + n, + k, + weight.dtype, + weight_scale.dtype, + output_dtype, + ), + compile=compile, + ) + + def apply_block_scaled_mm( + self, + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + ) -> torch.Tensor: + return _run_b12x_fp8_block_scaled_mm( + A, + B, + As, + Bs, + self.config.out_dtype, + ) + + +def _apply_b12x_tensor_fp8_packed_linear( + layer: torch.nn.Module, + x_q: torch.Tensor, + bias: torch.Tensor | None, + out_dtype: torch.dtype, +) -> torch.Tensor: + packed_weight = layer.b12x_tensor_fp8_packed_weight + + tensor_fp8 = _import_b12x_tensor_fp8() + assert tensor_fp8 is not None + + input_2d = x_q.reshape(-1, x_q.shape[-1]).contiguous() + output_shape = [*x_q.shape[:-1], int(packed_weight.out_features)] + output = tensor_fp8.mm( + input_2d, + packed_weight, + bias=bias, + out_dtype=out_dtype, + expected_m=max(1, int(input_2d.shape[0])), + ) + return output.view(*output_shape) + + +class B12xTensorFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): + """Static per-tensor FP8 linear through the B12X SM12x dense GEMM.""" + + @classmethod + def is_supported( + cls, + compute_capability: int | None = None, + ) -> tuple[bool, str | None]: + del compute_capability + if not current_platform.is_cuda(): + return False, "b12x tensor FP8 kernels are only available on CUDA" + if not current_platform.is_device_capability_family(120): + return False, "b12x tensor FP8 kernels require a Blackwell 12x device" + tensor_fp8 = _import_b12x_tensor_fp8() + if tensor_fp8 is None: + return False, "Install the B12X backend with `pip install vllm[b12x]`" + if not tensor_fp8.is_supported(): + return False, "b12x.gemm.tensor_fp8_linear is not supported" + return True, None + + @classmethod + def can_implement( + cls, + config: FP8ScaledMMLinearLayerConfig, + ) -> tuple[bool, str | None]: + activation_scale = config.activation_quant_key.scale + weight_scale = config.weight_quant_key.scale + if ( + not activation_scale.static + or not activation_scale.group_shape.is_per_tensor() + ): + return False, "requires static per-tensor activation scales" + if not weight_scale.static or not weight_scale.group_shape.is_per_tensor(): + return False, "requires static per-tensor weight scales" + if config.input_dtype not in (torch.bfloat16, torch.float16): + return False, "supports only bf16/fp16 input dtype" + if config.out_dtype not in (torch.bfloat16, torch.float16): + return False, "supports only bf16/fp16 output dtype" + out_features, in_features = config.weight_shape + if out_features <= 0 or in_features <= 0 or in_features % 32 != 0: + return False, "weight dimensions must be positive with K divisible by 32" + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + weight, weight_scale, input_scale, _ = self._get_layer_params(layer) + assert weight.dtype == torch.float8_e4m3fn + assert input_scale is not None + assert weight_scale.numel() == input_scale.numel() == 1 + + out_features, in_features = map(int, self.config.weight_shape) + assert tuple(weight.shape) == (in_features, out_features) + + tensor_fp8 = _import_b12x_tensor_fp8() + assert tensor_fp8 is not None + output_scale = ( + input_scale.detach().to(torch.float32).reshape(1) + * weight_scale.detach().to(torch.float32).reshape(1) + ).contiguous() + packed_weight = tensor_fp8.pack_weight( + weight.detach().T.contiguous(), + output_scale, + ) + layer.b12x_tensor_fp8_packed_weight = reuse_packed_weight_storage( + getattr(layer, "b12x_tensor_fp8_packed_weight", None), + packed_weight, + ) + weight_name, weight_scale_name, _, _ = self.layer_param_names + replace_parameter(layer, weight_name, weight.new_empty((0,))) + replace_parameter(layer, weight_scale_name, weight_scale.new_empty((0,))) + layer.b12x_warmup_provider = self + + def get_b12x_warmup_unit( + self, + layer: torch.nn.Module, + token_counts: tuple[int, ...], + output_dtype: torch.dtype, + ) -> B12xWarmupUnit: + packed_weight = layer.b12x_tensor_fp8_packed_weight + device = torch.device(packed_weight.values.device) + + def compile() -> None: + tensor_fp8 = _import_b12x_tensor_fp8() + assert tensor_fp8 is not None + tensor_fp8.prewarm( + packed_weight, + token_counts, + out_dtype=output_dtype, + stream=current_stream().cuda_stream, + ) + + return B12xWarmupUnit( + name="tensor FP8", + key=( + type(self), + device, + int(packed_weight.in_features), + int(packed_weight.padded_in_features), + int(packed_weight.out_features), + output_dtype, + ), + compile=compile, + ) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + assert isinstance(x, torch.Tensor) + _, _, input_scale, input_scale_ub = self._get_layer_params(layer) + input_2d = x.reshape(-1, x.shape[-1]) + x_q, _ = self.quant_fp8(input_2d, input_scale, input_scale_ub) + out_dtype = self.config.out_dtype + output = _apply_b12x_tensor_fp8_packed_linear( + layer, + x_q, + bias, + out_dtype, + ) + return output.view(*x.shape[:-1], output.shape[-1]) + + def apply_scaled_mm( + self, + *, + A: torch.Tensor, + B: torch.Tensor, + out_dtype: torch.dtype, + As: torch.Tensor, + Bs: torch.Tensor, + bias: torch.Tensor | None, + output_shape: list, + ) -> torch.Tensor: + del A, B, out_dtype, As, Bs, bias, output_shape + raise NotImplementedError("b12x tensor FP8 linear overrides apply_weights") + + +__all__ = [ + "B12xFp8BlockScaledMMKernel", + "B12xTensorFP8ScaledMMLinearKernel", +] diff --git a/vllm/model_executor/kernels/linear/scaled_mm/b12x_block.py b/vllm/model_executor/kernels/linear/scaled_mm/b12x_block.py deleted file mode 100644 index c2d12549fbde..000000000000 --- a/vllm/model_executor/kernels/linear/scaled_mm/b12x_block.py +++ /dev/null @@ -1,214 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from __future__ import annotations - -from collections.abc import Iterable -from typing import Any - -import torch - -from vllm.model_executor.layers.quantization.utils.fp8_utils import ( - _upcast_e8m0_to_fp32, -) -from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape -from vllm.model_executor.utils import replace_parameter -from vllm.platforms import current_platform -from vllm.utils.b12x import ( - b12x_warmup_token_counts, -) -from vllm.utils.b12x import ( - get_b12x_blockscaled as _import_b12x_blockscaled, -) - -from .BlockScaledMMLinearKernel import ( - Fp8BlockScaledMMLinearKernel, - FP8ScaledMMLinearLayerConfig, -) - - -def _run_b12x_fp8_block_scaled_mm( - a: torch.Tensor, - weight: torch.Tensor, - a_scale: torch.Tensor, - weight_scale: torch.Tensor, - out_dtype: torch.dtype, -) -> torch.Tensor: - blockscaled = _import_b12x_blockscaled() - assert blockscaled is not None - - return blockscaled.mm_block_fp8( - a, - a_scale, - weight, - weight_scale, - out_dtype=out_dtype, - ) - - -def warmup_b12x_block_fp8_linear( - model: torch.nn.Module, - *, - max_tokens: int, - cudagraph_capture_sizes: Iterable[int] = (), - output_dtype: torch.dtype = torch.bfloat16, -) -> int: - if not current_platform.is_cuda(): - return 0 - if not current_platform.is_device_capability_family(120): - return 0 - if output_dtype not in (torch.bfloat16, torch.float16): - output_dtype = torch.bfloat16 - - layer_map: dict[tuple[Any, ...], torch.nn.Module] = {} - for layer in model.modules(): - if not getattr(layer, "b12x_block_fp8_linear", False): - continue - weight = layer.weight - weight_scale = getattr(layer, "weight_scale_inv", None) - if weight_scale is None: - weight_scale = layer.weight_scale - n, k = map(int, weight.shape) - signature = ( - weight.device, - n, - k, - weight.dtype, - weight_scale.dtype, - output_dtype, - ) - layer_map.setdefault(signature, layer) - if not layer_map: - return 0 - - blockscaled = _import_b12x_blockscaled() - if blockscaled is None: - return 0 - token_counts = b12x_warmup_token_counts( - max_tokens=max_tokens, - cudagraph_capture_sizes=cudagraph_capture_sizes, - ) - warmed = 0 - - with torch.inference_mode(): - for signature, layer in layer_map.items(): - weight = layer.weight - weight_scale = getattr(layer, "weight_scale_inv", None) - if weight_scale is None: - weight_scale = layer.weight_scale - k = signature[2] - for tokens in token_counts: - a = torch.empty( - (tokens, k), - dtype=weight.dtype, - device=weight.device, - ) - a_scale = torch.empty( - (tokens, k // 128), - dtype=torch.float32, - device=weight.device, - ) - _run_b12x_fp8_block_scaled_mm( - a, - weight, - a_scale, - weight_scale, - output_dtype, - ) - warmed += 1 - return warmed - - -class B12xFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): - """K128 block-FP8 linear through the native B12X SM120 dense GEMM.""" - - @classmethod - def is_supported( - cls, - compute_capability: int | None = None, - ) -> tuple[bool, str | None]: - del compute_capability - if not current_platform.is_cuda(): - return False, "B12X FP8 kernels are only available on CUDA" - if not current_platform.is_device_capability_family(120): - return False, "B12X FP8 kernels require a Blackwell 12x device" - blockscaled = _import_b12x_blockscaled() - if blockscaled is None: - return False, "Install the B12X backend with `pip install vllm[b12x]`" - if not blockscaled.is_supported(): - return False, "B12X regular block-FP8 GEMM is not supported" - return True, None - - @classmethod - def can_implement( - cls, - config: FP8ScaledMMLinearLayerConfig, - ) -> tuple[bool, str | None]: - can_implement_base, reason = super().can_implement(config) - if not can_implement_base: - return can_implement_base, reason - - if config.input_dtype not in (torch.bfloat16, torch.float16): - return False, "Supports only bf16/fp16 input dtype" - if config.input_dtype != config.out_dtype: - return False, "Input and output dtype must match" - - act_group_shape = config.activation_quant_key.scale.group_shape - if act_group_shape != GroupShape(1, 128): - return ( - False, - "Supports only dynamic per-token group activation quantization " - "with group_shape=(1,128)", - ) - weight_group_shape = config.weight_quant_key.scale.group_shape - if weight_group_shape != GroupShape(128, 128): - return False, "Supports only 128x128 block-scaled FP8 weights" - - out_features, in_features = config.weight_shape - if in_features <= 0 or in_features % 128 != 0: - return False, "Input features must be a positive multiple of 128" - if out_features <= 0 or out_features % 128 != 0: - return False, "Output features must be a positive multiple of 128" - return True, None - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - super().process_weights_after_loading(layer) - params = self._get_layer_params(layer) - if params.weight_scale_inv is not None: - weight_scale = params.weight_scale_inv - scale_attr = params.WEIGHT_SCALE_INV - else: - weight_scale = params.weight_scale - scale_attr = params.WEIGHT_SCALE - if weight_scale is not None and weight_scale.dtype in ( - torch.float8_e8m0fnu, - torch.uint8, - ): - # TODO: Remove once B12X supports 128x128 UE8M0 block scales. - replace_parameter( - layer, - scale_attr, - _upcast_e8m0_to_fp32(weight_scale).contiguous(), - ) - layer.b12x_block_fp8_linear = True - - def apply_block_scaled_mm( - self, - A: torch.Tensor, - B: torch.Tensor, - As: torch.Tensor, - Bs: torch.Tensor, - ) -> torch.Tensor: - return _run_b12x_fp8_block_scaled_mm( - A, - B, - As, - Bs, - self.config.out_dtype, - ) - - -__all__ = [ - "B12xFp8BlockScaledMMKernel", - "warmup_b12x_block_fp8_linear", -] diff --git a/vllm/model_executor/kernels/linear/scaled_mm/b12x_tensor.py b/vllm/model_executor/kernels/linear/scaled_mm/b12x_tensor.py deleted file mode 100644 index 46b27a7d4b4e..000000000000 --- a/vllm/model_executor/kernels/linear/scaled_mm/b12x_tensor.py +++ /dev/null @@ -1,238 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from __future__ import annotations - -from collections.abc import Iterable -from typing import Any - -import torch - -from vllm.model_executor.utils import replace_parameter -from vllm.platforms import current_platform -from vllm.utils.b12x import ( - b12x_warmup_token_counts, - reuse_packed_weight_storage, -) -from vllm.utils.b12x import ( - get_b12x_tensor_fp8_linear as _import_b12x_tensor_fp8, -) -from vllm.utils.torch_utils import current_stream - -from .ScaledMMLinearKernel import ( - FP8ScaledMMLinearKernel, - FP8ScaledMMLinearLayerConfig, -) - - -def _apply_b12x_tensor_fp8_packed_linear( - layer: torch.nn.Module, - x_q: torch.Tensor, - bias: torch.Tensor | None, - out_dtype: torch.dtype, -) -> torch.Tensor: - packed_weight = getattr(layer, "b12x_tensor_fp8_packed_weight", None) - if packed_weight is None: - raise RuntimeError( - "b12x tensor FP8 packed weights are missing; " - "process_weights_after_loading did not run for this layer" - ) - - tensor_fp8 = _import_b12x_tensor_fp8() - assert tensor_fp8 is not None - - input_2d = x_q.reshape(-1, x_q.shape[-1]).contiguous() - output_shape = [*x_q.shape[:-1], int(packed_weight.out_features)] - output = tensor_fp8.mm( - input_2d, - packed_weight, - bias=bias, - out_dtype=out_dtype, - expected_m=max(1, int(input_2d.shape[0])), - ) - return output.view(*output_shape) - - -def warmup_b12x_tensor_fp8_linear( - model: torch.nn.Module, - *, - max_tokens: int, - cudagraph_capture_sizes: Iterable[int] = (), - output_dtype: torch.dtype = torch.bfloat16, -) -> int: - if not current_platform.is_cuda(): - return 0 - if not current_platform.is_device_capability_family(120): - return 0 - - if output_dtype not in (torch.bfloat16, torch.float16): - output_dtype = torch.bfloat16 - - layer_map: dict[tuple[Any, ...], Any] = {} - for layer in model.modules(): - packed_weight = getattr( - layer, - "b12x_tensor_fp8_packed_weight", - None, - ) - if packed_weight is None: - continue - device = torch.device(packed_weight.values.device) - signature = ( - device, - int(packed_weight.in_features), - int(packed_weight.padded_in_features), - int(packed_weight.out_features), - output_dtype, - ) - layer_map.setdefault(signature, packed_weight) - if not layer_map: - return 0 - - tensor_fp8 = _import_b12x_tensor_fp8() - if tensor_fp8 is None: - return 0 - - token_counts = b12x_warmup_token_counts( - max_tokens=max_tokens, - cudagraph_capture_sizes=cudagraph_capture_sizes, - ) - warmed = 0 - last_device: torch.device | None = None - - with torch.inference_mode(): - for signature, packed_weight in layer_map.items(): - last_device = signature[0] - warmed += int( - tensor_fp8.prewarm( - packed_weight, - token_counts, - out_dtype=output_dtype, - stream=current_stream().cuda_stream, - ) - ) - - if warmed > 0 and last_device is not None and last_device.type == "cuda": - torch.accelerator.synchronize(last_device) - - return warmed - - -class B12xTensorFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): - """Static per-tensor FP8 linear through the B12X SM12x dense GEMM.""" - - @classmethod - def is_supported( - cls, - compute_capability: int | None = None, - ) -> tuple[bool, str | None]: - del compute_capability - if not current_platform.is_cuda(): - return False, "b12x tensor FP8 kernels are only available on CUDA" - if not current_platform.is_device_capability_family(120): - return False, "b12x tensor FP8 kernels require a Blackwell 12x device" - tensor_fp8 = _import_b12x_tensor_fp8() - if tensor_fp8 is None: - return False, "Install the B12X backend with `pip install vllm[b12x]`" - if not tensor_fp8.is_supported(): - return False, "b12x.gemm.tensor_fp8_linear is not supported" - return True, None - - @classmethod - def can_implement( - cls, - config: FP8ScaledMMLinearLayerConfig, - ) -> tuple[bool, str | None]: - activation_scale = config.activation_quant_key.scale - weight_scale = config.weight_quant_key.scale - if ( - not activation_scale.static - or not activation_scale.group_shape.is_per_tensor() - ): - return False, "requires static per-tensor activation scales" - if not weight_scale.static or not weight_scale.group_shape.is_per_tensor(): - return False, "requires static per-tensor weight scales" - if config.input_dtype not in (torch.bfloat16, torch.float16): - return False, "supports only bf16/fp16 input dtype" - if config.out_dtype not in (torch.bfloat16, torch.float16): - return False, "supports only bf16/fp16 output dtype" - out_features, in_features = config.weight_shape - if out_features <= 0 or in_features <= 0 or in_features % 32 != 0: - return False, "weight dimensions must be positive with K divisible by 32" - return True, None - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - weight, weight_scale, input_scale, _ = self._get_layer_params(layer) - if weight.dtype != torch.float8_e4m3fn: - raise ValueError( - f"b12x tensor FP8 requires float8_e4m3fn weight, got {weight.dtype}" - ) - if weight_scale.numel() != 1 or input_scale is None or input_scale.numel() != 1: - raise ValueError( - "b12x tensor FP8 requires scalar weight and activation scales" - ) - - out_features, in_features = map(int, self.config.weight_shape) - if tuple(weight.shape) != (in_features, out_features): - raise ValueError( - "b12x tensor FP8 expects the processed weight in [K,N] layout, " - f"got {tuple(weight.shape)} for N={out_features}, K={in_features}" - ) - - tensor_fp8 = _import_b12x_tensor_fp8() - assert tensor_fp8 is not None - output_scale = ( - input_scale.detach().to(torch.float32).reshape(1) - * weight_scale.detach().to(torch.float32).reshape(1) - ).contiguous() - packed_weight = tensor_fp8.pack_weight( - weight.detach().T.contiguous(), - output_scale, - ) - layer.b12x_tensor_fp8_packed_weight = reuse_packed_weight_storage( - getattr(layer, "b12x_tensor_fp8_packed_weight", None), - packed_weight, - ) - weight_name, weight_scale_name, _, _ = self.layer_param_names - replace_parameter(layer, weight_name, weight.new_empty((0,))) - replace_parameter(layer, weight_scale_name, weight_scale.new_empty((0,))) - - def apply_weights( - self, - layer: torch.nn.Module, - x: torch.Tensor, - bias: torch.Tensor | None = None, - ) -> torch.Tensor: - if not isinstance(x, torch.Tensor): - raise TypeError("b12x tensor FP8 linear requires a Tensor input") - _, _, input_scale, input_scale_ub = self._get_layer_params(layer) - input_2d = x.reshape(-1, x.shape[-1]) - x_q, _ = self.quant_fp8(input_2d, input_scale, input_scale_ub) - out_dtype = self.config.out_dtype - output = _apply_b12x_tensor_fp8_packed_linear( - layer, - x_q, - bias, - out_dtype, - ) - return output.view(*x.shape[:-1], output.shape[-1]) - - def apply_scaled_mm( - self, - *, - A: torch.Tensor, - B: torch.Tensor, - out_dtype: torch.dtype, - As: torch.Tensor, - Bs: torch.Tensor, - bias: torch.Tensor | None, - output_shape: list, - ) -> torch.Tensor: - del A, B, out_dtype, As, Bs, bias, output_shape - raise NotImplementedError("b12x tensor FP8 linear overrides apply_weights") - - -__all__ = [ - "B12xTensorFP8ScaledMMLinearKernel", - "warmup_b12x_tensor_fp8_linear", -] diff --git a/vllm/model_executor/warmup/b12x_warmup.py b/vllm/model_executor/warmup/b12x_warmup.py index daaa0786286f..5e0e4a0fa498 100644 --- a/vllm/model_executor/warmup/b12x_warmup.py +++ b/vllm/model_executor/warmup/b12x_warmup.py @@ -2,26 +2,15 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Warm B12X JIT kernels used by a loaded model.""" +from collections import Counter +from collections.abc import Iterable from typing import TYPE_CHECKING import torch from vllm.logger import init_logger -from vllm.model_executor.kernels.linear.mxfp4.b12x import ( - warmup_b12x_mxfp4_linear, -) -from vllm.model_executor.kernels.linear.mxfp8.b12x import ( - warmup_b12x_mxfp8_linear, -) -from vllm.model_executor.kernels.linear.nvfp4.b12x import ( - warmup_b12x_nvfp4_linear, -) -from vllm.model_executor.kernels.linear.scaled_mm.b12x_block import ( - warmup_b12x_block_fp8_linear, -) -from vllm.model_executor.kernels.linear.scaled_mm.b12x_tensor import ( - warmup_b12x_tensor_fp8_linear, -) +from vllm.platforms import current_platform +from vllm.utils.b12x import B12xWarmupUnit, b12x_warmup_token_counts if TYPE_CHECKING: from vllm.v1.worker.gpu_worker import Worker @@ -29,32 +18,61 @@ logger = init_logger(__name__) +def _collect_warmup_units( + model: torch.nn.Module, + token_counts: tuple[int, ...], + output_dtype: torch.dtype, +) -> Iterable[B12xWarmupUnit]: + units: dict[object, B12xWarmupUnit] = {} + for layer in model.modules(): + provider = getattr(layer, "b12x_warmup_provider", None) + get_unit = getattr(provider, "get_b12x_warmup_unit", None) + if not callable(get_unit): + continue + unit = get_unit(layer, token_counts, output_dtype) + assert isinstance(unit, B12xWarmupUnit) + units.setdefault(unit.key, unit) + return units.values() + + +def _compile_warmup_units( + units: Iterable[B12xWarmupUnit], +) -> Counter[str]: + warmed: Counter[str] = Counter() + with torch.inference_mode(): + for unit in units: + unit.compile() + warmed[unit.name] += 1 + if warmed: + torch.accelerator.synchronize() + return warmed + + def b12x_warmup(worker: "Worker", cudagraph_capture_sizes: list[int]) -> None: - model = worker.get_model() - max_tokens = worker.scheduler_config.max_num_batched_tokens + if not current_platform.is_cuda(): + return + if not current_platform.is_device_capability_family(120): + return + output_dtype = getattr( getattr(worker, "model_config", None), "dtype", torch.bfloat16, ) - - warmup_kwargs = { - "max_tokens": max_tokens, - "cudagraph_capture_sizes": cudagraph_capture_sizes, - "output_dtype": output_dtype, - } - providers = ( - ("block-FP8", warmup_b12x_block_fp8_linear), - ("MXFP8", warmup_b12x_mxfp8_linear), - ("tensor FP8", warmup_b12x_tensor_fp8_linear), - ("MXFP4", warmup_b12x_mxfp4_linear), - ("NVFP4", warmup_b12x_nvfp4_linear), + if output_dtype not in (torch.bfloat16, torch.float16): + output_dtype = torch.bfloat16 + token_counts = b12x_warmup_token_counts( + max_tokens=worker.scheduler_config.max_num_batched_tokens, + cudagraph_capture_sizes=cudagraph_capture_sizes, ) - for name, warmup in providers: - warmed = warmup(model, **warmup_kwargs) - if warmed: - logger.info_once( - "Warmed up %d B12X %s linear GEMM signatures.", - warmed, - name, - ) + units = _collect_warmup_units( + worker.get_model(), + token_counts, + output_dtype, + ) + for name, count in _compile_warmup_units(units).items(): + logger.info_once( + "Warmed up %d B12X %s linear GEMM signatures.", + count, + name, + ) diff --git a/vllm/utils/b12x.py b/vllm/utils/b12x.py index 55afa35f093f..8e2d7e938b6d 100644 --- a/vllm/utils/b12x.py +++ b/vllm/utils/b12x.py @@ -5,14 +5,21 @@ import functools import importlib import importlib.util -from collections.abc import Iterable -from dataclasses import fields, is_dataclass +from collections.abc import Callable, Hashable, Iterable +from dataclasses import dataclass, fields, is_dataclass from types import ModuleType from typing import Any import torch +@dataclass(frozen=True) +class B12xWarmupUnit: + name: str + key: Hashable + compile: Callable[[], None] + + @functools.cache def has_b12x() -> bool: """Return whether the B12X package is installed.""" From 7075ddac28c25d4fd2b84bc2a9a6c5ffde0345c8 Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Mon, 17 Aug 2026 12:48:30 -0400 Subject: [PATCH 049/839] Support DSpark configs with `architectures=DSparkDraftModel` + `model_type=qwen3` (#52197) Signed-off-by: mgoin --- vllm/config/speculative.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 2b93113b7ed3..a85b895d8647 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -953,6 +953,10 @@ def __post_init__(self): "dspark" in self.draft_model_config.model.lower() or "Qwen3DSparkModel" in self.draft_model_config.architectures or "Gemma4DSparkModel" in self.draft_model_config.architectures + or ( + "DSparkDraftModel" in self.draft_model_config.architectures + and self.draft_model_config.hf_config.model_type == "qwen3" + ) ): self.method = "dspark" elif self.draft_model_config.hf_config.model_type == "medusa": @@ -1010,7 +1014,16 @@ def __post_init__(self): self.draft_model_config.hf_config = eagle_config self.update_arch_() - if self.method == "dspark" and ( + if ( + self.method == "dspark" + and "DSparkDraftModel" in self.draft_model_config.architectures + and self.draft_model_config.hf_config.model_type == "qwen3" + ): + self.draft_model_config.hf_config.architectures = [ + "Qwen3DSparkModel" + ] + self.update_arch_() + elif self.method == "dspark" and ( "Qwen3DSparkModel" not in self.draft_model_config.architectures and "Gemma4DSparkModel" not in self.draft_model_config.architectures and "K3DSparkModel" not in self.draft_model_config.architectures From 49905ad94dfc58ba94cca8b173dba31b820d6bc6 Mon Sep 17 00:00:00 2001 From: "Roberto L. Castro" <38211239+LopezCastroRoberto@users.noreply.github.com> Date: Mon, 17 Aug 2026 19:11:14 +0200 Subject: [PATCH 050/839] [3/N][Feat][Perf] Add new warmup infrastructure for JITs. Add provider registry and orchestration for JIT warmup (#50174) Signed-off-by: LopezCastroRoberto Co-authored-by: Codex --- docs/contributing/README.md | 1 + docs/contributing/jit_kernel_warmup.md | 308 ++++++++++++++++++ tests/model_executor/test_jit_warmup.py | 231 ++++++++++++- tests/v1/worker/test_jit_warmup_migration.py | 52 +++ vllm/model_executor/warmup/jit_warmup.py | 229 ++++++++++++- .../warmup/jit_warmup_triton_helper.py | 38 +++ vllm/model_executor/warmup/kernel_warmup.py | 23 +- .../warmup/v1_block_table_warmup.py | 29 -- vllm/v1/worker/block_table.py | 221 +++++++++---- vllm/v1/worker/cpu_model_runner.py | 2 +- vllm/v1/worker/gpu/model_runner.py | 2 + vllm/v1/worker/gpu_model_runner.py | 99 +++--- 12 files changed, 1067 insertions(+), 168 deletions(-) create mode 100644 docs/contributing/jit_kernel_warmup.md create mode 100644 tests/v1/worker/test_jit_warmup_migration.py delete mode 100644 vllm/model_executor/warmup/v1_block_table_warmup.py diff --git a/docs/contributing/README.md b/docs/contributing/README.md index 89acc6b7f5ae..25e7ecbdf1c7 100644 --- a/docs/contributing/README.md +++ b/docs/contributing/README.md @@ -71,6 +71,7 @@ uv pip install -e . --no-build-isolation For more details about installing from source and installing for other hardware, check out the [installation instructions](../getting_started/installation/README.md) for your hardware and head to the "Build wheel from source" section. For an optimized workflow when iterating on C++/CUDA kernels, see the [Incremental Compilation Workflow](./incremental_build.md) for recommendations. +For JIT kernel warmup conventions, see [JIT Kernel Warmup](./jit_kernel_warmup.md). !!! tip vLLM is compatible with Python versions 3.10 to 3.13. However, vLLM's default [Dockerfile](../../docker/Dockerfile) ships with Python 3.12 and tests in CI (except `mypy`) are run with Python 3.12. diff --git a/docs/contributing/jit_kernel_warmup.md b/docs/contributing/jit_kernel_warmup.md new file mode 100644 index 000000000000..4cf021279057 --- /dev/null +++ b/docs/contributing/jit_kernel_warmup.md @@ -0,0 +1,308 @@ +# JIT Kernel Warmup + +vLLM uses JIT-generated kernels from Triton, CuTeDSL, TileLang, and other backends. This contract makes their required specializations available during startup, before the first request, by warming the kernel's **compile-key space** without dummy runtime launches or real tensor allocation. + +Use it when adding a warmable JIT kernel or migrating an existing warmup path. + +## In This Guide + +- [1. Quickstart](#1-quickstart): for contributors adding or migrating a warmable kernel. +- [2. Search-Space Reference](#2-search-space-reference): additional details regarding warmup input expansion and traced dispatch rules. + +## 1. Quickstart + +Each warmable kernel defines its compile-key mapping and compile-only entry point beside its normal runtime implementation. The startup registry then warms only the wrappers selected by the current engine configuration. + +### Define the Kernel Wrapper + +Here, a **kernel wrapper** (or just **wrapper**) is an instance of a concrete `VllmJitKernel` subclass. + +Expose one wrapper near the kernel's normal runtime entry point. Prefer this shape: + +```python +class MyKernel(VllmJitKernel["MyKernel.CompileKey"]): + + @dataclass(frozen=True) + class CompileKey: + ... + + @staticmethod + def kernel(...): + ... + + def dispatch(self, ...) -> CompileKey: + return self.CompileKey(...) + + def get_warmup_keys(self, ...) -> list[CompileKey]: + return self._trace_dispatch(self.dispatch)(...) + + def compile(self, compile_key: CompileKey) -> None: + ... + + def __call__(self, ...): + return self.kernel(...) + + +MY_KERNEL = MyKernel() +``` + +`CompileKey`, `dispatch(...)`, and `get_warmup_keys(...)` are backend-agnostic. Backend-specific behavior belongs in `kernel(...)`, `compile(...)`, and `__call__(...)`. + +The module-level singleton should be used by warmup and by the runtime call path. This keeps dispatch behavior shared instead of duplicated. + +`VllmJitKernel.warmup(...)` compiles every key returned by `get_warmup_keys(...)`; wrappers should not reimplement it. + +### Choose Compile-Key Fields + +`CompileKey` must be frozen and hashable. Include only fields on which the backend specializes, such as tile sizes, head dimensions, dtypes, pointer alignment classes, or backend selectors; exclude runtime-only values. When unsure, inspect the backend cache key, specialization arguments, or verbose JIT-monitor output. + +### Generate Warmup Keys + +Use `_trace_dispatch(self.dispatch)` to describe representative inputs. The tracer maps them through the same specialization logic and deduplicates equal keys: + +```python +def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]: + max_tokens = vllm_config.scheduler_config.max_num_batched_tokens + return self._trace_dispatch(self.dispatch)( + num_tokens=WarmupIntRange(1, max_tokens + 1), + ) +``` + +Use independent ranges or alternatives for cartesian products, `zip_inputs(...)` for coupled rows, and `_when` for validity constraints. The complete syntax is documented in [Search-Space Reference](#2-search-space-reference). + +### Compile Without Launching + +`compile(compile_key)` means "make this specialization available". Depending on the backend, that may compile from source, call a compile-only API, load an already-built artifact, or compile on cache miss. + +`compile(...)` should not launch a real inference workload or allocate real tensors. Each DSL should expose fake tensor/spec descriptors suitable for compilation only. + +### Register the Selected Wrapper + +Register the wrapper where the runtime implementation is selected: + +```python +MY_KERNEL.register_warmup() +``` + +Registration records metadata only. It does not compile or launch the kernel. Repeated registrations from equivalent layers are allowed and deduplicated later. + +### Review Checklist + +- Warm actual compile keys rather than representative non-key inputs. +- Keep specialization mapping in `dispatch(...)` instead of duplicating it in warmup code. +- Use fake tensors or backend compile-only descriptors; never perform a dummy runtime launch. +- Keep registration metadata-only so model construction remains cheap and side-effect free. +- Compile registered kernels only through `kernel_warmup()`. +- Keep runtime execution and startup compilation separate and easy to review. +- Use one module-level wrapper instance for registration and runtime calls. + +## 2. Search-Space Reference + +### How Tracing Works + +`_trace_dispatch(...)` expands the inputs declared by `get_warmup_keys(...)`. Each concrete combination becomes a `dispatch_values` mapping from input names to selected values. `_when` may reject that mapping; otherwise the tracer evaluates `dispatch(...)` to construct one `CompileKey`. Equal keys are deduplicated after all combinations are evaluated. + +One call to `dispatch(...)` returns one key, but many input points may map to the same key. Prefer this traced mapping over manually reconstructing keys in warmup code; `dispatch(...)` should express the same specialization logic used by the runtime path. + +### Define Input Spaces + +Use ranges and alternatives for independent axes, `zip_inputs(...)` for coupled rows, and `_when` for validity constraints. + +#### Integer Ranges + +Use `WarmupIntRange` for integer ranges: + +```python +return self._trace_dispatch(self.dispatch)( + num_prefills=WarmupIntRange(1, max_prefills + 1), +) +``` + +`WarmupIntRange(start, stop, step)` follows Python `range(...)` semantics: `start` is inclusive, `stop` is exclusive, and `step` defaults to 1. + +For non-linear integer sequences, use `advance` to provide the action that computes each next value: + +```python +return self._trace_dispatch(self.dispatch)( + num_tokens=WarmupIntRange( + 1, + max_tokens + 1, + advance=lambda value: next_power_of_2(value) + 1, + ), +) +``` + +This is useful for traversing specialization boundaries without enumerating every integer. `advance` cannot be combined with a non-default `step`, and it must return a value greater than its input so expansion always makes forward progress. + +#### Independent Alternatives + +Use tuples or lists for independent alternatives. Multiple expanded inputs form a cartesian product: + +```python +return self._trace_dispatch(self.dispatch)( + query_slice_start=WarmupIntRange(0, 2), + query_slice_stop=(1, 2 * max_tokens - 1, 2 * max_tokens), + COMPRESS_RATIO=list(compress_ratios), +) +``` + +#### Coupled Inputs + +Use `zip_inputs(...)` when values must vary together row-by-row: + +```python +WARMUP_INPUTS = zip_inputs( + dict(compress_ratio=1, topk=0, topk_width=512), + dict(compress_ratio=4, topk=512, topk_width=512), +) + + +return self._trace_dispatch(self.dispatch)( + WARMUP_INPUTS, + WINDOW_SIZE=window_size, +) +``` + +Multiple `zip_inputs(...)` groups may be passed as positional arguments. The tracer forms the cartesian product across groups while preserving row-wise coupling inside each group. + +Every row in a `zip_inputs(...)` group must use the same string keys. A `zip_inputs(...)` group cannot specify a field that is also specified as a keyword input to `_trace_dispatch(...)`. + +#### Conditional Filtering + +Use `_when=...` to filter generated input points before they are passed to `dispatch(...)`. This is useful when independent ranges contain invalid combinations, but the validity rule belongs with the kernel warmup definition. + +```python +def _is_valid_warmup_input( + self, + *, + query_len: int, + num_reqs: int, + max_num_batched_tokens: int, +) -> bool: + return query_len + num_reqs - 1 <= max_num_batched_tokens + + +return self._trace_dispatch(self.dispatch)( + query_len=WarmupIntRange(1, max_tokens + 1), + num_reqs=WarmupIntRange(1, max_reqs + 1), + max_num_batched_tokens=max_tokens, + _when=self._is_valid_warmup_input, +) +``` + +`_when` accepts a function, bound method, or lambda and supports the same AST subset as `dispatch(...)`, including local assignments in function predicates. + +The predicate is evaluated on the expanded warmup inputs. If it returns `False`, that input point is skipped and no `CompileKey` is produced for it. + +### Write Dispatch Rules + +#### Local Assignments + +The traced body may contain local assignments, optionally annotated, followed by one `return self.CompileKey(...)` call. Local assignments let a kernel name intermediate specialization choices once and reuse them across fields: + +```python +def dispatch( + self, + *, + num_tokens: int, + vectorized: bool, +) -> CompileKey: + block_size = next_power_of_2(num_tokens) + return self.CompileKey( + BLOCK_SIZE=block_size, + VECTOR_WIDTH=4 if vectorized and block_size >= 4 else 1, + ) +``` + +#### Supported Expressions + +The evaluator supports these expressions inside local assignments and `CompileKey(...)` fields: + +| Feature | What It Allows | +| --- | --- | +| Names | Read dispatch inputs, local assignments, defaults, and module globals. | +| Constants | Use literals such as integers, strings, booleans, and `None`. | +| Attributes | Read structured values such as `cfg.block_size` or `mla_dims.v_head_dim`. | +| Subscriptions | Read sequence positions or mapping values such as `config[0]` and `config["block_size"]`. | +| Tuple/list literals | Build shapes, strides, and other small structured fields. | +| Conditional expressions | Select a field with `x if condition else y`. | +| Boolean expressions | Combine predicates with `and`, `or`, and `not`. | +| Comparisons | Use `==`, `!=`, `<`, `<=`, `>`, `>=`, `in`, `not in`, `is`, and `is not`. | +| Arithmetic | Use `+`, `-`, `*`, `//`, `%`, and `**`. | +| Unary minus | Build negative sentinel values or signed descriptors. | +| Helper calls | Call helpers with positional and explicit keyword arguments. | + +Python builtins such as `min(...)`, `max(...)`, and `len(...)` are resolved unless the name is overridden locally or globally. + +#### Helper Calls + +Helpers are useful for small specialization rules: + +```python +def dispatch(self, *, num_tokens: int, block_size: int) -> CompileKey: + return self.CompileKey( + PADDED_TOKENS=round_up(num_tokens, multiple=block_size), + ) +``` + +`_trace_dispatch(...)` does not inspect helper bodies. It evaluates the call arguments and invokes the helper as ordinary Python, so control flow inside that helper is outside the AST interpreter's scope. Keep helpers deterministic and side-effect free. + +#### Direct Keyword Forwarding + +For many direct pass-through fields, the dispatch `**kwargs` parameter may be unpacked into `CompileKey(...)`: + +```python +def dispatch( + self, + *, + num_tokens: int, + **compile_key_fields: int, +) -> CompileKey: + return self.CompileKey( + **compile_key_fields, + block_size=next_power_of_2(num_tokens), + ) +``` + +Unmatched dispatch arguments become compile-key fields and warmup inputs. Keep transformed inputs named and explicit. The unpacking must use the dispatch method's own `**kwargs` parameter directly and exactly once; arbitrary mappings, repeated unpacking, and helper-call `**kwargs` are rejected. The fully explicit form remains supported and is often clearer for non-trivial mappings. + +#### Unsupported Syntax + +Conditional expressions (`x if condition else y`) are supported, but statement-level `if` blocks are not supported directly inside traced `dispatch(...)` or `_when` bodies. The tracer expects a straight-line sequence of local assignments followed by one return expression. Small, pure helpers called by traced expressions execute as normal Python with concrete values and may use ordinary control flow, including `if` blocks. Do not put loops, mutation, side effects, or backend imports directly inside traced functions. Put environment and model gating in `get_warmup_keys(...)` or the outer warmup entry point. + +### Compile-Key Deduplication + +`_trace_dispatch(...)` deduplicates the resulting keys while preserving order. This is important when many runtime-like inputs map to the same static bucket. + +For example, this warmup range expands every token count, but the compile key only depends on the power-of-two bucket: + +```python +def dispatch( + self, + *, + num_tokens: int, +) -> CompileKey: + return self.CompileKey( + BLOCK_SIZE=next_power_of_2(num_tokens), + ) + + +def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]: + max_tokens = vllm_config.scheduler_config.max_num_batched_tokens + return self._trace_dispatch(self.dispatch)( + num_tokens=WarmupIntRange(1, max_tokens + 1), + ) +``` + +For `max_tokens == 8`, the expanded inputs are `1, 2, 3, 4, 5, 6, 7, 8`, but the returned keys are: + +```python +[ + CompileKey(BLOCK_SIZE=1), + CompileKey(BLOCK_SIZE=2), + CompileKey(BLOCK_SIZE=4), + CompileKey(BLOCK_SIZE=8), +] +``` + +Deduplication happens after `dispatch(...)` is evaluated, so the warmup system removes duplicate compile keys, not duplicate input values. `CompileKey` must be hashable for this to work; using `@dataclass(frozen=True)` is the standard pattern. diff --git a/tests/model_executor/test_jit_warmup.py b/tests/model_executor/test_jit_warmup.py index f93de50a3d30..b5c6f90ea491 100644 --- a/tests/model_executor/test_jit_warmup.py +++ b/tests/model_executor/test_jit_warmup.py @@ -9,11 +9,15 @@ import pytest from vllm.model_executor.warmup.jit_warmup import ( + JitWarmupRegistry, VllmJitKernel, WarmupIntRange, get_ast_full_name, zip_inputs, ) +from vllm.model_executor.warmup.jit_warmup_triton_helper import ( + triton_scalar_specialization_rep, +) def _next_power_of_2(value: int) -> int: @@ -94,6 +98,35 @@ def compile(self, compile_key: ToyKernel.CompileKey) -> None: self.compiled.append(compile_key) +@pytest.mark.parametrize( + ("value", "expected"), + [ + (-(1 << 63), 1 << 31), + (-(1 << 31) - 1, (1 << 31) + 1), + (-(1 << 31), 16), + (0, 16), + (1, 1), + (2, 2), + (16, 16), + ((1 << 31) - 1, 2), + (1 << 31, 1 << 31), + ((1 << 31) + 1, (1 << 31) + 1), + ((1 << 63) - 1, (1 << 31) + 1), + (1 << 63, 1 << 63), + ((1 << 63) + 1, (1 << 63) + 1), + ((1 << 64) - 1, (1 << 63) + 1), + ], +) +def test_triton_scalar_specialization_rep(value: int, expected: int) -> None: + assert triton_scalar_specialization_rep(value) == expected + + +@pytest.mark.parametrize("value", [-(1 << 63) - 1, 1 << 64]) +def test_triton_scalar_specialization_rep_rejects_out_of_range(value: int) -> None: + with pytest.raises(OverflowError, match="outside Triton's scalar range"): + triton_scalar_specialization_rep(value) + + def test_trace_dispatch_expands_ranges_dedupes_and_ignores_unused_inputs() -> None: cfg = _config() @@ -287,6 +320,31 @@ def compile(self, compile_key: CompileKey) -> None: StarKwargsKernel().compile_key({"tokens": 5, "block_size": 4}) +def test_dispatch_helper_calls_resolve_python_builtins() -> None: + class BuiltinKernel(VllmJitKernel["BuiltinKernel.CompileKey"]): + @dataclass(frozen=True) + class CompileKey: + value: int + + def dispatch( # type: ignore[override] + self, + *, + tokens: int, + limit: int, + ) -> CompileKey: + return self.CompileKey(value=max(1, min(tokens, limit))) + + def get_warmup_keys(self) -> list[CompileKey]: + return [] + + def compile(self, compile_key: CompileKey) -> None: + pass + + assert BuiltinKernel().compile_key({"tokens": 8, "limit": 4}) == ( + BuiltinKernel.CompileKey(value=4) + ) + + def test_dispatch_body_must_be_local_assignments_then_compile_key_return() -> None: class BranchKernel(VllmJitKernel["BranchKernel.CompileKey"]): @dataclass(frozen=True) @@ -320,10 +378,85 @@ def compile(self, compile_key: CompileKey) -> None: with pytest.raises(ValueError, match="local assignments"): BranchKernel() - with pytest.raises(ValueError, match=r"cannot use \*\*kwargs in CompileKey"): + with pytest.raises( + ValueError, + match=r"may unpack only its own \*\*kwargs parameter once", + ): KwargsReturnKernel() +def test_dispatch_can_forward_compile_key_fields() -> None: + class ForwardingKernel(VllmJitKernel["ForwardingKernel.CompileKey"]): + @dataclass(frozen=True) + class CompileKey: + mode: int + block_size: int + + def dispatch( # type: ignore[override] + self, + *, + tokens: int, + **compile_key_fields: int, + ) -> CompileKey: + return self.CompileKey( + **compile_key_fields, + block_size=_round_up(tokens, multiple=4), + ) + + def get_warmup_keys(self) -> list[CompileKey]: + return self._trace_dispatch(self.dispatch)( + tokens=(1, 5), + mode=(2, 3), + ) + + def compile(self, compile_key: CompileKey) -> None: + pass + + kernel = ForwardingKernel() + expected = kernel.CompileKey( + mode=2, + block_size=8, + ) + assert kernel.dispatch(tokens=5, mode=2) == expected + assert kernel.compile_key({"tokens": 5, "mode": 2}) == expected + assert kernel.get_warmup_keys() == [ + kernel.CompileKey(mode=2, block_size=4), + kernel.CompileKey(mode=3, block_size=4), + kernel.CompileKey(mode=2, block_size=8), + kernel.CompileKey(mode=3, block_size=8), + ] + with pytest.raises(TypeError, match="field 'block_size' is specified twice"): + kernel.compile_key({"tokens": 5, "mode": 2, "block_size": 4}) + with pytest.raises(TypeError, match="unexpected keyword argument 'extra'"): + kernel.compile_key({"tokens": 5, "mode": 2, "extra": 1}) + + +def test_dispatch_supports_tuple_and_mapping_subscriptions() -> None: + class SubscriptKernel(VllmJitKernel["SubscriptKernel.CompileKey"]): + @dataclass(frozen=True) + class CompileKey: + first: int + named: int + + def dispatch( # type: ignore[override] + self, + *, + values: tuple[int, ...], + config: dict[str, int], + ) -> CompileKey: + return self.CompileKey(first=values[0], named=config["named"]) + + def get_warmup_keys(self) -> list[CompileKey]: + return [] + + def compile(self, compile_key: CompileKey) -> None: + pass + + assert SubscriptKernel().compile_key( + {"values": (3, 5), "config": {"named": 7}} + ) == SubscriptKernel.CompileKey(first=3, named=7) + + def test_dispatch_reports_unsupported_expression_with_context() -> None: class UnsupportedKernel(VllmJitKernel["UnsupportedKernel.CompileKey"]): @dataclass(frozen=True) @@ -361,6 +494,102 @@ def test_warmup_compiles_all_returned_keys_in_order() -> None: ] +def test_runtime_cache_miss_compiles_and_caches_executor() -> None: + class CachedKernel(VllmJitKernel["CachedKernel.CompileKey"]): + @dataclass(frozen=True) + class CompileKey: + value: int + + def __init__(self) -> None: + self.compiled: list[CachedKernel.CompileKey] = [] + super().__init__() + + def dispatch(self, *, value: int) -> CompileKey: # type: ignore[override] + return self.CompileKey(value=value) + + def get_warmup_keys(self) -> list[CompileKey]: + return self._trace_dispatch(self.dispatch)(value=1) + + def compile(self, compile_key: CompileKey) -> None: + self.compiled.append(compile_key) + self._compiled_cache[compile_key] = object() + + def __call__(self, value: int) -> None: + compile_key = self.dispatch(value=value) + self._get_or_compile(compile_key) + + kernel = CachedKernel() + kernel.warmup() + kernel(1) + kernel(2) + + assert kernel.compiled == [ + CachedKernel.CompileKey(value=1), + CachedKernel.CompileKey(value=2), + ] + + +def test_registry_records_only_inside_model_setup_context() -> None: + registry = JitWarmupRegistry(_config()) + kernel = RecordingToyKernel() + + kernel.register_warmup(3, _config()) + with registry.activate(): + kernel.register_warmup(3, _config()) + + assert len(registry) == 1 + assert kernel.compiled == [] + + +def test_registry_expands_requests_and_deduplicates_owner_keys() -> None: + registry = JitWarmupRegistry(_config()) + kernel = RecordingToyKernel() + + with registry.activate(): + kernel.register_warmup(3, _config()) + kernel.register_warmup(5, _config()) + + registry.warmup() + + assert kernel.compiled == [ + ToyKernel.CompileKey(1, 1, 1, ("base", "default", -1, 1, 1), True), + ToyKernel.CompileKey(2, 2, 1, ("base", "default", -2, 2, 4), True), + ToyKernel.CompileKey(4, 4, 1, ("base", "default", -4, 1, 16), True), + ToyKernel.CompileKey(8, 8, 1, ("base", "default", -8, 2, 64), True), + ] + + +def test_registry_passes_vllm_config_to_default_requests() -> None: + class ConfigKernel(VllmJitKernel["ConfigKernel.CompileKey"]): + @dataclass(frozen=True) + class CompileKey: + value: int + + def __init__(self) -> None: + self.compiled: list[ConfigKernel.CompileKey] = [] + super().__init__() + + def dispatch(self, *, value: int) -> CompileKey: # type: ignore[override] + return self.CompileKey(value=value) + + def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: + return [self.dispatch(value=vllm_config.bias)] + + def compile(self, compile_key: CompileKey) -> None: + self.compiled.append(compile_key) + + registry = JitWarmupRegistry(_config(bias=7)) + kernel = ConfigKernel() + + with registry.activate(): + kernel.register_warmup() + kernel.register_warmup() + assert len(registry) == 1 + registry.warmup() + + assert kernel.compiled == [ConfigKernel.CompileKey(value=7)] + + def test_get_ast_full_name_handles_names_attributes_and_other_nodes() -> None: dotted_expr = ast.parse("foo.bar.baz").body[0] call_expr = ast.parse("foo()").body[0] diff --git a/tests/v1/worker/test_jit_warmup_migration.py b/tests/v1/worker/test_jit_warmup_migration.py new file mode 100644 index 000000000000..ef887c5351f6 --- /dev/null +++ b/tests/v1/worker/test_jit_warmup_migration.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Validate the registry reference kernel against runtime dispatch.""" + +import pytest + +from vllm.platforms import current_platform + +if not current_platform.is_cuda_alike(): + pytest.skip("NVIDIA dispatch tests require CUDA", allow_module_level=True) + +from vllm.v1.worker.block_table import ComputeSlotMappingKernel + + +@pytest.mark.parametrize( + ("kv_cache_block_size", "blocks_per_kv_block", "block_size", "block_size_rep"), + [ + (256, 1, 256, 16), + (256, 4, 64, 16), + (64, 1, 64, 16), + (8, 1, 8, 2), + (4, 1, 4, 2), + ], +) +def test_compute_slot_mapping_warmup_matches_runtime_specializations( + kv_cache_block_size: int, + blocks_per_kv_block: int, + block_size: int, + block_size_rep: int, +) -> None: + kernel = ComputeSlotMappingKernel() + kwargs = dict( + kv_cache_block_size=kv_cache_block_size, + blocks_per_kv_block=blocks_per_kv_block, + total_cp_world_size=1, + total_cp_rank=0, + cp_kv_cache_interleave_size=1, + block_table_stride=32768, + block_size=block_size, + ) + expected = kernel.CompileKey( + kv_cache_block_size=kv_cache_block_size, + blocks_per_kv_block=blocks_per_kv_block, + total_cp_world_size=1, + total_cp_rank=0, + cp_kv_cache_interleave_size=1, + block_table_stride=16, + block_size=block_size_rep, + ) + + assert kernel.dispatch(**kwargs) == expected + assert kernel.get_warmup_keys(**kwargs) == [expected] diff --git a/vllm/model_executor/warmup/jit_warmup.py b/vllm/model_executor/warmup/jit_warmup.py index d8c51409502e..9c43bbfc33d8 100644 --- a/vllm/model_executor/warmup/jit_warmup.py +++ b/vllm/model_executor/warmup/jit_warmup.py @@ -5,16 +5,20 @@ from __future__ import annotations import ast +import builtins import inspect import itertools import operator import textwrap from abc import ABC, abstractmethod -from collections.abc import Callable, Iterable, Mapping +from collections.abc import Callable, Iterable, Iterator, Mapping +from contextlib import contextmanager +from contextvars import ContextVar from dataclasses import dataclass from typing import Any, Generic, TypeVar __all__ = [ + "JitWarmupRegistry", "VllmJitKernel", "WarmupIntRange", "get_ast_full_name", @@ -136,8 +140,15 @@ class _CompileKeyDispatchTrace: field_exprs: tuple[tuple[str, ast.AST], ...] globals: Mapping[str, Any] input_names: frozenset[str] + # Named parameters are excluded from direct **kwargs forwarding. + named_parameters: frozenset[str] | None defaults: Mapping[str, Any] + def input_names_for(self, available_names: set[str]) -> frozenset[str]: + if self.named_parameters is None: + return self.input_names + return self.input_names | (available_names - self.named_parameters) + def compile_key( self, compile_key_type: type[CompileKeyT], @@ -146,12 +157,20 @@ def compile_key( dispatch_values = _eval_local_exprs( self.local_exprs, {**self.defaults, **kwargs}, self.globals ) - return compile_key_type( - **{ - field: _eval_dispatch_expr(expr, dispatch_values, self.globals) - for field, expr in self.field_exprs + named_parameters = self.named_parameters + # Materialize direct fields before evaluating named AST expressions. + fields: dict[str, Any] = {} + if named_parameters is not None: + fields = { + name: value + for name, value in kwargs.items() + if name not in named_parameters } - ) + for field, expr in self.field_exprs: + if field in fields: + raise TypeError(f"CompileKey field '{field}' is specified twice") + fields[field] = _eval_dispatch_expr(expr, dispatch_values, self.globals) + return compile_key_type(**fields) @dataclass(frozen=True) @@ -186,6 +205,10 @@ def matches(self, kwargs: Mapping[str, Any]) -> bool: ast.LtE: operator.le, ast.Gt: operator.gt, ast.GtE: operator.ge, + ast.In: lambda left, right: left in right, + ast.NotIn: lambda left, right: left not in right, + ast.Is: operator.is_, + ast.IsNot: operator.is_not, } @@ -200,8 +223,9 @@ def _dispatch_expr_error(node: ast.AST, reason: str) -> ValueError: return ValueError( f"{reason}: {_dispatch_expr_source(node)}. " "Supported dispatch expressions are names, constants, attributes, " - "tuple/list literals, conditional expressions, comparisons, boolean " - "operators, unary not/minus, arithmetic, and calls without **kwargs." + "subscriptions, tuple/list literals, conditional expressions, " + "comparisons, boolean operators, unary not/minus, arithmetic, and " + "calls without **kwargs." ) @@ -225,6 +249,8 @@ def visit_Name(self, node: ast.Name) -> Any: return self.values[node.id] if node.id in self.globals: return self.globals[node.id] + if hasattr(builtins, node.id): + return getattr(builtins, node.id) raise _dispatch_expr_error(node, f"Unknown dispatch name '{node.id}'") def visit_Constant(self, node: ast.Constant) -> Any: @@ -299,6 +325,9 @@ def visit_Call(self, node: ast.Call) -> Any: def visit_Attribute(self, node: ast.Attribute) -> Any: return getattr(self.visit(node.value), node.attr) + def visit_Subscript(self, node: ast.Subscript) -> Any: + return self.visit(node.value)[self.visit(node.slice)] + def get_ast_full_name(node: ast.AST) -> str | None: if isinstance(node, ast.Name): @@ -310,15 +339,23 @@ def get_ast_full_name(node: ast.AST) -> str | None: return None -def get_function_source_node(fn: Callable[..., Any]) -> ast.FunctionDef: +def get_function_source_node(fn: Callable[..., Any]) -> ast.FunctionDef | ast.Lambda: source_fn = getattr(fn, "fn", fn) source = textwrap.dedent(inspect.getsource(source_fn)) tree = ast.parse(source) function_defs = [node for node in tree.body if isinstance(node, ast.FunctionDef)] - if len(function_defs) != 1: - name = getattr(source_fn, "__name__", type(source_fn).__name__) - raise ValueError(f"Expected one function in {name}, found {len(function_defs)}") - return function_defs[0] + if len(function_defs) == 1: + return function_defs[0] + + lambdas = [node for node in ast.walk(tree) if isinstance(node, ast.Lambda)] + if len(lambdas) == 1: + return lambdas[0] + + name = getattr(source_fn, "__name__", type(source_fn).__name__) + raise ValueError( + f"Expected one function or lambda in {name}, found " + f"{len(function_defs)} functions and {len(lambdas)} lambdas" + ) def _eval_dispatch_expr( @@ -349,8 +386,11 @@ def _collect_input_names( def _collect_expression_body( fn: Callable[..., Any], - function_def: ast.FunctionDef, + function_def: ast.FunctionDef | ast.Lambda, ) -> tuple[list[tuple[str, ast.AST]], ast.AST]: + if isinstance(function_def, ast.Lambda): + return [], function_def.body + local_exprs: list[tuple[str, ast.AST]] = [] for statement in function_def.body: if ( @@ -421,6 +461,10 @@ def _trace_compile_key_dispatch( source_fn = getattr(fn, "__func__", fn) globals_ = source_fn.__globals__ function_def = get_function_source_node(fn) + if isinstance(function_def, ast.Lambda): + raise _dispatch_expr_error( + function_def, "Dispatch must be a function definition" + ) local_exprs, return_expr = _collect_expression_body(fn, function_def) if not isinstance(return_expr, ast.Call): @@ -431,13 +475,37 @@ def _trace_compile_key_dispatch( field_exprs: list[tuple[str, ast.AST]] = [] defaults, candidate_names = _function_trace_inputs(fn) + # Fields captured by dispatch **kwargs are forwarded, not AST-evaluated. + signature = inspect.signature(fn) + variadic_keyword = next( + ( + name + for name, parameter in signature.parameters.items() + if parameter.kind is inspect.Parameter.VAR_KEYWORD + ), + None, + ) + candidate_names.discard(variadic_keyword) input_names: set[str] = set() local_names = {name for name, _ in local_exprs} for _, expr in local_exprs: input_names.update(_collect_input_names(expr, candidate_names)) + named_parameters: frozenset[str] | None = None for keyword in return_expr.keywords: + # CompileKey may unpack only that **kwargs parameter, once. if keyword.arg is None: - raise ValueError(f"{fn.__name__} cannot use **kwargs in CompileKey") + if ( + named_parameters is not None + or variadic_keyword is None + or not isinstance(keyword.value, ast.Name) + or keyword.value.id != variadic_keyword + ): + raise ValueError( + f"{fn.__name__} may unpack only its own **kwargs parameter " + "once in CompileKey" + ) + named_parameters = frozenset(candidate_names) + continue field_exprs.append((keyword.arg, keyword.value)) input_names.update( _collect_input_names(keyword.value, candidate_names, local_names) @@ -448,6 +516,7 @@ def _trace_compile_key_dispatch( tuple(field_exprs), globals_, frozenset(input_names), + named_parameters, defaults, ) @@ -482,10 +551,32 @@ class VllmJitKernel(Generic[CompileKeyT], ABC): CompileKey: type[CompileKeyT] def __init__(self) -> None: - self.compile_key_dispatch_trace = _trace_compile_key_dispatch(self.dispatch) + self._dispatch_trace = _trace_compile_key_dispatch(self.dispatch) + self._compiled_cache: dict[Any, Any] = {} def compile_key(self, kwargs: Mapping[str, Any]) -> CompileKeyT: - return self.compile_key_dispatch_trace.compile_key(self.CompileKey, kwargs) + return self._dispatch_trace.compile_key(self.CompileKey, kwargs) + + def _get_or_compile( + self, + compile_key: CompileKeyT, + *, + runtime_context: Mapping[str, Any] | None = None, + ) -> Any: + """Return a cached executor, compiling it on a monitored cache miss.""" + if compile_key not in self._compiled_cache: + self.compile(compile_key) + + try: + return self._compiled_cache[compile_key] + except KeyError as exc: + details = [f"compile_key={compile_key!r}"] + if runtime_context: + details.append(f"runtime_context={dict(runtime_context)!r}") + raise RuntimeError( + f"{type(self).__name__}.compile(...) did not cache its JIT " + f"executor ({', '.join(details)})" + ) from exc def _trace_dispatch( self, dispatch: CompileKeyDispatchFn[CompileKeyT] @@ -506,7 +597,11 @@ def traced( predicate_trace = ( _trace_warmup_predicate(_when) if _when is not None else None ) - input_names = compile_key_dispatch_trace.input_names + # Unmatched **kwargs fields also belong to the expansion space. + available_names = set(kwargs).union( + *(group.rows[0] for group in input_groups) + ) + input_names = compile_key_dispatch_trace.input_names_for(available_names) if predicate_trace is not None: input_names = input_names | predicate_trace.input_names expanded_input_groups = tuple( @@ -565,7 +660,105 @@ def compile(self, compile_key: CompileKeyT) -> None: """Compile one warmup key.""" raise NotImplementedError + def register_warmup(self, *args: Any, **kwargs: Any) -> None: + """Register this kernel with the active runner's warmup registry.""" + JitWarmupRegistry.register(self, *args, **kwargs) + def warmup(self, *args: Any, **kwargs: Any) -> None: """Compile this kernel's warmup keys.""" for compile_key in self.get_warmup_keys(*args, **kwargs): self.compile(compile_key) + + +class JitWarmupRegistry: + """Collect and compile JIT kernels selected during runner setup.""" + + _active: ContextVar[JitWarmupRegistry | None] = ContextVar( + "active_jit_warmup_registry", + default=None, + ) + + def __init__(self, vllm_config: Any) -> None: + self.vllm_config = vllm_config + self._registrations: dict[ + VllmJitKernel[Any], + list[tuple[tuple[Any, ...], dict[str, Any]]], + ] = {} + + @contextmanager + def activate(self) -> Iterator[None]: + """Collect registrations made in this context.""" + token = self._active.set(self) + try: + yield + finally: + self._active.reset(token) + + @classmethod + def register( + cls, + kernel: VllmJitKernel[Any], + *args: Any, + **kwargs: Any, + ) -> None: + """Register a kernel with the active registry, if one exists.""" + registry = cls._active.get() + if registry is not None: + registry._add(kernel, args, kwargs) + + def _add( + self, + kernel: VllmJitKernel[Any], + args: tuple[Any, ...], + kwargs: dict[str, Any], + ) -> None: + registrations = self._registrations.setdefault(kernel, []) + if ( + not args + and not kwargs + and any( + not registered_args and not registered_kwargs + for registered_args, registered_kwargs in registrations + ) + ): + return + registrations.append((args, kwargs)) + + def __len__(self) -> int: + return sum(len(registrations) for registrations in self._registrations.values()) + + def warmup(self) -> None: + """Expand registrations and compile each wrapper/key pair once.""" + from tqdm import tqdm + + from vllm.distributed import is_global_first_rank + + kernel_items: list[tuple[VllmJitKernel[Any], dict[Any, None]]] = [] + for kernel, registrations in self._registrations.items(): + compile_keys: dict[Any, None] = {} + for args, kwargs in registrations: + if not args and not kwargs: + args = (self.vllm_config,) + for compile_key in kernel.get_warmup_keys(*args, **kwargs): + compile_keys[compile_key] = None + if compile_keys: + kernel_items.append((kernel, compile_keys)) + + if not kernel_items: + return + + total_keys = sum(len(compile_keys) for _, compile_keys in kernel_items) + with tqdm( + kernel_items, + desc=f"JIT kernel warmup ({total_keys} compile keys)", + disable=not is_global_first_rank(), + dynamic_ncols=True, + unit="kernel", + ) as progress: + for kernel, compile_keys in progress: + progress.set_postfix_str( + f"{kernel.__class__.__name__} ({len(compile_keys)} keys)", + refresh=False, + ) + for compile_key in compile_keys: + kernel.compile(compile_key) diff --git a/vllm/model_executor/warmup/jit_warmup_triton_helper.py b/vllm/model_executor/warmup/jit_warmup_triton_helper.py index b90975762175..29da528c93da 100644 --- a/vllm/model_executor/warmup/jit_warmup_triton_helper.py +++ b/vllm/model_executor/warmup/jit_warmup_triton_helper.py @@ -12,6 +12,42 @@ ) +def triton_scalar_specialization_rep(value: int) -> int: + """Return an integer with the same default Triton JIT specialization. + + For an ordinary integer argument, Triton's cache key contains its inferred + type (``i32``, ``i64``, or ``u64``) and one of three value classes: + + * ``1`` is specialized as the exact constant ``1``. + * Multiples of 16 receive a ``tt.divisibility = 16`` attribute. + * All other values have no value specialization. + + Warmup only needs one concrete value for each cache-key class. This helper + returns ``1`` for the exact-one class and otherwise returns a divisible or + generic representative while preserving the inferred integer type. + + This applies only to non-``constexpr`` integer arguments using Triton's + default specialization. Do not use it for arguments listed in + ``do_not_specialize`` or ``do_not_specialize_on_alignment``. + """ + if value == 1: + return 1 + + if -(1 << 31) <= value < (1 << 31): + divisible_rep = 16 + generic_rep = 2 + elif -(1 << 63) <= value < (1 << 63): + divisible_rep = 1 << 31 + generic_rep = (1 << 31) + 1 + elif 0 <= value < (1 << 64): + divisible_rep = 1 << 63 + generic_rep = (1 << 63) + 1 + else: + raise OverflowError(f"Integer {value} is outside Triton's scalar range") + + return divisible_rep if value % 16 == 0 else generic_rep + + @dataclass(frozen=True) class TritonWarmupTensor: # Compile-only tensor descriptor for Triton pointer specialization. @@ -167,6 +203,8 @@ def trace_triton_kernel_specialization_args( kernel: Callable[..., Any], ) -> tuple[str, ...]: function_def = get_function_source_node(kernel) + if not isinstance(function_def, ast.FunctionDef): + raise ValueError("Expected Triton kernel to be defined as a function") source_fn = getattr(kernel, "fn", kernel) arg_names = tuple(inspect.signature(source_fn).parameters) constexpr_args = _triton_constexpr_arg_names(kernel, function_def, arg_names) diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index cd16c63b589f..e25a60c8445a 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -6,6 +6,7 @@ happen during model execution. """ +import time from typing import TYPE_CHECKING import torch @@ -36,9 +37,6 @@ from vllm.model_executor.warmup.sparse_mla_triton_warmup import ( sparse_mla_triton_warmup, ) -from vllm.model_executor.warmup.v1_block_table_warmup import ( - warm_v1_block_table_kernels, -) from vllm.platforms import current_platform from vllm.utils.deep_gemm import is_deep_gemm_supported from vllm.utils.flashinfer import has_flashinfer @@ -103,15 +101,28 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False): ) if not worker.use_v2_model_runner: - # Pooling models do not use the generation slot-mapping path. - if not worker.model_runner.is_pooling_model: - warm_v1_block_table_kernels(worker.model_runner) # The KV-block zeroing kernel is driven by the scheduler's # `new_block_ids_to_zero`, so no dummy run ever reaches it. zeroer = getattr(worker.model_runner, "_kv_block_zeroer", None) if zeroer is not None: zeroer.warmup(worker.model_runner.kv_cache_config.num_blocks) + if worker.vllm_config.kernel_config.enable_jit_warmup: + logger.info("JIT kernel warmup starting.") + jit_warmup_start = time.perf_counter() + try: + worker.model_runner.jit_warmup_registry.warmup() + except Exception: + logger.exception( + "JIT kernel warmup failed after %.2fs.", + time.perf_counter() - jit_warmup_start, + ) + raise + logger.info( + "JIT kernel warmup finished in %.2fs.", + time.perf_counter() - jit_warmup_start, + ) + qwen_triton_warmup(worker.model_runner, worker.vllm_config.model_config) compilation_config = worker.vllm_config.compilation_config diff --git a/vllm/model_executor/warmup/v1_block_table_warmup.py b/vllm/model_executor/warmup/v1_block_table_warmup.py deleted file mode 100644 index d49e1ba7cc89..000000000000 --- a/vllm/model_executor/warmup/v1_block_table_warmup.py +++ /dev/null @@ -1,29 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Warm up v1 block-table Triton kernels.""" - -from typing import TYPE_CHECKING - -import torch - -if TYPE_CHECKING: - from vllm.v1.worker.gpu_model_runner import GPUModelRunner - -_SLOT_MAPPING_WARMUP_TOKENS = 8 - - -def warm_v1_block_table_kernels(runner: "GPUModelRunner") -> None: - """JIT-compile ``_compute_slot_mapping_kernel`` for the real block tables.""" - - device = runner.device - block_table = runner.input_batch.block_table - num_tokens = min( - _SLOT_MAPPING_WARMUP_TOKENS, - runner.scheduler_config.max_num_batched_tokens, - ) - if num_tokens <= 0: - return - - query_start_loc = torch.tensor([0, num_tokens], dtype=torch.int32, device=device) - positions = torch.arange(num_tokens, dtype=torch.int64, device=device) - block_table.compute_slot_mapping(1, query_start_loc, positions) diff --git a/vllm/v1/worker/block_table.py b/vllm/v1/worker/block_table.py index 85228364afcf..5bf5df3b3985 100644 --- a/vllm/v1/worker/block_table.py +++ b/vllm/v1/worker/block_table.py @@ -2,13 +2,22 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import math +from dataclasses import dataclass from enum import Enum +from typing import Any import numpy as np import torch from vllm.distributed import get_dcp_group, get_pcp_group from vllm.logger import init_logger +from vllm.model_executor.warmup.jit_warmup import ( + VllmJitKernel, +) +from vllm.model_executor.warmup.jit_warmup_triton_helper import ( + TritonWarmupTensor, + triton_scalar_specialization_rep, +) from vllm.triton_utils import tl, triton from vllm.utils.math_utils import cdiv from vllm.v1.attention.backends.utils import PAD_SLOT_ID @@ -134,6 +143,16 @@ def __init__( self.dcp_rank = 0 self.cp_kv_cache_interleave_size = cp_kv_cache_interleave_size self.slot_mapping_mode = slot_mapping_mode + if self.slot_mapping_mode == SlotMappingMode.TOKEN_TO_KV_SLOT: + _COMPUTE_SLOT_MAPPING_KERNEL.register_warmup( + kv_cache_block_size=self.kv_cache_block_size, + blocks_per_kv_block=self.blocks_per_kv_block, + total_cp_world_size=self.dcp_world_size, + total_cp_rank=self.dcp_rank, + cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size, + block_table_stride=self.block_table.gpu.stride(0), + block_size=self.block_size, + ) def append_row( self, @@ -192,7 +211,8 @@ def compute_slot_mapping( return assert self.slot_mapping_mode == SlotMappingMode.TOKEN_TO_KV_SLOT - _compute_slot_mapping_kernel[(num_reqs + 1,)]( + _COMPUTE_SLOT_MAPPING_KERNEL( + num_reqs, num_tokens, self.max_num_batched_tokens, query_start_loc, @@ -201,13 +221,11 @@ def compute_slot_mapping( self.block_table.gpu.stride(0), self.block_size, self.slot_mapping.gpu, - KV_CACHE_BLOCK_SIZE=self.kv_cache_block_size, - BLOCKS_PER_KV_BLOCK=self.blocks_per_kv_block, - TOTAL_CP_WORLD_SIZE=self.dcp_world_size, - TOTAL_CP_RANK=self.dcp_rank, - CP_KV_CACHE_INTERLEAVE_SIZE=self.cp_kv_cache_interleave_size, - PAD_ID=PAD_SLOT_ID, - BLOCK_SIZE=1024, + self.kv_cache_block_size, + self.blocks_per_kv_block, + self.dcp_world_size, + self.dcp_rank, + self.cp_kv_cache_interleave_size, ) def commit_block_table(self, num_reqs: int) -> None: @@ -376,67 +394,136 @@ def __getitem__(self, idx: int) -> "BlockTable": return self.block_tables[idx] -@triton.jit(do_not_specialize=["num_tokens", "max_num_tokens"]) -def _compute_slot_mapping_kernel( - num_tokens, - max_num_tokens, - query_start_loc_ptr, # [num_reqs + 1], int32 - positions_ptr, # [num_tokens], int64 - block_table_ptr, # [max_num_reqs, max_num_blocks_per_req], int32 (flat) - block_table_stride, # max_num_blocks_per_req - block_size, - slot_mapping_ptr, # [max_num_tokens], int64 - KV_CACHE_BLOCK_SIZE: tl.constexpr, - BLOCKS_PER_KV_BLOCK: tl.constexpr, - TOTAL_CP_WORLD_SIZE: tl.constexpr, - TOTAL_CP_RANK: tl.constexpr, - CP_KV_CACHE_INTERLEAVE_SIZE: tl.constexpr, - PAD_ID: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - req_idx = tl.program_id(0) - - if req_idx == tl.num_programs(0) - 1: - # Pad remaining slots for CUDA graph compatibility. - for i in range(num_tokens, max_num_tokens, BLOCK_SIZE): +class ComputeSlotMappingKernel(VllmJitKernel["ComputeSlotMappingKernel.CompileKey"]): + triton_block_size = 1024 + + @dataclass(frozen=True) + class CompileKey: + kv_cache_block_size: int + blocks_per_kv_block: int + total_cp_world_size: int + total_cp_rank: int + cp_kv_cache_interleave_size: int + block_table_stride: int + block_size: int + + @staticmethod + @triton.jit(do_not_specialize=["num_tokens", "max_num_tokens"]) + def kernel( + num_tokens, + max_num_tokens, + query_start_loc_ptr, # [num_reqs + 1], int32 + positions_ptr, # [num_tokens], int64 + block_table_ptr, # [max_num_reqs, max_num_blocks_per_req], int32 (flat) + block_table_stride, # max_num_blocks_per_req + block_size, + slot_mapping_ptr, # [max_num_tokens], int64 + KV_CACHE_BLOCK_SIZE: tl.constexpr, + BLOCKS_PER_KV_BLOCK: tl.constexpr, + TOTAL_CP_WORLD_SIZE: tl.constexpr, + TOTAL_CP_RANK: tl.constexpr, + CP_KV_CACHE_INTERLEAVE_SIZE: tl.constexpr, + PAD_ID: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + ): + req_idx = tl.program_id(0) + + if req_idx == tl.num_programs(0) - 1: + # Pad remaining slots for CUDA graph compatibility. + for i in range(num_tokens, max_num_tokens, BLOCK_SIZE): + offsets = i + tl.arange(0, BLOCK_SIZE) + tl.store( + slot_mapping_ptr + offsets, + PAD_ID, + mask=offsets < max_num_tokens, + ) + return + + start_idx = tl.load(query_start_loc_ptr + req_idx).to(tl.int64) + end_idx = tl.load(query_start_loc_ptr + req_idx + 1).to(tl.int64) + + virtual_block_size = KV_CACHE_BLOCK_SIZE * TOTAL_CP_WORLD_SIZE + row_offset = req_idx * block_table_stride + for i in range(start_idx, end_idx, BLOCK_SIZE): offsets = i + tl.arange(0, BLOCK_SIZE) - tl.store( - slot_mapping_ptr + offsets, - PAD_ID, - mask=offsets < max_num_tokens, + mask = offsets < end_idx + pos = tl.load(positions_ptr + offsets, mask=mask, other=0) + virtual_block_indices = pos // virtual_block_size + virtual_block_offsets = pos - virtual_block_indices * virtual_block_size + is_local = ( + virtual_block_offsets // CP_KV_CACHE_INTERLEAVE_SIZE + ) % TOTAL_CP_WORLD_SIZE == TOTAL_CP_RANK + local_block_offsets = ( + virtual_block_offsets + // (TOTAL_CP_WORLD_SIZE * CP_KV_CACHE_INTERLEAVE_SIZE) + ) * CP_KV_CACHE_INTERLEAVE_SIZE + ( + virtual_block_offsets % CP_KV_CACHE_INTERLEAVE_SIZE + ) + + block_indices = ( + virtual_block_indices * BLOCKS_PER_KV_BLOCK + + local_block_offsets // block_size ) - return - - start_idx = tl.load(query_start_loc_ptr + req_idx).to(tl.int64) - end_idx = tl.load(query_start_loc_ptr + req_idx + 1).to(tl.int64) - - virtual_block_size = KV_CACHE_BLOCK_SIZE * TOTAL_CP_WORLD_SIZE - row_offset = req_idx * block_table_stride - for i in range(start_idx, end_idx, BLOCK_SIZE): - offsets = i + tl.arange(0, BLOCK_SIZE) - mask = offsets < end_idx - pos = tl.load(positions_ptr + offsets, mask=mask, other=0) - virtual_block_indices = pos // virtual_block_size - virtual_block_offsets = pos - virtual_block_indices * virtual_block_size - is_local = ( - virtual_block_offsets // CP_KV_CACHE_INTERLEAVE_SIZE - ) % TOTAL_CP_WORLD_SIZE == TOTAL_CP_RANK - local_block_offsets = ( - virtual_block_offsets // (TOTAL_CP_WORLD_SIZE * CP_KV_CACHE_INTERLEAVE_SIZE) - ) * CP_KV_CACHE_INTERLEAVE_SIZE + ( - virtual_block_offsets % CP_KV_CACHE_INTERLEAVE_SIZE + block_numbers = tl.load( + block_table_ptr + row_offset + block_indices, + mask=mask & is_local, + other=0, + ).to(tl.int64) + slot_offsets = local_block_offsets % block_size + slot_ids = block_numbers * block_size + slot_offsets + slot_ids = tl.where(is_local, slot_ids, PAD_ID) + tl.store(slot_mapping_ptr + offsets, slot_ids, mask=mask) + + def dispatch( # type: ignore[override] + self, + *, + block_table_stride: int, + block_size: int, + **compile_key_fields: int, + ) -> CompileKey: + return self.CompileKey( + **compile_key_fields, + block_table_stride=triton_scalar_specialization_rep(block_table_stride), + block_size=triton_scalar_specialization_rep(block_size), ) - block_indices = ( - virtual_block_indices * BLOCKS_PER_KV_BLOCK - + local_block_offsets // block_size + def get_warmup_keys(self, **dispatch_kwargs: int) -> list[CompileKey]: + return self._trace_dispatch(self.dispatch)(**dispatch_kwargs) + + def compile(self, compile_key: CompileKey) -> None: + warmup = getattr(self.kernel, "warmup", None) + assert warmup is not None + int32_ptr = TritonWarmupTensor(torch.int32) + int64_ptr = TritonWarmupTensor(torch.int64) + warmup( + 2, # arbitrary, num_tokens in do_not_specialize + 2, # arbitrary, max_num_tokens in do_not_specialize + int32_ptr, + int64_ptr, + int32_ptr, + compile_key.block_table_stride, + compile_key.block_size, + int64_ptr, + KV_CACHE_BLOCK_SIZE=compile_key.kv_cache_block_size, + BLOCKS_PER_KV_BLOCK=compile_key.blocks_per_kv_block, + TOTAL_CP_WORLD_SIZE=compile_key.total_cp_world_size, + TOTAL_CP_RANK=compile_key.total_cp_rank, + CP_KV_CACHE_INTERLEAVE_SIZE=compile_key.cp_kv_cache_interleave_size, + PAD_ID=PAD_SLOT_ID, + BLOCK_SIZE=self.triton_block_size, + grid=(2,), ) - block_numbers = tl.load( - block_table_ptr + row_offset + block_indices, - mask=mask & is_local, - other=0, - ).to(tl.int64) - slot_offsets = local_block_offsets % block_size - slot_ids = block_numbers * block_size + slot_offsets - slot_ids = tl.where(is_local, slot_ids, PAD_ID) - tl.store(slot_mapping_ptr + offsets, slot_ids, mask=mask) + + def __call__( + self, + num_reqs: int, + *args: Any, + ) -> None: + self.kernel[(num_reqs + 1,)]( + *args, + PAD_ID=PAD_SLOT_ID, + BLOCK_SIZE=self.triton_block_size, + ) + + +_COMPUTE_SLOT_MAPPING_KERNEL = ComputeSlotMappingKernel() diff --git a/vllm/v1/worker/cpu_model_runner.py b/vllm/v1/worker/cpu_model_runner.py index 67cffa58454a..3e7f0d4c1e88 100644 --- a/vllm/v1/worker/cpu_model_runner.py +++ b/vllm/v1/worker/cpu_model_runner.py @@ -74,7 +74,7 @@ def _postprocess_triton(self) -> None: import vllm.v1.worker.block_table - vllm.v1.worker.block_table._compute_slot_mapping_kernel = ( + vllm.v1.worker.block_table._COMPUTE_SLOT_MAPPING_KERNEL.kernel = ( cpu_tl.compute_slot_mapping_kernel ) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index d68666cf4a0c..aff0c09cab58 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -51,6 +51,7 @@ get_offloader, set_offloader, ) +from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.encoder_budget import ( MultiModalBudget, @@ -170,6 +171,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.speculative_config is not None and self.speculative_config.use_dspark() ) self.observability_config = vllm_config.observability_config + self.jit_warmup_registry = JitWarmupRegistry(vllm_config) self.device = device self.dtype = self.model_config.dtype diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 9087f878546f..a8d08a7bf293 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -100,6 +100,7 @@ get_offloader, set_offloader, ) +from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.encoder_budget import MultiModalBudget from vllm.multimodal.inputs import ( @@ -516,6 +517,7 @@ def __init__( self.scheduler_config = vllm_config.scheduler_config self.speculative_config = vllm_config.speculative_config self.observability_config = vllm_config.observability_config + self.jit_warmup_registry = JitWarmupRegistry(vllm_config) model_config = self.model_config cache_config = self.cache_config @@ -750,34 +752,37 @@ def __init__( self._init_kernel_block_sizes = [placeholder_block_size] self._init_max_num_blocks = [placeholder_max_num_blocks] self._init_slot_mapping_modes = [SlotMappingMode.TOKEN_TO_KV_SLOT] - self.input_batch = InputBatch( - max_num_reqs=self.max_num_reqs, - # We need to use the encoder length for encoder-decoder - # because of KV cache for cross-attention. - max_model_len=max(self.max_model_len, self.max_encoder_len), - max_num_batched_tokens=self.max_num_tokens, - device=self.device, - vocab_size=self.model_config.get_vocab_size(), - block_sizes=[placeholder_block_size], - kernel_block_sizes=[placeholder_block_size], - max_num_blocks_per_req=[placeholder_max_num_blocks], - num_spec_tokens=self.num_spec_tokens, - logitsprocs=build_logitsprocs( - self.vllm_config, - self.device, - PIN_MEMORY, - self.is_pooling_model, - custom_logitsprocs, - ), - # We currently don't know whether a particular custom logits processor - # uses output token ids so we set this conservatively. Thinking-budget - # tracking is requested dynamically when a budgeted request is in the batch. - logitsprocs_need_output_token_ids=bool(custom_logitsprocs), - is_pooling_model=self.is_pooling_model, - cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size, - reasoning_config=self.vllm_config.reasoning_config, - use_replayssm=self.cache_config.use_replayssm, - ) + # Capture warmup providers registered by the initial placeholder InputBatch + with self.jit_warmup_registry.activate(): + self.input_batch = InputBatch( + max_num_reqs=self.max_num_reqs, + # We need to use the encoder length for encoder-decoder + # because of KV cache for cross-attention. + max_model_len=max(self.max_model_len, self.max_encoder_len), + max_num_batched_tokens=self.max_num_tokens, + device=self.device, + vocab_size=self.model_config.get_vocab_size(), + block_sizes=[placeholder_block_size], + kernel_block_sizes=[placeholder_block_size], + max_num_blocks_per_req=[placeholder_max_num_blocks], + num_spec_tokens=self.num_spec_tokens, + logitsprocs=build_logitsprocs( + self.vllm_config, + self.device, + PIN_MEMORY, + self.is_pooling_model, + custom_logitsprocs, + ), + # We currently don't know whether a particular custom logits processor + # uses output token ids so we set this conservatively. Thinking-budget + # tracking is requested dynamically when a budgeted request is in the + # batch. + logitsprocs_need_output_token_ids=bool(custom_logitsprocs), + is_pooling_model=self.is_pooling_model, + cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size, + reasoning_config=self.vllm_config.reasoning_config, + use_replayssm=self.cache_config.use_replayssm, + ) # Separate cuda stream for overlapping transfer of sampled token ids from # GPU to CPU when async scheduling is enabled. @@ -7416,24 +7421,26 @@ def may_reinitialize_input_batch( self._init_kernel_block_sizes = kernel_block_sizes self._init_max_num_blocks = max_num_blocks self._init_slot_mapping_modes = slot_mapping_modes - self.input_batch = InputBatch( - max_num_reqs=self.max_num_reqs, - max_model_len=max_model_len, - max_num_batched_tokens=self.max_num_tokens, - device=self.device, - vocab_size=self.model_config.get_vocab_size(), - block_sizes=block_sizes, - kernel_block_sizes=kernel_block_sizes, - max_num_blocks_per_req=max_num_blocks, - num_spec_tokens=self.num_spec_tokens, - logitsprocs=self.input_batch.logitsprocs, - logitsprocs_need_output_token_ids=self.input_batch.logitsprocs_need_output_token_ids, - is_pooling_model=self.is_pooling_model, - cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size, - reasoning_config=self.vllm_config.reasoning_config, - use_replayssm=self.cache_config.use_replayssm, - slot_mapping_modes=slot_mapping_modes, - ) + # Capture warmup providers registered after final KV-cache geometry is known + with self.jit_warmup_registry.activate(): + self.input_batch = InputBatch( + max_num_reqs=self.max_num_reqs, + max_model_len=max_model_len, + max_num_batched_tokens=self.max_num_tokens, + device=self.device, + vocab_size=self.model_config.get_vocab_size(), + block_sizes=block_sizes, + kernel_block_sizes=kernel_block_sizes, + max_num_blocks_per_req=max_num_blocks, + num_spec_tokens=self.num_spec_tokens, + logitsprocs=self.input_batch.logitsprocs, + logitsprocs_need_output_token_ids=self.input_batch.logitsprocs_need_output_token_ids, + is_pooling_model=self.is_pooling_model, + cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size, + reasoning_config=self.vllm_config.reasoning_config, + use_replayssm=self.cache_config.use_replayssm, + slot_mapping_modes=slot_mapping_modes, + ) assert self._init_block_sizes == block_sizes, ( f"InputBatch block_sizes {self._init_block_sizes} != " From cfbc5afbf7e56d825a3de98fc38ceb4be48d820e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=84=8D=F0=9D=95=A0=F0=9D=95=9D=F0=9D=95=9D=F0=9D=95=A0?= =?UTF-8?q?=F0=9D=95=A8=20=F0=9D=95=84=F0=9D=95=92=F0=9D=95=9F?= Date: Mon, 17 Aug 2026 20:40:51 +0300 Subject: [PATCH 051/839] [BugFix] lora_base_layer / routed_experts order in expert param mapping (#52552) Signed-off-by: Hollow Man --- vllm/model_executor/layers/fused_moe/routed_experts.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index 87fc3b1c3653..b5557f4b59e9 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -1060,8 +1060,8 @@ def build_expert_params_mapping( if routed_experts_prefix != "": routed_experts_prefix = f"{routed_experts_prefix}." - w13 = f"experts.{routed_experts_prefix}{lora_base_layer_prefix}w13_" - w2 = f"experts.{routed_experts_prefix}{lora_base_layer_prefix}w2_" + w13 = f"experts.{lora_base_layer_prefix}{routed_experts_prefix}w13_" + w2 = f"experts.{lora_base_layer_prefix}{routed_experts_prefix}w2_" fused_mapping = [] if include_fused: From ceb340e2eb5b4c8edee8d0dbd71195389ebd86b5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Juan=20P=C3=A9rez=20de=20Algaba?= <124347725+jperezdealgaba@users.noreply.github.com> Date: Mon, 17 Aug 2026 19:45:40 +0200 Subject: [PATCH 052/839] fix: prevent PyNvVideoCodec decoder slot limit bypass via ClassVar shadowing (#52126) Signed-off-by: jperezde --- tests/multimodal/test_video.py | 163 ++++++++++++++++++++++----------- vllm/multimodal/video.py | 70 ++++++++------ 2 files changed, 151 insertions(+), 82 deletions(-) diff --git a/tests/multimodal/test_video.py b/tests/multimodal/test_video.py index 8e27a5c58f88..17d29f850ea4 100644 --- a/tests/multimodal/test_video.py +++ b/tests/multimodal/test_video.py @@ -24,12 +24,14 @@ Molmo2VideoBackend, PyNvVideoCodecDecoderSlot, PyNvVideoCodecVideoBackend, + PyNvVideoCodecVideoBackendMixin, Qwen2VLVideoBackend, Qwen3VLVideoBackend, VideoBackend, VideoLoader, VideoSourceMetadata, VideoTargetMetadata, + _pynv_decoder_pool, get_video_loader_backend_for_processor, ) from vllm.platforms import current_platform @@ -47,6 +49,27 @@ FAKE_OUTPUT_2 = np.random.rand(NUM_FRAMES, 1280, 720, 3) +@contextmanager +def _fresh_decoder_pool(): + """Reset module-level decoder pool for isolated test runs.""" + pool = _pynv_decoder_pool + old_slots = pool.slots + old_active = pool.active + old_cond = pool.cond + old_max = pool.max_slots + pool.slots = [] + pool.active = 0 + pool.cond = threading.Condition() + pool.max_slots = None + try: + yield pool + finally: + pool.slots = old_slots + pool.active = old_active + pool.cond = old_cond + pool.max_slots = old_max + + @VIDEO_LOADER_REGISTRY.register("test_video_loader_1") class TestVideoLoader1(VideoLoader): @classmethod @@ -206,16 +229,7 @@ def test_pynvvideocodec_corrupted_videos_raise_value_error(): corrupted_video = (ASSETS_DIR / "corrupted.mp4").read_bytes() malformed_video = corrupted_video[:128] - old_slots = PyNvVideoCodecVideoBackend._decoder_slots - old_active_slots = PyNvVideoCodecVideoBackend._active_decoder_slots - old_cond = PyNvVideoCodecVideoBackend._decoder_slot_cond - old_max_slots = PyNvVideoCodecVideoBackend._max_decoder_slots - try: - PyNvVideoCodecVideoBackend._decoder_slots = [] - PyNvVideoCodecVideoBackend._active_decoder_slots = 0 - PyNvVideoCodecVideoBackend._decoder_slot_cond = threading.Condition() - PyNvVideoCodecVideoBackend._max_decoder_slots = None - + with _fresh_decoder_pool(): loader = VIDEO_LOADER_REGISTRY.load(PYNVVIDEOCODEC_VIDEO_BACKEND) with pytest.raises( ValueError, @@ -247,11 +261,6 @@ def test_pynvvideocodec_corrupted_videos_raise_value_error(): hw_decoders=1, ) assert frames.shape[0] == 1 - finally: - PyNvVideoCodecVideoBackend._decoder_slots = old_slots - PyNvVideoCodecVideoBackend._active_decoder_slots = old_active_slots - PyNvVideoCodecVideoBackend._decoder_slot_cond = old_cond - PyNvVideoCodecVideoBackend._max_decoder_slots = old_max_slots @pytest.mark.parametrize("hw_decoders", [1, 3]) @@ -263,15 +272,7 @@ class FakeSlot: pass create_count = 0 - old_slots = PyNvVideoCodecVideoBackend._decoder_slots - old_active_slots = PyNvVideoCodecVideoBackend._active_decoder_slots - old_cond = PyNvVideoCodecVideoBackend._decoder_slot_cond - old_max_slots = PyNvVideoCodecVideoBackend._max_decoder_slots - try: - PyNvVideoCodecVideoBackend._decoder_slots = [] - PyNvVideoCodecVideoBackend._active_decoder_slots = 0 - PyNvVideoCodecVideoBackend._decoder_slot_cond = threading.Condition() - PyNvVideoCodecVideoBackend._max_decoder_slots = None + with _fresh_decoder_pool(): PyNvVideoCodecVideoBackend._configure_decoder_slots(hw_decoders) def fake_create_slot(cls): @@ -309,17 +310,12 @@ def borrow_extra_slot(): assert seen_slots[0] in retained_slots assert create_count == hw_decoders - finally: - PyNvVideoCodecVideoBackend._decoder_slots = old_slots - PyNvVideoCodecVideoBackend._active_decoder_slots = old_active_slots - PyNvVideoCodecVideoBackend._decoder_slot_cond = old_cond - PyNvVideoCodecVideoBackend._max_decoder_slots = old_max_slots def test_pynvvideocodec_decoder_slots_are_configured_once( monkeypatch: pytest.MonkeyPatch, ): - monkeypatch.setattr(PyNvVideoCodecVideoBackend, "_max_decoder_slots", None) + monkeypatch.setattr(_pynv_decoder_pool, "max_slots", None) PyNvVideoCodecVideoBackend._configure_decoder_slots(2) PyNvVideoCodecVideoBackend._configure_decoder_slots(2) @@ -358,15 +354,16 @@ def SimpleDecoder(file_path: str, **kwargs): assert slot.source_path is None raise RuntimeError("construct failed") - old_slots = PyNvVideoCodecVideoBackend._decoder_slots - old_active_slots = PyNvVideoCodecVideoBackend._active_decoder_slots - old_cond = PyNvVideoCodecVideoBackend._decoder_slot_cond - old_max_slots = PyNvVideoCodecVideoBackend._max_decoder_slots + pool = _pynv_decoder_pool + old_slots = pool.slots + old_active = pool.active + old_cond = pool.cond + old_max = pool.max_slots try: - PyNvVideoCodecVideoBackend._decoder_slots = [slot] - PyNvVideoCodecVideoBackend._active_decoder_slots = 1 - PyNvVideoCodecVideoBackend._decoder_slot_cond = threading.Condition() - PyNvVideoCodecVideoBackend._max_decoder_slots = 1 + pool.slots = [slot] + pool.active = 1 + pool.cond = threading.Condition() + pool.max_slots = 1 with ( pytest.raises(RuntimeError, match="construct failed"), @@ -386,12 +383,12 @@ def SimpleDecoder(file_path: str, **kwargs): assert old_decoder.poisoned assert slot.decoder is None assert slot.source_path is None - assert PyNvVideoCodecVideoBackend._decoder_slots == [slot] + assert pool.slots == [slot] finally: - PyNvVideoCodecVideoBackend._decoder_slots = old_slots - PyNvVideoCodecVideoBackend._active_decoder_slots = old_active_slots - PyNvVideoCodecVideoBackend._decoder_slot_cond = old_cond - PyNvVideoCodecVideoBackend._max_decoder_slots = old_max_slots + pool.slots = old_slots + pool.active = old_active + pool.cond = old_cond + pool.max_slots = old_max @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") @@ -405,15 +402,15 @@ def test_pynvvideocodec_h200_recovers_after_unsupported_8k(): valid_video = create_long_gop_video(num_frames=2, width=64, height=64) unsupported_video = (ASSETS_DIR / "unsupported_8k_h264.mp4").read_bytes() - old_slots = PyNvVideoCodecVideoBackend._decoder_slots - old_active_slots = PyNvVideoCodecVideoBackend._active_decoder_slots - old_cond = PyNvVideoCodecVideoBackend._decoder_slot_cond - old_max_slots = PyNvVideoCodecVideoBackend._max_decoder_slots + old_slots = _pynv_decoder_pool.slots + old_active = _pynv_decoder_pool.active + old_cond = _pynv_decoder_pool.cond + old_max = _pynv_decoder_pool.max_slots try: - PyNvVideoCodecVideoBackend._decoder_slots = [] - PyNvVideoCodecVideoBackend._active_decoder_slots = 0 - PyNvVideoCodecVideoBackend._decoder_slot_cond = threading.Condition() - PyNvVideoCodecVideoBackend._max_decoder_slots = None + _pynv_decoder_pool.slots = [] + _pynv_decoder_pool.active = 0 + _pynv_decoder_pool.cond = threading.Condition() + _pynv_decoder_pool.max_slots = None loader = VIDEO_LOADER_REGISTRY.load(PYNVVIDEOCODEC_VIDEO_BACKEND) frames_before, _ = loader.load_bytes( @@ -443,12 +440,66 @@ def test_pynvvideocodec_h200_recovers_after_unsupported_8k(): assert frames_after.shape == frames_before.shape finally: - for slot in PyNvVideoCodecVideoBackend._decoder_slots: + for slot in _pynv_decoder_pool.slots: slot.invalidate() - PyNvVideoCodecVideoBackend._decoder_slots = old_slots - PyNvVideoCodecVideoBackend._active_decoder_slots = old_active_slots - PyNvVideoCodecVideoBackend._decoder_slot_cond = old_cond - PyNvVideoCodecVideoBackend._max_decoder_slots = old_max_slots + _pynv_decoder_pool.slots = old_slots + _pynv_decoder_pool.active = old_active + _pynv_decoder_pool.cond = old_cond + _pynv_decoder_pool.max_slots = old_max + + +def test_pynvvideocodec_cross_subclass_shares_single_pool(): + """Regression test for GHSA-j682-9xp5-rrf3. + + Multiple subclasses of PyNvVideoCodecVideoBackendMixin must share the + same process-wide decoder slot limit rather than getting independent + counters via ClassVar shadowing. + """ + + class FakeSlot: + pass + + create_count = 0 + + def fake_create_slot(cls): + nonlocal create_count + create_count += 1 + return FakeSlot() + + with _fresh_decoder_pool() as pool: + pool.max_slots = 2 + + orig_create = PyNvVideoCodecVideoBackendMixin._create_decoder_slot + PyNvVideoCodecVideoBackendMixin._create_decoder_slot = classmethod( + fake_create_slot + ) + try: + with ExitStack() as stack: + stack.enter_context(VideoBackend._borrow_decoder_slot()) + stack.enter_context(Qwen3VLVideoBackend._borrow_decoder_slot()) + assert pool.active == 2 + + blocked = threading.Event() + acquired = threading.Event() + + def try_borrow(): + blocked.set() + with Qwen2VLVideoBackend._borrow_decoder_slot(): + acquired.set() + + t = threading.Thread(target=try_borrow) + t.start() + blocked.wait(timeout=2.0) + assert not acquired.wait(timeout=0.3) + + assert acquired.wait(timeout=2.0) + t.join(timeout=2.0) + assert not t.is_alive() + + assert create_count == 2 + assert len(pool.slots) == 2 + finally: + PyNvVideoCodecVideoBackendMixin._create_decoder_slot = orig_create @pytest.mark.parametrize("hw_decoders", [0, -1, 1.5, True, "2"]) diff --git a/vllm/multimodal/video.py b/vllm/multimodal/video.py index 8d4d7090eded..c50ce7761438 100644 --- a/vllm/multimodal/video.py +++ b/vllm/multimodal/video.py @@ -671,13 +671,38 @@ def _pynvvc_frames_to_nhwc(frames: torch.Tensor) -> torch.Tensor: return frames.contiguous() +class _PyNvDecoderPool: + """Process-wide singleton managing PyNvVideoCodec decoder slot state. + + Prevents subclass counter shadowing (GHSA-j682-9xp5-rrf3) by storing + all mutable pool state in a single module-level instance rather than + in ClassVar attributes that get shadowed by Python's augmented + assignment semantics on subclasses. + """ + + def __init__(self) -> None: + self.slots: list[PyNvVideoCodecDecoderSlot] = [] + self.active: int = 0 + self.cond: threading.Condition = threading.Condition() + self.max_slots: int | None = None + + def configure(self, hw_decoders: int) -> None: + with self.cond: + if self.max_slots is None: + self.max_slots = hw_decoders + elif self.max_slots != hw_decoders: + raise RuntimeError( + "PyNvVideoCodec decoder count is already configured as " + f"{self.max_slots}, got {hw_decoders}" + ) + + +_pynv_decoder_pool = _PyNvDecoderPool() + + class PyNvVideoCodecVideoBackendMixin: """PyNvVideoCodec utilities for GPU-backed frame decode.""" - _decoder_slots: ClassVar[list[PyNvVideoCodecDecoderSlot]] = [] - _active_decoder_slots: ClassVar[int] = 0 - _decoder_slot_cond: ClassVar[threading.Condition] = threading.Condition() - _max_decoder_slots: ClassVar[int | None] = None _DEVICE_INDEX: ClassVar[int] = 0 @classmethod @@ -704,14 +729,7 @@ def _create_decoder_slot(cls) -> PyNvVideoCodecDecoderSlot: @classmethod def _configure_decoder_slots(cls, hw_decoders: object) -> None: hw_decoders = validate_pynvvideocodec_hw_decoders(hw_decoders) - with cls._decoder_slot_cond: - if cls._max_decoder_slots is None: - cls._max_decoder_slots = hw_decoders - elif cls._max_decoder_slots != hw_decoders: - raise RuntimeError( - "PyNvVideoCodec decoder count is already configured as " - f"{cls._max_decoder_slots}, got {hw_decoders}" - ) + _pynv_decoder_pool.configure(hw_decoders) @staticmethod @contextmanager @@ -729,28 +747,28 @@ def _torch_stream_context(stream): @classmethod @contextmanager def _borrow_decoder_slot(cls): + pool = _pynv_decoder_pool create_slot = False - with cls._decoder_slot_cond: - max_decoder_slots = cls._max_decoder_slots - if max_decoder_slots is None: + with pool.cond: + if pool.max_slots is None: raise RuntimeError("PyNvVideoCodec decoder slots are not configured") while True: - if cls._decoder_slots: - slot = cls._decoder_slots.pop() + if pool.slots: + slot = pool.slots.pop() break - if cls._active_decoder_slots < max_decoder_slots: - cls._active_decoder_slots += 1 + if pool.active < pool.max_slots: + pool.active += 1 create_slot = True break - cls._decoder_slot_cond.wait() + pool.cond.wait() if create_slot: try: slot = cls._create_decoder_slot() except Exception: - with cls._decoder_slot_cond: - cls._active_decoder_slots -= 1 - cls._decoder_slot_cond.notify() + with pool.cond: + pool.active -= 1 + pool.cond.notify() raise borrow_succeeded = False @@ -760,9 +778,9 @@ def _borrow_decoder_slot(cls): finally: if not borrow_succeeded: slot.invalidate() - with cls._decoder_slot_cond: - cls._decoder_slots.append(slot) - cls._decoder_slot_cond.notify() + with pool.cond: + pool.slots.append(slot) + pool.cond.notify() @staticmethod def _metadata_value(metadata, *names: str, default=None): From 402547d7f02bdbfc5dce5d27dc21f50dd4d627b6 Mon Sep 17 00:00:00 2001 From: stefankoncarevic Date: Mon, 17 Aug 2026 19:50:26 +0200 Subject: [PATCH 053/839] [Bugfix][CI] Release the shared ColBERT engine before `test_colbert_hf_comparison` (#52608) Signed-off-by: Stefan Koncarevic Co-authored-by: Andreas Karatzas --- tests/models/language/pooling/test_colbert.py | 153 ++++++++---------- 1 file changed, 71 insertions(+), 82 deletions(-) diff --git a/tests/models/language/pooling/test_colbert.py b/tests/models/language/pooling/test_colbert.py index 975c4d6e8fb0..bb7afb23365b 100644 --- a/tests/models/language/pooling/test_colbert.py +++ b/tests/models/language/pooling/test_colbert.py @@ -186,33 +186,33 @@ def _assert_embeddings_close(vllm_outputs, hf_embeddings): ) -@pytest.fixture(params=list(COLBERT_MODELS.keys()), scope="module") +@pytest.fixture(params=list(COLBERT_MODELS.keys()), scope="class") def colbert_spec(request): """Return the model spec dict for the current parametrization.""" return COLBERT_MODELS[request.param] -@pytest.fixture(scope="module") +@pytest.fixture(scope="class") def colbert_model_name(colbert_spec): return colbert_spec["model"] -@pytest.fixture(scope="module") +@pytest.fixture(scope="class") def colbert_dim(colbert_spec): return colbert_spec["colbert_dim"] -@pytest.fixture(scope="module") +@pytest.fixture(scope="class") def colbert_max_model_len(colbert_spec): return colbert_spec["max_model_len"] -@pytest.fixture(scope="module") +@pytest.fixture(scope="class") def colbert_extra_kwargs(colbert_spec): return colbert_spec["extra_kwargs"] -@pytest.fixture(scope="module") +@pytest.fixture(scope="class") def colbert_model( vllm_runner, colbert_model_name, @@ -230,103 +230,92 @@ def colbert_model( yield vllm_model -def test_colbert_token_embed( - colbert_model, - colbert_dim, -): - """Test that ColBERT model produces token embeddings.""" - outputs = colbert_model.token_embed([TEXTS_1[0]]) - - assert len(outputs) == 1 - emb = torch.as_tensor(outputs[0]) - assert emb.dim() == 2 - assert emb.shape[1] == colbert_dim - assert emb.shape[0] > 1 - - -def test_colbert_late_interaction_1_to_1( - colbert_model, -): - """Test ColBERT late interaction scoring with 1:1 query-document pair.""" - q_outputs = colbert_model.token_embed([TEXTS_1[0]]) - d_outputs = colbert_model.token_embed([TEXTS_2[0]]) +class TestColbertSharedEngine: + """Tests sharing one engine per model. - q_emb = torch.as_tensor(q_outputs[0]) - d_emb = torch.as_tensor(d_outputs[0]) + Class-scoped so the engine is released before `test_colbert_hf_comparison`, + which needs a runner of its own and cannot start while this one holds VRAM. + """ - manual_score = compute_maxsim_score(q_emb, d_emb).item() + def test_colbert_token_embed(self, colbert_model, colbert_dim): + """Test that ColBERT model produces token embeddings.""" + outputs = colbert_model.token_embed([TEXTS_1[0]]) - vllm_scores = colbert_model.score(TEXTS_1[0], TEXTS_2[0]) + assert len(outputs) == 1 + emb = torch.as_tensor(outputs[0]) + assert emb.dim() == 2 + assert emb.shape[1] == colbert_dim + assert emb.shape[0] > 1 - assert len(vllm_scores) == 1 - assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01) + def test_colbert_late_interaction_1_to_1(self, colbert_model): + """Test ColBERT late interaction scoring with 1:1 query-document pair.""" + q_outputs = colbert_model.token_embed([TEXTS_1[0]]) + d_outputs = colbert_model.token_embed([TEXTS_2[0]]) + q_emb = torch.as_tensor(q_outputs[0]) + d_emb = torch.as_tensor(d_outputs[0]) -def test_colbert_late_interaction_1_to_N( - colbert_model, -): - """Test ColBERT late interaction scoring with 1:N query-documents.""" - q_outputs = colbert_model.token_embed([TEXTS_1[0]]) - d_outputs = colbert_model.token_embed(TEXTS_2) + manual_score = compute_maxsim_score(q_emb, d_emb).item() - q_emb = torch.as_tensor(q_outputs[0]) + vllm_scores = colbert_model.score(TEXTS_1[0], TEXTS_2[0]) - manual_scores = [] - for d_out in d_outputs: - d_emb = torch.as_tensor(d_out) - manual_scores.append(compute_maxsim_score(q_emb, d_emb).item()) + assert len(vllm_scores) == 1 + assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01) - vllm_scores = colbert_model.score(TEXTS_1[0], TEXTS_2) + def test_colbert_late_interaction_1_to_N(self, colbert_model): + """Test ColBERT late interaction scoring with 1:N query-documents.""" + q_outputs = colbert_model.token_embed([TEXTS_1[0]]) + d_outputs = colbert_model.token_embed(TEXTS_2) - assert len(vllm_scores) == 2 - for i in range(2): - assert vllm_scores[i] == pytest.approx(manual_scores[i], rel=0.01) + q_emb = torch.as_tensor(q_outputs[0]) + manual_scores = [] + for d_out in d_outputs: + d_emb = torch.as_tensor(d_out) + manual_scores.append(compute_maxsim_score(q_emb, d_emb).item()) -def test_colbert_late_interaction_N_to_N( - colbert_model, -): - """Test ColBERT late interaction scoring with N:N query-documents.""" - q_outputs = colbert_model.token_embed(TEXTS_1) - d_outputs = colbert_model.token_embed(TEXTS_2) + vllm_scores = colbert_model.score(TEXTS_1[0], TEXTS_2) - manual_scores = [] - for q_out, d_out in zip(q_outputs, d_outputs): - q_emb = torch.as_tensor(q_out) - d_emb = torch.as_tensor(d_out) - manual_scores.append(compute_maxsim_score(q_emb, d_emb).item()) + assert len(vllm_scores) == 2 + for i in range(2): + assert vllm_scores[i] == pytest.approx(manual_scores[i], rel=0.01) - vllm_scores = colbert_model.score(TEXTS_1, TEXTS_2) + def test_colbert_late_interaction_N_to_N(self, colbert_model): + """Test ColBERT late interaction scoring with N:N query-documents.""" + q_outputs = colbert_model.token_embed(TEXTS_1) + d_outputs = colbert_model.token_embed(TEXTS_2) - assert len(vllm_scores) == 2 - for i in range(2): - assert vllm_scores[i] == pytest.approx(manual_scores[i], rel=0.01) + manual_scores = [] + for q_out, d_out in zip(q_outputs, d_outputs): + q_emb = torch.as_tensor(q_out) + d_emb = torch.as_tensor(d_out) + manual_scores.append(compute_maxsim_score(q_emb, d_emb).item()) + vllm_scores = colbert_model.score(TEXTS_1, TEXTS_2) -def test_colbert_relevance_ordering( - colbert_model, -): - """Test that ColBERT scores relevant documents higher than irrelevant.""" - query = "What is machine learning?" - documents = [ - "Machine learning is a subset of artificial intelligence.", - "Python is a programming language.", - "Deep learning uses neural networks.", - ] + assert len(vllm_scores) == 2 + for i in range(2): + assert vllm_scores[i] == pytest.approx(manual_scores[i], rel=0.01) - scores = colbert_model.score(query, documents) + def test_colbert_relevance_ordering(self, colbert_model): + """Test that ColBERT scores relevant documents higher than irrelevant.""" + query = "What is machine learning?" + documents = [ + "Machine learning is a subset of artificial intelligence.", + "Python is a programming language.", + "Deep learning uses neural networks.", + ] - assert len(scores) == 3 - assert scores[0] > scores[1], "ML doc should score higher than Python doc" - assert scores[2] > scores[1], "DL doc should score higher than Python doc" + scores = colbert_model.score(query, documents) + assert len(scores) == 3 + assert scores[0] > scores[1], "ML doc should score higher than Python doc" + assert scores[2] > scores[1], "DL doc should score higher than Python doc" -def test_colbert_embed_not_supported( - colbert_model, -): - """Test that ColBERT model does not support the embed task.""" - with pytest.raises(ValueError, match="Embedding API is not supported"): - colbert_model.embed([TEXTS_1[0]]) + def test_colbert_embed_not_supported(self, colbert_model): + """Test that ColBERT model does not support the embed task.""" + with pytest.raises(ValueError, match="Embedding API is not supported"): + colbert_model.embed([TEXTS_1[0]]) @pytest.mark.parametrize( From 75dde08d3fd0cf4217c90e8499881d60fe426aa0 Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Mon, 17 Aug 2026 14:21:32 -0400 Subject: [PATCH 054/839] [Perf][MoE] Optimize deepep_v2 receiver CPU Overhead (#51114) Signed-off-by: Lucas Wilkinson Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> --- .../fused_moe/prepare_finalize/deepep_v2.py | 41 +++++++++++-------- 1 file changed, 25 insertions(+), 16 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py index 129e3b5d5c27..5e950fdf606e 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py @@ -80,9 +80,23 @@ def __init__( # DBO microbatching: one handle slot per micro-batch. self.handles: list[deep_ep.EPHandle | None] = [None, None] + # arange(num_local_experts) + rank_expert_offset. Rank-constant, so it + # is built once per device instead of once per layer per step. + self._global_expert_ids_cache: torch.Tensor | None = None + def num_dispatchers(self) -> int: return self.num_dispatchers_ + def _global_expert_ids(self, num_local: int, device: torch.device) -> torch.Tensor: + ids = self._global_expert_ids_cache + if ids is None or ids.numel() != num_local or ids.device != device: + ids = ( + torch.arange(num_local, dtype=torch.int64, device=device) + + self.rank_expert_offset + ) + self._global_expert_ids_cache = ids + return ids + def output_is_reduced(self) -> bool: return True @@ -196,23 +210,23 @@ def _receiver( else: expert_x, expert_x_scale = recv_x, None + expert_tokens_meta = mk.ExpertTokensMetadata.make_from_list( + recv_expert_num_tokens, + device=expert_x.device, + ) + if recv_topk_idx is None: # do_expand=True (prefill mode): build topk_ids from # per-expert token counts. total_tokens = sum(recv_expert_num_tokens) if total_tokens > 0: - recv_topk_idx = torch.empty( - total_tokens, - dtype=torch.int64, - device=expert_x.device, + recv_topk_idx = torch.repeat_interleave( + self._global_expert_ids( + len(recv_expert_num_tokens), expert_x.device + ), + expert_tokens_meta.expert_num_tokens, + output_size=total_tokens, ) - offset = 0 - for i, count in enumerate(recv_expert_num_tokens): - if count > 0: - recv_topk_idx[offset : offset + count].fill_( - i + self.rank_expert_offset - ) - offset += count else: recv_topk_idx = torch.empty( 0, @@ -243,11 +257,6 @@ def _receiver( if recv_topk_weights is not None and recv_topk_weights.ndim == 1: recv_topk_weights = recv_topk_weights.unsqueeze(1) - expert_tokens_meta = mk.ExpertTokensMetadata.make_from_list( - recv_expert_num_tokens, - device=expert_x.device, - ) - if not quant_config.is_block_quantized and not defer_input_quant: expert_x_scale = None if expert_x.numel() != 0: From c1e438728c55281fda46c8baed755f4148c63660 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Mon, 17 Aug 2026 13:32:04 -0500 Subject: [PATCH 055/839] [ROCm][CI] Restore Torch defaults and type DSV4 scratch buffers (#52566) Signed-off-by: Andreas Karatzas Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/kernels/attention/test_mha_attn.py | 12 ++++++++++-- .../attention/test_rocm_triton_attn_dsv4.py | 14 ++++++++++++-- 2 files changed, 22 insertions(+), 4 deletions(-) diff --git a/tests/kernels/attention/test_mha_attn.py b/tests/kernels/attention/test_mha_attn.py index d73acfc0ee9c..b6578fa2f915 100644 --- a/tests/kernels/attention/test_mha_attn.py +++ b/tests/kernels/attention/test_mha_attn.py @@ -27,9 +27,17 @@ @pytest.fixture(autouse=True) -def clear_cache(): - """Clear lru cache to ensure each test case runs without caching.""" +def reset_test_state(): + """Clear cached selectors and restore process-wide torch defaults.""" + default_device = torch.get_default_device() + default_dtype = torch.get_default_dtype() _cached_get_attn_backend.cache_clear() + try: + yield + finally: + torch.set_default_device(default_device) + torch.set_default_dtype(default_dtype) + _cached_get_attn_backend.cache_clear() devices = ["cpu"] diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index 7a98f3ec7005..2f1e8d279356 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -730,9 +730,19 @@ def test_sparse_attn_decode_split_k_kernel( @torch.inference_mode() def test_sparse_attn_decode_gfx950_adaptive_reduce_ignores_stale_scratch() -> None: device = torch.device("cuda") - part_m = torch.full((1, 8, 1), torch.finfo(torch.float32).min, device=device) + part_m = torch.full( + (1, 8, 1), + torch.finfo(torch.float32).min, + dtype=torch.float32, + device=device, + ) part_l = torch.zeros_like(part_m) - part_acc = torch.full((1, 8, 1, HEAD_DIM), float("nan"), device=device) + part_acc = torch.full( + (1, 8, 1, HEAD_DIM), + float("nan"), + dtype=torch.float32, + device=device, + ) part_m[:, :2] = 0 part_l[:, :2] = 1 part_acc[:, 0] = 1 From 3fc28939099a3bad1bcd27e8eec20d9afe921879 Mon Sep 17 00:00:00 2001 From: crZhao Date: Mon, 17 Aug 2026 12:10:34 -0700 Subject: [PATCH 056/839] [Bugfix] Account for local DP workers in startup thread allocation (#52385) Signed-off-by: real-cpu --- tests/distributed/test_multiproc_executor.py | 47 ++++++++++++++++++++ vllm/v1/executor/multiproc_executor.py | 5 ++- 2 files changed, 51 insertions(+), 1 deletion(-) diff --git a/tests/distributed/test_multiproc_executor.py b/tests/distributed/test_multiproc_executor.py index 042349e4fbb1..599025d1c8ec 100644 --- a/tests/distributed/test_multiproc_executor.py +++ b/tests/distributed/test_multiproc_executor.py @@ -11,15 +11,62 @@ import os import socket +import pytest + from tests.utils import multi_gpu_test from vllm.config import VllmConfig from vllm.engine.arg_utils import EngineArgs from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.executor import multiproc_executor from vllm.v1.executor.multiproc_executor import MultiprocExecutor MODEL = "facebook/opt-125m" +@pytest.mark.parametrize( + ("local_world_size", "data_parallel_size_local", "expected_num_local_procs"), + [(1, 4, 4), (4, 1, 4), (2, 0, 2)], +) +def test_multiproc_executor_counts_all_local_dp_workers( + monkeypatch: pytest.MonkeyPatch, + local_world_size: int, + data_parallel_size_local: int, + expected_num_local_procs: int, +): + """All colocated DP workers share the node's startup CPU budget.""" + executor = object.__new__(MultiprocExecutor) + executor.world_size = local_world_size + executor.local_world_size = local_world_size + executor.parallel_config = type( + "ParallelConfig", + (), + {"data_parallel_size_local": data_parallel_size_local}, + )() + + monkeypatch.setattr( + executor, + "_get_parallel_sizes", + lambda: (local_world_size, 1, 1), + ) + + class StopExecutorInit(Exception): + pass + + def capture_num_local_procs(num_local_procs: int): + assert num_local_procs == expected_num_local_procs + raise StopExecutorInit + + monkeypatch.setattr( + multiproc_executor, + "set_multiprocessing_worker_envs", + capture_num_local_procs, + ) + + with pytest.raises(StopExecutorInit): + executor._init_executor() + executor._finalizer.detach() + + def create_vllm_config( tensor_parallel_size: int = 1, pipeline_parallel_size: int = 1, diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py index df20a90d51d3..8f0f638deb88 100644 --- a/vllm/v1/executor/multiproc_executor.py +++ b/vllm/v1/executor/multiproc_executor.py @@ -130,7 +130,10 @@ def _init_executor(self) -> None: f"_parallel_size ({pcp_size}). " ) - set_multiprocessing_worker_envs(self.local_world_size) + num_local_procs = self.local_world_size * max( + 1, self.parallel_config.data_parallel_size_local + ) + set_multiprocessing_worker_envs(num_local_procs) if aiter_requires_tcp_store(): distributed_init_method = get_distributed_init_method( From 9633933dd81228fbcae07969f20881ad0b7cb766 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Mon, 17 Aug 2026 12:41:09 -0700 Subject: [PATCH 057/839] Relax CuPy constraint to only exclude 14.1.0 (#44284) Signed-off-by: khluu Signed-off-by: Nick Hill Co-authored-by: khluu Co-authored-by: Nick Hill --- requirements/kv_connectors.txt | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/requirements/kv_connectors.txt b/requirements/kv_connectors.txt index d7a396bfa0ca..1760ed651734 100644 --- a/requirements/kv_connectors.txt +++ b/requirements/kv_connectors.txt @@ -1,7 +1,7 @@ lmcache >= 0.3.9 -# CuPy 14.1.0 imports pytest from cupy.testing._random. Use <14.1.0 -# until a fixed newer release is verified for runtime images. -cupy-cuda13x < 14.1.0 +# CuPy 14.1.0 imports pytest from cupy.testing._random, which breaks runtime +# images. 14.1.1 fixes the root cause, so only 14.1.0 is excluded. +cupy-cuda13x != 14.1.0 nixl == 1.3.2 # CUDA 12 build. On the CUDA 13 image install-kv-connectors.sh swaps this for # the mooncake-transfer-engine-cuda13 variant of the same version. From d1e3eee6fb8ed3623241ef5c8e3ac533f775bff9 Mon Sep 17 00:00:00 2001 From: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Date: Mon, 17 Aug 2026 16:08:45 -0400 Subject: [PATCH 058/839] [Spec decode] Support Kimi-K3 DCP with DSpark (#52188) Signed-off-by: wzhao18 --- .../test_dspark_mla_config.py | 23 ---- tests/v1/attention/test_flashinfer_mla_dcp.py | 1 + tests/v1/attention/test_mla_backends.py | 77 +++++++++++ .../spec_decode/test_dflash_prepare_inputs.py | 26 +++- vllm/config/speculative.py | 10 -- .../layers/attention/mla_attention.py | 29 +++++ vllm/models/kimi_k3/nvidia/mla.py | 5 - vllm/v1/attention/backend.py | 8 ++ .../attention/backends/mla/flashinfer_mla.py | 122 +++++++++++++++++- .../attention/backends/mla/tokenspeed_mla.py | 1 + vllm/v1/attention/selector.py | 5 +- vllm/v1/worker/gpu/cp_utils.py | 21 +++ .../gpu/spec_decode/dflash/cudagraph.py | 13 ++ .../gpu/spec_decode/dflash/speculator.py | 45 ++++++- vllm/v1/worker/gpu/spec_decode/speculator.py | 6 + 15 files changed, 341 insertions(+), 51 deletions(-) diff --git a/tests/transformers_utils/test_dspark_mla_config.py b/tests/transformers_utils/test_dspark_mla_config.py index f43d668b73cd..b81f06e14388 100644 --- a/tests/transformers_utils/test_dspark_mla_config.py +++ b/tests/transformers_utils/test_dspark_mla_config.py @@ -140,26 +140,3 @@ def test_dspark_mla_speculative_config_preserves_architecture(tmp_path): assert speculative_config.draft_model_config.architectures == ["K3DSparkModel"] assert speculative_config.draft_model_config.hf_config.model_type == "k3_dspark" assert speculative_config.draft_model_config.use_mla - - -def test_dspark_mla_rejects_decode_context_parallelism(tmp_path): - target_path = tmp_path / "target" - draft_path = tmp_path / "draft" - _write_target_config(target_path) - _write_dspark_config(draft_path) - target_config = ModelConfig( - model=str(target_path), tokenizer_mode="skip", max_model_len=32768 - ) - - with pytest.raises(ValueError, match="does not currently support decode context"): - SpeculativeConfig( - model=str(draft_path), - method="dspark", - num_speculative_tokens=8, - target_model_config=target_config, - target_parallel_config=ParallelConfig( - tensor_parallel_size=2, - decode_context_parallel_size=2, - distributed_executor_backend="external_launcher", - ), - ) diff --git a/tests/v1/attention/test_flashinfer_mla_dcp.py b/tests/v1/attention/test_flashinfer_mla_dcp.py index acc38919b802..f3a46bbb1900 100644 --- a/tests/v1/attention/test_flashinfer_mla_dcp.py +++ b/tests/v1/attention/test_flashinfer_mla_dcp.py @@ -45,6 +45,7 @@ def test_flashinfer_mla_forward_uses_gathered_head_count(monkeypatch): impl.bmm1_scale = 1.0 impl.bmm2_scale = 1.0 impl.need_to_return_lse_for_decode = True + impl.dcp_world_size = 2 impl.num_heads = 6 impl.qk_nope_head_dim = 128 impl.kv_lora_rank = 512 diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index cd291937241a..72773adcb674 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -837,9 +837,86 @@ class _AttnMeta: def test_tokenspeed_mla_noncausal_capability(): builder = tokenspeed_mla_module.TokenspeedMLAMetadataBuilder assert builder.supports_non_causal_multi_token_decode + assert builder.supports_non_causal_multi_token_dcp assert tokenspeed_mla_module.TokenspeedMLABackend.supports_non_causal() +def test_flashinfer_mla_dcp_multi_token_decode_uses_per_query_bounds(monkeypatch): + flashinfer_mla_module = pytest.importorskip( + "vllm.v1.attention.backends.mla.flashinfer_mla" + ) + + decode_call = None + + def fake_decode(**kwargs): + nonlocal decode_call + decode_call = kwargs + query = kwargs["query"] + output = torch.empty(*query.shape[:-1], 512, dtype=torch.bfloat16) + lse = torch.empty(query.shape[0], query.shape[-2], dtype=torch.float32) + return output, lse + + monkeypatch.setattr( + flashinfer_mla_module, + "trtllm_batch_decode_with_kv_cache_mla", + fake_decode, + ) + monkeypatch.setattr( + flashinfer_mla_module, + "_get_workspace_buffer", + lambda return_lse: torch.empty(1, dtype=torch.int8), + ) + + impl = object.__new__(flashinfer_mla_module.FlashInferMLAImpl) + impl.dcp_world_size = 2 + impl.dcp_rank = 1 + impl.cp_kv_cache_interleave_size = 1 + impl.need_to_return_lse_for_decode = True + impl.kv_lora_rank = 512 + impl.qk_nope_head_dim = 128 + impl.qk_rope_head_dim = 64 + impl.bmm1_scale = 1.0 + impl.bmm2_scale = 1.0 + + block_table = torch.tensor([[1], [2]], dtype=torch.int32) + metadata = SimpleNamespace( + num_decodes=2, + num_decode_tokens=6, + max_seq_len=7, + causal=True, + decode=SimpleNamespace( + block_table=block_table, + seq_lens=torch.tensor([5, 6], dtype=torch.int32), + dcp_tot_seq_lens=torch.tensor([10, 13], dtype=torch.int32), + flattened_block_table=None, + flattened_seq_lens=None, + query_len=0, + ), + ) + query = torch.empty(6, 2, 576, dtype=torch.bfloat16) + kv_cache = torch.empty(3, 16, 576, dtype=torch.bfloat16) + + output, lse = impl.forward_mqa( + query, + kv_cache, + metadata, + SimpleNamespace(), + ) + + assert output.shape == (6, 2, 512) + assert lse is not None + assert lse.shape == (6, 2) + assert decode_call is not None + assert decode_call["query"].shape == (6, 1, 2, 576) + torch.testing.assert_close( + decode_call["seq_lens"], + torch.tensor([4, 4, 5, 5, 6, 6], dtype=torch.int32), + ) + torch.testing.assert_close( + decode_call["block_tables"], block_table.repeat_interleave(3, dim=0) + ) + + @pytest.mark.parametrize( ("causal", "tokens_per_decode", "dcp_world_size", "dcp_rank"), [ diff --git a/tests/v1/spec_decode/test_dflash_prepare_inputs.py b/tests/v1/spec_decode/test_dflash_prepare_inputs.py index 50a414a83770..16d6d8e516df 100644 --- a/tests/v1/spec_decode/test_dflash_prepare_inputs.py +++ b/tests/v1/spec_decode/test_dflash_prepare_inputs.py @@ -17,7 +17,14 @@ ) -def _run_prepare(*, target_positions: list[int], block_table_values: list[int]): +def _run_prepare( + *, + target_positions: list[int], + block_table_values: list[int], + cp_rank: int = 0, + cp_size: int = 1, + cp_interleave: int = 1, +): device = torch.device("cuda") max_num_reqs = 4 max_num_tokens = 16 @@ -86,6 +93,9 @@ def _run_prepare(*, target_positions: list[int], block_table_values: list[int]): input_seeds, block_table, 4, + cp_rank, + cp_size, + cp_interleave, 123, num_speculative_steps, num_speculative_steps, @@ -131,6 +141,20 @@ def test_prepare_dflash_inputs_excludes_rejected_context_suffix(): assert out.seeds[2].item() == 17 +def test_prepare_dflash_inputs_excludes_rejected_context_suffix_with_dcp(): + out = _run_prepare( + target_positions=[10, 11, 12, 13], + block_table_values=[0, 7, 8, 9], + cp_rank=1, + cp_size=2, + cp_interleave=2, + ) + + assert out.context_positions[:4].tolist() == [10, 11, 0, 0] + assert out.context_slot_mapping[:4].tolist() == [28, 29, PAD_SLOT_ID, PAD_SLOT_ID] + assert out.query_slot_mapping[:3].tolist() == [PAD_SLOT_ID, PAD_SLOT_ID, 30] + + def test_prepare_dflash_inputs_never_writes_the_null_block(): # The valid context uses logical block 0 and the replacement query uses # logical block 1. Both map to the null block and must remain unwritable. diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index a85b895d8647..c3e0866b453d 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -1059,16 +1059,6 @@ def __post_init__(self): if self.method in ("dflash", "dspark"): self.parallel_drafting = True - if ( - self.method == "dspark" - and "K3DSparkModel" in self.draft_model_config.architectures - and self.target_parallel_config.decode_context_parallel_size > 1 - ): - raise ValueError( - "MLA DSpark does not currently support decode context " - "parallelism; set decode_context_parallel_size=1." - ) - if self.num_speculative_tokens is not None and hasattr( self.draft_model_config.hf_config, "num_lookahead_tokens" ): diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 839f4266cd97..72d65c15748d 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -1984,6 +1984,9 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]): # Whether this builder can flatten a non-causal query block into decode rows. supports_non_causal_multi_token_decode: ClassVar[bool] = False + # Whether can support non-causal multi-token decode with DCP KV cache. + supports_non_causal_multi_token_dcp: ClassVar[bool] = False + # The threshold for reordering the batch into decode and prefill requests. # If > 1, the batch will be reordered such that requests with # query length <= threshold are classified as decode requests. @@ -1991,6 +1994,31 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]): # when speculative decoding is enabled. reorder_batch_threshold: int = 1 + def _validate_dspark_dcp_support(self, supports_dcp_with_varlen: bool) -> None: + speculative_config = getattr(self.vllm_config, "speculative_config", None) + parallel_config = self.vllm_config.parallel_config + if ( + speculative_config is None + or getattr(speculative_config, "method", None) != "dspark" + or parallel_config.decode_context_parallel_size <= 1 + ): + return + + if self.non_causal_multi_token_decode: + supported = self.supports_non_causal_multi_token_dcp + query_mode = "non-causal draft" + else: + supported = supports_dcp_with_varlen + query_mode = "causal multi-token" + + if not supported: + raise ValueError( + f"{type(self).__name__} does not support {query_mode} MLA " + "attention for DSpark with decode context parallelism. Select " + "a backend with explicit DSpark DCP support or set " + "decode_context_parallel_size=1." + ) + @staticmethod def determine_chunked_prefill_workspace_size(vllm_config: VllmConfig) -> int: scheduler_config = vllm_config.scheduler_config @@ -2084,6 +2112,7 @@ def __init__( self.non_causal_multi_token_decode = getattr( kv_cache_spec, "non_causal_multi_token_decode", False ) + self._validate_dspark_dcp_support(supports_dcp_with_varlen) # A draft cache group can have a different head count from the target. self.num_heads = get_num_attention_heads_from_layers( diff --git a/vllm/models/kimi_k3/nvidia/mla.py b/vllm/models/kimi_k3/nvidia/mla.py index 71e9fa79f04d..2c67a8a4c4a6 100644 --- a/vllm/models/kimi_k3/nvidia/mla.py +++ b/vllm/models/kimi_k3/nvidia/mla.py @@ -341,11 +341,6 @@ def __init__( "parallelism." ) self.dcp_world_size = parallel_config.decode_context_parallel_size - assert self.dcp_world_size <= 1 or self.rotary_emb is None, ( - "Kimi-K3 MultiHeadLatentAttention does not support RoPE with decode " - "context parallelism because gathered queries require gathered " - "positions." - ) self.dcp_manager: MLADCPManager | None = None if self.dcp_world_size > 1: query_dtype = ( diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 672f9e63761e..daa217f8553d 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -330,6 +330,11 @@ def supports_pcp(cls) -> bool: except NotImplementedError: return False + @classmethod + def supports_non_causal_dcp(cls) -> bool: + builder_cls = cls.get_builder_cls() + return bool(getattr(builder_cls, "supports_non_causal_multi_token_dcp", False)) + @classmethod def supports_attn_type(cls, attn_type: str) -> bool: """Check if backend supports a given attention type. @@ -378,6 +383,7 @@ def validate_configuration( use_kv_connector: bool = False, use_pcp: bool = False, use_adaptive_verification: bool = False, + use_dcp: bool = False, ) -> list[str]: invalid_reasons = [] if not cls.supports_head_size(head_size): @@ -414,6 +420,8 @@ def validate_configuration( invalid_reasons.append("sliding window not supported") if use_non_causal and not cls.supports_non_causal(): invalid_reasons.append("non-causal attention not supported") + if use_mla and use_non_causal and use_dcp and not cls.supports_non_causal_dcp(): + invalid_reasons.append("non-causal MLA attention with DCP not supported") if use_batch_invariant and not cls.supports_batch_invariance(): invalid_reasons.append("batch invariance not supported") if use_kv_connector and not cls.supports_kv_connector(): diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla.py b/vllm/v1/attention/backends/mla/flashinfer_mla.py index bf92bf692bf0..79beb5332e2a 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import ClassVar +from dataclasses import dataclass +from typing import TYPE_CHECKING, ClassVar import torch from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla @@ -15,6 +16,7 @@ from vllm.logger import init_logger from vllm.model_executor.layers.attention.mla_attention import ( MLACommonBackend, + MLACommonDecodeMetadata, MLACommonImpl, MLACommonMetadata, MLACommonMetadataBuilder, @@ -30,6 +32,10 @@ ) from vllm.v1.attention.backends.utils import KVCacheLayoutType +if TYPE_CHECKING: + from vllm.config import VllmConfig + from vllm.v1.kv_cache_interface import AttentionSpec + logger = init_logger(__name__) @@ -112,12 +118,56 @@ def _get_multi_ctas_kv_counter_buffer( return _fi_multi_ctas_kv_counter -class FlashInferMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]): +@dataclass +class FlashInferMLADecodeMetadata(MLACommonDecodeMetadata): + flattened_block_table: torch.Tensor | None = None + flattened_seq_lens: torch.Tensor | None = None + query_len: int = 0 + + +@dataclass +class FlashInferMLAMetadata(MLACommonMetadata[FlashInferMLADecodeMetadata]): + pass + + +class FlashInferMLAMetadataBuilder(MLACommonMetadataBuilder[FlashInferMLAMetadata]): _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM # Non-causal DSpark blocks are flattened to single-token rows in forward_mqa. supports_non_causal_multi_token_decode: ClassVar[bool] = True + def __init__( + self, + kv_cache_spec: "AttentionSpec", + layer_names: list[str], + vllm_config: "VllmConfig", + device: torch.device, + ) -> None: + super().__init__( + kv_cache_spec, + layer_names, + vllm_config, + device, + FlashInferMLAMetadata, + supports_dcp_with_varlen=True, + ) + + def _build_decode( + self, + block_table_tensor: torch.Tensor, + seq_lens_device: torch.Tensor, + max_seq_len: int, + query_start_loc_cpu: torch.Tensor, + query_start_loc_device: torch.Tensor, + num_decode_tokens: int, + dcp_tot_seq_lens_device: torch.Tensor | None, + ) -> FlashInferMLADecodeMetadata: + return FlashInferMLADecodeMetadata( + block_table=block_table_tensor, + seq_lens=seq_lens_device, + dcp_tot_seq_lens=dcp_tot_seq_lens_device, + ) + class FlashInferMLABackend(MLACommonBackend): supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16] @@ -193,7 +243,7 @@ def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None": return "HND" -class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]): +class FlashInferMLAImpl(MLACommonImpl[FlashInferMLAMetadata]): can_return_lse_for_decode: bool = True # trtllm-gen MLA decode emits LSE in log2 (per flashinfer's own # reference at flashinfer/trace/templates/attention.py:81: @@ -264,7 +314,7 @@ def forward_mqa( self, q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], kv_c_and_k_pe_cache: torch.Tensor, - attn_metadata: MLACommonMetadata, + attn_metadata: FlashInferMLAMetadata, layer: AttentionLayer, ) -> tuple[torch.Tensor, torch.Tensor | None]: assert kv_c_and_k_pe_cache.numel() > 0 @@ -276,16 +326,28 @@ def forward_mqa( block_table = attn_metadata.decode.block_table seq_lens = attn_metadata.decode.seq_lens + query_len = attn_metadata.num_decode_tokens // attn_metadata.num_decodes if not attn_metadata.causal: # Non-causal DSpark block: flatten to single-token decode rows with # per-row context seq_lens (trtllm-gen has no causal flag and would # otherwise mask the block causally). - query_len = attn_metadata.num_decode_tokens // attn_metadata.num_decodes q = q.unsqueeze(1) if query_len > 1: - block_table = block_table.repeat_interleave(query_len, dim=0) - seq_lens = seq_lens.repeat_interleave(query_len) + block_table, seq_lens = self._prepare_flattened_decode_metadata( + attn_metadata, + query_len, + causal=False, + ) + elif self.dcp_world_size > 1 and query_len > 1: + # Causal DCP block: flatten to single-token decode rows with + # per-row rank-local seq_lens for each query's visible prefix. + block_table, seq_lens = self._prepare_flattened_decode_metadata( + attn_metadata, + query_len, + causal=True, + ) + q = q.unsqueeze(1) # trtllm API requires extra dimension q_len_per_request for MTP elif attn_metadata.num_decode_tokens % attn_metadata.num_decodes != 0: logger.warning_once( @@ -358,3 +420,49 @@ def forward_mqa( o = o.view(-1, o.shape[-2], o.shape[-1]) return o, lse + + def _prepare_flattened_decode_metadata( + self, + attn_metadata: FlashInferMLAMetadata, + query_len: int, + *, + causal: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Prepare flattened decode metadata once for all layers in the group.""" + decode = attn_metadata.decode + assert decode is not None + if decode.query_len: + assert decode.query_len == query_len + assert decode.flattened_block_table is not None + assert decode.flattened_seq_lens is not None + return decode.flattened_block_table, decode.flattened_seq_lens + + block_table = decode.block_table.repeat_interleave(query_len, dim=0) + if causal: + global_seq_lens = decode.dcp_tot_seq_lens + assert global_seq_lens is not None + offsets = torch.arange( + query_len - 1, + -1, + -1, + device=global_seq_lens.device, + dtype=global_seq_lens.dtype, + ) + per_query_global_lens = torch.clamp( + (global_seq_lens.unsqueeze(1) - offsets).reshape(-1), min=0 + ) + interleave = self.cp_kv_cache_interleave_size + dcp_span = self.dcp_world_size * interleave + remainder = torch.clamp( + per_query_global_lens % dcp_span - self.dcp_rank * interleave, + min=0, + max=interleave, + ) + seq_lens = per_query_global_lens // dcp_span * interleave + remainder + else: + seq_lens = decode.seq_lens.repeat_interleave(query_len) + + decode.flattened_block_table = block_table + decode.flattened_seq_lens = seq_lens + decode.query_len = query_len + return block_table, seq_lens diff --git a/vllm/v1/attention/backends/mla/tokenspeed_mla.py b/vllm/v1/attention/backends/mla/tokenspeed_mla.py index 8ab81c438f20..57a628b82b53 100644 --- a/vllm/v1/attention/backends/mla/tokenspeed_mla.py +++ b/vllm/v1/attention/backends/mla/tokenspeed_mla.py @@ -61,6 +61,7 @@ class TokenspeedMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]): # The kernel accepts an explicit causal mask, so a non-causal DSpark # block can remain fused instead of being flattened to single tokens. supports_non_causal_multi_token_decode: ClassVar[bool] = True + supports_non_causal_multi_token_dcp: ClassVar[bool] = True def __init__( self, diff --git a/vllm/v1/attention/selector.py b/vllm/v1/attention/selector.py index 00f2aed5c506..b10cb3497737 100644 --- a/vllm/v1/attention/selector.py +++ b/vllm/v1/attention/selector.py @@ -38,6 +38,7 @@ class AttentionSelectorConfig(NamedTuple): use_kv_connector: bool = False use_pcp: bool = False use_adaptive_verification: bool = False + use_dcp: bool = False def __repr__(self): return ( @@ -56,7 +57,8 @@ def __repr__(self): f"use_batch_invariant={self.use_batch_invariant}, " f"use_kv_connector={self.use_kv_connector}, " f"use_adaptive_verification={self.use_adaptive_verification}, " - f"use_pcp={self.use_pcp})" + f"use_pcp={self.use_pcp}, " + f"use_dcp={self.use_dcp})" ) @@ -168,6 +170,7 @@ def get_attn_backend( use_kv_connector=use_kv_connector, use_pcp=vllm_config.parallel_config.prefill_context_parallel_size > 1, use_adaptive_verification=use_adaptive_verification, + use_dcp=vllm_config.parallel_config.decode_context_parallel_size > 1, ) # A per-KV-group override (keyed by KVCacheSpecKind) takes precedence over diff --git a/vllm/v1/worker/gpu/cp_utils.py b/vllm/v1/worker/gpu/cp_utils.py index 6dd8fd34743e..77010990b8e3 100644 --- a/vllm/v1/worker/gpu/cp_utils.py +++ b/vllm/v1/worker/gpu/cp_utils.py @@ -59,3 +59,24 @@ def _dcp_local_seq_lens_kernel( # For [num_reqs, max_num_reqs), pad with 0 local_seq_lens = tl.where(block < num_reqs, local_seq_lens, 0) tl.store(out_ptr + block, local_seq_lens, mask=block < max_num_reqs) + + +@triton.jit +def cp_local_slot( + positions, + block_numbers, + block_size, + cp_rank, + CP_SIZE: tl.constexpr, + CP_INTERLEAVE: tl.constexpr, + PAD_ID: tl.constexpr, +): + """Return rank-local KV slots, or PAD_ID for positions not owned by this rank.""" + block_offsets = positions % (block_size * CP_SIZE) + if CP_SIZE == 1: + return block_numbers * block_size + block_offsets + is_local = block_offsets // CP_INTERLEAVE % CP_SIZE == cp_rank + rounds = block_offsets // (CP_INTERLEAVE * CP_SIZE) + remainder = block_offsets % CP_INTERLEAVE + local_offsets = rounds * CP_INTERLEAVE + remainder + return tl.where(is_local, block_numbers * block_size + local_offsets, PAD_ID) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py b/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py index e59202c87988..c68b5497458b 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py @@ -11,6 +11,7 @@ build_slot_mappings_by_layer, ) from vllm.v1.worker.gpu.block_table import BlockTables +from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens from vllm.v1.worker.gpu.cudagraph_utils import ( AttentionState, BatchExecutionDescriptor, @@ -41,6 +42,17 @@ def _prepare_dflash_inputs_to_capture( attn_metadata = None if not skip_attn: query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np) + dcp_local_seq_lens = None + if block_tables.cp_size > 1: + prepare_dcp_local_seq_lens( + input_buffers.dcp_local_seq_lens, + input_buffers.seq_lens, + num_reqs, + block_tables.cp_size, + block_tables.cp_rank, + block_tables.cp_interleave, + ) + dcp_local_seq_lens = input_buffers.dcp_local_seq_lens attn_metadata = build_attn_metadata( attn_groups=attn_groups, num_reqs=num_reqs, @@ -49,6 +61,7 @@ def _prepare_dflash_inputs_to_capture( query_start_loc_cpu=query_start_loc_cpu, max_query_len=num_tokens // num_reqs, seq_lens=input_batch.seq_lens, + dcp_local_seq_lens=dcp_local_seq_lens, max_seq_len=max_model_len, block_tables=input_block_tables, slot_mappings=slot_mappings, diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index eb5dc470b26c..92837afd68eb 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -17,6 +17,7 @@ from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer from vllm.v1.worker.gpu.block_table import BlockTables +from vllm.v1.worker.gpu.cp_utils import cp_local_slot, prepare_dcp_local_seq_lens from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers from vllm.v1.worker.gpu.model_states.interface import ModelState @@ -285,10 +286,21 @@ def _build_draft_attn_metadata( num_query_per_req: int | None = None, causal: bool | Mapping[int, bool] = False, query_start_loc_np: np.ndarray | None = None, + dcp_local_seq_lens: torch.Tensor | None = None, ) -> dict[str, Any] | None: if not self.draft_attn_layer_names: return None assert num_query_per_req is None # Omitted for DFlash, read from self instead + if dcp_local_seq_lens is None and self.block_tables.cp_size > 1: + prepare_dcp_local_seq_lens( + self.input_buffers.dcp_local_seq_lens, + self.input_buffers.seq_lens, + num_reqs, + self.block_tables.cp_size, + self.block_tables.cp_rank, + self.block_tables.cp_interleave, + ) + dcp_local_seq_lens = self.input_buffers.dcp_local_seq_lens return super()._build_draft_attn_metadata( num_reqs, num_reqs_padded, @@ -298,6 +310,7 @@ def _build_draft_attn_metadata( num_query_per_req=self.num_query_per_req, causal=causal, query_start_loc_np=query_start_loc_np, + dcp_local_seq_lens=dcp_local_seq_lens, ) @torch.inference_mode() @@ -393,6 +406,9 @@ def propose( seeds, self.block_tables.input_block_tables[gid], self.block_tables.kernel_block_sizes[gid], + self.block_tables.cp_rank, + self.block_tables.cp_size, + self.block_tables.cp_interleave, self.parallel_drafting_token_id, self.num_query_per_req, self.num_speculative_steps, @@ -507,8 +523,11 @@ def _prepare_dflash_inputs_kernel( max_num_reqs, max_num_tokens, max_model_len, + cp_rank, SAMPLE_FROM_ANCHOR: tl.constexpr, PAD_SLOT_ID: tl.constexpr, + CP_SIZE: tl.constexpr, + CP_INTERLEAVE: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): req_idx = tl.program_id(0) @@ -543,7 +562,7 @@ def _prepare_dflash_inputs_kernel( # --- Context positions / slots --- ctx_pos_idx = ctx_start + tl.where(is_ctx, j, 0) ctx_pos = tl.load(target_positions_ptr + ctx_pos_idx, mask=is_valid_ctx, other=0) - ctx_block_num = ctx_pos // block_size + ctx_block_num = ctx_pos // (block_size * CP_SIZE) ctx_block_num = tl.minimum(ctx_block_num, block_table_stride - 1) ctx_block_id = tl.load( block_table_ptr + req_idx * block_table_stride + ctx_block_num, @@ -554,9 +573,12 @@ def _prepare_dflash_inputs_kernel( # to it after eviction; rejected suffix rows are invalid context as well. # Neither kind of row may write draft KV into physical block 0. ctx_resident = is_valid_ctx & (ctx_block_id != 0) + local_ctx_slot = cp_local_slot( + ctx_pos, ctx_block_id, block_size, cp_rank, CP_SIZE, CP_INTERLEAVE, PAD_SLOT_ID + ) ctx_slot = tl.where( ctx_resident, - ctx_block_id * block_size + (ctx_pos % block_size), + local_ctx_slot, PAD_SLOT_ID, ) # Stored over the full [0, num_ctx) span while the loads above are masked to @@ -573,7 +595,7 @@ def _prepare_dflash_inputs_kernel( is_bonus = is_query & (query_off == 0) input_id = tl.where(is_bonus, bonus_token, parallel_drafting_token_id) - q_block_num = query_pos // block_size + q_block_num = query_pos // (block_size * CP_SIZE) q_block_num = tl.minimum(q_block_num, block_table_stride - 1) q_block_id = tl.load( block_table_ptr + req_idx * block_table_stride + q_block_num, @@ -583,9 +605,18 @@ def _prepare_dflash_inputs_kernel( # A null block is never a writable cache slot. This can occur when a # sliding-window block table contains evicted/global padding entries. q_resident = is_query & (q_block_id != 0) + local_q_slot = cp_local_slot( + query_pos, + q_block_id, + block_size, + cp_rank, + CP_SIZE, + CP_INTERLEAVE, + PAD_SLOT_ID, + ) q_slot = tl.where( q_resident, - q_block_id * block_size + (query_pos % block_size), + local_q_slot, PAD_SLOT_ID, ) @@ -681,6 +712,9 @@ def prepare_dflash_inputs( # [max_num_reqs, max_num_blocks] block_table: torch.Tensor, block_size: int, + cp_rank: int, + cp_size: int, + cp_interleave: int, parallel_drafting_token_id: int, num_query_per_req: int, num_speculative_steps: int, @@ -728,7 +762,10 @@ def prepare_dflash_inputs( max_num_reqs, max_num_tokens, max_model_len, + cp_rank, SAMPLE_FROM_ANCHOR=sample_from_anchor, PAD_SLOT_ID=PAD_SLOT_ID, + CP_SIZE=cp_size, + CP_INTERLEAVE=cp_interleave, BLOCK_SIZE=BLOCK_SIZE, ) diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 69e0cc109160..80de83109bf5 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -234,6 +234,7 @@ def _build_draft_attn_metadata( num_query_per_req: int = 1, causal: bool | Mapping[int, bool] = True, query_start_loc_np: np.ndarray | None = None, + dcp_local_seq_lens: torch.Tensor | None = None, ) -> dict[str, Any] | None: if query_start_loc_np is not None: # Non-uniform query layout (e.g. multi-module MTP's mixed @@ -278,6 +279,11 @@ def _build_draft_attn_metadata( query_start_loc_cpu=query_start_loc_cpu, max_query_len=max_query_len, seq_lens=self.input_buffers.seq_lens[:num_reqs_padded], + dcp_local_seq_lens=( + None + if dcp_local_seq_lens is None + else dcp_local_seq_lens[:num_reqs_padded] + ), max_seq_len=self.draft_max_seq_len, block_tables=block_tables, slot_mappings=slot_mappings, From f08a95f8d84260fb093975cbf2bbccee8d106b69 Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Mon, 17 Aug 2026 13:27:56 -0700 Subject: [PATCH 059/839] [Rust Frontend][gRPC] Advertise LoRA capabilities (#52031) Signed-off-by: Connor Carpenter Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- rust/proto/control.proto | 2 ++ rust/src/engine-core-client/src/client.rs | 27 +++++++++++++++++++ .../src/engine-core-client/src/mock_engine.rs | 2 ++ .../src/protocol/handshake.rs | 4 +++ .../engine-core-client/src/tests/client.rs | 2 ++ .../src/tests/python_compat.py | 4 +++ rust/src/server/src/grpc/control.rs | 2 ++ tests/v1/engine/test_engine_core_client.py | 2 ++ vllm/v1/engine/__init__.py | 2 ++ vllm/v1/engine/core.py | 6 +++++ 10 files changed, 53 insertions(+) diff --git a/rust/proto/control.proto b/rust/proto/control.proto index 428ad6460499..1903405e2749 100644 --- a/rust/proto/control.proto +++ b/rust/proto/control.proto @@ -38,6 +38,7 @@ message ServerInfo { uint64 total_kv_blocks = 7; uint64 max_running_requests = 8; uint64 max_batched_tokens = 9; + uint32 max_loras = 10; RlCapabilities rl_capabilities = 11; } @@ -65,6 +66,7 @@ message ModelInfo { bool supports_text_input = 20; bool supports_token_ids_input = 21; + bool supports_lora = 22; bool supports_multimodal = 23; string reasoning_parser = 24; string tool_call_parser = 25; diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index 7f8663f85268..1efca89621ff 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/rust/src/engine-core-client/src/client.rs @@ -286,6 +286,7 @@ impl EngineCoreClient { config: EngineCoreClientConfig, connected: transport::ConnectedTransport, ) -> Result { + validate_lora_capabilities(&connected.engines)?; let (output_tx, output_rx) = mpsc::channel(64); let (abort_tx, abort_rx) = mpsc::unbounded_channel(); let engines = connected.engines; @@ -476,6 +477,32 @@ impl EngineCoreClient { } } +fn validate_lora_capabilities(engines: &[ConnectedEngine]) -> Result<()> { + let first = engines.first().expect("engine core client requires at least one engine"); + for engine in engines { + let ready = &engine.ready_response; + if ready.supports_lora != (ready.max_loras > 0) { + return Err(Error::UnexpectedHandshakeMessage { + message: format!( + "engine {:?} reported inconsistent LoRA capability (supports_lora={}, max_loras={})", + engine.engine_id, ready.supports_lora, ready.max_loras + ), + }); + } + if ready.supports_lora != first.ready_response.supports_lora + || ready.max_loras != first.ready_response.max_loras + { + return Err(Error::UnexpectedHandshakeMessage { + message: format!( + "engine {:?} reported LoRA capability inconsistent with engine {:?}", + engine.engine_id, first.engine_id + ), + }); + } + } + Ok(()) +} + // Client API implementation. impl EngineCoreClient { /// Add a new request to the engine and return a per-request raw output diff --git a/rust/src/engine-core-client/src/mock_engine.rs b/rust/src/engine-core-client/src/mock_engine.rs index 28fb4d84fe05..8bd6d773607e 100644 --- a/rust/src/engine-core-client/src/mock_engine.rs +++ b/rust/src/engine-core-client/src/mock_engine.rs @@ -64,6 +64,8 @@ pub fn default_ready_response() -> EngineCoreReadyResponse { max_num_seqs: 256, max_num_batched_tokens: 8192, instance_id: "test-instance".to_string(), + supports_lora: false, + max_loras: 0, kv_cache_size_tokens: None, kv_cache_max_concurrency: None, kv_events_config: None, diff --git a/rust/src/engine-core-client/src/protocol/handshake.rs b/rust/src/engine-core-client/src/protocol/handshake.rs index b4214a06f27f..5cb4c9804678 100644 --- a/rust/src/engine-core-client/src/protocol/handshake.rs +++ b/rust/src/engine-core-client/src/protocol/handshake.rs @@ -80,6 +80,10 @@ pub struct EngineCoreReadyResponse { pub max_num_batched_tokens: u64, /// Unique identifier for this server instance. pub instance_id: String, + /// Whether the engine was started with LoRA support enabled. + pub supports_lora: bool, + /// Maximum number of LoRA adapters the engine may keep active. + pub max_loras: u32, /// Total KV cache capacity in tokens, if reported. pub kv_cache_size_tokens: Option, /// Maximum achievable request concurrency given the KV cache, if reported. diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index 8f07a8059e92..913b894ada7a 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -2686,6 +2686,8 @@ fn python_msgpack_fixtures_match_rust_encoding() { let ready_response: EngineCoreReadyResponse = rmp_serde::from_slice(&hex::decode(ready_response_hex).unwrap()).unwrap(); + assert!(ready_response.supports_lora); + assert_eq!(ready_response.max_loras, 8); assert_eq!( ready_response.weight_transfer_backend.as_deref(), Some("nccl") diff --git a/rust/src/engine-core-client/src/tests/python_compat.py b/rust/src/engine-core-client/src/tests/python_compat.py index 389e97cd9996..ac033914b0d1 100755 --- a/rust/src/engine-core-client/src/tests/python_compat.py +++ b/rust/src/engine-core-client/src/tests/python_compat.py @@ -409,6 +409,8 @@ class EngineCoreReadyResponse: max_num_seqs: int max_num_batched_tokens: int instance_id: str + supports_lora: bool + max_loras: int kv_cache_size_tokens: int | None = None kv_cache_max_concurrency: float | None = None kv_events_config: KVEventsConfig | None = None @@ -433,6 +435,8 @@ class EngineCoreReadyResponse: max_num_seqs=256, max_num_batched_tokens=8192, instance_id="test-instance", + supports_lora=True, + max_loras=8, weight_transfer_backend="nccl", enable_sleep_mode=True, supports_draft_weight_updates=True, diff --git a/rust/src/server/src/grpc/control.rs b/rust/src/server/src/grpc/control.rs index c7d9af88c82f..c08803306fc8 100644 --- a/rust/src/server/src/grpc/control.rs +++ b/rust/src/server/src/grpc/control.rs @@ -150,6 +150,7 @@ impl pb::control_server::Control for ControlServiceImpl { total_kv_blocks: self.state.engine_core_client().total_num_gpu_blocks(), max_running_requests: ready.max_num_seqs, max_batched_tokens: ready.max_num_batched_tokens, + max_loras: ready.max_loras, rl_capabilities: Some(self.rl_capabilities()), })) } @@ -166,6 +167,7 @@ impl pb::control_server::Control for ControlServiceImpl { // GenerateRequest accepts both prompt representations. supports_text_input: true, supports_token_ids_input: true, + supports_lora: self.ready().supports_lora, supports_multimodal: self.state.chat.supports_multimodal(), reasoning_parser: self .state diff --git a/tests/v1/engine/test_engine_core_client.py b/tests/v1/engine/test_engine_core_client.py index 74f7b803b890..1727c51b19e0 100644 --- a/tests/v1/engine/test_engine_core_client.py +++ b/tests/v1/engine/test_engine_core_client.py @@ -335,6 +335,8 @@ def test_apply_ready_response_syncs_block_size(): max_num_seqs=256, max_num_batched_tokens=8192, instance_id="test-instance", + supports_lora=False, + max_loras=0, ) ) client._apply_ready_response(payload) diff --git a/vllm/v1/engine/__init__.py b/vllm/v1/engine/__init__.py index d70778eb51e3..e07df08f3d9a 100644 --- a/vllm/v1/engine/__init__.py +++ b/vllm/v1/engine/__init__.py @@ -88,6 +88,8 @@ class EngineCoreReadyResponse: max_num_seqs: int max_num_batched_tokens: int instance_id: str + supports_lora: bool + max_loras: int # KV cache capacity (None for encoder-only/attention-free models). kv_cache_size_tokens: int | None = None kv_cache_max_concurrency: float | None = None diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 295c66b360ed..641c5ecc4e01 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -1636,6 +1636,12 @@ def _make_ready_response(self) -> EngineCoreReadyResponse: max_num_seqs=scheduler_config.max_num_seqs, max_num_batched_tokens=scheduler_config.max_num_batched_tokens, instance_id=self.vllm_config.instance_id, + supports_lora=self.vllm_config.lora_config is not None, + max_loras=( + self.vllm_config.lora_config.max_loras + if self.vllm_config.lora_config is not None + else 0 + ), kv_events_config=self.scheduler.get_kv_event_publisher_config(), weight_transfer_backend=( self.vllm_config.weight_transfer_config.backend From 8878ebd8fdfe286eb7cf00ccbb6919bc7a40b1dd Mon Sep 17 00:00:00 2001 From: Micah Williamson Date: Mon, 17 Aug 2026 15:46:22 -0500 Subject: [PATCH 060/839] [ROCm][CI] Expand AITER W4A4 MoE Coverage (#52647) Signed-off-by: Micah Williamson --- tests/kernels/moe/test_ocp_mx_moe.py | 25 ++++++++++++++-- tests/quantization/test_gfx950_moe.py | 41 +++++++++++++++++++++++---- 2 files changed, 59 insertions(+), 7 deletions(-) diff --git a/tests/kernels/moe/test_ocp_mx_moe.py b/tests/kernels/moe/test_ocp_mx_moe.py index 5b8dfab25bc4..2d3e48d435f3 100644 --- a/tests/kernels/moe/test_ocp_mx_moe.py +++ b/tests/kernels/moe/test_ocp_mx_moe.py @@ -125,7 +125,6 @@ def test_mxfp4_loading_and_execution_moe(vllm_runner, model_case: ModelCase): # if model_case.model_id == "fxmarty/qwen_1.5-moe-a2.7b-mxfp4": # llm.apply_model(check_model) - output = llm.generate_greedy("Today I am in the French Alps and", max_tokens=20) assert output @@ -187,6 +186,16 @@ def mxfp8_dequantize(x, scale): return x_float * scale +def mxfp4_quant_dequant(x: torch.Tensor) -> torch.Tensor: + shape = x.shape + quantized, scale = dynamic_mxfp4_quant(x.to(torch.bfloat16).flatten(0, -2)) + return ( + upcast_from_mxfp(quantized.view(torch.uint8), scale, torch.bfloat16, axis=-1) + .reshape(shape) + .float() + ) + + def reference_moe( roouting_logits, topk, @@ -219,6 +228,8 @@ def reference_moe( expert_weights = torch.nn.functional.softmax(experts.values, dim=1) expert_indices = experts.indices t = hidden_states.clone() + if act_type == "mxfp4": + t = mxfp4_quant_dequant(t) # MLP #1 mlp1_weight = w13[expert_indices, ...] mlp1_bias = bias13[expert_indices, ...] @@ -244,6 +255,8 @@ def reference_moe( t.to(torch.bfloat16), is_sf_swizzled_layout=False ) t = mxfp8_dequantize(t_quantized, t_scale) + elif act_type == "mxfp4": + t = mxfp4_quant_dequant(t) # MLP #2 mlp2_weight = w2[expert_indices, ...] mlp2_bias = bias2[expert_indices, ...] @@ -1392,6 +1405,14 @@ def test_trtllm_gen_mxfp8_block_scale_moe( "requires_aiter": True, "requires_gfx950": True, }, + "AITER_MXFP4_MXFP4": { + "activation": "SILU", + "act_type": "mxfp4", + "rtol": 1.0, + "percent": 0.8, + "requires_aiter": True, + "requires_gfx950": True, + }, } @@ -1662,7 +1683,7 @@ class MockLayer: alpha=1.702 if activation == MoEActivation.SWIGLUOAI else 1.0, beta=1.0 if activation == MoEActivation.SWIGLUOAI else 0.0, limit=7.0 if activation == MoEActivation.SWIGLUOAI else None, - act_type="bf16", + act_type=str(config.get("act_type", "bf16")), activation=act_name, use_interleaved_layout=use_interleaved, ) diff --git a/tests/quantization/test_gfx950_moe.py b/tests/quantization/test_gfx950_moe.py index c8d34bb0ab50..a4dac1c16340 100644 --- a/tests/quantization/test_gfx950_moe.py +++ b/tests/quantization/test_gfx950_moe.py @@ -25,14 +25,32 @@ ROCM_AVAILABLE = current_platform.is_rocm() ROCM_GFX950 = False -ROCM_AITER_AVAILABLE = False +ROCM_AITER_SUPPORTED = False if ROCM_AVAILABLE: - from vllm._aiter_ops import rocm_aiter_ops + from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops from vllm.platforms.rocm import on_gfx950 ROCM_GFX950 = on_gfx950() - ROCM_AITER_AVAILABLE = rocm_aiter_ops.is_fused_moe_enabled() + ROCM_AITER_SUPPORTED = is_aiter_found_and_supported() + + +def set_rocm_aiter(monkeypatch: pytest.MonkeyPatch, enabled: bool) -> None: + value = "1" if enabled else "0" + monkeypatch.setenv("VLLM_ROCM_USE_AITER", value) + monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", value) + monkeypatch.setattr(rocm_aiter_ops, "_AITER_ENABLED", enabled) + monkeypatch.setattr(rocm_aiter_ops, "_FMOE_ENABLED", enabled) + + +@pytest.fixture +def enable_rocm_aiter(monkeypatch: pytest.MonkeyPatch): + set_rocm_aiter(monkeypatch, True) + + +@pytest.fixture +def disable_rocm_aiter(monkeypatch: pytest.MonkeyPatch): + set_rocm_aiter(monkeypatch, False) def _make_w4a4_moe_config(moe_backend: str = "auto") -> FusedMoEConfig: @@ -68,8 +86,8 @@ def mxfp4_oracle_config(): @pytest.mark.skipif(not ROCM_GFX950, reason="Requires GFX950 (mi355x)") -@pytest.mark.skipif(not ROCM_AITER_AVAILABLE, reason="Requires AITER enabled") -def test_w4a4_dispatches_to_aiter(mxfp4_oracle_config): +@pytest.mark.skipif(not ROCM_AITER_SUPPORTED, reason="Requires supported AITER") +def test_w4a4_dispatches_to_aiter(mxfp4_oracle_config, enable_rocm_aiter): """With AITER enabled + GFX950, W4A4 selects AITER_MXFP4_MXFP4.""" config = _make_w4a4_moe_config() backend, experts_cls = select_mxfp4_moe_backend( @@ -79,6 +97,19 @@ def test_w4a4_dispatches_to_aiter(mxfp4_oracle_config): assert experts_cls is not None +@pytest.mark.skipif(not ROCM_GFX950, reason="Requires GFX950 (mi355x)") +def test_w4a4_falls_back_without_aiter( + mxfp4_oracle_config, + disable_rocm_aiter, +): + config = _make_w4a4_moe_config() + backend, experts_cls = select_mxfp4_moe_backend( + config, activation_key=kMxfp4Dynamic + ) + assert backend == Mxfp4MoeBackend.EMULATION + assert experts_cls is not None + + @pytest.mark.skipif(not ROCM_GFX950, reason="Requires GFX950 (mi355x)") def test_w4a4_dispatches_to_emulation_with_moe_backend(mxfp4_oracle_config): """With --moe-backend emulation, W4A4 selects EMULATION.""" From e68fb75b250e33f013162582ece39e766a424df8 Mon Sep 17 00:00:00 2001 From: Hongxia Yang <62075498+hongxiayang@users.noreply.github.com> Date: Mon, 17 Aug 2026 17:58:48 -0400 Subject: [PATCH 061/839] [ROCm][AMD][Installation] add LMCache kv-connector installation and runtime packages to docker image (#51208) Signed-off-by: Hongxia Yang Signed-off-by: Andreas Karatzas Co-authored-by: Andreas Karatzas --- .buildkite/release-pipeline.yaml | 1 + .buildkite/scripts/ci-bake-rocm.sh | 2 +- docker/Dockerfile.rocm | 87 +++++++++++++++++++++++++++++- 3 files changed, 88 insertions(+), 2 deletions(-) diff --git a/.buildkite/release-pipeline.yaml b/.buildkite/release-pipeline.yaml index 7b9cac7b5257..1dd8f0146c31 100644 --- a/.buildkite/release-pipeline.yaml +++ b/.buildkite/release-pipeline.yaml @@ -906,6 +906,7 @@ steps: --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \ --build-arg SCCACHE_REGION_NAME=us-west-2 \ --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \ + --build-arg INSTALL_LMCACHE=true \ --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm \ --target vllm-openai \ --progress plain \ diff --git a/.buildkite/scripts/ci-bake-rocm.sh b/.buildkite/scripts/ci-bake-rocm.sh index 45e45cf33a39..1e143113141f 100644 --- a/.buildkite/scripts/ci-bake-rocm.sh +++ b/.buildkite/scripts/ci-bake-rocm.sh @@ -17,7 +17,7 @@ DEFAULT_REPO_SLUG="vllm-project/vllm" DEFAULT_CI_HCL_SOURCE="docker/ci-rocm.hcl" DEFAULT_CI_BASE_CONTENT_FILES=".dockerignore requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt tools/install_torchcodec_rocm.sh tools/install_protoc.sh rust-toolchain.toml tests/vllm_test_utils" DEFAULT_CI_BASE_DOCKERFILE="docker/Dockerfile.rocm" -DEFAULT_CI_BASE_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust-toolchain-input rust-toolchain build_nixl build_rocshmem build_deepep mori_base ci_base" +DEFAULT_CI_BASE_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust-toolchain-input rust-toolchain build_nixl lmcache_source build_lmcache build_rocshmem build_deepep mori_base ci_base" DEFAULT_CI_BASE_METADATA_VERSION="3" # ROCm CI forces REMOTE_VLLM=0, so content identity covers only the selected # local-source stages rather than unreachable remote-fetch alternatives. diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index d3504d1a3dfb..0487c0c9165a 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -6,6 +6,7 @@ ARG COMMON_WORKDIR=/app ARG BASE_IMAGE=rocm/vllm-dev:base ARG CI_BASE_IMAGE=rocm/vllm-dev:ci_base ARG ROCM_TRITON_KERNELS_COMMIT=0f380657dbf3ee86eb57558ff71df24f03b5d4e7 +ARG INSTALL_LMCACHE=false # NIC backend for MoRI RDMA support. # By default (all), drivers and userspace libraries for all supported NIC types # (ainic and bnxt) are installed; MoRI selects the appropriate one at runtime. @@ -471,6 +472,71 @@ RUN cd /opt/nixl && \ /tmp/nixl_wheels/repaired/*.whl && \ cp /tmp/nixl_wheels/repaired/*.whl /app/install +# Fetch and verify pinned LMCache source. +FROM base AS lmcache_source +ARG LMCACHE_REPO="https://github.com/LMCache/LMCache.git" +ARG LMCACHE_REF="140819c9d57a975dbc5678a6459a218e544cb58b" +ARG LMCACHE_VERSION="0.5.3" + +# Remove NVIDIA-only dependencies and accept the release image's NumPy version. +RUN printf '%s\n' "${LMCACHE_REF}" | grep -Eq '^[0-9a-f]{40}$' \ + && git clone --depth=1 --branch "v${LMCACHE_VERSION}" \ + "${LMCACHE_REPO}" /app/lmcache \ + && test "$(git -C /app/lmcache rev-parse HEAD)" = "${LMCACHE_REF}" \ + && test "$(grep -Ec '^(cufile-python|nvtx)([<=>[:space:]]|$)' \ + /app/lmcache/requirements/common.txt)" -eq 2 \ + && test "$(grep -Fxc 'numpy<=2.2.6' \ + /app/lmcache/requirements/common.txt)" -eq 1 \ + && sed -i -E \ + -e '/^(cufile-python|nvtx)([<=>[:space:]]|$)/d' \ + -e 's/^numpy<=2\.2\.6$/numpy<=2.3.5/' \ + /app/lmcache/requirements/common.txt \ + && rm -rf /app/lmcache/.git + +FROM base AS build_lmcache +COPY --from=lmcache_source \ + /app/lmcache/requirements/build.txt /tmp/lmcache-build-requirements.txt +RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ + uv pip install --system -r /tmp/lmcache-build-requirements.txt + +ARG LMCACHE_VERSION="0.5.3" +ARG LMCACHE_ROCM_ARCH="gfx942;gfx950" +ARG USE_SCCACHE + +COPY --from=lmcache_source /app/lmcache /app/lmcache +WORKDIR /app/lmcache + +RUN --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ + set -eu; \ + printf '%s\n' "${LMCACHE_ROCM_ARCH}" \ + | grep -Eq '^gfx[0-9a-z]+(;gfx[0-9a-z]+)*$'; \ + offload_flags="$(printf '%s' "${LMCACHE_ROCM_ARCH}" \ + | sed 's/;/ --offload-arch=/g; s/^/--offload-arch=/')"; \ + if [ "${USE_SCCACHE:-0}" = "1" ]; then \ + export HIP_CLANG_PATH=/opt/sccache-wrappers; \ + fi; \ + mkdir -p /app/install; \ + SETUPTOOLS_SCM_PRETEND_VERSION="${LMCACHE_VERSION}" \ + PYTORCH_ROCM_ARCH="${LMCACHE_ROCM_ARCH}" \ + HIPCC_COMPILE_FLAGS_APPEND="${offload_flags}" \ + BUILD_WITH_HIP=1 CXX=hipcc \ + python3 setup.py bdist_wheel --dist-dir=/app/install; \ + test "$(find /app/install -maxdepth 1 -name '*.whl' | wc -l)" -eq 1 + +RUN --mount=type=tmpfs,target=/tmp/lmcache-wheel \ + set -eu; \ + python3 -m zipfile -e /app/install/*.whl /tmp/lmcache-wheel; \ + c_ops="$(find /tmp/lmcache-wheel -type f -name 'c_ops*.so')"; \ + test "$(find /tmp/lmcache-wheel -type f -name 'c_ops*.so' | wc -l)" -eq 1; \ + expected="$(printf '%s' "${LMCACHE_ROCM_ARCH}" | tr ';' '\n' | sort -u)"; \ + actual="$(/opt/rocm/llvm/bin/llvm-objdump --offloading "${c_ops}" \ + | grep -oE 'gfx[0-9a-z]+' | sort -u)"; \ + if [ "${actual}" != "${expected}" ]; then \ + printf 'LMCache code objects mismatch\nexpected:\n%s\nactual:\n%s\n' \ + "${expected}" "${actual}" >&2; \ + exit 1; \ + fi + # ROCShmem build stage - split from DeepEP so changing DEEPEP_BRANCH does not # invalidate the slow ROCShmem build. FROM base AS build_rocshmem @@ -797,6 +863,13 @@ RUN --mount=type=cache,target=/root/.cache/uv \ uv pip install --system /tmp/vllm_test_utils \ && rm -rf /tmp/vllm_test_utils +RUN --mount=type=bind,from=build_lmcache,src=/app/install,target=/lmcache_install \ + --mount=type=bind,source=requirements/test/rocm.txt,target=/tmp/lmcache-constraints.txt \ + --mount=type=cache,target=/root/.cache/uv \ + uv pip install --system --constraint /tmp/lmcache-constraints.txt \ + /lmcache_install/*.whl \ + && python3 -c "import lmcache, lmcache.c_ops; print('lmcache', lmcache.__version__)" + # ----------------------- # Test vLLM image (Tier 2) - vLLM-only layer on top of ci_base. FROM ${CI_BASE_IMAGE} AS test @@ -825,7 +898,7 @@ ENV VLLM_GPU_SYNC_CHECK=error # ----------------------- # Final vLLM image -FROM mori_base AS final +FROM mori_base AS final_common RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/* @@ -901,6 +974,18 @@ RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt \ CMD ["/bin/bash"] +FROM final_common AS final_lmcache_false + +FROM final_common AS final_lmcache_true +RUN --mount=type=bind,from=build_lmcache,src=/app/install,target=/lmcache_install \ + --mount=type=cache,target=/root/.cache/uv \ + uv pip install --system /lmcache_install/*.whl \ + && test "$(python3 -c 'import numpy; print(numpy.__version__)')" = "2.3.5" \ + && uv pip check --system \ + && python3 -c "import lmcache, lmcache.c_ops; print('lmcache', lmcache.__version__)" + +FROM final_lmcache_${INSTALL_LMCACHE} AS final + #Set entrypoint for vllm-openai official images FROM final AS vllm-openai ENTRYPOINT ["vllm", "serve"] From 455edc022b45bbb5a279cff21cd3a34e6371aadf Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Tue, 18 Aug 2026 05:58:50 +0800 Subject: [PATCH 062/839] [ModelRunnerV2] Support prompt embeds (#42963) Signed-off-by: gcanlin Signed-off-by: Canlin Guo <961750412@qq.com> Signed-off-by: Nick Hill Co-authored-by: Nick Hill --- tests/v1/worker/test_encoder_runner.py | 95 ++++++++++- tests/v1/worker/test_gpu_model_runner.py | 10 ++ tests/v1/worker/test_prompt_embeds_state.py | 149 +++++++++++++++++ vllm/config/model.py | 6 + vllm/config/vllm.py | 3 - vllm/model_executor/models/diffusion_gemma.py | 4 +- .../models/longcat_flash_ngram.py | 3 + vllm/v1/core/sched/output.py | 8 + vllm/v1/request.py | 18 ++- vllm/v1/worker/gpu/mm/encoder_runner.py | 13 +- vllm/v1/worker/gpu/model_runner.py | 24 +-- vllm/v1/worker/gpu/model_states/__init__.py | 26 ++- vllm/v1/worker/gpu/model_states/default.py | 77 ++++++--- .../gpu/model_states/encoder_decoder.py | 6 +- .../worker/gpu/model_states/encoder_only.py | 3 + vllm/v1/worker/gpu/model_states/interface.py | 8 +- .../worker/gpu/model_states/prompt_embeds.py | 152 ++++++++++++++++++ 17 files changed, 551 insertions(+), 54 deletions(-) create mode 100644 tests/v1/worker/test_prompt_embeds_state.py create mode 100644 vllm/v1/worker/gpu/model_states/prompt_embeds.py diff --git a/tests/v1/worker/test_encoder_runner.py b/tests/v1/worker/test_encoder_runner.py index 914985b82924..f8bd50c6d9c4 100644 --- a/tests/v1/worker/test_encoder_runner.py +++ b/tests/v1/worker/test_encoder_runner.py @@ -15,7 +15,13 @@ import pytest import torch -from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange +from vllm.multimodal.inputs import ( + MultiModalFeatureSpec, + MultiModalFieldElem, + MultiModalKwargsItem, + MultiModalSharedField, + PlaceholderRange, +) from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner from vllm.v1.worker.gpu.model_states.interface import ModelState @@ -25,6 +31,25 @@ HIDDEN = 4 +def _model_state(cache: EncoderCache) -> MagicMock: + """A mock ModelState backed by a real EncoderCache.""" + state = MagicMock() + state.encoder_cache = cache + state.device = torch.device("cpu") + return state + + +def _embeds_item(embeds: torch.Tensor) -> MultiModalKwargsItem: + """A `prompt_embeds` kwargs item, as the HF renderer builds it.""" + return MultiModalKwargsItem( + { + "embedding": MultiModalFieldElem( + data=embeds, field=MultiModalSharedField(batch_size=1) + ) + } + ) + + def _feature(identifier: str, offset: int, length: int) -> MultiModalFeatureSpec: return MultiModalFeatureSpec( data=None, @@ -197,8 +222,7 @@ def test_execute_mm_encoder_caches_outputs_without_gathering(): items the connector already holds, and a producer has no load path). """ cache = EncoderCache() - state = MagicMock() - state.encoder_cache = cache + state = _model_state(cache) embedding = torch.ones(2, HIDDEN) # (mm_hashes, [(modality, kwargs item), ...]), as prepare_mm_inputs returns. state.encoder_runner.prepare_mm_inputs.return_value = ( @@ -216,8 +240,7 @@ def test_execute_mm_encoder_caches_outputs_without_gathering(): def test_execute_mm_encoder_is_a_noop_without_scheduled_items(): """A step that schedules no encoder input must not touch the encoder.""" cache = EncoderCache() - state = MagicMock() - state.encoder_cache = cache + state = _model_state(cache) state.encoder_runner.prepare_mm_inputs.return_value = ([], []) ModelState.execute_mm_encoder(state, {}) @@ -226,6 +249,68 @@ def test_execute_mm_encoder_is_a_noop_without_scheduled_items(): state.encoder_runner.execute_mm_encoder.assert_not_called() +def _pe_feature(identifier: str, embeds: torch.Tensor, offset: int = 0): + return MultiModalFeatureSpec( + data=_embeds_item(embeds), + modality="prompt_embeds", + identifier=identifier, + mm_position=PlaceholderRange(offset=offset, length=embeds.shape[0]), + ) + + +def test_prepare_mm_inputs_passes_prompt_embeds_through(): + """`prompt_embeds` is already in embedding space, so no encoder may run. + + The renderer delivers prompt_embeds mixed with real media as an ordinary MM + modality. prepare_mm_inputs must cache the tensor directly and keep it out + of the encoder batch -- the vision encoder cannot consume it, and a missing + cache entry makes the subsequent gather raise "Encoder cache miss". + """ + prompt_embeds = torch.arange(2 * HIDDEN, dtype=torch.float32).view(2, HIDDEN) + image_feature = MultiModalFeatureSpec( + data=MagicMock(), + modality="image", + identifier="hash_img", + mm_position=PlaceholderRange(offset=2, length=2), + ) + runner = _make_runner( + [_pe_feature("hash_pe", prompt_embeds), image_feature], cached=[] + ) + + mm_hashes, mm_kwargs = runner.prepare_mm_inputs({"req0": [0, 1]}) + + # Only the image remains for the encoder; the embeds are already cached. + assert mm_hashes == ["hash_img"] + assert [modality for modality, _ in mm_kwargs] == ["image"] + assert torch.equal(runner.encoder_cache.encoder_outputs["hash_pe"], prompt_embeds) + + +def test_prepare_mm_inputs_skips_cached_prompt_embeds(): + """A prompt_embeds item already in the cache must not be re-uploaded.""" + prompt_embeds = torch.ones(3, HIDDEN) + feature = _pe_feature("hash_pe", prompt_embeds) + runner = _make_runner([feature], cached=[feature]) + sentinel = runner.encoder_cache.encoder_outputs["hash_pe"] + + mm_hashes, mm_kwargs = runner.prepare_mm_inputs({"req0": [0]}) + + assert mm_hashes == [] and mm_kwargs == [] + assert runner.encoder_cache.encoder_outputs["hash_pe"] is sentinel + + +def test_execute_mm_encoder_skips_encoder_for_prompt_embeds_only(): + """A batch of nothing but prompt_embeds must not invoke the encoder.""" + prompt_embeds = torch.ones(3, HIDDEN) + runner = _make_runner([_pe_feature("hash_pe", prompt_embeds)], cached=[]) + state = _model_state(runner.encoder_cache) + state.encoder_runner.prepare_mm_inputs.side_effect = runner.prepare_mm_inputs + + ModelState.execute_mm_encoder(state, {"req0": [0]}) + + state.encoder_runner.execute_mm_encoder.assert_not_called() + assert torch.equal(runner.encoder_cache.encoder_outputs["hash_pe"], prompt_embeds) + + def test_encoder_timing_stats_registry(): runner = _make_runner([], []) runner.enable_timing = True diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 8bbe15bb2a4e..6eaa074cc410 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -68,6 +68,16 @@ DEVICE_TYPE = current_platform.device_type +@pytest.fixture(autouse=True) +def _restore_default_dtype(): + """Several tests here set the process-wide default dtype to float16 and + previously leaked it, corrupting later float-sensitive tests in the same + pytest process (torch.randn silently produced fp16).""" + old = torch.get_default_dtype() + yield + torch.set_default_dtype(old) + + def initialize_kv_cache(runner: GPUModelRunner): """ Only perform necessary steps in GPUModelRunner.initialize_kv_cache() diff --git a/tests/v1/worker/test_prompt_embeds_state.py b/tests/v1/worker/test_prompt_embeds_state.py new file mode 100644 index 000000000000..1893d0f3470e --- /dev/null +++ b/tests/v1/worker/test_prompt_embeds_state.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the Model Runner V2 prompt-embeds overlay (PromptEmbedsState). + +The overlay kernel reads each request's GPU-resident prompt embeddings through +a per-request pointer table and writes the rows scheduled this step into +`inputs_embeds`, honoring chunked prefill (`num_computed_tokens` offset), the +prompt/decode boundary (rows past the embeds length untouched), and the +mixed-mode `prompt_is_token_ids` mask (token-id rows keep the base embedding). +""" + +from dataclasses import dataclass + +import pytest +import torch + +pytest.importorskip("triton") +if not torch.cuda.is_available(): + pytest.skip( + "CUDA required for prompt-embeds overlay tests", allow_module_level=True + ) + +from vllm.v1.worker.gpu.model_states.prompt_embeds import PromptEmbedsState + +HIDDEN = 24 +MAX_NUM_REQS = 8 +DEVICE = torch.device("cuda") + + +@dataclass +class _NewReqData: + req_id: str + prompt_embeds: torch.Tensor | None + prompt_is_token_ids: list[bool] | None = None + + +@dataclass +class _Batch: + num_reqs: int + num_scheduled_tokens: torch.Tensor # np-like, only .max() is used + idx_mapping: torch.Tensor + query_start_loc: torch.Tensor + + +def _make_state() -> PromptEmbedsState: + return PromptEmbedsState(MAX_NUM_REQS, HIDDEN, torch.float32, DEVICE) + + +def _batch(num_scheduled: list[int], idx_mapping: list[int]) -> _Batch: + query_start_loc = [0] + for n in num_scheduled: + query_start_loc.append(query_start_loc[-1] + n) + return _Batch( + num_reqs=len(num_scheduled), + num_scheduled_tokens=torch.tensor(num_scheduled, dtype=torch.int32), + idx_mapping=torch.tensor(idx_mapping, dtype=torch.int64, device=DEVICE), + query_start_loc=torch.tensor(query_start_loc, dtype=torch.int32, device=DEVICE), + ) + + +def _apply( + state: PromptEmbedsState, + batch: _Batch, + num_computed: list[int], + num_tokens: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run the overlay on a fresh base buffer; return (result, base).""" + num_computed_tokens = torch.zeros(MAX_NUM_REQS, dtype=torch.int32, device=DEVICE) + for batch_idx, req_index in enumerate(batch.idx_mapping.tolist()): + num_computed_tokens[req_index] = num_computed[batch_idx] + base = torch.randn(num_tokens, HIDDEN, dtype=torch.float32, device=DEVICE) + inputs_embeds = base.clone() + state.apply(batch, num_computed_tokens, inputs_embeds) + torch.accelerator.synchronize() + return inputs_embeds, base + + +def test_overlay_chunked_prefill_and_decode(): + """Rows within the embeds range come from prompt_embeds at the + num_computed offset; requests without embeds and requests past their + embeds length (decode) keep the base embedding.""" + state = _make_state() + embeds_a = torch.randn(6, HIDDEN, dtype=torch.float32) + embeds_b = torch.randn(5, HIDDEN, dtype=torch.float32) + state.add_request(0, _NewReqData("a", embeds_a)) + state.add_request(1, _NewReqData("b", embeds_b)) + state.add_request(2, _NewReqData("c", None)) + state.apply_staged_writes() + + # a: chunk [2, 6) of its embeds; b: fully decoded; c: no embeds. + batch = _batch(num_scheduled=[4, 1, 3], idx_mapping=[0, 1, 2]) + out, base = _apply(state, batch, num_computed=[2, 7, 1], num_tokens=8) + + torch.testing.assert_close(out[0:4], embeds_a[2:6].to(DEVICE)) + torch.testing.assert_close(out[4:8], base[4:8]) + + +def test_overlay_clamps_to_embeds_length(): + """A window straddling the end of the prompt embeds writes only the + in-range rows (e.g. final prefill chunk + sampled token).""" + state = _make_state() + embeds = torch.randn(4, HIDDEN, dtype=torch.float32) + state.add_request(3, _NewReqData("a", embeds)) + state.apply_staged_writes() + + batch = _batch(num_scheduled=[3], idx_mapping=[3]) + out, base = _apply(state, batch, num_computed=[2], num_tokens=3) + + torch.testing.assert_close(out[0:2], embeds[2:4].to(DEVICE)) + torch.testing.assert_close(out[2:3], base[2:3]) + + +def test_overlay_respects_is_token_ids_mask(): + """Mixed mode: positions marked as real token ids keep the base + embedding; only embed positions are overwritten.""" + state = _make_state() + embeds = torch.randn(5, HIDDEN, dtype=torch.float32) + is_token_ids = [True, False, False, True, False] + state.add_request(0, _NewReqData("a", embeds, is_token_ids)) + state.apply_staged_writes() + + batch = _batch(num_scheduled=[5], idx_mapping=[0]) + out, base = _apply(state, batch, num_computed=[0], num_tokens=5) + + embeds_gpu = embeds.to(DEVICE) + for pos, is_token in enumerate(is_token_ids): + expected = base[pos] if is_token else embeds_gpu[pos] + torch.testing.assert_close(out[pos], expected) + + +def test_index_reuse_clears_stale_entry(): + """A request added at a previously-used index without embeds must not + inherit the prior occupant's pointer-table entry.""" + state = _make_state() + state.add_request(0, _NewReqData("a", torch.randn(4, HIDDEN, dtype=torch.float32))) + # A second live embeds request so the kernel actually launches (the + # overlay is skipped entirely when no request holds embeds). + other = torch.randn(2, HIDDEN, dtype=torch.float32) + state.add_request(1, _NewReqData("other", other)) + state.apply_staged_writes() + state.remove_request("a") + state.add_request(0, _NewReqData("b", None)) + state.apply_staged_writes() + + batch = _batch(num_scheduled=[2, 2], idx_mapping=[0, 1]) + out, base = _apply(state, batch, num_computed=[0, 0], num_tokens=4) + + torch.testing.assert_close(out[0:2], base[0:2]) + torch.testing.assert_close(out[2:4], other.to(DEVICE)) diff --git a/vllm/config/model.py b/vllm/config/model.py index fcdce3ef5858..54d622743afe 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -926,6 +926,12 @@ def validate_model_config_after(self: "ModelConfig") -> "ModelConfig": f"got {type(self.max_model_len).__name__}: {self.max_model_len!r}. " "Example: max_model_len=2048" ) + if self.enable_prompt_embeds and self.is_encoder_decoder: + # No encoder-decoder model accepts `inputs_embeds`; their decoders + # embed `input_ids` internally. + raise ValueError( + "--enable-prompt-embeds is not supported with encoder-decoder models." + ) return self def _resolve_mm_device_do_normalize( diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index a2d6b293bb12..5faf09fc347e 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -2449,9 +2449,6 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]: ): unsupported.append("custom logits processors") - if model_config is not None and model_config.enable_prompt_embeds: - unsupported.append("prompt embeds") - if self.cache_config.kv_sharing_fast_prefill: # Will be added by https://github.com/vllm-project/vllm/pull/35045 unsupported.append("KV sharing fast prefill") diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 0fc32b99d8e9..52a3824860ea 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -776,7 +776,7 @@ def __init__( ) -> None: super().__init__(vllm_config, model, encoder_cache, device) - # Per-step MM data produced by get_mm_embeddings and consumed by + # Per-step MM data produced by prepare_inputs_embeds and consumed by # prepare_inputs. Stored as raw (mm_embeds, is_mm_embed) so that # prepare_inputs can call embed_input_ids directly into the # persistent _inputs_embeds_buf, avoiding the intermediate copy @@ -874,7 +874,7 @@ def remove_request(self, req_id: str) -> None: if idx is not None: self.diffusion_states.remove_request(idx) - def get_mm_embeddings( + def prepare_inputs_embeds( self, scheduled_encoder_inputs: dict[str, list[int]], input_batch: InputBatch, diff --git a/vllm/model_executor/models/longcat_flash_ngram.py b/vllm/model_executor/models/longcat_flash_ngram.py index 5aaa1aad9bb8..41f913cdfec1 100644 --- a/vllm/model_executor/models/longcat_flash_ngram.py +++ b/vllm/model_executor/models/longcat_flash_ngram.py @@ -273,6 +273,9 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: class LongcatNgramModelState(DefaultModelState): + # prepare_inputs builds its own inputs_embeds from n-gram token embeddings. + supports_prompt_embeds = False + """Per-request n-gram token history for LongCat-Flash-Lite. Maintains a small CPU-side per-slot context (last ``n-1`` processed tokens) diff --git a/vllm/v1/core/sched/output.py b/vllm/v1/core/sched/output.py index 8782061356f2..4d11be021ed4 100644 --- a/vllm/v1/core/sched/output.py +++ b/vllm/v1/core/sched/output.py @@ -68,6 +68,14 @@ def from_request( prefill_token_ids=prefill_token_ids, ) + @property + def prompt_len(self) -> int: + if self.prompt_token_ids is not None: + return len(self.prompt_token_ids) + if self.prompt_embeds is not None: + return self.prompt_embeds.shape[0] + return 0 + def __repr__(self) -> str: prompt_embeds_shape = ( self.prompt_embeds.shape if self.prompt_embeds is not None else None diff --git a/vllm/v1/request.py b/vllm/v1/request.py index 0b969c991d94..44de9ad1f7c8 100644 --- a/vllm/v1/request.py +++ b/vllm/v1/request.py @@ -142,11 +142,19 @@ def __init__( prompt_token_ids, prompt_embeds ) self._output_token_ids: list[int] = [] - self._all_token_ids: list[int] = ( - self.prompt_token_ids.copy() - if self.prompt_token_ids is not None - else [0] * self.num_prompt_tokens - ) + if self.prompt_token_ids is None: + self._all_token_ids: list[int] = [0] * self.num_prompt_tokens + elif self.prompt_is_token_ids is None: + self._all_token_ids = self.prompt_token_ids.copy() + else: + # Mixed-mode prompt: positions covered by prompt_embeds hold a sentinel + # special token id that may lie outside the embedding. Zero them, matching + # the no-token-ids case above, so embedding gathers over these placeholder + # ids stay in bounds; the actual inputs come from prompt_embeds. + self._all_token_ids = [ + t if is_tok else 0 + for t, is_tok in zip(self.prompt_token_ids, self.prompt_is_token_ids) + ] # Used in async scheduling. self.num_output_placeholders = 0 diff --git a/vllm/v1/worker/gpu/mm/encoder_runner.py b/vllm/v1/worker/gpu/mm/encoder_runner.py index a1c79a8eb7b7..49b4adee185b 100644 --- a/vllm/v1/worker/gpu/mm/encoder_runner.py +++ b/vllm/v1/worker/gpu/mm/encoder_runner.py @@ -18,7 +18,7 @@ group_and_batch_mm_kwargs, set_mm_embedding_modality, ) -from vllm.utils.torch_utils import PIN_MEMORY +from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache from vllm.v1.worker.utils import ( EncoderTimingStats, @@ -91,6 +91,17 @@ def prepare_mm_inputs( continue if mm_feature.identifier in self.encoder_cache.encoder_outputs: continue + if mm_feature.modality == "prompt_embeds": + # Passthrough modality: the tensor is already in the + # model's embedding space, so no encoder runs. Cache it + # directly so gather_mm_embeddings splices it via the + # standard is_mm_embed path. + embeds = mm_feature.data["embedding"].data + assert isinstance(embeds, torch.Tensor) + self.encoder_cache.encoder_outputs[mm_feature.identifier] = ( + async_tensor_h2d(embeds, device=self.device) + ) + continue mm_hashes.append(mm_feature.identifier) mm_kwargs.append((mm_feature.modality, mm_feature.data)) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index aff0c09cab58..77ae59f0ee38 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -46,6 +46,7 @@ initialize_mamba_ssu_backend, ) from vllm.model_executor.model_loader import get_model_loader +from vllm.model_executor.models.interfaces import requires_raw_input_tokens from vllm.model_executor.offloader import ( create_offloader, get_offloader, @@ -231,6 +232,9 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs( self.model_config ) + self.uses_inputs_embeds = ( + self.supports_mm_inputs or self.model_config.enable_prompt_embeds + ) self.encoder_cache = None if self.supports_mm_inputs and self.is_first_pp_rank: self.encoder_cache = EncoderCache() @@ -956,7 +960,6 @@ def update_pp_decode_requests(self): def add_requests(self, scheduler_output: SchedulerOutput) -> None: for new_req_data in scheduler_output.scheduled_new_reqs: - assert new_req_data.prompt_token_ids is not None assert new_req_data.prefill_token_ids is not None req_id = new_req_data.req_id @@ -965,7 +968,7 @@ def add_requests(self, scheduler_output: SchedulerOutput) -> None: # with the updated prompt_token_ids and mm_features. self._remove_request(req_id) - prompt_len = len(new_req_data.prompt_token_ids) + prompt_len = new_req_data.prompt_len sampling_params = new_req_data.sampling_params self.req_states.add_request( req_id=req_id, @@ -980,6 +983,7 @@ def add_requests(self, scheduler_output: SchedulerOutput) -> None: if self.pooling_runner is not None: assert new_req_data.pooling_params is not None + assert new_req_data.prompt_token_ids is not None self.pooling_runner.add_request( req_id, req_index, @@ -1550,17 +1554,17 @@ def execute_model( input_ids = input_batch.input_ids inputs_embeds = None ec_connector_output = None - if self.supports_mm_inputs and self.is_first_pp_rank: - # Run MM encoder (if needed) and get multimodal embeddings. - # Only first PP rank prepares multimodal embeddings. + if self.uses_inputs_embeds and self.is_first_pp_rank: + # Prepare inputs_embeds (MM encoder outputs and/or prompt_embeds + # overlay). Only first PP rank prepares them. if dummy_run: - # Obtain mm embeddings of correct shape for compiled model. + # Obtain embeddings of correct shape for compiled model. inputs_embeds = self.model_state.dummy_inputs_embeds( input_batch.num_tokens_after_padding ) else: scheduled_encoder_inputs = scheduler_output.scheduled_encoder_inputs - if self.lora_config is not None: + if self.supports_mm_inputs and self.lora_config is not None: set_active_mm_loras( model=self.model, lora_manager=self.lora_manager, @@ -1574,16 +1578,16 @@ def execute_model( ) as ec_connector_output: if self.is_encoder_only: # Encode and publish, nothing else: this instance runs no - # language model, so the gather inside get_mm_embeddings + # language model, so the gather inside prepare_inputs_embeds # would build an inputs_embeds nobody reads -- and it # raises "Encoder cache miss" for any scheduled item this # instance did not encode, taking the engine down with it. self.model_state.execute_mm_encoder(scheduled_encoder_inputs) else: - inputs_embeds = self.model_state.get_mm_embeddings( + inputs_embeds = self.model_state.prepare_inputs_embeds( scheduled_encoder_inputs, input_batch, self.req_states ) - if inputs_embeds is not None and not self.model.requires_raw_input_tokens: + if inputs_embeds is not None and not requires_raw_input_tokens(self.model): input_ids = None if self.is_encoder_only: diff --git a/vllm/v1/worker/gpu/model_states/__init__.py b/vllm/v1/worker/gpu/model_states/__init__.py index ed95d8285684..ceb1facb44f9 100644 --- a/vllm/v1/worker/gpu/model_states/__init__.py +++ b/vllm/v1/worker/gpu/model_states/__init__.py @@ -7,6 +7,7 @@ from vllm.model_executor.layers.attention import Attention, CrossAttention from vllm.v1.attention.backend import AttentionType from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache +from vllm.v1.worker.gpu.model_states.interface import ModelState def init_model_state( @@ -14,11 +15,22 @@ def init_model_state( model: nn.Module, encoder_cache: EncoderCache | None, device: torch.device, -): +) -> ModelState: + cls = resolve_model_state_cls(vllm_config, model) + + # Reject enable_prompt_embeds for states that would silently ignore it. + if vllm_config.model_config.enable_prompt_embeds and not cls.supports_prompt_embeds: + raise ValueError(f"--enable-prompt-embeds not supported with {cls.__name__}.") + + return cls(vllm_config, model, encoder_cache, device) + + +def resolve_model_state_cls( + vllm_config: VllmConfig, model: nn.Module +) -> type[ModelState]: # Let the model provide its own ModelState if it defines one. if hasattr(model, "get_model_state_cls"): - cls = model.get_model_state_cls() - return cls(vllm_config, model, encoder_cache, device) + return model.get_model_state_cls() # Cross-attention encoder-decoder models (Whisper, CohereASR, NemotronParse, ...) if any(isinstance(m, CrossAttention) for m in model.modules()): @@ -26,7 +38,7 @@ def init_model_state( EncoderDecoderModelState, ) - return EncoderDecoderModelState(vllm_config, model, encoder_cache, device) + return EncoderDecoderModelState # Encoder-only attention is non-causal and needs no KV cache. if any( @@ -35,13 +47,13 @@ def init_model_state( ): from vllm.v1.worker.gpu.model_states.encoder_only import EncoderOnlyModelState - return EncoderOnlyModelState(vllm_config, model, encoder_cache, device) + return EncoderOnlyModelState if vllm_config.model_config.is_hybrid or vllm_config.model_config.is_attention_free: from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState - return MambaHybridModelState(vllm_config, model, encoder_cache, device) + return MambaHybridModelState from vllm.v1.worker.gpu.model_states.default import DefaultModelState - return DefaultModelState(vllm_config, model, encoder_cache, device) + return DefaultModelState diff --git a/vllm/v1/worker/gpu/model_states/default.py b/vllm/v1/worker/gpu/model_states/default.py index 31802c017bca..e1b6d7d3de34 100644 --- a/vllm/v1/worker/gpu/model_states/default.py +++ b/vllm/v1/worker/gpu/model_states/default.py @@ -18,11 +18,14 @@ from vllm.v1.worker.gpu.mm.rope import get_rope_state from vllm.v1.worker.gpu.model_states.interface import ModelState from vllm.v1.worker.gpu.model_states.mm_pruning import maybe_create_mm_pruner +from vllm.v1.worker.gpu.model_states.prompt_embeds import PromptEmbedsState from vllm.v1.worker.gpu.states import RequestState from vllm.v1.worker.utils import AttentionGroup class DefaultModelState(ModelState): + supports_prompt_embeds = True + def __init__( self, vllm_config: VllmConfig, @@ -32,6 +35,18 @@ def __init__( ): super().__init__(vllm_config, model, encoder_cache, device) + self.prompt_embeds_state: PromptEmbedsState | None = None + if self.model_config.enable_prompt_embeds: + self.prompt_embeds_state = PromptEmbedsState( + self.max_num_reqs, self.inputs_embeds_size, self.dtype, self.device + ) + if not self.supports_mm_inputs: + # Persistent buffer analogous to encoder_runner.inputs_embeds. + embeds_buffer_size = (self.max_num_tokens, self.inputs_embeds_size) + self.inputs_embeds = torch.zeros( + embeds_buffer_size, dtype=self.dtype, device=self.device + ) + self.rope_state = get_rope_state( self.model_config, model, @@ -49,42 +64,69 @@ def __init__( def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: if self.rope_state is not None: assert new_req_data.prefill_token_ids is not None + # `prompt_embeds` is a passthrough modality with no grid info, but + # M-RoPE assumes per-feature grids. Filter it out. + mm_features = [ + f for f in new_req_data.mm_features if f.modality != "prompt_embeds" + ] self.rope_state.init_prefill_positions( req_index, self.model, new_req_data.prefill_token_ids, - mm_features=new_req_data.mm_features, + mm_features=mm_features, ) + if self.prompt_embeds_state is not None: + self.prompt_embeds_state.add_request(req_index, new_req_data) + + def remove_request(self, req_id: str) -> None: + if self.prompt_embeds_state is not None: + self.prompt_embeds_state.remove_request(req_id) def apply_staged_writes(self) -> None: if self.rope_state is not None: self.rope_state.apply_staged_writes() + if self.prompt_embeds_state is not None: + self.prompt_embeds_state.apply_staged_writes() def dummy_inputs_embeds(self, num_tokens: int) -> torch.Tensor: """Pre-allocated inputs_embeds buffer for dummy runs (contents unused).""" - return self.encoder_runner.inputs_embeds[:num_tokens] + if self.supports_mm_inputs: + return self.encoder_runner.inputs_embeds[:num_tokens] + return self.inputs_embeds[:num_tokens] - def get_mm_embeddings( + def prepare_inputs_embeds( self, scheduled_encoder_inputs: dict[str, list[int]], input_batch: InputBatch, req_states: RequestState, ) -> torch.Tensor: - self.execute_mm_encoder(scheduled_encoder_inputs) - - mm_embeds, is_mm_embed = super().gather_mm_embeddings(input_batch) - if self.mm_pruner is not None and mm_embeds: - # EVS: recompute mrope positions for pruned media. - mm_embeds = self.mm_pruner.recompute(mm_embeds, input_batch, req_states) - # We must flush the staged rope updates for prepare_inputs() to pick up. - self.apply_staged_writes() - # Use unpadded input_ids to match is_mm_embed size (num_tokens). # input_batch.input_ids may be padded for CUDA graphs. input_ids_unpadded = input_batch.input_ids[: input_batch.num_tokens] - inputs_embeds = self.encoder_runner.get_inputs_embeds( - input_ids_unpadded, mm_embeds, is_mm_embed - ) + + if self.supports_mm_inputs: + self.execute_mm_encoder(scheduled_encoder_inputs) + + mm_embeds, is_mm_embed = super().gather_mm_embeddings(input_batch) + if self.mm_pruner is not None and mm_embeds: + # EVS: recompute mrope positions for pruned media. + mm_embeds = self.mm_pruner.recompute(mm_embeds, input_batch, req_states) + # We must flush the staged rope updates for prepare_inputs() to pick up. + self.apply_staged_writes() + + inputs_embeds = self.encoder_runner.get_inputs_embeds( + input_ids_unpadded, mm_embeds, is_mm_embed + ) + else: + input_embeddings = self.model.embed_input_ids(input_ids_unpadded) + self.inputs_embeds[: input_embeddings.shape[0]] = input_embeddings + inputs_embeds = self.inputs_embeds + + if self.prompt_embeds_state is not None: + self.prompt_embeds_state.apply( + input_batch, req_states.num_computed_tokens.gpu, inputs_embeds + ) + return inputs_embeds[: input_batch.num_tokens_after_padding] def gather_mm_embeddings( @@ -115,9 +157,8 @@ def prepare_inputs( def prepare_dummy_inputs(self, num_reqs: int, num_tokens: int) -> dict[str, Any]: model_inputs = {} - if self.supports_mm_inputs: - inputs_embeds = self.encoder_runner.inputs_embeds[:num_tokens] - model_inputs["inputs_embeds"] = inputs_embeds + if self.supports_mm_inputs or self.prompt_embeds_state is not None: + model_inputs["inputs_embeds"] = self.dummy_inputs_embeds(num_tokens) if self.rope_state is not None: model_inputs["positions"] = self.rope_state.get_positions(num_tokens) return model_inputs diff --git a/vllm/v1/worker/gpu/model_states/encoder_decoder.py b/vllm/v1/worker/gpu/model_states/encoder_decoder.py index 7a461466d0aa..93c738fb226c 100644 --- a/vllm/v1/worker/gpu/model_states/encoder_decoder.py +++ b/vllm/v1/worker/gpu/model_states/encoder_decoder.py @@ -54,6 +54,10 @@ def __init__( device: torch.device, ) -> None: assert encoder_cache is not None + if vllm_config.model_config.enable_prompt_embeds: + raise ValueError( + "--enable-prompt-embeds is not supported with encoder-decoder models." + ) super().__init__(vllm_config, model, encoder_cache, device) self.max_encoder_len = getattr( @@ -67,7 +71,7 @@ def __init__( self.encoder_outputs: list[torch.Tensor] = [] - def get_mm_embeddings( + def prepare_inputs_embeds( self, scheduled_encoder_inputs: dict[str, list[int]], input_batch: InputBatch, diff --git a/vllm/v1/worker/gpu/model_states/encoder_only.py b/vllm/v1/worker/gpu/model_states/encoder_only.py index cf8918317964..10ba0a6a396e 100644 --- a/vllm/v1/worker/gpu/model_states/encoder_only.py +++ b/vllm/v1/worker/gpu/model_states/encoder_only.py @@ -34,6 +34,9 @@ class EncoderOnlyModelState(DefaultModelState): the normal KV-backed path untouched. """ + # The V2 pooling path is not wired for prompt embeds. + supports_prompt_embeds = False + def __init__( self, vllm_config: VllmConfig, diff --git a/vllm/v1/worker/gpu/model_states/interface.py b/vllm/v1/worker/gpu/model_states/interface.py index 980ec1dd8e1d..cd1b5b64da7d 100644 --- a/vllm/v1/worker/gpu/model_states/interface.py +++ b/vllm/v1/worker/gpu/model_states/interface.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from abc import ABC, abstractmethod -from typing import Any, cast +from typing import Any, ClassVar, cast import torch import torch.nn as nn @@ -43,6 +43,9 @@ def get_extra_attn_kwargs( class ModelState(ABC): + supports_prompt_embeds: ClassVar[bool] = False + """Whether this state implements user-provided prompt embeddings.""" + def __init__( self, vllm_config: VllmConfig, @@ -153,12 +156,13 @@ def postprocess_state( return None @abstractmethod - def get_mm_embeddings( + def prepare_inputs_embeds( self, scheduled_encoder_inputs: dict[str, list[int]], input_batch: InputBatch, req_states: RequestState, ) -> torch.Tensor | None: + """Prepare the ``inputs_embeds`` tensor for the current forward pass.""" raise NotImplementedError def dummy_inputs_embeds(self, num_tokens: int) -> torch.Tensor | None: diff --git a/vllm/v1/worker/gpu/model_states/prompt_embeds.py b/vllm/v1/worker/gpu/model_states/prompt_embeds.py new file mode 100644 index 000000000000..251558f9c1ba --- /dev/null +++ b/vllm/v1/worker/gpu/model_states/prompt_embeds.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch + +from vllm.triton_utils import tl, triton +from vllm.utils.torch_utils import async_tensor_h2d +from vllm.v1.core.sched.output import NewRequestData +from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor +from vllm.v1.worker.gpu.input_batch import InputBatch + +TOKEN_BLOCK = 16 + + +class PromptEmbedsState: + """GPU-side state for user-provided prompt embeddings. + + Each request's embeddings are copied to the GPU once at `add_request`, + off the per-step hot path. A per-request pointer table (UVA) then lets a + single triton kernel overlay all scheduled prompt-embeds rows onto + `inputs_embeds` each step, with no python loops or per-request H2D copies. + """ + + def __init__( + self, + max_num_reqs: int, + hidden_size: int, + dtype: torch.dtype, + device: torch.device, + ): + self.hidden_size = hidden_size + self.dtype = dtype + self.device = device + + # req_id -> (embeds, is_token_ids mask or None). Holds the references + # that keep the pointer table below valid. + self.gpu_tensors: dict[str, tuple[torch.Tensor, torch.Tensor | None]] = {} + + # Indexed by req_state index. Stale entries after removal are + # harmless: add_request rewrites all fields for every index it claims. + self.embeds_ptrs = UvaBackedTensor(max_num_reqs, dtype=torch.int64) + self.mask_ptrs = UvaBackedTensor(max_num_reqs, dtype=torch.int64) + self.embeds_lens = UvaBackedTensor(max_num_reqs, dtype=torch.int32) + + def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: + prompt_embeds = new_req_data.prompt_embeds + if prompt_embeds is None: + self.gpu_tensors.pop(new_req_data.req_id, None) + self.embeds_lens.np[req_index] = 0 + return + + embeds = async_tensor_h2d(prompt_embeds, device=self.device, dtype=self.dtype) + embeds = embeds.contiguous() + is_token_ids = new_req_data.prompt_is_token_ids + mask = None + if is_token_ids is not None: + mask = async_tensor_h2d(is_token_ids, device=self.device, dtype=torch.uint8) + self.gpu_tensors[new_req_data.req_id] = (embeds, mask) + self.embeds_ptrs.np[req_index] = embeds.data_ptr() + self.mask_ptrs.np[req_index] = 0 if mask is None else mask.data_ptr() + self.embeds_lens.np[req_index] = embeds.shape[0] + + def remove_request(self, req_id: str) -> None: + self.gpu_tensors.pop(req_id, None) + + def apply_staged_writes(self) -> None: + self.embeds_ptrs.copy_to_uva() + self.mask_ptrs.copy_to_uva() + self.embeds_lens.copy_to_uva() + + def apply( + self, + input_batch: InputBatch, + num_computed_tokens: torch.Tensor, + inputs_embeds: torch.Tensor, + ) -> None: + """Overlay prompt embeddings onto `inputs_embeds` for the batch.""" + if not self.gpu_tensors: + return + # The kernel reinterprets raw source pointers as inputs_embeds' dtype. + assert inputs_embeds.dtype == self.dtype + num_reqs = input_batch.num_reqs + max_query_len = int(input_batch.num_scheduled_tokens.max()) + grid = (num_reqs, triton.cdiv(max_query_len, TOKEN_BLOCK)) + _apply_prompt_embeds_kernel[grid]( + inputs_embeds, + inputs_embeds.stride(0), + self.embeds_ptrs.gpu, + self.mask_ptrs.gpu, + self.embeds_lens.gpu, + input_batch.idx_mapping, + input_batch.query_start_loc, + num_computed_tokens, + self.hidden_size, + TOKEN_BLOCK=TOKEN_BLOCK, + BLOCK_SIZE=1024, + ) + + +@triton.jit +def _apply_prompt_embeds_kernel( + inputs_embeds_ptr, + inputs_embeds_stride, + embeds_ptrs_ptr, # int64 [max_num_reqs], device pointers (0-len = unused) + mask_ptrs_ptr, # int64 [max_num_reqs], 0 = no is-token-ids mask + embeds_lens_ptr, # int32 [max_num_reqs] + idx_mapping_ptr, + query_start_loc_ptr, + num_computed_tokens_ptr, + hidden_size: tl.constexpr, + TOKEN_BLOCK: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + batch_idx = tl.program_id(0) + req_state_idx = tl.load(idx_mapping_ptr + batch_idx) + embeds_len = tl.load(embeds_lens_ptr + req_state_idx) + num_computed = tl.load(num_computed_tokens_ptr + req_state_idx) + if num_computed >= embeds_len: + # No prompt embeds for this request, or they are fully consumed. + return + + query_start = tl.load(query_start_loc_ptr + batch_idx) + query_end = tl.load(query_start_loc_ptr + batch_idx + 1) + num_rows = tl.minimum(query_end - query_start, embeds_len - num_computed) + + t_start = tl.program_id(1) * TOKEN_BLOCK + if t_start >= num_rows: + return + + src_ptr = tl.load(embeds_ptrs_ptr + req_state_idx).to( + tl.pointer_type(inputs_embeds_ptr.dtype.element_ty) + ) + mask_int = tl.load(mask_ptrs_ptr + req_state_idx) + mask_ptr = mask_int.to(tl.pointer_type(tl.int8)) + + for t_offset in tl.static_range(TOKEN_BLOCK): + t = t_start + t_offset + if t < num_rows: + src_row = (num_computed + t).to(tl.int64) + is_token_id = 0 + if mask_int != 0: + is_token_id = tl.load(mask_ptr + src_row).to(tl.int32) + if is_token_id == 0: + dst_row = (query_start + t).to(tl.int64) + for h in tl.range(0, hidden_size, BLOCK_SIZE): + offs = h + tl.arange(0, BLOCK_SIZE) + h_mask = offs < hidden_size + row = tl.load(src_ptr + src_row * hidden_size + offs, mask=h_mask) + tl.store( + inputs_embeds_ptr + dst_row * inputs_embeds_stride + offs, + row, + mask=h_mask, + ) From 60c3a31b12b0a735baa79684885391747771ec64 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Mon, 17 Aug 2026 18:09:12 -0500 Subject: [PATCH 063/839] [CI][AMD] Improve Kubernetes failure diagnostics (#52264) Signed-off-by: Andreas Karatzas Signed-off-by: Andreas Karatzas Co-authored-by: OpenAI Codex --- .../scripts/hardware_ci/run-amd-test.sh | 558 ++++++++++++++---- 1 file changed, 448 insertions(+), 110 deletions(-) diff --git a/.buildkite/scripts/hardware_ci/run-amd-test.sh b/.buildkite/scripts/hardware_ci/run-amd-test.sh index 86e4b34869bb..ff1fd0dd833c 100755 --- a/.buildkite/scripts/hardware_ci/run-amd-test.sh +++ b/.buildkite/scripts/hardware_ci/run-amd-test.sh @@ -44,11 +44,14 @@ amd_diagnostics_dir="${VLLM_CI_DIAGNOSTICS_DIR:-artifacts/amd-gpu-diagnostics}" amd_diagnostics_checkout_root="${BUILDKITE_BUILD_CHECKOUT_PATH:-$(pwd -P)}" amd_diagnostics_execution_mode="${VLLM_CI_EXECUTION_MODE:-single-node}" amd_diagnostics_test_group="${VLLM_TEST_GROUP_NAME:-${BUILDKITE_STEP_KEY:-unknown}}" -amd_diagnostics_k8s_node_name="${VLLM_CI_K8S_NODE_NAME:-}" +amd_diagnostics_workspace="${VLLM_CI_WORKSPACE:-/vllm-workspace}" +amd_diagnostics_pod_name="${VLLM_CI_K8S_POD_NAME:-${POD_NAME:-${HOSTNAME:-}}}" +amd_diagnostics_k8s_namespace="${VLLM_CI_K8S_NAMESPACE:-${POD_NAMESPACE:-unknown}}" +amd_diagnostics_k8s_node_name="${VLLM_CI_K8S_NODE_NAME:-${NODE_NAME:-}}" amd_diagnostics_collected=0 +amd_diagnostics_memory_events_path="" amd_diagnostics_probe_budget_seconds=25 amd_diagnostics_command_timeout_seconds=5 -amd_diagnostics_upload_timeout_seconds=20 if [[ " ${PYTEST_ADDOPTS:-} " != *" --color"* ]]; then PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }--color=yes" fi @@ -101,6 +104,8 @@ clear_ci_orchestration_env() { VLLM_CI_ARTIFACT_GLOB \ VLLM_CI_ARTIFACT_CHECKSUM_GLOB \ VLLM_CI_EXPECTED_GPU_COUNT \ + VLLM_CI_K8S_POD_NAME \ + VLLM_CI_K8S_NAMESPACE \ VLLM_CI_K8S_NODE_NAME \ VLLM_CI_USE_ARTIFACTS \ VLLM_CI_RESULTS_ROOT \ @@ -577,7 +582,24 @@ is_multi_node() { return 1 } -run_amd_diagnostic() { +append_failure_diagnostic_section() { + local log_file=$1 + local title=$2 + + printf '\n===============================================================================\n' \ + >> "${log_file}" + printf '%s\n' "${title}" >> "${log_file}" + printf '===============================================================================\n' \ + >> "${log_file}" +} + +append_failure_diagnostic_note() { + local log_file=$1 + shift + printf '%s\n' "$@" >> "${log_file}" +} + +run_failure_diagnostic() { local log_file=$1 local probe_deadline=$2 shift 2 @@ -589,23 +611,285 @@ run_amd_diagnostic() { printf '\n$' printf ' %q' "$@" printf '\n' - } | tee -a "${log_file}" + } >> "${log_file}" + if ! command -v timeout >/dev/null 2>&1; then + echo "[timeout unavailable; bounded diagnostic command skipped]" \ + >> "${log_file}" + return 0 + fi if [[ "${remaining_seconds}" -le 0 ]]; then echo "[diagnostic probe budget exhausted; command skipped]" \ - | tee -a "${log_file}" + >> "${log_file}" return 0 fi if [[ "${remaining_seconds}" -lt "${command_timeout_seconds}" ]]; then command_timeout_seconds=${remaining_seconds} fi - timeout --kill-after=1s "${command_timeout_seconds}s" "$@" 2>&1 \ - | tee -a "${log_file}" - command_status=${PIPESTATUS[0]} + timeout --kill-after=1s "${command_timeout_seconds}s" "$@" \ + >> "${log_file}" 2>&1 + command_status=$? if [[ "${command_status}" -ne 0 ]]; then echo "[diagnostic command exited with status ${command_status}]" \ - | tee -a "${log_file}" + >> "${log_file}" + fi + return 0 +} + +append_failure_diagnostic_file() { + local log_file=$1 + local label=$2 + local path=$3 + local command_status=0 + + printf '\n%s:\n' "${label}" >> "${log_file}" + printf '$ cat %q\n' "${path}" >> "${log_file}" + if [[ ! -r "${path}" ]]; then + echo "[file unavailable]" >> "${log_file}" + return 0 + fi + + cat "${path}" >> "${log_file}" 2>&1 + command_status=$? + if [[ "${command_status}" -ne 0 ]]; then + echo "[diagnostic file read exited with status ${command_status}]" \ + >> "${log_file}" + fi + return 0 +} + +decode_mountinfo_path() { + local value=$1 + + value=${value//\\040/ } + value=${value//\\011/$'\t'} + value=${value//\\012/$'\n'} + value=${value//\\134/\\} + printf '%s\n' "${value}" +} + +resolve_current_cgroup_v2_dir() { + local cgroup_path="" + local mount_root="" + local mount_point="" + local relative_path="" + local candidate="" + + cgroup_path=$(awk -F: '$2 == "" { print $3; exit }' \ + /proc/self/cgroup 2>/dev/null) + [[ -n "${cgroup_path}" ]] || return 1 + + while IFS=$'\t' read -r mount_root mount_point; do + mount_root=$(decode_mountinfo_path "${mount_root}") + mount_point=$(decode_mountinfo_path "${mount_point}") + if [[ "${cgroup_path}" == "/" ]]; then + relative_path="" + elif [[ "${mount_root}" == "/" ]]; then + relative_path="${cgroup_path}" + elif [[ "${cgroup_path}" == "${mount_root}" ]]; then + relative_path="" + elif [[ "${cgroup_path}" == "${mount_root}/"* ]]; then + relative_path="/${cgroup_path#"${mount_root}/"}" + else + continue + fi + + candidate="${mount_point%/}${relative_path}" + if [[ -d "${candidate}" ]]; then + printf '%s\n' "${candidate}" + return 0 + fi + done < <( + awk ' + { + for (i = 6; i <= NF; i++) { + if ($i == "-" && $(i + 1) == "cgroup2") { + print $4 "\t" $5 + break + } + } + } + ' /proc/self/mountinfo 2>/dev/null + ) + return 1 +} + +collect_cgroup_files() { + local log_file=$1 + local title=$2 + local cgroup_dir=$3 + shift 3 + local cgroup_file="" + local files_collected=0 + + append_failure_diagnostic_section "${log_file}" "${title}" + if [[ -z "${cgroup_dir}" || ! -d "${cgroup_dir}" ]]; then + append_failure_diagnostic_note "${log_file}" \ + "resolved_cgroup_path=unavailable" + return 0 + fi + + append_failure_diagnostic_note "${log_file}" \ + "resolved_cgroup_path=${cgroup_dir}" + for cgroup_file in "$@"; do + if [[ -r "${cgroup_dir}/${cgroup_file}" ]]; then + append_failure_diagnostic_file "${log_file}" "${cgroup_file}" \ + "${cgroup_dir}/${cgroup_file}" + files_collected=$((files_collected + 1)) + fi + done + if [[ "${files_collected}" -eq 0 ]]; then + append_failure_diagnostic_note "${log_file}" \ + "[no requested cgroup files were readable]" + fi +} + +collect_cgroup_diagnostics() { + local log_file=$1 + local cgroup_dir="" + + append_failure_diagnostic_section "${log_file}" "Cgroup membership" + append_failure_diagnostic_file "${log_file}" "current_process" \ + /proc/self/cgroup + append_failure_diagnostic_file "${log_file}" "pid_namespace_init_process" \ + /proc/1/cgroup + + cgroup_dir=$(resolve_current_cgroup_v2_dir 2>/dev/null || true) + collect_cgroup_files "${log_file}" \ + "Cgroup v2 current-process resource state" "${cgroup_dir}" \ + cgroup.type cgroup.events \ + memory.current memory.peak memory.max memory.high memory.low memory.min \ + memory.oom.group memory.swap.current memory.swap.max \ + memory.events memory.events.local memory.swap.events memory.stat \ + memory.pressure \ + cpu.max cpu.max.burst cpu.weight cpu.stat cpu.pressure \ + cpuset.cpus cpuset.cpus.effective cpuset.mems cpuset.mems.effective \ + pids.current pids.max pids.events pids.events.local \ + io.stat io.max io.weight io.pressure + if [[ -n "${cgroup_dir}" && -r "${cgroup_dir}/memory.events" ]]; then + amd_diagnostics_memory_events_path="${cgroup_dir}/memory.events" + fi +} + +collect_process_diagnostics() { + local log_file=$1 + local probe_deadline=$2 + local shell_proc_dir="/proc/$$" + + append_failure_diagnostic_section "${log_file}" \ + "Runner process constraints and PID-namespace-visible processes" + append_failure_diagnostic_file "${log_file}" "diagnostic_shell_limits" \ + "${shell_proc_dir}/limits" + if command -v awk >/dev/null 2>&1; then + # shellcheck disable=SC2016 # The expression is evaluated by awk. + run_failure_diagnostic "${log_file}" "${probe_deadline}" awk ' + $1 ~ /^(Pid:|PPid:|Threads:|VmPeak:|VmSize:|VmHWM:|VmRSS:|RssAnon:|RssFile:|RssShmem:|Cpus_allowed_list:|Mems_allowed_list:|voluntary_ctxt_switches:|nonvoluntary_ctxt_switches:)$/ { + print + } + ' "${shell_proc_dir}/status" + fi + if command -v nproc >/dev/null 2>&1; then + run_failure_diagnostic "${log_file}" "${probe_deadline}" nproc + fi + run_failure_diagnostic "${log_file}" "${probe_deadline}" \ + /bin/bash -c 'ulimit -a' + if command -v ps >/dev/null 2>&1 && command -v head >/dev/null 2>&1; then + append_failure_diagnostic_note "${log_file}" \ + "process_snapshot=top_100_by_rss" + run_failure_diagnostic "${log_file}" "${probe_deadline}" \ + /bin/bash -c \ + 'ps -eo pid,ppid,stat,etimes,nlwp,pcpu,rss,vsz,comm --sort=-rss | head -n 101' + fi + return 0 +} + +collect_mount_diagnostic() { + local log_file=$1 + local probe_deadline=$2 + local label=$3 + local path=$4 + + printf '\n%s (%s):\n' "${label}" "${path}" >> "${log_file}" + if [[ ! -e "${path}" ]]; then + echo "[path unavailable]" >> "${log_file}" + return 0 + fi + if command -v findmnt >/dev/null 2>&1; then + run_failure_diagnostic "${log_file}" "${probe_deadline}" findmnt \ + -n -T "${path}" -o TARGET,FSTYPE,VFS-OPTIONS + fi + if command -v df >/dev/null 2>&1; then + run_failure_diagnostic "${log_file}" "${probe_deadline}" df -h \ + --output=fstype,size,used,avail,pcent,itotal,iused,iavail,ipcent,target \ + -- "${path}" + fi + return 0 +} + +collect_execution_resource_diagnostics() { + local log_file=$1 + local probe_deadline=$2 + + # free and /proc/meminfo generally describe the node from a pod. Cgroups are + # the live source for this container's resource limits, usage, and events. + collect_cgroup_diagnostics "${log_file}" + collect_process_diagnostics "${log_file}" "${probe_deadline}" + + append_failure_diagnostic_section "${log_file}" \ + "Mount-visible filesystem capacity" + append_failure_diagnostic_note "${log_file}" \ + "These values are filesystem views, not Kubernetes ephemeral-storage quotas." + collect_mount_diagnostic "${log_file}" "${probe_deadline}" \ + "checkout" "${amd_diagnostics_checkout_root}" + if is_native_runtime; then + collect_mount_diagnostic "${log_file}" "${probe_deadline}" \ + "native workspace emptyDir" "${amd_diagnostics_workspace}" + else + collect_mount_diagnostic "${log_file}" "${probe_deadline}" \ + "test workspace (if mounted)" "${amd_diagnostics_workspace}" + fi + collect_mount_diagnostic "${log_file}" "${probe_deadline}" \ + "temporary directory" "${TMPDIR:-/tmp}" + if [[ "${TMPDIR:-/tmp}" != "/tmp" ]]; then + collect_mount_diagnostic "${log_file}" "${probe_deadline}" \ + "system temporary directory" /tmp + fi + if is_native_runtime; then + collect_mount_diagnostic "${log_file}" "${probe_deadline}" \ + "shared-memory emptyDir" /dev/shm + else + collect_mount_diagnostic "${log_file}" "${probe_deadline}" \ + "shared-memory mount" /dev/shm + fi + if [[ -n "${HF_HOME:-}" ]]; then + collect_mount_diagnostic "${log_file}" "${probe_deadline}" \ + "Hugging Face cache" "${HF_HOME}" + fi + return 0 +} + +collect_amd_device_nodes() { + local log_file=$1 + local probe_deadline=$2 + local device="" + local device_count=0 + + for device in /dev/kfd /dev/dri/renderD*; do + if [[ ! -e "${device}" ]]; then + continue + fi + device_count=$((device_count + 1)) + if command -v stat >/dev/null 2>&1; then + run_failure_diagnostic "${log_file}" "${probe_deadline}" stat -Lc \ + '%n type=%F mode=%a owner=%u:%g device=%t:%T' -- "${device}" + else + append_failure_diagnostic_note "${log_file}" "device=${device}" + fi + done + if [[ "${device_count}" -eq 0 ]]; then + append_failure_diagnostic_note "${log_file}" \ + "No /dev/kfd or /dev/dri/renderD* device nodes are visible." fi return 0 } @@ -617,30 +901,55 @@ collect_rocm_failure_diagnostics() { local parallel_job="${BUILDKITE_PARALLEL_JOB:-0}" local diagnostics_relative_path="" local diagnostics_path="" - local upload_root="${amd_diagnostics_checkout_root}" - local upload_path="" + local diagnostics_parent="" + local checkout_real="" + local diagnostics_parent_real="" local exit_signal="" local probe_deadline=0 + local runtime="single-node-docker" + local diagnostics_scope="outer-runner-after-test-container-exit" + local identity_label="runner" + local identity_value="${HOSTNAME:-unknown}" + local k8s_pod="unknown" + local k8s_namespace="${amd_diagnostics_k8s_namespace}" + local oom_kill_count="" + local summary_identity_label="Runner" + local -a summary_rows=() if [[ "${amd_diagnostics_collected}" == "1" ]]; then return 0 fi amd_diagnostics_collected=1 - if [[ "${amd_diagnostics_expected_gpu_count}" == "0" ]]; then - echo "Skipping AMD GPU diagnostics for a CPU-only job." - return 0 + if is_native_runtime; then + runtime="native-kubernetes" + diagnostics_scope="current-container-cgroup-and-namespaces" + if [[ -n "${amd_diagnostics_pod_name}" ]]; then + identity_label="pod" + identity_value="${amd_diagnostics_pod_name}" + k8s_pod="${amd_diagnostics_pod_name}" + else + identity_label="container_hostname" + fi + elif [[ "${amd_diagnostics_execution_mode}" == "multi-node" ]]; then + runtime="multi-node-docker" fi - - if ! command -v timeout >/dev/null 2>&1; then - echo "WARNING: timeout is unavailable; skipping failure diagnostics rather than risking a hung job." - return 0 + if [[ "${k8s_namespace}" == "unknown" \ + && -r /var/run/secrets/kubernetes.io/serviceaccount/namespace ]]; then + k8s_namespace=$( + tr -d '\r\n' \ + < /var/run/secrets/kubernetes.io/serviceaccount/namespace 2>/dev/null + ) fi job_id="${job_id//[^A-Za-z0-9_.-]/_}" retry_count="${retry_count//[^A-Za-z0-9_.-]/_}" parallel_job="${parallel_job//[^A-Za-z0-9_.-]/_}" + if [[ "${exit_code}" -gt 128 && "${exit_code}" -le 192 ]]; then + exit_signal=$(kill -l "$((exit_code - 128))" 2>/dev/null || true) + fi + # Buildkite keeps the supplied artifact path. Restrict the configurable # directory to a checkout-relative path so the UI shows a clean artifact key. if [[ -z "${amd_diagnostics_dir}" \ @@ -649,26 +958,32 @@ collect_rocm_failure_diagnostics() { echo "WARNING: ignoring unsafe VLLM_CI_DIAGNOSTICS_DIR" amd_diagnostics_dir="artifacts/amd-gpu-diagnostics" fi - diagnostics_relative_path="${amd_diagnostics_dir}/${job_id}/amd-gpu-diagnostics.log" + diagnostics_relative_path="${amd_diagnostics_dir}/${job_id}/diagnostics.log" diagnostics_path="${amd_diagnostics_checkout_root}/${diagnostics_relative_path}" - upload_path="${diagnostics_relative_path}" - - if ! mkdir -p "$(dirname "${diagnostics_path}")" \ - || ! : > "${diagnostics_path}"; then - diagnostics_path=$(mktemp -t amd-gpu-diagnostics.XXXXXX.log) || { - echo "WARNING: unable to create an AMD GPU diagnostics log." - return 0 - } - upload_root=$(dirname "${diagnostics_path}") - upload_path=$(basename "${diagnostics_path}") - fi - - if [[ "${exit_code}" -gt 128 && "${exit_code}" -le 192 ]]; then - exit_signal=$(kill -l "$((exit_code - 128))" 2>/dev/null || true) + diagnostics_parent=$(dirname "${diagnostics_path}") + checkout_real=$(readlink -m "${amd_diagnostics_checkout_root}" 2>/dev/null || true) + diagnostics_parent_real=$(readlink -m "${diagnostics_parent}" 2>/dev/null || true) + if [[ -z "${checkout_real}" || -z "${diagnostics_parent_real}" \ + || ("${diagnostics_parent_real}" != "${checkout_real}" \ + && "${diagnostics_parent_real}" != "${checkout_real}/"*) ]] \ + || ! mkdir -p "${diagnostics_parent}" \ + || [[ -L "${diagnostics_path}" ]] \ + || ! (set -o noclobber; : > "${diagnostics_path}") 2>/dev/null; then + echo "WARNING: unable to create AMD CI diagnostics at ${diagnostics_relative_path}." + printf '\nAMD CI failure summary\n' + printf '%-22s | %s\n' \ + "Field" "Value" \ + "----------------------" "-----" \ + "Exit code" "${exit_code}" \ + "Signal" "${exit_signal:-none}" \ + "Runtime" "${runtime}" \ + "Pod/container" "${identity_value}" \ + "Diagnostics artifact" "unavailable: ${diagnostics_relative_path}" + return 0 fi { - echo "AMD GPU failure diagnostics" + echo "AMD CI failure diagnostics" echo "timestamp_utc=$(date -u +%Y-%m-%dT%H:%M:%SZ)" echo "exit_code=${exit_code}" echo "exit_signal=${exit_signal:-none}" @@ -680,105 +995,130 @@ collect_rocm_failure_diagnostics() { echo "retry_count=${retry_count}" echo "parallel_job=${parallel_job}" echo "execution_mode=${amd_diagnostics_execution_mode}" + echo "runtime=${runtime}" + echo "diagnostics_scope=${diagnostics_scope}" echo "expected_gpu_count=${amd_diagnostics_expected_gpu_count}" echo "probe_budget_seconds=${amd_diagnostics_probe_budget_seconds}" echo "command_timeout_seconds=${amd_diagnostics_command_timeout_seconds}" - echo "upload_timeout_seconds=${amd_diagnostics_upload_timeout_seconds}" echo "agent_name=${BUILDKITE_AGENT_NAME:-unknown}" + echo "container_hostname=${HOSTNAME:-unknown}" + echo "k8s_pod=${k8s_pod}" + echo "k8s_namespace=${k8s_namespace:-unknown}" echo "k8s_node=${amd_diagnostics_k8s_node_name:-unknown}" echo "kernel=$(uname -srmo 2>/dev/null || echo unknown)" } > "${diagnostics_path}" - echo "--- :rotating_light: AMD GPU failure diagnostics" - cat "${diagnostics_path}" - probe_deadline=$((SECONDS + amd_diagnostics_probe_budget_seconds)) - # Collect the most useful evidence first so auxiliary probes cannot consume - # the shared deadline before GPU health is captured. - echo "AMD GPU diagnostics" | tee -a "${diagnostics_path}" - if command -v amd-smi >/dev/null 2>&1; then - run_amd_diagnostic "${diagnostics_path}" "${probe_deadline}" amd-smi version - # Bus data identifies a card within the public node without publishing its - # persistent UUID, serial number, or process list. - run_amd_diagnostic "${diagnostics_path}" "${probe_deadline}" amd-smi static -b -g all - run_amd_diagnostic "${diagnostics_path}" "${probe_deadline}" amd-smi metric -e -k -P -x -g all - run_amd_diagnostic "${diagnostics_path}" "${probe_deadline}" amd-smi bad-pages -p -r -u -g all - run_amd_diagnostic "${diagnostics_path}" "${probe_deadline}" amd-smi metric -p -t -u -m -v -g all - run_amd_diagnostic "${diagnostics_path}" "${probe_deadline}" amd-smi xgmi -l -g all - elif command -v rocm-smi >/dev/null 2>&1; then - run_amd_diagnostic "${diagnostics_path}" "${probe_deadline}" rocm-smi \ - --showbus --showreplaycount --showrasinfo --showpagesinfo - else - echo "Neither amd-smi nor rocm-smi is available." \ - | tee -a "${diagnostics_path}" - fi + # Cgroup, process, and mount state is fast and most useful for explaining OOM, + # throttling, PID exhaustion, and emptyDir pressure. Capture it before external + # probes can consume the shared deadline. + collect_execution_resource_diagnostics \ + "${diagnostics_path}" "${probe_deadline}" - echo "Node resource diagnostics (runner-visible scope)" \ - | tee -a "${diagnostics_path}" - if command -v df >/dev/null 2>&1; then - echo "checkout_filesystem:" | tee -a "${diagnostics_path}" - if [[ -d "${amd_diagnostics_checkout_root}" ]]; then - ( - cd "${amd_diagnostics_checkout_root}" || exit 0 - run_amd_diagnostic "${diagnostics_path}" "${probe_deadline}" df -h \ - --output=fstype,size,used,avail,pcent,itotal,iused,iavail,ipcent \ - -- . - ) + append_failure_diagnostic_section "${diagnostics_path}" \ + "AMD GPU diagnostics (tool-visible devices)" + if [[ "${amd_diagnostics_expected_gpu_count}" == "0" ]]; then + append_failure_diagnostic_note "${diagnostics_path}" \ + "CPU-only job; AMD GPU probes skipped." + else + if [[ "${runtime}" == "native-kubernetes" ]]; then + append_failure_diagnostic_note "${diagnostics_path}" \ + "SMI visibility may be broader than the Kubernetes device allocation." else - echo "[checkout directory unavailable]" | tee -a "${diagnostics_path}" + append_failure_diagnostic_note "${diagnostics_path}" \ + "Outer-runner SMI visibility may be broader than the test container." fi - echo "tmp_filesystem:" | tee -a "${diagnostics_path}" - run_amd_diagnostic "${diagnostics_path}" "${probe_deadline}" df -h \ - --output=fstype,size,used,avail,pcent,itotal,iused,iavail,ipcent \ - -- /tmp + collect_amd_device_nodes "${diagnostics_path}" "${probe_deadline}" fi - if command -v free >/dev/null 2>&1; then - run_amd_diagnostic "${diagnostics_path}" "${probe_deadline}" free -h - fi - if [[ -r /sys/fs/cgroup/memory.events ]]; then - run_amd_diagnostic "${diagnostics_path}" "${probe_deadline}" cat \ - /sys/fs/cgroup/memory.events + if [[ "${amd_diagnostics_expected_gpu_count}" != "0" ]] \ + && command -v amd-smi >/dev/null 2>&1; then + run_failure_diagnostic "${diagnostics_path}" "${probe_deadline}" \ + amd-smi version + # Bus data identifies a card within the public node without publishing its + # persistent UUID, serial number, or process list. + run_failure_diagnostic "${diagnostics_path}" "${probe_deadline}" \ + amd-smi static -b -g all + run_failure_diagnostic "${diagnostics_path}" "${probe_deadline}" \ + amd-smi metric -e -k -P -x -g all + run_failure_diagnostic "${diagnostics_path}" "${probe_deadline}" \ + amd-smi bad-pages -p -r -u -g all + run_failure_diagnostic "${diagnostics_path}" "${probe_deadline}" \ + amd-smi metric -p -t -u -m -v -g all + run_failure_diagnostic "${diagnostics_path}" "${probe_deadline}" \ + amd-smi xgmi -l -g all + elif [[ "${amd_diagnostics_expected_gpu_count}" != "0" ]] \ + && command -v rocm-smi >/dev/null 2>&1; then + run_failure_diagnostic "${diagnostics_path}" "${probe_deadline}" rocm-smi \ + --showbus --showreplaycount --showrasinfo --showpagesinfo + elif [[ "${amd_diagnostics_expected_gpu_count}" != "0" ]]; then + append_failure_diagnostic_note "${diagnostics_path}" \ + "Neither amd-smi nor rocm-smi is available." fi - echo "Unauthenticated outer-runner network reachability diagnostics" \ - | tee -a "${diagnostics_path}" + if [[ "${runtime}" == "native-kubernetes" ]]; then + append_failure_diagnostic_section "${diagnostics_path}" \ + "Unauthenticated pod network reachability" + else + append_failure_diagnostic_section "${diagnostics_path}" \ + "Unauthenticated outer-runner network reachability" + fi if command -v curl >/dev/null 2>&1; then - run_amd_diagnostic "${diagnostics_path}" "${probe_deadline}" curl -q \ + run_failure_diagnostic "${diagnostics_path}" "${probe_deadline}" curl -q \ --silent --location --proto '=https' --proto-redir '=https' \ --max-redirs 3 --output /dev/null --connect-timeout 2 --max-time 4 \ --write-out 'target=huggingface http_code=%{http_code} dns_done_s=%{time_namelookup} connect_done_s=%{time_connect} tls_done_s=%{time_appconnect} first_byte_s=%{time_starttransfer} total_s=%{time_total}\n' \ https://huggingface.co/api/models/gpt2 - run_amd_diagnostic "${diagnostics_path}" "${probe_deadline}" curl -q \ + run_failure_diagnostic "${diagnostics_path}" "${probe_deadline}" curl -q \ --silent --location --proto '=https' --proto-redir '=https' \ --max-redirs 3 --output /dev/null --connect-timeout 2 --max-time 4 \ --write-out 'target=github_git_smart_http http_code=%{http_code} dns_done_s=%{time_namelookup} connect_done_s=%{time_connect} tls_done_s=%{time_appconnect} first_byte_s=%{time_starttransfer} total_s=%{time_total}\n' \ 'https://github.com/vllm-project/ci-infra.git/info/refs?service=git-upload-pack' else - echo "curl unavailable; reachability probes skipped." \ - | tee -a "${diagnostics_path}" - fi - - if command -v buildkite-agent >/dev/null 2>&1 \ - && [[ -n "${BUILDKITE_JOB_ID:-}" ]]; then - if ( - cd "${upload_root}" \ - && BUILDKITE_AGENT_DEBUG=false \ - BUILDKITE_AGENT_DEBUG_HTTP=false \ - BUILDKITE_AGENT_TRACE_HTTP=false \ - BUILDKITE_AGENT_LOG_LEVEL=error \ - timeout --kill-after=2s "${amd_diagnostics_upload_timeout_seconds}s" \ - buildkite-agent artifact upload "${upload_path}" \ - >/dev/null 2>&1 - ); then - echo "Uploaded AMD GPU diagnostics artifact: ${upload_path}" - else - echo "WARNING: failed to upload AMD GPU diagnostics artifact: ${upload_path}" + append_failure_diagnostic_note "${diagnostics_path}" \ + "curl unavailable; reachability probes skipped." + fi + + if [[ "${runtime}" == "native-kubernetes" \ + && -n "${amd_diagnostics_memory_events_path}" \ + && -r "${amd_diagnostics_memory_events_path}" ]]; then + oom_kill_count=$( + awk '$1 == "oom_kill" { print $2; exit }' \ + "${amd_diagnostics_memory_events_path}" 2>/dev/null + ) + if [[ ! "${oom_kill_count}" =~ ^[0-9]+$ ]]; then + oom_kill_count="" fi - else - echo "Buildkite agent unavailable; diagnostics artifact was not uploaded." fi + case "${identity_label}" in + pod) + summary_identity_label="Kubernetes pod" + ;; + container_hostname) + summary_identity_label="Container hostname" + ;; + esac + + summary_rows=( + "Field" "Value" + "----------------------" "-----" + "Exit code" "${exit_code}" + "Signal" "${exit_signal:-none}" + "Runtime" "${runtime}" + "${summary_identity_label}" "${identity_value}" + ) + if [[ -n "${amd_diagnostics_k8s_node_name}" ]]; then + summary_rows+=("Kubernetes node" "${amd_diagnostics_k8s_node_name}") + fi + if [[ -n "${oom_kill_count}" ]]; then + summary_rows+=("Cgroup OOM kills" "${oom_kill_count}") + fi + summary_rows+=("Diagnostics artifact" "${diagnostics_relative_path}") + + printf '\nAMD CI failure summary\n' + printf '%-22s | %s\n' "${summary_rows[@]}" + return 0 } @@ -788,9 +1128,6 @@ handle_pytest_exit() { echo "Pytest exit code 5 (no tests collected) - treating as success." exit 0 fi - if [[ "${exit_code}" -ne 0 ]]; then - collect_rocm_failure_diagnostics "${exit_code}" - fi exit "$exit_code" } @@ -959,7 +1296,7 @@ handle_amd_runner_exit() { } # Catch both test failures and wrapper/setup failures. Runtime-specific cleanup -# traps below replace this trap but call the same handler before cleaning up. +# traps below replace this trap and call the same handler after cleanup. trap handle_amd_runner_exit EXIT ############################################################################### @@ -973,10 +1310,10 @@ if is_native_runtime; then # shellcheck disable=SC2317 # Called indirectly by the EXIT trap. cleanup_native_workspace() { local exit_code=$? - handle_amd_runner_exit "${exit_code}" if [[ -n "${artifact_work_dir}" ]]; then rm -rf "${artifact_work_dir}" fi + handle_amd_runner_exit "${exit_code}" } trap cleanup_native_workspace EXIT @@ -1020,6 +1357,7 @@ if is_native_runtime; then echo "Native test commands: $commands" run_native_preflight || exit 1 + echo "--- Test log" # Keep AMD CI orchestration variables out of vLLM's runtime environment. clear_ci_orchestration_env /bin/bash -o pipefail -c "${commands}" @@ -1042,7 +1380,6 @@ container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | hea # shellcheck disable=SC2317 # Called indirectly by the EXIT trap. remove_docker_container() { local exit_code=$? - handle_amd_runner_exit "${exit_code}" if docker container inspect "${container_name}" >/dev/null 2>&1; then docker rm -f "${container_name}" || true fi @@ -1055,6 +1392,7 @@ remove_docker_container() { if [[ -n "${artifact_work_dir}" ]]; then rm -rf "${artifact_work_dir}" fi + handle_amd_runner_exit "${exit_code}" } trap remove_docker_container EXIT From 5ae2d38821ba381ff2091b52b463296784b56217 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Tue, 18 Aug 2026 09:30:41 +1000 Subject: [PATCH 064/839] [Perf][Structured Output] Skip unused request-local reasoners (#52573) Signed-off-by: Bugen Zhao Co-authored-by: OpenAI Codex Co-authored-by: Nick Hill --- .../test_reasoning_structured_output.py | 65 +++++++++++++++++-- vllm/v1/structured_output/__init__.py | 40 ++++++------ 2 files changed, 79 insertions(+), 26 deletions(-) diff --git a/tests/v1/structured_output/test_reasoning_structured_output.py b/tests/v1/structured_output/test_reasoning_structured_output.py index ad5f1d5d7951..00a6cd87bcdf 100644 --- a/tests/v1/structured_output/test_reasoning_structured_output.py +++ b/tests/v1/structured_output/test_reasoning_structured_output.py @@ -84,17 +84,62 @@ def manager_with_reasoner(self, mock_vllm_config): return manager def test_should_fill_bitmask_with_enable_in_reasoning( - self, mock_vllm_config, mock_request_with_structured_output + self, manager_with_reasoner, mock_request_with_structured_output ): """Test should_fill_bitmask when enable_in_reasoning is True.""" - # Enable enable_in_reasoning - mock_vllm_config.structured_outputs_config.enable_in_reasoning = True - - manager = StructuredOutputManager(mock_vllm_config) + manager_with_reasoner.enable_in_reasoning = True # Should always return True when enable_in_reasoning is enabled - result = manager.should_fill_bitmask(mock_request_with_structured_output) + result = manager_with_reasoner.should_fill_bitmask( + mock_request_with_structured_output + ) + assert result is True + assert ( + mock_request_with_structured_output.structured_output_request.reasoner + is None + ) + + def test_should_fill_bitmask_reasoning_already_ended( + self, + manager_with_reasoner, + mock_request_with_structured_output, + ): + """An active grammar does not need a request-local reasoner.""" + structured_req = mock_request_with_structured_output.structured_output_request + structured_req.reasoning_ended = True + + result = manager_with_reasoner.should_fill_bitmask( + mock_request_with_structured_output + ) + assert result is True + assert structured_req.reasoner is None + + def test_grammar_bitmask_skips_reasoner_when_already_active( + self, + manager_with_reasoner, + mock_request_with_structured_output, + ): + """Bitmask generation skips reasoning-boundary detection.""" + manager_with_reasoner.vllm_config.num_speculative_tokens = 1 + manager_with_reasoner._grammar_bitmask = Mock() + manager_with_reasoner._grammar_bitmask.shape = (1, 1) + expected_bitmask = Mock() + manager_with_reasoner._grammar_bitmask.numpy.return_value = expected_bitmask + manager_with_reasoner._fill_bitmasks = Mock() + + structured_req = mock_request_with_structured_output.structured_output_request + structured_req.reasoning_ended = True + request_id = mock_request_with_structured_output.request_id + + bitmask = manager_with_reasoner.grammar_bitmask( + requests={request_id: mock_request_with_structured_output}, + structured_output_request_ids=[request_id], + scheduled_spec_decode_tokens={}, + ) + + assert bitmask is expected_bitmask + assert structured_req.reasoner is None def test_should_fill_bitmask_without_enable_in_reasoning( self, @@ -171,6 +216,10 @@ def test_should_advance_with_enable_in_reasoning( mock_request_with_structured_output ) assert result is True + assert ( + mock_request_with_structured_output.structured_output_request.reasoner + is None + ) def test_should_advance_reasoning_not_ended( self, @@ -233,6 +282,10 @@ def test_should_advance_reasoning_already_ended( # Should return True since reasoning has ended assert result is True + assert ( + mock_request_with_structured_output.structured_output_request.reasoner + is None + ) def test_should_advance_uses_new_token_ids_when_provided( self, diff --git a/vllm/v1/structured_output/__init__.py b/vllm/v1/structured_output/__init__.py index f2ba4ac368e4..6a2409c8a5a9 100644 --- a/vllm/v1/structured_output/__init__.py +++ b/vllm/v1/structured_output/__init__.py @@ -289,12 +289,8 @@ def grammar_bitmask( assert isinstance(grammar, StructuredOutputGrammar) apply_bitmask = self.should_fill_bitmask(request) - reasoner = self._get_reasoner(request) - detect_reasoning_end = ( - not apply_bitmask - and reasoner is not None - and not self.enable_in_reasoning - ) + reasoner = None if apply_bitmask else self._get_reasoner(request) + detect_reasoning_end = reasoner is not None simulated_buf: list[int] | None = None history_len = 0 @@ -369,20 +365,24 @@ def should_fill_bitmask(self, request: "Request") -> bool: # NOTE (Hanchen) if enable_in_reasoning is True, it means that # the model needs to be constrained in reasoning. So we should always # enable the bitmask filling. + structured_req = request.structured_output_request + if self.enable_in_reasoning or ( + structured_req is not None and structured_req.reasoning_ended is True + ): + return True + reasoner = self._get_reasoner(request) if reasoner is not None: - if self.enable_in_reasoning: - return True - assert request.structured_output_request is not None - if request.structured_output_request.reasoning_ended is None: + assert structured_req is not None + if structured_req.reasoning_ended is None: # This should be removed here, but since `openai_gptoss` # is an independent code path, it is kept for now. # After unifying the `openai_gptoss` and non-`openai_gptoss` styles, # it can be removed. - request.structured_output_request.reasoning_ended = ( - reasoner.is_reasoning_end(request.prompt_token_ids or []) + structured_req.reasoning_ended = reasoner.is_reasoning_end( + request.prompt_token_ids or [] ) - return request.structured_output_request.reasoning_ended + return structured_req.reasoning_ended return True def should_advance( @@ -398,19 +398,19 @@ def should_advance( if TYPE_CHECKING: assert request.structured_output_request is not None assert request.structured_output_request.grammar is not None + structured_req = request.structured_output_request + if self.enable_in_reasoning or ( + structured_req is not None and structured_req.reasoning_ended is True + ): + return True + # by default, we should always advance # for cases that don't use thinking mode. reasoner = self._get_reasoner(request) if reasoner is None: return True - # if the model needs structured in reasoning, we should advance - if self.enable_in_reasoning: - return True - - structured_req = request.structured_output_request - if structured_req.reasoning_ended: - return True + assert structured_req is not None # Check if reasoning ends in *this* step. # When the caller passes new_token_ids (the tokens that were just From 49fb2ee3481246ce88224220442c6ef30f88b209 Mon Sep 17 00:00:00 2001 From: Divakar Verma <137818590+divakar-amd@users.noreply.github.com> Date: Mon, 17 Aug 2026 19:33:22 -0400 Subject: [PATCH 065/839] [ROCm][CI] add Aiter ops tests (#52208) Signed-off-by: Divakar Verma --- .buildkite/test-amd.yaml | 18 + tests/kernels/core/test_rocm_aiter_ops.py | 625 ++++++++++++++++++++++ vllm/_aiter_ops.py | 29 + 3 files changed, 672 insertions(+) create mode 100644 tests/kernels/core/test_rocm_aiter_ops.py diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index b477b298bb10..e9726f73ff5c 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -3404,6 +3404,24 @@ steps: #---------------------------------------------------------- mi355 · kernels ----------------------------------------------------------# +- label: Kernels Core Operation Test %N # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_1 + parallelism: 3 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - csrc/ + - tests/kernels/core + - tests/kernels/test_top_k_per_row.py + - tests/kernels/test_concat_mla_q.py + - vllm/model_executor/layers/rotary_embedding/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - label: Kernels (B200-MI355) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] diff --git a/tests/kernels/core/test_rocm_aiter_ops.py b/tests/kernels/core/test_rocm_aiter_ops.py new file mode 100644 index 000000000000..cf5e96f49519 --- /dev/null +++ b/tests/kernels/core/test_rocm_aiter_ops.py @@ -0,0 +1,625 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""ROCm AITER helper-op tests for CI stability. + +These tests directly validate AITER kernel correctness against PyTorch/reference +implementations. They provide clear failure signals when AITER library updates +introduce regressions, independent of vLLM's fusion/compilation passes. + +Tested ops: +- ``rms_norm`` and ``rms_norm2d_with_add`` vs PyTorch reference +- ``triton_rotary_embedding`` vs manual NeoX RoPE reference (xfail - known issue) +- ``act_mul_and_fp8_group_quant`` (SiGLU + FP8 group quant) +- fused RMSNorm + quantization ops vs sequential composition + +Related coverage: +- per-token/per-tensor quant roundtrips: ``tests/rocm/aiter/test_quant_op_schema.py`` +- RMSNorm determinism: ``tests/rocm/aiter/test_quant_op_schema.py`` +- group_fp8_quant: ``tests/kernels/quantization/test_rocm_aiter_grouped_quant.py`` +""" + +import warnings + +import pytest +import torch + +import vllm._aiter_ops # noqa: F401 - ensure ops are registered +from tests.kernels.utils import bf16_ulp_distance +from vllm.platforms import current_platform +from vllm.utils.torch_utils import set_random_seed + +pytestmark = pytest.mark.skipif( + not current_platform.is_rocm(), reason="ROCm-specific tests" +) + + +@pytest.fixture(autouse=True) +def setup_cuda_device(): + """Set default device to CUDA and seed RNG for all tests.""" + torch.set_default_device("cuda") + set_random_seed(0) + + +def require_aiter(): + from vllm._aiter_ops import is_aiter_found_and_supported + + if not is_aiter_found_and_supported(): + pytest.skip("aiter is not found or not supported on this hardware") + + +def require_fp8(): + if not current_platform.supports_fp8(): + pytest.skip("FP8 is not supported on this hardware") + + +def _format_observed_rate(count: int, total: int) -> str: + return f"{count / total:.4%} ({count}/{total})" + + +def _format_allowed_rate(rate: float, total: int) -> str: + allowed_count = int(rate * total) + return f"{rate:.4%} (<= {allowed_count}/{total})" + + +def _quantile(values: torch.Tensor, q: float) -> float: + if values.numel() == 0: + return 0.0 + return torch.quantile(values, q).item() + + +def _dequantize_grouped( + quantized: torch.Tensor, scales: torch.Tensor, group_size: int +) -> torch.Tensor: + """Expand per-group scales and dequantize to float32.""" + M, N = quantized.shape + scales_expanded = scales.repeat_interleave(group_size, dim=1)[:, :N] + return quantized.float() * scales_expanded + + +def _rms_norm_reference( + x: torch.Tensor, weight: torch.Tensor, eps: float +) -> torch.Tensor: + """RMSNorm reference implementation in float32.""" + rms = x.float().pow(2).mean(-1, keepdim=True).add(eps).sqrt() + return x.float() / rms * weight.float() + + +def _assert_quant_shapes( + x_quant: torch.Tensor, + scale: torch.Tensor, + M: int, + N: int, + fp8_dtype: torch.dtype, + num_groups: int | None = None, +) -> None: + """Verify quantized output shapes and dtypes.""" + assert x_quant.shape == (M, N) + assert x_quant.dtype == fp8_dtype + if num_groups is not None: + assert scale.shape == (M, num_groups) + else: + assert scale.shape == (M, 1) + + +def _assert_close_budget( + actual: torch.Tensor, + expected: torch.Tensor, + *, + label: str, + atol: float, + rtol: float = 0.0, + pass_rate: float = 0.99999, + max_violation_factor: float = 3.0, +) -> None: + actual_f = actual.detach().float().flatten() + expected_f = expected.detach().float().flatten() + abs_diff = (actual_f - expected_f).abs() + allowed = atol + rtol * expected_f.abs() + + total = abs_diff.numel() + passed = int((abs_diff <= allowed).sum().item()) + failed = total - passed + allowed_fail_rate = 1.0 - pass_rate + max_abs = abs_diff.max().item() + mean_abs = abs_diff.mean().item() + p99_abs = _quantile(abs_diff, 0.99) + p999_abs = _quantile(abs_diff, 0.999) + worst_ratio = (abs_diff / allowed.clamp_min(1e-12)).max().item() + max_atol = max_violation_factor * atol + above_max_count = int((abs_diff > max_atol).sum().item()) + + msg = ( + "[rocm_aiter_ops] " + f"{label}: " + f"pass={passed / total:.4%} ({passed}/{total}) " + f"fail={_format_observed_rate(failed, total)} " + f"allowed_fail={_format_allowed_rate(allowed_fail_rate, total)} " + f"atol={atol:g} " + f"rtol={rtol:g} " + f"abs>{max_atol:g}={_format_observed_rate(above_max_count, total)} " + f"allowed_above_max={_format_allowed_rate(0.0, total)} " + f"max_abs={max_abs:.6g} " + f"mean_abs={mean_abs:.6g} " + f"p99_abs={p99_abs:.6g} " + f"p999_abs={p999_abs:.6g} " + f"worst_ratio={worst_ratio:.6g}" + ) + print(msg) + if failed > 0: + warnings.warn(msg, stacklevel=2) + + assert passed / total >= pass_rate, msg + assert max_abs <= max_atol, msg + assert mean_abs <= atol * 0.25, msg + + +def _assert_rel_error_quality( + actual: torch.Tensor, + expected: torch.Tensor, + *, + label: str, + mean_limit: float, + preferred_rel: float, + max_rel: float, + max_fail_rate: float, +) -> None: + rel = ( + (actual.float() - expected.float()).abs() + / expected.float().abs().clamp_min(1e-5) + ).flatten() + total = rel.numel() + within_preferred_count = int((rel <= preferred_rel).sum().item()) + fail_count = total - within_preferred_count + above_max_count = int((rel > max_rel).sum().item()) + mean_rel = rel.mean().item() + max_rel_err = rel.max().item() + p99 = _quantile(rel, 0.99) + p999 = _quantile(rel, 0.999) + + msg = ( + "[rocm_aiter_ops] " + f"{label}: " + f"rel<={preferred_rel:g} pass={within_preferred_count / total:.4%} " + f"({within_preferred_count}/{total}) " + f"fail={_format_observed_rate(fail_count, total)} " + f"mean_limit={mean_limit:.4%} " + f"rel>{max_rel:g}={_format_observed_rate(above_max_count, total)} " + f"allowed_above_max={_format_allowed_rate(max_fail_rate, total)} " + f"mean_rel={mean_rel:.6g} " + f"max_rel={max_rel_err:.6g} " + f"p99={p99:.6g} " + f"p999={p999:.6g}" + ) + print(msg) + if fail_count > 0: + warnings.warn(msg, stacklevel=2) + + assert mean_rel < mean_limit, msg + assert above_max_count / total <= max_fail_rate, msg + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_rocm_aiter_rms_norm_vs_torch(dtype): + """rocm_aiter_rms_norm matches PyTorch manual RMSNorm for float16 and bfloat16.""" + require_aiter() + from vllm._aiter_ops import rocm_aiter_ops + + M, N = 8, 128 + x = torch.randn(M, N, dtype=dtype) + weight = torch.ones(N, dtype=dtype) + eps = 1e-5 + + # Reference: float32 RMSNorm for precision + ref = _rms_norm_reference(x, weight, eps).to(dtype) + + out = rocm_aiter_ops.rms_norm(x, weight, eps) + + # Allow max 1 ULP difference due to different accumulation order. + ulp = bf16_ulp_distance(out, ref) + max_ulp = int(ulp.max().item()) + assert max_ulp <= 1, f"rms_norm dtype={dtype}: max ULP distance {max_ulp} > 1" + + +# -- Numerical accuracy tests for AITER custom ops ------------------------- + + +def test_rocm_aiter_rmsnorm_with_add_vs_torch(): + """rocm_aiter_rmsnorm2d_fwd_with_add matches manual residual+RMSNorm reference.""" + require_aiter() + from vllm._aiter_ops import rocm_aiter_ops + + M, N = 16, 256 + x = torch.randn(M, N, dtype=torch.bfloat16) + residual = torch.randn(M, N, dtype=torch.bfloat16) + weight = torch.ones(N, dtype=torch.bfloat16) + eps = 1e-5 + + # Reference: add residual, then RMSNorm + h = x.float() + residual.float() + rms = h.pow(2).mean(-1, keepdim=True).add(eps).sqrt() + ref_normed = (h / rms * weight.float()).to(torch.bfloat16) + ref_residual = h.to(torch.bfloat16) + + out, res_out = rocm_aiter_ops.rms_norm2d_with_add(x, residual, weight, eps) + + # Allow max 1 ULP difference due to different accumulation order. + ulp_normed = bf16_ulp_distance(out, ref_normed) + max_ulp_normed = int(ulp_normed.max().item()) + assert max_ulp_normed <= 1, ( + f"rmsnorm_with_add normed: max ULP distance {max_ulp_normed} > 1" + ) + + ulp_residual = bf16_ulp_distance(res_out, ref_residual) + max_ulp_residual = int(ulp_residual.max().item()) + assert max_ulp_residual <= 1, ( + f"rmsnorm_with_add residual: max ULP distance {max_ulp_residual} > 1" + ) + + +@pytest.mark.xfail( + strict=True, + raises=AssertionError, + reason="AITER Triton RoPE precision issue: https://github.com/ROCm/aiter/issues/4765", +) +def test_rocm_aiter_triton_rotary_embedding_vs_torch(): + """rocm_aiter_triton_rotary_embedding matches manual NeoX-style RoPE reference. + + The AITER kernel calls rope_cached_thd_positions_offsets_2c_fwd_inplace with + reuse_freqs_front_part=True. This means it reads only the first head_size//2 + entries of cos and sin from the cache and applies them to both halves of the + head (pairwise NeoX rotation). The cos_sin_cache must be built with the second + half mirroring the first so the reference and kernel agree. + """ + require_aiter() + + num_tokens = 8 + num_heads = 4 + head_size = 64 + half_dim = head_size // 2 + max_pos = 32 + + positions = torch.randint(0, max_pos, (num_tokens,), dtype=torch.long) + query = torch.randn(num_tokens, num_heads * head_size, dtype=torch.bfloat16) + key = torch.randn(num_tokens, num_heads * head_size, dtype=torch.bfloat16) + + # Build cos/sin cache with second half mirroring first half. + # The kernel uses reuse_freqs_front_part=True: it reads only cos/sin[:, :half_dim] + # and applies those same frequencies to both the first and second + # halves of the head. + cos_half = torch.randn(max_pos, half_dim, dtype=torch.bfloat16) + sin_half = torch.randn(max_pos, half_dim, dtype=torch.bfloat16) + cos_cache = torch.cat([cos_half, cos_half], dim=-1) # [max_pos, head_size] + sin_cache = torch.cat([sin_half, sin_half], dim=-1) # [max_pos, head_size] + cos_sin_cache = torch.cat([cos_cache, sin_cache], dim=-1) # [max_pos, 2*head_size] + + # Reference: NeoX-style pairwise rotation with front-half frequencies only. + # rotate_style=0 (NeoX): [x1*c - x2*s, x2*c + x1*s] + cos_pos = cos_half[positions] # [num_tokens, half_dim] + sin_pos = sin_half[positions] # [num_tokens, half_dim] + + def apply_rope_ref(t: torch.Tensor) -> torch.Tensor: + t_r = t.float().view(num_tokens, num_heads, head_size) + c = cos_pos.float().unsqueeze(1) # [num_tokens, 1, half_dim] + s = sin_pos.float().unsqueeze(1) + x1, x2 = t_r[..., :half_dim], t_r[..., half_dim:] + rotated = torch.cat([x1 * c - x2 * s, x2 * c + x1 * s], dim=-1) + return rotated.to(t.dtype).view(num_tokens, num_heads * head_size) + + ref_q = apply_rope_ref(query) + ref_k = apply_rope_ref(key) + + # AITER in-place RoPE (modifies query/key in-place) + q_aiter = query.clone() + k_aiter = key.clone() + torch.ops.vllm.rocm_aiter_triton_rotary_embedding( + positions, + q_aiter, + k_aiter, + head_size, + cos_sin_cache, + True, # is_neox style -> rotate_style=0 + ) + + _assert_close_budget( + q_aiter.float(), + ref_q.float(), + label="triton_rope query", + atol=1e-3, + rtol=1.6e-2, + ) + _assert_close_budget( + k_aiter.float(), + ref_k.float(), + label="triton_rope key", + atol=1e-3, + rtol=1.6e-2, + ) + + +def test_rocm_aiter_act_mul_fp8_group_quant_roundtrip(): + """act_mul_and_fp8_group_quant: dequantized output matches SiLU gate reference.""" + require_aiter() + require_fp8() + + M, N = 32, 512 # N even: N//2 gate, N//2 up + group_size = 128 + x = torch.randn(M, N, dtype=torch.bfloat16) + + x_quant, scale = torch.ops.vllm.rocm_aiter_act_mul_and_fp8_group_quant( + x, group_size + ) + + N_half = N // 2 + # Reference: SiLU(gate) * up + gate = x.float()[:, :N_half] + up = x.float()[:, N_half:] + ref = torch.sigmoid(gate) * gate * up # SiGLU + + x_dequant = _dequantize_grouped(x_quant, scale, group_size) + + _assert_rel_error_quality( + x_dequant, + ref, + label="act_mul_fp8_group_quant", + mean_limit=0.1, + preferred_rel=0.1, + max_rel=1.0, + max_fail_rate=0.01, + ) + + +# -- Fused RMSNorm + quantization accuracy tests --------------------------- + + +def test_rocm_aiter_rmsnorm_fused_dynamic_quant_vs_sequential(): + """Fused RMSNorm+per-token-FP8-quant matches sequential rms_norm->per_token_quant. + + Tests that the fused kernel produces the same result as the two-step + sequential composition: rms_norm(x) followed by per_token_quant. + The fused path is used in production for inference throughput. + """ + require_aiter() + require_fp8() + from vllm._aiter_ops import rocm_aiter_ops + + M, N = 32, 512 + x = torch.randn(M, N, dtype=torch.bfloat16) + weight = torch.ones(N, dtype=torch.bfloat16) + eps = 1e-5 + fp8_dtype = current_platform.fp8_dtype() + + # Sequential reference + normed = rocm_aiter_ops.rms_norm(x, weight, eps) + ref_q, ref_scale = rocm_aiter_ops.per_token_quant(normed, fp8_dtype) + ref_dequant = ref_q.float() * ref_scale.float() + + # Fused op + fused_q, fused_scale = torch.ops.vllm.rocm_aiter_rmsnorm_fused_dynamic_quant( + x, weight, eps, fp8_dtype + ) + _assert_quant_shapes(fused_q, fused_scale, M, N, fp8_dtype) + fused_dequant = fused_q.float() * fused_scale.float() + + # Fused vs sequential: both should recover RMSNorm output within FP8 error + _assert_rel_error_quality( + fused_dequant, + ref_dequant, + label="rmsnorm_fused_dynamic_quant", + mean_limit=0.05, + preferred_rel=0.05, + max_rel=0.5, + max_fail_rate=0.01, + ) + + +def test_rocm_aiter_rmsnorm_fused_add_dynamic_quant_vs_reference(): + """Fused (residual-add + RMSNorm + per-token-FP8-quant) + matches sequential reference. + + Production path: input + residual -> RMSNorm -> FP8 quantize, returning + both the quantized norm output and the residual sum for the next layer. + """ + require_aiter() + require_fp8() + from vllm._aiter_ops import rocm_aiter_ops + + M, N = 16, 256 + x = torch.randn(M, N, dtype=torch.bfloat16) + residual = torch.randn(M, N, dtype=torch.bfloat16) + weight = torch.ones(N, dtype=torch.bfloat16) + eps = 1e-5 + fp8_dtype = current_platform.fp8_dtype() + + # Sequential reference: add residual -> rms_norm -> per_token_quant + h = (x.float() + residual.float()).to(torch.bfloat16) + normed = rocm_aiter_ops.rms_norm(h, weight, eps) + ref_q, ref_scale = rocm_aiter_ops.per_token_quant(normed, fp8_dtype) + ref_residual_out = h + + # Fused op: returns (x_quant, residual_out, scale) + fused_q, fused_res_out, fused_scale = ( + torch.ops.vllm.rocm_aiter_rmsnorm_fused_add_dynamic_quant( + x, residual, weight, eps, fp8_dtype + ) + ) + _assert_quant_shapes(fused_q, fused_scale, M, N, fp8_dtype) + assert fused_res_out.shape == (M, N) + + torch.testing.assert_close(fused_res_out, ref_residual_out, atol=1e-2, rtol=1.6e-2) + # Dequantized output matches sequential path + fused_dequant = fused_q.float() * fused_scale.float() + ref_dequant = ref_q.float() * ref_scale.float() + _assert_rel_error_quality( + fused_dequant, + ref_dequant, + label="rmsnorm_fused_add_dynamic_quant", + mean_limit=0.05, + preferred_rel=0.05, + max_rel=0.5, + max_fail_rate=0.01, + ) + + +def test_rocm_aiter_rmsnorm_fp8_group_quant_vs_sequential(): + """Fused RMSNorm+FP8-group-quant matches sequential rms_norm->group_fp8_quant. + + Tests both output shapes and dequantized accuracy against the two-step + sequential composition. + """ + require_aiter() + require_fp8() + from vllm._aiter_ops import rocm_aiter_ops + + M, N = 32, 512 + group_size = 128 + x = torch.randn(M, N, dtype=torch.bfloat16) + weight = torch.ones(N, dtype=torch.bfloat16) + eps = 1e-5 + fp8_dtype = current_platform.fp8_dtype() + expected_groups = (N + group_size - 1) // group_size + + # Fused op: (x_quant, scales) + fused_q, fused_scales = torch.ops.vllm.rocm_aiter_rmsnorm_fp8_group_quant( + x, weight, eps, group_size + ) + _assert_quant_shapes(fused_q, fused_scales, M, N, fp8_dtype, expected_groups) + + # Dequantize and compare to reference: rms_norm -> group quant -> dequant + normed = rocm_aiter_ops.rms_norm(x, weight, eps) + ref_q, ref_scales = rocm_aiter_ops.group_fp8_quant(normed, group_size) + ref_dequant = _dequantize_grouped(ref_q, ref_scales, group_size) + fused_dequant = _dequantize_grouped(fused_q, fused_scales, group_size) + + _assert_rel_error_quality( + fused_dequant, + ref_dequant, + label="rmsnorm_fp8_group_quant", + mean_limit=0.05, + preferred_rel=0.05, + max_rel=0.5, + max_fail_rate=0.01, + ) + + +def test_rocm_aiter_rmsnorm_with_add_fp8_group_quant_residual_accuracy(): + """Fused rmsnorm_with_add_fp8_group_quant residual output matches x + residual.""" + require_aiter() + require_fp8() + from vllm._aiter_ops import rocm_aiter_ops + + M, N = 16, 256 + group_size = 128 + x = torch.randn(M, N, dtype=torch.bfloat16) + residual = torch.randn(M, N, dtype=torch.bfloat16) + weight = torch.ones(N, dtype=torch.bfloat16) + eps = 1e-5 + + fused_q, fused_res, fused_scales = ( + torch.ops.vllm.rocm_aiter_rmsnorm_with_add_fp8_group_quant( + x, residual, weight, eps, group_size + ) + ) + + # Residual output must equal x + residual + ref_residual = (x.float() + residual.float()).to(torch.bfloat16) + _assert_close_budget( + fused_res.float(), + ref_residual.float(), + label="rmsnorm_with_add_fp8_group_quant residual", + atol=1e-2, + rtol=1e-2, + ) + + # Dequantized quant output must match rms_norm(x + residual) + h = ref_residual + ref_normed = _rms_norm_reference(h, weight, eps).to(torch.bfloat16) + ref_q, ref_scales = rocm_aiter_ops.group_fp8_quant(ref_normed, group_size) + ref_dequant = _dequantize_grouped(ref_q, ref_scales, group_size) + fused_dequant = _dequantize_grouped(fused_q, fused_scales, group_size) + _assert_rel_error_quality( + fused_dequant, + ref_dequant, + label="rmsnorm_with_add_fp8_group_quant", + mean_limit=0.05, + preferred_rel=0.05, + max_rel=0.5, + max_fail_rate=0.01, + ) + + +# -- End-to-end inference chain test --------------------------------------- + + +def test_rocm_aiter_rms_norm_then_per_token_quant_e2e(): + """End-to-end: BF16 RMSNorm -> per-token FP8 quantization -> dequantize. + + Simulates the inference path through a transformer layer norm before a + linear projection: verifies the full chain produces accurate output + compared to a float32 reference. + """ + require_aiter() + require_fp8() + from vllm._aiter_ops import rocm_aiter_ops + + # Llama-style hidden dim + M, N = 32, 4096 + x = torch.randn(M, N, dtype=torch.bfloat16) + weight = torch.randn(N, dtype=torch.bfloat16) # learned scale + eps = 1e-5 + fp8_dtype = current_platform.fp8_dtype() + + # Float32 reference for the full chain + ref_normed_f32 = _rms_norm_reference(x, weight, eps) + + # AITER chain: RMSNorm -> per-token FP8 quant -> dequant + normed = rocm_aiter_ops.rms_norm(x, weight, eps) + x_q, scale = rocm_aiter_ops.per_token_quant(normed, fp8_dtype) + x_dequant = x_q.float() * scale.float() # scale: [M, 1] + + _assert_quant_shapes(x_q, scale, M, N, fp8_dtype) + # Dequantized result should match the float32 reference within FP8 quant error + _assert_rel_error_quality( + x_dequant, + ref_normed_f32, + label="rms_norm_then_per_token_quant_e2e", + mean_limit=0.05, + preferred_rel=0.05, + max_rel=0.5, + max_fail_rate=0.01, + ) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_rocm_aiter_rms_norm_then_group_fp8_quant_e2e(dtype): + """End-to-end: RMSNorm (fp16/bf16) -> FP8 group quantization -> dequantize. + + Covers both float16 and bfloat16 inputs to verify dtype-agnostic + behavior of the RMSNorm+group-quant pipeline. + """ + require_aiter() + require_fp8() + from vllm._aiter_ops import rocm_aiter_ops + + M, N = 16, 512 + group_size = 128 + x = torch.randn(M, N, dtype=dtype) + weight = torch.randn(N, dtype=dtype) + eps = 1e-5 + + # Float32 reference + ref_normed_f32 = _rms_norm_reference(x, weight, eps) + + # AITER chain: RMSNorm -> group FP8 quant -> dequant + normed = rocm_aiter_ops.rms_norm(x, weight, eps) + x_fp8, scales = rocm_aiter_ops.group_fp8_quant(normed.bfloat16(), group_size) + x_dequant = _dequantize_grouped(x_fp8, scales, group_size) + + _assert_rel_error_quality( + x_dequant, + ref_normed_f32, + label=f"rms_norm_then_group_fp8_quant_e2e dtype={dtype}", + mean_limit=0.1, + preferred_rel=0.1, + max_rel=1.0, + max_fail_rate=0.01, + ) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 2721dec0beaf..b6a504d0cfbe 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -2340,6 +2340,35 @@ def get_fused_mla_dual_rms_norm_op() -> OpOverload: def get_fused_mla_dual_rms_norm_per_token_quant_op() -> OpOverload: return torch.ops.vllm.fused_mla_dual_rms_norm_per_token_quant.default + @staticmethod + def rms_norm( + x: torch.Tensor, + weight: torch.Tensor, + epsilon: float, + ) -> torch.Tensor: + """RMSNorm via AITER kernel.""" + import aiter + + return aiter.rms_norm(x, weight, epsilon) + + @staticmethod + def rms_norm2d_with_add( + x: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + epsilon: float, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Fused residual-add + RMSNorm via AITER kernel. + + Returns (normalized_output, residual_sum). + """ + import aiter + + out = torch.empty_like(x) + residual_out = torch.empty_like(x) + aiter.rmsnorm2d_fwd_with_add(out, x, residual, residual_out, weight, epsilon, 0) + return out, residual_out + @staticmethod def w8a8_gemm( A: torch.Tensor, From c296cf8259ed834b7174c59a2ee1f2f25967a87b Mon Sep 17 00:00:00 2001 From: Jakub Byczkowski Date: Tue, 18 Aug 2026 01:54:05 +0200 Subject: [PATCH 066/839] [Bugfix] Add forward_xpu to XDRotaryEmbedding for HunyuanOCR on XPU (#52174) Signed-off-by: Jakub Byczkowski Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- .../model_executor/layers/rotary_embedding/xdrope.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/vllm/model_executor/layers/rotary_embedding/xdrope.py b/vllm/model_executor/layers/rotary_embedding/xdrope.py index dab7aad9759a..a799b10a190e 100644 --- a/vllm/model_executor/layers/rotary_embedding/xdrope.py +++ b/vllm/model_executor/layers/rotary_embedding/xdrope.py @@ -137,6 +137,18 @@ def forward_cuda( key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape) return query, key + def forward_xpu( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor | None = None, + offsets: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + # No fused XPU kernel for XD-sectioned RoPE; the base + # RotaryEmbedding.forward_xpu calls the generic C++ kernel, which + # rejects the 2-D [4, num_tokens] xdrope positions. Use the native path. + return self.forward_native(positions, query, key, offsets) + @staticmethod def get_next_input_positions( context_len: int, From 58aa1e3d2692250c6bb24a0f852e3cb8e809a14f Mon Sep 17 00:00:00 2001 From: Tommy Asai <147074725+tommy-asai-sonarsource@users.noreply.github.com> Date: Tue, 18 Aug 2026 01:58:02 +0200 Subject: [PATCH 067/839] [Bugfix][SM120][MLA] Disable dense prefill for FlashInfer sparse MLA (#51395) Signed-off-by: tommy-asai-sonarsource Co-authored-by: Roger Wang --- tests/v1/attention/test_flashinfer_sparse_mla_sm120_api.py | 7 +++++++ .../attention/backends/mla/flashinfer_mla_sparse_sm120.py | 1 + 2 files changed, 8 insertions(+) diff --git a/tests/v1/attention/test_flashinfer_sparse_mla_sm120_api.py b/tests/v1/attention/test_flashinfer_sparse_mla_sm120_api.py index 12a8cf5cbc39..724c8874b931 100644 --- a/tests/v1/attention/test_flashinfer_sparse_mla_sm120_api.py +++ b/tests/v1/attention/test_flashinfer_sparse_mla_sm120_api.py @@ -34,6 +34,13 @@ def test_sm120_backend_uses_dedicated_backend_name() -> None: ) +def test_sm120_backend_uses_sparse_mqa_for_prefill() -> None: + impl_cls = FlashInferMLASparseSM120Backend.get_impl_cls() + + assert impl_cls.is_sparse + assert not impl_cls.supports_dense_mha_prefill + + def test_v32_glm_sm120_backend_accepts_glm_block_size( monkeypatch, ) -> None: diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py index 9db696a39c43..4de2c2f53f3b 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py @@ -33,6 +33,7 @@ class FlashInferMLASparseSM120Impl(MLAAttentionImpl[FlashInferMLASparseMetadata] """SM120 FlashInfer sparse-MLA implementation.""" is_sparse = True + supports_dense_mha_prefill = False def __init__( self, From 0e8989b41648e99176a53480e42c447f8bf95620 Mon Sep 17 00:00:00 2001 From: Kunshang Ji Date: Tue, 18 Aug 2026 08:05:48 +0800 Subject: [PATCH 068/839] [ROCm] gaurd on_gfx1250 call with rocm platform (#52625) Signed-off-by: Kunshang Ji --- .../model_executor/layers/quantization/utils/fp8_utils.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/quantization/utils/fp8_utils.py b/vllm/model_executor/layers/quantization/utils/fp8_utils.py index 7910db192564..5df6bcb73421 100644 --- a/vllm/model_executor/layers/quantization/utils/fp8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/fp8_utils.py @@ -881,9 +881,13 @@ def w8a8_triton_block_scaled_mm( torch.Tensor: The result of matmul. """ - from vllm.platforms.rocm import on_gfx1250 + _on_gfx1250 = False + if current_platform.is_rocm(): + from vllm.platforms.rocm import on_gfx1250 - if on_gfx1250(): + _on_gfx1250 = on_gfx1250() + + if _on_gfx1250: # Torch upcast reference: dequantize A,B to fp32 and matmul in fp32. # Avoids the gfx1250 native-fp8 block GEMM NaN bug. Correct but slow. _bn, _bk = block_size[0], block_size[1] From 0db502c8d8a620f29dd3560e69a22f0a7ed3477c Mon Sep 17 00:00:00 2001 From: Summer Yang Date: Mon, 17 Aug 2026 17:05:51 -0700 Subject: [PATCH 069/839] [Kimi-K3] support DCP partial prefix cache hit (#50493) Signed-off-by: Summer Yang Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Yifan Qiao --- .../test_kimi_linear_context_parallel.py | 111 ++++++++++ tests/v1/attention/test_mla_backends.py | 43 ++++ .../test_partial_prefix_cache_hits.py | 204 +++++++++++++++++- .../test_partial_prefix_cache_primitives.py | 44 ++-- tests/v1/worker/test_gpu_model_runner_v2.py | 118 ++++++++++ vllm/v1/core/kv_cache_coordinator.py | 15 +- vllm/v1/worker/gpu/block_table.py | 9 +- vllm/v1/worker/gpu/model_runner.py | 16 +- 8 files changed, 523 insertions(+), 37 deletions(-) create mode 100644 tests/v1/worker/test_gpu_model_runner_v2.py diff --git a/tests/distributed/test_kimi_linear_context_parallel.py b/tests/distributed/test_kimi_linear_context_parallel.py index d37e19d0bd3b..06348074bb14 100644 --- a/tests/distributed/test_kimi_linear_context_parallel.py +++ b/tests/distributed/test_kimi_linear_context_parallel.py @@ -2,12 +2,15 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import math +from pathlib import Path import pytest from tests.conftest import VllmRunner from tests.utils import create_new_process_for_each_test from vllm import SamplingParams, TokensPrompt +from vllm.config import CUDAGraphMode +from vllm.transformers_utils.configs.kimi_k3 import KimiK3Config MODEL = "moonshotai/Kimi-Linear-48B-A3B-Instruct" @@ -126,3 +129,111 @@ def test_kimi_linear_dcp_tiny( logprob_drifts.append(abs(baseline_logprob - dcp_logprob)) assert max(logprob_drifts) <= 1e-2 + + +def _make_tiny_k3_config(model_dir: Path) -> str: + config = KimiK3Config( + text_config=_make_tiny_overrides(), + vision_config={ + "vt_num_attention_heads": 2, + "vt_num_hidden_layers": 1, + "vt_hidden_size": 32, + "vt_intermediate_size": 64, + "qkv_hidden_size": 48, + }, + architectures=["KimiK3ForConditionalGeneration"], + ) + config.save_pretrained(model_dir) + return str(model_dir) + + +def _run_k3_partial_prefix_reuse( + vllm_runner: type[VllmRunner], model_name: str, dcp_size: int +) -> tuple[list[int], list[float], int]: + block_size = 256 + common_prefix = [3 + (token_idx % 251) for token_idx in range(block_size)] + prime_prompt = TokensPrompt(prompt_token_ids=common_prefix + [17]) + replay_prompt = TokensPrompt(prompt_token_ids=common_prefix + [18, 19]) + sampling_params = SamplingParams( + temperature=0, + max_tokens=8, + seed=0, + ignore_eos=True, + logprobs=20, + ) + + with vllm_runner( + model_name=model_name, + skip_tokenizer_init=True, + load_format="dummy", + language_model_only=True, + tensor_parallel_size=2, + decode_context_parallel_size=dcp_size, + cp_kv_cache_interleave_size=1, + distributed_executor_backend="mp", + dtype="bfloat16", + seed=0, + max_model_len=512, + max_num_seqs=4, + max_num_batched_tokens=256, + gpu_memory_utilization=0.85, + block_size=block_size, + enable_chunked_prefill=True, + enable_prefix_caching=True, + mamba_cache_mode="align", + ) as runner: + assert ( + runner.llm.llm_engine.vllm_config.compilation_config.cudagraph_mode + == CUDAGraphMode.FULL_AND_PIECEWISE + ) + runner.llm.generate( + [prime_prompt], + sampling_params, + use_tqdm=False, + ) + replay_output = runner.llm.generate( + [replay_prompt], + sampling_params, + use_tqdm=False, + )[0] + + completion = replay_output.outputs[0] + token_ids = list(completion.token_ids) + assert len(token_ids) == sampling_params.max_tokens + assert completion.logprobs is not None + assert len(completion.logprobs) == sampling_params.max_tokens + selected_logprobs = [ + step_logprobs[token_id].logprob + for token_id, step_logprobs in zip(token_ids, completion.logprobs) + ] + return token_ids, selected_logprobs, replay_output.num_cached_tokens + + +@create_new_process_for_each_test() +@pytest.mark.distributed(num_gpus=2) +def test_kimi_k3_dcp_partial_prefix_reuse( + vllm_runner: type[VllmRunner], + num_gpus_available: int, + tmp_path: Path, +) -> None: + if num_gpus_available < 2: + pytest.skip("Need at least 2 GPUs") + + model_name = _make_tiny_k3_config(tmp_path / "tiny-kimi-k3") + baseline_tokens, baseline_logprobs, baseline_cached = _run_k3_partial_prefix_reuse( + vllm_runner, model_name, dcp_size=1 + ) + dcp_tokens, dcp_logprobs, dcp_cached = _run_k3_partial_prefix_reuse( + vllm_runner, model_name, dcp_size=2 + ) + + assert baseline_cached == 256 + assert dcp_cached == 256 + assert dcp_tokens == baseline_tokens + assert len(dcp_logprobs) == len(baseline_logprobs) == 8 + logprob_drifts = [] + for baseline_logprob, dcp_logprob in zip(baseline_logprobs, dcp_logprobs): + assert math.isfinite(baseline_logprob) + assert math.isfinite(dcp_logprob) + logprob_drifts.append(abs(baseline_logprob - dcp_logprob)) + assert max(logprob_drifts) <= 1e-2 diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index 72773adcb674..cd164250a2fc 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -27,6 +27,7 @@ MLAAttention, QueryLenSupport, _DecodeConcatQuantFP8, + build_mla_chunked_context_metadata, ) from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape from vllm.platforms import current_platform @@ -93,6 +94,48 @@ def test_mla_kv_cache_spec_uses_layer_cache_dtype( assert spec.page_size_bytes == 64 * 656 +@pytest.mark.cpu_test +def test_dcp_chunked_context_accepts_non_virtual_block_aligned_prefix(): + context_lens_cpu = torch.tensor([65, 96], dtype=torch.int32) + prefill_query_start_loc_cpu = torch.tensor([0, 1, 2], dtype=torch.int32) + + metadata = build_mla_chunked_context_metadata( + context_lens_cpu=context_lens_cpu, + prefill_query_start_loc_cpu=prefill_query_start_loc_cpu, + chunked_prefill_workspace=torch.empty(0), + chunked_prefill_workspace_size=128, + block_size=64, + align_chunk_to_block=True, + device=torch.device("cpu"), + dcp_world_size=4, + dcp_local_block_size=1, + dcp_virtual_block_size=4, + ) + + assert metadata is not None + assert metadata.context_lens.tolist() == [65, 96] + assert [chunk.seq_lens.tolist() for chunk in metadata.chunks] == [[65], [96]] + assert [chunk.starts.tolist() for chunk in metadata.chunks] == [[0], [0]] + assert [chunk.local_context_lens_allranks for chunk in metadata.chunks] == [ + [[17, 16, 16, 16]], + [[24, 24, 24, 24]], + ] + assert [chunk.padded_local_seq_lens for chunk in metadata.chunks] == [ + [17], + [24], + ] + assert [chunk.padded_local_cu_seq_lens.tolist() for chunk in metadata.chunks] == [ + [0, 17], + [0, 24], + ] + assert [chunk.cu_seq_lens.tolist() for chunk in metadata.chunks] == [ + [0, 65], + [0, 96], + ] + assert [chunk.num_context_tokens for chunk in metadata.chunks] == [65, 96] + assert [chunk.num_local_context_tokens for chunk in metadata.chunks] == [17, 24] + + # Remove sm100 backends from the list if not using sm100 if not torch.cuda.is_available() or torch.cuda.get_device_properties(0).major < 10: BACKENDS_TO_TEST.remove(AttentionBackendEnum.CUTLASS_MLA) diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py index d237ce9f345b..da5bd5b5a1f3 100644 --- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py @@ -4,6 +4,7 @@ "align") models: scheduler chunk splitting, partial tail registration, CoW on partial hits, and same-step deferral.""" +from math import lcm from types import SimpleNamespace from unittest.mock import MagicMock @@ -68,12 +69,62 @@ def test_capable_connector_uses_divergent_partial_hit_lookup(): manager.get_computed_blocks.assert_not_called() -def test_mamba_align_split_partial_tail_schedule(): +def make_full_mamba_manager( + *, + dcp_world_size: int, + hash_block_size: int = 2, + full_block_size: int = 4, + mamba_block_size: int = 4, + num_blocks: int = 32, + use_eagle: bool = False, +): + kv_cache_config = KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=full_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=mamba_block_size, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ], + ) + scheduler_block_size = lcm( + full_block_size * dcp_world_size, + mamba_block_size, + ) + return make_kv_cache_manager( + kv_cache_config=kv_cache_config, + max_model_len=8192, + enable_caching=True, + dcp_world_size=dcp_world_size, + scheduler_block_size=scheduler_block_size, + hash_block_size=hash_block_size, + use_eagle=use_eagle, + ) + + +@pytest.mark.parametrize("dcp_world_size", [1, 4]) +def test_mamba_align_split_partial_tail_schedule(dcp_world_size: int): """Chunk ends with partial hits on: block-aligned chunks, one extra stop at the prompt's last hash boundary (registering the partial tail), then the remaining tokens. block=512, hash=32, prompt=10000, budget=8192: 0 -> 8192 -> 9728 -> 9984 -> 10000.""" block_size = 512 + scheduler_block_size = block_size * dcp_world_size hash_block_size = 32 mock = SimpleNamespace( cache_config=SimpleNamespace(block_size=block_size), @@ -81,6 +132,8 @@ def test_mamba_align_split_partial_tail_schedule(): scheduler_config=SimpleNamespace(long_prefill_token_threshold=0), use_eagle=False, hash_block_size=hash_block_size, + dcp_world_size=dcp_world_size, + scheduler_block_size=scheduler_block_size, mamba_partial_cache_hit=True, ) split = Scheduler._mamba_block_aligned_split @@ -109,6 +162,7 @@ def test_mamba_align_split_partial_tail_schedule(): # stops at the next block boundary (10240), later chunk ends re-align. req2 = make_request("1", [0] * 12000, hash_block_size, sha256) req2.num_computed_tokens = 9984 + assert req2.num_computed_tokens % scheduler_block_size != 0 assert split(self=mock, request=req2, num_new_tokens=2016) == 256 req2.num_computed_tokens = 10240 assert split(self=mock, request=req2, num_new_tokens=1000) == 512 @@ -1499,3 +1553,151 @@ def test_hybrid_sliding_window_group_disables_partial_hash_hits(): assert num_computed == mamba_block_size assert len(computed_blocks.blocks[0]) * hash_block_size == num_computed + + +@pytest.mark.parametrize("dcp_world_size", [1, 2, 4]) +def test_hybrid_partial_hash_hit_uses_cow_under_dcp(dcp_world_size: int): + hash_block_size = 2 + physical_block_size = 4 + manager = make_full_mamba_manager( + dcp_world_size=dcp_world_size, + hash_block_size=hash_block_size, + full_block_size=physical_block_size, + mamba_block_size=physical_block_size, + ) + assert manager.coordinator.enable_partial_hash_hits + + req0 = make_request("dcp-owner", [0, 0, 1, 1, 2, 2], 2, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) + assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None + manager.free(req0) + manager.new_step_starts() + + partial_hash = req0.block_hashes[2] + partial_full_block = manager.block_pool.get_cached_block(partial_hash, [0]) + partial_mamba_block = manager.block_pool.get_cached_block(partial_hash, [1]) + assert partial_full_block is not None + assert partial_mamba_block is not None + + req1 = make_request("dcp-replay", [0, 0, 1, 1, 2, 2, 3, 3], 2, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req1) + assert num_computed == 6 + full_block_size = physical_block_size * dcp_world_size + assert [len(group) for group in computed_blocks.blocks] == [ + (6 + full_block_size - 1) // full_block_size, + 2, + ] + + new_blocks = manager.allocate_slots(req1, 2, num_computed, computed_blocks) + assert new_blocks is not None + full_new_block_id = new_blocks.get_block_ids()[0][0] + mamba_new_block_id = new_blocks.get_block_ids()[1][0] + copies, retained = manager.take_kv_cache_block_copies() + assert KVCacheBlockCopy(partial_full_block[0].block_id, full_new_block_id) in copies + assert ( + KVCacheBlockCopy(partial_mamba_block[0].block_id, mamba_new_block_id) in copies + ) + manager.block_pool.free_blocks(retained) + + +@pytest.mark.parametrize("dcp_world_size", [2, 4]) +def test_dcp_partial_hit_resumes_on_replicated_mamba_snapshot( + dcp_world_size: int, +): + block_size = 4 + manager = make_full_mamba_manager( + dcp_world_size=dcp_world_size, + hash_block_size=block_size, + full_block_size=block_size, + mamba_block_size=block_size, + ) + assert manager.coordinator.enable_partial_hash_hits + assert manager.coordinator._cache_hit_alignment_tokens == block_size + assert manager.coordinator.single_type_managers[0].block_size == ( + block_size * dcp_world_size + ) + assert manager.coordinator.single_type_managers[1].block_size == block_size + + prefix = list(range(12)) + req0 = make_request("snapshot-owner", prefix, block_size, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) + assert manager.allocate_slots(req0, 12, num_computed, computed_blocks) is not None + manager.free(req0) + manager.new_step_starts() + + req1 = make_request( + "snapshot-replay", prefix + list(range(12, 16)), block_size, sha256 + ) + + computed_blocks, num_computed, _ = manager.get_computed_blocks(req1) + assert num_computed == 12 + assert [len(group) for group in computed_blocks.blocks] == [ + (12 + block_size * dcp_world_size - 1) // (block_size * dcp_world_size), + 3, + ] + partial_full_block = computed_blocks.blocks[0][-1] + new_blocks = manager.allocate_slots(req1, 4, num_computed, computed_blocks) + assert new_blocks is not None + full_new_block_id = new_blocks.get_block_ids()[0][0] + copies, retained = manager.take_kv_cache_block_copies() + assert KVCacheBlockCopy(partial_full_block.block_id, full_new_block_id) in copies + assert all( + copy.src_block_id != computed_blocks.blocks[1][-1].block_id for copy in copies + ) + manager.block_pool.free_blocks(retained) + + +def test_dcp_joint_hit_is_bounded_by_replicated_mamba_snapshots(): + block_size = 4 + manager = make_full_mamba_manager( + dcp_world_size=2, + hash_block_size=block_size, + full_block_size=block_size, + mamba_block_size=block_size, + ) + prefix = list(range(12)) + req0 = make_request("joint-owner", prefix, block_size, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) + assert manager.allocate_slots(req0, 8, num_computed, computed_blocks) is not None + manager.new_step_starts() + + replay_tokens = prefix + list(range(12, 16)) + req1 = make_request("joint-before", replay_tokens, block_size, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req1) + assert num_computed == 8 + assert [len(group) for group in computed_blocks.blocks] == [1, 2] + + req0.num_computed_tokens = 8 + assert manager.allocate_slots(req0, 4) is not None + manager.new_step_starts() + + req2 = make_request("joint-after", replay_tokens, block_size, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req2) + assert num_computed == 12 + assert [len(group) for group in computed_blocks.blocks] == [2, 3] + + +def test_dcp_partial_hit_with_eagle_rewinds_one_hash_unit(): + hash_block_size = 2 + manager = make_full_mamba_manager( + dcp_world_size=2, + hash_block_size=hash_block_size, + full_block_size=4, + mamba_block_size=4, + use_eagle=True, + ) + + req0 = make_request("eagle-owner", [7] * 6, hash_block_size, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) + assert manager.allocate_slots(req0, 4, num_computed, computed_blocks) is not None + req0.num_computed_tokens = 4 + manager.new_step_starts() + assert manager.allocate_slots(req0, 2) is not None + req0.num_computed_tokens = 6 + manager.new_step_starts() + + req1 = make_request("eagle-replay", [7] * 6 + [9] * 2, 2, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req1) + assert num_computed == 4 + assert [len(group) for group in computed_blocks.blocks] == [1, 1] + assert manager.allocate_slots(req1, 4, num_computed, computed_blocks) is not None diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_primitives.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_primitives.py index 225d87e86796..bbb7d2626017 100644 --- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_primitives.py +++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_primitives.py @@ -55,9 +55,9 @@ def cache_full_block_and_partial_tail( token_ids: list[int], *, enable_kv_cache_events: bool = False, + block_size: int = 6, ) -> tuple[BlockPool, Request, list[KVCacheBlock], BlockHash]: hash_block_size = 2 - block_size = 6 kv_cache_group_id = 0 req = make_request("0", token_ids, hash_block_size, sha256) pool = BlockPool( @@ -389,27 +389,34 @@ def test_reset_prefix_cache_clears_partial_entry_metadata(): assert pool.cached_block_hashes_by_block == {} -def test_evict_cached_block_removes_full_hash_and_partial_entry(): - pool, req, blocks, partial_hash_10 = cache_full_block_and_partial_tail( - [0, 0, 1, 1, 2, 2, 3, 3, 4, 4] +@pytest.mark.parametrize("dcp_world_size", [1, 2, 4]) +def test_evict_cached_block_removes_full_hash_and_partial_entry( + dcp_world_size: int, +): + block_size = 6 * dcp_world_size + partial_num_tokens = 2 * block_size - 2 + pool, req, blocks, partial_hash = cache_full_block_and_partial_tail( + list(range(partial_num_tokens)), block_size=block_size ) - full_hash = BlockHashListWithBlockSize(req.block_hashes, 2, 6)[0] + full_hash = BlockHashListWithBlockSize(req.block_hashes, 2, block_size)[0] assert pool.get_cached_block(full_hash, [0]) == [blocks[0]] - assert pool.get_cached_block(partial_hash_10, [0]) == [blocks[1]] + assert pool.get_cached_block(partial_hash, [0]) == [blocks[1]] pool.evict_blocks({blocks[0].block_id, blocks[1].block_id}) assert pool.get_cached_block(full_hash, [0]) is None - assert pool.get_cached_block(partial_hash_10, [0]) is None + assert pool.get_cached_block(partial_hash, [0]) is None assert pool.cached_block_hashes_by_block == {} -def test_partial_block_promotes_to_direct_full_block_hash(): +@pytest.mark.parametrize("dcp_world_size", [1, 2, 4]) +def test_partial_block_promotes_to_direct_full_block_hash(dcp_world_size: int): hash_block_size = 2 - block_size = 6 + block_size = 6 * dcp_world_size kv_cache_group_id = 0 - token_ids = [0, 0, 1, 1, 2, 2, 3, 3, 4, 4] + partial_num_tokens = 2 * block_size - hash_block_size + token_ids = list(range(partial_num_tokens)) req = make_request("0", token_ids, hash_block_size, sha256) pool = BlockPool( num_gpu_blocks=3, @@ -426,27 +433,22 @@ def test_partial_block_promotes_to_direct_full_block_hash(): block_size=block_size, kv_cache_group_id=kv_cache_group_id, ) - partial_hash_10 = boundary_hash(req, hash_block_size, 10) + partial_hash = boundary_hash(req, hash_block_size, partial_num_tokens) assert pool.cache_partial_block( request=req, block=blocks[1], - num_tokens=10, + num_tokens=partial_num_tokens, kv_cache_group_id=kv_cache_group_id, block_size=block_size, ) - assert pool.get_cached_block(partial_hash_10, [kv_cache_group_id]) == [blocks[1]] + assert pool.get_cached_block(partial_hash, [kv_cache_group_id]) == [blocks[1]] - req.append_output_token_ids([5, 5]) + req.append_output_token_ids(list(range(partial_num_tokens, 2 * block_size))) full_hashes = BlockHashListWithBlockSize( req.block_hashes, hash_block_size, block_size ) promoted_full_hash = full_hashes[1] - # The promoted full-block hash is the fine hash at the 12-token boundary, - # not a concatenation of the fine hashes inside the block. - assert promoted_full_hash == req.block_hashes[12 // hash_block_size - 1] - assert promoted_full_hash != BlockHash( - req.block_hashes[3] + req.block_hashes[4] + req.block_hashes[5] - ) + assert promoted_full_hash == req.block_hashes[2 * block_size // hash_block_size - 1] pool.cache_full_blocks( request=req, @@ -457,4 +459,4 @@ def test_partial_block_promotes_to_direct_full_block_hash(): kv_cache_group_id=kv_cache_group_id, ) assert pool.get_cached_block(promoted_full_hash, [kv_cache_group_id]) == [blocks[1]] - assert pool.get_cached_block(partial_hash_10, [kv_cache_group_id]) is None + assert pool.get_cached_block(partial_hash, [kv_cache_group_id]) is None diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py new file mode 100644 index 000000000000..59c11e79c49e --- /dev/null +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +import vllm.v1.worker.gpu.model_runner as model_runner_module +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, + MambaSpec, +) +from vllm.v1.worker.gpu.block_table import BlockTables +from vllm.v1.worker.gpu.model_runner import GPUModelRunner + + +@pytest.mark.parametrize( + ("mamba_cache_mode", "num_speculative_blocks", "expected"), + [ + pytest.param("align", 0, 65_536, id="align-prefix-cache"), + pytest.param("none", 7, 8, id="no-prefix-cache-with-speculation"), + ], +) +def test_initialize_kv_cache_does_not_dcp_shard_mamba_block_table( + monkeypatch, + mamba_cache_mode: str, + num_speculative_blocks: int, + expected: int, +): + """Mamba/GDN block-table rows index global positions, unlike DCP KV.""" + + max_model_len = 1_048_576 + attention_block_size = 1_536 + mamba_block_size = 16 + dcp_size = 8 + full_attention_spec = FullAttentionSpec( + block_size=attention_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.bfloat16, + ) + mamba_spec = MambaSpec( + shapes=((1,),), + dtypes=(torch.bfloat16,), + block_size=mamba_block_size, + mamba_cache_mode=mamba_cache_mode, + num_speculative_blocks=num_speculative_blocks, + ) + kv_cache_config = KVCacheConfig( + num_blocks=1, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec(["attention"], full_attention_spec), + KVCacheGroupSpec(["kda"], mamba_spec), + ], + ) + vllm_config = SimpleNamespace( + parallel_config=SimpleNamespace(decode_context_parallel_size=dcp_size), + cache_config=SimpleNamespace(mamba_cache_mode=mamba_cache_mode), + ) + runner = SimpleNamespace( + max_model_len=max_model_len, + is_encoder_decoder=False, + vllm_config=vllm_config, + ) + + class _CapturedWidths(Exception): + pass + + captured: list[int] = [] + + def capture_width(max_num_blocks: int, *_args, **_kwargs) -> int: + captured.append(max_num_blocks) + if len(captured) == 2: + raise _CapturedWidths + return max_num_blocks + + monkeypatch.setattr(model_runner_module, "get_block_table_width", capture_width) + + with pytest.raises(_CapturedWidths): + GPUModelRunner.initialize_kv_cache(runner, kv_cache_config) + + # Attention KV is local to one of eight DCP ranks; KDA state is replicated + # and therefore needs one table entry for every global 16-token page. + assert captured == [86, expected] + + +def test_append_block_ids_rejects_write_past_row_capacity(): + """Reject an oversized staged write before it can corrupt the next row.""" + + class _BlockTable: + gpu = torch.empty((2, 4), dtype=torch.int32) + + def stage_write(self, *_args): + pytest.fail("an oversized write must not be staged") + + block_tables = BlockTables.__new__(BlockTables) + block_tables.num_kv_cache_groups = 1 + block_tables.blocks_per_kv_block = [1] + block_tables.block_tables = [_BlockTable()] + block_tables.num_blocks = SimpleNamespace( + np=torch.tensor([[0, 3]], dtype=torch.int32) + ) + + with pytest.raises( + RuntimeError, + match=r"request 1, group 0 exceeds row capacity \(5 > 4\)", + ): + block_tables.append_block_ids( + req_index=1, + new_block_ids=([4, 5],), + overwrite=False, + ) + + assert block_tables.num_blocks.np[0, 1] == 3 diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index 0829ec0d2be8..767ea0033483 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -619,15 +619,22 @@ def __init__( "full-attention and Mamba groups, got: " f"{type(g.kv_cache_spec).__name__}." ) - # Fine-grained hash hits require Mamba "align", no context - # parallelism, and compatible cache managers in every group. + # Fine-grained hash hits require Mamba "align" and compatible cache + # managers in every group. TP needs hashing finer than the Mamba block; + # DCP accepts equality because it scales the effective full-attention + # block instead. has_partial_mamba_group = any( isinstance(g.kv_cache_spec, MambaSpec) and g.kv_cache_spec.mamba_cache_mode == "align" - and g.kv_cache_spec.block_size > hash_block_size + and ( + (dcp_world_size == 1 and g.kv_cache_spec.block_size > hash_block_size) + or ( + dcp_world_size > 1 and g.kv_cache_spec.block_size >= hash_block_size + ) + ) for g in kv_cache_config.kv_cache_groups ) - self.enable_partial_hash_hits = dcp_world_size == 1 and has_partial_mamba_group + self.enable_partial_hash_hits = has_partial_mamba_group if self.enable_partial_hash_hits: unsupported_partial_hit_managers = { type(manager).__name__ diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py index 90fc104dc29f..fbc4f383cb35 100644 --- a/vllm/v1/worker/gpu/block_table.py +++ b/vllm/v1/worker/gpu/block_table.py @@ -116,8 +116,15 @@ def append_block_ids( bpk = self.blocks_per_kv_block[i] if bpk > 1: block_ids = [b * bpk + k for b in block_ids for k in range(bpk)] + end = start + len(block_ids) + row_capacity = self.block_tables[i].gpu.shape[1] + if end > row_capacity: + raise RuntimeError( + f"Block table write for request {req_index}, group {i} exceeds " + f"row capacity ({end} > {row_capacity})" + ) self.block_tables[i].stage_write(req_index, start, block_ids) - self.num_blocks.np[i, req_index] = start + len(block_ids) + self.num_blocks.np[i, req_index] = end def apply_staged_writes(self) -> None: if self.num_kv_cache_groups == 0: diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 77ae59f0ee38..4f65769c9580 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -60,7 +60,6 @@ ) from vllm.sequence import IntermediateTensors from vllm.tasks import SupportedTask -from vllm.utils.math_utils import cdiv from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput @@ -517,17 +516,14 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: for kv_cache_group in kv_cache_config.kv_cache_groups: spec = kv_cache_group.kv_cache_spec block_sizes.append(spec.block_size) - # When using DCP, each request's KV cache is sharded among different ranks. - # As a result, one block on the current rank covers `block_size * cp_size` - # tokens in the full, global (unsharded) sequence. - max_num_blocks = cdiv( - block_table_max_model_len, spec.block_size * self.dcp_size + # Let each cache type account for CP. Attention KV is DCP-sharded, + # while Mamba/GDN recurrent state is replicated across DCP ranks. + max_num_blocks = spec.max_num_blocks_per_req( + self.vllm_config, block_table_max_model_len ) - # For Mamba/Hybrid Model, KVCaches need extra blocks for speculative tokens + # Preserve each cache type's alignment requirements after applying + # its topology-aware block-table width. if isinstance(spec, MambaSpec): - max_num_blocks = ( - max_num_blocks if self.cache_config.enable_prefix_caching else 1 - ) + spec.num_speculative_blocks max_num_blocks = get_block_table_width( max_num_blocks, spec.block_size, token_alignment=None ) From c296851a7d173fa89d2eefbca0243be42ae9b5e0 Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Tue, 18 Aug 2026 08:18:14 +0800 Subject: [PATCH 070/839] [MoE] Refine FlashInfer one-sided All2All integration (#51924) Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- docs/design/moe_kernel_features.md | 2 +- tests/distributed/test_mnnvl_alltoall.py | 30 ++++---- tests/kernels/moe/test_moe_layer.py | 2 +- vllm/config/parallel.py | 1 + .../device_communicators/all2all.py | 12 +-- .../layers/fused_moe/all2all_utils.py | 74 +++++++++++++------ .../fused_moe/experts/trtllm_fp8_moe.py | 28 ++++++- .../flashinfer_nvlink_one_sided.py | 48 ++++++------ 8 files changed, 129 insertions(+), 68 deletions(-) diff --git a/docs/design/moe_kernel_features.md b/docs/design/moe_kernel_features.md index bdf407510294..0b1a9151ec69 100644 --- a/docs/design/moe_kernel_features.md +++ b/docs/design/moe_kernel_features.md @@ -36,7 +36,7 @@ th { | deepep_high_throughput | standard | fp8 | G(128),A,T2 | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht.DeepEPHTPrepareAndFinalize] | | deepep_low_latency | batched | fp8 | G(128),A,T3 | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll.DeepEPLLPrepareAndFinalize] | | flashinfer_nvlink_two_sided | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferNVLinkTwoSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided.FlashInferNVLinkTwoSidedPrepareAndFinalize] | -| flashinfer_nvlink_one_sided | standard | nvfp4,bf16,mxfp8 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] | +| flashinfer_nvlink_one_sided | standard | nvfp4,bf16,mxfp8,fp8 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] | !!! info "Table key" 1. All types: mxfp4, nvfp4, int4, int8, fp8 diff --git a/tests/distributed/test_mnnvl_alltoall.py b/tests/distributed/test_mnnvl_alltoall.py index f9b76ea1b030..636f7cb525f5 100644 --- a/tests/distributed/test_mnnvl_alltoall.py +++ b/tests/distributed/test_mnnvl_alltoall.py @@ -342,6 +342,8 @@ def _one_sided_lifecycle_worker(rank, world_size): top_k=2, num_experts=world_size * 8, hidden_size=4096, + x_bytes_per_token=4096 * 2, + x_sf_bytes_per_token=0, ) # Initialize @@ -409,12 +411,12 @@ def _one_sided_workspace_grow_worker(rank, world_size): hidden_size=4096, ) nvfp4_kwargs = dict( - dispatch_dtype_bytes_per_elem=0, - dispatch_scale_bytes_per_token=base_kwargs["hidden_size"] // 16, + x_bytes_per_token=base_kwargs["hidden_size"] // 2, + x_sf_bytes_per_token=base_kwargs["hidden_size"] // 16, ) bf16_kwargs = dict( - dispatch_dtype_bytes_per_elem=2, - dispatch_scale_bytes_per_token=0, + x_bytes_per_token=base_kwargs["hidden_size"] * 2, + x_sf_bytes_per_token=0, ) # First init: NVFP4-like (hidden_bytes = hidden // 2 + hidden // 16). @@ -780,11 +782,12 @@ def _one_sided_data_worker(rank, world_size): top_k=experts_per_token, num_experts=num_experts, hidden_size=hidden_size, - # Account for the fp8 block-scale payload (a1q_scale: hidden//16 bytes + x_bytes_per_token=hidden_size // 2, + # Account for the fp8 block-scale payload (x_sf: hidden//16 bytes # per token) that is dispatched alongside the nvfp4 hidden states. # Without this the dispatch region is under-reserved and the combine # payload overflows the per-rank workspace. - dispatch_scale_bytes_per_token=hidden_size // 16, + x_sf_bytes_per_token=hidden_size // 16, ) assert manager.initialized assert manager.moe_alltoall is not None @@ -798,17 +801,17 @@ def _one_sided_data_worker(rank, world_size): # Create test data with raw tensors matching the nvfp4 payload # sizes the workspace was allocated for: - # a1q: (tokens, hidden_size // 2) — nvfp4 hidden states - # a1q_scale: (tokens, hidden_size // 16) — fp8 scaling factors + # x: (tokens, hidden_size // 2) — nvfp4 hidden states + # x_sf: (tokens, hidden_size // 16) — fp8 scaling factors torch.manual_seed(rank + 42) - a1q = torch.randint( + x = torch.randint( 0, 256, (tokens_per_rank, hidden_size // 2), device=device, dtype=torch.uint8, ) - a1q_scale = torch.randint( + x_sf = torch.randint( 0, 256, (tokens_per_rank, hidden_size // 16), @@ -830,15 +833,16 @@ def _one_sided_data_worker(rank, world_size): ) # --- One-sided dispatch --- - payloads = [a1q, a1q_scale, topk_ids, topk_weights] + payloads = [x, x_sf, topk_ids, topk_weights] recv_payloads = manager.moe_alltoall.dispatch( token_selected_experts=topk_ids, input_payloads=payloads, runtime_max_tokens_per_rank=runtime_max_tokens, ) assert len(recv_payloads) == 4 - recv_a1q, recv_scale, recv_ids, recv_weights = recv_payloads - assert recv_a1q.numel() > 0 + recv_x, recv_x_sf, recv_ids, recv_weights = recv_payloads + assert recv_x.numel() > 0 + assert recv_x_sf.numel() > 0 assert recv_ids.numel() > 0 # --- Round-trip exact verification --- diff --git a/tests/kernels/moe/test_moe_layer.py b/tests/kernels/moe/test_moe_layer.py index 1942ce96bcca..74554944a314 100644 --- a/tests/kernels/moe/test_moe_layer.py +++ b/tests/kernels/moe/test_moe_layer.py @@ -130,7 +130,7 @@ "mori_high_throughput": {None, "fp8", "modelopt_fp8"}, "mori_low_latency": {None, "fp8", "modelopt_fp8"}, "flashinfer_nvlink_two_sided": {None, "fp8_blocked", "modelopt_fp4"}, # noqa: E501 - "flashinfer_nvlink_one_sided": {None, "modelopt_fp4"}, # noqa: E501 + "flashinfer_nvlink_one_sided": {None, "fp8_blocked", "modelopt_fp4"}, # noqa: E501 "deepep_low_latency": {None, "fp8_blocked", "modelopt_fp4"}, # noqa: E501 "deepep_high_throughput": {None, "fp8_blocked", "modelopt_fp8", "modelopt_fp4"}, # noqa: E501 "nixl_ep": {None, "fp8_blocked", "modelopt_fp8"}, diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 83c4e6fe53d7..eb346ed4a34a 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -677,6 +677,7 @@ def use_sequence_parallel_moe(self) -> bool: "allgather_reducescatter", "deepep_high_throughput", "deepep_low_latency", + "flashinfer_nvlink_one_sided", "mori_high_throughput", "mori_low_latency", "nixl_ep", diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index 5abe7568a292..657ad9fd0e13 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -734,17 +734,13 @@ def initialize( top_k: int, num_experts: int, hidden_size: int, - dispatch_dtype_bytes_per_elem: int = 0, - dispatch_scale_bytes_per_token: int = 0, + x_bytes_per_token: int, + x_sf_bytes_per_token: int, ): """Initialize (or grow) the MoeAlltoAll workspace.""" - if dispatch_dtype_bytes_per_elem == 0: - hidden_bytes = hidden_size // 2 - else: - hidden_bytes = hidden_size * dispatch_dtype_bytes_per_elem total_dispatch_payload_size_per_token = ( - hidden_bytes - + dispatch_scale_bytes_per_token + x_bytes_per_token + + x_sf_bytes_per_token + top_k * 4 # int32 topks ids + top_k * 4 # float32 topk weights ) diff --git a/vllm/model_executor/layers/fused_moe/all2all_utils.py b/vllm/model_executor/layers/fused_moe/all2all_utils.py index 29905927c106..7861345453a0 100644 --- a/vllm/model_executor/layers/fused_moe/all2all_utils.py +++ b/vllm/model_executor/layers/fused_moe/all2all_utils.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from dataclasses import dataclass from typing import Any import torch @@ -37,6 +38,52 @@ has_nixl_ep, ) + +@dataclass(frozen=True) +class FlashInferOneSidedDispatchLayout: + x_bytes_per_token: int + x_sf_bytes_per_token: int + + +def flashinfer_one_sided_dispatch_layout( + hidden_dim: int, quant_config: FusedMoEQuantConfig +) -> FlashInferOneSidedDispatchLayout: + """Return the one-sided activation payload layout.""" + if quant_config.quant_dtype is None: + return FlashInferOneSidedDispatchLayout(hidden_dim * 2, 0) + if quant_config.quant_dtype == "nvfp4": + scale_elems = hidden_dim // 16 + return FlashInferOneSidedDispatchLayout(hidden_dim // 2, scale_elems) + if quant_config.quant_dtype == "mxfp8": + align = quant_config.mx_alignment + padded_k = ( + ((hidden_dim + align - 1) // align) * align if align > 0 else hidden_dim + ) + scale_elems = padded_k // 32 + return FlashInferOneSidedDispatchLayout(hidden_dim, scale_elems) + if ( + quant_config.use_fp8_w8a8 + and quant_config.quant_dtype == current_platform.fp8_dtype() + and quant_config.block_shape == [128, 128] + ): + if hidden_dim % 128 != 0: + raise NotImplementedError( + "flashinfer_nvlink_one_sided DeepSeek Blockwise FP8 dispatch " + f"requires hidden_dim divisible by 128; got {hidden_dim}" + ) + scale_elems = hidden_dim // 128 + scale_bytes = scale_elems * torch.float32.itemsize + return FlashInferOneSidedDispatchLayout(hidden_dim, scale_bytes) + raise NotImplementedError( + "flashinfer_nvlink_one_sided dispatch supports nvfp4, mxfp8, " + "DeepSeek Blockwise FP8 (E4M3 with FP32 1x128 scales), and bf16 " + "(quant_dtype=None) today; got " + f"quant_dtype={quant_config.quant_dtype!r}, " + f"use_fp8_w8a8={quant_config.use_fp8_w8a8!r}, " + f"block_shape={quant_config.block_shape!r}" + ) + + logger = init_logger(__name__) if current_platform.is_cuda_alike(): @@ -282,34 +329,17 @@ def maybe_make_prepare_finalize( max_num_tokens = ( get_current_vllm_config().scheduler_config.max_num_batched_tokens ) - if quant_config.quant_dtype is None: - dispatch_dtype_bytes_per_elem = 2 - dispatch_scale_bytes_per_token = 0 - elif quant_config.quant_dtype == "nvfp4": - dispatch_dtype_bytes_per_elem = 0 - dispatch_scale_bytes_per_token = moe.hidden_dim // 16 - elif quant_config.quant_dtype == "mxfp8": - dispatch_dtype_bytes_per_elem = 1 - align = quant_config.mx_alignment - if align > 0: - padded_k = ((moe.hidden_dim + align - 1) // align) * align - else: - padded_k = moe.hidden_dim - dispatch_scale_bytes_per_token = padded_k // 32 - else: - raise NotImplementedError( - "flashinfer_nvlink_one_sided dispatch supports nvfp4, mxfp8, " - "and bf16 (quant_dtype=None) today; got " - f"quant_dtype={quant_config.quant_dtype!r}" - ) + dispatch_layout = flashinfer_one_sided_dispatch_layout( + moe.hidden_dim, quant_config + ) prepare_finalize = FlashInferNVLinkOneSidedPrepareAndFinalize( max_num_tokens=max_num_tokens, top_k=moe.experts_per_token, num_experts=moe.num_experts, hidden_size=moe.hidden_dim, num_dispatchers=all2all_manager.world_size, - dispatch_dtype_bytes_per_elem=dispatch_dtype_bytes_per_elem, - dispatch_scale_bytes_per_token=dispatch_scale_bytes_per_token, + x_bytes_per_token=dispatch_layout.x_bytes_per_token, + x_sf_bytes_per_token=dispatch_layout.x_sf_bytes_per_token, ) elif moe.use_ag_rs_all2all_kernels and allow_new_interface: diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py index d32877eb9a95..00ab06fe952b 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py @@ -36,6 +36,28 @@ logger = init_logger(__name__) +def prepare_deepseek_fp8_x_sf(x: torch.Tensor, x_sf: torch.Tensor) -> torch.Tensor: + """Validate DeepSeek Blockwise FP8 tensors and return TRTLLM layout.""" + if x.dtype != current_platform.fp8_dtype(): + raise ValueError( + f"DeepSeekFp8 activations must use the platform E4M3 dtype; got {x.dtype}" + ) + if x.ndim != 2 or x.shape[1] % 128 != 0: + raise ValueError( + "DeepSeekFp8 activations must be [M,K] with K divisible by 128; " + f"got {tuple(x.shape)}" + ) + expected_shape = (x.shape[0], x.shape[1] // 128) + if x_sf.dtype != torch.float32 or tuple(x_sf.shape) != expected_shape: + raise ValueError( + "DeepSeekFp8 activation scales must be FP32 [M,K/128]; " + f"expected {expected_shape}, got dtype={x_sf.dtype}, " + f"shape={tuple(x_sf.shape)}" + ) + # FlashInfer TRTLLM-gen consumes [K/128,M] for DeepSeekFp8/BlockMajorK. + return x_sf.t().contiguous() + + class TrtLlmFp8ExpertsBase: """ Fp8 TRTLLM-Gen MoE kernels. Shared base for modular and monolithic @@ -239,6 +261,10 @@ def apply( # Pack topk ids and weights into format expected by the kernel. packed_topk_ids = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights) + if a1q_scale is None: + raise RuntimeError( + "TRT-LLM FP8 experts require precomputed activation scales" + ) assert a1q_scale is not None is_mxfp8 = self.quant_config.block_shape == [1, 32] @@ -251,7 +277,7 @@ def apply( fp8_quant_type = Fp8QuantizationType.DeepSeekFp8 use_shuffled_weight = True weight_layout = WeightLayout.BlockMajorK - hidden_states_scale = a1q_scale.t().contiguous() + hidden_states_scale = prepare_deepseek_fp8_x_sf(hidden_states, a1q_scale) flashinfer.fused_moe.trtllm_fp8_block_scale_routed_moe( topk_ids=packed_topk_ids, diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py index 74341e7681f8..0b57afb6062c 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py @@ -30,9 +30,9 @@ def __init__( top_k: int, num_experts: int, hidden_size: int, + x_bytes_per_token: int, + x_sf_bytes_per_token: int, num_dispatchers: int = 1, - dispatch_dtype_bytes_per_elem: int = 0, - dispatch_scale_bytes_per_token: int = 0, ): super().__init__() self.max_num_tokens = max_num_tokens @@ -40,7 +40,6 @@ def __init__( self.num_experts = num_experts self.hidden_size = hidden_size self.num_dispatchers_ = num_dispatchers - self.scale_elems_per_token = dispatch_scale_bytes_per_token device_communicator = get_ep_group().device_communicator assert device_communicator is not None @@ -52,8 +51,8 @@ def __init__( top_k=self.top_k, num_experts=self.num_experts, hidden_size=self.hidden_size, - dispatch_dtype_bytes_per_elem=dispatch_dtype_bytes_per_elem, - dispatch_scale_bytes_per_token=dispatch_scale_bytes_per_token, + x_bytes_per_token=x_bytes_per_token, + x_sf_bytes_per_token=x_sf_bytes_per_token, ) @property @@ -98,9 +97,9 @@ def prepare( ) if defer_input_quant: - a1q, a1q_scale = a1, None + dispatch_x, dispatch_x_sf = a1, None else: - a1q, a1q_scale = moe_kernel_quantize_input( + dispatch_x, dispatch_x_sf = moe_kernel_quantize_input( a1, quant_config.a1_gscale, quant_config.quant_dtype, @@ -110,10 +109,9 @@ def prepare( mx_alignment=quant_config.mx_alignment, ) - payloads = [] - payloads.append(a1q) - if a1q_scale is not None: - payloads.append(a1q_scale) + payloads = [dispatch_x] + if dispatch_x_sf is not None: + payloads.append(dispatch_x_sf) topk_ids_payload_index = len(payloads) payloads.append(topk_ids) payloads.append(topk_weights) @@ -126,23 +124,29 @@ def prepare( invalid_token_expert_id=-1, # Follow TRTLLM Pattern expert_id_payload_index=topk_ids_payload_index, ) - if a1q_scale is not None: - a1q_recv, a1q_scale_recv, topk_ids_recv, topk_weights_recv = recv_payloads + if dispatch_x_sf is not None: + recv_x, recv_x_sf, topk_ids_recv, topk_weights_recv = recv_payloads + x_sf_width = recv_x_sf.shape[-1] # Apply scale interleaving only for CUTLASS (not TRT-LLM) if quant_config.quant_dtype == "nvfp4" and quant_config.is_scale_swizzled: - a1q_scale_recv = a1q_scale_recv.view(-1, a1q_scale_recv.shape[-1]) - a1q_scale_recv = a1q_scale_recv.view(torch.uint8) - a1q_scale_recv = nvfp4_block_scale_interleave(a1q_scale_recv) - assert self.scale_elems_per_token > 0 - a1q_scale_recv = a1q_scale_recv.view(-1, self.scale_elems_per_token) + recv_x_sf = recv_x_sf.view(-1, x_sf_width) + recv_x_sf = recv_x_sf.view(torch.uint8) + recv_x_sf = nvfp4_block_scale_interleave(recv_x_sf) + recv_x_sf = recv_x_sf.view(-1, x_sf_width) else: - a1q_recv, topk_ids_recv, topk_weights_recv = recv_payloads - a1q_scale_recv = None - a1q_recv = a1q_recv.view(-1, a1q_recv.shape[-1]) + recv_x, topk_ids_recv, topk_weights_recv = recv_payloads + recv_x_sf = None + recv_x = recv_x.view(-1, recv_x.shape[-1]) topk_ids_recv = topk_ids_recv.view(-1, topk_ids_recv.shape[-1]) topk_weights_recv = topk_weights_recv.view(-1, topk_weights_recv.shape[-1]) - return a1q_recv, a1q_scale_recv, None, topk_ids_recv, topk_weights_recv + return ( + recv_x, + recv_x_sf, + None, + topk_ids_recv, + topk_weights_recv, + ) def finalize( self, From cdb8545a91be16f2b80234e0801458ad2a18cf1b Mon Sep 17 00:00:00 2001 From: kiroxu <148877251+BabyDrangoner@users.noreply.github.com> Date: Tue, 18 Aug 2026 09:17:40 +0800 Subject: [PATCH 071/839] [Kernel][Perf] Support Qwen head ratios in fused GDN MTP (#52539) Signed-off-by: kiroxu <148877251+BabyDrangoner@users.noreply.github.com> --- .../gdn/fused_gdn_decode_kernel.cu | 72 ++++++++++++------- tests/kernels/mamba/test_gdn_fused_mtp.py | 42 +++++++++++ tests/kernels/test_fused_gdn_post_conv.py | 59 ++++++++++++--- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 3 +- 4 files changed, 142 insertions(+), 34 deletions(-) diff --git a/csrc/libtorch_stable/gdn/fused_gdn_decode_kernel.cu b/csrc/libtorch_stable/gdn/fused_gdn_decode_kernel.cu index 31a16fcf9965..1166b41868e0 100644 --- a/csrc/libtorch_stable/gdn/fused_gdn_decode_kernel.cu +++ b/csrc/libtorch_stable/gdn/fused_gdn_decode_kernel.cu @@ -146,7 +146,7 @@ __device__ __forceinline__ Sum2 warp_reduce_sum_pair(float x, float y) { return {x, y}; } -template +template __global__ __launch_bounds__(kThreads, 2) void gdn_decode_post_conv_mtp_kernel( const __nv_bfloat16* __restrict__ mixed_qkv, const __nv_bfloat16* __restrict__ a, const __nv_bfloat16* __restrict__ b, @@ -186,7 +186,7 @@ __global__ __launch_bounds__(kThreads, 2) void gdn_decode_post_conv_mtp_kernel( return; } - const int key_head = value_head / 8; + const int key_head = value_head / ValueHeadsPerKeyHead; __shared__ StateT shared_state[2][kChunkV][kDimK]; __shared__ float shared_q[kMaxMtpTokens][kDimK]; __shared__ float shared_k[kMaxMtpTokens][kDimK]; @@ -370,7 +370,7 @@ __global__ __launch_bounds__(kThreads, 2) void gdn_decode_post_conv_mtp_kernel( } } -template +template void launch_gdn_decode_post_conv_mtp( torch::stable::Tensor const& mixed_qkv, torch::stable::Tensor const& a_log, torch::stable::Tensor const& dt_bias, @@ -395,17 +395,19 @@ void launch_gdn_decode_post_conv_mtp( get_current_cuda_stream(mixed_qkv.get_device_index()); const int num_requests = static_cast(state_indices.size(0)); const dim3 grid(num_requests, num_value_heads); - gdn_decode_post_conv_mtp_kernel<<>>( - static_cast(mixed_qkv.data_ptr()), a, b, - static_cast(a_log.data_ptr()), dt_bias.data_ptr(), - static_cast(state_indices.data_ptr()), - static_cast(cu_seqlens.data_ptr()), - static_cast(num_accepted_tokens.data_ptr()), - static_cast(state.data_ptr()), output_gate, - norm_weight.data_ptr(), static_cast<__nv_bfloat16*>(out.data_ptr()), - num_key_heads, num_value_heads, static_cast(state_indices.size(1)), - dt_bias_type, norm_weight.scalar_type() == ScalarType::BFloat16, - static_cast(scale), static_cast(norm_eps), strides); + gdn_decode_post_conv_mtp_kernel + <<>>( + static_cast(mixed_qkv.data_ptr()), a, b, + static_cast(a_log.data_ptr()), dt_bias.data_ptr(), + static_cast(state_indices.data_ptr()), + static_cast(cu_seqlens.data_ptr()), + static_cast(num_accepted_tokens.data_ptr()), + static_cast(state.data_ptr()), output_gate, + norm_weight.data_ptr(), static_cast<__nv_bfloat16*>(out.data_ptr()), + num_key_heads, num_value_heads, + static_cast(state_indices.size(1)), dt_bias_type, + norm_weight.scalar_type() == ScalarType::BFloat16, + static_cast(scale), static_cast(norm_eps), strides); const cudaError_t error = cudaGetLastError(); STD_TORCH_CHECK(error == cudaSuccess, "GDN decode MTP post-conv kernel launch failed: ", @@ -479,8 +481,12 @@ void fused_gdn_decode_post_conv_mtp( STD_TORCH_CHECK(key_width > 0 && key_width % (2 * kDimK) == 0, "mixed_qkv width is inconsistent with state"); const int num_key_heads = static_cast(key_width / (2 * kDimK)); - STD_TORCH_CHECK(num_value_heads == 8 * num_key_heads, - "GDN decode MTP fusion requires HV/H=8"); + const int value_heads_per_key_head = num_value_heads / num_key_heads; + STD_TORCH_CHECK( + num_value_heads % num_key_heads == 0 && + ((value_heads_per_key_head >= 1 && value_heads_per_key_head <= 4) || + value_heads_per_key_head == 8), + "GDN decode MTP fusion requires HV/H in {1, 2, 3, 4, 8}"); STD_TORCH_CHECK(state_indices.dim() == 2 && state_indices.size(0) > 0 && state_indices.size(1) > 0 && @@ -541,17 +547,35 @@ void fused_gdn_decode_post_conv_mtp( const auto* b_ptr = static_cast(b.data_ptr()); const auto* output_gate_ptr = static_cast(output_gate.data_ptr()); - if (state_scalar_type == ScalarType::Float) { - launch_gdn_decode_post_conv_mtp( - mixed_qkv, a_log, dt_bias, state_indices, cu_seqlens, - num_accepted_tokens, state, norm_weight, out, a_ptr, b_ptr, - output_gate_ptr, num_key_heads, num_value_heads, scale, norm_eps, - strides); - } else { - launch_gdn_decode_post_conv_mtp<__nv_bfloat16>( + const auto launch = [&]() { + launch_gdn_decode_post_conv_mtp( mixed_qkv, a_log, dt_bias, state_indices, cu_seqlens, num_accepted_tokens, state, norm_weight, out, a_ptr, b_ptr, output_gate_ptr, num_key_heads, num_value_heads, scale, norm_eps, strides); + }; + const auto dispatch_state_type = [&]() { + if (state_scalar_type == ScalarType::Float) { + launch.template operator()(); + } else { + launch.template operator()<__nv_bfloat16, ValueHeadsPerKeyHead>(); + } + }; + switch (value_heads_per_key_head) { + case 1: + dispatch_state_type.template operator()<1>(); + break; + case 2: + dispatch_state_type.template operator()<2>(); + break; + case 3: + dispatch_state_type.template operator()<3>(); + break; + case 4: + dispatch_state_type.template operator()<4>(); + break; + default: + dispatch_state_type.template operator()<8>(); + break; } } diff --git a/tests/kernels/mamba/test_gdn_fused_mtp.py b/tests/kernels/mamba/test_gdn_fused_mtp.py index 97ed4c48e787..0d78f8ceeba0 100644 --- a/tests/kernels/mamba/test_gdn_fused_mtp.py +++ b/tests/kernels/mamba/test_gdn_fused_mtp.py @@ -5,6 +5,7 @@ from __future__ import annotations import types +from typing import cast from unittest.mock import patch import pytest @@ -37,6 +38,7 @@ ) from vllm.utils.torch_utils import _encode_layer_name # noqa: E402 from vllm.v1.attention.backends.gdn_attn import ( # noqa: E402 + GDNAttentionMetadata, GDNAttentionMetadataBuilder, ) from vllm.v1.kv_cache_interface import MambaSpec # noqa: E402 @@ -139,6 +141,46 @@ def _build_layer( return layer +@pytest.mark.parametrize( + "num_v_heads,expected", + [ + pytest.param(2, True, id="ratio1"), + pytest.param(4, True, id="ratio2"), + pytest.param(6, True, id="ratio3"), + pytest.param(8, True, id="ratio4"), + pytest.param(16, True, id="ratio8"), + pytest.param(5, False, id="non-integral-ratio"), + pytest.param(10, False, id="unsupported-ratio5"), + ], +) +def test_fused_mtp_head_ratio_guard(num_v_heads: int, expected: bool) -> None: + if not hasattr(torch.ops._C, "fused_gdn_decode_post_conv_mtp"): + pytest.skip("fused GDN decode MTP op is not built") + + layer = types.SimpleNamespace( + num_k_heads=2, + num_v_heads=num_v_heads, + kv_cache=(None, torch.empty(1, dtype=torch.float32, device="cuda")), + gdn_decode_kernel="cuda", + ) + attn_metadata = types.SimpleNamespace( + spec_state_indices_tensor=torch.ones( + 1, SPEC_TOKENS, dtype=torch.int32, device="cuda" + ), + spec_sequence_masks=object(), + num_decodes=0, + num_spec_decodes=1, + ) + + assert ( + QwenGatedDeltaNetAttention._can_use_fused_gdn_mtp_decode( + cast(QwenGatedDeltaNetAttention, layer), + cast(GDNAttentionMetadata, attn_metadata), + ) + is expected + ) + + @torch.inference_mode() def test_fused_forward_uses_packed_entrypoint() -> None: """Fused mode keeps projected QKVZ and BA packed through the model op.""" diff --git a/tests/kernels/test_fused_gdn_post_conv.py b/tests/kernels/test_fused_gdn_post_conv.py index b6dacf2e73f2..7cdeb55b9b57 100644 --- a/tests/kernels/test_fused_gdn_post_conv.py +++ b/tests/kernels/test_fused_gdn_post_conv.py @@ -217,18 +217,59 @@ def test_fused_post_conv_l0(): @pytest.mark.parametrize( - "tp_size,query_lengths,state_dtype,norm_dtype", + "head_ratio,tp_size,query_lengths,state_dtype,norm_dtype", [ - pytest.param(16, (4, 4), torch.bfloat16, torch.bfloat16, id="tp16-bf16"), - pytest.param(4, (4, 4), torch.float32, torch.float32, id="tp4-fp32"), - pytest.param(16, (4, 2, 0), torch.bfloat16, torch.float32, id="tp16-ragged"), - pytest.param(4, (4, 2, 0), torch.float32, torch.bfloat16, id="tp4-ragged"), - pytest.param(16, (8,), torch.float32, torch.bfloat16, id="tp16-max"), - pytest.param(4, (8,), torch.bfloat16, torch.float32, id="tp4-max"), + pytest.param(8, 16, (4, 4), torch.bfloat16, torch.bfloat16, id="tp16-bf16"), + pytest.param(8, 4, (4, 4), torch.float32, torch.float32, id="tp4-fp32"), + pytest.param(8, 16, (4, 2, 0), torch.bfloat16, torch.float32, id="tp16-ragged"), + pytest.param(8, 4, (4, 2, 0), torch.float32, torch.bfloat16, id="tp4-ragged"), + pytest.param(8, 16, (8,), torch.float32, torch.bfloat16, id="tp16-max"), + pytest.param(8, 4, (8,), torch.bfloat16, torch.float32, id="tp4-max"), + pytest.param( + 1, + 1, + (4, 4), + torch.float32, + torch.bfloat16, + id="ratio1-tp1-fp32", + ), + pytest.param( + 2, + 1, + (4, 2, 0), + torch.bfloat16, + torch.bfloat16, + id="ratio2-tp1-ragged-bf16", + ), + pytest.param( + 2, + 4, + (8,), + torch.float32, + torch.float32, + id="ratio2-tp4-max-fp32", + ), + pytest.param( + 3, + 2, + (4, 2, 0), + torch.float32, + torch.float32, + id="ratio3-tp2-ragged-fp32", + ), + pytest.param( + 4, + 4, + (8,), + torch.float32, + torch.bfloat16, + id="ratio4-tp4-max-fp32", + ), ], ) @torch.inference_mode() -def test_fused_gdn_decode_post_conv_mtp_ratio8( +def test_fused_gdn_decode_post_conv_mtp_head_ratios( + head_ratio: int, tp_size: int, query_lengths: tuple[int, ...], state_dtype: torch.dtype, @@ -242,7 +283,7 @@ def test_fused_gdn_decode_post_conv_mtp_ratio8( torch.manual_seed(0) device = "cuda" H = 16 // tp_size - HV = 128 // tp_size + HV = head_ratio * H K = V = 128 num_reqs = len(query_lengths) state_width = max(query_lengths) diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 27dc4d90d7e0..d83b076509ee 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -1805,7 +1805,8 @@ def _can_use_fused_gdn_mtp_decode( and attn_metadata.num_spec_decodes > 0 and self.kv_cache[1].dtype in FUSED_GDN_STATE_DTYPES and self.gdn_decode_kernel == "cuda" - and self.num_v_heads == 8 * self.num_k_heads + and self.num_v_heads % self.num_k_heads == 0 + and self.num_v_heads // self.num_k_heads in (1, 2, 3, 4, 8) and state_indices is not None and state_indices.size(1) <= MAX_FUSED_GDN_MTP_TOKENS and hasattr(torch.ops._C, "fused_gdn_decode_post_conv_mtp") From 69d3335066718b7e2b3d00713e02646874b424a9 Mon Sep 17 00:00:00 2001 From: Han <180024497+LH-and-FPGA@users.noreply.github.com> Date: Tue, 18 Aug 2026 03:38:42 +0200 Subject: [PATCH 072/839] [Bugfix][Quantization] Guard the MXFP8 FlashInfer path on FlashInfer availability (#52648) Signed-off-by: Han Li Signed-off-by: Misha Goin Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: Misha Goin --- .../kernels/linear/mxfp8/flashinfer.py | 12 ++++++++---- .../layers/quantization/utils/mxfp8_utils.py | 3 ++- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py b/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py index d26e5579edbb..3150d93f99e1 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py +++ b/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py @@ -11,7 +11,7 @@ ) from vllm.platforms import current_platform from vllm.utils import flashinfer as vllm_flashinfer -from vllm.utils.flashinfer import has_flashinfer_cutedsl +from vllm.utils.flashinfer import has_flashinfer, has_flashinfer_cutedsl from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig @@ -23,9 +23,13 @@ class FlashInferCutlassMxfp8LinearKernel(Mxfp8LinearKernel): def is_supported( cls, compute_capability: int | None = None ) -> tuple[bool, str | None]: - if current_platform.has_device_capability(100): - return True, None - return False, "requires >=sm_100 (Blackwell)" + if not ( + current_platform.is_cuda() and current_platform.has_device_capability(100) + ): + return False, "requires >=sm_100 (Blackwell)" + if not has_flashinfer(): + return False, "requires FlashInfer" + return True, None @classmethod def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]: diff --git a/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py b/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py index 932bf9235289..1e45ea9ab03b 100644 --- a/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py @@ -177,8 +177,9 @@ def _mxfp8_e4m3_quantize_impl( alignment: int = 0, ) -> tuple[torch.Tensor, torch.Tensor]: from vllm.platforms import current_platform + from vllm.utils.flashinfer import has_flashinfer - if current_platform.has_device_capability(100): + if current_platform.has_device_capability(100) and has_flashinfer(): from flashinfer import mxfp8_quantize as flashinfer_mxfp8_quantize x_q, x_scales = flashinfer_mxfp8_quantize( From f4b161d7fca438bfe29509984759be1943a5aa88 Mon Sep 17 00:00:00 2001 From: Qiming Zhang Date: Mon, 17 Aug 2026 18:47:35 -0700 Subject: [PATCH 073/839] [XPU][UT] Fix OOM and skip graph case (#49287) Signed-off-by: mayuyuace Signed-off-by: Qiming Zhang Co-authored-by: Kunshang Ji --- tests/conftest.py | 37 +++++++++++++------ .../models/language/generation/test_hybrid.py | 3 ++ tests/models/language/pooling/test_colbert.py | 4 +- .../generation/test_voxtral_realtime.py | 8 ++-- tests/utils.py | 15 ++++++-- .../nixl_integration/run_accuracy_test.sh | 2 +- 6 files changed, 47 insertions(+), 22 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index e277b35dd46b..bc6d9cb281e3 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -914,13 +914,13 @@ def __enter__(self): return self def __exit__(self, exc_type, exc_value, traceback): - from tests.utils import wait_for_rocm_memory_to_settle + from tests.utils import wait_for_memory_to_settle del self.model cleanup_dist_env_and_memory() # ROCm frees VRAM lazily; wait so a runner started right after this HF # model exits does not OOM on its startup memory guard. - wait_for_rocm_memory_to_settle( + wait_for_memory_to_settle( threshold_ratio=getattr(self, "threshold_ratios", None) ) if hasattr(self, "threshold_ratios"): @@ -998,9 +998,24 @@ def __init__( # V1 startup requires free_memory >= total * gpu_memory_utilization. # ROCm CI can hand a test a device that is still lazily releasing # VRAM from a previous process, so wait before constructing LLM. - from tests.utils import wait_for_rocm_memory_to_settle - - wait_for_rocm_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization) + from tests.utils import wait_for_memory_to_settle + + wait_for_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization) + elif current_platform.is_xpu(): + # The XPU/oneAPI runtime keeps ~1 GiB of context resident in the + # parent pytest process for its whole lifetime (grown by in-process + # HfRunner models), and distributed tests additionally allocate a + # CCL context in the engine subprocess. The default utilization of + # 0.92 leaves too little headroom for both, so lower it on XPU when + # the caller did not request an explicit value. + if "gpu_memory_utilization" not in kwargs: + kwargs["gpu_memory_utilization"] = 0.9 + gpu_memory_utilization = kwargs["gpu_memory_utilization"] + # XPU (Level Zero) can also release device memory lazily after a + # previous engine shuts down, so wait before constructing LLM. + from tests.utils import wait_for_memory_to_settle + + wait_for_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization) with init_ctx: self.llm = LLM( @@ -1326,14 +1341,14 @@ def collective_rpc(self, *args, **kwargs): def __enter__(self): return self - def _wait_for_rocm_memory_release(self, gpu_memory_utilization: float) -> None: - from tests.utils import wait_for_rocm_memory_to_settle + def _wait_for_memory_release(self, gpu_memory_utilization: float) -> None: + from tests.utils import wait_for_memory_to_settle # V1 startup requires free_memory >= total * gpu_memory_utilization. # Wait for the complementary used-memory ratio so the next runner does # not fail the startup guard immediately after this runner exits. The # wait is bounded so cleanup failures fail this test instead of hanging. - wait_for_rocm_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization) + wait_for_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization) def __exit__(self, exc_type, exc_value, traceback): # Explicitly shutdown the engine core to release GPU resources @@ -1359,7 +1374,7 @@ def __exit__(self, exc_type, exc_value, traceback): del self.llm torch._dynamo.reset() cleanup_dist_env_and_memory() - self._wait_for_rocm_memory_release(gpu_memory_utilization) + self._wait_for_memory_release(gpu_memory_utilization) @pytest.fixture(scope="session") @@ -1746,7 +1761,7 @@ def clean_gpu_memory_between_tests(): import gc - from tests.utils import wait_for_gpu_memory_to_clear, wait_for_rocm_memory_to_settle + from tests.utils import wait_for_gpu_memory_to_clear, wait_for_memory_to_settle num_gpus = torch.accelerator.device_count() @@ -1755,7 +1770,7 @@ def _wait_for_settled_gpu_memory() -> None: return try: if current_platform.is_rocm(): - wait_for_rocm_memory_to_settle() + wait_for_memory_to_settle() else: wait_for_gpu_memory_to_clear( devices=list(range(num_gpus)), diff --git a/tests/models/language/generation/test_hybrid.py b/tests/models/language/generation/test_hybrid.py index bb73b1dea825..c5406fae48e0 100644 --- a/tests/models/language/generation/test_hybrid.py +++ b/tests/models/language/generation/test_hybrid.py @@ -9,6 +9,7 @@ from tests.models.registry import HF_EXAMPLE_MODELS from tests.utils import multi_gpu_test from vllm import LLM +from vllm.config import CUDAGraphMode from vllm.engine.arg_utils import EngineArgs from vllm.platforms import current_platform from vllm.sampling_params import SamplingParams @@ -210,6 +211,8 @@ def test_mamba_cache_cg_padding( cudagraph_dispatcher.initialize_cudagraph_keys( vllm_config.compilation_config.cudagraph_mode ) + if cudagraph_dispatcher.cudagraph_mode == CUDAGraphMode.NONE: + pytest.skip("CUDA/XPU graph is disabled.Please enable it to run this test. ") while ( len(example_prompts) == cudagraph_dispatcher.dispatch(len(example_prompts))[1].num_tokens diff --git a/tests/models/language/pooling/test_colbert.py b/tests/models/language/pooling/test_colbert.py index bb7afb23365b..dd4e1dea9772 100644 --- a/tests/models/language/pooling/test_colbert.py +++ b/tests/models/language/pooling/test_colbert.py @@ -11,7 +11,7 @@ import pytest import torch -from tests.utils import wait_for_rocm_memory_to_settle +from tests.utils import wait_for_memory_to_settle from vllm.distributed import cleanup_dist_env_and_memory from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score @@ -165,7 +165,7 @@ def _hf_colbert_model(model_name: str, hf_spec: dict, device: torch.device): finally: del hf_model, linear_weight cleanup_dist_env_and_memory() - wait_for_rocm_memory_to_settle() + wait_for_memory_to_settle() def _assert_embeddings_close(vllm_outputs, hf_embeddings): diff --git a/tests/models/multimodal/generation/test_voxtral_realtime.py b/tests/models/multimodal/generation/test_voxtral_realtime.py index df2d63ad84be..86fe732378be 100644 --- a/tests/models/multimodal/generation/test_voxtral_realtime.py +++ b/tests/models/multimodal/generation/test_voxtral_realtime.py @@ -104,9 +104,9 @@ async def async_engine(): from vllm.platforms import current_platform if current_platform.is_rocm(): - from tests.utils import wait_for_rocm_memory_to_settle + from tests.utils import wait_for_memory_to_settle - wait_for_rocm_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization) + wait_for_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization) engine_args = AsyncEngineArgs(**ENGINE_CONFIG) llm = AsyncLLM.from_engine_args(engine_args) @@ -122,9 +122,9 @@ async def async_engine(): from vllm.distributed import cleanup_dist_env_and_memory cleanup_dist_env_and_memory() - from tests.utils import wait_for_rocm_memory_to_settle + from tests.utils import wait_for_memory_to_settle - wait_for_rocm_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization) + wait_for_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization) def test_voxtral_realtime_forward(audio_assets, tokenizer, vllm_runner, monkeypatch): diff --git a/tests/utils.py b/tests/utils.py index fac455ab738c..86774bf4bf82 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -1540,6 +1540,13 @@ def record_gpu_memory_usage_stats( mem_info = amdsmi_get_gpu_vram_usage(dev_handle) gb_used = mem_info["vram_used"] / 2**10 gb_total = mem_info["vram_total"] / 2**10 + elif current_platform.is_xpu(): + # nvml/amdsmi are unavailable on XPU. Query device memory through + # torch.accelerator.get_memory_info, which the XPU platform patches + # to return (free, total) bytes via Level Zero. + free_b, total_b = torch.accelerator.get_memory_info(device) + gb_used = (total_b - free_b) / 2**30 + gb_total = total_b / 2**30 else: dev_handle = get_nvml_device_handle(device) mem_info = nvmlDeviceGetMemoryInfo(dev_handle) @@ -1680,19 +1687,19 @@ def wait_for_gpu_memory_to_clear( time.sleep(poll_interval_s) -def wait_for_rocm_memory_to_settle( +def wait_for_memory_to_settle( *, threshold_ratio: float | dict[int, float] | None = 0.1, timeout_s: float = 240, ) -> None: - """Block until ROCm device VRAM usage drops below ``threshold_ratio``. + """Block until ROCm or XPU device VRAM usage drops below ``threshold_ratio``. - ROCm reclaims GPU memory more lazily than CUDA, so back-to-back model + ROCm and XPU reclaims GPU memory more lazily than CUDA, so back-to-back model loads in a single test process can OOM the *next* engine/model startup even after ``cleanup_dist_env_and_memory``. This gives the driver time to actually release VRAM before the next allocation. No-op off ROCm. """ - if not current_platform.is_rocm(): + if not current_platform.is_rocm() and not current_platform.is_xpu(): return num_gpus = current_platform.device_count() diff --git a/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh b/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh index fc5c04a1ad07..b8a37e1e96fa 100755 --- a/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh +++ b/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh @@ -112,7 +112,7 @@ wait_for_server() { # Function to clean up previous instances wait_for_gpu_memory_release() { if [[ "$SMI_BIN" == *"rocm"* ]]; then - PYTHONPATH="${GIT_ROOT}" python3 -c "from tests.utils import wait_for_rocm_memory_to_settle; wait_for_rocm_memory_to_settle()" + PYTHONPATH="${GIT_ROOT}" python3 -c "from tests.utils import wait_for_memory_to_settle; wait_for_memory_to_settle()" fi } From d5f5de7a7db4e13ea9f94ce8662b55b6b5e97d02 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=84=8D=F0=9D=95=A0=F0=9D=95=9D=F0=9D=95=9D=F0=9D=95=A0?= =?UTF-8?q?=F0=9D=95=A8=20=F0=9D=95=84=F0=9D=95=92=F0=9D=95=9F?= Date: Tue, 18 Aug 2026 05:55:37 +0300 Subject: [PATCH 074/839] [Bugfix] Fix DeepSeek V4 mHC broadcast buffer for weight sync (#52626) Signed-off-by: Hollow Man --- tests/kernels/test_mhc_kernels.py | 61 +++++++++++++++++++++++++ vllm/models/deepseek_v4/nvidia/model.py | 6 ++- 2 files changed, 66 insertions(+), 1 deletion(-) diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index 2bdce9f9c14d..a2070dd595e3 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/tests/kernels/test_mhc_kernels.py @@ -1,7 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + import pytest import torch +import torch.nn as nn import vllm.model_executor.kernels.mhc # noqa: F401 from vllm.model_executor.kernels.mhc.tilelang import ( @@ -9,6 +12,10 @@ _torch_hc_prenorm_gemm, ) from vllm.model_executor.layers.mhc import HAS_TILELANG_MHC +from vllm.models.deepseek_v4.nvidia.model import ( + DeepseekV4DecoderLayer, + DeepseekV4Model, +) from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed @@ -355,3 +362,57 @@ def test_hc_head_tilelang(num_tokens, hidden_size, hc_mult): out_ref = hc_head_ref(residual, fn, hc_scale, hc_base, rms_eps, hc_eps) torch.testing.assert_close(out, out_ref, atol=5e-2, rtol=1e-2) + + +def _make_mhc_decoder_layer(hc_mult: int, hidden_size: int) -> DeepseekV4DecoderLayer: + layer = DeepseekV4DecoderLayer.__new__(DeepseekV4DecoderLayer) + nn.Module.__init__(layer) + layer.hc_mult = hc_mult + layer.hidden_size = hidden_size + mix_hc = (2 + hc_mult) * hc_mult + layer.hc_attn_fn = nn.Parameter( + torch.randn(mix_hc, hc_mult * hidden_size, dtype=torch.float32), + requires_grad=False, + ) + layer.hc_attn_fn_broadcast = None + return layer + + +def _patch_first_rank_pp_group(monkeypatch): + monkeypatch.setattr( + "vllm.models.deepseek_v4.nvidia.model.get_pp_group", + lambda: SimpleNamespace(is_first_rank=True), + ) + + +def test_deepseek_v4_mhc_broadcast_finalize_sums_hc_streams(monkeypatch): + """First finalize (at the end of load_weights) allocates + hc_attn_fn_broadcast as hc_attn_fn summed over hc streams.""" + _patch_first_rank_pp_group(monkeypatch) + layer = _make_mhc_decoder_layer(hc_mult=2, hidden_size=8) + model = SimpleNamespace(start_layer=0, end_layer=1, layers=[layer]) + + DeepseekV4Model.finalize_mhc_broadcast_weights(model) + + assert layer.hc_attn_fn_broadcast is not None + expected = layer.hc_attn_fn.detach().view(-1, 2, 8).sum(dim=1) + assert torch.equal(layer.hc_attn_fn_broadcast, expected) + + +def test_deepseek_v4_mhc_broadcast_refit_refreshes_in_place(monkeypatch): + """Re-finalizing after a weight refit must copy into the existing + broadcast tensor so its address stays stable for captured CUDA graphs, + while picking up the new hc_attn_fn values.""" + _patch_first_rank_pp_group(monkeypatch) + layer = _make_mhc_decoder_layer(hc_mult=2, hidden_size=8) + model = SimpleNamespace(start_layer=0, end_layer=1, layers=[layer]) + + DeepseekV4Model.finalize_mhc_broadcast_weights(model) + buffer = layer.hc_attn_fn_broadcast + + layer.hc_attn_fn.add_(1.0) + DeepseekV4Model.finalize_mhc_broadcast_weights(model) + + assert layer.hc_attn_fn_broadcast is buffer + expected = layer.hc_attn_fn.detach().view(-1, 2, 8).sum(dim=1) + assert torch.equal(layer.hc_attn_fn_broadcast, expected) diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index ed7ee11ce575..17f29f4a1ae1 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -1372,11 +1372,15 @@ def finalize_mhc_broadcast_weights(self) -> None: return layer = self.layers[self.start_layer] if isinstance(layer, DeepseekV4DecoderLayer): - layer.hc_attn_fn_broadcast = ( + broadcast = ( layer.hc_attn_fn.detach() .view(-1, layer.hc_mult, layer.hidden_size) .sum(dim=1) ) + if layer.hc_attn_fn_broadcast is None: + layer.hc_attn_fn_broadcast = broadcast + else: + layer.hc_attn_fn_broadcast.copy_(broadcast) def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper: From e0e5a7fb2808504ba86c94f7b379e38496002fd0 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Mon, 17 Aug 2026 20:16:54 -0700 Subject: [PATCH 075/839] [Rust Frontend] Fix GLM-5.2 chat template rendering parity (#51426) Co-authored-by: Bugen Zhao Co-authored-by: OpenAI Codex Signed-off-by: Bugen Zhao Signed-off-by: Woosuk Kwon --- rust/Cargo.lock | 8 +-- rust/Cargo.toml | 4 +- rust/src/chat/src/renderer/hf/mod.rs | 66 +++++++++++++++++++++++ rust/src/chat/src/renderer/hf/template.rs | 20 +++++++ rust/src/chat/tests/roundtrip.rs | 14 +++++ 5 files changed, 106 insertions(+), 6 deletions(-) diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 092a79843124..6262bf2e783d 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -2402,9 +2402,9 @@ dependencies = [ [[package]] name = "minijinja" -version = "2.22.0" +version = "2.24.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ef84a52be188a1d4124bd717903fdde96ca4705f2b56adfe2d91fcc57fcb6987" +checksum = "86886cf6dbf4e614b19c9a1eec9775f021869d7eadde0fc73921a81b90c9b4c9" dependencies = [ "indexmap 2.13.0", "memo-map", @@ -2414,9 +2414,9 @@ dependencies = [ [[package]] name = "minijinja-contrib" -version = "2.22.0" +version = "2.24.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fd6e279dc925840c2d9ebc1e9f85410611d01292561297463ba3e516c3ad94dd" +checksum = "bd3e5f077bc2379f0f7d911e7cfdd921114ed99fc884533dca502944cb355b11" dependencies = [ "minijinja", "serde", diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 405d6775bcf4..723e34806f78 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -61,8 +61,8 @@ itertools = "0.14.0" libc = "0.2.177" llm-multimodal = { git = "https://github.com/smg-project/llm-multimodal", rev = "15adba5e025d8636ba4a334fb379b1371f6196a1", default-features = false, features = ["native-tls"] } mimalloc = "0.1.52" -minijinja = { version = "2.22", features = ["unstable_machinery", "json", "builtins", "loader", "loop_controls", "preserve_order"] } -minijinja-contrib = { version = "2.22", features = ["pycompat"] } +minijinja = { version = "2.24", features = ["unstable_machinery", "json", "builtins", "loader", "loop_controls", "preserve_order"] } +minijinja-contrib = { version = "2.24", features = ["pycompat"] } native-tls-vendored = { package = "native-tls", version = "0.2.18", features = ["vendored"] } ndarray = { version = "0.17", features = ["serde"] } openai-harmony = { package = "oss-harmony", git = "https://github.com/oss-harmony/harmony", tag = "v0.0.11", default-features = false } diff --git a/rust/src/chat/src/renderer/hf/mod.rs b/rust/src/chat/src/renderer/hf/mod.rs index df0d94331a8b..7461acf2955a 100644 --- a/rust/src/chat/src/renderer/hf/mod.rs +++ b/rust/src/chat/src/renderer/hf/mod.rs @@ -290,6 +290,7 @@ struct TemplateToolDefinition { name: String, description: Option, parameters: JsonValue, + #[serde(skip_serializing_if = "Option::is_none")] strict: Option, } @@ -1280,6 +1281,71 @@ mod tests { assert_eq!(rendered, "get_weather|city"); } + #[test] + fn chat_template_preserves_openai_tool_field_order() { + let mut request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]); + let tools = vec![ChatTool { + name: "get_weather".to_string(), + description: Some("Get weather".to_string()), + parameters: serde_json::json!({"type": "object"}), + strict: None, + }]; + request.tool_context = crate::request::ResolvedToolContext::new( + &request.messages, + tools, + Some(ChatToolChoice::Auto), + true, + ) + .expect("tool context should resolve"); + + let rendered = render( + Some("{% for key, value in tools[0].function.items() %}{{ key }}|{% endfor %}"), + &request, + ) + .unwrap(); + + assert_eq!(rendered, "name|description|parameters|"); + } + + #[test] + fn chat_template_preserves_python_optional_tool_fields() { + let mut request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]); + let tools = vec![ + ChatTool { + name: "without_strict".to_string(), + description: None, + parameters: Value::Null, + strict: None, + }, + ChatTool { + name: "with_strict".to_string(), + description: Some("description".to_string()), + parameters: serde_json::json!({"type": "object"}), + strict: Some(false), + }, + ]; + request.tool_context = crate::request::ResolvedToolContext::new( + &request.messages, + tools, + Some(ChatToolChoice::Auto), + true, + ) + .expect("tool context should resolve"); + + let rendered = render( + Some( + "{% for tool in tools %}{% for key, value in tool.function.items() %}{{ key }}={{ value|tojson }}|{% endfor %};{% endfor %}", + ), + &request, + ) + .unwrap(); + + assert_eq!( + rendered, + "name=\"without_strict\"|description=null|parameters=null|;name=\"with_strict\"|description=\"description\"|parameters={\"type\": \"object\"}|strict=false|;" + ); + } + #[test] fn chat_template_exposes_assistant_tool_calls_and_tool_messages() { let request = sample_request(vec![ diff --git a/rust/src/chat/src/renderer/hf/template.rs b/rust/src/chat/src/renderer/hf/template.rs index 5a609621cc6e..b7de0172f36f 100644 --- a/rust/src/chat/src/renderer/hf/template.rs +++ b/rust/src/chat/src/renderer/hf/template.rs @@ -152,6 +152,26 @@ mod tests { assert_eq!(result, "[]"); } + #[test] + fn test_midchain_dotted_integer_lookup() { + let template = CompiledChatTemplate::new( + "{{ values.0.name }}".to_string(), + ChatTemplateContentFormatOption::Auto, + ) + .unwrap(); + let mut kwargs = HashMap::new(); + kwargs.insert("values".to_string(), serde_json::json!([{"name": "first"}])); + + let result = template + .apply(TemplateContext { + template_kwargs: Some(&kwargs), + ..Default::default() + }) + .unwrap(); + + assert_eq!(result, "first"); + } + #[test] fn test_chat_template_state_invalid_template() { let result = CompiledChatTemplate::new( diff --git a/rust/src/chat/tests/roundtrip.rs b/rust/src/chat/tests/roundtrip.rs index f8c3519bbfdd..049dc6a40308 100644 --- a/rust/src/chat/tests/roundtrip.rs +++ b/rust/src/chat/tests/roundtrip.rs @@ -185,6 +185,19 @@ impl RoundtripCase { } } + /// GLM-5.2 XML-like argument format with `` reasoning tags. + fn glm52() -> Self { + Self { + model_id: "zai-org/GLM-5.2-FP8", + assistant_stop_suffix: "", + tool_call_parser: ParserSelection::Auto, + reasoning_parser: ParserSelection::Auto, + thinking_behavior: ThinkingBehavior::Toggleable { default: true }, + json_fmt: compact_json_fmt(), + sort_json_keys: false, + } + } + /// Gemma4 channel reasoning with custom function-call arguments. fn gemma4() -> Self { Self { @@ -324,6 +337,7 @@ roundtrip_tests! { deepseek_v32 => [tool_call_mix], glm45 => [reasoning_and_content, tool_call_mix], glm47 => [reasoning_and_content, tool_call_mix], + glm52 => [reasoning_and_content, tool_call_mix], seed_oss => [reasoning_and_content, tool_call_mix], step3p5 => [reasoning_and_content], nemotron_v3 => [reasoning_and_content], From 101c4477dd4538f6fc32161d680ff2fba942b6d0 Mon Sep 17 00:00:00 2001 From: Do_it_now_! Date: Tue, 18 Aug 2026 12:46:08 +0800 Subject: [PATCH 076/839] [Bugfix] Handle DeepseekV4ForCausalLM in benchmark_moe get_model_params (#52044) Co-authored-by: Do_it_now_! <23432123@users.noreply.github.com> --- benchmarks/kernels/benchmark_moe.py | 1 + 1 file changed, 1 insertion(+) diff --git a/benchmarks/kernels/benchmark_moe.py b/benchmarks/kernels/benchmark_moe.py index 9ea22dd1304f..5d86a7599fca 100644 --- a/benchmarks/kernels/benchmark_moe.py +++ b/benchmarks/kernels/benchmark_moe.py @@ -766,6 +766,7 @@ def get_model_params(config): "DeepseekV2ForCausalLM", "DeepseekV3ForCausalLM", "DeepseekV32ForCausalLM", + "DeepseekV4ForCausalLM", "GlmMoeDsaForCausalLM", "Glm4MoeForCausalLM", "Glm4MoeLiteForCausalLM", From aa9903490c616dc6871e5acc62cec7bb1e5e9434 Mon Sep 17 00:00:00 2001 From: Komal Kumar Teru <162363718+kkt-cohere@users.noreply.github.com> Date: Tue, 18 Aug 2026 00:59:57 -0400 Subject: [PATCH 077/839] [Cohere] Misc changes to cohere model definitions (#50156) Co-authored-by: Cursor --- vllm/model_executor/models/cohere2_moe.py | 2 +- vllm/model_executor/models/commandr.py | 111 +++++++++++++++++----- 2 files changed, 87 insertions(+), 26 deletions(-) diff --git a/vllm/model_executor/models/cohere2_moe.py b/vllm/model_executor/models/cohere2_moe.py index 12291fb1423d..3873b6ee004f 100644 --- a/vllm/model_executor/models/cohere2_moe.py +++ b/vllm/model_executor/models/cohere2_moe.py @@ -208,7 +208,7 @@ def __init__( layer_types is not None and layer_types[self.layer_idx] == "sliding_attention" ): - self.sliding_window = config.sliding_window + self.sliding_window = config.sliding_window + 1 # Prefix-dense layers have full attention (no sliding window). When # prefix_dense_sliding_window_pattern == 1, they keep RoPE even though diff --git a/vllm/model_executor/models/commandr.py b/vllm/model_executor/models/commandr.py index 3c755c65e10c..5481fb0abcee 100644 --- a/vllm/model_executor/models/commandr.py +++ b/vllm/model_executor/models/commandr.py @@ -32,7 +32,11 @@ from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig -from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size +from vllm.distributed import ( + get_pp_group, + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_reduce, +) from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.linear import ( @@ -44,9 +48,7 @@ from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding -from vllm.model_executor.model_loader.weight_utils import ( - row_parallel_weight_loader, -) +from vllm.model_executor.model_loader.weight_utils import row_parallel_weight_loader from vllm.model_executor.utils import set_weight_attrs from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors @@ -87,18 +89,61 @@ def forward(self, hidden_states, residuals=None): return hidden_states, residuals +@torch.compile(backend=current_platform.simple_compile_backend) +def rms_norm_func(hidden_states, weight, variance_epsilon): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + variance_epsilon) + + hidden_states = weight.to(torch.float32) * hidden_states + return hidden_states.to(input_dtype) + + +class RMSNorm(nn.Module): + def __init__(self, param_shape=None, eps=1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(param_shape)) + self.variance_epsilon = eps + set_weight_attrs(self.weight, {"weight_loader": row_parallel_weight_loader}) + + def forward(self, hidden_states, residuals=None): + hidden_states = rms_norm_func(hidden_states, self.weight, self.variance_epsilon) + return hidden_states, residuals + + +def select_norm_impl(config: CohereConfig) -> tuple[type[nn.Module], float]: + """ + Returns the normalization layer class and epsilon value to use. + If `config.rms_norm_eps` is present, use RMSNorm. + Otherwise default to LayerNorm with `config.layer_norm_eps`. + """ + rms_eps = getattr(config, "rms_norm_eps", None) + if rms_eps is not None: + return RMSNorm, rms_eps + + return LayerNorm, config.layer_norm_eps + + # Copied from transformers.models.llama.modeling_llama.LlamaMLP Llama->Cohere class CohereMLP(nn.Module): def __init__( self, - config: CohereConfig | Cohere2Config, + config: CohereConfig, + intermediate_size: int | None = None, quant_config: QuantizationConfig | None = None, + # so we can reduce the attention and MLP outputs together + reduce_results: bool = False, prefix: str = "", ): super().__init__() self.config = config self.hidden_size = config.hidden_size - self.intermediate_size = config.intermediate_size + if intermediate_size is None: + self.intermediate_size = config.intermediate_size + else: + self.intermediate_size = intermediate_size self.gate_up_proj = MergedColumnParallelLinear( self.hidden_size, [self.intermediate_size] * 2, @@ -111,6 +156,7 @@ def __init__( self.hidden_size, bias=False, quant_config=quant_config, + reduce_results=reduce_results, prefix=f"{prefix}.down_proj", ) self.act_fn = SiluAndMul() @@ -133,11 +179,15 @@ def __init__( super().__init__() tp_size = get_tensor_model_parallel_world_size() self.config = config + self.layer_idx = extract_layer_index(prefix) self.attention_dropout = config.attention_dropout self.hidden_size = config.hidden_size self.total_num_heads = config.num_attention_heads self.num_heads = self.total_num_heads // tp_size - self.head_dim = self.hidden_size // self.total_num_heads + if hasattr(config, "head_dim"): + self.head_dim = config.head_dim + else: + self.head_dim = self.hidden_size // self.total_num_heads self.total_num_kv_heads = config.num_key_value_heads if self.total_num_kv_heads >= tp_size: # Number of KV heads is greater than TP size, so we partition @@ -171,6 +221,9 @@ def __init__( bias=False, quant_config=quant_config, prefix=f"{prefix}.o_proj", + # NOTE: reduction will happen in the decoder layer forward + # so we can combine the attention and MLP outputs together + reduce_results=False, ) self.rotary_emb = get_rope( self.head_dim, @@ -182,11 +235,14 @@ def __init__( # Model v2 has interleaved sliding windows, v1 does not self.v1 = isinstance(config, CohereConfig) + # cohere swa layer sees [pos - sliding_window, pos], i.e. sliding_window + 1 + # tokens. vLLM's FlashAttention backend does (value - 1, 0), so we pass + # sliding_window + 1 here to match the training convention. The same +1 + # propagates into the KV-cache eviction formula + # (single_type_kv_cache_manager.py), keeping both paths consistent. self.sliding_window = None - if not self.v1: - layer_idx = extract_layer_index(prefix) - if config.layer_types[layer_idx] == "sliding_attention": - self.sliding_window = config.sliding_window + if not self.v1 and config.layer_types[self.layer_idx] == "sliding_attention": + self.sliding_window = config.sliding_window + 1 self.attn = Attention( self.num_heads, @@ -198,13 +254,13 @@ def __init__( per_layer_sliding_window=self.sliding_window, prefix=f"{prefix}.attn", ) + norm_cls, norm_eps = select_norm_impl(config) if self.use_qk_norm: - self.q_norm = LayerNorm( - param_shape=(self.num_heads, self.head_dim), eps=config.layer_norm_eps + self.q_norm = norm_cls( + param_shape=(self.num_heads, self.head_dim), eps=norm_eps ) - self.k_norm = LayerNorm( - param_shape=(self.num_kv_heads, self.head_dim), - eps=config.layer_norm_eps, + self.k_norm = norm_cls( + param_shape=(self.num_kv_heads, self.head_dim), eps=norm_eps ) def _apply_qk_norm(self, q, k): @@ -242,18 +298,16 @@ def __init__( ): super().__init__() self.hidden_size = config.hidden_size - + self.tp_size = get_tensor_model_parallel_world_size() self.self_attn = CohereAttention( config, cache_config, quant_config=quant_config, prefix=f"{prefix}.self_attn", ) - self.mlp = CohereMLP(config, quant_config=quant_config, prefix=f"{prefix}.mlp") - self.input_layernorm = LayerNorm( - param_shape=(config.hidden_size), eps=config.layer_norm_eps - ) + norm_cls, norm_eps = select_norm_impl(config) + self.input_layernorm = norm_cls(param_shape=(config.hidden_size), eps=norm_eps) def forward( self, @@ -269,8 +323,14 @@ def forward( hidden_states=hidden_states, ) hidden_states_mlp = self.mlp(hidden_states) + parallel_block_output = hidden_states_attention + hidden_states_mlp # Add everything together - hidden_states = residual + hidden_states_attention + hidden_states_mlp + if self.tp_size > 1: + # do the reduction in 1 shot instead of 2 separate all reduce + parallel_block_output = tensor_model_parallel_all_reduce( + parallel_block_output + ) + hidden_states = residual + parallel_block_output return hidden_states, residual @@ -299,9 +359,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ), prefix=f"{prefix}.layers", ) - self.norm = LayerNorm( - param_shape=(config.hidden_size), eps=config.layer_norm_eps - ) + + norm_cls, norm_eps = select_norm_impl(config) + self.norm = norm_cls(param_shape=(config.hidden_size), eps=norm_eps) + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( ["hidden_states", "residual"], config.hidden_size ) From b0e9cff5e7f7dc4180ff0cd9380f9dcc2933c2f7 Mon Sep 17 00:00:00 2001 From: Yan Ma Date: Tue, 18 Aug 2026 14:27:41 +0800 Subject: [PATCH 078/839] [XPU] update xpu-manager to v2.1.0 (#52569) Signed-off-by: Yan Ma Co-authored-by: Kunshang Ji --- docker/Dockerfile.xpu | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/docker/Dockerfile.xpu b/docker/Dockerfile.xpu index 31efca624151..ea7159f8c453 100644 --- a/docker/Dockerfile.xpu +++ b/docker/Dockerfile.xpu @@ -87,11 +87,8 @@ RUN apt-get update -y && \ apt-get install -y --no-install-recommends libhwloc15 libpciaccess0 libigsc1 intel-metrics-library && \ mkdir xpu-m && \ cd xpu-m && \ - wget https://github.com/intel/xpumanager/releases/download/v2.0.0/libxpum2_2.0.0-238.24.04_amd64.deb && \ - wget https://github.com/intel/xpumanager/releases/download/v2.0.0/xpu-smi_2.0.0-238.24.04_amd64.deb && \ - dpkg-deb -x libxpum2_2.0.0-238.24.04_amd64.deb / && \ - dpkg-deb -x xpu-smi_2.0.0-238.24.04_amd64.deb / && \ - ldconfig && \ + wget https://github.com/intel/xpumanager/releases/download/v2.1.0/xpu-smi_2.1.0+26.33.6468cec-1.24.04_amd64.deb && \ + apt install -y ./xpu-smi_2.1.0+26.33.6468cec-1.24.04_amd64.deb --install-suggests && \ cd .. && \ rm -rf xpu-m && \ add-apt-repository --remove -y ppa:kobuk-team/intel-graphics && \ From c89d692cb5fa430ddca06a9e16682977a18b55dc Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Tue, 18 Aug 2026 16:37:37 +1000 Subject: [PATCH 079/839] [Build] Propagate vLLM version to Rust binaries (#52593) Signed-off-by: Bugen Zhao Co-authored-by: Nick Hill Co-authored-by: Andreas Karatzas --- .buildkite/scripts/ci-bake-rocm.sh | 8 ++--- .buildkite/test-amd.yaml | 2 +- .buildkite/test_areas/rust_frontend.yaml | 2 +- docker/Dockerfile | 5 +++ docker/Dockerfile.cpu | 5 +++ docker/Dockerfile.rocm | 35 ++++++++++++------- docker/Dockerfile.rocm_gfx1250 | 35 ++++++++++++------- docker/Dockerfile.xpu | 5 +++ requirements/build/rust.txt | 1 + rust/Cargo.lock | 7 ++++ rust/Cargo.toml | 2 ++ rust/src/bench/Cargo.toml | 1 + rust/src/bench/src/main.rs | 2 +- rust/src/build-info/Cargo.toml | 8 +++++ rust/src/build-info/src/lib.rs | 10 ++++++ rust/src/cmd/Cargo.toml | 1 + rust/src/cmd/src/cli.rs | 3 +- rust/src/server/Cargo.toml | 1 + rust/src/server/src/routes/tests.rs | 2 +- rust/src/server/src/routes/version.rs | 2 +- setup.py | 7 +++- tests/entrypoints/openai/test_uds.py | 7 ++-- .../serve/instrumentator/test_basic.py | 7 ++-- tools/build_rust.py | 16 +++++++++ 24 files changed, 135 insertions(+), 39 deletions(-) create mode 100644 rust/src/build-info/Cargo.toml create mode 100644 rust/src/build-info/src/lib.rs diff --git a/.buildkite/scripts/ci-bake-rocm.sh b/.buildkite/scripts/ci-bake-rocm.sh index 1e143113141f..e3681e214a42 100644 --- a/.buildkite/scripts/ci-bake-rocm.sh +++ b/.buildkite/scripts/ci-bake-rocm.sh @@ -23,8 +23,8 @@ DEFAULT_CI_BASE_METADATA_VERSION="3" # local-source stages rather than unreachable remote-fetch alternatives. DEFAULT_ROCM_CSRC_CONTENT_FILES=".dockerignore requirements/common.txt requirements/rocm.txt pyproject.toml setup.py CMakeLists.txt cmake csrc vllm/envs.py vllm/__init__.py tools/build_rust.py" DEFAULT_ROCM_CSRC_DOCKERFILE_STAGES="base fetch_vllm_0 fetch_vllm build_vllm_dependencies rocm-triton-kernels csrc-build" -DEFAULT_ROCM_RUST_CONTENT_FILES=".dockerignore requirements/build/rust.txt rust/Cargo.lock rust/Cargo.toml rust/proto rust/src rust-toolchain.toml tools/build_rust.py tools/install_protoc.sh build_rust.sh" -DEFAULT_ROCM_RUST_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust-toolchain-input rust_input_0 rust-input rust-toolchain rust-build" +DEFAULT_ROCM_RUST_CONTENT_FILES=".dockerignore .git_archival.txt pyproject.toml requirements/build/rust.txt rust/Cargo.lock rust/Cargo.toml rust/proto rust/src rust-toolchain.toml tools/build_rust.py tools/install_protoc.sh build_rust.sh" +DEFAULT_ROCM_RUST_DOCKERFILE_STAGES="base fetch_vllm_0 fetch_vllm vllm-version rust_toolchain_input_0 rust-toolchain-input rust_input_0 rust-input rust-toolchain rust-build" # Docker's 128-character tag limit minus the longest cache prefix # ("csrc-rocm-branch-" and "rust-rocm-branch-", both 17 characters). ROCM_CACHE_BRANCH_TAG_MAX_LEN=111 @@ -523,8 +523,8 @@ prepare_ci_build_context() { return 1 fi - # setuptools-scm understands Git's stable archive format, so full-source - # wheels retain their exact version without copying mutable Git history. + # setuptools-scm understands Git's stable archive format, so wheels and + # Rust artifacts retain their exact version without copying Git history. if ! is_ci_base_target; then write_ci_git_archival_metadata "${source_root}" "${context_root}" \ || return $? diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index e9726f73ff5c..41dbd629bc9b 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -2356,7 +2356,7 @@ steps: - export VLLM_USE_RUST_FRONTEND=1 - export VLLM_WORKER_MULTIPROC_METHOD=spawn - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py - - pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not show_version and not server_load" + - pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not server_load" - pytest -v -s entrypoints/scale_out/token_in_token_out/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow" - pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist" - pytest -v -s entrypoints/serve/tokenize/test_tokenization.py -k "not tokenizer_info" diff --git a/.buildkite/test_areas/rust_frontend.yaml b/.buildkite/test_areas/rust_frontend.yaml index e9e66c46c0c0..6486df1c00a1 100644 --- a/.buildkite/test_areas/rust_frontend.yaml +++ b/.buildkite/test_areas/rust_frontend.yaml @@ -58,7 +58,7 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py # server_load can be flaky under the Rust frontend; keep it excluded for now. - - pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not show_version and not server_load" + - pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not server_load" # test_generate_logprobs expects Python-style top_logprobs truncation (dedup sampled + cap at max(k, 1)). - pytest -v -s entrypoints/scale_out/token_in_token_out/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow" - pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist" diff --git a/docker/Dockerfile b/docker/Dockerfile index 804a3a4bb3ef..b2c949d2d155 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -279,11 +279,16 @@ ENV CARGO_BUILD_JOBS=4 # separate local sccache daemon while sharing the same remote cache backend. ENV SCCACHE_SERVER_PORT=4227 +# Only Rust inputs are present, so suppress the artificial dirty state caused +# by other tracked files being absent from this build stage. +ENV SETUPTOOLS_SCM_PRETEND_METADATA={dirty=false} + # Build the release artifacts. Cache cargo registry/git, but not target/, # because stale target metadata can outlive source updates across BuildKit # cache reuse. RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \ --mount=type=cache,target=/root/.cargo/git,sharing=locked \ + --mount=type=bind,source=.git,target=.git \ --mount=type=secret,id=aws-credentials,target=/root/.aws/credentials,required=false \ if [ "$USE_SCCACHE" = "1" ]; then \ if [ -n "${SCCACHE_ENDPOINT}" ]; then export SCCACHE_ENDPOINT="${SCCACHE_ENDPOINT}"; fi; \ diff --git a/docker/Dockerfile.cpu b/docker/Dockerfile.cpu index 69abab3b774f..43fab4bccb16 100644 --- a/docker/Dockerfile.cpu +++ b/docker/Dockerfile.cpu @@ -166,11 +166,16 @@ COPY build_rust.sh build_rust.sh # (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise). ENV CARGO_BUILD_JOBS=4 +# Only Rust inputs are present, so suppress the artificial dirty state caused +# by other tracked files being absent from this build stage. +ENV SETUPTOOLS_SCM_PRETEND_METADATA={dirty=false} + # Build the release artifacts. Cache cargo registry/git, but not target/, # because stale target metadata can outlive source updates across BuildKit # cache reuse. RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \ --mount=type=cache,target=/root/.cargo/git,sharing=locked \ + --mount=type=bind,source=.git,target=.git \ --mount=type=secret,id=aws-credentials,target=/root/.aws/credentials,required=false \ if [ "$USE_SCCACHE" = "1" ]; then \ export RUSTC_WRAPPER=sccache; \ diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 0487c0c9165a..0b02b9ec2153 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -125,6 +125,18 @@ ONBUILD RUN git clone ${VLLM_REPO} \ && git fetch upstream ; fi FROM fetch_vllm_${REMOTE_VLLM} AS fetch_vllm +# Resolve the device-independent source version once for every downstream +# artifact. setuptools-scm reads Git metadata for local/remote sources and the +# archive metadata generated for canonical CI contexts. +FROM fetch_vllm AS vllm-version +ARG COMMON_WORKDIR +RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ + cd ${COMMON_WORKDIR}/vllm \ + && uv pip install --system -r requirements/build/rust.txt \ + && python3 -c \ + "from setuptools_scm import get_version; print(get_version())" \ + > /vllm-version.txt + # ----------------------- # Rust/protoc toolchain inputs # @@ -238,6 +250,7 @@ COPY --from=rust-input /rust-input/vllm/rust rust COPY --from=rust-input /rust-input/vllm/rust-toolchain.toml rust-toolchain.toml COPY --from=rust-input /rust-input/vllm/tools/build_rust.py tools/build_rust.py COPY --from=rust-input /rust-input/vllm/build_rust.sh build_rust.sh +COPY --from=vllm-version /vllm-version.txt /tmp/vllm-version.txt # Build the release binary. Cargo's registry/git caches can be written by # concurrent BuildKit jobs on shared workers, so lock those cache mounts while @@ -249,6 +262,7 @@ RUN --mount=type=cache,id=vllm-rocm-cargo-registry,target=/root/.cargo/registry, export RUSTC_WRAPPER=sccache \ && sccache --show-stats; \ fi \ + && export SETUPTOOLS_SCM_PRETEND_VERSION="$(cat /tmp/vllm-version.txt)" \ && bash build_rust.sh \ && test -x vllm/vllm-rs \ && if [ "$USE_SCCACHE" = "1" ]; then \ @@ -334,6 +348,7 @@ RUN --mount=type=bind,source=pyproject.toml,target=${COMMON_WORKDIR}/vllm/pyproj FROM build_vllm_dependencies AS build_vllm COPY --from=fetch_vllm ${COMMON_WORKDIR}/vllm ${COMMON_WORKDIR}/vllm COPY --from=csrc-build ${COMMON_WORKDIR}/vllm/dist /precompiled-wheels +COPY --from=vllm-version /vllm-version.txt /tmp/vllm-version.txt # Drop the pre-built Rust artifacts into the source tree. setup.py detects # them and ships them as-is, skipping the local Rust build. @@ -341,6 +356,7 @@ COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vll COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/ RUN cd "${COMMON_WORKDIR}/vllm" \ + && export SETUPTOOLS_SCM_PRETEND_VERSION="$(cat /tmp/vllm-version.txt)" \ && export VLLM_USE_PRECOMPILED=1 \ && export VLLM_PRECOMPILED_WHEEL_LOCATION="$(ls /precompiled-wheels/*.whl)" \ && export VLLM_DOCKER_BUILD_CONTEXT=1 \ @@ -640,6 +656,7 @@ FROM fetch_vllm AS build_vllm_wheel_release ARG COMMON_WORKDIR ARG ROCM_TRITON_KERNELS_COMMIT +COPY --from=vllm-version /vllm-version.txt /tmp/vllm-version.txt COPY --from=rocm-triton-kernels \ /opt/rocm-triton-kernels /opt/rocm-triton-kernels ENV TRITON_KERNELS_SRC_DIR=/opt/rocm-triton-kernels @@ -678,17 +695,12 @@ RUN if [ "$GIT_REPO_CHECK" != "0" ]; then \ cd vllm && bash tools/check_repo.sh; \ fi -# Extract version from git BEFORE any modifications (pin_rocm_dependencies.py modifies requirements/rocm.txt) -# This ensures setuptools_scm sees clean repo state for version detection -RUN --mount=type=bind,source=.git,target=vllm/.git \ - --mount=type=cache,target=/root/.cache/uv \ +# Install the packaging helpers before modifying requirements/rocm.txt. +RUN --mount=type=cache,target=/root/.cache/uv \ grep -Fq "set(TRITON_KERNELS_TAG \"${ROCM_TRITON_KERNELS_COMMIT}\")" \ vllm/cmake/external_projects/triton_kernels.cmake \ && cd vllm \ - && uv pip install --system setuptools_scm regex \ - && VLLM_VERSION=$(python3 -c "import setuptools_scm; print(setuptools_scm.get_version())") \ - && echo "Detected vLLM version: ${VLLM_VERSION}" \ - && echo "${VLLM_VERSION}" > /tmp/vllm_version.txt + && uv pip install --system setuptools_scm regex # Fail if git-based package dependencies are found in requirements files # (uv doesn't handle git+ URLs well, and packages should be distributed on PyPI) @@ -726,13 +738,12 @@ RUN --mount=type=cache,target=/root/.cache/uv \ && echo "Building vLLM with custom wheels from /install" \ && uv pip install --system --find-links /install -r requirements/rocm.txt -# Build wheel using pre-extracted version to avoid dirty state from modified requirements/rocm.txt +# Build wheel using the shared source version to avoid dirty state from modified requirements/rocm.txt # (setup.py auto-detects ccache/sccache in PATH) -RUN --mount=type=bind,source=.git,target=vllm/.git \ - --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ +RUN --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ cd vllm \ && export CCACHE_BASEDIR="$PWD" \ - && export SETUPTOOLS_SCM_PRETEND_VERSION=$(cat /tmp/vllm_version.txt) \ + && export SETUPTOOLS_SCM_PRETEND_VERSION="$(cat /tmp/vllm-version.txt)" \ && echo "Building wheel with version: ${SETUPTOOLS_SCM_PRETEND_VERSION}" \ && MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py bdist_wheel --dist-dir=dist diff --git a/docker/Dockerfile.rocm_gfx1250 b/docker/Dockerfile.rocm_gfx1250 index 4c0008f6ecb7..6e88eb9fed9e 100644 --- a/docker/Dockerfile.rocm_gfx1250 +++ b/docker/Dockerfile.rocm_gfx1250 @@ -127,6 +127,17 @@ ONBUILD RUN git clone ${VLLM_REPO} \ && git fetch upstream ; fi FROM fetch_vllm_${REMOTE_VLLM} AS fetch_vllm +# Resolve the device-independent source version once for every downstream +# artifact. setuptools-scm reads Git metadata from local and remote sources. +FROM fetch_vllm AS vllm-version +ARG COMMON_WORKDIR +RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ + cd ${COMMON_WORKDIR}/vllm \ + && uv pip install --system -r requirements/build/rust.txt \ + && python3 -c \ + "from setuptools_scm import get_version; print(get_version())" \ + > /vllm-version.txt + # ----------------------- # Rust build stage # Builds the `vllm-rs` frontend in a dedicated stage so the wheel build stages @@ -152,6 +163,8 @@ RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ cd ${COMMON_WORKDIR}/vllm \ && uv pip install --system -r requirements/build/rust.txt +COPY --from=vllm-version /vllm-version.txt /tmp/vllm-version.txt + # Build the release binary. Cargo's registry/git caches can be written by # concurrent BuildKit jobs on shared workers, so lock those cache mounts while # keeping the cache benefit. Do not cache target/, because stale target metadata @@ -159,6 +172,7 @@ RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ RUN --mount=type=cache,id=vllm-rocm-cargo-registry,target=/root/.cargo/registry,sharing=locked \ --mount=type=cache,id=vllm-rocm-cargo-git,target=/root/.cargo/git,sharing=locked \ cd ${COMMON_WORKDIR}/vllm \ + && export SETUPTOOLS_SCM_PRETEND_VERSION="$(cat /tmp/vllm-version.txt)" \ && bash build_rust.sh \ && test -x vllm/vllm-rs @@ -211,6 +225,7 @@ ARG COMMON_WORKDIR ENV VLLM_TARGET_DEVICE=rocm COPY --from=csrc-build ${COMMON_WORKDIR}/vllm/dist /precompiled-wheels +COPY --from=vllm-version /vllm-version.txt /tmp/vllm-version.txt # Drop the pre-built Rust artifacts into the source tree. setup.py detects # them and ships them as-is, skipping the local Rust build. @@ -220,6 +235,7 @@ COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/ RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ cd vllm \ && uv pip install --system -r requirements/rocm.txt \ + && export SETUPTOOLS_SCM_PRETEND_VERSION="$(cat /tmp/vllm-version.txt)" \ && export VLLM_USE_PRECOMPILED=1 \ && export VLLM_PRECOMPILED_WHEEL_LOCATION="$(ls /precompiled-wheels/*.whl)" \ && export VLLM_DOCKER_BUILD_CONTEXT=1 \ @@ -426,6 +442,7 @@ RUN /bin/bash -lc 'set -euo pipefail; \ FROM fetch_vllm AS build_vllm_wheel_release ARG COMMON_WORKDIR +COPY --from=vllm-version /vllm-version.txt /tmp/vllm-version.txt # Drop the pre-built Rust artifacts into the source tree. setup.py detects # them and ships them as-is, skipping the local Rust build. @@ -461,15 +478,10 @@ RUN if [ "$GIT_REPO_CHECK" != "0" ]; then \ cd vllm && bash tools/check_repo.sh; \ fi -# Extract version from git BEFORE any modifications (pin_rocm_dependencies.py modifies requirements/rocm.txt) -# This ensures setuptools_scm sees clean repo state for version detection -RUN --mount=type=bind,source=.git,target=vllm/.git \ - --mount=type=cache,target=/root/.cache/uv \ +# Install the packaging helpers before modifying requirements/rocm.txt. +RUN --mount=type=cache,target=/root/.cache/uv \ cd vllm \ - && uv pip install --system setuptools_scm regex \ - && VLLM_VERSION=$(python3 -c "import setuptools_scm; print(setuptools_scm.get_version())") \ - && echo "Detected vLLM version: ${VLLM_VERSION}" \ - && echo "${VLLM_VERSION}" > /tmp/vllm_version.txt + && uv pip install --system setuptools_scm regex # Fail if git-based package dependencies are found in requirements files # (uv doesn't handle git+ URLs well, and packages should be distributed on PyPI) @@ -507,13 +519,12 @@ RUN --mount=type=cache,target=/root/.cache/uv \ && echo "Building vLLM with custom wheels from /install" \ && uv pip install --system --find-links /install -r requirements/rocm.txt -# Build wheel using pre-extracted version to avoid dirty state from modified requirements/rocm.txt +# Build wheel using the shared source version to avoid dirty state from modified requirements/rocm.txt # (setup.py auto-detects ccache/sccache in PATH) -RUN --mount=type=bind,source=.git,target=vllm/.git \ - --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ +RUN --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ cd vllm \ && export CCACHE_BASEDIR="$PWD" \ - && export SETUPTOOLS_SCM_PRETEND_VERSION=$(cat /tmp/vllm_version.txt) \ + && export SETUPTOOLS_SCM_PRETEND_VERSION="$(cat /tmp/vllm-version.txt)" \ && echo "Building wheel with version: ${SETUPTOOLS_SCM_PRETEND_VERSION}" \ && MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py bdist_wheel --dist-dir=dist diff --git a/docker/Dockerfile.xpu b/docker/Dockerfile.xpu index ea7159f8c453..456be41f6f2a 100644 --- a/docker/Dockerfile.xpu +++ b/docker/Dockerfile.xpu @@ -28,8 +28,13 @@ COPY build_rust.sh build_rust.sh # (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise). ENV CARGO_BUILD_JOBS=4 +# Only Rust inputs are present, so suppress the artificial dirty state caused +# by other tracked files being absent from this build stage. +ENV SETUPTOOLS_SCM_PRETEND_METADATA={dirty=false} + RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \ --mount=type=cache,target=/root/.cargo/git,sharing=locked \ + --mount=type=bind,source=.git,target=.git \ bash build_rust.sh FROM ubuntu:24.04 AS vllm-base diff --git a/requirements/build/rust.txt b/requirements/build/rust.txt index e2874dee0aba..f507deb8037b 100644 --- a/requirements/build/rust.txt +++ b/requirements/build/rust.txt @@ -1,4 +1,5 @@ # Dependencies for building Rust artifacts through setuptools-rust. setuptools>=77.0.3,<81.0.0 +setuptools-scm>=9.2.0 setuptools-rust>=1.9.0 wheel diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 6262bf2e783d..a8608e34f2fd 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -5505,9 +5505,14 @@ dependencies = [ "tracing", "url", "uuid", + "vllm-build-info", "vllm-tracing", ] +[[package]] +name = "vllm-build-info" +version = "0.1.0" + [[package]] name = "vllm-chat" version = "0.1.0" @@ -5574,6 +5579,7 @@ dependencies = [ "tracing", "uuid", "vllm-bench", + "vllm-build-info", "vllm-chat", "vllm-engine-core-client", "vllm-managed-engine", @@ -5755,6 +5761,7 @@ dependencies = [ "url", "uuid", "validator", + "vllm-build-info", "vllm-chat", "vllm-engine-core-client", "vllm-llm", diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 723e34806f78..5c4b47d5817e 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -1,6 +1,7 @@ [workspace] members = [ "src/bench", + "src/build-info", "src/chat", "src/cmd", "src/engine-core-client", @@ -135,6 +136,7 @@ url = "2.5.7" uuid = { version = "1.22.0", features = ["v4"] } validator = { version = "0.20.0", features = ["derive"] } vllm-bench = { path = "src/bench" } +vllm-build-info = { path = "src/build-info" } vllm-chat = { path = "src/chat" } vllm-engine-core-client = { path = "src/engine-core-client" } vllm-llm = { path = "src/llm" } diff --git a/rust/src/bench/Cargo.toml b/rust/src/bench/Cargo.toml index 960e7a62f7dd..24a8f12e910c 100644 --- a/rust/src/bench/Cargo.toml +++ b/rust/src/bench/Cargo.toml @@ -34,6 +34,7 @@ tokio-stream.workspace = true tracing.workspace = true url.workspace = true uuid.workspace = true +vllm-build-info.workspace = true vllm-tracing.workspace = true [lints] diff --git a/rust/src/bench/src/main.rs b/rust/src/bench/src/main.rs index 1764983e4487..3485e22bb1e5 100644 --- a/rust/src/bench/src/main.rs +++ b/rust/src/bench/src/main.rs @@ -12,7 +12,7 @@ use clap::Parser; #[command( name = "vllm-bench", about = "Benchmark online serving throughput", - version + version = vllm_build_info::VERSION )] struct Cli { #[command(flatten)] diff --git a/rust/src/build-info/Cargo.toml b/rust/src/build-info/Cargo.toml new file mode 100644 index 000000000000..41d3092af0f2 --- /dev/null +++ b/rust/src/build-info/Cargo.toml @@ -0,0 +1,8 @@ +[package] +name = "vllm-build-info" +version.workspace = true +edition.workspace = true +license.workspace = true + +[lints] +workspace = true diff --git a/rust/src/build-info/src/lib.rs b/rust/src/build-info/src/lib.rs new file mode 100644 index 000000000000..962fbe2f7d92 --- /dev/null +++ b/rust/src/build-info/src/lib.rs @@ -0,0 +1,10 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +/// The vLLM package version supplied by the build system. +/// +/// Direct Cargo builds fall back to the internal crate version. +pub const VERSION: &str = match option_env!("VLLM_RS_BUILD_VERSION") { + Some(version) => version, + None => env!("CARGO_PKG_VERSION"), +}; diff --git a/rust/src/cmd/Cargo.toml b/rust/src/cmd/Cargo.toml index a326a0f9992a..bf92e05ee186 100644 --- a/rust/src/cmd/Cargo.toml +++ b/rust/src/cmd/Cargo.toml @@ -28,6 +28,7 @@ tokio-util.workspace = true tracing.workspace = true uuid.workspace = true vllm-bench.workspace = true +vllm-build-info.workspace = true vllm-chat.workspace = true vllm-engine-core-client.workspace = true vllm-managed-engine.workspace = true diff --git a/rust/src/cmd/src/cli.rs b/rust/src/cmd/src/cli.rs index 833b93ee43e3..e851fad54996 100644 --- a/rust/src/cmd/src/cli.rs +++ b/rust/src/cmd/src/cli.rs @@ -39,7 +39,8 @@ use crate::cli::unsupported::UnsupportedArgs; #[derive(Debug, Parser)] #[command( name = "vllm-rs", - about = "Rust frontend and managed-engine CLI for vLLM." + about = "Rust frontend and managed-engine CLI for vLLM.", + version = vllm_build_info::VERSION )] pub struct Cli { #[command(subcommand)] diff --git a/rust/src/server/Cargo.toml b/rust/src/server/Cargo.toml index 8b8f57a0dac5..f6e6f08b530f 100644 --- a/rust/src/server/Cargo.toml +++ b/rust/src/server/Cargo.toml @@ -45,6 +45,7 @@ tracing-subscriber.workspace = true url.workspace = true uuid.workspace = true validator.workspace = true +vllm-build-info.workspace = true vllm-chat.workspace = true vllm-engine-core-client.workspace = true vllm-llm.workspace = true diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index ac65ae269afd..260851b4efe5 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -1746,7 +1746,7 @@ async fn version_returns_engine_vllm_version() { json, json!({ "version": "test-vllm-version", - "rust_frontend_version": env!("CARGO_PKG_VERSION"), + "rust_frontend_version": vllm_build_info::VERSION, }) ); } diff --git a/rust/src/server/src/routes/version.rs b/rust/src/server/src/routes/version.rs index e4f6f3f1ea5e..2398224c5c98 100644 --- a/rust/src/server/src/routes/version.rs +++ b/rust/src/server/src/routes/version.rs @@ -21,6 +21,6 @@ pub async fn version(State(state): State>) -> Json None: ): cmdclass["build_rust"] = precompiled_build_rust +# Resolve the Python version first because get_vllm_version() may set +# SETUPTOOLS_SCM_PRETEND_VERSION, which the Rust version should inherit. +vllm_version = get_vllm_version() +rust_build.prepare_build_environment() + # Rust artifacts, built via setuptools-rust and installed into the package # directory alongside the Python modules. rust_extensions = rust_build.rust_extensions( @@ -1501,7 +1506,7 @@ def add_vllm_package_data(filename: str) -> None: setup( # static metadata should rather go in pyproject.toml - version=get_vllm_version(), + version=vllm_version, ext_modules=ext_modules, rust_extensions=rust_extensions, install_requires=get_requirements(), diff --git a/tests/entrypoints/openai/test_uds.py b/tests/entrypoints/openai/test_uds.py index f79e40ee4132..7c8557141087 100644 --- a/tests/entrypoints/openai/test_uds.py +++ b/tests/entrypoints/openai/test_uds.py @@ -6,6 +6,7 @@ import httpx import pytest +from vllm import envs from vllm.version import __version__ as VLLM_VERSION from ...utils import RemoteOpenAIServer @@ -40,5 +41,7 @@ async def test_show_version(server: RemoteOpenAIServer): response = client.get(server.url_for("version")) response.raise_for_status() - # Tolerate additive fields (e.g. the Rust frontend reports its own version). - assert response.json()["version"] == VLLM_VERSION + payload = response.json() + assert payload["version"] == VLLM_VERSION + if envs.VLLM_USE_RUST_FRONTEND: + assert payload["rust_frontend_version"] == VLLM_VERSION diff --git a/tests/entrypoints/serve/instrumentator/test_basic.py b/tests/entrypoints/serve/instrumentator/test_basic.py index 5b00d2e578e5..73a97c4fa264 100644 --- a/tests/entrypoints/serve/instrumentator/test_basic.py +++ b/tests/entrypoints/serve/instrumentator/test_basic.py @@ -12,6 +12,7 @@ from fastapi import Request from tests.utils import RemoteOpenAIServer +from vllm import envs from vllm.v1.engine.exceptions import EngineDeadError from vllm.version import __version__ as VLLM_VERSION @@ -83,8 +84,10 @@ async def test_show_version(server: RemoteOpenAIServer): response = requests.get(server.url_for("version")) response.raise_for_status() - # Tolerate additive fields (e.g. the Rust frontend reports its own version). - assert response.json()["version"] == VLLM_VERSION + payload = response.json() + assert payload["version"] == VLLM_VERSION + if envs.VLLM_USE_RUST_FRONTEND: + assert payload["rust_frontend_version"] == VLLM_VERSION @pytest.mark.asyncio diff --git a/tools/build_rust.py b/tools/build_rust.py index b5951bfe576b..a000ec8169fd 100644 --- a/tools/build_rust.py +++ b/tools/build_rust.py @@ -11,8 +11,23 @@ from setuptools import setup from setuptools_rust import Binding, RustExtension +from setuptools_scm import get_version ROOT_DIR = Path(__file__).resolve().parents[1] +VLLM_RS_BUILD_VERSION = "VLLM_RS_BUILD_VERSION" + + +def prepare_build_environment() -> str | None: + """Set the device-independent vLLM source version for Rust artifacts.""" + version = os.getenv(VLLM_RS_BUILD_VERSION) or None + if version is None: + try: + version = get_version(root=ROOT_DIR) + except LookupError: + return None + + os.environ[VLLM_RS_BUILD_VERSION] = version + return version def rust_extensions(*, optional: bool = False) -> list[RustExtension]: @@ -61,6 +76,7 @@ def build_binary(build_rust_args: list[str]) -> None: def main() -> None: + prepare_build_environment() build_binary(sys.argv[1:]) From d785eb51ce24210ec6da82140729d30822332032 Mon Sep 17 00:00:00 2001 From: flb_ Date: Tue, 18 Aug 2026 14:46:44 +0800 Subject: [PATCH 080/839] [Test] Add pause/resume E2E tests (#52144) Signed-off-by: flb_ --- tests/entrypoints/serve/dev/rlhf/conftest.py | 414 ++++++++++++++++++ .../dev/rlhf/state_transitions/__init__.py | 0 .../state_transitions/test_pause_resume.py | 168 +++++++ 3 files changed, 582 insertions(+) create mode 100644 tests/entrypoints/serve/dev/rlhf/conftest.py create mode 100644 tests/entrypoints/serve/dev/rlhf/state_transitions/__init__.py create mode 100644 tests/entrypoints/serve/dev/rlhf/state_transitions/test_pause_resume.py diff --git a/tests/entrypoints/serve/dev/rlhf/conftest.py b/tests/entrypoints/serve/dev/rlhf/conftest.py new file mode 100644 index 000000000000..b2050997905a --- /dev/null +++ b/tests/entrypoints/serve/dev/rlhf/conftest.py @@ -0,0 +1,414 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Shared fixtures and helpers for the RL lifecycle test suite. + +All test modules under this directory import from here to avoid duplication. + +RFC: https://github.com/vllm-project/vllm/issues/45585 +PR: https://github.com/vllm-project/vllm/pull/45586 +""" + +import contextlib +import json +import os +import subprocess +import sys +import threading +import time +from collections.abc import Callable +from contextlib import contextmanager +from dataclasses import dataclass, field +from typing import Any + +import requests + +# --------------------------------------------------------------------------- +# Model / server defaults +# --------------------------------------------------------------------------- + + +MODEL_NAME = os.environ.get("VLLM_TEST_MODEL", "Qwen/Qwen3-0.6B") + + +_BASE_ARGS = [ + "--dtype", + "bfloat16", + "--max-model-len", + "2048", + "--max-num-seqs", + "32", + "--gpu-memory-utilization", + "0.75", + "--enable-sleep-mode", + "--enforce-eager", +] + + +# Lightweight args for state-machine / protocol tests that don't need real +# weights (avoids spending time downloading a 1B checkpoint in T0 tests). +_DUMMY_ARGS = [ + "--dtype", + "bfloat16", + "--max-model-len", + "128", + "--max-num-seqs", + "8", + "--gpu-memory-utilization", + "0.5", + "--enable-sleep-mode", + "--enforce-eager", + "--load-format", + "dummy", +] + + +# --------------------------------------------------------------------------- +# Server harness +# --------------------------------------------------------------------------- + + +@contextmanager +def server( + extra_args=None, + port: int = 8770, + timeout: float = 180.0, + dummy_weights: bool = False, +): + """Launch a vLLM server with the dev router; yield its base URL. + + Args: + extra_args: Additional CLI flags appended after the base args. + port: HTTP port to bind (caller is responsible for uniqueness). + timeout: Seconds to wait for /health before giving up. + dummy_weights: If True, use --load-format dummy (fast, no real weights). + """ + env = {**os.environ, "VLLM_SERVER_DEV_MODE": "1"} + base = _DUMMY_ARGS if dummy_weights else _BASE_ARGS + cmd = [ + sys.executable, + "-m", + "vllm.entrypoints.openai.api_server", + "--model", + MODEL_NAME, + "--port", + str(port), + "--served-model-name", + "m", + *(base + (extra_args or [])), + ] + proc = subprocess.Popen( + cmd, env=env, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE + ) + url = f"http://localhost:{port}" + try: + deadline = time.time() + timeout + while time.time() < deadline: + if proc.poll() is not None: + err = ( + proc.stderr.read(4000).decode(errors="replace") + if proc.stderr + else "" + ) + raise RuntimeError(f"vllm server exited during startup:\n{err}") + with contextlib.suppress(Exception): + if requests.get(f"{url}/health", timeout=3).status_code == 200: + break + time.sleep(1) + else: + proc.terminate() + raise RuntimeError("vllm server did not start in time") + yield url + finally: + proc.terminate() + with contextlib.suppress(subprocess.TimeoutExpired): + proc.wait(timeout=10) + if proc.poll() is None: + proc.kill() + + +# --------------------------------------------------------------------------- +# Polling helper (200-lie workaround) +# --------------------------------------------------------------------------- + + +def poll_until( + predicate: Callable[[], bool], + timeout: float = 10.0, + interval: float = 0.5, +) -> bool: + """Poll predicate() until it returns True or timeout expires. + + Workaround for the vLLM sleep/wake "200-lie" — the HTTP endpoints may + return 200 before the underlying operation is complete, so callers that + need to verify state *after* an operation can use this helper instead of + assuming the 200 means completion. + + Returns True if predicate became true within timeout, False otherwise. + """ + deadline = time.time() + timeout + while time.time() < deadline: + try: + if predicate(): + return True + except Exception: + pass + time.sleep(interval) + return False + + +# --------------------------------------------------------------------------- +# HTTP helpers — generation +# --------------------------------------------------------------------------- + + +def gen(url, prompt="The capital of France is", max_tokens=8, timeout=30): + """Fire a /v1/completions request; return JSON or None on any error.""" + try: + r = requests.post( + f"{url}/v1/completions", + json={ + "model": "m", + "prompt": prompt, + "max_tokens": max_tokens, + "temperature": 0, + }, + timeout=timeout, + ) + return r.json() + except Exception: + return None + + +def gen_with_logprobs( + url, prompt="The capital of France is", max_tokens=8, logprobs=5, timeout=30 +): + """Fire a /v1/completions request with logprobs; return JSON or None.""" + try: + r = requests.post( + f"{url}/v1/completions", + json={ + "model": "m", + "prompt": prompt, + "max_tokens": max_tokens, + "temperature": 0, + "logprobs": logprobs, + }, + timeout=timeout, + ) + return r.json() + except Exception: + return None + + +def ok(resp) -> bool: + """True iff resp is a successful completion (has choices, no error key).""" + return ( + resp is not None + and "choices" in resp + and bool(resp["choices"]) + and "error" not in resp + ) + + +# --------------------------------------------------------------------------- +# HTTP helpers — stream generation +# --------------------------------------------------------------------------- + + +@dataclass +class StreamResult: + started: threading.Event = field(default_factory=threading.Event) + done: threading.Event = field(default_factory=threading.Event) + chunks: list[dict[str, Any]] = field(default_factory=list) + finish_reason: str | None = None + error: Exception | None = None + + +def stream_completion(url: str, result: StreamResult, max_tokens: int) -> None: + try: + with requests.post( + f"{url}/v1/completions", + json={ + "model": "m", + "prompt": "Count upward slowly: one, two, three,", + "max_tokens": max_tokens, + "temperature": 0, + "ignore_eos": True, + "stream": True, + }, + stream=True, + timeout=(5, 60), + ) as response: + response.raise_for_status() + for line in response.iter_lines(decode_unicode=True): + if not line or line == "data: [DONE]": + continue + assert line.startswith("data: ") + chunk = json.loads(line.removeprefix("data: ")) + result.chunks.append(chunk) + choice = chunk["choices"][0] + if choice.get("text"): + result.started.set() + if choice.get("finish_reason") is not None: + result.finish_reason = choice["finish_reason"] + except Exception as error: + result.error = error + finally: + result.done.set() + + +def start_stream(url: str, max_tokens: int) -> tuple[StreamResult, threading.Thread]: + result = StreamResult() + thread = threading.Thread( + target=stream_completion, + args=(url, result, max_tokens), + ) + thread.start() + started = result.started.wait(timeout=10) + if not started or result.done.is_set(): + pause(url, mode="abort") + resume(url) + thread.join(timeout=10) + assert started, "request did not start generating" + assert not result.done.is_set(), "request completed before it could be paused" + return result, thread + + +# --------------------------------------------------------------------------- +# HTTP helpers — pause / resume +# --------------------------------------------------------------------------- + + +def pause(url, mode="abort", clear_cache=True): + return requests.post( + f"{url}/pause", + params={"mode": mode, "clear_cache": clear_cache}, + timeout=15, + ).status_code + + +def resume(url): + return requests.post(f"{url}/resume", timeout=10).status_code + + +def completion_with_cache_details(url: str, prompt: str) -> dict[str, Any]: + response = requests.post( + f"{url}/v1/completions", + json={ + "model": "m", + "prompt": prompt, + "max_tokens": 8, + "temperature": 0, + "logprobs": 1, + }, + timeout=30, + ) + response.raise_for_status() + return response.json() + + +def golden_output(response: dict[str, Any]) -> dict[str, Any]: + choice = response["choices"][0] + usage = response["usage"] + return { + "text": choice["text"], + "finish_reason": choice["finish_reason"], + "tokens": choice["logprobs"]["tokens"], + "prompt_tokens": usage["prompt_tokens"], + "completion_tokens": usage["completion_tokens"], + } + + +def cached_tokens(response: dict[str, Any]) -> int: + return response["usage"]["prompt_tokens_details"]["cached_tokens"] + + +# --------------------------------------------------------------------------- +# HTTP helpers — sleep / wake +# --------------------------------------------------------------------------- + + +def sleep(url, level=1, mode="abort"): + return requests.post( + f"{url}/sleep", params={"level": level, "mode": mode}, timeout=15 + ).status_code + + +def wake(url, tags=None): + params = {"tags": tags} if tags else {} + return requests.post(f"{url}/wake_up", params=params, timeout=20).status_code + + +def is_sleeping(url) -> bool: + return requests.get(f"{url}/is_sleeping", timeout=5).json()["is_sleeping"] + + +def is_paused(url) -> bool: + return requests.get(f"{url}/is_paused", timeout=5).json()["is_paused"] + + +def health(url) -> int: + try: + return requests.get(f"{url}/health", timeout=5).status_code + except Exception: + return 0 + + +# --------------------------------------------------------------------------- +# HTTP helpers — weight transfer +# --------------------------------------------------------------------------- + + +def start_weight_update(url, is_checkpoint_format=True): + return requests.post( + f"{url}/start_weight_update", + json={"is_checkpoint_format": is_checkpoint_format}, + timeout=10, + ) + + +def finish_weight_update(url): + return requests.post(f"{url}/finish_weight_update", timeout=10) + + +def get_world_size(url, include_dp=True): + return requests.get( + f"{url}/get_world_size", + params={"include_dp": include_dp}, + timeout=5, + ) + + +# --------------------------------------------------------------------------- +# GPU / metrics helpers +# --------------------------------------------------------------------------- + + +def gpu_free_bytes(device: int = 0) -> int: + """Read GPU free bytes via subprocess to avoid import-time torch init.""" + out = subprocess.check_output( + [ + sys.executable, + "-c", + f"import torch; f,_=torch.accelerator.get_memory_info({device}); print(f)", + ], + timeout=10, + ) + return int(out.strip()) + + +def sleep_metrics(url): + """Return (awake, weights_offloaded, discard_all) from /metrics.""" + try: + from prometheus_client.parser import text_string_to_metric_families + except ImportError: + return None, None, None + + r = requests.get(f"{url}/metrics", timeout=5) + vals: dict = {} + for family in text_string_to_metric_families(r.text): + if family.name == "vllm:engine_sleep_state": + for s in family.samples: + vals[s.labels.get("sleep_state", "")] = s.value + return vals.get("awake"), vals.get("weights_offloaded"), vals.get("discard_all") diff --git a/tests/entrypoints/serve/dev/rlhf/state_transitions/__init__.py b/tests/entrypoints/serve/dev/rlhf/state_transitions/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/entrypoints/serve/dev/rlhf/state_transitions/test_pause_resume.py b/tests/entrypoints/serve/dev/rlhf/state_transitions/test_pause_resume.py new file mode 100644 index 000000000000..3de9e0054963 --- /dev/null +++ b/tests/entrypoints/serve/dev/rlhf/state_transitions/test_pause_resume.py @@ -0,0 +1,168 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end tests for the vLLM RL pause/resume lifecycle.""" + +import os +import threading +from typing import Any +from unittest.mock import patch + +import pytest +import requests + +from tests.entrypoints.serve.dev.rlhf.conftest import ( + cached_tokens, + completion_with_cache_details, + gen, + golden_output, + is_paused, + ok, + pause, + resume, + server, + start_stream, +) + + +@pytest.fixture(scope="module", params=[False, True], ids=["MRV1", "MRV2"]) +def use_v2(request): + return request.param + + +@pytest.fixture(scope="module") +def server_url(use_v2): + env_vars = { + "VLLM_USE_V2_MODEL_RUNNER": "1" if use_v2 else "0", + } + + with ( + patch.dict(os.environ, env_vars), + server( + extra_args=[ + "--enable-prefix-caching", + "--enable-prompt-tokens-details", + ] + ) as url, + ): + yield url + + +@pytest.fixture(autouse=True) +def restore_unpaused_state(server_url): + assert resume(server_url) == 200 + yield + assert resume(server_url) == 200 + + +class TestPauseResume: + def test_state_and_idempotency_across_cycles(self, server_url): + assert not is_paused(server_url) + + assert resume(server_url) == 200 + assert resume(server_url) == 200 + assert not is_paused(server_url) + + for _ in range(2): + for mode in ("abort", "wait", "keep"): + assert pause(server_url, mode=mode) == 200 + assert pause(server_url, mode=mode) == 200 + assert is_paused(server_url) + + assert resume(server_url) == 200 + assert resume(server_url) == 200 + assert not is_paused(server_url) + + def test_invalid_mode_preserves_state(self, server_url): + for paused in (False, True): + if paused: + assert pause(server_url) == 200 + assert is_paused(server_url) is paused + + response = requests.post( + f"{server_url}/pause", + params={"mode": "invalid"}, + timeout=10, + ) + assert response.status_code == 400 + assert response.json()["error"]["param"] == "query.mode" + assert is_paused(server_url) is paused + + assert resume(server_url) == 200 + + @pytest.mark.parametrize( + ("mode", "max_tokens", "inflight_finish_reason"), + [ + pytest.param("abort", 256, "abort", id="abort"), + pytest.param("wait", 256, "length", id="wait"), + pytest.param("keep", 256, "length", id="keep"), + ], + ) + def test_mode_request_lifecycle( + self, + server_url, + mode, + max_tokens, + inflight_finish_reason, + ): + inflight, inflight_thread = start_stream(server_url, max_tokens) + new_result: dict[str, Any] = {} + new_done = threading.Event() + + def _new_request(): + new_result["response"] = gen(server_url, max_tokens=4, timeout=60) + new_done.set() + + new_thread = threading.Thread(target=_new_request) + try: + assert pause(server_url, mode=mode) == 200 + assert is_paused(server_url) + + if mode in ("abort", "wait"): + assert inflight.done.is_set() + else: + chunks_after_pause = len(inflight.chunks) + assert not inflight.done.wait(timeout=5) + assert len(inflight.chunks) == chunks_after_pause, ( + "in-flight request continued generating in keep mode" + ) + + new_thread.start() + assert not new_done.wait(timeout=0.3), ( + "new request completed while generation was paused" + ) + finally: + assert resume(server_url) == 200 + inflight_thread.join(timeout=30) + if new_thread.ident is not None: + new_thread.join(timeout=30) + + assert not inflight_thread.is_alive() + assert inflight.error is None + assert inflight.finish_reason == inflight_finish_reason + assert not new_thread.is_alive() + assert ok(new_result.get("response")) + + def test_clear_cache_preserves_output_and_controls_prefix_cache(self, server_url): + prompt = ( + "Paris is the capital of France. Berlin is the capital of Germany. " + ) * 20 + + assert pause(server_url, clear_cache=True) == 200 + assert resume(server_url) == 200 + baseline = completion_with_cache_details(server_url, prompt) + warmed = completion_with_cache_details(server_url, prompt) + + assert cached_tokens(baseline) == 0 + assert cached_tokens(warmed) > 0 + + assert pause(server_url, clear_cache=False) == 200 + assert resume(server_url) == 200 + preserved = completion_with_cache_details(server_url, prompt) + assert golden_output(preserved) == golden_output(baseline) + assert cached_tokens(preserved) > 0 + + assert pause(server_url, clear_cache=True) == 200 + assert resume(server_url) == 200 + cleared = completion_with_cache_details(server_url, prompt) + assert golden_output(cleared) == golden_output(baseline) + assert cached_tokens(cleared) == 0 From 5fa8ca971a7a3f19806808f72987a06187bdce55 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Tue, 18 Aug 2026 01:50:35 -0500 Subject: [PATCH 081/839] [ROCm][CI] Move ROCm AITER quantization tests (#40938) Signed-off-by: Andreas Karatzas Signed-off-by: Andreas Karatzas Co-authored-by: OpenAI Codex --- .buildkite/test-amd.yaml | 24 +-- .buildkite/test_areas/kernels.yaml | 1 + .../test_rocm_aiter_mla_fp8_support.py | 137 ++++++++++++++++++ .../test_aiter_hipb_mm_linear_kernel.py | 0 .../quantization}/test_quant_op_schema.py | 0 .../test_rocm_aiter_grouped_quant.py} | 0 .../rocm/aiter/test_mla_fp8_support_check.py | 118 --------------- 7 files changed, 144 insertions(+), 136 deletions(-) create mode 100644 tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py rename tests/{rocm/aiter => kernels/quantization}/test_aiter_hipb_mm_linear_kernel.py (100%) rename tests/{rocm/aiter => kernels/quantization}/test_quant_op_schema.py (100%) rename tests/{rocm/aiter/test_grouped_quant.py => kernels/quantization/test_rocm_aiter_grouped_quant.py} (100%) delete mode 100644 tests/rocm/aiter/test_mla_fp8_support_check.py diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 41dbd629bc9b..25eb3aa0637b 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -1544,11 +1544,13 @@ steps: - vllm/_aiter_ops.py - tests/kernels/attention/test_rocm_aiter_mla_decode.py - tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py + - tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py - tests/kernels/attention/test_rocm_aiter_mla_op_registration.py commands: - rocm-smi - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode.py + - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_op_registration.py - label: Kernels Attention Test %N # TBD @@ -1664,6 +1666,7 @@ steps: - tests/kernels/quant_utils.py - tests/kernels/utils.py - vllm/_aiter_ops.py + - vllm/kernels/aiter_ops.py - vllm/_custom_ops.py - vllm/envs.py - vllm/platforms/rocm.py @@ -2452,24 +2455,6 @@ steps: - uv pip install --system conch-triton-kernels - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py -#----------------------------------------------------------- mi300 · rocm ------------------------------------------------------------# - -- label: ROCm AITER Ops Test # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/_aiter_ops.py - - vllm/envs.py - - vllm/platforms/rocm.py - - tests/rocm/aiter/ - - vllm/v1/attention/backends/mla/rocm_aiter_mla.py - - vllm/v1/attention/selector.py - commands: - - pytest -v -s rocm/aiter/ - #--------------------------------------------------------- mi300 · samplers ----------------------------------------------------------# - label: Samplers Test # TBD @@ -3461,11 +3446,13 @@ steps: - vllm/_aiter_ops.py - tests/kernels/attention/test_rocm_aiter_mla_decode.py - tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py + - tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py - tests/kernels/attention/test_rocm_aiter_mla_op_registration.py commands: - rocm-smi - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode.py + - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_op_registration.py - label: Kernels Attention Test %N # TBD @@ -3525,6 +3512,7 @@ steps: - tests/kernels/quant_utils.py - tests/kernels/utils.py - vllm/_aiter_ops.py + - vllm/kernels/aiter_ops.py - vllm/_custom_ops.py - vllm/envs.py - vllm/platforms/rocm.py diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 887ed03e372a..4dec86e929d9 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -183,6 +183,7 @@ steps: - tests/kernels/quantization - tests/kernels/quantization/test_rocm_skinny_gemms.py - vllm/_aiter_ops.py + - vllm/kernels/aiter_ops.py - vllm/platforms/rocm.py - vllm/model_executor/kernels/ depends_on: diff --git a/tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py b/tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py new file mode 100644 index 000000000000..d67df8ee2b4d --- /dev/null +++ b/tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for ROCm AITER MLA FP8 support detection.""" + +import sys +import types +from typing import Any +from unittest.mock import patch + +import pytest + +from vllm.platforms import current_platform + +_SKIP_UNSUPPORTED_AITER_HARDWARE = True +if current_platform.is_rocm(): + from vllm.platforms.rocm import get_cdna_version + + _SKIP_UNSUPPORTED_AITER_HARDWARE = get_cdna_version() <= 2 + +pytestmark = pytest.mark.skipif( + not current_platform.is_rocm(), reason="ROCm-specific tests" +) + + +@pytest.fixture(autouse=True) +def reset_aiter_mla_support_cache(monkeypatch: pytest.MonkeyPatch) -> None: + import vllm._aiter_ops as aiter_ops + + monkeypatch.setattr(aiter_ops, "_AITER_MLA_SUPPORTS_FP8", None) + + +def _install_fake_aiter_modules( + monkeypatch: pytest.MonkeyPatch, *, supports_fp8: bool +) -> None: + aiter_mod: Any = types.ModuleType("aiter") + mla_mod: Any = types.ModuleType("aiter.mla") + + if supports_fp8: + + def mla_decode_fwd_with_fp8( + q, + kv_buffer, + kv_indptr, + kv_indices, + o, + sm_scale, + q_scale=None, + kv_scale=None, + ): + return None + + mla_decode_fwd: Any = mla_decode_fwd_with_fp8 + + else: + + def mla_decode_fwd_without_fp8( + q, + kv_buffer, + kv_indptr, + kv_indices, + o, + sm_scale, + ): + return None + + mla_decode_fwd = mla_decode_fwd_without_fp8 + + mla_mod.mla_decode_fwd = mla_decode_fwd + aiter_mod.mla = mla_mod + + monkeypatch.setitem(sys.modules, "aiter", aiter_mod) + monkeypatch.setitem(sys.modules, "aiter.mla", mla_mod) + + +def test_aiter_mla_fp8_support_detects_fp8_signature(monkeypatch): + """The support check should detect q_scale and kv_scale parameters.""" + from vllm._aiter_ops import _check_aiter_mla_fp8_support + + _install_fake_aiter_modules(monkeypatch, supports_fp8=True) + + assert _check_aiter_mla_fp8_support() is True + + +def test_aiter_mla_fp8_support_rejects_missing_fp8_signature(monkeypatch): + """The support check should return False when FP8 params are absent.""" + from vllm._aiter_ops import _check_aiter_mla_fp8_support + + _install_fake_aiter_modules(monkeypatch, supports_fp8=False) + + assert _check_aiter_mla_fp8_support() is False + + +@pytest.mark.skipif( + _SKIP_UNSUPPORTED_AITER_HARDWARE, + reason="Installed AITER MLA FP8 check requires CDNA 3 or newer", +) +def test_installed_aiter_mla_supports_fp8(): + """Supported ROCm CI must provide AITER with MLA FP8 scaling.""" + from vllm._aiter_ops import ( + _check_aiter_mla_fp8_support, + is_aiter_found_and_supported, + ) + + assert is_aiter_found_and_supported(), ( + "AITER must be installed on supported ROCm hardware" + ) + assert _check_aiter_mla_fp8_support() is True + + +@pytest.mark.parametrize( + "error_type", + [ImportError, ModuleNotFoundError, AttributeError, ValueError, TypeError], +) +def test_aiter_mla_fp8_support_handles_signature_errors(monkeypatch, error_type): + """The support check should fail closed on import or signature problems.""" + import vllm._aiter_ops as aiter_ops + from vllm._aiter_ops import _check_aiter_mla_fp8_support + + _install_fake_aiter_modules(monkeypatch, supports_fp8=True) + + with patch("inspect.signature", side_effect=error_type("boom")): + assert _check_aiter_mla_fp8_support() is False + assert aiter_ops._AITER_MLA_SUPPORTS_FP8 is False + + +def test_aiter_mla_fp8_support_result_is_cached(monkeypatch): + """The support check should reuse the cached result on later calls.""" + import inspect + + from vllm._aiter_ops import _check_aiter_mla_fp8_support + + _install_fake_aiter_modules(monkeypatch, supports_fp8=True) + + with patch("inspect.signature", wraps=inspect.signature) as signature_mock: + assert _check_aiter_mla_fp8_support() is True + assert _check_aiter_mla_fp8_support() is True + assert signature_mock.call_count == 1 diff --git a/tests/rocm/aiter/test_aiter_hipb_mm_linear_kernel.py b/tests/kernels/quantization/test_aiter_hipb_mm_linear_kernel.py similarity index 100% rename from tests/rocm/aiter/test_aiter_hipb_mm_linear_kernel.py rename to tests/kernels/quantization/test_aiter_hipb_mm_linear_kernel.py diff --git a/tests/rocm/aiter/test_quant_op_schema.py b/tests/kernels/quantization/test_quant_op_schema.py similarity index 100% rename from tests/rocm/aiter/test_quant_op_schema.py rename to tests/kernels/quantization/test_quant_op_schema.py diff --git a/tests/rocm/aiter/test_grouped_quant.py b/tests/kernels/quantization/test_rocm_aiter_grouped_quant.py similarity index 100% rename from tests/rocm/aiter/test_grouped_quant.py rename to tests/kernels/quantization/test_rocm_aiter_grouped_quant.py diff --git a/tests/rocm/aiter/test_mla_fp8_support_check.py b/tests/rocm/aiter/test_mla_fp8_support_check.py deleted file mode 100644 index 28da59a1aefc..000000000000 --- a/tests/rocm/aiter/test_mla_fp8_support_check.py +++ /dev/null @@ -1,118 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit tests for AITER MLA FP8 support detection. - -These tests verify that the _check_aiter_mla_fp8_support() function -correctly handles various error conditions without crashing. -""" - -from unittest.mock import patch - -import pytest - - -class TestAiterMlaFp8SupportCheck: - """Test cases for _check_aiter_mla_fp8_support() function.""" - - def setup_method(self): - """Reset the global cache before each test.""" - import vllm._aiter_ops as aiter_ops - - aiter_ops._AITER_MLA_SUPPORTS_FP8 = None - - @patch("vllm._aiter_ops.is_aiter_found_and_supported", return_value=True) - def test_import_error_handling(self, mock_supported): - """Test that ImportError is handled gracefully.""" - import vllm._aiter_ops as aiter_ops - from vllm._aiter_ops import _check_aiter_mla_fp8_support - - aiter_ops._AITER_MLA_SUPPORTS_FP8 = None - - # Should return False without raising - with patch( - "inspect.signature", - side_effect=ImportError("No module"), - ): - result = _check_aiter_mla_fp8_support() - assert result is False - - @patch("vllm._aiter_ops.is_aiter_found_and_supported", return_value=True) - def test_module_not_found_error_handling(self, mock_supported): - """Test that ModuleNotFoundError is handled gracefully.""" - import vllm._aiter_ops as aiter_ops - from vllm._aiter_ops import _check_aiter_mla_fp8_support - - aiter_ops._AITER_MLA_SUPPORTS_FP8 = None - - with patch( - "inspect.signature", - side_effect=ModuleNotFoundError("Module not found"), - ): - # Should return False without raising - assert _check_aiter_mla_fp8_support() is False - # Cache should be set to False - assert aiter_ops._AITER_MLA_SUPPORTS_FP8 is False - - @patch("vllm._aiter_ops.is_aiter_found_and_supported", return_value=True) - def test_attribute_error_handling(self, mock_supported): - """Test that AttributeError is handled gracefully.""" - import vllm._aiter_ops as aiter_ops - from vllm._aiter_ops import _check_aiter_mla_fp8_support - - aiter_ops._AITER_MLA_SUPPORTS_FP8 = None - - with patch( - "inspect.signature", - side_effect=AttributeError("No attribute"), - ): - assert _check_aiter_mla_fp8_support() is False - assert aiter_ops._AITER_MLA_SUPPORTS_FP8 is False - - @patch("vllm._aiter_ops.is_aiter_found_and_supported", return_value=True) - def test_value_error_handling(self, mock_supported): - """Test that ValueError is handled gracefully (no signature).""" - import vllm._aiter_ops as aiter_ops - from vllm._aiter_ops import _check_aiter_mla_fp8_support - - aiter_ops._AITER_MLA_SUPPORTS_FP8 = None - - with patch( - "inspect.signature", - side_effect=ValueError("No signature"), - ): - assert _check_aiter_mla_fp8_support() is False - assert aiter_ops._AITER_MLA_SUPPORTS_FP8 is False - - @patch("vllm._aiter_ops.is_aiter_found_and_supported", return_value=True) - def test_type_error_handling(self, mock_supported): - """Test that TypeError is handled gracefully (not callable).""" - import vllm._aiter_ops as aiter_ops - from vllm._aiter_ops import _check_aiter_mla_fp8_support - - aiter_ops._AITER_MLA_SUPPORTS_FP8 = None - - with patch( - "inspect.signature", - side_effect=TypeError("Not a callable"), - ): - assert _check_aiter_mla_fp8_support() is False - assert aiter_ops._AITER_MLA_SUPPORTS_FP8 is False - - @patch("vllm._aiter_ops.is_aiter_found_and_supported", return_value=True) - def test_result_caching(self, mock_supported): - """Test that the result is cached after first check.""" - import vllm._aiter_ops as aiter_ops - - # Set cache to True - aiter_ops._AITER_MLA_SUPPORTS_FP8 = True - - from vllm._aiter_ops import _check_aiter_mla_fp8_support - - # Should return cached value without re-checking - result = _check_aiter_mla_fp8_support() - assert result is True - - -if __name__ == "__main__": - pytest.main([__file__, "-v"]) From 2687fec6ef744f7d8ac174303e9a2ee6fe8d7153 Mon Sep 17 00:00:00 2001 From: Jee Jee Li Date: Tue, 18 Aug 2026 14:53:21 +0800 Subject: [PATCH 082/839] Replicated embedding and norm fusion for DSV3 flat model (#48484) Signed-off-by: Jee Jee Li --- tests/kernels/core/test_fused_embed_norm.py | 75 ++++++ vllm/envs.py | 4 + .../model_executor/layers/fused_embed_norm.py | 231 ++++++++++++++++++ vllm/models/deepseek_v32/nvidia/model.py | 38 ++- vllm/models/deepseek_v32/nvidia/mtp.py | 58 +++-- 5 files changed, 385 insertions(+), 21 deletions(-) create mode 100644 tests/kernels/core/test_fused_embed_norm.py create mode 100644 vllm/model_executor/layers/fused_embed_norm.py diff --git a/tests/kernels/core/test_fused_embed_norm.py b/tests/kernels/core/test_fused_embed_norm.py new file mode 100644 index 000000000000..9da43dba8361 --- /dev/null +++ b/tests/kernels/core/test_fused_embed_norm.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the replicated-embedding fused gather/norm kernels +(``vllm.model_executor.layers.fused_embed_norm``). + +The guarantee: enabling ``VLLM_REPLICATE_EMBED`` (replicated table + fused +kernels) must not change model outputs. The gathered residual is bit-exact and +the fused norms match the unfused reference. +""" + +import pytest +import torch + +from vllm.model_executor.layers.fused_embed_norm import ( + fused_embed_eh_norm, + fused_embed_norm, +) + +# The model-local (untouched) eh-norm the replicate path must match. +from vllm.models.deepseek_v32.common.kernels import fused_eh_norm +from vllm.platforms import current_platform +from vllm.utils.torch_utils import set_random_seed + +DTYPE = torch.bfloat16 +VOCAB, HIDDEN, NUM_TOKENS, EPS = 8192, 4096, 129, 1e-6 + +requires_cuda = pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="fused embed/norm Triton kernels require a CUDA/ROCm device", +) + + +def _rmsnorm(x: torch.Tensor, w: torch.Tensor, eps: float) -> torch.Tensor: + # Full-precision (fp32) reference RMSNorm. + var = x.float().pow(2).mean(dim=-1, keepdim=True) + return x.float() * torch.rsqrt(var + eps) * w.float() + + +@requires_cuda +@torch.inference_mode() +def test_fused_embed_norm_matches_reference(): + """Main-model fusion: the residual is the exact gather and the second output + is a correct RMSNorm. The norm matches a full-precision reference to ~2 bf16 + ulp (rtol 1e-2) -- that gap is bf16 rounding, not the kernel.""" + set_random_seed(13) + table = torch.randn(VOCAB, HIDDEN, dtype=DTYPE, device="cuda") + ids = torch.randint(0, VOCAB, (NUM_TOKENS,), dtype=torch.int32, device="cuda") + weight = torch.empty(HIDDEN, dtype=DTYPE, device="cuda").normal_(1.0, 0.1) + + residual, normed = fused_embed_norm(ids, table, chain_weight=weight, eps=EPS) + + embeds = table[ids.long()] + torch.testing.assert_close(residual, embeds, atol=0.0, rtol=0.0) + torch.testing.assert_close( + normed.float(), _rmsnorm(embeds, weight, EPS), atol=1e-3, rtol=1e-2 + ) + + +@requires_cuda +@torch.inference_mode() +def test_fused_embed_eh_norm_matches_reference(): + """MTP fusion (folded gather) is bit-exact vs gathering the embeds and + feeding the untouched model-local ``fused_eh_norm``.""" + set_random_seed(13) + table = torch.randn(VOCAB, HIDDEN, dtype=DTYPE, device="cuda") + ids = torch.randint(0, VOCAB, (NUM_TOKENS,), dtype=torch.int32, device="cuda") + prev = torch.randn(NUM_TOKENS, HIDDEN, dtype=DTYPE, device="cuda") + enorm_w = torch.randn(HIDDEN, dtype=DTYPE, device="cuda") + hnorm_w = torch.randn(HIDDEN, dtype=DTYPE, device="cuda") + positions = torch.arange(NUM_TOKENS, device="cuda") # includes pos 0 + + fused = fused_embed_eh_norm(positions, ids, table, prev, enorm_w, hnorm_w, EPS) + ref = fused_eh_norm(positions, table[ids.long()], prev, enorm_w, hnorm_w, EPS) + + torch.testing.assert_close(fused, ref, atol=0.0, rtol=0.0) diff --git a/vllm/envs.py b/vllm/envs.py index 15705b7c1779..38b64f8c3863 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -156,6 +156,7 @@ VLLM_ENABLE_V1_MULTIPROCESSING: bool = True VLLM_LOG_BATCHSIZE_INTERVAL: float = -1 VLLM_DISABLE_COMPILE_CACHE: bool = False + VLLM_REPLICATE_EMBED: bool = False VLLM_USE_LAYERNAME: bool = True Q_SCALE_CONSTANT: int = 200 K_SCALE_CONSTANT: int = 200 @@ -620,6 +621,9 @@ def _resolve_rust_cli_path() -> str | None: # Enable batch-invariant mode: deterministic results regardless of # batch composition. Requires NVIDIA GPU with compute capability >= 9.0. "VLLM_BATCH_INVARIANT": lambda: bool(int(os.getenv("VLLM_BATCH_INVARIANT", "0"))), + "VLLM_REPLICATE_EMBED": lambda: ( + os.getenv("VLLM_REPLICATE_EMBED", "0").strip().lower() in ("1", "true") + ), # Use tensor descriptors for Q/K/V loads and output stores in the # Triton unified-attention kernel. Enables HW 2D block reads on # Intel XPU; the non-TD branch is dead-code-eliminated at Triton diff --git a/vllm/model_executor/layers/fused_embed_norm.py b/vllm/model_executor/layers/fused_embed_norm.py new file mode 100644 index 000000000000..ddeb5156ae18 --- /dev/null +++ b/vllm/model_executor/layers/fused_embed_norm.py @@ -0,0 +1,231 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Replicated input embedding + its fused gather/norm kernels. + +Groups the ``VLLM_REPLICATE_EMBED`` path in one place: the embedding factory, +the predicate that says whether the fusions apply, and the two Triton fusions +the full on-rank table unlocks -- + + * ``fused_embed_norm``: gather + a chained RMSNorm (e.g. the first decoder + layer's ``input_layernorm``), and + * ``fused_embed_eh_norm``: gather + pos-0 zeroing + enorm/hnorm + cat, the + embed/previous-hidden input norm for a speculative (MTP/eagle) depth layer + (the replicated-table analogue of the model-local ``fused_eh_norm``, which + takes precomputed embeds). + +Self-contained (no model-local imports) so it can live under ``layers/``. +""" + +import torch + +import vllm.envs as envs +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.model_executor.layers.vocab_parallel_embedding import ( + UnquantizedEmbeddingMethod, + VocabParallelEmbedding, +) +from vllm.triton_utils import tl, triton + + +@triton.jit +def _rms_norm(x, w, eps, HIDDEN_SIZE: tl.constexpr): + x = x.to(tl.float32) + mean_sq = tl.sum(x * x, axis=0) / HIDDEN_SIZE + rrms = tl.rsqrt(mean_sq + eps) + w = w.to(tl.float32) + return (x * rrms) * w + + +def make_input_embedding( + num_embeddings: int, + embedding_dim: int, + *, + params_dtype: torch.dtype | None = None, + quant_config=None, + prefix: str = "", + tie_word_embeddings: bool = False, +) -> VocabParallelEmbedding: + """Input token embedding with an optional replicated escape hatch. + + ``VLLM_REPLICATE_EMBED=1`` builds the embedding with ``disable_tp``: the full + table lives on every rank and the lookup is a local gather with no mask and + no all-reduce, which unlocks the fused gather+norm path. The cost is a full + table per rank at TP>1 (no extra memory at TP=1, where vocab-parallel is + already unsharded). A replicated, unsharded table cannot be tied to a + vocab-parallel ``ParallelLMHead``, so tied word embeddings are rejected at + TP>1 (at TP=1 ``disable_tp`` is a no-op and tying still works). + """ + disable_tp = envs.VLLM_REPLICATE_EMBED + if disable_tp and tie_word_embeddings: + assert get_tensor_model_parallel_world_size() == 1, ( + "VLLM_REPLICATE_EMBED is unsupported with tied word embeddings " + "(the replicated table cannot tie to a vocab-parallel lm_head)" + ) + return VocabParallelEmbedding( + num_embeddings, + embedding_dim, + params_dtype=params_dtype, + quant_config=quant_config, + prefix=prefix, + disable_tp=disable_tp, + ) + + +def has_full_vocab_on_rank(embedding: torch.nn.Module) -> bool: + """Whether ``embedding.weight`` is the whole vocab as a plain [V, H] table. + + The fused gather kernels index the table directly, so they need every row + on-rank (``disable_tp``, or any TP=1 run) and an unquantized weight. + """ + return getattr(embedding, "tp_size", 0) == 1 and isinstance( + getattr(embedding, "quant_method", None), UnquantizedEmbeddingMethod + ) + + +@triton.jit +def _fused_embed_norm_kernel( + ids_ptr, # [T] token ids + table_ptr, # [V, H] embedding table (full vocab, replicated on-rank) + table_stride_0, + out_ptr, # [T, H] gathered embedding (the residual stream) + normed_ptr, # [T, H] rmsnorm(out, chain_w) (HAS_NORM only) + chain_w_ptr, # [H] next norm weight (HAS_NORM only) + eps, + H: tl.constexpr, + BLOCK: tl.constexpr, + HAS_NORM: tl.constexpr, +): + tok = tl.program_id(0).to(tl.int64) + off = tl.arange(0, BLOCK) + mask = off < H + row = tl.load(ids_ptr + tok).to(tl.int64) + x = tl.load(table_ptr + row * table_stride_0 + off, mask=mask, other=0.0) + tl.store(out_ptr + tok * H + off, x, mask=mask) + if HAS_NORM: + w = tl.load(chain_w_ptr + off, mask=mask) + y = _rms_norm(x, w, eps, H).to(normed_ptr.dtype.element_ty) + tl.store(normed_ptr + tok * H + off, y, mask=mask) + + +# Base model fusion +def fused_embed_norm( + input_ids: torch.Tensor, + embed_table: torch.Tensor, + chain_weight: torch.Tensor | None = None, + eps: float = 0.0, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Fused embedding row gather (``embed_table[input_ids]``). + + Requires the full vocab on-rank (replicated embedding). When + ``chain_weight`` is given, also emits ``rmsnorm(gathered, chain_weight)`` + (the first decoder layer's ``input_layernorm``) as a second output in the + same launch, so the returned pair is ``(residual, normed_input)``. Bit-exact + vs a plain gather followed by an ``RMSNorm``. + """ + assert embed_table.ndim == 2, embed_table.shape + ids = input_ids.view(-1) + (t,) = ids.shape + h = embed_table.shape[1] + if chain_weight is not None: + assert chain_weight.shape == (h,), (chain_weight.shape, h) + out = torch.empty((t, h), dtype=embed_table.dtype, device=embed_table.device) + normed = torch.empty_like(out) if chain_weight is not None else None + if t > 0: + block = triton.next_power_of_2(h) + _fused_embed_norm_kernel[(t,)]( + ids, + embed_table, + embed_table.stride(0), + out, + normed if normed is not None else out, + chain_weight if chain_weight is not None else embed_table, + eps, + h, + block, + HAS_NORM=chain_weight is not None, + num_warps=min(32, max(4, block // 512)), + ) + if normed is not None: + return out, normed + return out + + +@triton.jit +def _fused_embed_eh_norm_kernel( + pos_ptr, + ids_ptr, # [T] token ids + table_ptr, # [V, H] embedding table (full vocab, replicated on-rank) + table_stride, + prev_ptr, # [T, H] previous-step hidden + prev_stride, + enorm_w_ptr, + hnorm_w_ptr, + eps, + out_ptr, # [T, 2H] + out_stride, + H: tl.constexpr, + BLOCK: tl.constexpr, +): + """MTP input fusion with a folded embedding gather: gather + ``table[ids]``, zero it at position 0, RMSNorm(embed) with enorm and + RMSNorm(prev_hidden) with hnorm, written side-by-side into ``out`` ([N, 2H]) + ready for the eh_proj GEMM. Replaces embedding lookup + where + 2x RMSNorm + + cat. Requires the full table on-rank (replicated embedding).""" + tok = tl.program_id(0) + off = tl.arange(0, BLOCK) + mask = off < H + + pos = tl.load(pos_ptr + tok) + row = tl.load(ids_ptr + tok).to(tl.int64) + e = tl.load(table_ptr + row * table_stride + off, mask=mask, other=0.0) + e = tl.where(pos == 0, 0.0, e.to(tl.float32)) + ew = tl.load(enorm_w_ptr + off, mask=mask) + e_normed = _rms_norm(e, ew, eps, H) + tl.store(out_ptr + tok * out_stride + off, e_normed, mask=mask) + + p = tl.load(prev_ptr + tok * prev_stride + off, mask=mask, other=0.0) + hw = tl.load(hnorm_w_ptr + off, mask=mask) + p_normed = _rms_norm(p, hw, eps, H) + tl.store(out_ptr + tok * out_stride + H + off, p_normed, mask=mask) + + +# MTP fusion +def fused_embed_eh_norm( + positions: torch.Tensor, + input_ids: torch.Tensor, + embed_table: torch.Tensor, + previous_hidden: torch.Tensor, + enorm_w: torch.Tensor, + hnorm_w: torch.Tensor, + eps: float, +) -> torch.Tensor: + """Fused ``cat([enorm(masked embed_table[ids]), hnorm(prev_hidden)])`` -> [N, 2H]. + + Folds the embedding row gather into the MTP eh-norm launch; requires the full + table on-rank (replicated embedding). Bit-exact vs gathering ``embed_table[ + input_ids]`` and passing it to the model-local ``fused_eh_norm``. + """ + assert previous_hidden.ndim == 2 and embed_table.ndim == 2 + n, h = previous_hidden.shape + assert positions.shape == (n,) and input_ids.view(-1).shape == (n,) + assert embed_table.shape[1] == h, (embed_table.shape, h) + assert enorm_w.shape == (h,) and hnorm_w.shape == (h,) + out = torch.empty( + n, 2 * h, dtype=previous_hidden.dtype, device=previous_hidden.device + ) + _fused_embed_eh_norm_kernel[(n,)]( + positions, + input_ids, + embed_table, + embed_table.stride(0), + previous_hidden, + previous_hidden.stride(0), + enorm_w, + hnorm_w, + eps, + out, + out.stride(0), + h, + triton.next_power_of_2(h), + ) + return out diff --git a/vllm/models/deepseek_v32/nvidia/model.py b/vllm/models/deepseek_v32/nvidia/model.py index 0173d1ba6c40..76aa6ca8e161 100644 --- a/vllm/models/deepseek_v32/nvidia/model.py +++ b/vllm/models/deepseek_v32/nvidia/model.py @@ -10,13 +10,15 @@ from vllm.config import VllmConfig from vllm.distributed import get_pp_group from vllm.forward_context import get_forward_context, is_forward_context_available +from vllm.model_executor.layers.fused_embed_norm import ( + fused_embed_norm, + has_full_vocab_on_rank, + make_input_embedding, +) from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, ) from vllm.model_executor.layers.layernorm import RMSNorm -from vllm.model_executor.layers.vocab_parallel_embedding import ( - VocabParallelEmbedding, -) from vllm.model_executor.model_loader.weight_utils import ( default_weight_loader, maybe_remap_kv_scale_name, @@ -116,12 +118,18 @@ def forward( positions: torch.Tensor, hidden_states: torch.Tensor, residual: torch.Tensor | None, + attn_in: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: full_num_tokens = positions.shape[0] if residual is None: + # First layer: hidden_states is the embedding (the residual). residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) + # ``attn_in`` is input_layernorm(embedding) already computed fused + # with the embedding gather; otherwise apply it here. + hidden_states = ( + attn_in if attn_in is not None else self.input_layernorm(hidden_states) + ) elif self.use_sequence_parallel: hidden_states, residual = self.input_layernorm(hidden_states, residual) else: @@ -182,14 +190,17 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) if get_pp_group().is_first_rank: - self.embed_tokens = VocabParallelEmbedding( + self.embed_tokens = make_input_embedding( config.vocab_size, config.hidden_size, quant_config=quant_config, prefix=f"{prefix}.embed_tokens", + tie_word_embeddings=getattr(config, "tie_word_embeddings", False), ) else: self.embed_tokens = PPMissingLayer() + # The fused embed+norm gather needs the full table on-rank. + self.replicated_embed = has_full_vocab_on_rank(self.embed_tokens) self.start_layer, self.end_layer, self.layers = make_layers( config.num_hidden_layers, @@ -222,9 +233,21 @@ def forward( intermediate_tensors: IntermediateTensors | None = None, inputs_embeds: torch.Tensor | None = None, ) -> torch.Tensor | IntermediateTensors: + attn_in = None if get_pp_group().is_first_rank: if inputs_embeds is not None: hidden_states = inputs_embeds + elif self.replicated_embed: + assert input_ids is not None + # Full table on-rank: gather the embedding and the first layer's + # input_layernorm in one launch. ``attn_in`` is the pre-normed + # attention input; ``hidden_states`` is the (residual) embedding. + hidden_states, attn_in = fused_embed_norm( + input_ids, + self.embed_tokens.weight, + chain_weight=self.layers[self.start_layer].input_layernorm.weight, + eps=self.config.rms_norm_eps, + ) else: assert input_ids is not None hidden_states = self.embed_input_ids(input_ids) @@ -242,6 +265,8 @@ def forward( forward_context.is_padding, hidden_states ) hidden_states = sp_shard(hidden_states) + if attn_in is not None: + attn_in = sp_shard(attn_in) assert residual is None, "Currently, SP is not supported with PP" aux_hidden_states = [] @@ -253,7 +278,8 @@ def forward( aux_hidden_states.append( hidden_states if residual is None else hidden_states + residual ) - hidden_states, residual = layer(positions, hidden_states, residual) + hidden_states, residual = layer(positions, hidden_states, residual, attn_in) + attn_in = None if not get_pp_group().is_last_rank: assert not self.use_sequence_parallel, ( diff --git a/vllm/models/deepseek_v32/nvidia/mtp.py b/vllm/models/deepseek_v32/nvidia/mtp.py index fe2085af1109..c2dcbdd64d43 100644 --- a/vllm/models/deepseek_v32/nvidia/mtp.py +++ b/vllm/models/deepseek_v32/nvidia/mtp.py @@ -10,14 +10,16 @@ from vllm._aiter_ops import rocm_aiter_ops from vllm.config import VllmConfig from vllm.forward_context import get_forward_context, is_forward_context_available +from vllm.model_executor.layers.fused_embed_norm import ( + fused_embed_eh_norm, + has_full_vocab_on_rank, + make_input_embedding, +) from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, ) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.logits_processor import LogitsProcessor -from vllm.model_executor.layers.vocab_parallel_embedding import ( - VocabParallelEmbedding, -) from vllm.model_executor.model_loader.mtp_validation import ( is_mtp_completeness_check_enabled, ) @@ -93,18 +95,33 @@ def forward( positions: torch.Tensor, previous_hidden_states: torch.Tensor, inputs_embeds: torch.Tensor | None = None, + embed_table: torch.Tensor | None = None, spec_step_index: int = 0, ) -> torch.Tensor: - assert inputs_embeds is not None - # Fused: zero pos-0 embeds + enorm(embeds) + hnorm(prev) + cat -> [N, 2H]. - eh_input = fused_eh_norm( - positions, - inputs_embeds, - previous_hidden_states, - self.enorm.weight, - self.hnorm.weight, - self.enorm.variance_epsilon, - ) + # Fused zero pos-0 + enorm(embeds) + hnorm(prev) + cat -> [N, 2H]. With a + # replicated table the caller passes ``embed_table`` so the embedding + # lookup is folded in too (fused_embed_eh_norm); otherwise the embeds are + # precomputed and go through the model-local fused_eh_norm. + if embed_table is not None: + eh_input = fused_embed_eh_norm( + positions, + input_ids, + embed_table, + previous_hidden_states, + self.enorm.weight, + self.hnorm.weight, + self.enorm.variance_epsilon, + ) + else: + assert inputs_embeds is not None + eh_input = fused_eh_norm( + positions, + inputs_embeds, + previous_hidden_states, + self.enorm.weight, + self.hnorm.weight, + self.enorm.variance_epsilon, + ) is_sequence_parallel = self.mtp_block.use_sequence_parallel if is_sequence_parallel: if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available(): @@ -158,11 +175,15 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) } ) - self.embed_tokens = VocabParallelEmbedding( + self.embed_tokens = make_input_embedding( config.vocab_size, config.hidden_size, + quant_config=vllm_config.quant_config, prefix=maybe_prefix(prefix, "embed_tokens"), + tie_word_embeddings=getattr(config, "tie_word_embeddings", False), ) + # A full on-rank table lets the eh_norm fusion fold in the embedding gather. + self.replicated_embed = has_full_vocab_on_rank(self.embed_tokens) self.logits_processor = LogitsProcessor(config.vocab_size) def set_skip_topk(self, skip: bool): @@ -192,14 +213,21 @@ def forward( inputs_embeds: torch.Tensor | None = None, spec_step_idx: int = 0, ) -> torch.Tensor: + # With a replicated table, defer the embedding gather to fused_eh_norm + # (folded into the enorm/hnorm/cat launch); otherwise gather it here. + embed_table = None if inputs_embeds is None: - inputs_embeds = self.embed_tokens(input_ids) + if self.replicated_embed: + embed_table = self.embed_tokens.weight + else: + inputs_embeds = self.embed_tokens(input_ids) current_step_idx = spec_step_idx % self.num_mtp_layers return self.layers[str(self.mtp_start_layer_idx + current_step_idx)]( input_ids, positions, previous_hidden_states, inputs_embeds, + embed_table, current_step_idx, ) From b01728b0880ca419bb41199523535457f4ab0010 Mon Sep 17 00:00:00 2001 From: Rajath Pai <82278285+rajathpi@users.noreply.github.com> Date: Tue, 18 Aug 2026 09:07:15 +0200 Subject: [PATCH 083/839] [Bugfix] Return 4xx for client-caused errors in /detokenize (#52622) Signed-off-by: rajathpi --- .../tokenize/test_serving_tokenization.py | 113 ++++++++++++++++++ vllm/entrypoints/serve/tokenize/api_router.py | 12 +- 2 files changed, 115 insertions(+), 10 deletions(-) diff --git a/tests/entrypoints/serve/tokenize/test_serving_tokenization.py b/tests/entrypoints/serve/tokenize/test_serving_tokenization.py index 99267d857557..a74962ab0ddd 100644 --- a/tests/entrypoints/serve/tokenize/test_serving_tokenization.py +++ b/tests/entrypoints/serve/tokenize/test_serving_tokenization.py @@ -1,20 +1,27 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from argparse import Namespace from dataclasses import dataclass, field +from http import HTTPStatus from typing import Any from unittest.mock import AsyncMock, MagicMock import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient from vllm.config.multimodal import MultiModalConfig from vllm.entrypoints.openai.models.protocol import BaseModelPath from vllm.entrypoints.openai.models.serving import OpenAIServingModels +from vllm.entrypoints.serve.exception_handling.register import init_exception_handler +from vllm.entrypoints.serve.tokenize.api_router import attach_router from vllm.entrypoints.serve.tokenize.protocol import ( TokenizeChatRequest, TokenizeCompletionRequest, ) from vllm.entrypoints.serve.tokenize.serving import ServingTokenization +from vllm.exceptions import VLLMNotFoundError, VLLMValidationError from vllm.renderers.online_renderer import OnlineRenderer from vllm.v1.engine.async_llm import AsyncLLM @@ -133,3 +140,109 @@ async def test_tokenize_completion_skips_mm_cache_for_renderer_only_path(): serving.online_renderer.preprocess_completion.call_args.kwargs["skip_mm_cache"] is True ) + + +class TestDetokenizeClientErrorResponses: + """Client-caused errors from /detokenize are 4xx, not 500. + + /detokenize used to wrap ``create_detokenize`` in a blanket + ``except Exception`` that reported every failure as 500, while its + sibling /tokenize let exceptions propagate to the global exception + handlers, which map client-caused errors to 4xx and keep genuine + server errors at 500. These tests pin the shared contract for both + endpoints. See #52246 for the same fix in the Anthropic entrypoint. + """ + + @staticmethod + def _make_api_app(handler: MagicMock) -> FastAPI: + app = FastAPI() + app.state.args = Namespace(log_error_stack=False) + app.state.serving_tokenization = handler + attach_router(app) + init_exception_handler(app) + return app + + def _post(self, handler: MagicMock, path: str = "/detokenize"): + app = self._make_api_app(handler) + body: dict[str, Any] = {"model": "test-model"} + if path == "/detokenize": + body["tokens"] = [1, 2, 3] + else: + body["prompt"] = "Hello" + with TestClient(app, raise_server_exceptions=False) as client: + return client.post(path, json=body) + + def test_vllm_validation_error_returns_bad_request(self): + handler = MagicMock(spec=ServingTokenization) + handler.create_detokenize.side_effect = VLLMValidationError( + "invalid token ids", parameter="tokens" + ) + + response = self._post(handler) + + assert response.status_code == HTTPStatus.BAD_REQUEST + error = response.json()["error"] + assert error["type"] == "BadRequestError" + assert error["param"] == "tokens" + + def test_vllm_not_found_error_returns_not_found(self): + handler = MagicMock(spec=ServingTokenization) + handler.create_detokenize.side_effect = VLLMNotFoundError( + "LoRA adapter nonexistent-lora not found" + ) + + response = self._post(handler) + + assert response.status_code == HTTPStatus.NOT_FOUND + assert response.json()["error"]["type"] == "NotFoundError" + + def test_value_error_returns_bad_request(self): + handler = MagicMock(spec=ServingTokenization) + handler.create_detokenize.side_effect = ValueError("bad input") + + response = self._post(handler) + + assert response.status_code == HTTPStatus.BAD_REQUEST + assert response.json()["error"]["type"] == "BadRequestError" + + def test_overflow_error_returns_bad_request(self): + """Fast tokenizers raise OverflowError for out-of-range token ids. + + This used to be special-cased into a RequestValidationError; the + global OverflowError handler keeps the same 400 status code. + """ + handler = MagicMock(spec=ServingTokenization) + handler.create_detokenize.side_effect = OverflowError( + "out of range integral type conversion attempted" + ) + + response = self._post(handler) + + assert response.status_code == HTTPStatus.BAD_REQUEST + assert response.json()["error"]["type"] == "BadRequestError" + + def test_server_error_still_returns_internal_server_error(self): + """Genuine server bugs keep the existing 500 behaviour.""" + handler = MagicMock(spec=ServingTokenization) + handler.create_detokenize.side_effect = RuntimeError("boom") + + response = self._post(handler) + + assert response.status_code == HTTPStatus.INTERNAL_SERVER_ERROR + assert response.json()["error"]["type"] == "InternalServerError" + + def test_tokenize_and_detokenize_agree_on_client_errors(self): + """The two sibling endpoints report the same client error alike.""" + statuses = {} + for path, method in ( + ("/tokenize", "create_tokenize"), + ("/detokenize", "create_detokenize"), + ): + handler = MagicMock(spec=ServingTokenization) + getattr(handler, method).side_effect = VLLMValidationError( + "invalid input", parameter="model" + ) + statuses[path] = self._post(handler, path=path).status_code + + assert statuses["/tokenize"] == HTTPStatus.BAD_REQUEST + assert statuses["/detokenize"] == HTTPStatus.BAD_REQUEST diff --git a/vllm/entrypoints/serve/tokenize/api_router.py b/vllm/entrypoints/serve/tokenize/api_router.py index 9695e6cceafe..f7a60fe4cf3d 100644 --- a/vllm/entrypoints/serve/tokenize/api_router.py +++ b/vllm/entrypoints/serve/tokenize/api_router.py @@ -4,8 +4,7 @@ from http import HTTPStatus -from fastapi import APIRouter, Depends, FastAPI, HTTPException, Request -from fastapi.exceptions import RequestValidationError +from fastapi import APIRouter, Depends, FastAPI, Request from fastapi.responses import JSONResponse from typing_extensions import assert_never @@ -72,14 +71,7 @@ async def tokenize(request: TokenizeRequest, raw_request: Request): async def detokenize(request: DetokenizeRequest, raw_request: Request): handler = tokenization(raw_request) - try: - generator = await handler.create_detokenize(request, raw_request) - except OverflowError as e: - raise RequestValidationError(errors=[str(e)]) from e - except Exception as e: - raise HTTPException( - status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=str(e) - ) from e + generator = await handler.create_detokenize(request, raw_request) if isinstance(generator, ErrorResponse): return JSONResponse( From 5c9ff5366b039a69b344773bdfead8466ed9a097 Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Tue, 18 Aug 2026 16:22:13 +0800 Subject: [PATCH 084/839] [Bugfix] Accept logprobs=-1 in the Completion API (#46175) Signed-off-by: Yufeng He <40085740+he-yufeng@users.noreply.github.com> --- .../completion/test_completion_error.py | 24 +++++++++++++++++++ .../entrypoints/openai/completion/protocol.py | 8 +++++-- 2 files changed, 30 insertions(+), 2 deletions(-) diff --git a/tests/entrypoints/openai/completion/test_completion_error.py b/tests/entrypoints/openai/completion/test_completion_error.py index 2f3db9f697d0..79fd6206d610 100644 --- a/tests/entrypoints/openai/completion/test_completion_error.py +++ b/tests/entrypoints/openai/completion/test_completion_error.py @@ -476,6 +476,30 @@ def test_negative_prompt_token_ids_flat(): ) +def test_logprobs_minus_one_allowed(): + """logprobs=-1 means "return all logprobs". The sampling layer and the chat + top_logprobs / prompt_logprobs validators all accept -1, so the completion + logprobs validator must accept it too instead of rejecting it as negative.""" + request = CompletionRequest( + model=MODEL_NAME, + prompt="Test prompt", + max_tokens=10, + logprobs=-1, + ) + assert request.logprobs == -1 + + +def test_logprobs_below_minus_one_rejected(): + """Values more negative than -1 stay invalid.""" + with pytest.raises(Exception, match="must be a positive value or -1"): + CompletionRequest( + model=MODEL_NAME, + prompt="Test prompt", + max_tokens=10, + logprobs=-2, + ) + + class TestCompletionPromptListLimit: """Regression tests for CVE: unbounded prompt list fan-out.""" diff --git a/vllm/entrypoints/openai/completion/protocol.py b/vllm/entrypoints/openai/completion/protocol.py index 0cb830294687..2fc6ce1cd9cf 100644 --- a/vllm/entrypoints/openai/completion/protocol.py +++ b/vllm/entrypoints/openai/completion/protocol.py @@ -526,9 +526,13 @@ def check_logprobs(cls, data): parameter="prompt_logprobs", value=prompt_logprobs, ) - if (logprobs := data.get("logprobs")) is not None and logprobs < 0: + if ( + (logprobs := data.get("logprobs")) is not None + and logprobs < 0 + and logprobs != -1 + ): raise VLLMValidationError( - "`logprobs` must be a positive value.", + "`logprobs` must be a positive value or -1.", parameter="logprobs", value=logprobs, ) From e8ad2855e7f5d40665250eb468bfd567e3a4b3c1 Mon Sep 17 00:00:00 2001 From: qli88 Date: Tue, 18 Aug 2026 03:31:51 -0500 Subject: [PATCH 085/839] [Bugfix][ROCm] Fix a few int4/int8 quantization errors (#52112) Signed-off-by: Qiang Li --- .../layers/fused_moe/experts/triton_moe.py | 9 +++-- .../layers/fused_moe/oracle/int_wna16.py | 36 +++++++++++++++++++ .../compressed_tensors_moe_wna16.py | 30 +++++++++++++--- .../layers/quantization/moe_wna16.py | 1 - 4 files changed, 69 insertions(+), 7 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py index 7d600dea69f2..d0ece26cfa01 100644 --- a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py @@ -48,6 +48,8 @@ kFp8StaticTensorSym, kInt4Static, kInt4Static32, + kInt4Static32Asym, + kInt4StaticAsym, kInt8DynamicTensorSym, kInt8DynamicTokenSym, kInt8Static, @@ -186,8 +188,8 @@ def activation( swiglu_limit_func(output, input, activation_config.clamp_limit) return - # SWIGLUOAI_UNINTERLEAVE routes to the silu_and_mul_with_clamp kernel and - # requires a clamp limit. Other activations ignore these parameters. + # SWIGLUOAI_UNINTERLEAVE routes to torch.ops._C.silu_and_mul_with_clamp + # via apply_moe_activation() and requires clamp_limit to be set. if activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE: assert activation_config.clamp_limit is not None, ( "SWIGLUOAI_UNINTERLEAVE requires gemm1_clamp_limit" @@ -580,6 +582,8 @@ def _supports_quant_scheme( kInt4Static, kInt8Static, kInt4Static32, + kInt4StaticAsym, + kInt4Static32Asym, # other group sizes? ] return weight_key in SUPPORTED_W @@ -591,6 +595,7 @@ def _supports_activation(activation: MoEActivation) -> bool: MoEActivation.GELU, MoEActivation.GELU_TANH, MoEActivation.SWIGLUOAI, + MoEActivation.SWIGLUOAI_UNINTERLEAVE, MoEActivation.SWIGLUSTEP, MoEActivation.SILU_NO_MUL, MoEActivation.GELU_NO_MUL, diff --git a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py index 5ab0277afc43..5e4f30b9d197 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -1649,6 +1649,42 @@ def convert_to_wna16_moe_kernel_format( w2_uint8 = w2.transpose(1, 2).contiguous().view(torch.uint8) w13_scale = w13_scale.transpose(1, 2).contiguous() w2_scale = w2_scale.transpose(1, 2).contiguous() + # Zero points from compressed-tensors checkpoints are K-first int32 + # with 8 int4 ZPs packed per element: shape (E, K//gs, N//8). + # fused_moe_kernel_gptq_awq expects N-first uint8 with 2 int4 ZPs + # per byte: shape (E, N//2, K//gs), indexed as + # (offs_bn // 2) * stride_bzn + offs_k_group * stride_bzk. + # Conversion steps: + # (E, K//gs, N//8) int32 + # → transpose(1,2) → (E, N//8, K//gs) int32 + # → view(uint8) → (E, N//8, K//gs*4) [each int32 → 4 bytes] + # → reshape(…,4) → (E, N//8, K//gs, 4) [isolate byte index] + # → permute(0,1,3,2) → (E, N//8, 4, K//gs) [byte index before K] + # → reshape → (E, N//2, K//gs) [kernel expected layout] + # After this, element [e, offs_bn//2, k_group] is the uint8 byte + # holding the two int4 ZPs for output channels offs_bn and offs_bn+1. + if w13_qzeros is not None: + E13, Kg13, Np13 = w13_qzeros.shape + w13_qzeros = ( + w13_qzeros.transpose(1, 2) + .contiguous() + .view(torch.uint8) + .reshape(E13, Np13, Kg13, 4) + .permute(0, 1, 3, 2) + .reshape(E13, Np13 * 4, Kg13) + .contiguous() + ) + if w2_qzeros is not None: + E2, Kg2, Np2 = w2_qzeros.shape + w2_qzeros = ( + w2_qzeros.transpose(1, 2) + .contiguous() + .view(torch.uint8) + .reshape(E2, Np2, Kg2, 4) + .permute(0, 1, 3, 2) + .reshape(E2, Np2 * 4, Kg2) + .contiguous() + ) else: # MoeWNA16 uses N-first uint8 weights and scales. w13_uint8 = w13.view(torch.uint8) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py index eaec0fec4540..bd7ad4234a5f 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py @@ -123,8 +123,6 @@ def __init__( # grouped actorder isn't supported by this kernel assert weight_quant.actorder != "group" - assert self.symmetric, "Only symmetric quantization is supported for MoE" - # Non-Marlin WNA16 always uses bf16/fp16 inputs self.input_dtype = torch.bfloat16 @@ -317,6 +315,27 @@ def create_weights( num_groups_w2 = w2_scales_size // self.group_size num_groups_w13 = hidden_size // self.group_size + if not self.symmetric: + # For asymmetric int4 quantization, packed_factor (=8) int4 ZPs are + # stored per int32. Each TP shard must contain a whole number of + # int32 elements so the packed ZP can be sliced without straddling + # an int32 boundary. + w13_n = 2 * intermediate_size_per_partition # gate+up output channels + if w13_n % self.packed_factor != 0: + raise ValueError( + f"CompressedTensors WNA16 MoE: gate+up output channels per " + f"TP rank (2 * intermediate_size_per_partition = {w13_n}) " + f"must be divisible by packed_factor ({self.packed_factor}). " + f"Use a TP size where 2 * intermediate_size is divisible by " + f"{self.packed_factor}." + ) + if hidden_size % self.packed_factor != 0: + raise ValueError( + f"CompressedTensors WNA16 MoE: hidden_size ({hidden_size}) " + f"must be divisible by packed_factor ({self.packed_factor}) " + f"for correct ZP unpacking." + ) + layer.num_groups_w13 = num_groups_w13 layer.num_groups_w2 = num_groups_w2 @@ -513,8 +532,11 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: replace_parameter(layer, "w13_weight_scale", w13_scales) replace_parameter(layer, "w2_weight_scale", w2_scales) - # CPU fused_experts_cpu requires zero points even for symmetric quant - if not self.symmetric or self.wna16_backend == WNA16MoEBackend.CPU: + # CPU fused_experts_cpu requires zero points even for symmetric quant. + # EMULATION bakes ZP into the dequantized bf16 weights — ZP is None. + if ( + not self.symmetric or self.wna16_backend == WNA16MoEBackend.CPU + ) and self.wna16_backend != WNA16MoEBackend.EMULATION: assert w13_qzeros is not None and w2_qzeros is not None replace_parameter(layer, "w13_weight_zero_point", w13_qzeros) replace_parameter(layer, "w2_weight_zero_point", w2_qzeros) diff --git a/vllm/model_executor/layers/quantization/moe_wna16.py b/vllm/model_executor/layers/quantization/moe_wna16.py index 02207abf3d2d..d2e17d4b9bcf 100644 --- a/vllm/model_executor/layers/quantization/moe_wna16.py +++ b/vllm/model_executor/layers/quantization/moe_wna16.py @@ -222,7 +222,6 @@ def __init__(self, quant_config: MoeWNA16Config, moe: "FusedMoEConfig") -> None: else: scale = kInt4StaticGroupScale elif num_bits == 8: - assert group_size == -1 quant_type = INT8_DTYPE scale = kInt8StaticGroupScale else: From eab1cff5b0ca87ec415c3ca555f060933d0a6b72 Mon Sep 17 00:00:00 2001 From: yimdev <4mengy@gmail.com> Date: Tue, 18 Aug 2026 16:43:12 +0800 Subject: [PATCH 086/839] Harden DeepSeek V3.2 fused kernel grids (#52381) Signed-off-by: yimdev <5779256+yimdev@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../test_fused_deepseek_v32_norm_rope.py | 73 +++++++++++++++++++ vllm/models/deepseek_v32/common/kernels.py | 12 +-- 2 files changed, 79 insertions(+), 6 deletions(-) diff --git a/tests/kernels/test_fused_deepseek_v32_norm_rope.py b/tests/kernels/test_fused_deepseek_v32_norm_rope.py index 48a1f78edbc4..d72c1b31c4d7 100644 --- a/tests/kernels/test_fused_deepseek_v32_norm_rope.py +++ b/tests/kernels/test_fused_deepseek_v32_norm_rope.py @@ -362,6 +362,46 @@ def test_fused_norm_rope_ds_mla(num_tokens: int): assert (topk == 7).all(), "topk buffer should be untouched (no indexer)" +def test_fused_norm_rope_supports_large_token_count(): + """Keep the token count off CUDA grid-y at its 65,536-block boundary.""" + num_tokens = 65536 + dev = "cuda" + dtype = torch.bfloat16 + positions = torch.zeros(num_tokens, device=dev, dtype=torch.int64) + q_c = torch.ones((num_tokens, 1), device=dev, dtype=dtype) + kv_c = torch.ones((num_tokens, 1), device=dev, dtype=dtype) + k_pe = torch.ones((num_tokens, 2), device=dev, dtype=dtype) + norm_w = torch.ones(1, device=dev, dtype=dtype) + cos_sin = torch.tensor([[1.0, 0.0]], device=dev, dtype=torch.float32) + topk = torch.empty((num_tokens, 1), device=dev, dtype=torch.int32) + slot_mapping = torch.arange(num_tokens, device=dev, dtype=torch.int64) + mla_cache = torch.empty((1, num_tokens, 3), device=dev, dtype=dtype) + + q_out = K.fused_norm_rope( + positions, + q_c, + norm_w, + EPS, + kv_c, + norm_w, + EPS, + k_pe, + cos_sin, + None, + None, + None, + EPS, + None, + topk, + slot_mapping=slot_mapping, + mla_kv_cache=mla_cache, + has_indexer=False, + ) + + rows = torch.tensor([0, num_tokens - 1], device=dev) + assert_bf16(q_out[rows], rms_norm(q_c[rows], norm_w), "large-token q norm") + + # ── fused_q ────────────────────────────────────────────────────────────────── @@ -539,6 +579,39 @@ def test_fused_q_bf16_query(num_tokens: int, has_indexer: bool): torch.testing.assert_close(iw_out, iw_ref, rtol=1e-3, atol=1e-3) +def test_fused_q_triton_supports_large_token_count(): + """Keep the token count off CUDA grid-y in the Triton fallback. + + The minimal dimensions also bypass CuTeDSL on SM100. + """ + num_tokens = 65536 + dev = "cuda" + dtype = torch.bfloat16 + positions = torch.zeros(num_tokens, device=dev, dtype=torch.int64) + q_pe = torch.ones((num_tokens, 1, 2), device=dev, dtype=dtype) + ql_nope = torch.ones((num_tokens, 1, 1), device=dev, dtype=dtype) + cos_sin = torch.tensor([[1.0, 0.0]], device=dev, dtype=torch.float32) + q_scale = torch.ones(1, device=dev, dtype=torch.float32) + + _, _, mqa_q = K.fused_q( + positions, + q_pe, + cos_sin, + None, + None, + ql_nope, + q_scale, + None, + 0.0, + 0.0, + has_indexer=False, + ) + + rows = torch.tensor([0, num_tokens - 1], device=dev) + ref = torch.cat([ql_nope, q_pe], dim=-1).to(FP8) + assert_fp8(mqa_q[rows], ref[rows], "large-token fused Q") + + # ── fused_eh_norm (MTP) ────────────────────────────────────────────────────── diff --git a/vllm/models/deepseek_v32/common/kernels.py b/vllm/models/deepseek_v32/common/kernels.py index 2c6e895fdb7e..ec642409cbdb 100644 --- a/vllm/models/deepseek_v32/common/kernels.py +++ b/vllm/models/deepseek_v32/common/kernels.py @@ -152,8 +152,8 @@ def _fused_norm_rope_kernel( INDEX_ROPE_INTERLEAVE: tl.constexpr, USE_PDL: tl.constexpr, ): - pid = tl.program_id(0) - tok_idx = tl.program_id(1) + tok_idx = tl.program_id(0).to(tl.int64) + pid = tl.program_id(1) if USE_PDL: tl.extra.cuda.gdc_wait() tl.extra.cuda.gdc_launch_dependents() @@ -479,7 +479,7 @@ def fused_norm_rope( if q_c_out is None: q_c_out = torch.empty_like(q_c) use_pdl = current_platform.is_arch_support_pdl() - _fused_norm_rope_kernel[(4, num_tokens)]( + _fused_norm_rope_kernel[(num_tokens, 4)]( positions, # Q RMS norm q_c, @@ -594,8 +594,8 @@ def _fused_q_kernel( QUANTIZE_MQA: tl.constexpr, USE_PDL: tl.constexpr, ): - pid = tl.program_id(0) - tok_idx = tl.program_id(1) + tok_idx = tl.program_id(0).to(tl.int64) + pid = tl.program_id(1) head_idx = tl.program_id(2) if USE_PDL: tl.extra.cuda.gdc_wait() @@ -869,7 +869,7 @@ def fused_q( return index_q_fp8, index_weights_out, mqa_q use_pdl = current_platform.is_arch_support_pdl() - _fused_q_kernel[(3, num_tokens, grid_heads)]( + _fused_q_kernel[(num_tokens, 3, grid_heads)]( positions, q_pe, q_pe.stride(0), From 41f179b57aa8ab6f634f508128ce1f1efadd0eb1 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Tue, 18 Aug 2026 19:06:05 +1000 Subject: [PATCH 087/839] [Rust Frontend] Simplify data-parallel size ownership (#52575) Signed-off-by: Bugen Zhao Co-authored-by: Nick Hill --- .buildkite/test-amd.yaml | 3 + .buildkite/test_areas/distributed.yaml | 1 + .buildkite/test_areas/rust_frontend.yaml | 2 + rust/src/cmd/src/cli.rs | 3 +- rust/src/cmd/src/cli/tests.rs | 6 +- rust/src/engine-core-client/src/client.rs | 74 ++++++++++++++++++- rust/src/engine-core-client/src/error.rs | 2 + .../src/engine-core-client/src/mock_engine.rs | 2 +- .../src/protocol/handshake.rs | 7 +- .../engine-core-client/src/tests/client.rs | 33 +++++++++ .../examples/external_engine_openai_qwen.rs | 1 - rust/src/server/src/config.rs | 43 +---------- rust/src/server/src/grpc/control.rs | 2 +- rust/src/server/src/grpc/tests.rs | 8 +- rust/src/server/src/lib.rs | 1 - rust/src/server/src/routes/tests.rs | 65 ++-------------- rust/src/server/src/routes/world_size.rs | 2 +- rust/src/server/src/state.rs | 15 ---- .../distributed/test_dense_dp_world_size.py | 52 +++++++++++++ 19 files changed, 188 insertions(+), 134 deletions(-) create mode 100644 tests/v1/distributed/test_dense_dp_world_size.py diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 25eb3aa0637b..2f13ff83e9f2 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -2417,6 +2417,7 @@ steps: - vllm/v1/engine/ - vllm/v1/worker/ - tests/utils.py + - tests/v1/distributed/test_dense_dp_world_size.py - tests/v1/distributed/test_external_lb_dp.py - tests/v1/distributed/test_hybrid_lb_dp.py - tests/v1/distributed/test_internal_lb_dp.py @@ -2424,6 +2425,7 @@ steps: commands: - export VLLM_USE_RUST_FRONTEND=1 - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_dense_dp_world_size.py - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py -k "not 4 and not server_info" - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py -k "not 4 and not server_info" - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py -k "not 4 and not server_info" @@ -2882,6 +2884,7 @@ steps: - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py + - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_dense_dp_world_size.py - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py - pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index fb0329d58355..32d1ff0342c7 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -154,6 +154,7 @@ steps: - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py + - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_dense_dp_world_size.py - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py - pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp diff --git a/.buildkite/test_areas/rust_frontend.yaml b/.buildkite/test_areas/rust_frontend.yaml index 6486df1c00a1..d56271670172 100644 --- a/.buildkite/test_areas/rust_frontend.yaml +++ b/.buildkite/test_areas/rust_frontend.yaml @@ -108,6 +108,7 @@ steps: - vllm/v1/engine/ - vllm/v1/worker/ - tests/utils.py + - tests/v1/distributed/test_dense_dp_world_size.py - tests/v1/distributed/test_external_lb_dp.py - tests/v1/distributed/test_hybrid_lb_dp.py - tests/v1/distributed/test_internal_lb_dp.py @@ -115,6 +116,7 @@ steps: - export VLLM_USE_RUST_FRONTEND=1 - export VLLM_WORKER_MULTIPROC_METHOD=spawn - export NCCL_CUMEM_HOST_ENABLE=0 + - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_dense_dp_world_size.py - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py -k "not 4 and not server_info" - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py -k "not 4 and not server_info" - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py -k "not 4 and not server_info" diff --git a/rust/src/cmd/src/cli.rs b/rust/src/cmd/src/cli.rs index e851fad54996..64b6dc9800cf 100644 --- a/rust/src/cmd/src/cli.rs +++ b/rust/src/cmd/src/cli.rs @@ -474,9 +474,9 @@ impl SharedRuntimeArgs { output_address, engine_start_index, engine_count, + data_parallel_size, ready_timeout, }, - data_parallel_size, coordinator_mode: match coordinator_address { Some(address) => CoordinatorMode::External { address }, None => CoordinatorMode::None, @@ -533,7 +533,6 @@ impl SharedRuntimeArgs { local_input_address, local_output_address, }, - data_parallel_size: engine_count, coordinator_mode: CoordinatorMode::MaybeInProc, model: self.model, served_model_name: self.served_model_name, diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs index 98a22ad3244c..80436ec423e6 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -864,7 +864,7 @@ fn frontend_args_json_applies_defaults() { assert_eq!(args.runtime.renderer, RendererSelection::Auto); assert_eq!(args.runtime.max_logprobs, None); assert_eq!(args.runtime.shutdown_timeout, 0); - assert_eq!(args.into_config().data_parallel_size, 4); + assert_eq!(args.into_config().transport_mode.data_parallel_size(), 4); } #[test] @@ -1521,7 +1521,6 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() { "", ), }, - data_parallel_size: 4, coordinator_mode: MaybeInProc, model: "Qwen/Qwen3-0.6B", served_model_name: [], @@ -1607,7 +1606,6 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() { local_input_address: None, local_output_address: None, }, - data_parallel_size: 4, coordinator_mode: MaybeInProc, model: "Qwen/Qwen3-0.6B", served_model_name: [], @@ -1710,9 +1708,9 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present output_address: "ipc:///tmp/output.sock", engine_start_index: 3, engine_count: 1, + data_parallel_size: 4, ready_timeout: 600s, }, - data_parallel_size: 4, coordinator_mode: External { address: "tcp://127.0.0.1:7000", }, diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index 1efca89621ff..e7b6e42b6a0b 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/rust/src/engine-core-client/src/client.rs @@ -15,7 +15,7 @@ use tracing::{debug, info, trace}; use crate::client::imp::{ClientInner, run_abort_loop, run_output_dispatcher_loop}; use crate::coordinator::CoordinatorHandle; -use crate::error::{Error, Result}; +use crate::error::{Error, Result, bail_invalid_client_config}; use crate::protocol::dtype::ModelDtype; use crate::protocol::handshake::EngineCoreReadyResponse; use crate::protocol::lora::LoraRequest; @@ -68,11 +68,70 @@ pub enum TransportMode { engine_start_index: u32, /// Total number of engines expected to register on this transport. engine_count: usize, + /// Deployment-wide data-parallel size. This may be larger than + /// `engine_count` when a supervisor partitions ranks across frontends. + data_parallel_size: usize, /// Maximum time to wait for all expected engines to register. ready_timeout: Duration, }, } +impl TransportMode { + /// Return the deployment-wide data-parallel size for this transport. + pub fn data_parallel_size(&self) -> usize { + match self { + Self::HandshakeOwner { engine_count, .. } => *engine_count, + Self::Bootstrapped { + data_parallel_size, .. + } => *data_parallel_size, + } + } + + /// Validate the transport topology before opening any sockets. + pub fn validate(&self) -> Result<()> { + let data_parallel_size = self.data_parallel_size(); + if data_parallel_size == 0 { + bail_invalid_client_config!("data parallel size must be at least 1"); + } + if data_parallel_size > usize::from(u16::MAX) + 1 { + bail_invalid_client_config!( + "data parallel size ({data_parallel_size}) exceeds the two-byte engine identity limit" + ); + } + + match self { + Self::HandshakeOwner { .. } => {} + Self::Bootstrapped { + engine_start_index, + engine_count, + .. + } => { + if *engine_count == 0 { + bail_invalid_client_config!("engine count must be at least 1"); + } + let engine_start_index = usize::try_from(*engine_start_index).map_err(|_| { + Error::InvalidClientConfig { + message: "engine start index does not fit usize".to_string(), + } + })?; + let engine_end_index = + engine_start_index.checked_add(*engine_count).ok_or_else(|| { + Error::InvalidClientConfig { + message: "engine start index + engine count overflows".to_string(), + } + })?; + if engine_end_index > data_parallel_size { + bail_invalid_client_config!( + "connected engine range [{engine_start_index}, {engine_end_index}) exceeds data parallel size ({data_parallel_size})" + ); + } + } + } + + Ok(()) + } +} + /// Which coordinator implementation should be active when one is present for a /// frontend client. #[derive(Debug, Clone, PartialEq, Eq)] @@ -117,6 +176,11 @@ impl EngineCoreClientConfig { } } + /// Validate the client topology before opening any transport sockets. + pub fn validate(&self) -> Result<()> { + self.transport_mode.validate() + } + /// Set the model name used by frontend-side metrics and diagnostics. pub fn with_model_name(mut self, model_name: impl Into) -> Self { self.model_name = model_name.into(); @@ -226,6 +290,7 @@ impl EngineCoreClient { /// handshake. In bootstrapped mode it binds the provided frontend /// sockets and waits for the expected engine registration frames. pub async fn connect(config: EngineCoreClientConfig) -> Result { + config.validate()?; let connected = match &config.transport_mode { TransportMode::HandshakeOwner { handshake_address, @@ -261,6 +326,7 @@ impl EngineCoreClient { engine_start_index, engine_count, ready_timeout, + .. } => { if let Some(CoordinatorMode::InProc) = config.coordinator_mode { panic!("cannot use in-process coordinator with bootstrapped transport mode") @@ -378,6 +444,12 @@ impl EngineCoreClient { self.engines.len() } + /// Return the deployment-wide data-parallel size configured for this + /// client. + pub fn data_parallel_size(&self) -> usize { + self.config.transport_mode.data_parallel_size() + } + /// Return the engine-side indices connected to this client. pub fn engine_indices(&self) -> Vec { self.engines diff --git a/rust/src/engine-core-client/src/error.rs b/rust/src/engine-core-client/src/error.rs index 98a578061efe..d43952f0239f 100644 --- a/rust/src/engine-core-client/src/error.rs +++ b/rust/src/engine-core-client/src/error.rs @@ -49,6 +49,8 @@ pub enum Error { UnexpectedHandshakeIdentity { expected: Vec, actual: Vec }, #[error("unexpected startup handshake message: {message}")] UnexpectedHandshakeMessage { message: String }, + #[error("invalid engine-core client configuration: {message}")] + InvalidClientConfig { message: String }, #[error("unexpected non-control output on coordinator path: {message}")] UnexpectedCoordinatorOutput { message: String }, #[error("unexpected output on main dispatcher path: {message}")] diff --git a/rust/src/engine-core-client/src/mock_engine.rs b/rust/src/engine-core-client/src/mock_engine.rs index 8bd6d773607e..77d72b525fbd 100644 --- a/rust/src/engine-core-client/src/mock_engine.rs +++ b/rust/src/engine-core-client/src/mock_engine.rs @@ -56,7 +56,7 @@ pub fn default_ready_response() -> EngineCoreReadyResponse { dtype: ModelDtype::Float32, vllm_version: "test-vllm-version".to_string(), world_size: 1, - data_parallel_size: 1, + effective_data_parallel_size: 1, tensor_parallel_size: 1, pipeline_parallel_size: 1, decode_context_parallel_size: 1, diff --git a/rust/src/engine-core-client/src/protocol/handshake.rs b/rust/src/engine-core-client/src/protocol/handshake.rs index 5cb4c9804678..c7146abdbb6a 100644 --- a/rust/src/engine-core-client/src/protocol/handshake.rs +++ b/rust/src/engine-core-client/src/protocol/handshake.rs @@ -63,8 +63,11 @@ pub struct EngineCoreReadyResponse { pub vllm_version: String, /// World size (TP * PP) from the parallel config. pub world_size: u64, - /// Data parallelism size from the parallel config. - pub data_parallel_size: u64, + /// Data-parallel size from this EngineCore's effective parallel config. + /// Dense independent-DP ranks are reconfigured to report `1`; the client + /// transport owns the deployment-wide data-parallel size. + #[serde(rename = "data_parallel_size")] + pub effective_data_parallel_size: u64, // Required discovery metadata; EngineCore and client versions must match. /// Tensor-parallel size of this engine. pub tensor_parallel_size: u32, diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index 913b894ada7a..2f988e2bfe94 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -303,6 +303,7 @@ fn bootstrapped_test_config( output_address, engine_start_index: 0, engine_count, + data_parallel_size: engine_count, ready_timeout, }, coordinator_mode, @@ -330,15 +331,47 @@ fn bootstrapped_test_config_with_start_index( ); let TransportMode::Bootstrapped { engine_start_index: start, + data_parallel_size, .. } = &mut config.transport_mode else { unreachable!("bootstrapped_test_config returns bootstrapped transport") }; *start = engine_start_index; + *data_parallel_size = usize::try_from(engine_start_index) + .expect("test start index fits usize") + .checked_add(engine_count) + .expect("test engine range fits usize"); config } +#[test] +fn client_config_validates_bootstrapped_dp_range() { + let mut config = bootstrapped_test_config_with_start_index( + "ipc://unused-input".to_string(), + "ipc://unused-output".to_string(), + 1, + 2, + Duration::from_secs(1), + 0, + None, + ); + config.validate().expect("frontend may own a subset of global DP ranks"); + + let TransportMode::Bootstrapped { + data_parallel_size, .. + } = &mut config.transport_mode + else { + unreachable!("expected bootstrapped transport") + }; + *data_parallel_size = 2; + let error = config.validate().expect_err("engine range above DP size must fail"); + expect_test::expect![[ + "invalid engine-core client configuration: connected engine range [1, 3) exceeds data parallel size (2)" + ]] + .assert_eq(&error.to_string()); +} + async fn recv_xpub_message(xpub: &mut XPubSocket) -> Vec { xpub.recv().await.unwrap().into_vec() } diff --git a/rust/src/server/examples/external_engine_openai_qwen.rs b/rust/src/server/examples/external_engine_openai_qwen.rs index 591252699a10..387776263059 100644 --- a/rust/src/server/examples/external_engine_openai_qwen.rs +++ b/rust/src/server/examples/external_engine_openai_qwen.rs @@ -58,7 +58,6 @@ async fn main() -> Result<()> { local_input_address: None, local_output_address: None, }, - data_parallel_size: args.engine_count, coordinator_mode: CoordinatorMode::MaybeInProc, model: args.model, served_model_name: vec![], diff --git a/rust/src/server/src/config.rs b/rust/src/server/src/config.rs index b67e87f45f3c..a33d7e25fed8 100644 --- a/rust/src/server/src/config.rs +++ b/rust/src/server/src/config.rs @@ -161,8 +161,6 @@ impl TlsConfig { pub struct Config { /// Frontend-to-engine transport setup. pub transport_mode: TransportMode, - /// Deployment-wide data-parallel size retained by the frontend. - pub data_parallel_size: usize, /// Requested frontend-side coordinator behavior. pub coordinator_mode: CoordinatorMode, /// Backend model identifier used for engine-core loading. @@ -239,46 +237,7 @@ impl Config { max_logprobs ); } - if self.data_parallel_size == 0 { - bail!("data parallel size must be at least 1"); - } - if self.data_parallel_size > usize::from(u16::MAX) + 1 { - bail!( - "data parallel size ({}) exceeds the two-byte engine identity limit", - self.data_parallel_size - ); - } - match &self.transport_mode { - TransportMode::HandshakeOwner { engine_count, .. } => { - if *engine_count != self.data_parallel_size { - bail!( - "managed frontend engine count ({engine_count}) must equal data parallel size ({})", - self.data_parallel_size - ); - } - } - TransportMode::Bootstrapped { - engine_start_index, - engine_count, - .. - } => { - if *engine_count == 0 { - bail!("engine count must be at least 1"); - } - let engine_start_index = usize::try_from(*engine_start_index) - .map_err(|_| anyhow::anyhow!("engine start index does not fit usize"))?; - let engine_end_index = - engine_start_index.checked_add(*engine_count).ok_or_else(|| { - anyhow::anyhow!("engine start index + engine count overflows") - })?; - if engine_end_index > self.data_parallel_size { - bail!( - "connected engine range [{engine_start_index}, {engine_end_index}) exceeds data parallel size {}", - self.data_parallel_size - ); - } - } - } + self.transport_mode.validate()?; Ok(()) } diff --git a/rust/src/server/src/grpc/control.rs b/rust/src/server/src/grpc/control.rs index c08803306fc8..9637dcbf8efa 100644 --- a/rust/src/server/src/grpc/control.rs +++ b/rust/src/server/src/grpc/control.rs @@ -43,7 +43,7 @@ impl ControlServiceImpl { pb::ParallelismInfo { tensor_parallel_size: ready.tensor_parallel_size, pipeline_parallel_size: ready.pipeline_parallel_size, - data_parallel_size: self.state.data_parallel_size().min(u32::MAX as usize) as u32, + data_parallel_size: self.client().data_parallel_size().min(u32::MAX as usize) as u32, data_parallel_rank: ready.data_parallel_rank, decode_context_parallel_size: ready.decode_context_parallel_size, } diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index d5ee87a87f28..9c1ea33b6b0e 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -1625,7 +1625,7 @@ async fn control_aggregates_multi_engine_capacity() { let mut ready_0 = default_ready_response(); ready_0.max_model_len = 8_192; ready_0.num_gpu_blocks = 10; - ready_0.data_parallel_size = 2; + ready_0.effective_data_parallel_size = 2; ready_0.weight_transfer_backend = Some("nccl".to_string()); ready_0.enable_sleep_mode = true; ready_0.supports_draft_weight_updates = true; @@ -1633,7 +1633,7 @@ async fn control_aggregates_multi_engine_capacity() { let mut ready_1 = default_ready_response(); ready_1.max_model_len = 4_096; ready_1.num_gpu_blocks = 20; - ready_1.data_parallel_size = 2; + ready_1.effective_data_parallel_size = 2; ready_1.data_parallel_rank = 1; let engine_tasks = [ready_0, ready_1].map(|ready| { @@ -1667,7 +1667,7 @@ async fn control_aggregates_multi_engine_capacity() { Llm::new(client), Arc::new(FakeTextBackend) as Arc, ); - let state = AppState::new(vec!["test-model".to_string()], chat).with_data_parallel_size(4); + let state = AppState::new(vec!["test-model".to_string()], chat); let service = ControlServiceImpl::new(Arc::new(state)); let server = pb::control_server::Control::get_server_info( @@ -1684,7 +1684,7 @@ async fn control_aggregates_multi_engine_capacity() { assert!(rl.weight_transfer_backend.is_empty()); assert!(!rl.sleep_mode_enabled); assert!(!rl.draft_weight_updates_enabled); - assert_eq!(server.parallelism.unwrap().data_parallel_size, 4); + assert_eq!(server.parallelism.unwrap().data_parallel_size, 2); drop(engine_tasks); } diff --git a/rust/src/server/src/lib.rs b/rust/src/server/src/lib.rs index 4423287e5175..f9f5c423506f 100644 --- a/rust/src/server/src/lib.rs +++ b/rust/src/server/src/lib.rs @@ -141,7 +141,6 @@ async fn build_state(config: &Config) -> Result> { .with_model_path(config.model.clone()) .with_api_server_options(config.api_server_options) .with_server_info(ServerInfoSnapshot::from_config(config)) - .with_data_parallel_size(config.data_parallel_size) .with_api_keys(config.api_keys.clone()) .with_cors(config.cors.clone()) .with_profiler(config.profiler.clone()), diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 260851b4efe5..9ac21ae07f9f 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -704,14 +704,13 @@ async fn test_app_with_dev_mode(dev_mode_enabled: bool) -> axum::Router { /// stays alive for the duration of the test. async fn test_dev_mode_app_with_ready( ready_response: vllm_engine_core_client::protocol::handshake::EngineCoreReadyResponse, - data_parallel_size: usize, ) -> (axum::Router, MockEngineTask) { let ipc = IpcNamespace::new().expect("create ipc namespace"); let handshake_address = ipc.handshake_endpoint(); let engine_id = b"engine-world-size".to_vec(); let engine_task = MockEngineTask::new(spawn_mock_engine_task_with_ready( handshake_address.clone(), - engine_id.clone(), + engine_id, ready_response, |_dealer, _push| boxed_test_future(async {}), )); @@ -729,10 +728,10 @@ async fn test_dev_mode_app_with_ready( let chat = ChatLlm::from_shared_backend(test_llm(client), Arc::new(FakeChatBackend::new())); let app = build_router_with_dev_mode( - Arc::new( - AppState::new(vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()], chat) - .with_data_parallel_size(data_parallel_size), - ), + Arc::new(AppState::new( + vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()], + chat, + )), true, ); (app, engine_task) @@ -6292,7 +6291,7 @@ async fn tokenize_allows_prompts_at_or_above_max_model_len() { max_model_len: 4, ..default_ready_response() }; - let (mut app, _engine_task) = test_dev_mode_app_with_ready(ready, 1).await; + let (mut app, _engine_task) = test_dev_mode_app_with_ready(ready).await; let (completion_status, completion_json) = post_json( &mut app, @@ -6543,58 +6542,6 @@ async fn world_size_endpoint_is_dev_mode_only() { assert_eq!(response.status(), StatusCode::NOT_FOUND); } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -#[serial] -async fn world_size_includes_data_parallelism_by_default() { - let ready = vllm_engine_core_client::protocol::handshake::EngineCoreReadyResponse { - world_size: 2, - data_parallel_size: 1, - ..default_ready_response() - }; - let (mut app, _engine_task) = test_dev_mode_app_with_ready(ready, 4).await; - - let response = app - .call( - Request::builder() - .uri("/get_world_size") - .body(Body::empty()) - .expect("build request"), - ) - .await - .expect("call app"); - - assert_eq!(response.status(), StatusCode::OK); - let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body"); - let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json"); - assert_eq!(json, json!({"world_size": 8})); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -#[serial] -async fn world_size_excludes_data_parallelism_when_include_dp_false() { - let ready = vllm_engine_core_client::protocol::handshake::EngineCoreReadyResponse { - world_size: 2, - data_parallel_size: 4, - ..default_ready_response() - }; - let (mut app, _engine_task) = test_dev_mode_app_with_ready(ready, 4).await; - - let response = app - .call( - Request::builder() - .uri("/get_world_size?include_dp=false") - .body(Body::empty()) - .expect("build request"), - ) - .await - .expect("call app"); - - assert_eq!(response.status(), StatusCode::OK); - let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body"); - let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json"); - assert_eq!(json, json!({"world_size": 2})); -} - // ========================= Profiler route tests ========================= async fn test_profiling_app_with_engine_script(script: F) -> (axum::Router, MockEngineTask) diff --git a/rust/src/server/src/routes/world_size.rs b/rust/src/server/src/routes/world_size.rs index 2be1bf76b7e7..8f5d3b204cd0 100644 --- a/rust/src/server/src/routes/world_size.rs +++ b/rust/src/server/src/routes/world_size.rs @@ -47,7 +47,7 @@ pub async fn get_world_size( let ws = client.world_size(); let world_size = if params.include_dp { - let dp = state.data_parallel_size() as u64; + let dp = client.data_parallel_size() as u64; ws * dp } else { ws diff --git a/rust/src/server/src/state.rs b/rust/src/server/src/state.rs index f745cf7c0042..722dede8dda0 100644 --- a/rust/src/server/src/state.rs +++ b/rust/src/server/src/state.rs @@ -40,8 +40,6 @@ pub struct AppState { pub cors: CorsConfig, /// Runtime server information returned by `/server_info`, when available. server_info: Option, - /// Deployment-wide data-parallel size retained by the frontend. - data_parallel_size: usize, /// SHA-256 hashes of API keys accepted as bearer tokens for guarded routes. api_key_hashes: Vec, /// Number of in-flight inference requests currently owned by this frontend. @@ -71,14 +69,12 @@ impl AppState { !served_model_names.is_empty(), "served_model_names must not be empty" ); - let data_parallel_size = chat.engine_core_client().engine_count(); Self { served_model_names, chat, api_server_options: ApiServerOptions::default(), cors: CorsConfig::default(), server_info: None, - data_parallel_size, api_key_hashes: Vec::new(), server_load: AtomicU64::new(0), lora_manager: LoraManager::new(), @@ -118,17 +114,6 @@ impl AppState { self } - /// Set the deployment-wide data-parallel size reported by frontend APIs. - pub(crate) fn with_data_parallel_size(mut self, size: usize) -> Self { - self.data_parallel_size = size; - self - } - - /// Return the deployment-wide data-parallel size. - pub(crate) fn data_parallel_size(&self) -> usize { - self.data_parallel_size - } - /// Build a `/server_info` response payload. pub(crate) fn server_info_response( &self, diff --git a/tests/v1/distributed/test_dense_dp_world_size.py b/tests/v1/distributed/test_dense_dp_world_size.py new file mode 100644 index 000000000000..1fc37936463a --- /dev/null +++ b/tests/v1/distributed/test_dense_dp_world_size.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import os + +import requests + +from tests.utils import RemoteOpenAIServer +from vllm.platforms import current_platform + +MODEL_NAME = "Qwen/Qwen3-0.6B" +DP_SIZE = int(os.getenv("DP_SIZE", "2")) +TP_SIZE = int(os.getenv("TP_SIZE", "1")) + + +def test_dense_dp_world_size(): + server_args = [ + "--dtype", + "bfloat16", + "--max-model-len", + "2048", + "--max-num-seqs", + "128", + "--enforce-eager", + "--data-parallel-size", + str(DP_SIZE), + "--data-parallel-size-local", + str(DP_SIZE), + "--tensor-parallel-size", + str(TP_SIZE), + ] + env_dict = { + "VLLM_SERVER_DEV_MODE": "1", + current_platform.device_control_env_var: ",".join( + str(current_platform.device_id_to_physical_device_id(i)) + for i in range(DP_SIZE * TP_SIZE) + ), + } + + with RemoteOpenAIServer( + MODEL_NAME, + server_args, + env_dict=env_dict, + ) as server: + response = requests.get(server.url_for("get_world_size")) + response.raise_for_status() + assert response.json() == {"world_size": TP_SIZE * DP_SIZE} + + response = requests.get( + server.url_for("get_world_size"), params={"include_dp": "false"} + ) + response.raise_for_status() + assert response.json() == {"world_size": TP_SIZE} From be3f614ff158ccecc99546958bf5e4893f8e0581 Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Tue, 18 Aug 2026 07:28:00 -0500 Subject: [PATCH 088/839] [CI] Register CPU CI "VLLM_CPU_CI_ENV" environment variable (#52633) Signed-off-by: Taneem Ibrahim --- vllm/envs.py | 3 +++ vllm/platforms/cpu.py | 6 ++---- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/vllm/envs.py b/vllm/envs.py index 38b64f8c3863..1a7458802a90 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -51,6 +51,7 @@ VLLM_CPU_KVCACHE_SPACE: int | None = 0 VLLM_CPU_OMP_THREADS_BIND: str = "auto" VLLM_CPU_NUM_OF_RESERVED_CPU: int | None = None + VLLM_CPU_CI_ENV: bool = False VLLM_CPU_ATTN_SPLIT_KV: bool = True VLLM_ZENTORCH_WEIGHT_PREPACK: bool = True VLLM_CPU_INT4_W4A8: bool = True @@ -885,6 +886,8 @@ def _resolve_rust_cli_path() -> str | None: if "VLLM_CPU_NUM_OF_RESERVED_CPU" in os.environ else None ), + # (CPU backend only) whether vLLM is running in a CI environment. + "VLLM_CPU_CI_ENV": lambda: bool(int(os.getenv("VLLM_CPU_CI_ENV", "0"))), # (CPU backend only) whether to enable attention spilt KV. "VLLM_CPU_ATTN_SPLIT_KV": lambda: bool( int(os.getenv("VLLM_CPU_ATTN_SPLIT_KV", "1")) diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index 5633e160cc2c..dd18363e0d26 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -10,6 +10,7 @@ import torch +from vllm import envs from vllm.logger import init_logger from vllm.utils.cpu_resource_utils import ( DEVICE_CONTROL_ENV_VAR, @@ -203,10 +204,7 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None: # cache. So use VLLM_CPU_CI_ENV to indicate the CI environment, # and just execute model with dynamo + eager mode to save time. # VLLM_CPU_CI_ENV is only used as an internal variable. - if os.environ.get("VLLM_CPU_CI_ENV", "0") != "0": - backend = "eager" - else: - backend = "inductor" + backend = "eager" if envs.VLLM_CPU_CI_ENV else "inductor" compilation_config.mode = CompilationMode.DYNAMO_TRACE_ONCE compilation_config.backend = backend From d29dc3ab87840aef42129b30825176295ea73b07 Mon Sep 17 00:00:00 2001 From: Alexander Lee Date: Tue, 18 Aug 2026 08:41:26 -0400 Subject: [PATCH 089/839] [Bugfix][Gemma4] Align parser enable_thinking default with template (#52430) Signed-off-by: lxy-alexander Co-authored-by: Chauncey Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../engine/test_gemma4_streaming_reasoning.py | 19 ++++++++++++++----- .../reasoning/test_gemma4_reasoning_parser.py | 19 ++++++++++++++++--- vllm/parser/gemma4.py | 2 +- 3 files changed, 31 insertions(+), 9 deletions(-) diff --git a/tests/parser/engine/test_gemma4_streaming_reasoning.py b/tests/parser/engine/test_gemma4_streaming_reasoning.py index 5f6b0a414e14..10681786afbb 100644 --- a/tests/parser/engine/test_gemma4_streaming_reasoning.py +++ b/tests/parser/engine/test_gemma4_streaming_reasoning.py @@ -170,7 +170,7 @@ def mock_tokenizer(): @pytest.fixture def parser(mock_tokenizer): - return Gemma4Parser(mock_tokenizer) + return Gemma4Parser(mock_tokenizer, chat_template_kwargs={"enable_thinking": True}) @pytest.fixture @@ -289,7 +289,10 @@ def open_reasoning_tokenizer(self): @pytest.fixture def open_reasoning_parser(self, open_reasoning_tokenizer): - return Gemma4Parser(open_reasoning_tokenizer) + return Gemma4Parser( + open_reasoning_tokenizer, + chat_template_kwargs={"enable_thinking": True}, + ) @staticmethod def _prompt_ids_open_channel() -> list[int]: @@ -394,7 +397,9 @@ def pre_init_tokenizer(self): @pytest.fixture def pre_init_parser(self, pre_init_tokenizer): - return Gemma4Parser(pre_init_tokenizer) + return Gemma4Parser( + pre_init_tokenizer, chat_template_kwargs={"enable_thinking": True} + ) def test_model_emitted_channel_open_after_new_turn( self, pre_init_parser, pre_init_tokenizer, request_obj @@ -489,7 +494,9 @@ def plain_tokenizer(self): @pytest.fixture def plain_parser(self, plain_tokenizer): - return Gemma4Parser(plain_tokenizer) + return Gemma4Parser( + plain_tokenizer, chat_template_kwargs={"enable_thinking": True} + ) def test_streaming_channel_less_output_is_content( self, plain_parser, plain_tokenizer, request_obj @@ -521,7 +528,9 @@ def test_streaming_matches_non_streaming( assert reasoning is None assert content == _PLAIN_ANSWER_TEXT - stream_parser = Gemma4Parser(plain_tokenizer) + stream_parser = Gemma4Parser( + plain_tokenizer, chat_template_kwargs={"enable_thinking": True} + ) results = _stream_tokens_batched( stream_parser, plain_tokenizer, diff --git a/tests/reasoning/test_gemma4_reasoning_parser.py b/tests/reasoning/test_gemma4_reasoning_parser.py index b92d84b195c8..825cc7d0f707 100644 --- a/tests/reasoning/test_gemma4_reasoning_parser.py +++ b/tests/reasoning/test_gemma4_reasoning_parser.py @@ -262,14 +262,27 @@ def test_gemma4_adjust_request(generic_tokenizer): assert result is request -def test_gemma4_previous_turn_reasoning_is_reasoning_end(generic_tokenizer): +@pytest.mark.parametrize( + ("chat_template_kwargs", "expected_is_reasoning_end"), + [ + pytest.param({}, True, id="omitted_enable_thinking"), + pytest.param({"enable_thinking": False}, True, id="thinking_disabled"), + pytest.param({"enable_thinking": True}, False, id="thinking_enabled"), + ], +) +def test_gemma4_new_turn_reasoning_end_matches_enable_thinking( + generic_tokenizer, + chat_template_kwargs, + expected_is_reasoning_end, +): output = ( "<|channel>thought\n1st thought1st content\n" "<|turn>user\nThanks<|turn>model\n" ) output_tokens = gemma4_encode_output(generic_tokenizer, output) parser: ReasoningParser = ReasoningParserManager.get_reasoning_parser(parser_name)( - generic_tokenizer + generic_tokenizer, + chat_template_kwargs=chat_template_kwargs, ) is_reasoning_end = parser.is_reasoning_end(output_tokens) - assert not is_reasoning_end + assert is_reasoning_end is expected_is_reasoning_end diff --git a/vllm/parser/gemma4.py b/vllm/parser/gemma4.py index 801bb190bcda..47a97cc49037 100644 --- a/vllm/parser/gemma4.py +++ b/vllm/parser/gemma4.py @@ -400,7 +400,7 @@ def __init__( **kwargs, ) -> None: chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {} - self._thinking_enabled = chat_kwargs.get("enable_thinking", True) + self._thinking_enabled = chat_kwargs.get("enable_thinking", False) super().__init__( tokenizer, tools, From 689be2bcd37a3a862d1330c2aa91727406e64085 Mon Sep 17 00:00:00 2001 From: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Date: Tue, 18 Aug 2026 11:04:36 -0400 Subject: [PATCH 090/839] Upgrade Flashinfer version to 0.6.17 (#52681) Signed-off-by: wzhao18 --- docker/Dockerfile | 2 +- docker/versions.json | 2 +- requirements/cuda.txt | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/docker/Dockerfile b/docker/Dockerfile index b2c949d2d155..550cc6f4dca7 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -758,7 +758,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \ # Install FlashInfer JIT cache (requires CUDA-version-specific index URL) # https://docs.flashinfer.ai/installation.html # From versions.json: .flashinfer.version -ARG FLASHINFER_VERSION=0.6.16.post3 +ARG FLASHINFER_VERSION=0.6.17 RUN --mount=type=cache,target=/opt/uv/cache \ uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \ --index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') diff --git a/docker/versions.json b/docker/versions.json index ea51010c90b2..7ac2bada5e44 100644 --- a/docker/versions.json +++ b/docker/versions.json @@ -68,7 +68,7 @@ "default": "true" }, "FLASHINFER_VERSION": { - "default": "0.6.16.post3" + "default": "0.6.17" }, "GDRCOPY_CUDA_VERSION": { "default": "12.8" diff --git a/requirements/cuda.txt b/requirements/cuda.txt index 3aa9016fd806..530c223bea53 100644 --- a/requirements/cuda.txt +++ b/requirements/cuda.txt @@ -14,8 +14,8 @@ PyNvVideoCodec==2.0.4 # flashinfer-cubin is not on PyPI since 0.6.14; setup.py excludes it from # install_requires so the published wheel does not carry an unresolvable pin --extra-index-url https://flashinfer.ai/whl/ -flashinfer-python==0.6.16.post3 -flashinfer-cubin==0.6.16.post3 +flashinfer-python==0.6.17 +flashinfer-cubin==0.6.17 apache-tvm-ffi==0.1.11 tilelang==0.1.12 nvidia-cudnn-frontend>=1.19.1 From 3bb9c18f0c845170679498f6aa39c6a92f2f68cb Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Tue, 18 Aug 2026 23:06:18 +0800 Subject: [PATCH 091/839] [Multimodal] Reorganize video decoder backends (#49155) Signed-off-by: Isotr0py --- tests/multimodal/media/test_video.py | 10 +- tests/multimodal/test_gpu_ipc_memory.py | 35 +- tests/multimodal/test_video.py | 173 ++- vllm/multimodal/gpu_ipc_memory.py | 8 +- vllm/multimodal/video.py | 1111 +---------------- vllm/multimodal/video_decoders/__init__.py | 110 ++ vllm/multimodal/video_decoders/base.py | 37 + vllm/multimodal/video_decoders/deepstream.py | 151 +++ vllm/multimodal/video_decoders/opencv.py | 295 +++++ vllm/multimodal/video_decoders/pyav.py | 111 ++ .../video_decoders/pynvvideocodec.py | 372 ++++++ vllm/multimodal/video_decoders/torchcodec.py | 100 ++ 12 files changed, 1416 insertions(+), 1097 deletions(-) create mode 100644 vllm/multimodal/video_decoders/__init__.py create mode 100644 vllm/multimodal/video_decoders/base.py create mode 100644 vllm/multimodal/video_decoders/deepstream.py create mode 100644 vllm/multimodal/video_decoders/opencv.py create mode 100644 vllm/multimodal/video_decoders/pyav.py create mode 100644 vllm/multimodal/video_decoders/pynvvideocodec.py create mode 100644 vllm/multimodal/video_decoders/torchcodec.py diff --git a/tests/multimodal/media/test_video.py b/tests/multimodal/media/test_video.py index 186678f3c19a..d0402a2dfa29 100644 --- a/tests/multimodal/media/test_video.py +++ b/tests/multimodal/media/test_video.py @@ -21,8 +21,10 @@ from vllm.multimodal.video import ( PYNVVIDEOCODEC_VIDEO_BACKEND, VIDEO_LOADER_REGISTRY, - PyNvVideoCodecVideoBackend, VideoLoader, +) +from vllm.multimodal.video_decoders.pynvvideocodec import ( + PyNvVideoCodecVideoBackendMixin, _pynvvc_frames_to_nhwc, ) @@ -381,13 +383,15 @@ def raise_unrelated_error(cls, file_path, nvc): raise original_error monkeypatch.setattr( - PyNvVideoCodecVideoBackend, + PyNvVideoCodecVideoBackendMixin, "_read_source_metadata", classmethod(raise_unrelated_error), ) with pytest.raises(RuntimeError) as exc_info: - PyNvVideoCodecVideoBackend.decode_frames_pynvvideocodec(b"video", None) + PyNvVideoCodecVideoBackendMixin.decode_frames_pynvvideocodec( + None, b"video", None + ) assert exc_info.value is original_error diff --git a/tests/multimodal/test_gpu_ipc_memory.py b/tests/multimodal/test_gpu_ipc_memory.py index c98fd7673cbf..9afbb3a1c76b 100644 --- a/tests/multimodal/test_gpu_ipc_memory.py +++ b/tests/multimodal/test_gpu_ipc_memory.py @@ -6,7 +6,6 @@ import pytest -import vllm.config.multimodal as multimodal_config_module from vllm.config.multimodal import MultiModalConfig from vllm.multimodal.gpu_ipc_memory import ( MultiModalGPUMemoryPool, @@ -15,10 +14,10 @@ reserve_mm_ipc_gpu_memory, set_mm_gpu_ipc_pool, ) -from vllm.multimodal.video import ( +from vllm.multimodal.video_decoders import PYNVVIDEOCODEC_VIDEO_BACKEND +from vllm.multimodal.video_decoders.pynvvideocodec import ( PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES, PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES, - PYNVVIDEOCODEC_VIDEO_BACKEND, ) from vllm.utils.mem_constants import GiB_bytes @@ -186,11 +185,7 @@ def test_reserve_mm_ipc_gpu_memory_raw_frame_budget_only( monkeypatch: pytest.MonkeyPatch, video_backend: str | None, ): - monkeypatch.setattr( - multimodal_config_module.envs, - "VLLM_VIDEO_LOADER_BACKEND", - "opencv", - ) + monkeypatch.setenv("VLLM_VIDEO_LOADER_BACKEND", "opencv") mm_config = _mm_config( mm_ipc_gpu_memory_gb=0.25, video_backend=video_backend, @@ -202,11 +197,7 @@ def test_reserve_mm_ipc_gpu_memory_raw_frame_budget_only( def test_reserve_mm_ipc_gpu_memory_includes_pynvvideocodec_decoder_budget( monkeypatch: pytest.MonkeyPatch, ): - monkeypatch.setattr( - multimodal_config_module.envs, - "VLLM_VIDEO_LOADER_BACKEND", - "opencv", - ) + monkeypatch.setenv("VLLM_VIDEO_LOADER_BACKEND", "opencv") mm_config = _mm_config( mm_ipc_gpu_memory_gb=0.25, video_backend=PYNVVIDEOCODEC_VIDEO_BACKEND, @@ -221,11 +212,7 @@ def test_reserve_mm_ipc_gpu_memory_includes_pynvvideocodec_decoder_budget( def test_reserve_mm_ipc_gpu_memory_uses_env_video_backend( monkeypatch: pytest.MonkeyPatch, ): - monkeypatch.setattr( - multimodal_config_module.envs, - "VLLM_VIDEO_LOADER_BACKEND", - PYNVVIDEOCODEC_VIDEO_BACKEND, - ) + monkeypatch.setenv("VLLM_VIDEO_LOADER_BACKEND", PYNVVIDEOCODEC_VIDEO_BACKEND) available_bytes = 4 * GiB_bytes assert reserve_mm_ipc_gpu_memory(available_bytes, _mm_config()) == ( @@ -236,11 +223,7 @@ def test_reserve_mm_ipc_gpu_memory_uses_env_video_backend( def test_reserve_mm_ipc_gpu_memory_scales_decoder_budget_by_api_servers( monkeypatch: pytest.MonkeyPatch, ): - monkeypatch.setattr( - multimodal_config_module.envs, - "VLLM_VIDEO_LOADER_BACKEND", - PYNVVIDEOCODEC_VIDEO_BACKEND, - ) + monkeypatch.setenv("VLLM_VIDEO_LOADER_BACKEND", PYNVVIDEOCODEC_VIDEO_BACKEND) available_bytes = 8 * GiB_bytes assert reserve_mm_ipc_gpu_memory( @@ -253,11 +236,7 @@ def test_reserve_mm_ipc_gpu_memory_scales_decoder_budget_by_api_servers( def test_reserve_mm_ipc_gpu_memory_uses_configured_hw_decoders( monkeypatch: pytest.MonkeyPatch, ): - monkeypatch.setattr( - multimodal_config_module.envs, - "VLLM_VIDEO_LOADER_BACKEND", - "opencv", - ) + monkeypatch.setenv("VLLM_VIDEO_LOADER_BACKEND", "opencv") available_bytes = 4 * GiB_bytes mm_config = _mm_config( video_backend=PYNVVIDEOCODEC_VIDEO_BACKEND, diff --git a/tests/multimodal/test_video.py b/tests/multimodal/test_video.py index 17d29f850ea4..9ab718d9457d 100644 --- a/tests/multimodal/test_video.py +++ b/tests/multimodal/test_video.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import itertools +import subprocess import sys import threading from contextlib import ExitStack, contextmanager @@ -16,24 +17,26 @@ from vllm.assets.base import get_vllm_public_assets from vllm.models.minimax_m3.common.mm_preprocess import MiniMaxM3VideoBackend from vllm.multimodal.video import ( - PYNVVIDEOCODEC_DECODER_CACHE_SIZE, PYNVVIDEOCODEC_VIDEO_BACKEND, VIDEO_LOADER_REGISTRY, DynamicVideoBackend, GLM46VVideoBackend, Molmo2VideoBackend, - PyNvVideoCodecDecoderSlot, - PyNvVideoCodecVideoBackend, - PyNvVideoCodecVideoBackendMixin, Qwen2VLVideoBackend, Qwen3VLVideoBackend, VideoBackend, VideoLoader, VideoSourceMetadata, VideoTargetMetadata, - _pynv_decoder_pool, get_video_loader_backend_for_processor, ) +from vllm.multimodal.video_decoders import decode_video, resolve_video_backend_kwargs +from vllm.multimodal.video_decoders.pynvvideocodec import ( + PYNVVIDEOCODEC_DECODER_CACHE_SIZE, + PyNvVideoCodecDecoderSlot, + PyNvVideoCodecVideoBackendMixin, + _pynv_decoder_pool, +) from vllm.platforms import current_platform from vllm.transformers_utils.processor import get_video_processor_cls_name_from_config @@ -99,6 +102,126 @@ def test_video_loader_type_doesnt_exist(): VIDEO_LOADER_REGISTRY.load("non_existing_video_loader") +def test_video_decoder_backends_are_lazy_imported(): + code = """ +import sys +import vllm.multimodal.video # noqa: F401 + +backend_modules = { + f"vllm.multimodal.video_decoders.{backend}" + for backend in ("opencv", "pyav", "torchcodec", "pynvvideocodec", "deepstream") +} +loaded = sorted(backend_modules & sys.modules.keys()) +assert not loaded, loaded +""" + result = subprocess.run( + [sys.executable, "-c", code], + capture_output=True, + text=True, + ) + assert result.returncode == 0, result.stderr + + +@pytest.mark.parametrize( + "backend", + ["opencv", "pyav", "torchcodec", "pynvvideocodec", "deepstream"], +) +def test_decode_video_imports_only_selected_backend( + backend: str, + monkeypatch: pytest.MonkeyPatch, +): + imports = [] + decoded = object() + + def fake_decoder(*args, **kwargs): + return decoded + + class FakeBackendModule: + pass + + setattr(FakeBackendModule, f"decode_{backend}", fake_decoder) + + def fake_import_module(name: str, package: str): + imports.append((name, package)) + return FakeBackendModule + + monkeypatch.setattr( + "vllm.multimodal.video_decoders.import_module", fake_import_module + ) + result = decode_video( + backend, + loader_cls=None, + data=b"", + target=VideoTargetMetadata(-1, -1, 300), + sampling_kwargs={}, + backend_kwargs={}, + frame_recovery=False, + ) + + assert result is decoded + assert imports == [(f".{backend}", "vllm.multimodal.video_decoders")] + + +@pytest.mark.parametrize( + ("backend", "kwargs", "expected_sampling", "expected_backend"), + [ + ( + "torchcodec", + {"min_frames": 4, "num_ffmpeg_threads": 2, "seek_mode": "approximate"}, + {"min_frames": 4}, + {"num_ffmpeg_threads": 2, "seek_mode": "approximate"}, + ), + ( + "deepstream", + {"max_frames": 16, "pool_size": 3, "timeout_sec": 10.0}, + {"max_frames": 16}, + {"pool_size": 3, "timeout_sec": 10.0}, + ), + ], +) +def test_video_backend_kwargs_are_separated_from_sampling_kwargs( + backend: str, + kwargs: dict, + expected_sampling: dict, + expected_backend: dict, +): + original_kwargs = dict(kwargs) + sampling_kwargs, backend_kwargs = resolve_video_backend_kwargs(backend, kwargs) + + assert sampling_kwargs == expected_sampling + assert backend_kwargs == expected_backend + assert kwargs == original_kwargs + + +def test_video_backend_rejects_options_for_another_decoder(): + with pytest.raises( + ValueError, match="num_ffmpeg_threads is not supported by the 'pyav' backend" + ): + resolve_video_backend_kwargs("pyav", {"num_ffmpeg_threads": 2}) + + +@pytest.mark.parametrize( + ("backend", "error"), + [ + ("pyav", AssertionError), + (PYNVVIDEOCODEC_VIDEO_BACKEND, ValueError), + ], +) +def test_video_decoder_spec_validates_frame_recovery( + backend: str, error: type[Exception] +): + with pytest.raises(error, match="frame_recovery is not supported"): + decode_video( + backend, + loader_cls=None, + data=b"", + target=VideoTargetMetadata(-1, -1, 300), + sampling_kwargs={}, + backend_kwargs={}, + frame_recovery=True, + ) + + @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") def test_pynvvideocodec_backend_accounts_raw_decoded_frames( monkeypatch: pytest.MonkeyPatch, @@ -145,7 +268,9 @@ def fake_decode(cls, file_path: str, frame_idx: list[int], nvc): "vllm.multimodal.gpu_ipc_memory.get_mm_gpu_ipc_pool", lambda: pool ) monkeypatch.setattr( - PyNvVideoCodecVideoBackend, "_decode_to_pinned_host", classmethod(fake_decode) + PyNvVideoCodecVideoBackendMixin, + "_decode_to_pinned_host", + classmethod(fake_decode), ) loader = VIDEO_LOADER_REGISTRY.load(PYNVVIDEOCODEC_VIDEO_BACKEND) @@ -205,7 +330,9 @@ def fake_decode(cls, file_path: str, frame_idx: list[int], nvc): "vllm.multimodal.gpu_ipc_memory.get_mm_gpu_ipc_pool", lambda: pool ) monkeypatch.setattr( - DynamicVideoBackend, "_decode_to_pinned_host", classmethod(fake_decode) + PyNvVideoCodecVideoBackendMixin, + "_decode_to_pinned_host", + classmethod(fake_decode), ) loader = VIDEO_LOADER_REGISTRY.load("opencv_dynamic") @@ -273,7 +400,7 @@ class FakeSlot: create_count = 0 with _fresh_decoder_pool(): - PyNvVideoCodecVideoBackend._configure_decoder_slots(hw_decoders) + PyNvVideoCodecVideoBackendMixin._configure_decoder_slots(hw_decoders) def fake_create_slot(cls): nonlocal create_count @@ -281,7 +408,7 @@ def fake_create_slot(cls): return FakeSlot() monkeypatch.setattr( - PyNvVideoCodecVideoBackend, + PyNvVideoCodecVideoBackendMixin, "_create_decoder_slot", classmethod(fake_create_slot), ) @@ -291,12 +418,16 @@ def fake_create_slot(cls): with ExitStack() as stack: retained_slots = [ - stack.enter_context(PyNvVideoCodecVideoBackend._borrow_decoder_slot()) + stack.enter_context( + PyNvVideoCodecVideoBackendMixin._borrow_decoder_slot() + ) for _ in range(hw_decoders) ] def borrow_extra_slot(): - with PyNvVideoCodecVideoBackend._borrow_decoder_slot() as extra_slot: + with ( + PyNvVideoCodecVideoBackendMixin._borrow_decoder_slot() + ) as extra_slot: seen_slots.append(extra_slot) borrowed.set() @@ -317,11 +448,11 @@ def test_pynvvideocodec_decoder_slots_are_configured_once( ): monkeypatch.setattr(_pynv_decoder_pool, "max_slots", None) - PyNvVideoCodecVideoBackend._configure_decoder_slots(2) - PyNvVideoCodecVideoBackend._configure_decoder_slots(2) + PyNvVideoCodecVideoBackendMixin._configure_decoder_slots(2) + PyNvVideoCodecVideoBackendMixin._configure_decoder_slots(2) with pytest.raises(RuntimeError, match="already configured as 2, got 3"): - PyNvVideoCodecVideoBackend._configure_decoder_slots(3) + PyNvVideoCodecVideoBackendMixin._configure_decoder_slots(3) def test_pynvvideocodec_failed_rebuild_invalidates_decoder_slot(): @@ -367,7 +498,7 @@ def SimpleDecoder(file_path: str, **kwargs): with ( pytest.raises(RuntimeError, match="construct failed"), - PyNvVideoCodecVideoBackend._borrow_decoder_slot() as borrowed, + PyNvVideoCodecVideoBackendMixin._borrow_decoder_slot() as borrowed, ): assert borrowed is slot borrowed.get_decoder( @@ -456,6 +587,12 @@ def test_pynvvideocodec_cross_subclass_shares_single_pool(): counters via ClassVar shadowing. """ + class MixinSubclassA(PyNvVideoCodecVideoBackendMixin): + pass + + class MixinSubclassB(PyNvVideoCodecVideoBackendMixin): + pass + class FakeSlot: pass @@ -475,8 +612,8 @@ def fake_create_slot(cls): ) try: with ExitStack() as stack: - stack.enter_context(VideoBackend._borrow_decoder_slot()) - stack.enter_context(Qwen3VLVideoBackend._borrow_decoder_slot()) + stack.enter_context(MixinSubclassA._borrow_decoder_slot()) + stack.enter_context(MixinSubclassB._borrow_decoder_slot()) assert pool.active == 2 blocked = threading.Event() @@ -484,7 +621,7 @@ def fake_create_slot(cls): def try_borrow(): blocked.set() - with Qwen2VLVideoBackend._borrow_decoder_slot(): + with MixinSubclassB._borrow_decoder_slot(): acquired.set() t = threading.Thread(target=try_borrow) diff --git a/vllm/multimodal/gpu_ipc_memory.py b/vllm/multimodal/gpu_ipc_memory.py index 58cb98888021..58f807303700 100644 --- a/vllm/multimodal/gpu_ipc_memory.py +++ b/vllm/multimodal/gpu_ipc_memory.py @@ -188,11 +188,13 @@ def reserve_mm_ipc_gpu_memory( return available_kv_cache_memory_bytes from vllm import envs - from vllm.multimodal.video import ( - PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES, - PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES, + from vllm.multimodal.video_decoders import ( PYNVVIDEOCODEC_DEFAULT_HW_DECODERS, PYNVVIDEOCODEC_VIDEO_BACKEND, + ) + from vllm.multimodal.video_decoders.pynvvideocodec import ( + PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES, + PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES, validate_pynvvideocodec_hw_decoders, ) diff --git a/vllm/multimodal/video.py b/vllm/multimodal/video.py index c50ce7761438..4f750c5cfa1b 100644 --- a/vllm/multimodal/video.py +++ b/vllm/multimodal/video.py @@ -1,42 +1,29 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import math -import os -import tempfile -import threading from abc import abstractmethod -from contextlib import contextmanager, suppress -from io import BytesIO -from typing import Any, ClassVar, Literal, NamedTuple, cast +from typing import Any, ClassVar, Literal, cast import numpy as np import numpy.typing as npt import torch -from vllm import envs from vllm.logger import init_logger -from vllm.utils.import_utils import PlaceholderModule, check_torchcodec_available -from vllm.utils.mem_constants import MiB_bytes +from vllm.multimodal.video_decoders import ( + PYNVVIDEOCODEC_VIDEO_BACKEND, + VideoDecoderBackend, + VideoSourceMetadata, + VideoTargetMetadata, + decode_video, + resolve_video_backend_kwargs, +) +from vllm.utils.import_utils import PlaceholderModule from vllm.utils.registry import ExtensionManager try: import cv2 - import cv2.videoio_registry as vr except ImportError: cv2 = PlaceholderModule("cv2") - vr = PlaceholderModule("cv2").placeholder_attr("videoio_registry") - -try: - import av -except ImportError: - av = PlaceholderModule("av") # type: ignore[assignment] - -try: - from torchcodec.decoders import VideoDecoder -except (ImportError, RuntimeError): - VideoDecoder = PlaceholderModule("torchcodec").placeholder_attr( # type: ignore[assignment] - "decoders.VideoDecoder" - ) logger = init_logger(__name__) @@ -106,18 +93,6 @@ def get_video_loader_backend_for_processor( return VIDEO_LOADER_REGISTRY.get_backend_for_video_processor(video_processor) -def _check_frame_pixel_limit(width: int, height: int) -> None: - """Reject video frames exceeding VLLM_MAX_IMAGE_PIXELS before decoding.""" - max_pixels = envs.VLLM_MAX_IMAGE_PIXELS - if max_pixels > 0 and width * height > max_pixels: - raise ValueError( - f"Video frame dimensions {width}x{height} " - f"({width * height} pixels) exceed the maximum of " - f"{max_pixels} pixels. Set VLLM_MAX_IMAGE_PIXELS to " - f"increase this limit." - ) - - def resize_video(frames: npt.NDArray, size: tuple[int, int]) -> npt.NDArray: num_frames, _, _, channels = frames.shape new_height, new_width = size @@ -149,30 +124,6 @@ def sample_frames_from_video(frames: npt.NDArray, num_frames: int) -> npt.NDArra return sampled_frames -class VideoTargetMetadata(NamedTuple): - """Metadata represents target video.""" - - num_frames: int - fps: float - max_duration: float - - -class VideoSourceMetadata(NamedTuple): - """Metadata represents source video.""" - - total_frames_num: int - original_fps: float - duration: float - - -class PyNvVideoCodecSourceMetadata(NamedTuple): - """Metadata needed before GPU video decode.""" - - source: VideoSourceMetadata - width: int - height: int - - class VideoLoader: @classmethod def compute_frames_index_to_sample( @@ -184,6 +135,11 @@ def compute_frames_index_to_sample( """Return the list of frame indices to sample from the video.""" raise NotImplementedError + @classmethod + def _prepare_source(cls, source: VideoSourceMetadata) -> VideoSourceMetadata: + """Sampling-algorithm-specific metadata adjustment hook.""" + return source + @classmethod @abstractmethod def load_bytes( @@ -214,835 +170,9 @@ def create_hf_metadata( VIDEO_LOADER_REGISTRY = VideoLoaderRegistry() VIDEO_LOADER_REGISTRY.register_gpu_codec("deepstream") -PYNVVIDEOCODEC_VIDEO_BACKEND: Literal["pynvvideocodec"] = "pynvvideocodec" -# Per-decoder upper bound reserved for persistent PyNvVideoCodec surfaces. -PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES = 128 * MiB_bytes -PYNVVIDEOCODEC_DECODER_CACHE_SIZE = 2 -PYNVVIDEOCODEC_DEFAULT_HW_DECODERS = 2 -# Per-API-server CUDA context and driver allocation, measured with -# PyNvVideoCodec 2.0.4 on H100. -PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES = int(1.8 * 1024 * MiB_bytes) - - -def validate_pynvvideocodec_hw_decoders(hw_decoders: object) -> int: - if ( - isinstance(hw_decoders, bool) - or not isinstance(hw_decoders, int) - or hw_decoders < 1 - ): - raise ValueError("hw_decoders must be a positive integer") - return hw_decoders - - -def _pynvvideocodec_exception_types(nvc) -> tuple[type[Exception], ...]: - return tuple( - exception_type - for name in dir(nvc) - if name.startswith("PyNvVCException") - and isinstance((exception_type := getattr(nvc, name)), type) - and issubclass(exception_type, Exception) - ) - - -class PyNvVideoCodecDecoderSlot: - """A retained PyNv decoder slot and its CUDA stream. - - The decoder is reused across requests: ``reconfigure_decoder`` repoints the - existing decoder at each new source instead of paying a fresh - ``SimpleDecoder`` construction per request. Construction (CUVID parser + - decoder + surface-pool allocation) is the dominant per-request cost, so - reconfiguring is far cheaper. A single decoder serves both metadata - (``len``/``get_stream_metadata``) and frame decode -- no separate - metadata decoder. - """ - - def __init__(self, stream) -> None: - self.stream = stream - self.decoder = None - self.source_path: str | None = None - - def invalidate(self) -> None: - self.decoder = None - self.source_path = None - - def _construct(self, file_path: str, nvc, device_index: int) -> None: - self.invalidate() - decoder = nvc.SimpleDecoder( - file_path, - output_color_type=nvc.OutputColorType.RGB, - use_device_memory=True, - need_scanned_stream_metadata=True, - gpu_id=device_index, - cuda_stream=self.stream.cuda_stream, - decoder_cache_size=PYNVVIDEOCODEC_DECODER_CACHE_SIZE, - ) - self.decoder = decoder - self.source_path = file_path - - def get_decoder(self, file_path: str, nvc, device_index: int): - if self.decoder is None: - self._construct(file_path, nvc, device_index) - elif self.source_path != file_path: - try: - self.decoder.reconfigure_decoder(file_path) - self.source_path = file_path - except Exception: - # reconfigure unsupported/unsafe for this source -> rebuild. - self._construct(file_path, nvc, device_index) - return self.decoder - - -class OpenCVVideoBackendMixin: - @staticmethod - def get_cv2_video_api(): - api_pref = None - for backend in vr.getStreamBufferedBackends(): - if not vr.hasBackend(backend): - continue - if not vr.isBackendBuiltIn(backend): - _, abi, api = vr.getStreamBufferedBackendPluginVersion(backend) - if abi < 1 or (abi == 1 and api < 2): - continue - api_pref = backend - break - return api_pref - - @classmethod - def open_video_capture(cls, data: bytes) -> "cv2.VideoCapture": - backend = cls.get_cv2_video_api() - cap = cv2.VideoCapture(BytesIO(data), backend, []) - if not cap.isOpened(): - raise ValueError("Could not open video stream") - return cap - - @staticmethod - def get_video_metadata(cap: "cv2.VideoCapture") -> VideoSourceMetadata: - total_frames_num = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) - original_fps = cap.get(cv2.CAP_PROP_FPS) - duration = total_frames_num / original_fps if original_fps > 0 else 0 - return VideoSourceMetadata( - total_frames_num=total_frames_num, - original_fps=original_fps, - duration=duration, - ) - - @classmethod - def _can_use_for_recovery( - cls, - idx: int, - failed_frames: list[int], - next_target_map: dict[int, int], - total_frames: int, - ) -> bool: - """Check if current frame can recover the oldest failed frame.""" - if not failed_frames: - return False - oldest_failed = failed_frames[0] - limit = next_target_map.get(oldest_failed, total_frames) - return idx < limit - - @classmethod - def _read_frames_with_recovery( - cls, - cap: "cv2.VideoCapture", - frame_indices: list[int], - total_frames: int, - ) -> tuple[npt.NDArray, list[int], dict[int, int]]: - """ - Read frames with dynamic window forward-scan recovery. - - When a target frame fails to load, the next successfully grabbed - frame (before the next target frame) will be used to recover it. - - Args: - cap: OpenCV VideoCapture object - frame_indices: Sorted list of target frame indices to load - total_frames: Total number of frames in the video - - Returns: - Tuple of (frames_array, valid_frame_indices, recovered_map) - - frames_array: Array of loaded frames - - valid_frame_indices: List of frame indices that were loaded - - recovered_map: Dict mapping recovered_idx -> source_idx - """ - width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) - height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) - - assert width > 0 and height > 0, ( - f"Invalid video frame size: width={width}, height={height}" - ) - - frame_idx_set = set(frame_indices) - max_frame_idx = frame_indices[-1] if frame_indices else 0 - - # Build map: target_idx -> next_target_idx (for recovery window) - next_target_map: dict[int, int] = {} - for k in range(len(frame_indices) - 1): - next_target_map[frame_indices[k]] = frame_indices[k + 1] - next_target_map[frame_indices[-1]] = total_frames - - frames_list: list[npt.NDArray] = [] - valid_frame_indices: list[int] = [] - failed_frames_idx: list[int] = [] - recovered_map: dict[int, int] = {} - - i = 0 - for idx in range(max_frame_idx + 1): - is_target_frame = idx in frame_idx_set - - # Attempt to grab the current frame - ok = cap.grab() - - if not ok: - if is_target_frame: - logger.debug( - "Failed to grab frame %d during video loading.", - idx, - ) - failed_frames_idx.append(idx) - continue - - # Check if we should retrieve: target frame OR can recover a failed one - can_recover = cls._can_use_for_recovery( - idx, failed_frames_idx, next_target_map, total_frames - ) - - if is_target_frame or can_recover: - ret, frame = cap.retrieve() - - if ret and frame is not None and frame.size > 0: - rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) - frames_list.append(rgb_frame) - valid_frame_indices.append(idx) - i += 1 - - if can_recover: - recovered_idx = failed_frames_idx.pop(0) - recovered_map[recovered_idx] = idx - logger.info( - "Recovered frame %d using frame %d (delay: %d)", - recovered_idx, - idx, - idx - recovered_idx, - ) - elif is_target_frame: - logger.debug( - "Failed to retrieve frame %d during video loading.", - idx, - ) - failed_frames_idx.append(idx) - - # Log any remaining failed frames - for failed_idx in failed_frames_idx: - logger.debug( - "Frame %d could not be recovered (end of video).", - failed_idx, - ) - - # Stack frames - if frames_list: - frames = np.stack(frames_list) - else: - frames = np.empty((0, height, width, 3), dtype=np.uint8) - - return frames, valid_frame_indices, recovered_map - - @classmethod - def _read_frames_no_recovery( - cls, - cap, - frame_indices: set[int], - max_frame_idx: int, - ) -> tuple[npt.NDArray, list[int]]: - num_expected_frames = len(frame_indices) - width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) - height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) - frames = np.empty((num_expected_frames, height, width, 3), dtype=np.uint8) - - i = 0 - valid_frame_indices = [] - for idx in range(max_frame_idx + 1): - ok = cap.grab() - if not ok: - # Frame is broken/unreadable, skip it - if idx in frame_indices: - logger.debug( - "Failed to grab frame %d during video loading. " - "This frame will be skipped.", - idx, - ) - continue - if idx in frame_indices: - ret, frame = cap.retrieve() - if ret: - frames[i] = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) - valid_frame_indices.append(idx) - i += 1 - else: - # retrieve() failed even though grab() succeeded - logger.debug( - "Failed to retrieve frame %d during video loading. " - "This frame will be skipped.", - idx, - ) - - valid_num_frames = len(valid_frame_indices) - if valid_num_frames < num_expected_frames: - logger.warning( - "Video loading completed with %d broken/unreadable frames. " - "Expected %d frames but only loaded %d frames.", - num_expected_frames - valid_num_frames, - num_expected_frames, - valid_num_frames, - ) - - return frames[:valid_num_frames], valid_frame_indices - - @classmethod - def read_frames( - cls, - cap: "cv2.VideoCapture", - frame_idx: list[int], - total_frames_num: int, - *, - frame_recovery: bool = False, - ) -> tuple[npt.NDArray, list[int]]: - if frame_recovery: - num_frames_to_sample = len(frame_idx) - frames, valid_frame_indices, recovered_map = cls._read_frames_with_recovery( - cap, frame_idx, total_frames_num - ) - - if recovered_map: - logger.info( - "Frame recovery: %d frames recovered using forward scan.", - len(recovered_map), - ) - else: - frame_idx_set = set(frame_idx) - num_frames_to_sample = len(frame_idx_set) - frames, valid_frame_indices = cls._read_frames_no_recovery( - cap, frame_idx_set, max(frame_idx) - ) - valid_num_frames = len(valid_frame_indices) - if valid_num_frames < num_frames_to_sample: - logger.warning( - "Video loading completed with %d broken/unreadable frames. " - "Expected to sample %d frames but only loaded %d frames.", - num_frames_to_sample - valid_num_frames, - num_frames_to_sample, - valid_num_frames, - ) - return frames, valid_frame_indices - - -class PyAVVideoBackendMixin: - """PyAV (in-process FFmpeg bindings) codec utilities. - - Reads stream metadata and decodes target frames via per-frame - ``container.seek()``. The seek releases the GIL between frames and - scales with the number of sampled frames rather than the video - length, enabling concurrent decoding under serving load. - """ - - @staticmethod - def get_metadata( - container: "av.container.InputContainer", - ) -> VideoSourceMetadata: - if not container.streams.video: - raise ValueError("No video streams found in container") - stream = container.streams.video[0] - total_frames = stream.frames or 0 - fps = float(stream.average_rate) if stream.average_rate else 0.0 - duration = float(stream.duration * stream.time_base) if stream.duration else 0.0 - if total_frames == 0 and duration > 0 and fps > 0: - total_frames = int(duration * fps) - return VideoSourceMetadata(total_frames, fps, duration) - - @staticmethod - def decode_frames( - container: "av.container.InputContainer", - frame_indices: list[int], - fps: float, - duration: float, - ) -> tuple[npt.NDArray, list[int]]: - """Decode target frames via per-frame seek + forward decode to PTS.""" - stream = container.streams.video[0] - # SLICE parallelizes within a single frame without the - # one-frame-per-thread latency penalty of FRAME threading. - stream.thread_type = "SLICE" - time_base = stream.time_base - - frames_list: list[npt.NDArray] = [] - valid_indices: list[int] = [] - frame_interval = 1.0 / fps if fps > 0 else 0.1 - max_ts = max(0.0, duration - frame_interval) if duration > 0 else float("inf") - - decoder = None - last_pts = None - for idx in frame_indices: - ts = min(idx / fps, max_ts) if fps > 0 else 0.0 - pts = int(ts / time_base) - # seek() snaps backward to a keyframe; reuse the running decoder - # while targets advance monotonically to avoid re-decoding the - # GOP prefix once per requested frame. - if decoder is None or last_pts is None or pts <= last_pts: - container.seek(pts, stream=stream) - decoder = container.decode(video=0) - chosen = None - for frame in decoder: - if frame.pts is not None and frame.pts >= pts: - chosen = frame - last_pts = frame.pts - break - if chosen is not None: - frames_list.append(chosen.to_ndarray(format="rgb24")) - valid_indices.append(idx) - else: - decoder = None - - if not frames_list: - return np.empty((0,), dtype=np.uint8), valid_indices - return np.stack(frames_list), valid_indices - - -class TorchCodecVideoBackendMixin: - """TorchCodec (FFmpeg-backed, PyTorch-native) codec utilities. - - Builds a :class:`~torchcodec.decoders.VideoDecoder` over the in-memory - bytes and extracts the sampled indices with a single batched - ``get_frames_at`` call, while releasing the GIL during decode. - """ - - @staticmethod - def make_torchcodec_decoder( - data: bytes, - *, - num_ffmpeg_threads: int = 0, - seek_mode: Literal["exact", "approximate"] = "exact", - ) -> "VideoDecoder": - # NHWC matches the (num_frames, H, W, 3) uint8 RGB layout the rest - # of the pipeline expects, avoiding a transpose. - return VideoDecoder( - data, - dimension_order="NHWC", - num_ffmpeg_threads=num_ffmpeg_threads, - seek_mode=seek_mode, - ) - - @staticmethod - def get_torchcodec_metadata(decoder: "VideoDecoder") -> VideoSourceMetadata: - md = decoder.metadata - total_frames = md.num_frames or 0 - fps = float(md.average_fps) if md.average_fps else 0.0 - duration = float(md.duration_seconds) if md.duration_seconds else 0.0 - if total_frames == 0 and duration > 0 and fps > 0: - total_frames = int(duration * fps) - return VideoSourceMetadata(total_frames, fps, duration) - - @staticmethod - def decode_torchcodec_frames( - decoder: "VideoDecoder", - frame_indices: list[int], - ) -> tuple[npt.NDArray, list[int]]: - """Decode the requested indices in one batched, index-exact call.""" - if not frame_indices: - return np.empty((0,), dtype=np.uint8), [] - # Note: torchcodec releases the GIL for the entire call - batch = decoder.get_frames_at(frame_indices) - return batch.data.numpy(), list(frame_indices) - - -def _pynvvc_frames_to_nhwc(frames: torch.Tensor) -> torch.Tensor: - """Return a stacked PyNvVideoCodec frame batch as contiguous NHWC. - - PyNvVideoCodec's per-frame layout has varied across versions (HWC vs CHW), - so detect the channel axis rather than assuming a fixed order. NHWC is the - layout the other video backends return and the HF video processors expect. - - Args: - frames: A ``(N, ?, ?, ?)`` uint8 tensor in either NHWC or NCHW order. - - Returns: - The same frames as a contiguous ``(N, H, W, C)`` tensor. - """ - if frames.shape[-1] != 3 and frames.shape[-3] == 3: - frames = frames.permute(0, 2, 3, 1) # NCHW -> NHWC - return frames.contiguous() - - -class _PyNvDecoderPool: - """Process-wide singleton managing PyNvVideoCodec decoder slot state. - - Prevents subclass counter shadowing (GHSA-j682-9xp5-rrf3) by storing - all mutable pool state in a single module-level instance rather than - in ClassVar attributes that get shadowed by Python's augmented - assignment semantics on subclasses. - """ - - def __init__(self) -> None: - self.slots: list[PyNvVideoCodecDecoderSlot] = [] - self.active: int = 0 - self.cond: threading.Condition = threading.Condition() - self.max_slots: int | None = None - - def configure(self, hw_decoders: int) -> None: - with self.cond: - if self.max_slots is None: - self.max_slots = hw_decoders - elif self.max_slots != hw_decoders: - raise RuntimeError( - "PyNvVideoCodec decoder count is already configured as " - f"{self.max_slots}, got {hw_decoders}" - ) - - -_pynv_decoder_pool = _PyNvDecoderPool() - - -class PyNvVideoCodecVideoBackendMixin: - """PyNvVideoCodec utilities for GPU-backed frame decode.""" - - _DEVICE_INDEX: ClassVar[int] = 0 - - @classmethod - @abstractmethod - def compute_frames_index_to_sample( - cls, - source: VideoSourceMetadata, - target: VideoTargetMetadata, - **kwargs, - ) -> list[int]: - raise NotImplementedError - - @classmethod - @abstractmethod - def _prepare_source(cls, source: VideoSourceMetadata) -> VideoSourceMetadata: - raise NotImplementedError - - @classmethod - def _create_decoder_slot(cls) -> PyNvVideoCodecDecoderSlot: - import torch - - return PyNvVideoCodecDecoderSlot(torch.cuda.Stream(device=cls._DEVICE_INDEX)) - - @classmethod - def _configure_decoder_slots(cls, hw_decoders: object) -> None: - hw_decoders = validate_pynvvideocodec_hw_decoders(hw_decoders) - _pynv_decoder_pool.configure(hw_decoders) - - @staticmethod - @contextmanager - def _torch_stream_context(stream): - import torch - - torch.accelerator.set_device_index(stream.device.index) - previous_stream = torch.accelerator.current_stream() - torch.accelerator.set_stream(stream) - try: - yield - finally: - torch.accelerator.set_stream(previous_stream) - - @classmethod - @contextmanager - def _borrow_decoder_slot(cls): - pool = _pynv_decoder_pool - create_slot = False - with pool.cond: - if pool.max_slots is None: - raise RuntimeError("PyNvVideoCodec decoder slots are not configured") - while True: - if pool.slots: - slot = pool.slots.pop() - break - if pool.active < pool.max_slots: - pool.active += 1 - create_slot = True - break - pool.cond.wait() - - if create_slot: - try: - slot = cls._create_decoder_slot() - except Exception: - with pool.cond: - pool.active -= 1 - pool.cond.notify() - raise - - borrow_succeeded = False - try: - yield slot - borrow_succeeded = True - finally: - if not borrow_succeeded: - slot.invalidate() - with pool.cond: - pool.slots.append(slot) - pool.cond.notify() - - @staticmethod - def _metadata_value(metadata, *names: str, default=None): - for name in names: - value = getattr(metadata, name, None) - if value is not None: - return value - return default - - @classmethod - def _read_source_metadata( - cls, - file_path: str, - nvc, - ) -> PyNvVideoCodecSourceMetadata: - with cls._borrow_decoder_slot() as decoder_slot: - with cls._torch_stream_context(decoder_slot.stream): - decoder = decoder_slot.get_decoder( - file_path, nvc, device_index=cls._DEVICE_INDEX - ) - metadata = decoder.get_stream_metadata() - total_frames_num = len(decoder) - width = int(cls._metadata_value(metadata, "width", default=0)) - height = int(cls._metadata_value(metadata, "height", default=0)) - original_fps = float( - cls._metadata_value( - metadata, - "average_fps", - "avg_frame_rate", - "frame_rate", - "frameRate", - default=0.0, - ) - ) - duration = float( - cls._metadata_value(metadata, "duration", default=0.0) - or (total_frames_num / original_fps if original_fps > 0 else 0.0) - ) - if total_frames_num <= 0: - raise ValueError("Could not determine video frame count") - if width <= 0 or height <= 0: - raise ValueError("Could not determine video dimensions") - return PyNvVideoCodecSourceMetadata( - source=VideoSourceMetadata(total_frames_num, original_fps, duration), - width=width, - height=height, - ) - - @classmethod - def _decode_to_pinned_host( - cls, - file_path: str, - frame_idx: list[int], - nvc, - ) -> npt.NDArray: - import torch - - if not frame_idx: - return np.empty((0,), dtype=np.uint8) - - with cls._borrow_decoder_slot() as decoder_slot: - stream = decoder_slot.stream - with cls._torch_stream_context(stream): - try: - decoder = decoder_slot.get_decoder( - file_path, nvc, device_index=cls._DEVICE_INDEX - ) - decoded_frames = decoder.get_batch_frames_by_index(frame_idx) - except Exception as exc: - if not isinstance( - exc, - _pynvvideocodec_exception_types(nvc) + (IndexError,), - ): - raise - raise ValueError("Invalid or unsupported video file.") from exc - if len(decoded_frames) < len(frame_idx): - logger.warning( - "pynvvideocodec video loading: expected %d frames but got %d.", - len(frame_idx), - len(decoded_frames), - ) - torch_frames = [torch.from_dlpack(frame) for frame in decoded_frames] - if not torch_frames: - return np.empty((0,), dtype=np.uint8) - device_frames = torch.stack(torch_frames) - if device_frames.ndim != 4: - raise ValueError( - "PyNvVideoCodec returned frames with unexpected shape " - f"{tuple(device_frames.shape)}" - ) - device_frames = _pynvvc_frames_to_nhwc(device_frames) - host_frames = torch.empty( - device_frames.shape, - dtype=device_frames.dtype, - device="cpu", - pin_memory=True, - ) - host_frames.copy_(device_frames, non_blocking=True) - stream.synchronize() - host_array = host_frames.numpy() - del decoded_frames, torch_frames, device_frames - return host_array - - @classmethod - def decode_frames_pynvvideocodec( - cls, - data: bytes, - target: VideoTargetMetadata, - **kwargs, - ) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]: - import PyNvVideoCodec as nvc - - from vllm.multimodal.gpu_ipc_memory import get_mm_gpu_ipc_pool - - temp_fd, temp_path = tempfile.mkstemp(suffix=".mp4") - try: - with os.fdopen(temp_fd, "wb") as temp_file: - temp_file.write(data) - - try: - gpu_source = cls._read_source_metadata(temp_path, nvc) - except Exception as exc: - if not isinstance(exc, _pynvvideocodec_exception_types(nvc)): - raise - raise ValueError("Invalid or unsupported video file.") from exc - _check_frame_pixel_limit(gpu_source.width, gpu_source.height) - source = cls._prepare_source(gpu_source.source) - frame_idx = cls.compute_frames_index_to_sample( - source=source, target=target, **kwargs - ) - raw_frame_bytes = len(frame_idx) * gpu_source.height * gpu_source.width * 3 - pool = get_mm_gpu_ipc_pool() - if pool is None or raw_frame_bytes == 0: - frames = cls._decode_to_pinned_host(temp_path, frame_idx, nvc) - else: - with pool.acquire(raw_frame_bytes): - frames = cls._decode_to_pinned_host(temp_path, frame_idx, nvc) - finally: - with suppress(FileNotFoundError): - os.unlink(temp_path) - - valid_frame_indices = frame_idx[: int(frames.shape[0])] - return frames, source, frame_idx, valid_frame_indices - - -class DeepStreamVideoBackendMixin: - """NVIDIA DeepStream (NVDEC) GPU-decode codec utilities. - - Decoding runs on a shared pool of daemon threads inside one CUDA - context (see the ``nvidia-deepstream-videodecode-cu13`` package). The - container bytes are pushed into an ``appsrc`` GStreamer pipeline, so no - local file path is required — HTTP and base64 sources decode identically - to local files. - - Like the OpenCV/PyAV mixins, this provides only the codec layer. - Frame *selection* lives in the loader's - ``compute_frames_index_to_sample`` and arrives here as an explicit - list of frame indices. - """ - - # Process-wide lazy decode pool, shared across all DeepStream backends. - _pool: ClassVar[Any] = None - _pool_lock: ClassVar[Any] = None - - @classmethod - def _get_pool(cls, pool_size: int | None = None): - """Lazy-initialize the shared decode pool on first use. - - ``pool_size`` (number of decode worker threads) comes from - ``--media-io-kwargs`` (``{"video": {"pool_size": N}}``); when unset it - defaults to the existing ``VLLM_MEDIA_LOADING_THREAD_COUNT`` so no - DeepStream-specific env var is needed. The pool is a process-wide - singleton, so the first decode's value wins. - """ - if cls._pool is not None: - return cls._pool - if cls._pool_lock is None: - cls._pool_lock = threading.Lock() - with cls._pool_lock: - if cls._pool is not None: - return cls._pool - import os - - from nvidia.deepstream_videodecode import DecodePool - - if pool_size is None: - pool_size = int(os.environ.get("VLLM_MEDIA_LOADING_THREAD_COUNT", 8)) - pool_size = max(1, min(int(pool_size), 16)) - logger.info( - "[DeepStream] initializing decode pool with %d workers", - pool_size, - ) - cls._pool = DecodePool(num_workers=pool_size) - return cls._pool - - @classmethod - def decode_indices( - cls, - data: bytes, - frame_indices: list[int], - source: VideoSourceMetadata, - codec: str = "", - pool_size: int | None = None, - timeout_sec: float = 120.0, - ) -> tuple[npt.NDArray, list[int]]: - """Decode the requested frame indices from raw container bytes. - - The whole stream is decoded; the pool keeps exactly the frames whose - decode-order index is in ``frame_indices`` (1:1, frame-exact) and - sends EOS once the last one is matched. - - ``codec`` (e.g. ``"h264"``/``"hevc"``) lets the pool keep its NVDEC - session warm across same-codec streams and rebuild only on a codec - change. Frames are returned as a CPU NHWC uint8 array so the - upstream multimodal parser sees the same shape as the other - backends. - """ - if not frame_indices: - raise ValueError("DeepStream backend received no frame indices") - - result = cls._get_pool(pool_size).decode( - data, - target_indices=frame_indices, - codec=codec, - max_frames=len(frame_indices), - timeout_sec=timeout_sec, - ) - if result.error: - raise ValueError(f"DeepStream decode failed: {result.error}") - if result.frames is None or result.n_kept == 0: - raise ValueError("DeepStream decode produced no frames") - - valid = frame_indices[: result.n_kept] - # GPU -> CPU NHWC uint8 at the codec boundary (one PCIe copy); keeps - # the array shape identical to the OpenCV/PyAV backends. Copy into - # PINNED host memory (reused across calls by PyTorch's pinned caching - # allocator) so the D2H runs at full PCIe bandwidth (~13 GB/s) rather - # than the ~1 GB/s pageable path that plain ``.cpu()`` takes — ~12x - # faster for a 1080p x8 frame batch (~46ms -> ~4ms). ``numpy()`` keeps - # the pinned tensor alive via the array's base. - import torch - - gpu = result.frames - if gpu.is_cuda: - host = torch.empty(gpu.shape, dtype=gpu.dtype, pin_memory=True) - host.copy_(gpu, non_blocking=True) - torch.cuda.current_stream().synchronize() - arr = host.numpy() - else: - arr = gpu.numpy() - return arr, valid - @VIDEO_LOADER_REGISTRY.register("opencv") -class VideoBackend( - VideoLoader, - OpenCVVideoBackendMixin, - PyAVVideoBackendMixin, - TorchCodecVideoBackendMixin, - PyNvVideoCodecVideoBackendMixin, - DeepStreamVideoBackendMixin, -): +class VideoBackend(VideoLoader): """Uniform-sampling video backend. Samples ``num_frames`` uniformly across the video (or one frame every @@ -1080,11 +210,6 @@ def compute_frames_index_to_sample( 0, total_frames_num - 1, num_frames_to_sample, dtype=int ).tolist() - @classmethod - def _prepare_source(cls, source: VideoSourceMetadata) -> VideoSourceMetadata: - """Sampling-algorithm-specific metadata adjustment hook.""" - return source - @classmethod def load_bytes( cls, @@ -1094,12 +219,7 @@ def load_bytes( max_duration: int = 300, frame_recovery: bool = False, *, - backend: Literal[ - "opencv", "pyav", "torchcodec", "pynvvideocodec", "deepstream" - ] = "opencv", - num_ffmpeg_threads: int = 0, - seek_mode: Literal["exact", "approximate"] = "exact", - hw_decoders: int = PYNVVIDEOCODEC_DEFAULT_HW_DECODERS, + backend: VideoDecoderBackend = "opencv", **kwargs, ) -> tuple[npt.NDArray, dict[str, Any]]: """Load sampled frames from raw video bytes. @@ -1114,20 +234,25 @@ def load_bytes( Only honored by the OpenCV codec. backend: Decoding codec — ``"opencv"``, ``"pyav"``, ``"torchcodec"``, ``"pynvvideocodec"`` or ``"deepstream"``. - num_ffmpeg_threads: Number of FFmpeg decoding threads, only used by - TorchCodec: ``0`` (default) relies on the FFmpeg default value - which is ``min(cpu_count + 1, 16)``. - OpenCV will always use ``min(cpu_count, 16)`` while pyav will - always use ``min(cpu_count, (height + 15) / 16)``. - seek_mode: Seek mode for the TorchCodec decoder, only used by - TorchCodec: ``"exact"`` (default) guarantees frame-accurate - sampling by scanning the file on creation, while - ``"approximate"`` skips that scan for faster decoder creation - at the cost of relying on the file's metadata. See - https://meta-pytorch.org/torchcodec/stable/generated_examples/decoding/approximate_mode.html - for details. - hw_decoders: Maximum number of concurrent PyNvVideoCodec decoder - slots. Defaults to 2 and must be a positive integer. + kwargs: Codec-specific options, validated against and forwarded to + ``backend``: + + - ``num_ffmpeg_threads`` (TorchCodec): number of FFmpeg + decoding threads; ``0`` (default) relies on the FFmpeg + default value which is ``min(cpu_count + 1, 16)``. + OpenCV will always use ``min(cpu_count, 16)`` while pyav + will always use ``min(cpu_count, (height + 15) / 16)``. + - ``seek_mode`` (TorchCodec): ``"exact"`` (default) guarantees + frame-accurate sampling by scanning the file on creation, + while ``"approximate"`` skips that scan for faster decoder + creation at the cost of relying on the file's metadata. See + https://meta-pytorch.org/torchcodec/stable/generated_examples/decoding/approximate_mode.html + for details. + - ``hw_decoders`` (PyNvVideoCodec): maximum number of + concurrent decoder slots. Defaults to 2 and must be a + positive integer. + - ``pool_size`` / ``timeout_sec`` (DeepStream): decoder pool + size and pool acquisition timeout in seconds. Returns: Tuple of ``(frames_array, metadata_dict)``. @@ -1135,101 +260,16 @@ def load_bytes( target = VideoTargetMetadata( num_frames=num_frames, fps=fps, max_duration=max_duration ) - - if backend == "opencv": - cap = cls.open_video_capture(data) - _check_frame_pixel_limit( - int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), - int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)), - ) - source = cls._prepare_source(cls.get_video_metadata(cap)) - frame_idx = cls.compute_frames_index_to_sample( - source=source, target=target, **kwargs - ) - frames, valid = cls.read_frames( - cap, - frame_idx, - total_frames_num=source.total_frames_num, - frame_recovery=frame_recovery, - ) - elif backend == "pyav": - assert not frame_recovery, ( - "frame_recovery is only available for `opencv` backend" - ) - with av.open(BytesIO(data)) as container: - stream = container.streams.video[0] - _check_frame_pixel_limit(stream.width, stream.height) - source = cls._prepare_source(cls.get_metadata(container)) - frame_idx = cls.compute_frames_index_to_sample( - source=source, target=target, **kwargs - ) - frames, valid = cls.decode_frames( - container, frame_idx, source.original_fps, source.duration - ) - elif backend == "torchcodec": - assert not frame_recovery, ( - "frame_recovery is only available for `opencv` backend" - ) - check_torchcodec_available() - decoder = cls.make_torchcodec_decoder( - data, - num_ffmpeg_threads=num_ffmpeg_threads, - seek_mode=seek_mode, - ) - _check_frame_pixel_limit( - decoder.metadata.width or 0, - decoder.metadata.height or 0, - ) - source = cls._prepare_source(cls.get_torchcodec_metadata(decoder)) - frame_idx = cls.compute_frames_index_to_sample( - source=source, target=target, **kwargs - ) - frames, valid = cls.decode_torchcodec_frames(decoder, frame_idx) - elif backend == PYNVVIDEOCODEC_VIDEO_BACKEND: - if frame_recovery: - raise ValueError( - "frame_recovery is not supported for " - f"`{PYNVVIDEOCODEC_VIDEO_BACKEND}` backend" - ) - cls._configure_decoder_slots(hw_decoders) - frames, source, frame_idx, valid = cls.decode_frames_pynvvideocodec( - data, - target, - **kwargs, - ) - elif backend == "deepstream": - assert not frame_recovery, ( - "frame_recovery is only available for `opencv` backend" - ) - # Decode-pool size comes from media-io-kwargs (no env var); the - # pool is a process-wide singleton so the first decode's value - # wins. Pop it so it isn't forwarded to the frame sampler. - pool_size = kwargs.pop("pool_size", None) - # Probe container metadata from the bytes via GStreamer (in - # the deepstream video-decode wheel) — no PyAV/pymediainfo, no path. - from nvidia.deepstream_videodecode import probe_metadata - - total_frames, original_fps, duration, _w, _h, codec = probe_metadata(data) - _check_frame_pixel_limit(_w, _h) - source = cls._prepare_source( - VideoSourceMetadata( - total_frames_num=total_frames, - original_fps=original_fps, - duration=duration, - ) - ) - frame_idx = cls.compute_frames_index_to_sample( - source=source, target=target, **kwargs - ) - frames, valid = cls.decode_indices( - data, frame_idx, source, codec=codec, pool_size=pool_size - ) - else: - raise ValueError( - f"Unknown video codec backend {backend!r}; " - "valid options: 'opencv', 'pyav', 'torchcodec', " - "'pynvvideocodec' and 'deepstream'." - ) + sampling_kwargs, backend_kwargs = resolve_video_backend_kwargs(backend, kwargs) + frames, source, frame_idx, valid = decode_video( + backend, + cls, + data, + target, + sampling_kwargs, + backend_kwargs, + frame_recovery=frame_recovery, + ) if len(valid) < len(frame_idx): logger.warning( @@ -1742,7 +782,7 @@ def load_bytes( "molmo2", video_processor="Molmo2VideoProcessor", ) -class Molmo2VideoBackend(VideoLoader, OpenCVVideoBackendMixin): +class Molmo2VideoBackend(VideoLoader): @classmethod def get_candidate_target_fps( cls, @@ -1975,30 +1015,21 @@ def load_bytes_opencv( frame_recovery: bool = False, **kwargs, ) -> tuple[npt.NDArray, dict[str, Any]]: - cap = cls.open_video_capture(data) - _check_frame_pixel_limit( - int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), - int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)), - ) - - source = OpenCVVideoBackendMixin.get_video_metadata(cap) target = VideoTargetMetadata( num_frames=num_frames, fps=sampling_fps, - max_duration=source.duration, + max_duration=-1, ) - - frame_idx = cls.compute_frames_index_to_sample( - source=source, - target=target, - frame_sample_mode=frame_sample_mode, - max_fps=max_fps, - ) - - frames, valid_frame_indices = cls.read_frames( - cap, - frame_idx, - total_frames_num=source.total_frames_num, + frames, source, _, valid_frame_indices = decode_video( + "opencv", + cls, + data, + target, + { + "frame_sample_mode": frame_sample_mode, + "max_fps": max_fps, + }, + {}, frame_recovery=frame_recovery, ) @@ -2064,7 +1095,7 @@ def load_bytes( @VIDEO_LOADER_REGISTRY.register("openpangu") -class OpenCVDynamicOpenPanguVideoBackend(VideoLoader, OpenCVVideoBackendMixin): +class OpenCVDynamicOpenPanguVideoBackend(VideoLoader): @classmethod def compute_frames_index_to_sample( cls, @@ -2130,14 +1161,6 @@ def load_bytes( Returns: Tuple of (frames_array, metadata_dict) """ - cap = cls.open_video_capture(data) - _check_frame_pixel_limit( - int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), - int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)), - ) - - source = OpenCVVideoBackendMixin.get_video_metadata(cap) - # recompute source metadata with adjusted duration to ensure correct # sampling indices computation target = VideoTargetMetadata( @@ -2146,15 +1169,13 @@ def load_bytes( max_duration=max_duration, ) - frame_indices_list = cls.compute_frames_index_to_sample( - source=source, - target=target, - ) - - frames, valid_frame_indices = cls.read_frames( - cap, - frame_indices_list, - total_frames_num=source.total_frames_num, + frames, source, _, valid_frame_indices = decode_video( + "opencv", + cls, + data, + target, + {}, + {}, frame_recovery=frame_recovery, ) diff --git a/vllm/multimodal/video_decoders/__init__.py b/vllm/multimodal/video_decoders/__init__.py new file mode 100644 index 000000000000..c62310431652 --- /dev/null +++ b/vllm/multimodal/video_decoders/__init__.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from importlib import import_module +from typing import Any, Literal + +from .base import ( + PYNVVIDEOCODEC_DEFAULT_HW_DECODERS, + PYNVVIDEOCODEC_VIDEO_BACKEND, + VideoSourceMetadata, + VideoTargetMetadata, + check_frame_pixel_limit, +) + +VideoDecoderBackend = Literal[ + "opencv", "pyav", "torchcodec", "pynvvideocodec", "deepstream" +] + +_BACKEND_OPTION_DEFAULTS: dict[str, dict[str, Any]] = { + "opencv": {}, + "pyav": {}, + "torchcodec": { + "num_ffmpeg_threads": 0, + "seek_mode": "exact", + }, + PYNVVIDEOCODEC_VIDEO_BACKEND: { + "hw_decoders": PYNVVIDEOCODEC_DEFAULT_HW_DECODERS, + }, + "deepstream": { + "pool_size": None, + "timeout_sec": 120.0, + }, +} + + +def _get_backend_option_defaults(backend: str) -> dict[str, Any]: + try: + return _BACKEND_OPTION_DEFAULTS[backend] + except KeyError: + valid = ", ".join(repr(name) for name in _BACKEND_OPTION_DEFAULTS) + raise ValueError( + f"Unknown video codec backend {backend!r}; valid options: {valid}." + ) from None + + +def resolve_video_backend_kwargs( + backend: str, + kwargs: dict[str, Any], +) -> tuple[dict[str, Any], dict[str, Any]]: + """Split frame-sampling kwargs from options owned by a decoder backend.""" + defaults = _get_backend_option_defaults(backend) + sampling_kwargs = dict(kwargs) + backend_kwargs = dict(defaults) + backend_option_names = { + name for options in _BACKEND_OPTION_DEFAULTS.values() for name in options + } + misplaced = (sampling_kwargs.keys() & backend_option_names) - defaults.keys() + if misplaced: + names = ", ".join(sorted(misplaced)) + raise ValueError(f"{names} is not supported by the {backend!r} backend") + + for name in defaults: + if name in sampling_kwargs: + backend_kwargs[name] = sampling_kwargs.pop(name) + + return sampling_kwargs, backend_kwargs + + +def decode_video( + backend: str, + loader_cls, + data: bytes, + target: VideoTargetMetadata, + sampling_kwargs: dict[str, Any], + backend_kwargs: dict[str, Any], + *, + frame_recovery: bool, +): + """Decode a sampled video, importing only the selected backend.""" + _get_backend_option_defaults(backend) + if frame_recovery and backend != "opencv": + error = ( + ValueError if backend == PYNVVIDEOCODEC_VIDEO_BACKEND else AssertionError + ) + raise error(f"frame_recovery is not supported by the {backend!r} backend") + + decoder_kwargs = dict(backend_kwargs) + if backend == "opencv": + decoder_kwargs["frame_recovery"] = frame_recovery + module = import_module(f".{backend}", __name__) + decoder = getattr(module, f"decode_{backend}") + return decoder( + loader_cls, + data, + target, + sampling_kwargs, + **decoder_kwargs, + ) + + +__all__ = [ + "PYNVVIDEOCODEC_DEFAULT_HW_DECODERS", + "PYNVVIDEOCODEC_VIDEO_BACKEND", + "VideoDecoderBackend", + "VideoSourceMetadata", + "VideoTargetMetadata", + "check_frame_pixel_limit", + "decode_video", + "resolve_video_backend_kwargs", +] diff --git a/vllm/multimodal/video_decoders/base.py b/vllm/multimodal/video_decoders/base.py new file mode 100644 index 000000000000..7d354f68fffe --- /dev/null +++ b/vllm/multimodal/video_decoders/base.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Literal, NamedTuple + +from vllm import envs + +PYNVVIDEOCODEC_VIDEO_BACKEND: Literal["pynvvideocodec"] = "pynvvideocodec" +PYNVVIDEOCODEC_DEFAULT_HW_DECODERS = 2 + + +class VideoTargetMetadata(NamedTuple): + """Metadata describing the requested video sample.""" + + num_frames: int + fps: float + max_duration: float + + +class VideoSourceMetadata(NamedTuple): + """Metadata describing the encoded video source.""" + + total_frames_num: int + original_fps: float + duration: float + + +def check_frame_pixel_limit(width: int, height: int) -> None: + """Reject video frames exceeding ``VLLM_MAX_IMAGE_PIXELS``.""" + max_pixels = envs.VLLM_MAX_IMAGE_PIXELS + if max_pixels > 0 and width * height > max_pixels: + raise ValueError( + f"Video frame dimensions {width}x{height} " + f"({width * height} pixels) exceed the maximum of " + f"{max_pixels} pixels. Set VLLM_MAX_IMAGE_PIXELS to " + f"increase this limit." + ) diff --git a/vllm/multimodal/video_decoders/deepstream.py b/vllm/multimodal/video_decoders/deepstream.py new file mode 100644 index 000000000000..05262f55863f --- /dev/null +++ b/vllm/multimodal/video_decoders/deepstream.py @@ -0,0 +1,151 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import threading +from typing import Any, ClassVar + +import numpy.typing as npt + +from vllm.logger import init_logger + +from .base import VideoSourceMetadata, VideoTargetMetadata, check_frame_pixel_limit + +logger = init_logger(__name__) + + +def decode_deepstream( + loader_cls, + data: bytes, + target: VideoTargetMetadata, + sampling_kwargs: dict, + *, + pool_size: int | None = None, + timeout_sec: float = 120.0, +) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]: + from nvidia.deepstream_videodecode import probe_metadata + + total_frames, original_fps, duration, _width, _height, codec = probe_metadata(data) + check_frame_pixel_limit(_width, _height) + source = loader_cls._prepare_source( + VideoSourceMetadata(total_frames, original_fps, duration) + ) + frame_idx = loader_cls.compute_frames_index_to_sample( + source=source, target=target, **sampling_kwargs + ) + frames, valid = DeepStreamVideoBackendMixin.decode_indices( + data, + frame_idx, + source, + codec=codec, + pool_size=pool_size, + timeout_sec=timeout_sec, + ) + return frames, source, frame_idx, valid + + +class DeepStreamVideoBackendMixin: + """NVIDIA DeepStream (NVDEC) GPU-decode codec utilities. + + Decoding runs on a shared pool of daemon threads inside one CUDA + context (see the ``nvidia-deepstream-videodecode-cu13`` package). The + container bytes are pushed into an ``appsrc`` GStreamer pipeline, so no + local file path is required — HTTP and base64 sources decode identically + to local files. + + Like the OpenCV/PyAV mixins, this provides only the codec layer. + Frame *selection* lives in the loader's + ``compute_frames_index_to_sample`` and arrives here as an explicit + list of frame indices. + """ + + # Process-wide lazy decode pool, shared across all DeepStream backends. + _pool: ClassVar[Any] = None + _pool_lock: ClassVar[Any] = None + + @classmethod + def _get_pool(cls, pool_size: int | None = None): + """Lazy-initialize the shared decode pool on first use. + + ``pool_size`` (number of decode worker threads) comes from + ``--media-io-kwargs`` (``{"video": {"pool_size": N}}``); when unset it + defaults to the existing ``VLLM_MEDIA_LOADING_THREAD_COUNT`` so no + DeepStream-specific env var is needed. The pool is a process-wide + singleton, so the first decode's value wins. + """ + if cls._pool is not None: + return cls._pool + if cls._pool_lock is None: + cls._pool_lock = threading.Lock() + with cls._pool_lock: + if cls._pool is not None: + return cls._pool + import os + + from nvidia.deepstream_videodecode import DecodePool + + if pool_size is None: + pool_size = int(os.environ.get("VLLM_MEDIA_LOADING_THREAD_COUNT", 8)) + pool_size = max(1, min(int(pool_size), 16)) + logger.info( + "[DeepStream] initializing decode pool with %d workers", + pool_size, + ) + cls._pool = DecodePool(num_workers=pool_size) + return cls._pool + + @classmethod + def decode_indices( + cls, + data: bytes, + frame_indices: list[int], + source: VideoSourceMetadata, + codec: str = "", + pool_size: int | None = None, + timeout_sec: float = 120.0, + ) -> tuple[npt.NDArray, list[int]]: + """Decode the requested frame indices from raw container bytes. + + The whole stream is decoded; the pool keeps exactly the frames whose + decode-order index is in ``frame_indices`` (1:1, frame-exact) and + sends EOS once the last one is matched. + + ``codec`` (e.g. ``"h264"``/``"hevc"``) lets the pool keep its NVDEC + session warm across same-codec streams and rebuild only on a codec + change. Frames are returned as a CPU NHWC uint8 array so the + upstream multimodal parser sees the same shape as the other + backends. + """ + if not frame_indices: + raise ValueError("DeepStream backend received no frame indices") + + result = cls._get_pool(pool_size).decode( + data, + target_indices=frame_indices, + codec=codec, + max_frames=len(frame_indices), + timeout_sec=timeout_sec, + ) + if result.error: + raise ValueError(f"DeepStream decode failed: {result.error}") + if result.frames is None or result.n_kept == 0: + raise ValueError("DeepStream decode produced no frames") + + valid = frame_indices[: result.n_kept] + # GPU -> CPU NHWC uint8 at the codec boundary (one PCIe copy); keeps + # the array shape identical to the OpenCV/PyAV backends. Copy into + # PINNED host memory (reused across calls by PyTorch's pinned caching + # allocator) so the D2H runs at full PCIe bandwidth (~13 GB/s) rather + # than the ~1 GB/s pageable path that plain ``.cpu()`` takes — ~12x + # faster for a 1080p x8 frame batch (~46ms -> ~4ms). ``numpy()`` keeps + # the pinned tensor alive via the array's base. + import torch + + gpu = result.frames + if gpu.is_cuda: + host = torch.empty(gpu.shape, dtype=gpu.dtype, pin_memory=True) + host.copy_(gpu, non_blocking=True) + torch.cuda.current_stream().synchronize() + arr = host.numpy() + else: + arr = gpu.numpy() + return arr, valid diff --git a/vllm/multimodal/video_decoders/opencv.py b/vllm/multimodal/video_decoders/opencv.py new file mode 100644 index 000000000000..e5e474eab606 --- /dev/null +++ b/vllm/multimodal/video_decoders/opencv.py @@ -0,0 +1,295 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from io import BytesIO + +import numpy as np +import numpy.typing as npt + +from vllm.logger import init_logger +from vllm.utils.import_utils import PlaceholderModule + +from .base import ( + VideoSourceMetadata, + VideoTargetMetadata, + check_frame_pixel_limit, +) + +try: + import cv2 + import cv2.videoio_registry as vr +except ImportError: + cv2 = PlaceholderModule("cv2") + vr = PlaceholderModule("cv2").placeholder_attr("videoio_registry") + +logger = init_logger(__name__) + + +def decode_opencv( + loader_cls, + data: bytes, + target: VideoTargetMetadata, + sampling_kwargs: dict, + *, + frame_recovery: bool = False, +) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]: + cap = OpenCVVideoBackendMixin.open_video_capture(data) + check_frame_pixel_limit( + int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), + int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)), + ) + source = loader_cls._prepare_source(OpenCVVideoBackendMixin.get_video_metadata(cap)) + frame_idx = loader_cls.compute_frames_index_to_sample( + source=source, target=target, **sampling_kwargs + ) + frames, valid = OpenCVVideoBackendMixin.read_frames( + cap, + frame_idx, + total_frames_num=source.total_frames_num, + frame_recovery=frame_recovery, + ) + return frames, source, frame_idx, valid + + +class OpenCVVideoBackendMixin: + @staticmethod + def get_cv2_video_api(): + api_pref = None + for backend in vr.getStreamBufferedBackends(): + if not vr.hasBackend(backend): + continue + if not vr.isBackendBuiltIn(backend): + _, abi, api = vr.getStreamBufferedBackendPluginVersion(backend) + if abi < 1 or (abi == 1 and api < 2): + continue + api_pref = backend + break + return api_pref + + @classmethod + def open_video_capture(cls, data: bytes) -> "cv2.VideoCapture": + backend = cls.get_cv2_video_api() + cap = cv2.VideoCapture(BytesIO(data), backend, []) + if not cap.isOpened(): + raise ValueError("Could not open video stream") + return cap + + @staticmethod + def get_video_metadata(cap: "cv2.VideoCapture") -> VideoSourceMetadata: + total_frames_num = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + original_fps = cap.get(cv2.CAP_PROP_FPS) + duration = total_frames_num / original_fps if original_fps > 0 else 0 + return VideoSourceMetadata( + total_frames_num=total_frames_num, + original_fps=original_fps, + duration=duration, + ) + + @classmethod + def _can_use_for_recovery( + cls, + idx: int, + failed_frames: list[int], + next_target_map: dict[int, int], + total_frames: int, + ) -> bool: + """Check if current frame can recover the oldest failed frame.""" + if not failed_frames: + return False + oldest_failed = failed_frames[0] + limit = next_target_map.get(oldest_failed, total_frames) + return idx < limit + + @classmethod + def _read_frames_with_recovery( + cls, + cap: "cv2.VideoCapture", + frame_indices: list[int], + total_frames: int, + ) -> tuple[npt.NDArray, list[int], dict[int, int]]: + """ + Read frames with dynamic window forward-scan recovery. + + When a target frame fails to load, the next successfully grabbed + frame (before the next target frame) will be used to recover it. + + Args: + cap: OpenCV VideoCapture object + frame_indices: Sorted list of target frame indices to load + total_frames: Total number of frames in the video + + Returns: + Tuple of (frames_array, valid_frame_indices, recovered_map) + - frames_array: Array of loaded frames + - valid_frame_indices: List of frame indices that were loaded + - recovered_map: Dict mapping recovered_idx -> source_idx + """ + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + + assert width > 0 and height > 0, ( + f"Invalid video frame size: width={width}, height={height}" + ) + + frame_idx_set = set(frame_indices) + max_frame_idx = frame_indices[-1] if frame_indices else 0 + + # Build map: target_idx -> next_target_idx (for recovery window) + next_target_map: dict[int, int] = {} + for k in range(len(frame_indices) - 1): + next_target_map[frame_indices[k]] = frame_indices[k + 1] + next_target_map[frame_indices[-1]] = total_frames + + frames_list: list[npt.NDArray] = [] + valid_frame_indices: list[int] = [] + failed_frames_idx: list[int] = [] + recovered_map: dict[int, int] = {} + + i = 0 + for idx in range(max_frame_idx + 1): + is_target_frame = idx in frame_idx_set + + # Attempt to grab the current frame + ok = cap.grab() + + if not ok: + if is_target_frame: + logger.debug( + "Failed to grab frame %d during video loading.", + idx, + ) + failed_frames_idx.append(idx) + continue + + # Check if we should retrieve: target frame OR can recover a failed one + can_recover = cls._can_use_for_recovery( + idx, failed_frames_idx, next_target_map, total_frames + ) + + if is_target_frame or can_recover: + ret, frame = cap.retrieve() + + if ret and frame is not None and frame.size > 0: + rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + frames_list.append(rgb_frame) + valid_frame_indices.append(idx) + i += 1 + + if can_recover: + recovered_idx = failed_frames_idx.pop(0) + recovered_map[recovered_idx] = idx + logger.info( + "Recovered frame %d using frame %d (delay: %d)", + recovered_idx, + idx, + idx - recovered_idx, + ) + elif is_target_frame: + logger.debug( + "Failed to retrieve frame %d during video loading.", + idx, + ) + failed_frames_idx.append(idx) + + # Log any remaining failed frames + for failed_idx in failed_frames_idx: + logger.debug( + "Frame %d could not be recovered (end of video).", + failed_idx, + ) + + # Stack frames + if frames_list: + frames = np.stack(frames_list) + else: + frames = np.empty((0, height, width, 3), dtype=np.uint8) + + return frames, valid_frame_indices, recovered_map + + @classmethod + def _read_frames_no_recovery( + cls, + cap, + frame_indices: set[int], + max_frame_idx: int, + ) -> tuple[npt.NDArray, list[int]]: + num_expected_frames = len(frame_indices) + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + frames = np.empty((num_expected_frames, height, width, 3), dtype=np.uint8) + + i = 0 + valid_frame_indices = [] + for idx in range(max_frame_idx + 1): + ok = cap.grab() + if not ok: + # Frame is broken/unreadable, skip it + if idx in frame_indices: + logger.debug( + "Failed to grab frame %d during video loading. " + "This frame will be skipped.", + idx, + ) + continue + if idx in frame_indices: + ret, frame = cap.retrieve() + if ret: + frames[i] = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + valid_frame_indices.append(idx) + i += 1 + else: + # retrieve() failed even though grab() succeeded + logger.debug( + "Failed to retrieve frame %d during video loading. " + "This frame will be skipped.", + idx, + ) + + valid_num_frames = len(valid_frame_indices) + if valid_num_frames < num_expected_frames: + logger.warning( + "Video loading completed with %d broken/unreadable frames. " + "Expected %d frames but only loaded %d frames.", + num_expected_frames - valid_num_frames, + num_expected_frames, + valid_num_frames, + ) + + return frames[:valid_num_frames], valid_frame_indices + + @classmethod + def read_frames( + cls, + cap: "cv2.VideoCapture", + frame_idx: list[int], + total_frames_num: int, + *, + frame_recovery: bool = False, + ) -> tuple[npt.NDArray, list[int]]: + if frame_recovery: + num_frames_to_sample = len(frame_idx) + frames, valid_frame_indices, recovered_map = cls._read_frames_with_recovery( + cap, frame_idx, total_frames_num + ) + + if recovered_map: + logger.info( + "Frame recovery: %d frames recovered using forward scan.", + len(recovered_map), + ) + else: + frame_idx_set = set(frame_idx) + num_frames_to_sample = len(frame_idx_set) + frames, valid_frame_indices = cls._read_frames_no_recovery( + cap, frame_idx_set, max(frame_idx) + ) + valid_num_frames = len(valid_frame_indices) + if valid_num_frames < num_frames_to_sample: + logger.warning( + "Video loading completed with %d broken/unreadable frames. " + "Expected to sample %d frames but only loaded %d frames.", + num_frames_to_sample - valid_num_frames, + num_frames_to_sample, + valid_num_frames, + ) + return frames, valid_frame_indices diff --git a/vllm/multimodal/video_decoders/pyav.py b/vllm/multimodal/video_decoders/pyav.py new file mode 100644 index 000000000000..dc82b01230d1 --- /dev/null +++ b/vllm/multimodal/video_decoders/pyav.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from io import BytesIO + +import numpy as np +import numpy.typing as npt + +from vllm.utils.import_utils import PlaceholderModule + +from .base import ( + VideoSourceMetadata, + VideoTargetMetadata, + check_frame_pixel_limit, +) + +try: + import av +except ImportError: + av = PlaceholderModule("av") # type: ignore[assignment] + + +def decode_pyav( + loader_cls, + data: bytes, + target: VideoTargetMetadata, + sampling_kwargs: dict, +) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]: + with av.open(BytesIO(data)) as container: + stream = container.streams.video[0] + check_frame_pixel_limit(stream.width, stream.height) + source = loader_cls._prepare_source( + PyAVVideoBackendMixin.get_metadata(container) + ) + frame_idx = loader_cls.compute_frames_index_to_sample( + source=source, target=target, **sampling_kwargs + ) + frames, valid = PyAVVideoBackendMixin.decode_frames( + container, frame_idx, source.original_fps, source.duration + ) + return frames, source, frame_idx, valid + + +class PyAVVideoBackendMixin: + """PyAV (in-process FFmpeg bindings) codec utilities. + + Reads stream metadata and decodes target frames via per-frame + ``container.seek()``. The seek releases the GIL between frames and + scales with the number of sampled frames rather than the video + length, enabling concurrent decoding under serving load. + """ + + @staticmethod + def get_metadata( + container: "av.container.InputContainer", + ) -> VideoSourceMetadata: + if not container.streams.video: + raise ValueError("No video streams found in container") + stream = container.streams.video[0] + total_frames = stream.frames or 0 + fps = float(stream.average_rate) if stream.average_rate else 0.0 + duration = float(stream.duration * stream.time_base) if stream.duration else 0.0 + if total_frames == 0 and duration > 0 and fps > 0: + total_frames = int(duration * fps) + return VideoSourceMetadata(total_frames, fps, duration) + + @staticmethod + def decode_frames( + container: "av.container.InputContainer", + frame_indices: list[int], + fps: float, + duration: float, + ) -> tuple[npt.NDArray, list[int]]: + """Decode target frames via per-frame seek + forward decode to PTS.""" + stream = container.streams.video[0] + # SLICE parallelizes within a single frame without the + # one-frame-per-thread latency penalty of FRAME threading. + stream.thread_type = "SLICE" + time_base = stream.time_base + + frames_list: list[npt.NDArray] = [] + valid_indices: list[int] = [] + frame_interval = 1.0 / fps if fps > 0 else 0.1 + max_ts = max(0.0, duration - frame_interval) if duration > 0 else float("inf") + + decoder = None + last_pts = None + for idx in frame_indices: + ts = min(idx / fps, max_ts) if fps > 0 else 0.0 + pts = int(ts / time_base) + # seek() snaps backward to a keyframe; reuse the running decoder + # while targets advance monotonically to avoid re-decoding the + # GOP prefix once per requested frame. + if decoder is None or last_pts is None or pts <= last_pts: + container.seek(pts, stream=stream) + decoder = container.decode(video=0) + chosen = None + for frame in decoder: + if frame.pts is not None and frame.pts >= pts: + chosen = frame + last_pts = frame.pts + break + if chosen is not None: + frames_list.append(chosen.to_ndarray(format="rgb24")) + valid_indices.append(idx) + else: + decoder = None + + if not frames_list: + return np.empty((0,), dtype=np.uint8), valid_indices + return np.stack(frames_list), valid_indices diff --git a/vllm/multimodal/video_decoders/pynvvideocodec.py b/vllm/multimodal/video_decoders/pynvvideocodec.py new file mode 100644 index 000000000000..794c7a86d9d6 --- /dev/null +++ b/vllm/multimodal/video_decoders/pynvvideocodec.py @@ -0,0 +1,372 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import os +import tempfile +import threading +from contextlib import contextmanager, suppress +from typing import ClassVar, NamedTuple + +import numpy as np +import numpy.typing as npt + +from vllm.logger import init_logger +from vllm.utils.mem_constants import MiB_bytes + +from .base import ( + PYNVVIDEOCODEC_DEFAULT_HW_DECODERS, + VideoSourceMetadata, + VideoTargetMetadata, + check_frame_pixel_limit, +) + +logger = init_logger(__name__) + + +def decode_pynvvideocodec( + loader_cls, + data: bytes, + target: VideoTargetMetadata, + sampling_kwargs: dict, + *, + hw_decoders: int = PYNVVIDEOCODEC_DEFAULT_HW_DECODERS, +) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]: + PyNvVideoCodecVideoBackendMixin._configure_decoder_slots(hw_decoders) + return PyNvVideoCodecVideoBackendMixin.decode_frames_pynvvideocodec( + loader_cls, + data, + target, + **sampling_kwargs, + ) + + +class PyNvVideoCodecSourceMetadata(NamedTuple): + """Metadata needed before GPU video decode.""" + + source: VideoSourceMetadata + width: int + height: int + + +# Per-decoder upper bound reserved for persistent PyNvVideoCodec surfaces. +PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES = 128 * MiB_bytes +PYNVVIDEOCODEC_DECODER_CACHE_SIZE = 2 +# Per-API-server CUDA context and driver allocation, measured with +# PyNvVideoCodec 2.0.4 on H100. +PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES = int(1.8 * 1024 * MiB_bytes) + + +def validate_pynvvideocodec_hw_decoders(hw_decoders: object) -> int: + if ( + isinstance(hw_decoders, bool) + or not isinstance(hw_decoders, int) + or hw_decoders < 1 + ): + raise ValueError("hw_decoders must be a positive integer") + return hw_decoders + + +def _pynvvideocodec_exception_types(nvc) -> tuple[type[Exception], ...]: + return tuple( + exception_type + for name in dir(nvc) + if name.startswith("PyNvVCException") + and isinstance((exception_type := getattr(nvc, name)), type) + and issubclass(exception_type, Exception) + ) + + +def _pynvvc_frames_to_nhwc(frames): + """Return a stacked PyNvVideoCodec frame batch as contiguous NHWC.""" + if frames.shape[-1] != 3 and frames.shape[-3] == 3: + frames = frames.permute(0, 2, 3, 1) + return frames.contiguous() + + +class PyNvVideoCodecDecoderSlot: + """A retained PyNv decoder slot and its CUDA stream. + + The decoder is reused across requests: ``reconfigure_decoder`` repoints the + existing decoder at each new source instead of paying a fresh + ``SimpleDecoder`` construction per request. Construction (CUVID parser + + decoder + surface-pool allocation) is the dominant per-request cost, so + reconfiguring is far cheaper. A single decoder serves both metadata + (``len``/``get_stream_metadata``) and frame decode -- no separate + metadata decoder. + """ + + def __init__(self, stream) -> None: + self.stream = stream + self.decoder = None + self.source_path: str | None = None + + def invalidate(self) -> None: + self.decoder = None + self.source_path = None + + def _construct(self, file_path: str, nvc, device_index: int) -> None: + self.invalidate() + decoder = nvc.SimpleDecoder( + file_path, + output_color_type=nvc.OutputColorType.RGB, + use_device_memory=True, + need_scanned_stream_metadata=True, + gpu_id=device_index, + cuda_stream=self.stream.cuda_stream, + decoder_cache_size=PYNVVIDEOCODEC_DECODER_CACHE_SIZE, + ) + self.decoder = decoder + self.source_path = file_path + + def get_decoder(self, file_path: str, nvc, device_index: int): + if self.decoder is None: + self._construct(file_path, nvc, device_index) + elif self.source_path != file_path: + try: + self.decoder.reconfigure_decoder(file_path) + self.source_path = file_path + except Exception: + # reconfigure unsupported/unsafe for this source -> rebuild. + self._construct(file_path, nvc, device_index) + return self.decoder + + +class _PyNvDecoderPool: + """Process-wide singleton managing PyNvVideoCodec decoder slot state. + + Prevents subclass counter shadowing (GHSA-j682-9xp5-rrf3) by storing + all mutable pool state in a single module-level instance rather than + in ClassVar attributes that get shadowed by Python's augmented + assignment semantics on subclasses. + """ + + def __init__(self) -> None: + self.slots: list[PyNvVideoCodecDecoderSlot] = [] + self.active: int = 0 + self.cond: threading.Condition = threading.Condition() + self.max_slots: int | None = None + + def configure(self, hw_decoders: int) -> None: + with self.cond: + if self.max_slots is None: + self.max_slots = hw_decoders + elif self.max_slots != hw_decoders: + raise RuntimeError( + "PyNvVideoCodec decoder count is already configured as " + f"{self.max_slots}, got {hw_decoders}" + ) + + +_pynv_decoder_pool = _PyNvDecoderPool() + + +class PyNvVideoCodecVideoBackendMixin: + """PyNvVideoCodec utilities for GPU-backed frame decode.""" + + _DEVICE_INDEX: ClassVar[int] = 0 + + @classmethod + def _create_decoder_slot(cls) -> PyNvVideoCodecDecoderSlot: + import torch + + return PyNvVideoCodecDecoderSlot(torch.cuda.Stream(device=cls._DEVICE_INDEX)) + + @classmethod + def _configure_decoder_slots(cls, hw_decoders: object) -> None: + hw_decoders = validate_pynvvideocodec_hw_decoders(hw_decoders) + _pynv_decoder_pool.configure(hw_decoders) + + @staticmethod + @contextmanager + def _torch_stream_context(stream): + import torch + + torch.accelerator.set_device_index(stream.device.index) + previous_stream = torch.accelerator.current_stream() + torch.accelerator.set_stream(stream) + try: + yield + finally: + torch.accelerator.set_stream(previous_stream) + + @classmethod + @contextmanager + def _borrow_decoder_slot(cls): + pool = _pynv_decoder_pool + create_slot = False + with pool.cond: + if pool.max_slots is None: + raise RuntimeError("PyNvVideoCodec decoder slots are not configured") + while True: + if pool.slots: + slot = pool.slots.pop() + break + if pool.active < pool.max_slots: + pool.active += 1 + create_slot = True + break + pool.cond.wait() + + if create_slot: + try: + slot = cls._create_decoder_slot() + except Exception: + with pool.cond: + pool.active -= 1 + pool.cond.notify() + raise + + borrow_succeeded = False + try: + yield slot + borrow_succeeded = True + finally: + if not borrow_succeeded: + slot.invalidate() + with pool.cond: + pool.slots.append(slot) + pool.cond.notify() + + @staticmethod + def _metadata_value(metadata, *names: str, default=None): + for name in names: + value = getattr(metadata, name, None) + if value is not None: + return value + return default + + @classmethod + def _read_source_metadata( + cls, + file_path: str, + nvc, + ) -> PyNvVideoCodecSourceMetadata: + with cls._borrow_decoder_slot() as decoder_slot: + with cls._torch_stream_context(decoder_slot.stream): + decoder = decoder_slot.get_decoder( + file_path, nvc, device_index=cls._DEVICE_INDEX + ) + metadata = decoder.get_stream_metadata() + total_frames_num = len(decoder) + width = int(cls._metadata_value(metadata, "width", default=0)) + height = int(cls._metadata_value(metadata, "height", default=0)) + original_fps = float( + cls._metadata_value( + metadata, + "average_fps", + "avg_frame_rate", + "frame_rate", + "frameRate", + default=0.0, + ) + ) + duration = float( + cls._metadata_value(metadata, "duration", default=0.0) + or (total_frames_num / original_fps if original_fps > 0 else 0.0) + ) + if total_frames_num <= 0: + raise ValueError("Could not determine video frame count") + if width <= 0 or height <= 0: + raise ValueError("Could not determine video dimensions") + return PyNvVideoCodecSourceMetadata( + source=VideoSourceMetadata(total_frames_num, original_fps, duration), + width=width, + height=height, + ) + + @classmethod + def _decode_to_pinned_host( + cls, + file_path: str, + frame_idx: list[int], + nvc, + ) -> npt.NDArray: + import torch + + if not frame_idx: + return np.empty((0,), dtype=np.uint8) + + with cls._borrow_decoder_slot() as decoder_slot: + stream = decoder_slot.stream + with cls._torch_stream_context(stream): + try: + decoder = decoder_slot.get_decoder( + file_path, nvc, device_index=cls._DEVICE_INDEX + ) + decoded_frames = decoder.get_batch_frames_by_index(frame_idx) + except Exception as exc: + if not isinstance( + exc, + _pynvvideocodec_exception_types(nvc) + (IndexError,), + ): + raise + raise ValueError("Invalid or unsupported video file.") from exc + if len(decoded_frames) < len(frame_idx): + logger.warning( + "pynvvideocodec video loading: expected %d frames but got %d.", + len(frame_idx), + len(decoded_frames), + ) + torch_frames = [torch.from_dlpack(frame) for frame in decoded_frames] + if not torch_frames: + return np.empty((0,), dtype=np.uint8) + device_frames = torch.stack(torch_frames) + if device_frames.ndim != 4: + raise ValueError( + "PyNvVideoCodec returned frames with unexpected shape " + f"{tuple(device_frames.shape)}" + ) + device_frames = _pynvvc_frames_to_nhwc(device_frames) + host_frames = torch.empty( + device_frames.shape, + dtype=device_frames.dtype, + device="cpu", + pin_memory=True, + ) + host_frames.copy_(device_frames, non_blocking=True) + stream.synchronize() + host_array = host_frames.numpy() + del decoded_frames, torch_frames, device_frames + return host_array + + @classmethod + def decode_frames_pynvvideocodec( + cls, + loader_cls, + data: bytes, + target: VideoTargetMetadata, + **kwargs, + ) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]: + import PyNvVideoCodec as nvc + + from vllm.multimodal.gpu_ipc_memory import get_mm_gpu_ipc_pool + + temp_fd, temp_path = tempfile.mkstemp(suffix=".mp4") + try: + with os.fdopen(temp_fd, "wb") as temp_file: + temp_file.write(data) + + try: + gpu_source = cls._read_source_metadata(temp_path, nvc) + except Exception as exc: + if not isinstance(exc, _pynvvideocodec_exception_types(nvc)): + raise + raise ValueError("Invalid or unsupported video file.") from exc + check_frame_pixel_limit(gpu_source.width, gpu_source.height) + source = loader_cls._prepare_source(gpu_source.source) + frame_idx = loader_cls.compute_frames_index_to_sample( + source=source, target=target, **kwargs + ) + raw_frame_bytes = len(frame_idx) * gpu_source.height * gpu_source.width * 3 + pool = get_mm_gpu_ipc_pool() + if pool is None or raw_frame_bytes == 0: + frames = cls._decode_to_pinned_host(temp_path, frame_idx, nvc) + else: + with pool.acquire(raw_frame_bytes): + frames = cls._decode_to_pinned_host(temp_path, frame_idx, nvc) + finally: + with suppress(FileNotFoundError): + os.unlink(temp_path) + + valid_frame_indices = frame_idx[: int(frames.shape[0])] + return frames, source, frame_idx, valid_frame_indices diff --git a/vllm/multimodal/video_decoders/torchcodec.py b/vllm/multimodal/video_decoders/torchcodec.py new file mode 100644 index 000000000000..f1f3290629a2 --- /dev/null +++ b/vllm/multimodal/video_decoders/torchcodec.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Literal + +import numpy as np +import numpy.typing as npt + +from vllm.utils.import_utils import PlaceholderModule, check_torchcodec_available + +from .base import ( + VideoSourceMetadata, + VideoTargetMetadata, + check_frame_pixel_limit, +) + +try: + from torchcodec.decoders import VideoDecoder +except (ImportError, RuntimeError): + VideoDecoder = PlaceholderModule("torchcodec").placeholder_attr( # type: ignore[assignment] + "decoders.VideoDecoder" + ) + + +def decode_torchcodec( + loader_cls, + data: bytes, + target: VideoTargetMetadata, + sampling_kwargs: dict, + *, + num_ffmpeg_threads: int = 0, + seek_mode: Literal["exact", "approximate"] = "exact", +) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]: + check_torchcodec_available() + decoder = TorchCodecVideoBackendMixin.make_torchcodec_decoder( + data, + num_ffmpeg_threads=num_ffmpeg_threads, + seek_mode=seek_mode, + ) + check_frame_pixel_limit( + decoder.metadata.width or 0, + decoder.metadata.height or 0, + ) + source = loader_cls._prepare_source( + TorchCodecVideoBackendMixin.get_torchcodec_metadata(decoder) + ) + frame_idx = loader_cls.compute_frames_index_to_sample( + source=source, target=target, **sampling_kwargs + ) + frames, valid = TorchCodecVideoBackendMixin.decode_torchcodec_frames( + decoder, frame_idx + ) + return frames, source, frame_idx, valid + + +class TorchCodecVideoBackendMixin: + """TorchCodec (FFmpeg-backed, PyTorch-native) codec utilities. + + Builds a :class:`~torchcodec.decoders.VideoDecoder` over the in-memory + bytes and extracts the sampled indices with a single batched + ``get_frames_at`` call, while releasing the GIL during decode. + """ + + @staticmethod + def make_torchcodec_decoder( + data: bytes, + *, + num_ffmpeg_threads: int = 0, + seek_mode: Literal["exact", "approximate"] = "exact", + ) -> "VideoDecoder": + # NHWC matches the (num_frames, H, W, 3) uint8 RGB layout the rest + # of the pipeline expects, avoiding a transpose. + return VideoDecoder( + data, + dimension_order="NHWC", + num_ffmpeg_threads=num_ffmpeg_threads, + seek_mode=seek_mode, + ) + + @staticmethod + def get_torchcodec_metadata(decoder: "VideoDecoder") -> VideoSourceMetadata: + md = decoder.metadata + total_frames = md.num_frames or 0 + fps = float(md.average_fps) if md.average_fps else 0.0 + duration = float(md.duration_seconds) if md.duration_seconds else 0.0 + if total_frames == 0 and duration > 0 and fps > 0: + total_frames = int(duration * fps) + return VideoSourceMetadata(total_frames, fps, duration) + + @staticmethod + def decode_torchcodec_frames( + decoder: "VideoDecoder", + frame_indices: list[int], + ) -> tuple[npt.NDArray, list[int]]: + """Decode the requested indices in one batched, index-exact call.""" + if not frame_indices: + return np.empty((0,), dtype=np.uint8), [] + # Note: torchcodec releases the GIL for the entire call + batch = decoder.get_frames_at(frame_indices) + return batch.data.numpy(), list(frame_indices) From 241ff8c443175ae3e5e2b42b866e4c9275bd7744 Mon Sep 17 00:00:00 2001 From: karthik <56480632+gangula-karthik@users.noreply.github.com> Date: Tue, 18 Aug 2026 23:53:49 +0800 Subject: [PATCH 092/839] [Model] Enable LoRA support for tower and connector in LlavaNextForConditionalGeneration (#49788) Signed-off-by: gangula-karthik --- docs/models/supported_models.md | 2 +- vllm/model_executor/models/llava_next.py | 56 ++++++++++++++++++++++-- 2 files changed, 53 insertions(+), 5 deletions(-) diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index f9d55cb97e7a..0ed82796290b 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -571,7 +571,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen | `Llama4ForConditionalGeneration` | Llama 4 | T + I+ | `meta-llama/Llama-4-Scout-17B-16E-Instruct`, `meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8`, `meta-llama/Llama-4-Maverick-17B-128E-Instruct`, etc. | ✅︎ | ✅︎ | | `Llama_Nemotron_Nano_VL` | Llama Nemotron Nano VL | T + IE+ | `nvidia/Llama-3.1-Nemotron-Nano-VL-8B-V1` | ✅︎ | ✅︎ | | `LlavaForConditionalGeneration` | LLaVA-1.5, Pixtral (HF Transformers) | T + IE+ | `llava-hf/llava-1.5-7b-hf`, `mistral-community/pixtral-12b`, etc. | ✅︎ | ✅︎ | -| `LlavaNextForConditionalGeneration` | LLaVA-NeXT, Granite Vision | T + IE+ | `llava-hf/llava-v1.6-mistral-7b-hf`, `llava-hf/llava-v1.6-vicuna-7b-hf`, `ibm-granite/granite-vision-3.3-2b`, etc. | | ✅︎ | +| `LlavaNextForConditionalGeneration` | LLaVA-NeXT, Granite Vision | T + IE+ | `llava-hf/llava-v1.6-mistral-7b-hf`, `llava-hf/llava-v1.6-vicuna-7b-hf`, `ibm-granite/granite-vision-3.3-2b`, etc. | ✅︎ | ✅︎ | | `LlavaNextVideoForConditionalGeneration` | LLaVA-NeXT-Video | T + V | `llava-hf/LLaVA-NeXT-Video-7B-hf`, etc. | ✅︎ | ✅︎ | | `LlavaOnevision2ForConditionalGeneration` | LLaVA-OneVision-2 | T + I+ + V+ | `lmms-lab-encoder/LLaVA-OneVision-2-8B-Instruct` | | | | `LlavaOnevisionForConditionalGeneration` | LLaVA-Onevision | T + I+ + V+ | `llava-hf/llava-onevision-qwen2-7b-ov-hf`, `llava-hf/llava-onevision-qwen2-0.5b-ov-hf`, etc. | | ✅︎ | diff --git a/vllm/model_executor/models/llava_next.py b/vllm/model_executor/models/llava_next.py index 44d50f434a4c..42a2df982430 100644 --- a/vllm/model_executor/models/llava_next.py +++ b/vllm/model_executor/models/llava_next.py @@ -15,13 +15,18 @@ from vllm.config import VllmConfig from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.multimodal.inputs import MultiModalFieldConfig +from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItem from vllm.multimodal.parse import ImageSize from vllm.sequence import IntermediateTensors from vllm.utils.tensor_schema import TensorSchema, TensorShape from .clip import CLIPVisionModel -from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP +from .interfaces import ( + MultiModalEmbeddings, + SupportsLoRA, + SupportsMultiModal, + SupportsPP, +) from .llava import ( BaseLlavaMultiModalProcessor, BaseLlavaProcessingInfo, @@ -30,6 +35,7 @@ LlavaMultiModalProjector, init_vision_tower_for_llava, ) +from .module_mapping import MultiModelKeys from .siglip import SiglipVisionModel from .utils import ( AutoWeightsLoader, @@ -37,7 +43,7 @@ init_vllm_registered_model, maybe_prefix, ) -from .vision import get_num_selected_vision_tokens +from .vision import get_num_selected_vision_tokens, get_vision_encoder_info class LlavaNextImagePixelInputs(TensorSchema): @@ -222,7 +228,14 @@ def _get_mm_fields_config( info=LlavaNextProcessingInfo, dummy_inputs=LlavaDummyInputsBuilder, ) -class LlavaNextForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP): +class LlavaNextForConditionalGeneration( + nn.Module, SupportsLoRA, SupportsMultiModal, SupportsPP +): + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], + } + hf_to_vllm_mapper = WeightsMapper( orig_to_new_prefix={ # mapping for new names in checkpoint saved after transformers v4.52 @@ -582,3 +595,38 @@ def compute_logits( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + + def get_mm_mapping(self) -> MultiModelKeys: + """ + Get the module prefix in multimodal models + """ + return MultiModelKeys.from_string_field( + language_model="language_model", + connector="multi_modal_projector", + tower_model="vision_tower", + ) + + def get_mm_lora_token_counts( + self, + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality + + pixel_values = mm_kwargs.get("pixel_values") if mm_kwargs else None + if pixel_values is None or not isinstance(pixel_values.data, torch.Tensor): + return num_mm_embeds, num_mm_embeds + + # Unpad runs after the connector, so `num_mm_embeds` is not invertible. + num_tiles = pixel_values.data.shape[0] + encoder_info = get_vision_encoder_info(self.config) + tile_size = encoder_info.get_image_size() + tokens_per_tile = encoder_info.get_num_image_tokens( + image_width=tile_size, image_height=tile_size + ) + selected_per_tile = get_num_selected_vision_tokens( + tokens_per_tile, self.config.vision_feature_select_strategy + ) + return num_tiles * tokens_per_tile, num_tiles * selected_per_tile From f9f066d195ca079c7403d9d9447c6b1d740c348c Mon Sep 17 00:00:00 2001 From: Jikui Xie Date: Wed, 19 Aug 2026 00:00:55 +0800 Subject: [PATCH 093/839] [Bugfix][PaliGemma] Remove stale image embedding scaling (#52692) Signed-off-by: jikuixie Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: OpenAI Codex Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- vllm/model_executor/models/paligemma.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/vllm/model_executor/models/paligemma.py b/vllm/model_executor/models/paligemma.py index d7b8e77c63b6..c6c2eb64eb42 100644 --- a/vllm/model_executor/models/paligemma.py +++ b/vllm/model_executor/models/paligemma.py @@ -378,8 +378,6 @@ def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: if image_input is None: return [] vision_embeddings = self._process_image_input(image_input) - # https://github.com/huggingface/transformers/blob/main/src/transformers/models/paligemma/modeling_paligemma.py#L294 # noqa - vision_embeddings = vision_embeddings * (self.config.hidden_size**-0.5) return vision_embeddings def forward( From 88b2bff2c63d0f28396451f1199d09ee0f3e2d88 Mon Sep 17 00:00:00 2001 From: fxmarty-amd Date: Tue, 18 Aug 2026 18:56:56 +0200 Subject: [PATCH 094/839] [MOE] Standardize and abstract fused shared expert optimization selection (#51695) Signed-off-by: Felix Marty --- .../layers/test_fused_shared_expert.py | 700 ++++++++++++++++++ vllm/model_executor/layers/fused_moe/layer.py | 19 +- vllm/model_executor/layers/fused_moe/utils.py | 85 +++ .../layers/quantization/quark/quark.py | 102 ++- .../layers/quantization/utils/config_utils.py | 155 ++++ vllm/model_executor/models/AXK1.py | 32 +- vllm/model_executor/models/deepseek_mtp.py | 16 +- vllm/model_executor/models/deepseek_v2.py | 32 +- vllm/model_executor/models/glm4_moe.py | 32 +- vllm/model_executor/models/glm4_moe_lite.py | 17 +- .../models/glm4_moe_lite_mtp.py | 16 +- vllm/model_executor/models/glm4_moe_mtp.py | 17 +- vllm/model_executor/models/qwen3_5.py | 13 +- vllm/model_executor/models/qwen3_5_mtp.py | 18 +- vllm/model_executor/models/qwen3_next.py | 54 +- vllm/model_executor/models/qwen3_next_mtp.py | 10 + vllm/models/deepseek_v32/amd/mtp.py | 16 +- vllm/models/deepseek_v32/nvidia/mtp.py | 46 +- vllm/models/deepseek_v4/amd/model.py | 95 ++- vllm/models/minimax_m3/amd/model.py | 94 ++- 20 files changed, 1288 insertions(+), 281 deletions(-) create mode 100644 tests/model_executor/layers/test_fused_shared_expert.py create mode 100644 vllm/model_executor/layers/quantization/utils/config_utils.py diff --git a/tests/model_executor/layers/test_fused_shared_expert.py b/tests/model_executor/layers/test_fused_shared_expert.py new file mode 100644 index 000000000000..7080c66f128c --- /dev/null +++ b/tests/model_executor/layers/test_fused_shared_expert.py @@ -0,0 +1,700 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import importlib +from copy import deepcopy +from types import SimpleNamespace +from typing import Any, cast +from unittest.mock import patch + +import pytest +import torch +from torch import nn + +import vllm.config as vllm_config_module +from vllm.config import VllmConfig, set_current_vllm_config +from vllm.model_executor.layers.fused_moe import utils as fused_moe_utils +from vllm.model_executor.layers.fused_moe.layer import determine_expert_counts +from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig +from vllm.model_executor.layers.quantization.utils.config_utils import ( + is_shared_expert_quant_fse_compatible, +) +from vllm.model_executor.models.utils import PPMissingLayer +from vllm.models.deepseek_v4 import quant_config as deepseek_v4_quant_config +from vllm.models.minimax_m3.amd import model as minimax_m3_model +from vllm.transformers_utils.configs.minimax_m3 import MiniMaxM3TextConfig +from vllm.transformers_utils.configs.qwen3_5_moe import Qwen3_5MoeTextConfig + +_QUARK_FSE_CONFIG: dict[str, Any] = { + "global_quant_config": { + "input_tensors": { + "dtype": "fp4", + "is_dynamic": True, + "qscheme": "per_group", + "ch_axis": -1, + "group_size": 32, + "block_size": None, + "symmetric": None, + "round_method": "half_even", + "scale_type": "float", + "scale_format": "e8m0", + "scale_calculation_mode": "even", + "mx_element_dtype": None, + "observer_cls": "PerBlockMXObserver", + "is_scale_quant": False, + "enable_buffer_reuse": False, + "max_input_numel": 4194304, + }, + "output_tensors": None, + "weight": { + "dtype": "fp4", + "is_dynamic": False, + "qscheme": "per_group", + "ch_axis": -1, + "group_size": 32, + "block_size": None, + "symmetric": None, + "round_method": "half_even", + "scale_type": "float", + "scale_format": "e8m0", + "scale_calculation_mode": "even", + "mx_element_dtype": None, + "observer_cls": "PerBlockMXObserver", + "is_scale_quant": False, + "enable_buffer_reuse": False, + "max_input_numel": 4194304, + }, + "bias": None, + "target_device": None, + }, + "algo_config": None, + "softmax_quant_spec": None, + "quant_method": "quark", + "layer_type_quant_config": {}, + "layer_quant_config": {}, + "kv_cache_quant_config": {}, + "kv_cache_post_rope": False, + "quant_mode": "eager_mode", + "version": "0.12+9d3d471cdf1", + "export": { + "kv_cache_group": [], + "min_kv_scale": 0.0, + "pack_method": "reorder", + "weight_format": "real_quantized", + "weight_merge_groups": None, + }, +} + + +def get_deepseek_v4_quark_config(exclude: list[str]) -> dict[str, Any]: + """Return the DeepSeek-V4-Pro-MXFP4 FSE quantization layout.""" + # Mimics https://huggingface.co/amd/DeepSeek-V4-Pro-MXFP4. + quantization_config: dict[str, Any] = deepcopy(_QUARK_FSE_CONFIG) + quantization_config["exclude"] = exclude + mxfp4_config = cast( + dict[str, Any], deepcopy(quantization_config["global_quant_config"]) + ) + fp8_config = deepcopy(mxfp4_config) + fp8_config["input_tensors"].update( + { + "dtype": "fp8_e4m3", + "group_size": 128, + "symmetric": True, + "scale_type": None, + "scale_format": None, + "scale_calculation_mode": None, + "observer_cls": "PerGroupMinMaxObserver", + } + ) + fp8_config["weight"].update( + { + "dtype": "fp8_e4m3", + "qscheme": "per_block", + "ch_axis": None, + "group_size": None, + "block_size": [128, 128], + "symmetric": True, + "scale_type": "float8_e8m0fnu", + "scale_format": None, + "scale_calculation_mode": None, + "observer_cls": "PerBlock2DMinMaxObserver", + } + ) + quantization_config["layer_quant_config"] = { + r"re:mtp\.0\.ffn\.experts\.\d+\.w1": mxfp4_config, + r"re:mtp\.0\.ffn\.shared_experts\.w1": fp8_config, + } + return quantization_config + + +def get_fse_test_model_config( + model_type: str, + quantization_config: dict[str, Any], +) -> tuple[object, type[nn.Module]]: + if model_type == "minimax_m3": + config = MiniMaxM3TextConfig( + hidden_size=128, + intermediate_size=32, + dense_intermediate_size=32, + num_hidden_layers=2, + num_attention_heads=1, + num_key_value_heads=1, + head_dim=128, + num_local_experts=2, + num_experts_per_tok=1, + moe_layer_freq=[0, 1], + sparse_attention_config=None, + rotary_dim=64, + quantization_config=quantization_config, + ) + return config, minimax_m3_model.MiniMaxM3Model + if model_type == "deepseek_v4": + from vllm.models.deepseek_v4.amd.model import DeepseekV4Model + + config = SimpleNamespace( + vocab_size=256, + hidden_size=128, + num_hidden_layers=1, + num_attention_heads=1, + head_dim=128, + max_position_embeddings=128, + sliding_window=None, + compress_ratios=[1], + rope_theta=10000.0, + compress_rope_theta=10000.0, + rope_parameters={"rope_type": "default"}, + rms_norm_eps=1e-6, + hidden_act="silu", + q_lora_rank=0, + o_lora_rank=0, + o_groups=1, + qk_rope_head_dim=64, + index_head_dim=64, + index_n_heads=1, + index_topk=1, + n_routed_experts=2, + n_shared_experts=1, + num_experts_per_tok=1, + n_group=1, + topk_group=1, + moe_intermediate_size=32, + norm_topk_prob=False, + num_hash_layers=0, + swiglu_limit=0.0, + hc_eps=1e-6, + hc_mult=1, + hc_sinkhorn_iters=1, + expert_dtype="fp4", + num_nextn_predict_layers=1, + quantization_config=quantization_config, + ) + return config, DeepseekV4Model + if model_type == "qwen3_5": + from vllm.model_executor.models.qwen3_5 import Qwen3_5Model + + config = Qwen3_5MoeTextConfig( + vocab_size=256, + hidden_size=128, + num_hidden_layers=1, + num_attention_heads=1, + num_key_value_heads=1, + head_dim=128, + linear_key_head_dim=64, + linear_value_head_dim=64, + linear_num_key_heads=1, + linear_num_value_heads=1, + moe_intermediate_size=32, + shared_expert_intermediate_size=32, + num_experts_per_tok=1, + num_experts=2, + layer_types=["full_attention"], + quantization_config=quantization_config, + ) + return config, Qwen3_5Model + if model_type == "deepseek_v2": + from transformers import DeepseekV2Config + + from vllm.model_executor.models.deepseek_v2 import DeepseekV2Model + + config = DeepseekV2Config( + vocab_size=256, + hidden_size=128, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=1, + num_key_value_heads=1, + head_dim=128, + first_k_dense_replace=0, + n_routed_experts=2, + n_shared_experts=1, + num_experts_per_tok=1, + n_group=1, + topk_group=1, + moe_intermediate_size=32, + q_lora_rank=None, + kv_lora_rank=0, + qk_nope_head_dim=0, + qk_rope_head_dim=0, + v_head_dim=0, + quantization_config=quantization_config, + ) + return config, DeepseekV2Model + if model_type == "glm4_moe": + from transformers.models.glm4_moe import Glm4MoeConfig + + from vllm.model_executor.models.glm4_moe import Glm4MoeModel + + config = Glm4MoeConfig( + vocab_size=256, + hidden_size=128, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=1, + num_key_value_heads=1, + head_dim=128, + moe_intermediate_size=32, + n_routed_experts=2, + n_shared_experts=1, + num_experts_per_tok=1, + first_k_dense_replace=0, + quantization_config=quantization_config, + ) + return config, Glm4MoeModel + raise ValueError(f"Unsupported FSE test model: {model_type}") + + +def test_determine_expert_counts_fuse_shared_experts_override( + monkeypatch: pytest.MonkeyPatch, +) -> None: + common_args = (8, 0, 2) + assert determine_expert_counts(*common_args, True)[2] == 2 + assert determine_expert_counts(*common_args, False)[2] == 0 + + +def test_resolve_layer_fused_shared_expert_skips_compatibility_when_disabled( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + fused_moe_utils.rocm_aiter_ops, + "is_fusion_moe_shared_experts_enabled", + lambda: False, + ) + monkeypatch.setattr( + fused_moe_utils, + "is_shared_expert_quant_fse_compatible", + lambda *_: pytest.fail( + "compatibility must not be checked when FSE is disabled" + ), + ) + + assert not fused_moe_utils.resolve_layer_fused_shared_expert( + object(), "model.layers.0.mlp" + ) + + +def test_resolve_layer_fused_shared_expert_normalizes_unavailable_aiter( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + fused_moe_utils.rocm_aiter_ops, + "is_fusion_moe_shared_experts_enabled", + lambda: None, + ) + + assert ( + fused_moe_utils.resolve_layer_fused_shared_expert( + object(), "model.layers.0.mlp" + ) + is False + ) + + +def test_resolve_layer_fused_shared_expert_passes_module_prefixes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + quant_config = object() + monkeypatch.setattr( + fused_moe_utils.rocm_aiter_ops, + "is_fusion_moe_shared_experts_enabled", + lambda: True, + ) + + calls: list[tuple[object, str, str]] = [] + + def check_compatibility( + config: object, expert_prefix: str, shared_expert_prefix: str + ) -> tuple[bool, str | None]: + calls.append((config, expert_prefix, shared_expert_prefix)) + return True, None + + monkeypatch.setattr( + fused_moe_utils, "is_shared_expert_quant_fse_compatible", check_compatibility + ) + + assert fused_moe_utils.resolve_layer_fused_shared_expert( + quant_config, + "model.layers.0.mlp", + shared_expert_name="shared_expert", + ) + assert calls == [ + ( + quant_config, + "model.layers.0.mlp.experts", + "model.layers.0.mlp.shared_expert", + ) + ] + + +def test_resolve_layer_fused_shared_expert_rejects_incompatible_quantization( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + monkeypatch.setattr( + fused_moe_utils.rocm_aiter_ops, + "is_fusion_moe_shared_experts_enabled", + lambda: True, + ) + monkeypatch.setattr( + fused_moe_utils, + "is_shared_expert_quant_fse_compatible", + lambda *_: (False, "shared experts are excluded"), + ) + + assert not fused_moe_utils.resolve_layer_fused_shared_expert( + object(), "model.layers.0.mlp" + ) + assert "shared experts are excluded" in caplog.text + + +def test_deepseek_v4_shared_expert_fse_uses_mtp_quantization_config_prefix( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class DeepseekV4Config: + expert_dtype = "fp4" + + def _is_quark_mxfp4_ocp(self, hf_config: object) -> bool: + return True + + hf_config = SimpleNamespace( + num_hidden_layers=2, + quantization_config={ + "layer_quant_config": { + r"re:mtp\.0\.ffn\.shared_experts\.w1": {"weight": {"dtype": "fp4"}} + }, + "global_quant_config": {"weight": {"dtype": "fp8"}}, + }, + ) + monkeypatch.setattr( + deepseek_v4_quant_config, "DeepseekV4FP8Config", DeepseekV4Config + ) + monkeypatch.setattr( + vllm_config_module, + "get_current_vllm_config", + lambda: SimpleNamespace(model_config=SimpleNamespace(hf_config=hf_config)), + ) + + compatible, reason = is_shared_expert_quant_fse_compatible( + DeepseekV4Config(), + "model.layers.2.ffn.experts", + "model.layers.2.ffn.shared_experts", + ) + + assert compatible + assert reason is None + + +def test_is_model_fused_shared_expert_compatible() -> None: + class MoE(nn.Module): + def __init__(self, enabled: bool) -> None: + super().__init__() + self.is_fused_shared_expert_enabled = enabled + + class Layer(nn.Module): + def __init__(self, enabled: bool) -> None: + super().__init__() + self.mlp = MoE(enabled) + + enabled_layers = nn.ModuleList([Layer(True)]) + disabled_layers = nn.ModuleList([Layer(False)]) + mixed_layers = nn.ModuleList([Layer(True), Layer(False)]) + empty_layers = nn.ModuleList() + pipeline_layers = nn.ModuleList([Layer(True), PPMissingLayer()]) + + assert fused_moe_utils.is_model_fused_shared_expert_compatible( + enabled_layers, MoE, "mlp" + ) + assert not fused_moe_utils.is_model_fused_shared_expert_compatible( + disabled_layers, MoE, "mlp" + ) + assert not fused_moe_utils.is_model_fused_shared_expert_compatible( + empty_layers, MoE, "mlp" + ) + assert fused_moe_utils.is_model_fused_shared_expert_compatible( + pipeline_layers, MoE, "mlp" + ) + with pytest.raises(NotImplementedError, match="1 enabled and 1 disabled layers"): + fused_moe_utils.is_model_fused_shared_expert_compatible( + mixed_layers, MoE, "mlp" + ) + + +@pytest.mark.parametrize( + "model_type", + ["minimax_m3", "deepseek_v4", "qwen3_5", "glm4_moe", "deepseek_v2"], +) +@pytest.mark.parametrize( + ("use_fse", "exclude"), + [ + (False, []), + (True, []), + (True, ["*.shared_experts.*"]), + ], +) +def test_models_fse_init( + model_type: str, + use_fse: bool, + exclude: list[str], + dist_init: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Model construction resolves FSE consistently with Quark quantization.""" + + quantization_config: dict[str, Any] = ( + get_deepseek_v4_quark_config(["layers.0.ffn.shared_experts"] if exclude else []) + if model_type == "deepseek_v4" + else {**_QUARK_FSE_CONFIG, "exclude": exclude} + ) + + config, model_constructor = get_fse_test_model_config( + model_type, quantization_config + ) + vllm_config = VllmConfig() + vllm_config.model_config = SimpleNamespace( + hf_config=config, + hf_text_config=config, + dtype=torch.bfloat16, + max_model_len=128, + is_diffusion=False, + head_dtype=None, + is_mm_prefix_lm=False, + multimodal_config=None, + quantization_config=None, + ) + vllm_config.parallel_config.enable_expert_parallel = False + if model_type == "deepseek_v4": + from vllm.models.deepseek_v4.quant_config import DeepseekV4FP8Config + + vllm_config.cache_config.cache_dtype = "fp8_ds_mla" + vllm_config.quant_config = DeepseekV4FP8Config( + is_checkpoint_fp8_serialized=True, + weight_block_size=[128, 128], + ) + else: + vllm_config.quant_config = QuarkConfig(quantization_config) + + import vllm.envs as envs + from vllm._aiter_ops import rocm_aiter_ops + + if model_type == "deepseek_v4": + from vllm.models.deepseek_v4.amd import model as deepseek_v4_model + + warning_logger = deepseek_v4_model.logger + elif model_type == "minimax_m3": + warning_logger = minimax_m3_model.logger + else: + warning_logger = fused_moe_utils.logger + + with monkeypatch.context() as mp: + mp.setenv("VLLM_ROCM_USE_AITER", str(use_fse)) + mp.setenv("VLLM_ROCM_USE_AITER_MOE", str(use_fse)) + mp.setenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", str(use_fse)) + mp.setattr( + "vllm.model_executor.layers.fused_moe.experts." + "ocp_mx_emulation_moe.has_quark", + lambda: True, + ) + importlib.reload(envs) + mp.setattr(envs, "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", use_fse) + rocm_aiter_ops.refresh_env_variables() + aiter_fse_enabled = bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()) + # These AMD-specific models currently use the raw environment flag. + # and do not rely on `rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()`. + fse_enabled = use_fse and ( + aiter_fse_enabled or model_type in {"deepseek_v4", "minimax_m3"} + ) + + with ( + patch.object(warning_logger, "warning") as warning, + set_current_vllm_config(vllm_config), + ): + model = model_constructor(vllm_config=vllm_config) + mtp = None + if model_type == "deepseek_v4" and use_fse and not exclude: + from vllm.models.deepseek_v4.amd.mtp import ( + DeepSeekV4MTP, + ) + + vllm_config.speculative_config = SimpleNamespace( + draft_model_config=SimpleNamespace(hf_config=config) + ) + mtp = DeepSeekV4MTP(vllm_config=vllm_config) + assert model.is_fused_shared_expert_enabled is (fse_enabled and not exclude) + + # The dummy quant config here uses mixed mxfp4/fp8 for experts/shared_expert + # so should just raise a warning. + if mtp is not None: + assert not mtp.model.layers[ + "1" + ].mtp_block.ffn.is_fused_shared_expert_enabled + warning.assert_called_once() + assert ( + "DeepSeek-V4 shared experts at mtp.0.ffn.shared_experts" + in (warning.call_args.args[1]) + ) + if aiter_fse_enabled and exclude: + warning.assert_called_once() + assert ( + "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled" + in (warning.call_args.args[0]) + ) + assert "excludes shared experts" in warning.call_args.args[1] + + importlib.reload(envs) + rocm_aiter_ops.refresh_env_variables() + + +@pytest.mark.parametrize( + ("exclude", "expected"), + [([], True), (["*.shared_expert.*"], False)], +) +def test_quark_shared_expert_fse_compatibility( + exclude: list[str], expected: bool +) -> None: + compatible, reason = is_shared_expert_quant_fse_compatible( + QuarkConfig( + { + "exclude": exclude, + "global_quant_config": {}, + "layer_quant_config": {}, + } + ), + "model.layers.0.mlp.experts", + "model.layers.0.mlp.shared_expert", + ) + + assert compatible is expected + if expected: + assert reason is None + else: + assert ( + reason + == "Quark excludes shared experts at model.layers.0.mlp.shared_expert" + ) + + +def test_quark_shared_expert_fse_requires_matching_layer_quant_configs() -> None: + global_quant_config = {"weight": {"dtype": "fp4"}} + quant_config = QuarkConfig( + { + "exclude": [], + "global_quant_config": global_quant_config, + "layer_quant_config": { + "model.layers.0.mlp.shared_expert.gate_up_proj": { + "weight": {"dtype": "fp8"} + }, + "model.layers.0.mlp.shared_expert.down_proj": { + "weight": {"dtype": "fp8"} + }, + }, + } + ) + + compatible, reason = is_shared_expert_quant_fse_compatible( + quant_config, + "model.layers.0.mlp.experts", + "model.layers.0.mlp.shared_expert", + ) + + assert not compatible + assert reason == ( + "Quark uses different quantization configurations for routed and " + "shared experts at model.layers.0.mlp.shared_expert" + ) + + +def test_quark_shared_expert_fse_rejects_partial_packed_projection_override() -> None: + fp4_config = {"weight": {"dtype": "fp4"}} + fp8_config = {"weight": {"dtype": "fp8"}} + quant_config = QuarkConfig( + { + "exclude": [], + "global_quant_config": fp4_config, + "layer_quant_config": { + "model.layers.0.mlp.experts": fp8_config, + "model.layers.0.mlp.shared_expert.gate_proj": fp8_config, + "model.layers.0.mlp.shared_expert.down_proj": fp8_config, + }, + } + ) + quant_config.packed_modules_mapping = {"gate_up_proj": ["gate_proj", "up_proj"]} + + compatible, reason = is_shared_expert_quant_fse_compatible( + quant_config, + "model.layers.0.mlp.experts", + "model.layers.0.mlp.shared_expert", + ) + + assert not compatible + assert reason == ( + "Quark uses different quantization configurations for routed and " + "shared experts at model.layers.0.mlp.shared_expert" + ) + + +def test_quark_layer_config_from_name_checks_packed_projections() -> None: + quant_config = QuarkConfig( + { + "global_quant_config": {"weight": {"dtype": "fp4"}}, + "layer_quant_config": { + "model.layers.0.mlp.shared_expert.w1": {"weight": {"dtype": "fp4"}}, + "model.layers.0.mlp.shared_expert.w3": {"weight": {"dtype": "fp8"}}, + }, + } + ) + quant_config.packed_modules_mapping = {"gate_up_proj": ["w1", "w3"]} + + with pytest.raises(ValueError, match="requires all to use the same scheme"): + quant_config.get_layer_quant_config_from_name( + "model.layers.0.mlp.shared_expert.gate_up_proj" + ) + + +def test_quark_packed_layer_config_must_match_global_config() -> None: + quant_config = QuarkConfig( + { + "global_quant_config": {"weight": {"dtype": "fp4"}}, + "layer_type_quant_config": {}, + "layer_quant_config": { + "model.layers.0.mlp.shared_expert.w1": {"weight": {"dtype": "fp8"}}, + }, + } + ) + quant_config.packed_modules_mapping = {"gate_up_proj": ["w1", "w3"]} + + with pytest.raises(ValueError, match="requires all to use the same scheme"): + quant_config._find_matched_config( + "model.layers.0.mlp.shared_expert.gate_up_proj", nn.Module() + ) + + +def test_non_quark_shared_expert_fse_is_incompatible() -> None: + compatible, reason = is_shared_expert_quant_fse_compatible( + object(), + "model.layers.0.mlp.experts", + "model.layers.0.mlp.shared_experts", + ) + + assert not compatible + assert reason == ( + "shared-expert FSE quantization compatibility is not implemented for object" + ) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 28b994057fa1..d9fd1385a91f 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -6,7 +6,6 @@ import torch -import vllm.envs as envs from vllm._aiter_ops import rocm_aiter_ops from vllm.config import ParallelConfig, get_current_vllm_config from vllm.distributed import ( @@ -74,23 +73,13 @@ def determine_expert_counts( num_experts: int, num_redundant_experts: int, n_shared_experts: int | None, - is_act_and_mul: bool, + fuse_shared_experts: bool, ) -> tuple[int, int, int]: global_num_experts = num_experts + num_redundant_experts logical_num_experts = num_experts - # Shared-expert fusion: append the shared expert(s) as routed-expert slots - # so they run in the same grouped GEMM. Gated by - # VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: either the native aiter fused-MoE - # path (env + master switch, via is_fusion_moe_shared_experts_enabled) or the - # backend-neutral router-append path (env alone, independent of the master - # switch; e.g. the MM3 triton/flydsl mxfp8 MoE). Gated activations only. - fuse_shared_enabled = ( - rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() - or envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS - ) and is_act_and_mul num_fused_shared_experts = ( - n_shared_experts if n_shared_experts is not None and fuse_shared_enabled else 0 + n_shared_experts if n_shared_experts is not None and fuse_shared_experts else 0 ) return global_num_experts, logical_num_experts, num_fused_shared_experts @@ -132,6 +121,7 @@ def FusedMoEFactory( ckpt_names: tuple[str, str, str] = ("gate_proj", "down_proj", "up_proj"), is_fused_checkpoint_transposed: bool = False, n_shared_experts: int | None = None, + fuse_shared_experts: bool = False, router_logits_dtype: torch.dtype | None = None, gate: torch.nn.Module | None = None, shared_experts: torch.nn.Module | None = None, @@ -198,6 +188,7 @@ def FusedMoEFactory( block scales use transposed storage. n_shared_experts: Number of shared experts to fuse into the routed grouped GEMM (ROCm; requires aiter FSE or the router-append path) + fuse_shared_experts: Whether to enable shared-expert fusion. router_logits_dtype: Data type for router logits buffers gate: Pre-configured gate module shared_experts: Pre-configured shared experts module @@ -244,7 +235,7 @@ def FusedMoEFactory( num_experts, num_redundant_experts, n_shared_experts, - is_act_and_mul, + fuse_shared_experts, ) ) diff --git a/vllm/model_executor/layers/fused_moe/utils.py b/vllm/model_executor/layers/fused_moe/utils.py index cce8ccd073fc..b873d8321d8b 100644 --- a/vllm/model_executor/layers/fused_moe/utils.py +++ b/vllm/model_executor/layers/fused_moe/utils.py @@ -1,15 +1,21 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import functools +from collections.abc import Iterable from math import prod from typing import TYPE_CHECKING import torch import torch.nn.functional as F +from torch import nn import vllm.envs as envs from vllm import _custom_ops as ops +from vllm._aiter_ops import rocm_aiter_ops from vllm.logger import init_logger +from vllm.model_executor.layers.quantization.utils.config_utils import ( + is_shared_expert_quant_fse_compatible, +) from vllm.model_executor.layers.quantization.utils.fp8_utils import ( per_token_group_quant_fp8, ) @@ -34,16 +40,95 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( per_tensor_dequantize, ) +from vllm.model_executor.models.utils import PPMissingLayer from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.utils.math_utils import cdiv if TYPE_CHECKING: from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig + from vllm.model_executor.layers.quantization import QuantizationConfig logger = init_logger(__name__) +def resolve_layer_fused_shared_expert( + quant_config: "QuantizationConfig | None", + prefix: str, + shared_expert_name: str = "shared_experts", +) -> bool: + """Resolve whether AITER fused shared-expert execution is enabled. + + Args: + quant_config: Model quantization configuration. + prefix: MoE module prefix. + shared_expert_name: Shared-expert module name under ``prefix``. + + Returns: + Whether AITER fused shared experts are enabled. + + Raises: + ValueError: If requested shared-expert fusion is quantization-incompatible. + """ + # NOTE: is_fusion_moe_shared_experts_enabled is decorated with @if_aiter_supported + # that returns None if AITER is not available. + fse_requested = bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()) + fse_compatible, fse_reason = ( + is_shared_expert_quant_fse_compatible( + quant_config, + f"{prefix}.experts", + f"{prefix}.{shared_expert_name}", + ) + if fse_requested + else (True, None) + ) + is_fused_shared_expert_enabled = fse_requested and fse_compatible + if fse_requested and not is_fused_shared_expert_enabled: + logger.warning( + "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but " + "cannot be enabled: %s.", + fse_reason, + ) + return is_fused_shared_expert_enabled + + +def is_model_fused_shared_expert_compatible( + layers: nn.ModuleList | Iterable[nn.Module], + moe_cls: type[nn.Module], + moe_name: str, +) -> bool: + """Resolve one fused-shared-expert state for a model's MoE layers.""" + + def get_moe_layer(layer: nn.Module) -> nn.Module | None: + for name in moe_name.split("."): + layer = getattr(layer, name, None) + if layer is None: + return None + return layer + + moe_layers = ( + moe_layer + for layer in layers + if not isinstance(layer, PPMissingLayer) + and (moe_layer := get_moe_layer(layer)) is not None + and isinstance(moe_layer, moe_cls) + ) + + enabled = [ + getattr(layer, "is_fused_shared_expert_enabled", False) for layer in moe_layers + ] + enabled_count = sum(enabled) + disabled_count = len(enabled) - enabled_count + if enabled_count > 0 and disabled_count > 0: + raise NotImplementedError( + "Fused shared experts must be enabled for all MoE layers; found " + f"{enabled_count} enabled and {disabled_count} disabled layers. " + "Per-layer fused shared experts is not yet supported. Please open " + "an issue." + ) + return enabled_count > 0 and disabled_count == 0 + + @triton.jit def _count_expert_num_tokens( topk_ids_ptr, diff --git a/vllm/model_executor/layers/quantization/quark/quark.py b/vllm/model_executor/layers/quantization/quark/quark.py index 9a65432c5c42..912f61ac386a 100644 --- a/vllm/model_executor/layers/quantization/quark/quark.py +++ b/vllm/model_executor/layers/quantization/quark/quark.py @@ -576,63 +576,91 @@ def _is_w_ocp_mx_a_x( return True - def _find_matched_config( - self, layer_name: str, module: torch.nn.Module - ) -> dict[str, Any]: + def get_layer_quant_config_from_name( + self, layer_name: str + ) -> dict[str, Any] | None: proj_name = layer_name.split(".")[-1] if proj_name in self.packed_modules_mapping: shard_proj_names = self.packed_modules_mapping[proj_name] - - # Convert fused_name --> [shard_names] - shard_names = [ - layer_name.replace(proj_name, shard_proj_name) - for shard_proj_name in shard_proj_names - ] - shard_configs = [] - for shard_name in shard_names: - if shard_name == layer_name: - config = cast( - dict[str, Any], self.quant_config.get("global_quant_config") - ) + for shard_proj_name in shard_proj_names: + shard_name = layer_name.replace(proj_name, shard_proj_name) + if shard_name != layer_name: + config = self.get_layer_quant_config_from_name(shard_name) else: - config = self._find_matched_config(shard_name, module) + config = None shard_configs.append(config) - if not all( - deep_compare(q_config, shard_configs[0]) for q_config in shard_configs + matched_configs = [config for config in shard_configs if config is not None] + if matched_configs and not all( + deep_compare(config, matched_configs[0]) for config in matched_configs ): raise ValueError( f"Found a different quantization configuration for " f"{shard_proj_names} in {layer_name}. vLLM " "requires all to use the same scheme." ) - return shard_configs[0] + if matched_configs: + return matched_configs[0] + return None else: layer_quant_config = cast( - dict[str, Any], self.quant_config.get("layer_quant_config") + dict[str, Any], self.quant_config.get("layer_quant_config") or {} ) - - def _matches_pattern(layer_name, pattern): - if "*" not in pattern: - return layer_name in pattern - return fnmatch.fnmatch(layer_name, pattern) - for name_pattern, config in layer_quant_config.items(): - if _matches_pattern(layer_name, name_pattern): + if "*" not in name_pattern: + matches = layer_name in name_pattern + else: + matches = fnmatch.fnmatch(layer_name, name_pattern) + if matches: return config + return None - layer_type = cast(str, type(module)) - layer_type_quant_config = cast( - dict[str, Any], self.quant_config.get("layer_type_quant_config") - ) - if layer_type in layer_type_quant_config: - return layer_type_quant_config[layer_type] + def _find_matched_config( + self, layer_name: str, module: torch.nn.Module + ) -> dict[str, Any]: + # Priority order: + # 1. layer_quant_config, + # 2. layer_type_quant_config, + # 3. global_quant_config. + + layer_type = cast(str, type(module)) + layer_type_quant_config = cast( + dict[str, Any], self.quant_config.get("layer_type_quant_config") + ) + global_quant_config = cast( + dict[str, Any], self.quant_config.get("global_quant_config") + ) + fallback_config = layer_type_quant_config.get(layer_type, global_quant_config) - global_quant_config = cast( - dict[str, Any], self.quant_config.get("global_quant_config") - ) - return global_quant_config + proj_name = layer_name.split(".")[-1] + if proj_name in self.packed_modules_mapping: + shard_proj_names = self.packed_modules_mapping[proj_name] + shard_configs = [] + for shard_proj_name in shard_proj_names: + shard_name = layer_name.replace(proj_name, shard_proj_name) + if shard_name == layer_name: + config = fallback_config + else: + config = self.get_layer_quant_config_from_name(shard_name) + if config is None: + config = fallback_config + shard_configs.append(config) + + if not all( + deep_compare(config, shard_configs[0]) for config in shard_configs + ): + raise ValueError( + f"Found a different quantization configuration for " + f"{shard_proj_names} in {layer_name}. vLLM requires all " + "to use the same scheme." + ) + return shard_configs[0] + else: + layer_quant_config = self.get_layer_quant_config_from_name(layer_name) + if layer_quant_config is not None: + return layer_quant_config + return fallback_config def _get_scheme_from_config( self, config: dict[str, Any], dynamic_mxfp4_quant: bool = False diff --git a/vllm/model_executor/layers/quantization/utils/config_utils.py b/vllm/model_executor/layers/quantization/utils/config_utils.py new file mode 100644 index 000000000000..00cdd0c87dc4 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/config_utils.py @@ -0,0 +1,155 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import TYPE_CHECKING + +import regex as re + +if TYPE_CHECKING: + from vllm.model_executor.layers.quantization.base_config import QuantizationConfig + + +def is_shared_expert_quant_fse_compatible( + quant_config: "QuantizationConfig | None", + expert_prefix: str, + shared_expert_prefix: str, + projection_names: list[str] | None = None, +) -> tuple[bool, str | None]: + """Check whether quantization permits fused shared-expert execution. + + Args: + quant_config: Model quantization configuration. + expert_prefix: Routed-expert module prefix. + shared_expert_prefix: Shared-expert module prefix. + projection_names: Shared-expert projection names. + + Returns: + A compatibility flag and, when incompatible, the reason. + """ + if projection_names is None: + projection_names = ["gate_up_proj", "down_proj"] + + if quant_config is None: + return True, None + + from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig + from vllm.models.deepseek_v4.quant_config import DeepseekV4FP8Config + + if isinstance(quant_config, DeepseekV4FP8Config): + from vllm.config import get_current_vllm_config + from vllm.model_executor.models.utils import extract_layer_index + + if quant_config.expert_dtype != "fp4": + return False, "DeepSeek-V4 routed experts are not MXFP4" + + hf_config = get_current_vllm_config().model_config.hf_config + + # TODO: This is adapted from former `_shared_experts_are_fp4`, and + # needs to be cleaned up this . There should not be Quark-specific + # logic in DeepseekV4FP8Config. + quantization_config = getattr(hf_config, "quantization_config", None) + if quantization_config is None: + return False, "DeepSeek-V4 has no quantization configuration" + + if not quant_config._is_quark_mxfp4_ocp(quantization_config): + return False, "DeepSeek-v4 FSE is only implemented/tested with Quark MXFP4" + + layer_idx = extract_layer_index(shared_expert_prefix) + if layer_idx >= hf_config.num_hidden_layers: + shared_expert_prefix = ( + f"mtp.{layer_idx - hf_config.num_hidden_layers}.ffn.shared_experts" + ) + else: + shared_expert_prefix = f"layers.{layer_idx}.ffn.shared_experts" + + if any( + entry.startswith(shared_expert_prefix) + for entry in quantization_config.get("exclude") or [] + if isinstance(entry, str) + ): + return ( + False, + f"DeepSeek-V4 excludes shared experts at {shared_expert_prefix}", + ) + + shared_expert_weight_name = f"{shared_expert_prefix}.w1" + layer_quant_config = quantization_config.get("layer_quant_config") or {} + layer_config = layer_quant_config.get(shared_expert_weight_name) + if layer_config is None: + layer_config = next( + ( + config + for pattern, config in layer_quant_config.items() + if isinstance(pattern, str) + and pattern.startswith("re:") + and re.fullmatch( + pattern.removeprefix("re:"), shared_expert_weight_name + ) + ), + None, + ) + shared_weight_config = ( + layer_config or quantization_config.get("global_quant_config") or {} + ).get("weight") or {} + if shared_weight_config.get("dtype") == "fp4": + return True, None + return ( + False, + f"DeepSeek-V4 shared experts at {shared_expert_prefix} are not MXFP4", + ) + + if isinstance(quant_config, QuarkConfig): + # TODO: layer_type_quant_config is not taken into account here. + assert "exclude" in quant_config.quant_config + assert "global_quant_config" in quant_config.quant_config + + is_compatible = not any( + "shared_expert" in str(entry) + for entry in quant_config.quant_config["exclude"] + ) + if not is_compatible: + return False, f"Quark excludes shared experts at {shared_expert_prefix}" + + global_quant_config = quant_config.quant_config["global_quant_config"] + + def get_projection_quant_configs(layer_name: str) -> list[object]: + module_prefix, _, projection_name = layer_name.rpartition(".") + packed_projection_names = quant_config.packed_modules_mapping.get( + projection_name, [projection_name] + ) + return [ + quant_config.get_layer_quant_config_from_name( + f"{module_prefix}.{packed_projection_name}" + ) + or global_quant_config + for packed_projection_name in packed_projection_names + ] + + expert_quant_config = ( + quant_config.get_layer_quant_config_from_name(expert_prefix) + or global_quant_config + ) + shared_expert_quant_configs = [ + config + for projection_name in projection_names + for config in get_projection_quant_configs( + f"{shared_expert_prefix}.{projection_name}" + ) + ] + if all(config == expert_quant_config for config in shared_expert_quant_configs): + return True, None + return ( + False, + "Quark uses different quantization configurations for routed and " + f"shared experts at {shared_expert_prefix}", + ) + + # TODO: Extend FSE support detection to other quantization methods. Typically, + # one would check that the experts and shared_experts use the same + # quantization config. This may be refactored as part of QuantizationConfig later. + + return ( + False, + "shared-expert FSE quantization compatibility is not implemented for " + f"{type(quant_config).__name__}", + ) diff --git a/vllm/model_executor/models/AXK1.py b/vllm/model_executor/models/AXK1.py index 6000e62828f9..031808ee12b8 100644 --- a/vllm/model_executor/models/AXK1.py +++ b/vllm/model_executor/models/AXK1.py @@ -46,6 +46,10 @@ FusedMoEFactory, fused_moe_make_expert_params_mapping, ) +from vllm.model_executor.layers.fused_moe.utils import ( + is_model_fused_shared_expert_compatible, + resolve_layer_fused_shared_expert, +) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( ColumnParallelLinear, @@ -148,10 +152,14 @@ def __init__( self.n_local_physical_experts = self.n_physical_experts // self.ep_size self.is_rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled() - self.is_fusion_moe_shared_experts_enabled = ( - rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() - ) - if config.n_shared_experts is None or self.is_fusion_moe_shared_experts_enabled: + + self.is_fused_shared_expert_enabled = False + if config.n_shared_experts is not None: + self.is_fused_shared_expert_enabled = resolve_layer_fused_shared_expert( + quant_config, prefix + ) + + if config.n_shared_experts is None or self.is_fused_shared_expert_enabled: self.shared_experts = None else: intermediate_size = config.moe_intermediate_size * config.n_shared_experts @@ -189,8 +197,9 @@ def __init__( num_redundant_experts=self.n_redundant_experts, is_sequence_parallel=self.is_sequence_parallel, n_shared_experts=config.n_shared_experts - if self.is_fusion_moe_shared_experts_enabled + if self.is_fused_shared_expert_enabled else None, + fuse_shared_experts=self.is_fused_shared_expert_enabled, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: @@ -725,6 +734,12 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=f"{prefix}.layers", ) + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( + self.layers, + AXK1MoE, + "mlp", + ) + if get_pp_group().is_last_rank: self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) else: @@ -789,9 +804,6 @@ def forward( return hidden_states def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - rocm_aiter_moe_shared_expert_enabled = ( - rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() - ) stacked_params_mapping: list[tuple[str, str, int | str]] = [ # (param_name, shard_name, shard_id) ("gate_up_proj", "gate_proj", 0), @@ -822,7 +834,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: num_experts=self.config.n_routed_experts + ( (self.config.n_shared_experts or 0) - if rocm_aiter_moe_shared_expert_enabled + if self.is_fused_shared_expert_enabled else 0 ), num_redundant_experts=self.num_redundant_experts, @@ -839,7 +851,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: continue # skip spec decode layers for main model is_fusion_moe_shared_experts_layer = ( - rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name) + self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name) ) for param_name, weight_name, shard_id in stacked_params_mapping: diff --git a/vllm/model_executor/models/deepseek_mtp.py b/vllm/model_executor/models/deepseek_mtp.py index 7fc78f13ff5a..8ac43e0bb2a1 100644 --- a/vllm/model_executor/models/deepseek_mtp.py +++ b/vllm/model_executor/models/deepseek_mtp.py @@ -7,13 +7,15 @@ import torch.nn as nn from transformers import PretrainedConfig -from vllm._aiter_ops import rocm_aiter_ops from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig from vllm.distributed import tensor_model_parallel_all_gather from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, ) +from vllm.model_executor.layers.fused_moe.utils import ( + is_model_fused_shared_expert_compatible, +) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization import QuantizationConfig @@ -255,6 +257,11 @@ def set_moe_parameters(self): self.moe_mlp_layers.append(layer.mlp) self.moe_layers.append(layer.mlp.experts) self.extract_moe_parameters(example_moe) + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( + self.model.layers.values(), + DeepseekV2MoE, + "mtp_block.mlp", + ) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.model.embed_input_ids(input_ids) @@ -285,9 +292,6 @@ def compute_logits( return self.model.compute_logits(hidden_states, spec_step_idx) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - rocm_aiter_moe_shared_expert_enabled = ( - rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() - ) stacked_params_mapping = [ ("gate_up_proj", "gate_proj", 0), ("gate_up_proj", "up_proj", 1), @@ -310,7 +314,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: num_experts=self.config.n_routed_experts + ( self.config.n_shared_experts - if rocm_aiter_moe_shared_expert_enabled + if self.is_fused_shared_expert_enabled else 0 ), ) @@ -326,7 +330,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: if spec_layer is None: continue is_fusion_moe_shared_experts_layer = ( - rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name) + self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name) ) name = self._rewrite_spec_layer_name(spec_layer, name) diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 81e941d28450..d5521a1960b8 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -54,6 +54,10 @@ GateLinear, fused_moe_make_expert_params_mapping, ) +from vllm.model_executor.layers.fused_moe.utils import ( + is_model_fused_shared_expert_compatible, + resolve_layer_fused_shared_expert, +) from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm from vllm.model_executor.layers.linear import ( ColumnParallelLinear, @@ -329,9 +333,13 @@ def __init__( self.n_local_physical_experts = self.n_physical_experts // self.ep_size self.is_rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled() - self.is_fusion_moe_shared_experts_enabled = ( - rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() - ) + + self.is_fused_shared_expert_enabled = False + if config.n_shared_experts is not None: + self.is_fused_shared_expert_enabled = resolve_layer_fused_shared_expert( + quant_config, prefix + ) + if ( self.is_rocm_aiter_moe_enabled and self.gate.e_score_correction_bias is not None @@ -340,7 +348,7 @@ def __init__( # Accumulates in fp32; avoids bf16->fp32 cast. self.gate.set_out_dtype(self.gate.weight.dtype) - if config.n_shared_experts is None or self.is_fusion_moe_shared_experts_enabled: + if config.n_shared_experts is None or self.is_fused_shared_expert_enabled: self.shared_experts = None else: intermediate_size = config.moe_intermediate_size * config.n_shared_experts @@ -377,8 +385,9 @@ def __init__( is_sequence_parallel=self.is_sequence_parallel, reduce_results=reduce_results, n_shared_experts=config.n_shared_experts - if self.is_fusion_moe_shared_experts_enabled + if self.is_fused_shared_expert_enabled else None, + fuse_shared_experts=self.is_fused_shared_expert_enabled, router_logits_dtype=self.gate.out_dtype, ) @@ -1401,6 +1410,12 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=f"{prefix}.layers", ) + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( + self.layers, + DeepseekV2MoE, + "mlp", + ) + if get_pp_group().is_last_rank: self.norm = RMSNorm(self.hidden_size, eps=config.rms_norm_eps) else: @@ -1511,9 +1526,6 @@ def forward( return hidden_states def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - rocm_aiter_moe_shared_expert_enabled = ( - rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() - ) stacked_params_mapping: list[tuple[str, str, int | str]] = [ # (param_name, shard_name, shard_id) ("gate_up_proj", "gate_proj", 0), @@ -1554,7 +1566,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: num_experts=self.config.n_routed_experts + ( self.config.n_shared_experts - if rocm_aiter_moe_shared_expert_enabled + if self.is_fused_shared_expert_enabled else 0 ), num_redundant_experts=self.num_redundant_experts, @@ -1582,7 +1594,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: continue # this layer has no indexer; drop its checkpoint weights is_fusion_moe_shared_experts_layer = ( - rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name) + self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name) ) if _try_load_fp8_indexer_wk( diff --git a/vllm/model_executor/models/glm4_moe.py b/vllm/model_executor/models/glm4_moe.py index 6272db2a9015..2192f74b5111 100644 --- a/vllm/model_executor/models/glm4_moe.py +++ b/vllm/model_executor/models/glm4_moe.py @@ -39,12 +39,15 @@ get_pp_group, get_tensor_model_parallel_world_size, ) -from vllm.logger import init_logger from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.fused_moe import ( FusedMoEFactory, ) +from vllm.model_executor.layers.fused_moe.utils import ( + is_model_fused_shared_expert_compatible, + resolve_layer_fused_shared_expert, +) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, @@ -72,8 +75,6 @@ skip_spec_layers, ) -logger = init_logger(__name__) - class Glm4MoeMLP(nn.Module): def __init__( @@ -159,14 +160,15 @@ def __init__( self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts self.n_local_physical_experts = self.n_physical_experts // self.ep_size - # AITER fused shared-expert (FSE) gate; mirrors the deepseek_v2.py - # pattern (see Glm4MoE / MoERunner wiring there). self.is_rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled() - self.is_fusion_moe_shared_experts_enabled = ( - rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() - ) - if config.n_shared_experts is None or self.is_fusion_moe_shared_experts_enabled: + self.is_fused_shared_expert_enabled = False + if config.n_shared_experts is not None: + self.is_fused_shared_expert_enabled = resolve_layer_fused_shared_expert( + quant_config, prefix + ) + + if config.n_shared_experts is None or self.is_fused_shared_expert_enabled: self.shared_experts = None else: intermediate_size = config.moe_intermediate_size * config.n_shared_experts @@ -200,10 +202,9 @@ def __init__( num_redundant_experts=self.n_redundant_experts, router_logits_dtype=torch.float32, n_shared_experts=( - config.n_shared_experts - if self.is_fusion_moe_shared_experts_enabled - else None + config.n_shared_experts if self.is_fused_shared_expert_enabled else None ), + fuse_shared_experts=self.is_fused_shared_expert_enabled, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: @@ -443,6 +444,12 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=f"{prefix}.layers", ) + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( + self.layers, + Glm4MoE, + "mlp", + ) + if get_pp_group().is_last_rank: self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) else: @@ -488,6 +495,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: skip_spec_layers(weights, self.config), n_routed_experts=self.config.n_routed_experts, n_shared_experts=self.config.n_shared_experts or 1, + enabled=self.is_fused_shared_expert_enabled, ) loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/glm4_moe_lite.py b/vllm/model_executor/models/glm4_moe_lite.py index 7f63130d883d..e51b0e231cc9 100644 --- a/vllm/model_executor/models/glm4_moe_lite.py +++ b/vllm/model_executor/models/glm4_moe_lite.py @@ -34,7 +34,6 @@ if TYPE_CHECKING: from transformers.models.glm4_moe_lite import Glm4MoeLiteConfig -from vllm._aiter_ops import rocm_aiter_ops from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig from vllm.distributed import ( @@ -44,6 +43,9 @@ from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, ) +from vllm.model_executor.layers.fused_moe.utils import ( + is_model_fused_shared_expert_compatible, +) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ( @@ -254,6 +256,12 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=f"{prefix}.layers", ) + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( + self.layers, + Glm4MoeLite, + "mlp", + ) + if get_pp_group().is_last_rank: self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) else: @@ -320,9 +328,6 @@ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - rocm_aiter_moe_shared_expert_enabled = ( - rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() - ) stacked_params_mapping = [ # (param_name, shard_name, shard_id) ("gate_up_proj", "gate_proj", 0), @@ -345,7 +350,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: num_experts=self.config.n_routed_experts + ( self.config.n_shared_experts - if rocm_aiter_moe_shared_expert_enabled + if self.is_fused_shared_expert_enabled else 0 ), ) @@ -361,7 +366,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: continue # skip spec decode layers for main model is_fusion_moe_shared_experts_layer = ( - rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name) + self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name) ) for param_name, weight_name, shard_id in stacked_params_mapping: diff --git a/vllm/model_executor/models/glm4_moe_lite_mtp.py b/vllm/model_executor/models/glm4_moe_lite_mtp.py index bbede18c40ad..6b28c38b5877 100644 --- a/vllm/model_executor/models/glm4_moe_lite_mtp.py +++ b/vllm/model_executor/models/glm4_moe_lite_mtp.py @@ -30,12 +30,14 @@ import torch.nn as nn from transformers import PretrainedConfig -from vllm._aiter_ops import rocm_aiter_ops from vllm.config import VllmConfig from vllm.model_executor.layers.fused_moe import ( MoERunner, fused_moe_make_expert_params_mapping, ) +from vllm.model_executor.layers.fused_moe.utils import ( + is_model_fused_shared_expert_compatible, +) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization import QuantizationConfig @@ -224,6 +226,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.moe_mlp_layers.append(layer.mlp) self.moe_layers.append(layer.mlp.experts) self.extract_moe_parameters(example_moe) + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( + self.model.layers.values(), + Glm4MoeLite, + "mtp_block.mlp", + ) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.model.embed_input_ids(input_ids) @@ -250,9 +257,6 @@ def compute_logits( return self.model.compute_logits(hidden_states, spec_step_idx) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - rocm_aiter_moe_shared_expert_enabled = ( - rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() - ) stacked_params_mapping = [ ("gate_up_proj", "gate_proj", 0), ("gate_up_proj", "up_proj", 1), @@ -268,7 +272,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: num_experts=self.config.n_routed_experts + ( self.config.n_shared_experts - if rocm_aiter_moe_shared_expert_enabled + if self.is_fused_shared_expert_enabled else 0 ), ) @@ -282,7 +286,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: if spec_layer is None: continue is_fusion_moe_shared_experts_layer = ( - rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name) + self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name) ) name = self._rewrite_spec_layer_name(spec_layer, name) for param_name, weight_name, shard_id in stacked_params_mapping: diff --git a/vllm/model_executor/models/glm4_moe_mtp.py b/vllm/model_executor/models/glm4_moe_mtp.py index 6708ba8139a7..8af40cf60bd5 100644 --- a/vllm/model_executor/models/glm4_moe_mtp.py +++ b/vllm/model_executor/models/glm4_moe_mtp.py @@ -31,12 +31,14 @@ import torch.nn as nn from transformers import PretrainedConfig -from vllm._aiter_ops import rocm_aiter_ops from vllm.config import CacheConfig, ParallelConfig, VllmConfig from vllm.model_executor.layers.fused_moe import ( MoERunner, fused_moe_make_expert_params_mapping, ) +from vllm.model_executor.layers.fused_moe.utils import ( + is_model_fused_shared_expert_compatible, +) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization import QuantizationConfig @@ -212,6 +214,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.moe_mlp_layers.append(layer.mlp) self.moe_layers.append(layer.mlp.experts) self.extract_moe_parameters(example_moe) + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( + self.model.layers.values(), + Glm4MoE, + "mtp_block.mlp", + ) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.model.embed_input_ids(input_ids) @@ -238,10 +245,6 @@ def compute_logits( return self.model.compute_logits(hidden_states, spec_step_idx) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - # FSE weight loading mirrors glm4_moe.py / deepseek_mtp.py. - rocm_aiter_moe_shared_expert_enabled = ( - rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() - ) stacked_params_mapping = [ # (param_name, shard_name, shard_id) ("qkv_proj", "q_proj", "q"), @@ -254,7 +257,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # Params for weights, fp8 weight scales, fp8 activation scales # (param_name, weight_name, expert_id, shard_id) num_experts = self.config.n_routed_experts - if rocm_aiter_moe_shared_expert_enabled and self.config.n_shared_experts: + if self.is_fused_shared_expert_enabled and self.config.n_shared_experts: num_experts += self.config.n_shared_experts expert_params_mapping = fused_moe_make_expert_params_mapping( self, @@ -279,7 +282,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: name = self._rewrite_spec_layer_name(spec_layer, name) is_fusion_moe_shared_experts_layer = ( - rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name) + self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name) ) for param_name, weight_name, shard_id in stacked_params_mapping: diff --git a/vllm/model_executor/models/qwen3_5.py b/vllm/model_executor/models/qwen3_5.py index 28c6a1189937..3fb3b7a782c8 100644 --- a/vllm/model_executor/models/qwen3_5.py +++ b/vllm/model_executor/models/qwen3_5.py @@ -29,13 +29,15 @@ import torch from torch import nn -from vllm._aiter_ops import rocm_aiter_ops from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig from vllm.distributed import ( get_pp_group, ) from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.utils import ( + is_model_fused_shared_expert_compatible, +) from vllm.model_executor.layers.layernorm import GemmaRMSNorm as Qwen3_5RMSNorm from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( @@ -78,7 +80,6 @@ Qwen3NextModel, Qwen3NextSparseMoeBlock, QwenNextMixtureOfExperts, - _is_shared_expert_fse_compatible, ) from .qwen3_vl import ( Qwen3_VisionTransformer, @@ -256,6 +257,11 @@ def get_layer(prefix: str): self.start_layer, self.end_layer, self.layers = make_layers( config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers" ) + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( + self.layers, + Qwen3NextSparseMoeBlock, + "mlp", + ) self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( ["hidden_states", "residual"], config.hidden_size ) @@ -274,8 +280,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: if "moe" in self.config.model_type: weights = maybe_fuse_shared_experts( weights, - enabled=rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() - and _is_shared_expert_fse_compatible(self.quant_config), + enabled=self.is_fused_shared_expert_enabled, n_routed_experts=self.config.num_experts, n_shared_experts=1, ckpt_prefix="mlp.shared_expert", diff --git a/vllm/model_executor/models/qwen3_5_mtp.py b/vllm/model_executor/models/qwen3_5_mtp.py index 67f06f4f5f24..4c064bead747 100644 --- a/vllm/model_executor/models/qwen3_5_mtp.py +++ b/vllm/model_executor/models/qwen3_5_mtp.py @@ -7,11 +7,13 @@ import torch from torch import nn -from vllm._aiter_ops import rocm_aiter_ops from vllm.compilation.decorators import support_torch_compile -from vllm.config import VllmConfig, get_current_vllm_config +from vllm.config import VllmConfig from vllm.distributed import get_pp_group, tensor_model_parallel_all_gather from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.utils import ( + is_model_fused_shared_expert_compatible, +) from vllm.model_executor.layers.linear import ColumnParallelLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ( @@ -25,8 +27,8 @@ Qwen3_5RMSNorm, ) from vllm.model_executor.models.qwen3_next import ( + Qwen3NextSparseMoeBlock, QwenNextMixtureOfExperts, - _is_shared_expert_fse_compatible, ) from vllm.model_executor.models.utils import sequence_parallel_chunk from vllm.sequence import IntermediateTensors @@ -122,6 +124,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): for idx in range(self.num_mtp_layers) ) vllm_config.quant_config = original_quant + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( + self.layers, + Qwen3NextSparseMoeBlock, + "mlp", + ) self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( ["hidden_states", "residual"], config.hidden_size ) @@ -185,10 +192,7 @@ def forward( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: weights = maybe_fuse_shared_experts( weights, - enabled=rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() - and _is_shared_expert_fse_compatible( - get_current_vllm_config().quant_config - ), + enabled=self.is_fused_shared_expert_enabled, n_routed_experts=getattr(self.config, "num_experts", 0), n_shared_experts=1, ckpt_prefix="mlp.shared_expert", diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py index b241baca6cbc..9511fcee9e7f 100644 --- a/vllm/model_executor/models/qwen3_next.py +++ b/vllm/model_executor/models/qwen3_next.py @@ -8,7 +8,6 @@ import torch from torch import nn -from vllm._aiter_ops import rocm_aiter_ops from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, ModelConfig, VllmConfig from vllm.distributed import ( @@ -18,9 +17,12 @@ tensor_model_parallel_all_gather, tensor_model_parallel_reduce_scatter, ) -from vllm.logger import init_logger from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.fused_moe import FusedMoEFactory +from vllm.model_executor.layers.fused_moe.utils import ( + is_model_fused_shared_expert_compatible, + resolve_layer_fused_shared_expert, +) from vllm.model_executor.layers.fused_qk_norm_rope import fused_qk_rmsnorm_rope_gate from vllm.model_executor.layers.layernorm import ( GemmaRMSNorm as Qwen3NextRMSNorm, @@ -73,8 +75,6 @@ maybe_prefix, ) -logger = init_logger(__name__) - KVCache = tuple[torch.Tensor, torch.Tensor] @@ -90,26 +90,6 @@ def _should_use_sequence_parallel(vllm_config: VllmConfig) -> bool: ) -def _is_shared_expert_fse_compatible(quant_config) -> bool: - """Check if shared expert can be fused with routed experts. - - FSE requires that shared and routed expert weights use the same - quantization format. Returns False when the shared expert is - excluded from quantization (e.g. float32 shared in an MXFP4 model) - or has a different quant spec than routed experts. - """ - if quant_config is None: - return True - # Quark stores its full config dict in quant_config.quant_config - raw_config = getattr(quant_config, "quant_config", None) - if not isinstance(raw_config, dict): - return True - exclude = raw_config.get("exclude", []) - if not exclude: - return True - return not any("shared_expert." in str(e) for e in exclude) - - class Qwen3NextSparseMoeBlock(nn.Module): def __init__(self, vllm_config: VllmConfig, prefix: str = ""): super().__init__() @@ -157,15 +137,18 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = ""): prefix=f"{prefix}.shared_expert_gate", ) - _fse_requested = rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() - _fse_enabled = _fse_requested and _is_shared_expert_fse_compatible(quant_config) - if _fse_requested and not _fse_enabled: - logger.warning( - "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but " - "shared expert has a different quantization spec than routed " - "experts. Falling back to non-fused shared expert path." + self.is_fused_shared_expert_enabled = False + if config.shared_expert_intermediate_size > 0: + self.is_fused_shared_expert_enabled = resolve_layer_fused_shared_expert( + quant_config, + prefix, + shared_expert_name="shared_expert", ) - if _fse_enabled or config.shared_expert_intermediate_size <= 0: + + if ( + self.is_fused_shared_expert_enabled + or config.shared_expert_intermediate_size <= 0 + ): self.shared_expert = None else: self.shared_expert = Qwen3NextMLP( @@ -193,6 +176,7 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = ""): num_redundant_experts=self.n_redundant_experts, is_sequence_parallel=self.is_sequence_parallel, n_shared_experts=1 if self.shared_expert is None else None, + fuse_shared_experts=self.is_fused_shared_expert_enabled, shared_expert_gate=self.shared_expert_gate if self.shared_expert is None else None, @@ -599,6 +583,11 @@ def get_layer(prefix: str): self.start_layer, self.end_layer, self.layers = make_layers( config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers" ) + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( + self.layers, + Qwen3NextSparseMoeBlock, + "mlp", + ) self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( ["hidden_states", "residual"], config.hidden_size ) @@ -678,6 +667,7 @@ def forward( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: weights = maybe_fuse_shared_experts( weights, + enabled=self.is_fused_shared_expert_enabled, n_routed_experts=getattr(self.config, "num_experts", 0), n_shared_experts=1, ckpt_prefix="mlp.shared_expert", diff --git a/vllm/model_executor/models/qwen3_next_mtp.py b/vllm/model_executor/models/qwen3_next_mtp.py index 35f550700a54..6b527acd62d7 100644 --- a/vllm/model_executor/models/qwen3_next_mtp.py +++ b/vllm/model_executor/models/qwen3_next_mtp.py @@ -11,6 +11,9 @@ from vllm.config import VllmConfig from vllm.distributed import get_pp_group, tensor_model_parallel_all_gather from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.utils import ( + is_model_fused_shared_expert_compatible, +) from vllm.model_executor.layers.linear import ColumnParallelLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ( @@ -21,6 +24,7 @@ Qwen3NextDecoderLayer, Qwen3NextModel, Qwen3NextRMSNorm, + Qwen3NextSparseMoeBlock, QwenNextMixtureOfExperts, ) from vllm.model_executor.models.utils import sequence_parallel_chunk @@ -91,6 +95,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): for idx in range(self.num_mtp_layers) ) + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( + self.layers, + Qwen3NextSparseMoeBlock, + "mlp", + ) self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( ["hidden_states", "residual"], config.hidden_size ) @@ -155,6 +164,7 @@ def forward( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: weights = maybe_fuse_shared_experts( weights, + enabled=self.is_fused_shared_expert_enabled, n_routed_experts=self.config.num_experts, n_shared_experts=1, ckpt_prefix="mlp.shared_expert", diff --git a/vllm/models/deepseek_v32/amd/mtp.py b/vllm/models/deepseek_v32/amd/mtp.py index 109bd2662265..9ea86ee26b0f 100644 --- a/vllm/models/deepseek_v32/amd/mtp.py +++ b/vllm/models/deepseek_v32/amd/mtp.py @@ -8,12 +8,14 @@ import torch import torch.nn as nn -from vllm._aiter_ops import rocm_aiter_ops from vllm.config import VllmConfig from vllm.distributed import tensor_model_parallel_all_reduce from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, ) +from vllm.model_executor.layers.fused_moe.utils import ( + is_model_fused_shared_expert_compatible, +) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ( @@ -167,6 +169,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) self.set_moe_parameters() + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( + self.model.layers.values(), + DeepseekV2MoE, + "mtp_block.mlp", + ) def set_moe_parameters(self): self.num_moe_layers = self.config.num_nextn_predict_layers @@ -231,9 +238,6 @@ def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: return name def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - rocm_aiter_moe_shared_expert_enabled = ( - rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() - ) stacked_params_mapping = [ ("gate_up_proj", "gate_proj", 0), ("gate_up_proj", "up_proj", 1), @@ -250,7 +254,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: num_experts=self.config.n_routed_experts + ( self.config.n_shared_experts - if rocm_aiter_moe_shared_expert_enabled + if self.is_fused_shared_expert_enabled else 0 ), ) @@ -266,7 +270,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: if spec_layer is None: continue is_fusion_moe_shared_experts_layer = ( - rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name) + self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name) ) name = self._rewrite_spec_layer_name(spec_layer, name) diff --git a/vllm/models/deepseek_v32/nvidia/mtp.py b/vllm/models/deepseek_v32/nvidia/mtp.py index c2dcbdd64d43..2ea01dbdcb9c 100644 --- a/vllm/models/deepseek_v32/nvidia/mtp.py +++ b/vllm/models/deepseek_v32/nvidia/mtp.py @@ -7,7 +7,6 @@ import torch.nn as nn import vllm.envs as envs -from vllm._aiter_ops import rocm_aiter_ops from vllm.config import VllmConfig from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.model_executor.layers.fused_embed_norm import ( @@ -344,9 +343,6 @@ def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: return name def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - rocm_aiter_moe_shared_expert_enabled = ( - rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() - ) stacked_params_mapping = [ ("gate_up_proj", "gate_proj", 0), ("gate_up_proj", "up_proj", 1), @@ -360,12 +356,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", - num_experts=self.config.n_routed_experts - + ( - self.config.n_shared_experts - if rocm_aiter_moe_shared_expert_enabled - else 0 - ), + num_experts=self.config.n_routed_experts, ) pp_missing_layer_names = get_pp_missing_layer_names(self) @@ -378,9 +369,6 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) if spec_layer is None: continue - is_fusion_moe_shared_experts_layer = ( - rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name) - ) name = self._rewrite_spec_layer_name(spec_layer, name) if _try_load_fp8_indexer_wk( @@ -398,8 +386,6 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: continue if ("mlp.experts." in name) and name not in params_dict: continue - if is_fusion_moe_shared_experts_layer: - continue name_mapped = name.replace(weight_name, param_name) if ( param_name == "fused_qkv_a_proj" @@ -415,32 +401,10 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: break else: num_chunks = 1 - if is_fusion_moe_shared_experts_layer: - num_chunks = getattr(self.config, "n_shared_experts", 1) or 1 - split_dim = ( - 1 - if ("down_proj.weight" in name and loaded_weight.ndim > 1) - else 0 - ) - total = loaded_weight.shape[split_dim] - assert total % num_chunks == 0 - chunk_size = total // num_chunks for j in range(num_chunks): chunk_name = name weight_to_load = loaded_weight - if is_fusion_moe_shared_experts_layer: - chunk_slice = slice(j * chunk_size, (j + 1) * chunk_size) - if loaded_weight.ndim == 1: - weight_to_load = loaded_weight[chunk_slice] - elif split_dim == 0: - weight_to_load = loaded_weight[chunk_slice, :] - else: - weight_to_load = loaded_weight[:, chunk_slice] - chunk_name = name.replace( - "mlp.shared_experts", - f"mlp.experts.{self.config.n_routed_experts + j}", - ) is_expert_weight = False for mapping in expert_params_mapping: @@ -462,10 +426,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return_success=True, ) if success: - if not is_fusion_moe_shared_experts_layer: - name = name_mapped - else: - loaded_params.add(name_mapped) + name = name_mapped break else: if is_expert_weight: @@ -485,8 +446,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: param, "weight_loader", default_weight_loader ) weight_loader(param, loaded_weight) - if not is_fusion_moe_shared_experts_layer: - loaded_params.add(name) + loaded_params.add(name) loaded_layers: set[int] = set() for param_name in loaded_params: diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index 5888ae563cb8..145640860572 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -16,12 +16,16 @@ get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, ) +from vllm.logger import init_logger from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp from vllm.model_executor.layers.fused_moe import ( FusedMoEFactory, GateLinear, fused_moe_make_expert_params_mapping, ) +from vllm.model_executor.layers.fused_moe.utils import ( + is_model_fused_shared_expert_compatible, +) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, @@ -37,6 +41,9 @@ MHCPreOp, ) from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.quantization.utils.config_utils import ( + is_shared_expert_quant_fse_compatible, +) from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, @@ -60,6 +67,8 @@ from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors +logger = init_logger(__name__) + class DeepseekV4MLP(nn.Module): def __init__( @@ -154,48 +163,11 @@ def forward(self, x): return x -def _shared_experts_are_fp4(config, layer_idx: int | None = None) -> bool: - """Whether the shared experts are MXFP4 and thus fusable. - - ``layer_idx=None`` resolves the model-wide default (global scheme), used by - the main-model weight loader / mapper callers that operate per-model. - """ - quant_cfg = getattr(config, "quantization_config", None) - if quant_cfg is None: - return False - if layer_idx is None: - base = None - elif layer_idx >= config.num_hidden_layers: - base = f"mtp.{layer_idx - config.num_hidden_layers}.ffn.shared_experts" - else: - base = f"layers.{layer_idx}.ffn.shared_experts" - if base and any(e.startswith(base) for e in (quant_cfg.get("exclude") or [])): - return False - entry = ( - (quant_cfg.get("layer_quant_config") or {}).get(f"{base}.w1") if base else None - ) - if entry is None: - entry = quant_cfg.get("global_quant_config") - return ((entry or {}).get("weight") or {}).get("dtype") == "fp4" - - -def _fuse_shared_experts_enabled(config, prefix: str = "") -> bool: - """Whether to fuse the shared expert into the routed MXFP4 grouped GEMM. - - Fusion fuses the shared expert into the routed experts' MXFP4 grouped GEMM, - so it only applies where the shared expert is the same precision as the - routed experts. Some layers may carry a shared expert in a different quantization - than the routed experts; when so, it runs as its own linear and must not be fused. - """ - if not ( - current_platform.is_rocm() - and getattr(config, "n_shared_experts", None) +def _fuse_shared_experts_enabled(config) -> bool: + return bool( + getattr(config, "n_shared_experts", None) and envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS and not get_current_vllm_config().parallel_config.enable_expert_parallel - ): - return False - return _shared_experts_are_fp4( - config, extract_layer_index(prefix) if prefix else None ) @@ -255,9 +227,25 @@ def __init__( self.n_shared_experts = config.n_shared_experts - self.fuse_shared_experts = _fuse_shared_experts_enabled(config, prefix) + # TODO: Historically, only `VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1` + # is checked to enable FSE for DeepSeek-v4, despite AITER not being used. + # This should be cleaned up and use `resolve_layer_fused_shared_expert`. + fse_requested = _fuse_shared_experts_enabled(config) + if fse_requested: + fse_compatible, fse_reason = is_shared_expert_quant_fse_compatible( + quant_config, + f"{prefix}.experts", + f"{prefix}.shared_experts", + ) + if not fse_compatible: + logger.warning( + "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but " + "cannot be enabled: %s.", + fse_reason, + ) + self.is_fused_shared_expert_enabled = fse_requested and fse_compatible - if config.n_shared_experts is None or self.fuse_shared_experts: + if config.n_shared_experts is None or self.is_fused_shared_expert_enabled: self.shared_experts = None else: intermediate_size = config.moe_intermediate_size * config.n_shared_experts @@ -282,8 +270,9 @@ def __init__( self.experts = FusedMoEFactory( shared_experts=self.shared_experts, n_shared_experts=( - config.n_shared_experts if self.fuse_shared_experts else None + config.n_shared_experts if self.is_fused_shared_expert_enabled else None ), + fuse_shared_experts=self.is_fused_shared_expert_enabled, gate=self.gate, num_experts=config.n_routed_experts, top_k=config.num_experts_per_tok, @@ -581,6 +570,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ), prefix=f"{prefix}.layers", ) + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( + self.layers, + DeepseekV4MoE, + "ffn", + ) if get_pp_group().is_last_rank: self.norm = RMSNorm(config.hidden_size, self.rms_norm_eps) @@ -763,7 +757,7 @@ def _resolve_param_name(name: str) -> str: # diverge from how the module was built if per-layer quantization ever # mixes fused and non-fused layers. fuse_by_layer = { - extract_layer_index(mod_name): mod.fuse_shared_experts + extract_layer_index(mod_name): mod.is_fused_shared_expert_enabled for mod_name, mod in self.named_modules() if isinstance(mod, DeepseekV4MoE) } @@ -874,7 +868,7 @@ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: # redirected shared-expert weights route through the expert loader. n_shared = getattr(self.config, "n_shared_experts", 0) or 0 num_experts = self.config.n_routed_experts + ( - n_shared if _fuse_shared_experts_enabled(self.config) else 0 + n_shared if self.is_fused_shared_expert_enabled else 0 ) return fused_moe_make_expert_params_mapping( self, @@ -956,15 +950,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): config = vllm_config.model_config.hf_config self.config = config expert_dtype = getattr(config, "expert_dtype", "fp4") - fuse_shared_experts = _fuse_shared_experts_enabled(config) - if expert_dtype != "fp4" or fuse_shared_experts: - self.hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper( - expert_dtype, fuse_shared_experts=fuse_shared_experts - ) - self.model = self.model_cls( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) + if expert_dtype != "fp4" or self.model.is_fused_shared_expert_enabled: + self.hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper( + expert_dtype, + fuse_shared_experts=self.model.is_fused_shared_expert_enabled, + ) if get_pp_group().is_last_rank: self.lm_head = ParallelLMHead( config.vocab_size, diff --git a/vllm/models/minimax_m3/amd/model.py b/vllm/models/minimax_m3/amd/model.py index 6aa172f141c7..798f9a475bc6 100644 --- a/vllm/models/minimax_m3/amd/model.py +++ b/vllm/models/minimax_m3/amd/model.py @@ -34,6 +34,7 @@ ) from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size from vllm.forward_context import get_forward_context +from vllm.logger import init_logger from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.attention.attention import set_default_quant_scales from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase @@ -45,6 +46,9 @@ GateLinear, fused_moe_make_expert_params_mapping, ) +from vllm.model_executor.layers.fused_moe.utils import ( + is_model_fused_shared_expert_compatible, +) from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, MinimaxM3QKVParallelLinearWithIndexer, @@ -53,6 +57,9 @@ ) from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.quantization.utils.config_utils import ( + is_shared_expert_quant_fse_compatible, +) from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, @@ -108,23 +115,7 @@ is_quantized_kv_cache, ) - -def _fuse_shared_experts_enabled(config: PretrainedConfig) -> bool: - """Whether to fuse the shared expert with routed experts. - - ROCm only. Opt-in via ``VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS`` (the - router-append fusion runs on both aiter and non-aiter MoE); - it is disabled under expert parallelism (the shared slot is appended to - the routed top-k, which the EP expert-mapping path does not handle). - """ - from vllm.platforms import current_platform - - return bool( - current_platform.is_rocm() - and getattr(config, "n_shared_experts", None) - and envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS - and not get_current_vllm_config().parallel_config.enable_expert_parallel - ) +logger = init_logger(__name__) def _sparse_attention_layer_ids(config: PretrainedConfig) -> set[int]: @@ -303,10 +294,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x -def _aiter_moe_fused_shared_experts_enabled(config: PretrainedConfig) -> bool: +def _aiter_moe_fused_shared_experts_enabled( + is_fused_shared_expert_enabled: bool, +) -> bool: """Whether the fused shared expert routes through aiter's grouped top-k MoE. - A strict sub-case of :func:`_fuse_shared_experts_enabled`: shared-expert + A strict sub-case of `is_fused_shared_expert_enabled`: shared-expert fusion must already be opted in (``VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS``) and allowed (not under expert parallelism). When additionally on gfx950 with an active aiter MoE backend, the shared expert is appended inside aiter's @@ -314,11 +307,13 @@ def _aiter_moe_fused_shared_experts_enabled(config: PretrainedConfig) -> bool: vLLM router's torch concat. Otherwise FSE still runs via the vLLM top-k bias router. """ - if not _fuse_shared_experts_enabled(config): - return False from vllm.platforms.rocm import on_gfx950 - return on_gfx950() and rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() + return ( + on_gfx950() + and is_fused_shared_expert_enabled + and rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() + ) class MiniMaxM3MoE(nn.Module): @@ -371,11 +366,46 @@ def __init__( # MoE call as the last expert slot, so we don't build a separate module. # On gfx950 with aiter MoE the append is fused inside aiter's grouped # top-k kernel; otherwise it goes through the vLLM top-k bias router. - self.fuse_shared_experts = _fuse_shared_experts_enabled(config) - self.use_aiter_moe_fse = _aiter_moe_fused_shared_experts_enabled(config) + # It is disabled under expert parallelism (the shared slot is appended to + # the routed top-k, which the EP expert-mapping path does not handle). + # TODO: Historically, only `VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1` + # is checked to enable FSE for MiniMax-M3, despite AITER not being used. + # This should be cleaned up and use `resolve_layer_fused_shared_expert`. + fse_requested = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS + + self.is_fused_shared_expert_enabled = False + if ( + fse_requested + and bool(getattr(config, "n_shared_experts", None)) + and not get_current_vllm_config().parallel_config.enable_expert_parallel + ): + fse_compatible, fse_reason = is_shared_expert_quant_fse_compatible( + quant_config, + f"{prefix}.experts", + f"{prefix}.shared_experts", + ) + if not fse_compatible: + logger.warning( + "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but " + "cannot be enabled: %s.", + fse_reason, + ) + else: + self.is_fused_shared_expert_enabled = True + + # When additionally on gfx950 with an active aiter MoE backend, the shared + # expert is appended inside aiter's an active aiter MoE backend, the shared + # expert is appended inside aiter's biased grouped top-k kernel + # (``num_fused_shared_experts``) instead of the vLLM router's torch concat. + # Otherwise FSE still runs via the vLLM top-k bias router. + # TODO: `on_gfx950()` check here should not be MiniMax-M3 specific, and + # the check should be done on resolved MOE backend directly. + self.use_aiter_moe_fse = _aiter_moe_fused_shared_experts_enabled( + self.is_fused_shared_expert_enabled + ) self.shared_experts: MiniMaxM3MLP | None = None - if self.n_shared_experts and not self.fuse_shared_experts: + if self.n_shared_experts and not self.is_fused_shared_expert_enabled: self.shared_experts = MiniMaxM3MLP( config=config, intermediate_size=config.intermediate_size * self.n_shared_experts, @@ -412,8 +442,9 @@ def __init__( router_logits_dtype=self.gate.out_dtype, shared_experts=self.shared_experts, n_shared_experts=( - self.n_shared_experts if self.fuse_shared_experts else None + self.n_shared_experts if self.is_fused_shared_expert_enabled else None ), + fuse_shared_experts=self.is_fused_shared_expert_enabled, quant_config=quant_config, prefix=f"{prefix}.experts", ) @@ -1101,6 +1132,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ), prefix=f"{prefix}.layers", ) + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( + self.layers, + MiniMaxM3MoE, + "block_sparse_moe", + ) if get_pp_group().is_last_rank: self.norm = MiniMAXGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -1155,7 +1191,7 @@ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: # expert, include the appended slot (id == num_local_experts). n_shared = getattr(self.config, "n_shared_experts", 0) or 0 num_experts = self.config.num_local_experts + ( - n_shared if _fuse_shared_experts_enabled(self.config) else 0 + n_shared if self.is_fused_shared_expert_enabled else 0 ) return fused_moe_make_expert_params_mapping( self, @@ -1187,8 +1223,6 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # (param_name, weight_name, expert_id, shard_id) expert_params_mapping = self.get_expert_mapping() - _fuse_shared = _fuse_shared_experts_enabled(self.config) - params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() for name, loaded_weight in weights: @@ -1206,7 +1240,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # down->w2) so it loads via the routed expert loader. Runs before the # stacked/dense mappings so shared_experts.gate_proj/up_proj are not # captured by the dense gate_up_proj mapping. - if _fuse_shared and ".shared_experts." in name: + if self.is_fused_shared_expert_enabled and ".shared_experts." in name: sid = self.config.num_local_experts name = name.replace(".shared_experts.gate_proj.", f".experts.{sid}.w1.") name = name.replace(".shared_experts.up_proj.", f".experts.{sid}.w3.") From bca7bea2405127bd5291bb6fffa679bdcd8f6dd9 Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Tue, 18 Aug 2026 12:58:04 -0400 Subject: [PATCH 095/839] Remove VLLM_TEST_FORCE_FP8_MARLIN to replace with linear_backend/moe_backend (#52182) Signed-off-by: mgoin Co-authored-by: OpenAI Codex --- tests/compile/passes/test_fusion.py | 4 +++ .../passes/test_mla_attn_quant_fusion.py | 6 ++++ .../Llama-4-Scout-Fp8-ModelOpt-marlin.yaml | 4 +-- .../Qwen3-30B-A3B-Fp8-AutoFp8-marlin.yaml | 4 +-- .../Qwen3-30B-A3B-Fp8-CT-Block-marlin.yaml | 4 +-- .../Qwen3-30B-A3B-Fp8-CT-Channel-marlin.yaml | 4 +-- .../Qwen3-30B-A3B-NvFp4-CT-marlin.yaml | 4 +-- .../Qwen3-30B-A3B-NvFp4-ModelOpt-marlin.yaml | 4 +-- tests/quantization/test_fp8.py | 33 +++++++++++-------- tests/utils.py | 11 +++++-- vllm/envs.py | 6 ---- .../model_executor/kernels/linear/__init__.py | 4 +-- .../kernels/linear/scaled_mm/marlin.py | 10 ------ .../layers/fused_moe/oracle/fp8.py | 7 ---- .../layers/fused_moe/oracle/nvfp4.py | 7 ---- 15 files changed, 46 insertions(+), 66 deletions(-) diff --git a/tests/compile/passes/test_fusion.py b/tests/compile/passes/test_fusion.py index 591b014d9e25..50c957fceb53 100644 --- a/tests/compile/passes/test_fusion.py +++ b/tests/compile/passes/test_fusion.py @@ -371,6 +371,10 @@ def test_fusion_rmsnorm_quant( use_aiter_fusion=False, use_aiter_quant=False, ) + if any( + type(layer.kernel) is not force_kernel for layer in model.fp8_linear_layers + ): + pytest.skip(f"{force_kernel.__name__} is not supported on this platform") backend, _ = _run_fusion_test( model, fusion_pass, vllm_config, dtype, hidden_size, num_tokens diff --git a/tests/compile/passes/test_mla_attn_quant_fusion.py b/tests/compile/passes/test_mla_attn_quant_fusion.py index 5b2ef5bfd95c..106bb8e80abc 100644 --- a/tests/compile/passes/test_mla_attn_quant_fusion.py +++ b/tests/compile/passes/test_mla_attn_quant_fusion.py @@ -481,6 +481,12 @@ def test_mla_attention_quant_pattern( device=device, vllm_config=vllm_config_unfused, ) + if ( + model_class is TestMLAAttentionFp8GroupQuantPatternModel + and type(model_unfused.block_fp8_linear.kernel) + is not CutlassFp8BlockScaledMMKernel + ): + pytest.skip("CUTLASS FP8 block kernel is not supported on this platform") model_unfused = model_unfused.to(device) # HACK: See #131044 result_unfused_0 = model_unfused(q, kv_c_normed, k_pe) # noqa: F841 diff --git a/tests/evals/gsm8k/configs/moe-refactor/Llama-4-Scout-Fp8-ModelOpt-marlin.yaml b/tests/evals/gsm8k/configs/moe-refactor/Llama-4-Scout-Fp8-ModelOpt-marlin.yaml index be8192f2a89a..23ceed6803d9 100644 --- a/tests/evals/gsm8k/configs/moe-refactor/Llama-4-Scout-Fp8-ModelOpt-marlin.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor/Llama-4-Scout-Fp8-ModelOpt-marlin.yaml @@ -2,6 +2,4 @@ model_name: "nvidia/Llama-4-Scout-17B-16E-Instruct-FP8" accuracy_threshold: 0.92 num_questions: 1319 num_fewshot: 5 -server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2" -env: - VLLM_TEST_FORCE_FP8_MARLIN: "1" +server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --linear-backend marlin --moe-backend marlin" diff --git a/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-AutoFp8-marlin.yaml b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-AutoFp8-marlin.yaml index c3d86e6bfbcb..0b8c90a37e2b 100644 --- a/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-AutoFp8-marlin.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-AutoFp8-marlin.yaml @@ -2,6 +2,4 @@ model_name: "Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8" accuracy_threshold: 0.88 num_questions: 1319 num_fewshot: 5 -server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2" -env: - VLLM_TEST_FORCE_FP8_MARLIN: "1" +server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --linear-backend marlin --moe-backend marlin" diff --git a/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-CT-Block-marlin.yaml b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-CT-Block-marlin.yaml index 46eee742131d..d22ee2086dbc 100644 --- a/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-CT-Block-marlin.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-CT-Block-marlin.yaml @@ -2,6 +2,4 @@ model_name: "RedHatAI/Qwen3-30B-A3B-FP8-block" accuracy_threshold: 0.85 num_questions: 1319 num_fewshot: 5 -server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2" -env: - VLLM_TEST_FORCE_FP8_MARLIN: "1" +server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --linear-backend marlin --moe-backend marlin" diff --git a/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-CT-Channel-marlin.yaml b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-CT-Channel-marlin.yaml index 8ed6410c36b5..dd86ee528b70 100644 --- a/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-CT-Channel-marlin.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-CT-Channel-marlin.yaml @@ -2,6 +2,4 @@ model_name: "RedHatAI/Qwen3-30B-A3B-FP8-dynamic" accuracy_threshold: 0.85 num_questions: 1319 num_fewshot: 5 -server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2" -env: - VLLM_TEST_FORCE_FP8_MARLIN: "1" +server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --linear-backend marlin --moe-backend marlin" diff --git a/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-CT-marlin.yaml b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-CT-marlin.yaml index 8199e6563495..ffe6829ccba3 100644 --- a/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-CT-marlin.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-CT-marlin.yaml @@ -2,6 +2,4 @@ model_name: "RedHatAI/Qwen3-30B-A3B-NVFP4" accuracy_threshold: 0.88 num_questions: 1319 num_fewshot: 5 -server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2" -env: - VLLM_TEST_FORCE_FP8_MARLIN: "1" +server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --linear-backend marlin --moe-backend marlin" diff --git a/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-ModelOpt-marlin.yaml b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-ModelOpt-marlin.yaml index 4156cec89761..b6b71f4d3dd4 100644 --- a/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-ModelOpt-marlin.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-ModelOpt-marlin.yaml @@ -2,6 +2,4 @@ model_name: "nvidia/Qwen3-30B-A3B-NVFP4" accuracy_threshold: 0.88 num_questions: 1319 num_fewshot: 5 -server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2" -env: - VLLM_TEST_FORCE_FP8_MARLIN: "1" +server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --linear-backend marlin --moe-backend marlin" diff --git a/tests/quantization/test_fp8.py b/tests/quantization/test_fp8.py index a1c74259f8b8..5ad3b666c2ea 100644 --- a/tests/quantization/test_fp8.py +++ b/tests/quantization/test_fp8.py @@ -64,10 +64,12 @@ def test_model_load_and_run( if use_rocm_aiter: monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") + kwargs = {} if force_marlin: - monkeypatch.setenv("VLLM_TEST_FORCE_FP8_MARLIN", "1") + kwargs["linear_backend"] = "marlin" + kwargs["moe_backend"] = "marlin" - with vllm_runner(model_id, enforce_eager=True) as llm: + with vllm_runner(model_id, enforce_eager=True, **kwargs) as llm: # note: this does not test accuracy, just that we can run through # see lm-eval tests for accuracy outputs = llm.generate_greedy(["Hello my name is"], max_tokens=4) @@ -98,8 +100,10 @@ def test_online_quantization( # `LLM.apply_model` requires pickling a function. monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") + kwargs = {} if force_marlin: - monkeypatch.setenv("VLLM_TEST_FORCE_FP8_MARLIN", "1") + kwargs["linear_backend"] = "marlin" + kwargs["moe_backend"] = "marlin" model_dtype = "auto" if kv_cache_dtype == "fp8" and current_platform.is_device_capability_family(90): @@ -112,6 +116,7 @@ def test_online_quantization( dtype=model_dtype, enforce_eager=True, kv_cache_dtype=kv_cache_dtype, + **kwargs, ) as llm: def check_model(model): @@ -340,7 +345,7 @@ def per_tensor_dequantize(tensor, inv_scale, dtype): @pytest.mark.parametrize("method_cls", [Fp8LinearMethod, Fp8MoEMethod]) # FP8 weight reloading does not support online quantization @pytest.mark.parametrize("is_checkpoint_fp8_serialized", [True]) # skip False -@pytest.mark.parametrize("weight_block_size", [None, [1, 1]]) +@pytest.mark.parametrize("weight_block_size", [None, [128, 128]]) # any postprocessing that is applied to the weights such as padding and repacking # (excluding device sharding) must also be applied to the reloaded weights # @@ -371,6 +376,8 @@ def test_fp8_reloading( # Set model config as model_config.dtype is required in Fp8LinearMethod. default_vllm_config.model_config = ModelConfig() + default_vllm_config.kernel_config.moe_backend = "triton" + layer_size = 128 if weight_block_size is not None else 1 with torch.device(f"{DEVICE_TYPE}:0"): config = Fp8Config( is_checkpoint_fp8_serialized=is_checkpoint_fp8_serialized, @@ -378,14 +385,14 @@ def test_fp8_reloading( ) if method_cls is Fp8LinearMethod: - layer = torch.nn.Linear(1, 1) + layer = torch.nn.Linear(layer_size, layer_size) method = method_cls(config) method.create_weights( layer=layer, - input_size_per_partition=1, - output_partition_sizes=[1], - input_size=1, - output_size=1, + input_size_per_partition=layer_size, + output_partition_sizes=[layer_size], + input_size=layer_size, + output_size=layer_size, params_dtype=torch.bfloat16, weight_loader=default_weight_loader, ) @@ -395,16 +402,16 @@ def test_fp8_reloading( layer = FusedMoEFactory( num_experts=1, top_k=1, - hidden_size=1, - intermediate_size=1, + hidden_size=layer_size, + intermediate_size=layer_size, ) layer = layer.routed_experts method = method_cls(config, layer) method.create_weights( layer=layer, num_experts=1, - hidden_size=1, - intermediate_size_per_partition=1, + hidden_size=layer_size, + intermediate_size_per_partition=layer_size, params_dtype=torch.bfloat16, weight_loader=default_weight_loader, ) diff --git a/tests/utils.py b/tests/utils.py index 86774bf4bf82..07601b74e486 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -2339,6 +2339,10 @@ def __init__( force_kernel: type[_KernelT] | None = None, ): super().__init__() + self.input_size_per_partition = weight_shape[1] + self.output_size_per_partition = weight_shape[0] + self.logical_widths = [self.output_size_per_partition] + self.orig_dtype = input_dtype act_scale_desc = activation_quant_key.scale weight_scale_desc = weight_quant_key.scale is_block_wise = act_scale_desc.group_shape.is_per_group() @@ -2346,11 +2350,12 @@ def __init__( block_size = weight_scale_desc.group_shape.col weight_scale_shape = weight_shape[0] // block_size self.weight_scale_inv = torch.rand( - (weight_scale_shape, weight_scale_shape), dtype=torch.float32 + (weight_scale_shape, weight_scale_shape), + dtype=torch.float32, + device=device, ) - self.weight = torch.rand(weight_shape).to(dtype=FP8_DTYPE) + self.weight = torch.rand(weight_shape, device=device).to(dtype=FP8_DTYPE) self.input_scale = None - self.weight_scale = None self.weight_block_size = [block_size, block_size] if transpose_weights: self.weight = self.weight.t() diff --git a/vllm/envs.py b/vllm/envs.py index 1a7458802a90..f80c51207fd6 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -1099,12 +1099,6 @@ def _resolve_rust_cli_path() -> str | None: os.environ.get("VLLM_ALLOW_LONG_MAX_MODEL_LEN", "0").strip().lower() in ("1", "true") ), - # If set, forces FP8 Marlin to be used for FP8 quantization regardless - # of the hardware support for FP8 compute. - "VLLM_TEST_FORCE_FP8_MARLIN": lambda: ( - os.environ.get("VLLM_TEST_FORCE_FP8_MARLIN", "0").strip().lower() - in ("1", "true") - ), "VLLM_TEST_FORCE_LOAD_FORMAT": lambda: os.getenv( "VLLM_TEST_FORCE_LOAD_FORMAT", "dummy" ), diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 3ba499155b6c..b94c8f7d58d0 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -394,12 +394,12 @@ def _resolve_backend_kernels( # in priority/performance order (when available) _POSSIBLE_FP8_KERNELS: dict[PlatformEnum, list[type[FP8ScaledMMLinearKernel]]] = { PlatformEnum.CUDA: [ - MarlinFP8ScaledMMLinearKernel, FlashInferFP8ScaledMMLinearKernel, CutlassFP8ScaledMMLinearKernel, B12xTensorFP8ScaledMMLinearKernel, PerTensorTorchFP8ScaledMMLinearKernel, ChannelWiseTorchFP8ScaledMMLinearKernel, + MarlinFP8ScaledMMLinearKernel, HummingFP8ScaledMMLinearKernel, ], PlatformEnum.ROCM: [ @@ -434,8 +434,8 @@ def _resolve_backend_kernels( CutlassFp8BlockScaledMMKernel, B12xFp8BlockScaledMMKernel, MarlinFP8ScaledMMLinearKernel, - TritonFp8BlockScaledMMKernel, HummingFP8ScaledMMLinearKernel, + TritonFp8BlockScaledMMKernel, BlockWiseTorchFP8ScaledMMLinearKernel, ], PlatformEnum.ROCM: [ diff --git a/vllm/model_executor/kernels/linear/scaled_mm/marlin.py b/vllm/model_executor/kernels/linear/scaled_mm/marlin.py index 66a03b4d205b..53bdc724bb90 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/marlin.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/marlin.py @@ -43,16 +43,6 @@ def is_supported( return False, "FP8 Marlin requires compute capability 7.5 or higher" if envs.VLLM_BATCH_INVARIANT: return False, "FP8 Marlin not supported for batch invariant execution." - if ( - compute_capability is not None - and compute_capability >= 89 - and not envs.VLLM_TEST_FORCE_FP8_MARLIN - ): - return ( - False, - "To apply FP8 Marlin on high-capability GPUs, please set " - "VLLM_TEST_FORCE_FP8_MARLIN=1", - ) return True, None @classmethod diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index 6153ed28c158..b2d76e4ef924 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -370,13 +370,6 @@ def _return_or_raise( backend, config, weight_key, activation_key, activation_format ) - # Handle explicit MARLIN FP8 configuration. - if envs.VLLM_TEST_FORCE_FP8_MARLIN: - backend = Fp8MoeBackend.MARLIN - return _return_or_raise( - backend, config, weight_key, activation_key, activation_format - ) - # Handle explicit AITER FP8 configuration. if envs.is_set("VLLM_ROCM_USE_AITER") or envs.is_set("VLLM_ROCM_USE_AITER_MOE"): skip_aiter_moe = ( diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index f48e79e13da3..3f23ecd61906 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -4,7 +4,6 @@ import torch -import vllm.envs as envs import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm.config.kernel import MoEBackend from vllm.logger import init_logger @@ -271,12 +270,6 @@ def _return_or_raise( requested_backend, config, weight_key, activation_key, activation_format ) - if envs.VLLM_TEST_FORCE_FP8_MARLIN: - backend = NvFp4MoeBackend.MARLIN - return _return_or_raise( - backend, config, weight_key, activation_key, activation_format - ) - # Select kernels in order of backend. for backend in AVAILABLE_BACKENDS: for k_cls in backend_to_kernel_cls(backend): From ddbf826bee965ace7d2a68599fb45a69e79745d2 Mon Sep 17 00:00:00 2001 From: Strahinja Stamenkovic Date: Tue, 18 Aug 2026 19:54:51 +0200 Subject: [PATCH 096/839] [ROCm] Gate Torch FP8 scaled-MM on architecture support (#51021) Signed-off-by: sstamenk Signed-off-by: Strahinja Stamenkovic Signed-off-by: Strahinja Stamenkovic Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../kernels/linear/scaled_mm/pytorch.py | 43 +++++++++++-------- 1 file changed, 26 insertions(+), 17 deletions(-) diff --git a/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py b/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py index 0d36b8079ddc..7cc9d8c5144a 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py @@ -24,6 +24,27 @@ def _get_num_tokens(output_shape: list) -> int: return math.prod(output_shape[:-1]) +def _rocm_torch_fp8_scaled_mm_supported() -> bool: + from vllm.platforms.rocm import on_gfx12x, on_gfx942, on_gfx950, on_gfx1250 + + return on_gfx942() or on_gfx950() or on_gfx12x() or on_gfx1250() + + +def _supports_torch_fp8_scaled_mm() -> bool: + if current_platform.is_cpu(): + return True + if current_platform.is_xpu(): + return True + if not current_platform.is_cuda_alike(): + return False + + # TODO: Use torch.cuda.is_scaled_mm_supported once it is in the supported + # PyTorch baseline. + if current_platform.is_rocm(): + return _rocm_torch_fp8_scaled_mm_supported() + return current_platform.supports_fp8() + + class TorchFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): """ Base class for FP8 linear kernels using Torch. @@ -35,17 +56,10 @@ class TorchFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): def is_supported( cls, compute_capability: int | None = None ) -> tuple[bool, str | None]: - if not ( - current_platform.is_cuda_alike() - or current_platform.is_cpu() - or current_platform.is_xpu() - ): - return False, "requires ROCm, CUDA, CPU or XPU." - - if compute_capability is not None and compute_capability < 89: - return False, "requires compute capability 89 and above." + if _supports_torch_fp8_scaled_mm(): + return True, None - return True, None + return False, "requires a platform with torch FP8 scaled-MM support." def get_output_padding(self) -> int | None: # Note: we pad the input because torch._scaled_mm is more performant @@ -111,13 +125,8 @@ def is_supported( if not current_platform.is_rocm(): return False, "requires ROCm." - from vllm.platforms.rocm import get_cdna_version, on_rdna4 - - if get_cdna_version() <= 2 and not on_rdna4(): - return False, "requires CDNA3+ or RDNA4" - - if compute_capability is not None and compute_capability < 94: - return False, "requires compute capability 94 and above." + if not _supports_torch_fp8_scaled_mm(): + return False, "requires platform with torch FP8 scaled-MM support." return True, None From ad5e71b276f6a8d9642560205736bee800464937 Mon Sep 17 00:00:00 2001 From: Matvei Pashkovskii Date: Tue, 18 Aug 2026 21:01:00 +0300 Subject: [PATCH 097/839] [ROCm][Perf] Enable fused KDA decode on gfx942 (MI325X) (#52293) Signed-off-by: Matvei Pashkovskii Co-authored-by: Claude --- CMakeLists.txt | 12 ++++++------ tests/models/kimi_k3/test_amd_kda_decode.py | 10 +++++----- vllm/models/kimi_k3/amd/ops/kda_decode.py | 7 ++++--- 3 files changed, 15 insertions(+), 14 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 64de3baa9fdf..bb0f51b43ef4 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1180,14 +1180,14 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() # HIP counterpart of the FUSED_KDA_DECODE block above: same `fused_kda_decode` - # op, same VLLM_ENABLE_FUSED_KDA_DECODE guard in torch_bindings.cpp. Only - # measured on gfx950, so it is only built when gfx950 is in the arch list. - # A multi-arch build that includes gfx950 still emits this source for every - # arch in that list. The runtime gate in kda_decode.py keeps the kernel off - # non-gfx950 devices. + # op, same VLLM_ENABLE_FUSED_KDA_DECODE guard in torch_bindings.cpp. Built for + # the CDNA archs the kernel supports (gfx942 / gfx950), so it is only emitted + # when one of those is in the arch list. A multi-arch build that includes them + # still emits this source for every arch in that list. The runtime gate in + # kda_decode.py keeps the kernel off unsupported devices. if(VLLM_GPU_LANG STREQUAL "HIP") set(FUSED_KDA_DECODE_HIP_ARCHS ${VLLM_GPU_ARCHES}) - list(FILTER FUSED_KDA_DECODE_HIP_ARCHS INCLUDE REGEX "gfx950") + list(FILTER FUSED_KDA_DECODE_HIP_ARCHS INCLUDE REGEX "gfx942|gfx950") if(FUSED_KDA_DECODE_HIP_ARCHS) set(FUSED_KDA_DECODE_HIP_SRC "csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel_rocm.cu") diff --git a/tests/models/kimi_k3/test_amd_kda_decode.py b/tests/models/kimi_k3/test_amd_kda_decode.py index e39acc72f246..7978e4fa7b5a 100644 --- a/tests/models/kimi_k3/test_amd_kda_decode.py +++ b/tests/models/kimi_k3/test_amd_kda_decode.py @@ -14,17 +14,17 @@ from vllm.platforms import current_platform -def _on_gfx950() -> bool: +def _on_supported_arch() -> bool: if not current_platform.is_rocm(): return False - from vllm.platforms.rocm import on_gfx950 + from vllm.platforms.rocm import on_gfx942, on_gfx950 - return on_gfx950() + return on_gfx950() or on_gfx942() pytestmark = pytest.mark.skipif( - not _on_gfx950(), - reason="The fused KDA decode kernel is only built for gfx950", + not _on_supported_arch(), + reason="The fused KDA decode kernel is only built for gfx942 / gfx950", ) # Kimi-K3 KDA: 96 heads x 128, conv width 4, gate_lower_bound -5.0. diff --git a/vllm/models/kimi_k3/amd/ops/kda_decode.py b/vllm/models/kimi_k3/amd/ops/kda_decode.py index eee6e3a0cff2..132cd58cf816 100644 --- a/vllm/models/kimi_k3/amd/ops/kda_decode.py +++ b/vllm/models/kimi_k3/amd/ops/kda_decode.py @@ -35,7 +35,7 @@ def is_fused_kda_decode_supported( conv_state_dtype: torch.dtype, ) -> bool: """Whether the fused decode kernel can serve this layer on this device.""" - from vllm.platforms.rocm import on_gfx950 + from vllm.platforms.rocm import on_gfx942, on_gfx950 if ( num_heads not in SUPPORTED_NUM_HEADS @@ -48,8 +48,9 @@ def is_fused_kda_decode_supported( or not hasattr(torch.ops._C, "fused_kda_decode") ): return False - # TODO: Verify on other archs; only measured on gfx950 for now - return on_gfx950() + # gfx950 (MI355X) and gfx942 (MI325X): both CDNA, sharing the wave64 / DPP / + # bf16 primitives the kernel relies on. + return on_gfx950() or on_gfx942() def make_decode_conv1d_weight_loader( From 01e56caaf2b2d6f62be5fd78e0a7733fc9c9ed5f Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Tue, 18 Aug 2026 11:20:48 -0700 Subject: [PATCH 098/839] [Bugfix][MLA] Do not use Dense MHA for GLM-5.2 (#52512) Signed-off-by: Woosuk Kwon Co-authored-by: OpenAI Codex --- .../layers/test_mla_short_prefill_indexer.py | 18 ++++++++++++++++-- .../layers/attention/mla_attention.py | 8 ++++++-- vllm/models/deepseek_v32/attention.py | 2 ++ 3 files changed, 24 insertions(+), 4 deletions(-) diff --git a/tests/model_executor/layers/test_mla_short_prefill_indexer.py b/tests/model_executor/layers/test_mla_short_prefill_indexer.py index 6e1e10e8b45c..47bfcd0b38d9 100644 --- a/tests/model_executor/layers/test_mla_short_prefill_indexer.py +++ b/tests/model_executor/layers/test_mla_short_prefill_indexer.py @@ -8,6 +8,7 @@ import vllm.model_executor.layers.sparse_attn_indexer as sparse_indexer from vllm.config import CUDAGraphMode +from vllm.models.deepseek_v32.attention import DeepseekV32Attention from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadata INDEXER_LAYER = "model.layers.0.self_attn.indexer.k_cache" @@ -45,7 +46,15 @@ def make_mla_metadata(*, use_dense_mha: bool = True, num_decode_tokens: int = 0) @pytest.mark.parametrize( "batch_kind", - ["short", "threshold_mismatch", "force_mqa", "mla_decode", "capture", "full"], + [ + "short", + "threshold_mismatch", + "mqa_only_layer", + "force_mqa", + "mla_decode", + "capture", + "full", + ], ) def test_short_prefill_updates_k_cache_before_scoring_decision( monkeypatch: pytest.MonkeyPatch, @@ -132,6 +141,11 @@ def scoring_decode(*args): topk_indices = torch.full((7, 2048), 17, dtype=torch.int32) def run_indexer(): + if batch_kind == "mqa_only_layer": + assert not DeepseekV32Attention.supports_dense_mha_prefill + dense_mha_layer = "" + else: + dense_mha_layer = MLA_LAYER return sparse_indexer.sparse_attn_indexer( hidden_states, INDEXER_LAYER, @@ -149,7 +163,7 @@ def run_indexer(): topk_indices, False, False, - MLA_LAYER, + dense_mha_layer, ) if should_skip: diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 72d65c15748d..27ccd8ca0ccb 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -390,6 +390,8 @@ class MLAAttention(nn.Module, AttentionLayerBase): 3. Return the output tensor. """ + supports_dense_mha_prefill: ClassVar[bool] = True + def __init__( self, num_heads: int, @@ -549,9 +551,11 @@ def __init__( compilation_config.static_forward_context[prefix] = self self.prefill_backend: MLAPrefillBackend | None - if self.impl.is_sparse and not self.impl.supports_dense_mha_prefill: + if self.impl.is_sparse and not ( + self.impl.supports_dense_mha_prefill and self.supports_dense_mha_prefill + ): logger.warning_once( - "Sparse MLA impl has no dense-MHA prefill path; using the top-k " + "Sparse MLA layer has no dense-MHA prefill path; using the top-k " "MQA path only." ) self.prefill_backend = None diff --git a/vllm/models/deepseek_v32/attention.py b/vllm/models/deepseek_v32/attention.py index 87458316c04d..cc270619895b 100644 --- a/vllm/models/deepseek_v32/attention.py +++ b/vllm/models/deepseek_v32/attention.py @@ -160,6 +160,7 @@ class DeepseekV32Attention(MLAAttention): indexer_cls: "type[DeepseekV32Indexer]" = DeepseekV32Indexer require_fp8_kv_cache: bool = True + supports_dense_mha_prefill = False def __init__( self, @@ -265,6 +266,7 @@ def __init__( and not skip_topk and not self.use_pcp and current_platform.is_cuda() + and self.supports_dense_mha_prefill ) self._dense_mha_metadata_layer_name = ( self.layer_name if enable_short_prefill_scoring_skip else "" From d75136c030cc62973dc470d1981199be8de47d62 Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Tue, 18 Aug 2026 11:23:34 -0700 Subject: [PATCH 099/839] [Rust Frontend] Wait for all utility calls to finish (#52671) Signed-off-by: Connor Carpenter Signed-off-by: Bugen Zhao Co-authored-by: Bugen Zhao --- rust/src/engine-core-client/src/client.rs | 118 +++++++------- rust/src/engine-core-client/src/client/imp.rs | 5 + .../engine-core-client/src/client/state.rs | 5 + .../engine-core-client/src/tests/client.rs | 144 +++++++++++++++--- 4 files changed, 187 insertions(+), 85 deletions(-) diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index e7b6e42b6a0b..63689ea3db83 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/rust/src/engine-core-client/src/client.rs @@ -5,7 +5,7 @@ use std::collections::BTreeMap; use std::sync::Arc; use std::time::Duration; -use futures::future::{join_all, try_join_all}; +use futures::future::join_all; use itertools::Itertools; use serde::Serialize; use serde_json::Value as JsonValue; @@ -450,6 +450,11 @@ impl EngineCoreClient { self.config.transport_mode.data_parallel_size() } + #[cfg(test)] + pub(crate) fn pending_utility_call_count(&self) -> usize { + self.inner.pending_utility_call_count() + } + /// Return the engine-side indices connected to this client. pub fn engine_indices(&self) -> Vec { self.engines @@ -596,6 +601,16 @@ impl EngineCoreClient { let (engine_id, rx) = self.inner.register_request(request_id.clone(), lora_name, data_parallel_rank)?; + // Construct the output stream first before actually sending the request to the engine. + // This ensures that cancelling the future will properly clean up the request registration + // via `Drop` of the stream. + let stream = EngineCoreOutputStream::new( + request_id.clone(), + engine_id.engine_index().unwrap_or(0), + self.abort_tx.clone(), + rx, + ); + let result: Result<()> = async { if let Some(coordinator) = self.coordinator.as_ref() { let snapshot = coordinator.snapshot(); @@ -622,12 +637,7 @@ impl EngineCoreClient { return Err(error); } - Ok(EngineCoreOutputStream::new( - request_id, - engine_id.engine_index().unwrap_or(0), - self.abort_tx.clone(), - rx, - )) + Ok(stream) } /// Abort currently in-flight requests by request ID. @@ -651,8 +661,9 @@ impl EngineCoreClient { } /// Call a typed utility method on all connected engines, returning one - /// decoded result per connected engine if all calls succeed or an error - /// if any call fails. + /// decoded result per connected engine if all calls succeed. The client + /// waits for every engine outcome before returning an error so callers can + /// safely compensate partially applied mutations. /// /// Callers should pass utility arguments using Rust tuple semantics so the /// encoded payload matches Python's `(client_index, call_id, @@ -669,62 +680,47 @@ impl EngineCoreClient { "sending utility request" ); - // Phase 1: allocate one call id per engine and build the per-engine - // request payloads up-front. Any failure here (registry closed, encode - // error) must roll back the call ids already allocated so they do not - // leak in the utility registry until shutdown. - let mut pending_calls = Vec::with_capacity(self.engines.len()); - let mut prepared_sends = Vec::with_capacity(self.engines.len()); - for engine in &self.engines { - let (call_id, rx) = match self.inner.allocate_and_register_utility_call() { - Ok(pair) => pair, - Err(err) => { - self.inner.unregister_utility_calls(pending_calls.iter().map(|(id, _)| *id)); - return Err(err); - } - }; - let request = match EngineCoreUtilityRequest::new( - self.config.client_index, - call_id, - method, - &args, - ) { - Ok(request) => request, - Err(err) => { - self.inner.unregister_utility_calls( - pending_calls.iter().map(|(id, _)| *id).chain(std::iter::once(call_id)), - ); - return Err(err); - } - }; - pending_calls.push((call_id, rx)); - prepared_sends.push((&engine.engine_id, request)); + /// Removes utility waiters if a call future is cancelled before completion. + struct UtilityCallGuard<'a> { + inner: &'a ClientInner, + call_ids: Vec, } - // Phase 2: dispatch every utility request concurrently. `try_join_all` - // fails fast on the first transport error and drops the remaining send - // futures; any engines that already received the request will reply, - // but those replies are simply dropped because we roll back the call - // ids below. - let send_futures = prepared_sends.iter().map(|(engine_id, request)| { - self.inner.send_to_engine(engine_id, EngineCoreRequestType::Utility, request) - }); - if let Err(err) = try_join_all(send_futures).await { - self.inner.unregister_utility_calls(pending_calls.iter().map(|(id, _)| *id)); - return Err(err); + impl Drop for UtilityCallGuard<'_> { + fn drop(&mut self) { + self.inner.unregister_utility_calls(self.call_ids.drain(..)); + } } - // Phase 3: wait for all engines to respond and preserve the per-engine - // result list. - let futures = pending_calls.into_iter().map(|(call_id, rx)| async move { - rx.await - .map_err(|_| Error::UtilityCallClosed { - method: method.to_string(), - call_id, - })?? - .into_typed_result(method) - }); - try_join_all(futures).await + let mut call_guard = UtilityCallGuard { + inner: self.inner.as_ref(), + call_ids: Vec::with_capacity(self.engines.len()), + }; + + let mut prepared_calls = Vec::with_capacity(self.engines.len()); + for engine in &self.engines { + let (call_id, rx) = self.inner.allocate_and_register_utility_call()?; + call_guard.call_ids.push(call_id); + let request = + EngineCoreUtilityRequest::new(self.config.client_index, call_id, method, &args)?; + prepared_calls.push((&engine.engine_id, call_id, rx, request)); + } + + let outcomes = join_all(prepared_calls.into_iter().map( + |(engine_id, call_id, rx, request)| async move { + self.inner + .send_to_engine(engine_id, EngineCoreRequestType::Utility, &request) + .await?; + rx.await + .map_err(|_| Error::UtilityCallClosed { + method: method.to_string(), + call_id, + })?? + .into_typed_result(method) + }, + )) + .await; + outcomes.into_iter().collect() } /// Call a utility method on all connected engines and return the shared diff --git a/rust/src/engine-core-client/src/client/imp.rs b/rust/src/engine-core-client/src/client/imp.rs index 9c539bf180e9..117473ce47cc 100644 --- a/rust/src/engine-core-client/src/client/imp.rs +++ b/rust/src/engine-core-client/src/client/imp.rs @@ -118,6 +118,11 @@ impl ClientInner { self.utility_reg.lock().unregister_many(call_ids); } + #[cfg(test)] + pub fn pending_utility_call_count(&self) -> usize { + self.utility_reg.lock().len() + } + /// Undo a request registration when `add_request()` fails. pub fn rollback_request(&self, request_id: &str) { let _ = self.request_reg.lock().remove(request_id); diff --git a/rust/src/engine-core-client/src/client/state.rs b/rust/src/engine-core-client/src/client/state.rs index 59f77bf71cbd..724ae5b64ce1 100644 --- a/rust/src/engine-core-client/src/client/state.rs +++ b/rust/src/engine-core-client/src/client/state.rs @@ -441,6 +441,11 @@ impl UtilityRegistry { self.utility_calls.contains_key(&call_id) } + #[cfg(test)] + pub fn len(&self) -> usize { + self.utility_calls.len() + } + pub fn is_closed(&self) -> bool { self.closed } diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index 2f988e2bfe94..df9eec90b361 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -1431,23 +1431,20 @@ async fn is_sleeping_wrapper_sends_typed_request_and_returns_typed_response() { } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn call_utility_failure_message_surfaces_as_error() { +async fn call_utility_waits_for_all_engines_before_returning_error() { init_tracing(); let ipc = IpcNamespace::new().unwrap(); let handshake_address = ipc.handshake_endpoint(); - let engine_id = b"engine-utility-fail".to_vec(); - - let (shutdown_tx, engine_task) = spawn_mock_engine_task( + let (failure_sent_tx, failure_sent_rx) = oneshot::channel(); + let (shutdown_tx_0, engine_task_0) = spawn_mock_engine_task( handshake_address.clone(), - engine_id.clone(), - |dealer, push| { + EngineId::from_engine_index(0).into_frame().to_vec(), + move |dealer, push| { Box::pin(async move { let utility = recv_engine_message(dealer).await; - assert_eq!(utility[0].as_ref(), &[0x03]); let payload = decode_value(&utility[1]); let call_id = payload.as_array().and_then(|array| array[1].as_u64()).expect("call_id"); - send_outputs( push, UtilityCallOutput { @@ -1462,35 +1459,134 @@ async fn call_utility_failure_message_surfaces_as_error() { .into(), ) .await; + let _ = failure_sent_tx.send(()); + }) + }, + ); + let (second_received_tx, second_received_rx) = oneshot::channel(); + let (release_second_tx, release_second_rx) = oneshot::channel(); + let (shutdown_tx_1, engine_task_1) = spawn_mock_engine_task( + handshake_address.clone(), + EngineId::from_engine_index(1).into_frame().to_vec(), + move |dealer, push| { + Box::pin(async move { + let utility = recv_engine_message(dealer).await; + let payload = decode_value(&utility[1]); + let call_id = + payload.as_array().and_then(|array| array[1].as_u64()).expect("call_id"); + let _ = second_received_tx.send(()); + let _ = release_second_rx.await; + send_outputs( + push, + UtilityCallOutput { + engine_index: 1, + timestamp: 0.0, + output: UtilityOutput { + call_id: call_id.into(), + failure_message: None, + result: Some(utility_result_value(true)), + }, + } + .into(), + ) + .await; }) }, ); - let client = connect_client_with_ipc( - handshake_test_config( - handshake_address, - 1, - "test-model", - Duration::from_secs(2), - 0, - None, - ), - &ipc, - ) - .await; + let client = std::sync::Arc::new( + connect_client_with_ipc( + handshake_test_config( + handshake_address, + 2, + "test-model", + Duration::from_secs(2), + 0, + None, + ), + &ipc, + ) + .await, + ); + let call_client = client.clone(); + let call = + tokio::spawn(async move { call_client.call_utility::("test_mutation", ()).await }); - let error = client.call_utility::("is_sleeping", ()).await.unwrap_err(); + failure_sent_rx.await.unwrap(); + second_received_rx.await.unwrap(); + timeout(Duration::from_secs(2), async { + while client.pending_utility_call_count() != 1 { + tokio::task::yield_now().await; + } + }) + .await + .expect("wait for first engine failure"); + assert!(!call.is_finished()); + + let _ = release_second_tx.send(()); + let error = call.await.unwrap().unwrap_err(); assert!(matches!( error, Error::UtilityCallFailed { method, message, .. - } if method == "is_sleeping" && message == "boom" + } if method == "test_mutation" && message == "boom" )); - let _ = shutdown_tx.send(()); - engine_task.await.unwrap(); + let _ = shutdown_tx_0.send(()); + let _ = shutdown_tx_1.send(()); + engine_task_0.await.unwrap(); + engine_task_1.await.unwrap(); + let client = std::sync::Arc::try_unwrap(client) + .unwrap_or_else(|_| panic!("utility task retained client after completion")); + client.shutdown().await.unwrap(); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn cancelled_utility_call_unregisters_waiter() { + init_tracing(); + let ipc = IpcNamespace::new().unwrap(); + let handshake_address = ipc.handshake_endpoint(); + let (received_tx, received_rx) = oneshot::channel(); + let (_shutdown, engine_task) = spawn_mock_engine_task( + handshake_address.clone(), + vec![0x00, 0x00], + |dealer, _push| { + Box::pin(async move { + let _utility = recv_engine_message(dealer).await; + let _ = received_tx.send(()); + std::future::pending::<()>().await; + }) + }, + ); + let client = std::sync::Arc::new( + connect_client_with_ipc( + handshake_test_config( + handshake_address, + 1, + "test-model", + Duration::from_secs(2), + 0, + None, + ), + &ipc, + ) + .await, + ); + let call_client = client.clone(); + let call = + tokio::spawn(async move { call_client.call_utility::("add_lora", ()).await }); + + received_rx.await.unwrap(); + assert_eq!(client.pending_utility_call_count(), 1); + call.abort(); + assert!(call.await.unwrap_err().is_cancelled()); + assert_eq!(client.pending_utility_call_count(), 0); + + engine_task.abort(); + let client = std::sync::Arc::try_unwrap(client) + .unwrap_or_else(|_| panic!("utility task retained client after cancellation")); client.shutdown().await.unwrap(); } From 90984ddbed27a09409506d6d6c0eea87f54b04b5 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Tue, 18 Aug 2026 13:35:44 -0500 Subject: [PATCH 100/839] [CI] Upgrade huggingface-hub to 1.28.0 (#52797) Signed-off-by: Andreas Karatzas Co-authored-by: OpenAI Codex --- requirements/common.txt | 2 +- requirements/test/cpu.txt | 2 +- requirements/test/cuda.txt | 2 +- requirements/test/rocm.txt | 2 +- requirements/test/xpu.txt | 2 +- 5 files changed, 5 insertions(+), 5 deletions(-) diff --git a/requirements/common.txt b/requirements/common.txt index 7e7a8ce32604..385dd28dacb1 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -8,7 +8,7 @@ tqdm blake3 py-cpuinfo transformers >= 5.5.3 -huggingface_hub >= 1.27.0 +huggingface_hub >= 1.28.0 tokenizers >= 0.21.1 # Required for fast incremental detokenization. safetensors >= 0.6.2 # MXFP4/MXFP6 dtype support (F8_E8M0, F4) added in 0.6.0: https://github.com/huggingface/safetensors/pull/611 protobuf >= 5.29.6, !=6.30.*, !=6.31.*, !=6.32.*, !=6.33.0.*, !=6.33.1.*, !=6.33.2.*, !=6.33.3.*, !=6.33.4.* # Required by LlamaTokenizer, gRPC. CVE-2026-0994 diff --git a/requirements/test/cpu.txt b/requirements/test/cpu.txt index 741e33dc0245..02a7b2a0c534 100644 --- a/requirements/test/cpu.txt +++ b/requirements/test/cpu.txt @@ -339,7 +339,7 @@ httpx==0.27.2 # schemathesis httpx-sse==0.4.3 # via mcp -huggingface-hub==1.27.0 +huggingface-hub==1.28.0 # via # -r requirements/test/../common.txt # accelerate diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index fa3597b77017..2af564acb806 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -358,7 +358,7 @@ httpx==0.27.2 # schemathesis httpx-sse==0.4.3 # via mcp -huggingface-hub==1.27.0 +huggingface-hub==1.28.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt index a7b7172ac22c..526a2490fce5 100644 --- a/requirements/test/rocm.txt +++ b/requirements/test/rocm.txt @@ -349,7 +349,7 @@ httpx==0.27.2 # schemathesis httpx-sse==0.4.3 # via mcp -huggingface-hub==1.27.0 +huggingface-hub==1.28.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt diff --git a/requirements/test/xpu.txt b/requirements/test/xpu.txt index f05459a3d5f5..6e8ac715acc3 100644 --- a/requirements/test/xpu.txt +++ b/requirements/test/xpu.txt @@ -224,7 +224,7 @@ httpx==0.28.1 # schemathesis httpx-sse==0.4.3 # via mcp -huggingface-hub==1.27.0 +huggingface-hub==1.28.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt From 6948a43fbbf427e69ca0d325cd7ab1daee8d131b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Micha=C5=82=20Ganczarenko?= Date: Tue, 18 Aug 2026 20:43:51 +0200 Subject: [PATCH 101/839] [Bugfix] Detect all attention-spelling variants in ModelConfig.is_hybrid (#52161) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Michal Ganczarenko Signed-off-by: Michał Ganczarenko Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> --- vllm/config/model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/config/model.py b/vllm/config/model.py index 54d622743afe..a5f98d8c5cca 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -1850,7 +1850,7 @@ def is_hybrid(self) -> bool: # actually contain any non-attention layers. layer_types = getattr(self.hf_config, "layer_types", None) return layer_types is None or not all( - layer == "attention" for layer in layer_types + layer in ("attention", "full_attention") for layer in layer_types ) @property From 9842d701450214d4b78cd9aefb8eee0c616bce33 Mon Sep 17 00:00:00 2001 From: Sage Moore Date: Tue, 18 Aug 2026 12:00:04 -0700 Subject: [PATCH 102/839] [DBO][CI] Increase the coverage of prefill DBO in test_dbo.py (#48628) Signed-off-by: Sage Moore --- tests/v1/distributed/test_dbo.py | 31 ++++++++++++++++++++----------- 1 file changed, 20 insertions(+), 11 deletions(-) diff --git a/tests/v1/distributed/test_dbo.py b/tests/v1/distributed/test_dbo.py index e5cbe1ce85e9..a4e3d92be6c4 100644 --- a/tests/v1/distributed/test_dbo.py +++ b/tests/v1/distributed/test_dbo.py @@ -11,7 +11,6 @@ import pytest import torch -from tests.evals.gsm8k.gsm8k_eval import evaluate_gsm8k from tests.utils import RemoteOpenAIServer from vllm.utils.import_utils import has_deep_ep @@ -58,6 +57,7 @@ def test_dbo_dp_ep_gsm8k(all2all_backend: str, num_gpus_available): """ Test DBO with DP+EP using GSM8K evaluation. """ + lm_eval = pytest.importorskip("lm_eval") required_gpus = DP_SIZE if num_gpus_available < required_gpus: @@ -79,7 +79,7 @@ def test_dbo_dp_ep_gsm8k(all2all_backend: str, num_gpus_available): "--dbo-decode-token-threshold", "16", "--dbo-prefill-token-threshold", - "256", + "32", "--all2all-backend", all2all_backend, ] @@ -90,19 +90,28 @@ def test_dbo_dp_ep_gsm8k(all2all_backend: str, num_gpus_available): max_wait_seconds=600, # Allow time for model loading with DP+EP ) as remote_server: # Use host and port directly from RemoteOpenAIServer - host = f"http://{remote_server.host}" - port = remote_server.port + + base_url = f"http://{remote_server.host}:{remote_server.port}/v1/completions" # Run GSM8K evaluation - results = evaluate_gsm8k( - num_questions=NUM_QUESTIONS, - num_shots=NUM_SHOTS, - host=host, - port=port, + results = lm_eval.simple_evaluate( + model="local-completions", + model_args=( + f"pretrained={MODEL_NAME}," + f"base_url={base_url}," + "num_concurrent=512,max_retries=3" + ), + tasks=["gsm8k"], + num_fewshot=NUM_SHOTS, + limit=NUM_QUESTIONS, ) - # Validate accuracy is reasonable - accuracy = results["accuracy"] + gsm8k = results["results"]["gsm8k"] + accuracy = gsm8k.get( + "exact_match,strict-match", + gsm8k.get("exact_match,flexible-extract"), + ) + assert accuracy is not None, f"gsm8k exact_match missing: {gsm8k}" assert accuracy >= MIN_ACCURACY, ( f"DBO+DP+EP accuracy too low ({all2all_backend}): " f"{accuracy:.3f} < {MIN_ACCURACY:.3f} " From 8d6b18329a303c3d4a8f1a7b515cd69a4ef64800 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Tue, 18 Aug 2026 12:50:32 -0700 Subject: [PATCH 103/839] [CI] Standardize test job labels by device (#52659) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Co-authored-by: OpenAI Codex --- .buildkite/test_areas/attention.yaml | 5 +- .buildkite/test_areas/basic_correctness.yaml | 3 +- .buildkite/test_areas/benchmarks.yaml | 5 +- .buildkite/test_areas/compile.yaml | 22 ++++---- .buildkite/test_areas/cuda.yaml | 4 +- .buildkite/test_areas/disaggregated.yaml | 31 ++++++----- .../test_areas/disaggregated_mooncake.yaml | 2 +- .buildkite/test_areas/distributed.yaml | 29 +++++----- .buildkite/test_areas/docker.yaml | 2 +- .buildkite/test_areas/e2e_integration.yaml | 10 ++-- .buildkite/test_areas/engine.yaml | 17 +++--- .buildkite/test_areas/entrypoints.yaml | 28 ++++++---- .buildkite/test_areas/expert_parallelism.yaml | 7 +-- .buildkite/test_areas/fault_tolerance.yaml | 2 +- .buildkite/test_areas/jit_monitor.yaml | 2 +- .buildkite/test_areas/kernels.yaml | 45 ++++++++-------- .buildkite/test_areas/lm_eval.yaml | 54 ++++++++++--------- .buildkite/test_areas/lora.yaml | 5 +- .buildkite/test_areas/misc.yaml | 35 ++++++------ .buildkite/test_areas/model_executor.yaml | 3 +- .buildkite/test_areas/model_runner_v2.yaml | 10 ++-- .buildkite/test_areas/models_basic.yaml | 13 ++--- .buildkite/test_areas/models_distributed.yaml | 2 +- .buildkite/test_areas/models_language.yaml | 20 ++++--- .buildkite/test_areas/models_multimodal.yaml | 32 ++++++----- .buildkite/test_areas/plugins.yaml | 8 +-- .buildkite/test_areas/pytorch.yaml | 13 ++--- .buildkite/test_areas/quantization.yaml | 8 +-- .buildkite/test_areas/ray_compat.yaml | 2 +- .buildkite/test_areas/rust_frontend.yaml | 10 ++-- .../test_areas/rust_frontend_cargo.yaml | 4 +- .buildkite/test_areas/samplers.yaml | 3 +- .buildkite/test_areas/spec_decode.yaml | 37 ++++++++----- .buildkite/test_areas/torch_abi.yaml | 2 +- .buildkite/test_areas/weight_loading.yaml | 3 +- 35 files changed, 268 insertions(+), 210 deletions(-) diff --git a/.buildkite/test_areas/attention.yaml b/.buildkite/test_areas/attention.yaml index 71d2975a5f41..e11108d92593 100644 --- a/.buildkite/test_areas/attention.yaml +++ b/.buildkite/test_areas/attention.yaml @@ -2,7 +2,7 @@ group: Attention depends_on: - image-build steps: -- label: V1 attention (H100-MI300) +- label: ":nvidia: (H100) V1 Attention Shard %N" key: v1-attention-h100-mi300 timeout_in_minutes: 85 device: h100 @@ -16,6 +16,7 @@ steps: parallelism: 2 mirror: amd: + label: ":amd: (MI300) V1 Attention Shard %N" dind: false device: mi300_1 timeout_in_minutes: 125 @@ -30,7 +31,7 @@ steps: - vllm/envs.py - vllm/platforms/rocm.py -- label: V1 attention (B200) +- label: ":nvidia: (B200) V1 Attention Shard %N" key: v1-attention-b200 timeout_in_minutes: 80 device: b200-k8s diff --git a/.buildkite/test_areas/basic_correctness.yaml b/.buildkite/test_areas/basic_correctness.yaml index d0dac3a112c9..1ce9dcbe8842 100644 --- a/.buildkite/test_areas/basic_correctness.yaml +++ b/.buildkite/test_areas/basic_correctness.yaml @@ -2,7 +2,7 @@ group: Basic Correctness depends_on: - image-build steps: -- label: Basic Correctness +- label: ":nvidia: (H200) Basic Correctness" key: basic-correctness timeout_in_minutes: 68 device: h200_18gb @@ -19,6 +19,7 @@ steps: - pytest -v -s basic_correctness/test_cpu_offload.py mirror: amd: + label: ":amd: (MI300) Basic Correctness" dind: false device: mi300_1 timeout_in_minutes: 60 diff --git a/.buildkite/test_areas/benchmarks.yaml b/.buildkite/test_areas/benchmarks.yaml index ccf3aa9cb87c..4010b8d0b8d8 100644 --- a/.buildkite/test_areas/benchmarks.yaml +++ b/.buildkite/test_areas/benchmarks.yaml @@ -2,7 +2,7 @@ group: Benchmarks depends_on: - image-build steps: -- label: Benchmarks CLI Test +- label: ":nvidia: (H200) Benchmarks CLI" key: benchmarks-cli-test timeout_in_minutes: 45 device: h200_18gb @@ -14,13 +14,14 @@ steps: - pytest -v -s benchmarks/ mirror: amd: + label: ":amd: (MI300) Benchmarks CLI" dind: false device: mi300_1 timeout_in_minutes: 40 depends_on: - image-build-amd -- label: Attention Benchmarks Smoke Test (B200) +- label: ":nvidia: (B200) Attention Benchmark Smoke" key: attention-benchmarks-smoke-test-b200 device: b200-k8s num_gpus: 2 diff --git a/.buildkite/test_areas/compile.yaml b/.buildkite/test_areas/compile.yaml index 3ce4d7709e30..dc847bf2c33a 100644 --- a/.buildkite/test_areas/compile.yaml +++ b/.buildkite/test_areas/compile.yaml @@ -2,7 +2,7 @@ group: Compile depends_on: - image-build steps: -- label: Sequence Parallel Correctness Tests (2xB200) +- label: ":nvidia: (B200) Sequence Parallel Correctness" key: sequence-parallel-correctness-tests-2xb200 timeout_in_minutes: 45 working_dir: "/vllm-workspace/" @@ -13,7 +13,7 @@ steps: - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py -- label: AsyncTP Correctness Tests (2xB200) +- label: ":nvidia: (B200) AsyncTP Correctness" key: asynctp-correctness-tests-b200 timeout_in_minutes: 30 working_dir: "/vllm-workspace/" @@ -24,7 +24,7 @@ steps: - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py -- label: Distributed Compile Unit Tests (2xH100) +- label: ":nvidia: (H100) Distributed Compile" key: distributed-compile-unit-tests-2xh100 timeout_in_minutes: 45 working_dir: "/vllm-workspace/" @@ -38,7 +38,7 @@ steps: - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - pytest -s -v tests/compile/passes/distributed -- label: Fusion and Compile Unit Tests (2xB200) +- label: ":nvidia: (B200) Fusion and Compile" key: fusion-and-compile-unit-tests-2xb200 timeout_in_minutes: 30 working_dir: "/vllm-workspace/" @@ -65,7 +65,7 @@ steps: # this runner has 2 GPUs available even though num_devices=2 is not set - pytest -v -s tests/compile/passes/distributed/test_fusion_all_reduce.py -- label: Fusion E2E Quick (H100) +- label: ":nvidia: (H100) Fusion E2E Quick" key: fusion-e2e-quick-h100 timeout_in_minutes: 25 working_dir: "/vllm-workspace/" @@ -84,7 +84,7 @@ steps: # Qwen/Deepseek requires +quant_fp8 as -quant_fp8 rms+quant fusion is not supported - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "inductor_partition and not +rms_norm and +quant_fp8 and (qwen3 or deepseek)" -- label: Fusion E2E Config Sweep (H100) +- label: ":nvidia: (H100) Fusion E2E Config Sweep" key: fusion-e2e-config-sweep-h100 timeout_in_minutes: 25 working_dir: "/vllm-workspace/" @@ -104,7 +104,7 @@ steps: # Run just llama3 (fp8) for all config combinations - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "llama-3" -- label: Fusion E2E Config Sweep (B200) +- label: ":nvidia: (B200) Fusion E2E Config Sweep" key: fusion-e2e-config-sweep-b200 timeout_in_minutes: 30 working_dir: "/vllm-workspace/" @@ -118,7 +118,7 @@ steps: # Run just llama3 (fp8 & fp4) for all config combinations (only inductor partition) - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "inductor_partition and (FLASHINFER and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek)) or llama-3)" -- label: Fusion E2E TP2 Quick (H100) +- label: ":nvidia: (H100) Fusion E2E TP2 Quick" key: fusion-e2e-tp2-quick-h100 timeout_in_minutes: 35 working_dir: "/vllm-workspace/" @@ -136,7 +136,7 @@ steps: - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek))" - pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek))" -- label: Fusion E2E TP2 AR-RMS Config Sweep (H100) +- label: ":nvidia: (H100) Fusion E2E TP2 AR-RMS Config Sweep" key: fusion-e2e-tp2-ar-rms-config-sweep-h100 timeout_in_minutes: 30 working_dir: "/vllm-workspace/" @@ -156,7 +156,7 @@ steps: # Run just llama3 (fp8 & bf16) for all config combinations - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "llama-3" -- label: Fusion E2E TP2 AsyncTP Config Sweep (H100) +- label: ":nvidia: (H100) Fusion E2E TP2 AsyncTP Config Sweep" key: fusion-e2e-tp2-asynctp-config-sweep-h100 timeout_in_minutes: 40 working_dir: "/vllm-workspace/" @@ -176,7 +176,7 @@ steps: # Run just llama3 (fp8 & bf16) for all config combinations - pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "llama-3" -- label: Fusion E2E TP2 (B200) +- label: ":nvidia: (B200) Fusion E2E TP2" key: fusion-e2e-tp2-b200 timeout_in_minutes: 45 working_dir: "/vllm-workspace/" diff --git a/.buildkite/test_areas/cuda.yaml b/.buildkite/test_areas/cuda.yaml index 431ce07af4d1..7ddbaeec6554 100644 --- a/.buildkite/test_areas/cuda.yaml +++ b/.buildkite/test_areas/cuda.yaml @@ -2,7 +2,7 @@ group: CUDA depends_on: - image-build steps: -- label: Platform Tests +- label: ":nvidia: (H200) CUDA Platform" key: platform-tests timeout_in_minutes: 20 device: h200_18gb @@ -18,7 +18,7 @@ steps: - pytest -v -s cuda/test_platform_no_cuda_init.py - pytest -v -s cuda/test_cuda_compatibility_path.py -- label: Cudagraph +- label: ":nvidia: (H200) CUDAGraph" device: h200_35gb key: cudagraph timeout_in_minutes: 30 diff --git a/.buildkite/test_areas/disaggregated.yaml b/.buildkite/test_areas/disaggregated.yaml index dd662168d2bc..dec438a7becb 100644 --- a/.buildkite/test_areas/disaggregated.yaml +++ b/.buildkite/test_areas/disaggregated.yaml @@ -2,7 +2,7 @@ group: Disaggregated depends_on: - image-build steps: -- label: Distributed NixlConnector PD accuracy (4 GPUs) +- label: ":nvidia: (L4) Distributed NixlConnector PD accuracy" key: distributed-nixlconnector-pd-accuracy-4-gpus timeout_in_minutes: 55 working_dir: "/vllm-workspace/tests" @@ -15,6 +15,7 @@ steps: - bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh mirror: amd: + label: ":amd: (MI300) Distributed NixlConnector PD accuracy" dind: false device: mi300_4 timeout_in_minutes: 60 @@ -28,7 +29,7 @@ steps: - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh -- label: Distributed FlashInfer NixlConnector PD accuracy (4 GPUs) +- label: ":nvidia: (L4) Distributed FlashInfer NixlConnector PD accuracy" key: distributed-flashinfer-nixlconnector-pd-accuracy-4-gpus timeout_in_minutes: 55 working_dir: "/vllm-workspace/tests" @@ -40,7 +41,7 @@ steps: - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - FLASHINFER=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh -- label: Push NixlConnector PP prefill PD accuracy (4 GPUs) +- label: ":nvidia: (L4) Push NixlConnector PP prefill PD accuracy" key: push-nixlconnector-pp-prefill-pd-accuracy-4-gpus timeout_in_minutes: 30 working_dir: "/vllm-workspace/tests" @@ -53,7 +54,7 @@ steps: - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - bash v1/kv_connector/nixl_push_integration/config_sweep_accuracy_test.sh -- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs) +- label: ":nvidia: (L4) DP EP Distributed NixlConnector PD accuracy" key: dp-ep-distributed-nixlconnector-pd-accuracy-tests-4-gpus timeout_in_minutes: 30 working_dir: "/vllm-workspace/tests" @@ -66,6 +67,7 @@ steps: - DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh mirror: amd: + label: ":amd: (MI300) DP EP Distributed NixlConnector PD accuracy" dind: false device: mi300_4 timeout_in_minutes: 40 @@ -79,7 +81,7 @@ steps: - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - DP_EP=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh -- label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs) +- label: ":nvidia: (L4) CrossLayer KV layout Distributed NixlConnector PD accuracy" key: crosslayer-kv-layout-distributed-nixlconnector-pd-accuracy-tests-4-gpus timeout_in_minutes: 55 working_dir: "/vllm-workspace/tests" @@ -92,6 +94,7 @@ steps: - CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh mirror: amd: + label: ":amd: (MI300) CrossLayer KV layout Distributed NixlConnector PD accuracy" dind: false device: mi300_4 timeout_in_minutes: 60 @@ -105,7 +108,7 @@ steps: - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - CROSS_LAYERS_BLOCKS=True ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh -- label: Hybrid SSM NixlConnector PD accuracy tests (4 GPUs) +- label: ":nvidia: (L4) Hybrid SSM NixlConnector PD accuracy" key: hybrid-ssm-nixlconnector-pd-accuracy-tests-4-gpus timeout_in_minutes: 60 working_dir: "/vllm-workspace/tests" @@ -118,6 +121,7 @@ steps: - HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh mirror: amd: + label: ":amd: (MI300) Hybrid SSM NixlConnector PD accuracy" dind: false device: mi300_4 timeout_in_minutes: 55 @@ -131,7 +135,7 @@ steps: - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - HYBRID_SSM=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh -- label: NixlConnector PD edge case test (2 GPUs) +- label: ":nvidia: (L4) NixlConnector PD edge case" key: nixlconnector-pd-edge-cases-2-gpus timeout_in_minutes: 40 working_dir: "/vllm-workspace/tests" @@ -147,7 +151,7 @@ steps: - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - bash v1/kv_connector/nixl_integration/run_edge_case_test.sh -- label: Hybrid SSM NixlConnector PD prefix cache test (2 GPUs) +- label: ":nvidia: (L4) Hybrid SSM NixlConnector PD prefix cache" key: hybrid-ssm-nixlconnector-pd-prefix-cache-2-gpus timeout_in_minutes: 25 working_dir: "/vllm-workspace/tests" @@ -161,7 +165,7 @@ steps: - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - bash v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh -- label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs) +- label: ":nvidia: (L4) MultiConnector (Nixl+Offloading) PD accuracy" key: multiconnector-nixl-offloading-pd-accuracy-2-gpus timeout_in_minutes: 40 working_dir: "/vllm-workspace/tests" @@ -176,7 +180,7 @@ steps: - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh -- label: NixlConnector PD + Spec Decode acceptance (2 GPUs) +- label: ":nvidia: (A100) NixlConnector PD + Spec Decode acceptance" key: nixlconnector-pd-spec-decode-acceptance-2-gpus timeout_in_minutes: 45 device: a100 @@ -191,6 +195,7 @@ steps: - bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh mirror: amd: + label: ":amd: (MI300) NixlConnector PD + Spec Decode acceptance" dind: false device: mi300_2 timeout_in_minutes: 45 @@ -205,7 +210,7 @@ steps: - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - KV_CACHE_MEMORY_BYTES=8G ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh -- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs) +- label: ":nvidia: (L4) MultiConnector (Nixl+Offloading) PD edge cases" key: multiconnector-nixl-offloading-pd-edge-cases-2-gpus timeout_in_minutes: 25 working_dir: "/vllm-workspace/tests" @@ -221,7 +226,7 @@ steps: - bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh # P TP 4 - D DPEP 4 test case for DSv4-Flash -- label: DSv4-Flash Disaggregated DP EP +- label: ":nvidia: (H200) DSv4-Flash Disaggregated DP EP" key: dsv4-flash-disaggregated timeout_in_minutes: 60 device: h200 @@ -242,7 +247,7 @@ steps: - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh -- label: Kimi-Linear-48B-A3B Disaggregated DP EP +- label: ":nvidia: (H200) Kimi-Linear-48B-A3B Disaggregated DP EP" key: kimi-linear-disaggregated timeout_in_minutes: 60 device: h200 diff --git a/.buildkite/test_areas/disaggregated_mooncake.yaml b/.buildkite/test_areas/disaggregated_mooncake.yaml index 92a9e625d1b8..0911507eb22c 100644 --- a/.buildkite/test_areas/disaggregated_mooncake.yaml +++ b/.buildkite/test_areas/disaggregated_mooncake.yaml @@ -2,7 +2,7 @@ group: Disaggregated Mooncake depends_on: - image-build steps: -- label: Distributed MooncakeConnector PD accuracy (4 GPUs) +- label: ":nvidia: (B200) Distributed MooncakeConnector PD accuracy" key: distributed-mooncakeconnector-pd-accuracy-4-gpus timeout_in_minutes: 10 working_dir: "/vllm-workspace/tests" diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index 32d1ff0342c7..b4b2b5ddc605 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -2,7 +2,7 @@ group: Distributed depends_on: - image-build steps: -- label: Distributed Comm Ops +- label: ":nvidia: (L4) Distributed Comm Ops" key: distributed-comm-ops timeout_in_minutes: 25 working_dir: "/vllm-workspace/tests" @@ -17,7 +17,7 @@ steps: - pytest -v -s distributed/test_shm_buffer.py - pytest -v -s distributed/test_shm_storage.py -- label: Distributed DP Tests (2 GPUs) +- label: ":nvidia: (L4) Distributed DP Basic" key: distributed-dp-tests-2-gpus timeout_in_minutes: 35 working_dir: "/vllm-workspace/tests" @@ -41,6 +41,7 @@ steps: - DP_SIZE=2 pytest -v -s entrypoints/openai/test_multi_api_servers.py mirror: amd: + label: ":amd: (MI300) Distributed DP Basic" dind: false device: mi300_2 timeout_in_minutes: 45 @@ -58,7 +59,7 @@ steps: - tests/entrypoints/openai/test_multi_api_servers.py - vllm/platforms/rocm.py -- label: Distributed Compile + RPC Tests (2 GPUs) +- label: ":nvidia: (L4) Distributed Compile + RPC" key: distributed-compile-rpc-tests-2-gpus timeout_in_minutes: 65 working_dir: "/vllm-workspace/tests" @@ -80,7 +81,7 @@ steps: - pytest -v -s entrypoints/llm/test_collective_rpc.py - pytest -v -s ./compile/test_wrapper.py -- label: Distributed Torchrun + Shutdown Tests (2 GPUs) +- label: ":nvidia: (L4) Distributed Torchrun + Shutdown" key: distributed-torchrun-shutdown-tests-2-gpus timeout_in_minutes: 30 working_dir: "/vllm-workspace/tests" @@ -104,7 +105,7 @@ steps: - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown - pytest -v -s v1/worker/test_worker_memory_snapshot.py -- label: Distributed Torchrun + Examples (4 GPUs) +- label: ":nvidia: (L4) Distributed Torchrun + Examples" key: distributed-torchrun-examples-4-gpus timeout_in_minutes: 30 working_dir: "/vllm-workspace" @@ -137,7 +138,7 @@ steps: - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_http_nccl.py - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_http_ipc.py -- label: Distributed DP Tests (4 GPUs) +- label: ":nvidia: (L4) Distributed DP Extended" key: distributed-dp-tests-4-gpus timeout_in_minutes: 45 working_dir: "/vllm-workspace/tests" @@ -160,7 +161,7 @@ steps: - pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp - pytest -v -s distributed/test_utils.py -- label: Distributed Compile + Comm (4 GPUs) +- label: ":nvidia: (L4) Distributed Compile + Comm" key: distributed-compile-comm-4-gpus timeout_in_minutes: 70 working_dir: "/vllm-workspace/tests" @@ -181,7 +182,7 @@ steps: # test multi-node TP with multiproc executor (simulated on single node) - pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node -- label: Distributed Tests (8xH100) +- label: ":nvidia: (H100) Distributed DP + EP" key: distributed-tests-8xh100 timeout_in_minutes: 20 device: h100 @@ -203,7 +204,7 @@ steps: # test with torchrun tp=2 and dp=4 with ep - torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep -- label: Distributed Tests (4xA100) +- label: ":nvidia: (A100) Distributed" key: distributed-tests-4xa100 device: a100 optional: true @@ -219,7 +220,7 @@ steps: - TARGET_TEST_SUITE=A100 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' - pytest -v -s -x lora/test_mixtral.py -- label: Distributed Tests (2xH100-2xMI300) +- label: ":nvidia: (H100) Distributed Features" key: distributed-tests-2xh100-2xmi300 timeout_in_minutes: 30 device: h100 @@ -234,7 +235,7 @@ steps: - VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py - pytest -v -s tests/distributed/test_packed_tensor.py -- label: Distributed Tests (2xB200) +- label: ":nvidia: (B200) Distributed" key: distributed-tests-2xb200 device: b200-k8s optional: true @@ -249,7 +250,7 @@ steps: -- label: 2 Node Test (4 GPUs) +- label: ":nvidia: (L4) Distributed 2-Node" key: 2-node-test-4-gpus timeout_in_minutes: 30 working_dir: "/vllm-workspace/tests" @@ -268,7 +269,7 @@ steps: commands: - ./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 2 $IMAGE_TAG "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/features/data_parallel/data_parallel_offline.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=0 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_multi_node_assignment.py && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_pipeline_parallel.py" "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/features/data_parallel/data_parallel_offline.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=1 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code" -- label: Pipeline + Context Parallelism (4 GPUs) +- label: ":nvidia: (L4) Pipeline + Context Parallelism" key: pipeline-context-parallelism-4-gpus timeout_in_minutes: 55 working_dir: "/vllm-workspace/tests" @@ -284,7 +285,7 @@ steps: - pytest -v -s distributed/test_pp_cudagraph.py - pytest -v -s distributed/test_pipeline_parallel.py -- label: RayExecutorV2 (4 GPUs) +- label: ":nvidia: (L4) RayExecutorV2" key: rayexecutorv2-4-gpus timeout_in_minutes: 45 working_dir: "/vllm-workspace/tests" diff --git a/.buildkite/test_areas/docker.yaml b/.buildkite/test_areas/docker.yaml index 9f0562ca3bcf..4bd36fad18e1 100644 --- a/.buildkite/test_areas/docker.yaml +++ b/.buildkite/test_areas/docker.yaml @@ -2,7 +2,7 @@ group: Docker depends_on: - image-build-cpu steps: -- label: Docker Build Metadata +- label: ":computer: (CPU) Docker Build Metadata" timeout_in_minutes: 20 device: cpu-small source_file_dependencies: diff --git a/.buildkite/test_areas/e2e_integration.yaml b/.buildkite/test_areas/e2e_integration.yaml index 6655ae781e8d..8f0e2042bf10 100644 --- a/.buildkite/test_areas/e2e_integration.yaml +++ b/.buildkite/test_areas/e2e_integration.yaml @@ -2,7 +2,7 @@ group: E2E Integration depends_on: - image-build steps: -- label: DeepSeek V2-Lite Sync EPLB Accuracy (4xH100) +- label: ":nvidia: (H100) DeepSeek V2-Lite Sync EPLB Accuracy" key: deepseek-v2-lite-sync-eplb-accuracy-4xh100 timeout_in_minutes: 25 device: h100 @@ -12,7 +12,7 @@ steps: commands: - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh 0.25 200 8010 -- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (4xH100) +- label: ":nvidia: (H100) Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy" key: qwen3-30b-a3b-fp8-block-sync-eplb-accuracy-4xh100 timeout_in_minutes: 25 device: h100 @@ -22,7 +22,7 @@ steps: commands: - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 -- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (2xB200) +- label: ":nvidia: (B200) Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy" key: qwen3-30b-a3b-fp8-block-sync-eplb-accuracy-2xb200 timeout_in_minutes: 20 device: b200-k8s @@ -32,7 +32,7 @@ steps: commands: - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 2 1 -- label: Qwen3-30B-A3B-FP8 DP4 Async EPLB Accuracy +- label: ":nvidia: (H100) Qwen3-30B-A3B-FP8 DP4 Async EPLB Accuracy" key: qwen3-30b-a3b-fp8-dp4-async-eplb-accuracy timeout_in_minutes: 25 device: h100 @@ -42,7 +42,7 @@ steps: commands: - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_dp4_async_eplb.sh 0.8 200 8050 -- label: DeepSeek V2-Lite Prefetch Offload Accuracy (H100) +- label: ":nvidia: (H100) DeepSeek V2-Lite Prefetch Offload Accuracy" key: deepseek-v2-lite-prefetch-offload-accuracy-h100 timeout_in_minutes: 20 device: h100 diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml index 4cb661ac162c..d7bb76c69eff 100644 --- a/.buildkite/test_areas/engine.yaml +++ b/.buildkite/test_areas/engine.yaml @@ -2,7 +2,7 @@ group: Engine depends_on: - image-build steps: -- label: Engine +- label: ":nvidia: (H200) Engine Core" key: engine timeout_in_minutes: 30 device: h200_18gb @@ -29,13 +29,14 @@ steps: - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py jit_monitor/test_hooks.py jit_monitor/test_hooks_gpu.py mirror: amd: + label: ":amd: (MI300) Engine Core" dind: false device: mi300_1 timeout_in_minutes: 40 depends_on: - image-build-amd -- label: Engine (1 GPU) +- label: ":nvidia: (L4) V1 Engine" key: engine-1-gpu timeout_in_minutes: 45 source_file_dependencies: @@ -48,12 +49,13 @@ steps: - pytest -v -s v1/test_tensor_ipc_queue.py mirror: amd: + label: ":amd: (MI250) V1 Engine" device: mi250_1 timeout_in_minutes: 45 depends_on: - image-build-amd -- label: e2e Scheduling (1 GPU) +- label: ":nvidia: (H200) E2E Scheduling" key: e2e-scheduling-1-gpu timeout_in_minutes: 53 device: h200_18gb @@ -64,12 +66,13 @@ steps: - pytest -v -s v1/e2e/general/test_async_scheduling.py mirror: amd: + label: ":amd: (MI250) E2E Scheduling" device: mi250_1 timeout_in_minutes: 55 depends_on: - image-build-amd -- label: e2e Core (1 GPU) +- label: ":nvidia: (H200) E2E Core" device: h200_35gb key: e2e-core-1-gpu timeout_in_minutes: 40 @@ -80,6 +83,7 @@ steps: - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py mirror: amd: + label: ":amd: (MI250) E2E Core" device: mi250_1 timeout_in_minutes: 50 depends_on: @@ -89,7 +93,7 @@ steps: - tests/v1/e2e/general/ - vllm/platforms/rocm.py -- label: V1 e2e (2 GPUs) +- label: ":nvidia: (L4) V1 E2E" key: v1-e2e-2-gpus timeout_in_minutes: 25 # TODO: Fix timeout after we have more confidence in the test stability optional: true @@ -122,13 +126,14 @@ steps: v1/e2e/spec_decode/draft_model/test_draft_model.py::test_draft_model_engine_args_tensor_parallelism mirror: amd: + label: ":amd: (MI300) V1 E2E" dind: false device: mi300_2 timeout_in_minutes: 30 depends_on: - image-build-amd -- label: V1 e2e (4xH100) +- label: ":nvidia: (H100) V1 E2E" key: v1-e2e-4xh100 timeout_in_minutes: 35 device: h100 diff --git a/.buildkite/test_areas/entrypoints.yaml b/.buildkite/test_areas/entrypoints.yaml index d36f98bd6783..ed143b881205 100644 --- a/.buildkite/test_areas/entrypoints.yaml +++ b/.buildkite/test_areas/entrypoints.yaml @@ -2,7 +2,7 @@ group: Entrypoints depends_on: - image-build steps: -- label: Entrypoints Unit Tests +- label: ":nvidia: (H200) Entrypoints Unit" device: h200_35gb key: entrypoints-unit-tests timeout_in_minutes: 25 @@ -15,7 +15,7 @@ steps: - pytest -v -s entrypoints/unit_tests - pytest -v -s entrypoints/weight_transfer -- label: Entrypoints Integration (LLM) +- label: ":nvidia: (H200) Entrypoints Integration (LLM)" device: h200_35gb key: entrypoints-integration-llm timeout_in_minutes: 60 @@ -31,13 +31,14 @@ steps: - pytest -v -s entrypoints/llm/offline_mode # Needs to avoid interference with other tests mirror: amd: + label: ":amd: (MI300) Entrypoints Integration (LLM)" dind: false device: mi300_1 timeout_in_minutes: 55 depends_on: - image-build-amd -- label: Entrypoints Integration (API Server) +- label: ":nvidia: (H200) Entrypoints Integration (API Server)" key: entrypoints-integration-api-server device: h200_35gb timeout_in_minutes: 75 @@ -54,13 +55,14 @@ steps: - pytest -v -s entrypoints/scale_out mirror: amd: + label: ":amd: (MI300) Entrypoints Integration (API Server)" dind: false device: mi300_1 timeout_in_minutes: 65 depends_on: - image-build-amd -- label: Entrypoints Integration (API Server OpenAI - Part 1) +- label: ":nvidia: (H200) Entrypoints Integration (API Server OpenAI - Part 1)" device: h200_35gb key: entrypoints-integration-api-server-openai-part-1 timeout_in_minutes: 68 @@ -75,13 +77,14 @@ steps: - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness mirror: amd: + label: ":amd: (MI300) Entrypoints Integration (API Server OpenAI - Part 1)" dind: false device: mi300_1 timeout_in_minutes: 65 depends_on: - image-build-amd -- label: Entrypoints Integration (API Server OpenAI - Part 2) +- label: ":nvidia: (H200) Entrypoints Integration (API Server OpenAI - Part 2)" device: h200_35gb key: entrypoints-integration-api-server-openai-part-2 timeout_in_minutes: 83 @@ -97,13 +100,14 @@ steps: - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py mirror: amd: + label: ":amd: (MI300) Entrypoints Integration (API Server OpenAI - Part 2)" dind: false device: mi300_1 timeout_in_minutes: 70 depends_on: - image-build-amd -- label: Entrypoints Integration (API Server Generate) +- label: ":nvidia: (H200) Entrypoints Integration (API Server Generate)" device: h200_35gb key: entrypoints-integration-api-server-generate timeout_in_minutes: 50 @@ -122,13 +126,14 @@ steps: - pytest -v -s entrypoints/anthropic mirror: amd: + label: ":amd: (MI300) Entrypoints Integration (API Server Generate)" dind: false device: mi300_1 timeout_in_minutes: 65 depends_on: - image-build-amd -- label: Entrypoints Integration (Responses API) +- label: ":nvidia: (H200) Entrypoints Integration (Responses API)" device: h200_35gb key: entrypoints-integration-responses-api timeout_in_minutes: 50 @@ -140,7 +145,7 @@ steps: commands: - pytest -v -s entrypoints/openai/responses -- label: Entrypoints Integration (Speech to Text) +- label: ":nvidia: (H200) Entrypoints Integration (Speech to Text)" device: h200_35gb key: entrypoints-integration-speech_to_text timeout_in_minutes: 45 @@ -153,7 +158,7 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/speech_to_text -- label: Entrypoints Integration (Multimodal) +- label: ":nvidia: (H200) Entrypoints Integration (Multimodal)" device: h200_35gb key: entrypoints-integration-multimodal timeout_in_minutes: 45 @@ -166,7 +171,7 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/multimodal -- label: Entrypoints Integration (Pooling) +- label: ":nvidia: (H200) Entrypoints Integration (Pooling)" device: h200_35gb key: entrypoints-integration-pooling timeout_in_minutes: 75 @@ -179,7 +184,7 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/pooling -- label: OpenAI API Correctness +- label: ":nvidia: (H200) OpenAI API Correctness" key: openai-api-correctness timeout_in_minutes: 20 device: h200_18gb @@ -190,6 +195,7 @@ steps: - pytest -s entrypoints/openai/correctness/ mirror: amd: + label: ":amd: (MI300) OpenAI API Correctness" dind: false device: mi300_1 timeout_in_minutes: 30 diff --git a/.buildkite/test_areas/expert_parallelism.yaml b/.buildkite/test_areas/expert_parallelism.yaml index 89c9fddf64a7..e0b931f5fe9f 100644 --- a/.buildkite/test_areas/expert_parallelism.yaml +++ b/.buildkite/test_areas/expert_parallelism.yaml @@ -2,7 +2,7 @@ group: Expert Parallelism depends_on: - image-build steps: -- label: EPLB Algorithm +- label: ":nvidia: (H200) EPLB Algorithm" key: eplb-algorithm timeout_in_minutes: 20 device: h200_18gb @@ -16,6 +16,7 @@ steps: - pytest -v -s distributed/test_eplb_utils.py mirror: amd: + label: ":amd: (MI300) EPLB Algorithm" dind: false device: mi300_1 timeout_in_minutes: 30 @@ -27,7 +28,7 @@ steps: - tests/distributed/test_eplb_utils.py - vllm/platforms/rocm.py -- label: EPLB Execution # 17min +- label: ":nvidia: (L4) EPLB Execution" key: eplb-execution timeout_in_minutes: 25 working_dir: "/vllm-workspace/tests" @@ -39,7 +40,7 @@ steps: - pytest -v -s distributed/test_eplb_execute.py - pytest -v -s distributed/test_eplb_spec_decode.py -- label: Elastic EP Scaling Test +- label: ":nvidia: (H100) Elastic EP Scaling" key: elastic-ep-scaling-test timeout_in_minutes: 30 device: h100 diff --git a/.buildkite/test_areas/fault_tolerance.yaml b/.buildkite/test_areas/fault_tolerance.yaml index e2f700a8bd86..2e85a1c3f856 100644 --- a/.buildkite/test_areas/fault_tolerance.yaml +++ b/.buildkite/test_areas/fault_tolerance.yaml @@ -2,7 +2,7 @@ group: Fault Tolerance depends_on: - image-build steps: -- label: Fault Tolerance E2E (2xH100) +- label: ":nvidia: (H100) Fault Tolerance E2E" key: fault-tolerance-e2e-2xh100 timeout_in_minutes: 35 device: h100 diff --git a/.buildkite/test_areas/jit_monitor.yaml b/.buildkite/test_areas/jit_monitor.yaml index 7119c489c016..941c19013882 100644 --- a/.buildkite/test_areas/jit_monitor.yaml +++ b/.buildkite/test_areas/jit_monitor.yaml @@ -2,7 +2,7 @@ group: JIT Monitor depends_on: - image-build steps: -- label: No Runtime JITs e2e tests +- label: ":nvidia: (H200) No Runtime JITs E2E" key: jit-monitor-no-runtime-jit device: h200_35gb timeout_in_minutes: 45 diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 4dec86e929d9..0887a6fb27a4 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -2,7 +2,7 @@ group: Kernels depends_on: - image-build steps: -- label: vLLM IR Tests +- label: ":nvidia: (H200) vLLM IR" key: vllm-ir-tests timeout_in_minutes: 35 device: h200_18gb @@ -14,7 +14,7 @@ steps: - pytest -v -s tests/ir - pytest -v -s tests/kernels/ir -- label: Kernels Core Operation Test +- label: ":nvidia: (H200) Core Operation Kernels Shard %N" device: h200_35gb key: kernels-core-operation-test timeout_in_minutes: 120 @@ -27,7 +27,7 @@ steps: - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_fused_qk_norm_rope_gate.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT parallelism: 3 -- label: Kernels MiniMax Reduce RMS Test (2 GPUs) +- label: ":nvidia: (H100) MiniMax Reduce RMS Kernels" key: kernels-minimax-reduce-rms-test-2-gpus timeout_in_minutes: 20 num_devices: 2 @@ -41,7 +41,7 @@ steps: commands: - pytest -v -s kernels/core/test_minimax_reduce_rms.py -- label: Deepseek V4 Kernel Test (H100) +- label: ":nvidia: (H100) DeepSeek V4 Kernels" key: deepseek-v4-kernel-test-h100 timeout_in_minutes: 30 device: h100 @@ -54,7 +54,7 @@ steps: - pytest -v -s kernels/test_fused_deepseek_v4_*.py - pytest -v -s kernels/test_top_k_per_row.py -- label: Deepseek V4 Kernel Test (B200) +- label: ":nvidia: (B200) DeepSeek V4 Kernels" key: deepseek-v4-kernel-test-b200 timeout_in_minutes: 20 device: b200-k8s @@ -73,7 +73,7 @@ steps: # Files with dedicated jobs elsewhere in this file are excluded via --ignore # (test_kda, test_bf16x3_router_gemm_cutedsl and test_ll_bf16_gemm run in # their own jobs / Kernels (B200)). -- label: Kernels Root Misc Test (B200) +- label: ":nvidia: (B200) Miscellaneous Kernels" key: kernels-root-misc-test-b200 timeout_in_minutes: 45 device: b200-k8s @@ -101,7 +101,7 @@ steps: # BROKEN on main, pending kernel fixes (B200): # test_shuffle_rows.py (1: test_shuffle_rows_edge_cases) -- label: Kernels Attention Test %N +- label: ":nvidia: (L4) Attention Kernels Shard %N" key: kernels-attention-test timeout_in_minutes: 65 source_file_dependencies: @@ -116,6 +116,7 @@ steps: parallelism: 2 mirror: amd: + label: ":amd: (MI300) Attention Kernels Shard %N" dind: false device: mi300_1 timeout_in_minutes: 90 @@ -130,7 +131,7 @@ steps: - vllm/envs.py - vllm/platforms/rocm.py -- label: Kernels Attention DiffKV Test (H100) +- label: ":nvidia: (H100) Attention DiffKV Kernels" key: kernels-attention-diffkv-test-h100 timeout_in_minutes: 20 device: h100 @@ -143,7 +144,7 @@ steps: commands: - pytest -v -s kernels/attention/test_triton_unified_attention_diffkv.py -- label: Kernels FlashMLA Test (H100) +- label: ":nvidia: (H100) FlashMLA Kernels" key: kernels-flashmla-test-h100 timeout_in_minutes: 25 device: h100 @@ -161,7 +162,7 @@ steps: - pytest -v -s kernels/attention/test_flashmla_sparse.py - pytest -v -s kernels/attention/test_mla_cross_layer_kernel_equivalence.py -- label: Kernels Quantization Test %N +- label: ":nvidia: (L4) Quantization Kernels Shard %N" key: kernels-quantization-test timeout_in_minutes: 60 source_file_dependencies: @@ -173,6 +174,7 @@ steps: parallelism: 2 mirror: amd: + label: ":amd: (MI300) Quantization Kernels Shard %N" dind: false device: mi300_1 timeout_in_minutes: 120 @@ -189,7 +191,7 @@ steps: depends_on: - image-build-amd -- label: Kernels MoE Test %N +- label: ":nvidia: (L4) MoE Kernels Shard %N" key: kernels-moe-test timeout_in_minutes: 50 source_file_dependencies: @@ -206,6 +208,7 @@ steps: parallelism: 5 mirror: amd: + label: ":amd: (MI300) MoE Kernels Shard %N" dind: false device: mi300_1 timeout_in_minutes: 55 @@ -222,7 +225,7 @@ steps: depends_on: - image-build-amd -- label: Kernels Mamba Test +- label: ":nvidia: (H200) Mamba Kernels" device: h200_35gb key: kernels-mamba-test timeout_in_minutes: 60 @@ -233,7 +236,7 @@ steps: commands: - pytest -v -s kernels/mamba -- label: Kernels DeepGEMM Test (H100) +- label: ":nvidia: (H100) DeepGEMM Kernels" key: kernels-deepgemm-test-h100 timeout_in_minutes: 35 device: h100 @@ -260,7 +263,7 @@ steps: - pytest -v -s kernels/attention/test_deepgemm_attention.py - pytest -v -s quantization/test_cutlass_w4a16.py -- label: Kernels (B200) +- label: ":nvidia: (B200) Kernels" key: kernels-b200 timeout_in_minutes: 80 working_dir: "/vllm-workspace/" @@ -340,7 +343,7 @@ steps: # e2e - pytest -v -s tests/models/quantization/test_nvfp4.py -- label: B12X Linear Kernels (DGX Spark) Nightly +- label: ":nvidia: (DGX) Spark B12X Linear Kernels Nightly" key: b12x-linear-kernels-dgx-spark-nightly timeout_in_minutes: 30 device: dgx-spark @@ -362,7 +365,7 @@ steps: model_executor/test_b12x_warmup.py kernels/quantization/test_block_fp8.py -k b12x -- label: Kernels Helion Test +- label: ":nvidia: (H100) Helion Kernels Shard %N" key: kernels-helion-test timeout_in_minutes: 115 device: h100 @@ -375,7 +378,7 @@ steps: parallelism: 2 -- label: Kernels FP8 MoE Test (1xH100) +- label: ":nvidia: (H100) FP8 MoE Kernels" key: kernels-fp8-moe-test-1xh100 timeout_in_minutes: 40 device: h100 @@ -392,7 +395,7 @@ steps: - pytest -v -s kernels/moe/test_triton_moe_no_act_mul.py - pytest -v -s kernels/moe/test_triton_moe_ptpc_fp8.py -- label: Kernels FP8 MoE Test (2xH100) +- label: ":nvidia: (H100) DeepEP FP8 MoE Kernels" key: kernels-fp8-moe-test-2xh100 timeout_in_minutes: 45 device: h100 @@ -402,7 +405,7 @@ steps: - pytest -v -s kernels/moe/test_deepep_deepgemm_moe.py - pytest -v -s kernels/moe/test_deepep_moe.py -- label: Kernels Fp4 MoE Test (B200) +- label: ":nvidia: (B200) FP4 MoE Kernels" key: kernels-fp4-moe-test-b200 timeout_in_minutes: 25 device: b200-k8s @@ -415,7 +418,7 @@ steps: - pytest -v -s kernels/moe/test_ocp_mx_moe.py -- label: Kernels FusedMoE Layer Test (2 H100s) +- label: ":nvidia: (H100) FusedMoE Layer Kernels" key: kernels-fusedmoe-layer-test-2-h100s timeout_in_minutes: 30 device: h100 @@ -432,7 +435,7 @@ steps: - pytest -v -s kernels/moe/test_moe_layer.py -- label: Kernels FusedMoE Layer Test (2 B200s) +- label: ":nvidia: (B200) FusedMoE Layer Kernels" key: kernels-fusedmoe-layer-test-2-b200s timeout_in_minutes: 90 device: b200-k8s diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index f9c21a54217b..58e6717acadc 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -2,7 +2,7 @@ group: LM Eval depends_on: - image-build steps: -- label: LM Eval Small Models +- label: ":nvidia: (H200) LM Eval Small Models" device: h200_35gb key: lm-eval-small-models timeout_in_minutes: 45 @@ -14,6 +14,7 @@ steps: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt mirror: amd: + label: ":amd: (MI300) LM Eval Small Models" dind: false device: mi300_1 timeout_in_minutes: 45 @@ -42,7 +43,7 @@ steps: # - export VLLM_WORKER_MULTIPROC_METHOD=spawn # - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4 -- label: LM Eval Large Models (4xH100) +- label: ":nvidia: (H100) LM Eval Large Models" key: lm-eval-large-models-4xh100 device: h100 optional: true @@ -55,7 +56,7 @@ steps: - export VLLM_USE_DEEP_GEMM=0 # We found Triton is faster than DeepGEMM for H100 - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-hopper.txt --tp-size=4 -- label: LM Eval Small Models (1xB200) +- label: ":nvidia: (B200) LM Eval Small Models" key: lm-eval-small-models-1xb200 timeout_in_minutes: 50 device: b200-k8s @@ -66,7 +67,7 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell.txt -- label: LM Eval Small Models Distributed (2xB200) +- label: ":nvidia: (B200) LM Eval Small Models Distributed" key: lm-eval-small-models-distributed-2xb200 timeout_in_minutes: 120 device: b200-k8s @@ -79,7 +80,7 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small-tp.txt -- label: LM Eval PCP (4xB200) +- label: ":nvidia: (B200) LM Eval PCP" key: lm-eval-pcp-4xb200 timeout_in_minutes: 360 device: b200-k8s @@ -101,7 +102,7 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-pcp.txt -- label: LM Eval Spec Decode (4xB200) +- label: ":nvidia: (B200) LM Eval Spec Decode" key: lm-eval-spec-decode-4xb200 timeout_in_minutes: 120 device: b200-k8s @@ -121,7 +122,7 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-spec-decode.txt -- label: LM Eval Large Models EP (2xB200) +- label: ":nvidia: (B200) LM Eval Large Models EP" key: lm-eval-large-models-ep-2xb200 timeout_in_minutes: 60 device: b200-k8s @@ -133,7 +134,7 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell-ep.txt -- label: LM Eval Qwen3.5 Models (2xB200) +- label: ":nvidia: (B200) LM Eval Qwen3.5 Models" key: lm-eval-qwen3-5-models-2xb200 timeout_in_minutes: 45 device: b200-k8s @@ -150,7 +151,7 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-qwen35-blackwell.txt -- label: LM Eval Large Models (8xH200) +- label: ":nvidia: (H200) LM Eval Large Models" key: lm-eval-large-models-8xh200 timeout_in_minutes: 50 device: h200 @@ -160,6 +161,7 @@ steps: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200.txt mirror: amd: + label: ":amd: (MI300) LM Eval Large Models" dind: false device: mi300_8 timeout_in_minutes: 40 @@ -170,7 +172,7 @@ steps: - export PYTORCH_ROCM_ARCH=gfx942 # Limit Quark compilation to save time - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt -- label: MoE Refactor Integration Test (H100 - TEMPORARY) +- label: ":nvidia: (H100) MoE Refactor Integration TEMPORARY" key: moe-refactor-integration-test-h100-temporary device: h100 optional: true @@ -178,7 +180,7 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-h100.txt -- label: MoE Refactor Integration Test (B200 - TEMPORARY) %N +- label: ":nvidia: (B200) MoE Refactor Integration TEMPORARY Shard %N" key: moe-refactor-integration-test-b200-temporary device: b200-k8s optional: true @@ -187,7 +189,7 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-b200-shard-$$BUILDKITE_PARALLEL_JOB.txt -- label: MoE Refactor Integration Test (B200 DP - TEMPORARY) +- label: ":nvidia: (B200) MoE Refactor DP Integration TEMPORARY" key: moe-refactor-integration-test-b200-dp-temporary device: b200-k8s optional: true @@ -195,7 +197,7 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt -- label: LM Eval Humming f16 (A100 - TEMPORARY) %N +- label: ":nvidia: (A100) LM Eval Humming FP16 TEMPORARY Shard %N" key: lm-eval-humming-f16-a100 timeout_in_minutes: 75 device: a100 @@ -211,7 +213,7 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-a100-shard-$$BUILDKITE_PARALLEL_JOB.txt -- label: LM Eval Humming Act int8 (A100 - TEMPORARY) +- label: ":nvidia: (A100) LM Eval Humming Activation INT8 TEMPORARY" key: lm-eval-humming-act-a100 timeout_in_minutes: 45 device: a100 @@ -226,7 +228,7 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-int8.txt -- label: LM Eval Humming f16 (H100 - TEMPORARY) %N +- label: ":nvidia: (H100) LM Eval Humming FP16 TEMPORARY Shard %N" key: lm-eval-humming-f16-h100 timeout_in_minutes: 70 device: h100 @@ -242,7 +244,7 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-h100-shard-$$BUILDKITE_PARALLEL_JOB.txt -- label: LM Eval Humming Act fp8/int8 (H100 - TEMPORARY) +- label: ":nvidia: (H100) LM Eval Humming Activation FP8/INT8 TEMPORARY" key: lm-eval-humming-act-h100 timeout_in_minutes: 70 device: h100 @@ -258,7 +260,7 @@ steps: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-fp8.txt # - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-int8.txt -- label: LM Eval Humming f16 (B200 - TEMPORARY) +- label: ":nvidia: (B200) LM Eval Humming FP16 TEMPORARY" key: lm-eval-humming-f16-b200 timeout_in_minutes: 50 device: b200-k8s @@ -273,7 +275,7 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt -- label: LM Eval Humming Act fp8/int8 (B200 - TEMPORARY) +- label: ":nvidia: (B200) LM Eval Humming Activation FP8/INT8 TEMPORARY" key: lm-eval-humming-act-b200 timeout_in_minutes: 50 device: b200-k8s @@ -289,7 +291,7 @@ steps: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-fp8.txt - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-int8.txt -- label: LM Eval TurboQuant KV Cache +- label: ":nvidia: (H200) LM Eval TurboQuant KV Cache" key: lm-eval-turboquant-kv-cache timeout_in_minutes: 55 device: h200_18gb @@ -301,7 +303,7 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant.txt -- label: GPQA Eval (GPT-OSS) (2xH100) +- label: ":nvidia: (H100) GPQA Eval (GPT-OSS)" key: gpqa-eval-gpt-oss-2xh100 timeout_in_minutes: 35 device: h100 @@ -315,7 +317,7 @@ steps: - uv pip install --system 'gpt-oss[eval]==0.0.5' - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-h100.txt -- label: GPQA Eval (GPT-OSS) (2xB200) +- label: ":nvidia: (B200) GPQA Eval (GPT-OSS)" key: gpqa-eval-gpt-oss-2xb200 timeout_in_minutes: 30 device: b200-k8s @@ -329,7 +331,7 @@ steps: - uv pip install --system 'gpt-oss[eval]==0.0.5' - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-b200.txt -- label: GPQA Eval (GPT-OSS) (DGX Spark) +- label: ":nvidia: (DGX) Spark GPQA Eval (GPT-OSS)" key: gpqa-eval-gpt-oss-spark timeout_in_minutes: 35 device: dgx-spark @@ -345,7 +347,7 @@ steps: - uv pip install --system 'gpt-oss[eval]==0.0.5' - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-spark.txt -- label: LM Eval KV-Offload (1xH200) +- label: ":nvidia: (H200) LM Eval KV-Offload" key: kv-offload-small timeout_in_minutes: 30 device: h200_35gb @@ -358,7 +360,7 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "nemotron-h-8b or gemma-4-e4b-it" -- label: LM Eval KV-Offload (2xH100) +- label: ":nvidia: (H100) LM Eval KV-Offload Medium" key: kv-offload-medium timeout_in_minutes: 45 device: h100 @@ -372,7 +374,7 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "qwen3.5-35b or deepseek-v2-lite" -- label: LM Eval KV-Offload (4xH100) +- label: ":nvidia: (H100) LM Eval KV-Offload Large" key: kv-offload-large timeout_in_minutes: 40 device: h100 @@ -386,7 +388,7 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "deepseek-v4-flash" -- label: MRCR Eval Small Models +- label: ":nvidia: (H200) MRCR Eval Small Models" device: h200_35gb timeout_in_minutes: 25 source_file_dependencies: diff --git a/.buildkite/test_areas/lora.yaml b/.buildkite/test_areas/lora.yaml index 4287816b90c2..5c3c7c0e8900 100644 --- a/.buildkite/test_areas/lora.yaml +++ b/.buildkite/test_areas/lora.yaml @@ -2,7 +2,7 @@ group: LoRA depends_on: - image-build steps: -- label: LoRA %N +- label: ":nvidia: (H200) LoRA Shard %N" device: h200_35gb key: lora timeout_in_minutes: 40 @@ -14,6 +14,7 @@ steps: parallelism: 4 mirror: amd: + label: ":amd: (MI300) LoRA Shard %N" dind: false device: mi300_1 working_dir: "/vllm-workspace/tests" @@ -26,7 +27,7 @@ steps: - image-build-amd -- label: LoRA TP (Distributed) +- label: ":nvidia: (L4) LoRA TP (Distributed)" key: lora-tp-distributed timeout_in_minutes: 60 num_devices: 4 diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 8f168cbd367f..750f51283abe 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -2,7 +2,7 @@ group: Miscellaneous depends_on: - image-build steps: -- label: V1 Sample + Logits +- label: ":nvidia: (H200) V1 Sample + Logits" key: v1-sample-logits timeout_in_minutes: 83 device: h200_18gb @@ -33,13 +33,14 @@ steps: - pytest -v -s v1/test_outputs.py mirror: amd: + label: ":amd: (MI300) V1 Sample + Logits" dind: false device: mi300_1 timeout_in_minutes: 70 depends_on: - image-build-amd -- label: V1 Core + KV + Metrics +- label: ":nvidia: (H200) V1 Core + KV + Metrics" device: h200_35gb key: v1-core-kv-metrics timeout_in_minutes: 80 @@ -90,13 +91,14 @@ steps: - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine mirror: amd: + label: ":amd: (MI300) V1 Core + KV + Metrics" dind: false device: mi300_1 timeout_in_minutes: 65 depends_on: - image-build-amd -- label: V1 Others (CPU) +- label: ":computer: (CPU) V1 Others" key: v1-others-cpu depends_on: - image-build-cpu @@ -130,7 +132,7 @@ steps: - pytest -v -s -m 'cpu_test' v1/ec_connector/unit - pytest -v -s -m 'cpu_test' v1/metrics -- label: Extract Hidden States Integration +- label: ":nvidia: (H200) Extract Hidden States Integration" key: extract-hidden-states-integration timeout_in_minutes: 20 device: h200_18gb @@ -144,7 +146,7 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s v1/kv_connector/extract_hidden_states_integration -- label: Extract Hidden States Integration (2 GPUs) +- label: ":nvidia: (L4) Extract Hidden States Integration" key: extract-hidden-states-integration-2-gpus timeout_in_minutes: 20 num_devices: 2 @@ -158,7 +160,7 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s -m 'distributed' v1/kv_connector/extract_hidden_states_integration -- label: Regression +- label: ":nvidia: (H200) Regression" key: regression timeout_in_minutes: 30 device: h200_18gb @@ -181,7 +183,7 @@ steps: - pytest -v -s test_regression.py working_dir: "/vllm-workspace/tests" # optional -- label: Examples +- label: ":nvidia: (H200) Examples" device: h200_35gb key: examples timeout_in_minutes: 40 @@ -215,6 +217,7 @@ steps: - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536 mirror: amd: + label: ":amd: (MI300) Examples" dind: false device: mi300_1 timeout_in_minutes: 75 @@ -226,7 +229,7 @@ steps: depends_on: - image-build-amd -- label: Metrics, Tracing (2 GPUs) +- label: ":nvidia: (L4) Metrics, Tracing" key: metrics-tracing-2-gpus timeout_in_minutes: 25 num_devices: 2 @@ -256,13 +259,14 @@ steps: - pytest -v -s tracing mirror: amd: + label: ":amd: (MI300) Metrics, Tracing" dind: false device: mi300_2 timeout_in_minutes: 30 depends_on: - image-build-amd -- label: Python-only Installation +- label: ":nvidia: (H200) Python-only Installation" key: python-only-installation depends_on: ~ optional: true @@ -275,6 +279,7 @@ steps: - bash standalone_tests/python_only_compile.sh mirror: amd: + label: ":amd: (MI300) Python-only Installation" dind: false device: mi300_1 timeout_in_minutes: 55 @@ -286,7 +291,7 @@ steps: - setup.py - vllm/platforms/rocm.py -- label: Async Engine, Inputs, Utils, Worker +- label: ":nvidia: (H200) Async Engine, Inputs, Utils, Worker" device: h200_35gb key: async-engine-inputs-utils-worker timeout_in_minutes: 25 @@ -313,7 +318,7 @@ steps: - pytest -v -s -m 'not cpu_test' multimodal - pytest -v -s utils_ -- label: Async Engine, Inputs, Utils, Worker, Config (CPU) +- label: ":computer: (CPU) Async Engine, Inputs, Utils, Worker, Config" key: async-engine-inputs-utils-worker-config-cpu depends_on: - image-build-cpu @@ -377,7 +382,7 @@ steps: - pytest -v -s transformers_utils - pytest -v -s config -- label: Batch Invariance (A100) +- label: ":nvidia: (A100) Batch Invariance" key: batch-invariance-a100 timeout_in_minutes: 60 device: a100 @@ -391,7 +396,7 @@ steps: - pytest -v -s v1/determinism/test_batch_invariance.py - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA -- label: Batch Invariance (H100) +- label: ":nvidia: (H100) Batch Invariance" key: batch-invariance-h100 timeout_in_minutes: 60 device: h100 @@ -407,7 +412,7 @@ steps: - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN -- label: Batch Invariance (B200) +- label: ":nvidia: (B200) Batch Invariance" key: batch-invariance-b200 timeout_in_minutes: 45 device: b200-k8s @@ -427,7 +432,7 @@ steps: - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py - pytest -v -s v1/determinism/test_online_batch_invariance.py -- label: Acceptance Length Test (Large Models) # optional +- label: ":nvidia: (H200) Acceptance Length (Large Models)" device: h200_35gb key: acceptance-length-test-large-models timeout_in_minutes: 20 diff --git a/.buildkite/test_areas/model_executor.yaml b/.buildkite/test_areas/model_executor.yaml index 07d4aea871c0..8228bbc83a58 100644 --- a/.buildkite/test_areas/model_executor.yaml +++ b/.buildkite/test_areas/model_executor.yaml @@ -2,7 +2,7 @@ group: Model Executor depends_on: - image-build steps: -- label: Model Executor +- label: ":nvidia: (H200) Model Executor" device: h200_35gb key: model-executor timeout_in_minutes: 60 @@ -28,6 +28,7 @@ steps: - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py --timeout=900 --timeout-method=thread mirror: amd: + label: ":amd: (MI300) Model Executor" dind: false device: mi300_1 timeout_in_minutes: 60 diff --git a/.buildkite/test_areas/model_runner_v2.yaml b/.buildkite/test_areas/model_runner_v2.yaml index 9b05f4b1e251..1b1bef43aa79 100644 --- a/.buildkite/test_areas/model_runner_v2.yaml +++ b/.buildkite/test_areas/model_runner_v2.yaml @@ -2,7 +2,7 @@ group: Model Runner V2 depends_on: - image-build steps: -- label: Model Runner V2 Core Tests +- label: ":nvidia: (H200) Model Runner V2 Core" device: h200_35gb key: model-runner-v2-core-tests timeout_in_minutes: 35 @@ -24,7 +24,7 @@ steps: # Temporary hack filter to exclude ngram spec decoding based tests. - pytest -v -s entrypoints/llm/test_struct_output_generate.py -k "xgrammar and not speculative_config6 and not speculative_config7 and not speculative_config8 and not speculative_config0" -- label: Model Runner V2 Examples +- label: ":nvidia: (H200) Model Runner V2 Examples" device: h200_35gb key: model-runner-v2-examples timeout_in_minutes: 35 @@ -61,7 +61,7 @@ steps: # https://github.com/vllm-project/vllm/pull/26682 uses slightly more memory in PyTorch 2.9+ causing this test to OOM in 1xL4 GPU - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536 -- label: Model Runner V2 Distributed (2 GPUs) +- label: ":nvidia: (L4) Model Runner V2 Distributed" key: model-runner-v2-distributed-2-gpus timeout_in_minutes: 30 working_dir: "/vllm-workspace/tests" @@ -82,7 +82,7 @@ steps: - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py -k "not ray" - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py -- label: Model Runner V2 Pipeline Parallelism (4 GPUs) +- label: ":nvidia: (L4) Model Runner V2 Pipeline Parallelism" key: model-runner-v2-pipeline-parallelism-4-gpus timeout_in_minutes: 50 working_dir: "/vllm-workspace/tests" @@ -100,7 +100,7 @@ steps: - pytest -v -s distributed/test_pp_cudagraph.py -k "not ray" - pytest -v -s v1/distributed/test_pp_dp_v2.py -- label: Model Runner V2 Spec Decode +- label: ":nvidia: (H200) Model Runner V2 Spec Decode" device: h200_35gb key: model-runner-v2-spec-decode timeout_in_minutes: 50 diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index b3d922a5e9f8..bdb3cab5a60f 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -2,7 +2,7 @@ group: Models - Basic depends_on: - image-build steps: -- label: Basic Models Tests (Initialization) +- label: ":nvidia: (H200) Basic Models (Initialization)" key: basic-models-tests-initialization timeout_in_minutes: 25 device: h200_18gb @@ -15,7 +15,7 @@ steps: # Run a subset of model initialization tests - pytest -v -s models/test_initialization.py::test_can_initialize_small_subset -- label: Basic Models Tests (Extra Initialization) %N +- label: ":nvidia: (H200) Basic Models (Extra Initialization) Shard %N" device: h200_35gb key: basic-models-tests-extra-initialization timeout_in_minutes: 100 @@ -30,7 +30,7 @@ steps: - pytest -v -s models/test_initialization.py -k 'not test_can_initialize_small_subset' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB parallelism: 4 -- label: Basic Models Tests (Other) +- label: ":nvidia: (H200) Basic Models (Other)" device: h200_35gb key: basic-models-tests-other timeout_in_minutes: 35 @@ -44,13 +44,14 @@ steps: - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py mirror: amd: + label: ":amd: (MI300) Basic Models (Other)" dind: false device: mi300_1 timeout_in_minutes: 50 depends_on: - image-build-amd -- label: Inkling Unit Tests (B200) +- label: ":nvidia: (B200) Inkling" key: inkling-unit-tests-b200 timeout_in_minutes: 40 device: b200-k8s @@ -63,7 +64,7 @@ steps: # FA4 kernel tests require SM100; the suite skips them elsewhere. - pytest -v -s models/inkling -- label: Kimi K3 Unit Tests (B200) +- label: ":nvidia: (B200) Kimi K3" key: kimi-k3-unit-tests-b200 timeout_in_minutes: 40 device: b200-k8s @@ -77,7 +78,7 @@ steps: # The native NVIDIA Kimi K3 kernels require the SM100 family. - pytest -v -s models/kimi_k3 kernels/attention/test_kimi_k3_mla_fused_epilogue.py kernels/test_bf16_skinny_gemm.py -- label: Basic Models Test (Other CPU) # 5min +- label: ":computer: (CPU) Basic Models Other" key: basic-models-test-other-cpu depends_on: - image-build-cpu diff --git a/.buildkite/test_areas/models_distributed.yaml b/.buildkite/test_areas/models_distributed.yaml index c1ec5eb00ae5..730c51b89b84 100644 --- a/.buildkite/test_areas/models_distributed.yaml +++ b/.buildkite/test_areas/models_distributed.yaml @@ -2,7 +2,7 @@ group: Models - Distributed depends_on: - image-build steps: -- label: Distributed Model Tests (2 GPUs) +- label: ":nvidia: (L4) Distributed Models" key: distributed-model-tests-2-gpus timeout_in_minutes: 60 working_dir: "/vllm-workspace/tests" diff --git a/.buildkite/test_areas/models_language.yaml b/.buildkite/test_areas/models_language.yaml index d0eff968f7c8..e8675141237e 100644 --- a/.buildkite/test_areas/models_language.yaml +++ b/.buildkite/test_areas/models_language.yaml @@ -2,7 +2,7 @@ group: Models - Language depends_on: - image-build steps: -- label: Language Models Tests (Standard) +- label: ":nvidia: (H200) Language Models (Standard)" key: language-models-tests-standard timeout_in_minutes: 30 device: h200_18gb @@ -16,13 +16,14 @@ steps: - pytest -v -s models/language -m 'core_model and (not slow_test)' mirror: amd: + label: ":amd: (MI300) Language Models (Standard)" dind: false device: mi300_1 timeout_in_minutes: 45 depends_on: - image-build-amd -- label: Language Models Tests (Extra Standard) %N +- label: ":nvidia: (H200) Language Models (Extra Standard) Shard %N" device: h200_35gb key: language-models-tests-extra-standard timeout_in_minutes: 40 @@ -39,6 +40,7 @@ steps: parallelism: 2 mirror: amd: + label: ":amd: (MI300) Language Models (Extra Standard) Shard %N" dind: false device: mi300_1 timeout_in_minutes: 40 @@ -55,7 +57,7 @@ steps: - tests/models/language/pooling/test_classification.py - vllm/_aiter_ops.py - vllm/platforms/rocm.py -- label: Language Models Tests (Hybrid) %N +- label: ":nvidia: (H200) Language Models (Hybrid) Shard %N" device: h200_35gb key: language-models-tests-hybrid timeout_in_minutes: 65 @@ -72,6 +74,7 @@ steps: parallelism: 2 mirror: amd: + label: ":amd: (MI300) Language Models (Hybrid) Shard %N" dind: false device: mi300_1 timeout_in_minutes: 60 @@ -85,7 +88,7 @@ steps: # Granite 4 hybrid generation is sensitive to hardware-specific Triton SSD # autotuning (https://github.com/vllm-project/vllm/issues/25194). Keep this one # correctness test on L4 until its H200 output matches the Transformers reference. -- label: Language Models Tests (Granite L4 Compatibility) +- label: ":nvidia: (L4) Granite Language Model Compatibility" key: language-models-tests-granite-l4-compatibility timeout_in_minutes: 65 source_file_dependencies: @@ -97,7 +100,7 @@ steps: - CAUSAL_CONV1D_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' - pytest -v -s models/language/generation -m hybrid_model -k 'granite-4.0-tiny-preview' -- label: Language Models Test (Extended Generation) # 80min +- label: ":nvidia: (H200) Language Models (Extended Generation)" device: h200_35gb key: language-models-test-extended-generation timeout_in_minutes: 80 @@ -112,7 +115,7 @@ steps: - CAUSAL_CONV1D_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)' -- label: Language Models Test (PPL) +- label: ":nvidia: (H200) Language Models (PPL)" key: language-models-test-ppl timeout_in_minutes: 30 device: h200_18gb @@ -124,7 +127,7 @@ steps: commands: - pytest -v -s models/language/generation_ppl_test -- label: Language Models Test (Extended Pooling) %N +- label: ":nvidia: (H200) Language Models (Extended Pooling) Shard %N" device: h200_35gb key: language-models-test-extended-pooling timeout_in_minutes: 120 @@ -138,6 +141,7 @@ steps: - pytest -v -s models/language/pooling -m 'not core_model' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB mirror: amd: + label: ":amd: (MI300) Language Models (Extended Pooling) Shard %N" dind: false device: mi300_1 timeout_in_minutes: 95 @@ -147,7 +151,7 @@ steps: commands: - pytest -v -s models/language/pooling -m 'not core_model' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB -- label: Language Models Test (MTEB) +- label: ":nvidia: (H200) Language Models (MTEB)" key: language-models-test-mteb timeout_in_minutes: 68 device: h200_18gb diff --git a/.buildkite/test_areas/models_multimodal.yaml b/.buildkite/test_areas/models_multimodal.yaml index 2ea3621342fc..1d97c3f5e622 100644 --- a/.buildkite/test_areas/models_multimodal.yaml +++ b/.buildkite/test_areas/models_multimodal.yaml @@ -2,7 +2,7 @@ group: Models - Multimodal depends_on: - image-build steps: -- label: "Multi-Modal Models (Standard) 1: qwen2" +- label: ":nvidia: (H200) Multimodal Models (Standard) 1: qwen2" key: multi-modal-models-standard-1-qwen2 timeout_in_minutes: 68 device: h200_18gb @@ -15,13 +15,14 @@ steps: - pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model mirror: amd: + label: ":amd: (MI300) Multimodal Models (Standard) 1: qwen2" dind: false device: mi300_1 timeout_in_minutes: 65 depends_on: - image-build-amd -- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma" +- label: ":nvidia: (H200) Multimodal Models (Standard) 2: qwen3 + gemma" key: multi-modal-models-standard-2-qwen3-gemma timeout_in_minutes: 75 device: h200_18gb @@ -35,13 +36,14 @@ steps: - pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model mirror: amd: + label: ":amd: (MI300) Multimodal Models (Standard) 2: qwen3 + gemma" dind: false device: mi300_1 timeout_in_minutes: 55 depends_on: - image-build-amd -- label: "Multi-Modal Models (Standard) 3: llava + qwen2_vl" +- label: ":nvidia: (H200) Multimodal Models (Standard) 3: llava + qwen2_vl" device: h200_35gb key: multi-modal-models-standard-3-llava-qwen2-vl timeout_in_minutes: 40 @@ -54,12 +56,13 @@ steps: - pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model mirror: amd: + label: ":amd: (MI250) Multimodal Models (Standard) 3: llava + qwen2_vl" device: mi250_1 timeout_in_minutes: 55 depends_on: - image-build-amd -- label: "Multi-Modal Models (Standard) 4: other + whisper" +- label: ":nvidia: (H200) Multimodal Models (Standard) 4: other + whisper" device: h200_35gb key: multi-modal-models-standard-4-other-whisper timeout_in_minutes: 75 @@ -74,13 +77,14 @@ steps: - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model # Otherwise, mp_method="spawn" doesn't work mirror: amd: + label: ":amd: (MI300) Multimodal Models (Standard) 4: other + whisper" dind: false device: mi300_1 timeout_in_minutes: 50 depends_on: - image-build-amd -- label: Multi-Modal Processor (CPU) %N +- label: ":computer: (CPU) Multimodal Processor Shard %N" key: multi-modal-processor-cpu depends_on: - image-build-cpu @@ -95,7 +99,7 @@ steps: - pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB parallelism: 4 -- label: Multi-Modal Processor # 44min +- label: ":nvidia: (H200) Multimodal Processor" key: multi-modal-processor timeout_in_minutes: 98 device: h200_18gb @@ -107,7 +111,7 @@ steps: commands: - pytest -v -s models/multimodal/processing/test_tensor_schema.py -- label: Multi-Modal Accuracy Eval (Small Models) # 50min +- label: ":nvidia: (H200) Multimodal Accuracy Eval (Small Models)" device: h200_35gb key: multi-modal-accuracy-eval-small-models timeout_in_minutes: 30 @@ -120,6 +124,7 @@ steps: - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-mm-small.txt --tp-size=1 mirror: amd: + label: ":amd: (MI300) Multimodal Accuracy Eval (Small Models)" dind: false device: mi300_1 timeout_in_minutes: 35 @@ -132,7 +137,7 @@ steps: - vllm/platforms/rocm.py - vllm/model_executor/model_loader/ -- label: Multi-Modal Models (Extended Generation 1) +- label: ":nvidia: (H200) Multimodal Models (Extended Generation 1)" device: h200_35gb key: multi-modal-models-extended-generation-1 optional: true @@ -146,13 +151,14 @@ steps: - pytest -v -s models/multimodal/test_mapping.py mirror: amd: + label: ":amd: (MI300) Multimodal Models (Extended Generation 1)" dind: false device: mi300_1 timeout_in_minutes: 90 depends_on: - image-build-amd -- label: Multi-Modal Models (PPL) +- label: ":nvidia: (H200) Multimodal Models (PPL)" device: h200_35gb key: multi-modal-models-extended-ppl optional: true @@ -163,13 +169,14 @@ steps: - pytest -v -s models/multimodal/generation_ppl_test/ mirror: amd: + label: ":amd: (MI300) Multimodal Models (PPL)" dind: false device: mi300_1 timeout_in_minutes: 90 depends_on: - image-build-amd -- label: Multi-Modal Models (Extended Generation 2) %N +- label: ":nvidia: (H200) Multimodal Models (Extended Generation 2) Shard %N" device: h200_35gb key: multi-modal-models-extended-generation-2 parallelism: 4 @@ -181,7 +188,7 @@ steps: commands: - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB -- label: Multi-Modal Models (Extended Generation 3) +- label: ":nvidia: (H200) Multimodal Models (Extended Generation 3)" device: h200_35gb key: multi-modal-models-extended-generation-3 optional: true @@ -192,7 +199,7 @@ steps: commands: - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model' -- label: Multi-Modal Models (Extended Pooling) +- label: ":nvidia: (H200) Multimodal Models (Extended Pooling)" key: multi-modal-models-extended-pooling optional: true device: h200_18gb @@ -204,6 +211,7 @@ steps: - pytest -v -s models/multimodal/pooling -m 'not core_model' mirror: amd: + label: ":amd: (MI300) Multimodal Models (Extended Pooling)" dind: false device: mi300_1 timeout_in_minutes: 60 diff --git a/.buildkite/test_areas/plugins.yaml b/.buildkite/test_areas/plugins.yaml index 4ba339207196..2cc4cb8e8bc3 100644 --- a/.buildkite/test_areas/plugins.yaml +++ b/.buildkite/test_areas/plugins.yaml @@ -2,7 +2,7 @@ group: Plugins depends_on: - image-build steps: -- label: Plugin Tests (2 GPUs) +- label: ":nvidia: (L4) Plugin Integration" key: plugin-tests-2-gpus timeout_in_minutes: 35 working_dir: "/vllm-workspace/tests" @@ -51,7 +51,7 @@ steps: - pytest -v -s plugins_tests/lora_resolvers # unit tests for in-tree lora resolver plugins -- label: GGUF Plugin +- label: ":nvidia: (H200) GGUF Plugin" key: gguf-plugin device: h200_18gb timeout_in_minutes: 30 @@ -64,7 +64,7 @@ steps: - pip install "vllm-gguf-plugin >= 0.0.2" - pytest -v -s plugins_tests/gguf -- label: BitsAndBytes Plugin +- label: ":nvidia: (H200) BitsAndBytes Plugin" key: bitsandbytes-plugin device: h200_18gb timeout_in_minutes: 30 @@ -81,7 +81,7 @@ steps: - pip install "vllm-bnb-plugin >= 0.0.1" - pytest -v -s plugins_tests/bitsandbytes -m 'not distributed' -- label: BitsAndBytes Plugin (2 GPUs) +- label: ":nvidia: (L4) BitsAndBytes Plugin" key: bitsandbytes-plugin-2-gpus timeout_in_minutes: 15 working_dir: "/vllm-workspace/tests" diff --git a/.buildkite/test_areas/pytorch.yaml b/.buildkite/test_areas/pytorch.yaml index b119e981f0e9..78cebb2e596d 100644 --- a/.buildkite/test_areas/pytorch.yaml +++ b/.buildkite/test_areas/pytorch.yaml @@ -2,7 +2,7 @@ group: PyTorch depends_on: - image-build steps: -- label: PyTorch Compilation Unit Tests +- label: ":nvidia: (H200) PyTorch Compilation" device: h200_35gb key: pytorch-compilation-unit-tests timeout_in_minutes: 60 @@ -43,7 +43,7 @@ steps: # (using -0 for proper path handling) - "find compile/ -maxdepth 1 -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'" -- label: PyTorch Compilation Unit Tests (H100) +- label: ":nvidia: (H200) PyTorch Compilation H100 Cases" key: pytorch-compilation-unit-tests-h100 timeout_in_minutes: 30 device: h200_18gb @@ -77,7 +77,7 @@ steps: commands: - "find compile/h100/ -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'" -- label: PyTorch Compilation Passes Unit Tests +- label: ":nvidia: (L4) PyTorch Compilation Passes" key: pytorch-compilation-passes-unit-tests timeout_in_minutes: 45 source_file_dependencies: @@ -110,7 +110,7 @@ steps: commands: - pytest -s -v compile/passes --ignore compile/passes/distributed -- label: PyTorch Fullgraph Test +- label: ":nvidia: (H200) PyTorch Fullgraph" device: h200_35gb key: pytorch-fullgraph-test timeout_in_minutes: 90 @@ -151,7 +151,7 @@ steps: # Hopper-only DeepSeek-V2-Lite cases in this file require two 29.3-GiB model # instances and cannot fit a 35GB MIG slice. L4 retains the original coverage: # those SM90 cases skip while the architecture-compatible cases still run. -- label: PyTorch Fullgraph CUDAGraph (L4 Compatibility) +- label: ":nvidia: (L4) PyTorch Fullgraph CUDAGraph Compatibility" key: pytorch-fullgraph-cudagraph-l4-compatibility timeout_in_minutes: 60 source_file_dependencies: @@ -184,7 +184,7 @@ steps: commands: - pytest -s -v compile/fullgraph/test_full_cudagraph.py -- label: Pytorch Nightly Dependency Override Check # 2min +- label: ":nvidia: (H200) PyTorch Nightly Dependency Override Check" key: pytorch-nightly-dependency-override-check # if this test fails, it means the nightly torch version is not compatible with some # of the dependencies. Please check the error message and add the package to whitelist @@ -197,6 +197,7 @@ steps: - bash standalone_tests/pytorch_nightly_dependency.sh mirror: amd: + label: ":amd: (MI300) PyTorch Nightly Dependency Override Check" dind: false device: mi300_1 timeout_in_minutes: 30 diff --git a/.buildkite/test_areas/quantization.yaml b/.buildkite/test_areas/quantization.yaml index b5c7c7a66dac..e659cd75f620 100644 --- a/.buildkite/test_areas/quantization.yaml +++ b/.buildkite/test_areas/quantization.yaml @@ -2,7 +2,7 @@ group: Quantization depends_on: - image-build steps: -- label: Quantization %N +- label: ":nvidia: (H200) Quantization Shard %N" device: h200_35gb key: quantization timeout_in_minutes: 40 @@ -21,7 +21,7 @@ steps: # parameter. It was not exercised by the previous L4 job. - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py -k 'not test_compressed_tensors_w4a8_fp8' --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT -- label: Quantized Fusions +- label: ":nvidia: (H200) Quantized Fusions" device: h200_35gb key: quantized-fusions timeout_in_minutes: 20 @@ -34,7 +34,7 @@ steps: commands: - pytest -v -s fusion/ -- label: Quantized MoE Test (B200) +- label: ":nvidia: (B200) Quantized MoE" key: quantized-moe-test-b200 timeout_in_minutes: 120 working_dir: "/vllm-workspace/" @@ -52,7 +52,7 @@ steps: commands: - pytest -s -v tests/quantization/test_blackwell_moe.py -- label: Quantized Models Test +- label: ":nvidia: (H200) Quantized Models" device: h200_35gb key: quantized-models-test timeout_in_minutes: 65 diff --git a/.buildkite/test_areas/ray_compat.yaml b/.buildkite/test_areas/ray_compat.yaml index 9207621a5830..243db5f21147 100644 --- a/.buildkite/test_areas/ray_compat.yaml +++ b/.buildkite/test_areas/ray_compat.yaml @@ -2,7 +2,7 @@ group: Ray Compatibility depends_on: - image-build steps: -- label: Ray Dependency Compatibility Check +- label: ":nvidia: (H200) Ray Dependency Compatibility Check" key: ray-dependency-compatibility-check # Informational only — does not block the pipeline. # If this fails, it means the PR introduces a dependency that diff --git a/.buildkite/test_areas/rust_frontend.yaml b/.buildkite/test_areas/rust_frontend.yaml index d56271670172..36b7d375d4ff 100644 --- a/.buildkite/test_areas/rust_frontend.yaml +++ b/.buildkite/test_areas/rust_frontend.yaml @@ -2,7 +2,7 @@ group: Rust Frontend E2E depends_on: - image-build steps: -- label: Rust Frontend OpenAI Coverage +- label: ":nvidia: (H200) Rust Frontend OpenAI Coverage" timeout_in_minutes: 30 device: h200_18gb working_dir: "/vllm-workspace/tests" @@ -37,7 +37,7 @@ steps: - pytest -v -s entrypoints/openai/test_uds.py - pytest -v -s v1/sample/test_logprobs_e2e.py -k "test_prompt_logprobs_e2e_server" -- label: Rust Frontend Serve/Admin Coverage +- label: ":nvidia: (H200) Rust Frontend Serve/Admin Coverage" timeout_in_minutes: 25 device: h200_18gb working_dir: "/vllm-workspace/tests" @@ -66,7 +66,7 @@ steps: # /tokenizer_info is not implemented in the Rust frontend (the CLI flag is accepted as a no-op). - pytest -v -s entrypoints/serve/tokenize/test_tokenization.py -k "not tokenizer_info" -- label: Rust Frontend Core Correctness +- label: ":nvidia: (H200) Rust Frontend Core Correctness" timeout_in_minutes: 20 device: h200_18gb working_dir: "/vllm-workspace/tests" @@ -80,7 +80,7 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine -- label: Rust Frontend Tool Use +- label: ":nvidia: (H200) Rust Frontend Tool Use" device: h200_35gb timeout_in_minutes: 25 working_dir: "/vllm-workspace/tests" @@ -95,7 +95,7 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s tool_use --ignore=tool_use/mistral --models llama3.2 -k "not test_response_format_with_tool_choice_required and not test_parallel_tool_calls_false and not test_tool_call_and_choice" -- label: Rust Frontend Distributed +- label: ":nvidia: (L4) Rust Frontend Distributed" timeout_in_minutes: 25 num_devices: 4 working_dir: "/vllm-workspace/tests" diff --git a/.buildkite/test_areas/rust_frontend_cargo.yaml b/.buildkite/test_areas/rust_frontend_cargo.yaml index 21d4c2ac2192..11bca6252667 100644 --- a/.buildkite/test_areas/rust_frontend_cargo.yaml +++ b/.buildkite/test_areas/rust_frontend_cargo.yaml @@ -1,7 +1,7 @@ group: Rust Frontend Cargo depends_on: [] steps: -- label: Rust Frontend Cargo Style + Clippy +- label: ":computer: (CPU) Rust Frontend Cargo Style + Clippy" key: rust-frontend-cargo-style-clippy depends_on: [] timeout_in_minutes: 20 @@ -15,7 +15,7 @@ steps: commands: - .buildkite/scripts/run-rust-frontend-cargo-ci.sh style-clippy -- label: Rust Frontend Cargo Tests +- label: ":computer: (CPU) Rust Frontend Cargo" key: rust-frontend-cargo-tests depends_on: [] timeout_in_minutes: 20 diff --git a/.buildkite/test_areas/samplers.yaml b/.buildkite/test_areas/samplers.yaml index 929cbec2aeb5..c0297d79e57b 100644 --- a/.buildkite/test_areas/samplers.yaml +++ b/.buildkite/test_areas/samplers.yaml @@ -2,7 +2,7 @@ group: Samplers depends_on: - image-build steps: -- label: Samplers Test +- label: ":nvidia: (H200) Samplers" device: h200_35gb key: samplers-test timeout_in_minutes: 40 @@ -19,6 +19,7 @@ steps: - VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers mirror: amd: + label: ":amd: (MI250) Samplers" device: mi250_1 timeout_in_minutes: 40 depends_on: diff --git a/.buildkite/test_areas/spec_decode.yaml b/.buildkite/test_areas/spec_decode.yaml index 852e5a463491..7486a467fbbe 100644 --- a/.buildkite/test_areas/spec_decode.yaml +++ b/.buildkite/test_areas/spec_decode.yaml @@ -2,7 +2,7 @@ group: Spec Decode depends_on: - image-build steps: -- label: V1 Spec Decode +- label: ":nvidia: (H200) V1 Spec Decode" device: h200_35gb key: v1-spec-decode timeout_in_minutes: 40 @@ -23,13 +23,14 @@ steps: - pytest -v -s -m 'not slow_test' v1/spec_decode mirror: amd: + label: ":amd: (MI300) V1 Spec Decode" dind: false device: mi300_1 timeout_in_minutes: 50 depends_on: - image-build-amd -- label: "Spec Decode Eagle 1: DeepSeek + Qwen" +- label: ":nvidia: (H200) Spec Decode Eagle 1: DeepSeek + Qwen" key: spec-decode-eagle-1-deepseek-qwen timeout_in_minutes: 25 device: h200_35gb @@ -41,6 +42,7 @@ steps: - pytest -v -s v1/e2e/spec_decode/eagle/ -k "deepseek_eagle or qwen3_eagle3" mirror: amd: + label: ":amd: (MI300) Spec Decode Eagle 1: DeepSeek + Qwen" dind: false device: mi300_1 timeout_in_minutes: 25 @@ -55,7 +57,7 @@ steps: - tests/v1/e2e/spec_decode/ - vllm/platforms/rocm.py -- label: "Spec Decode Eagle 2: Llama 3 + Qwen VL + Other" +- label: ":nvidia: (H200) Spec Decode Eagle 2: Llama 3 + Qwen VL + Other" key: spec-decode-eagle-2-llama3-qwen-vl-other timeout_in_minutes: 25 device: h200_35gb @@ -67,6 +69,7 @@ steps: - pytest -v -s v1/e2e/spec_decode/eagle/ -k "not deepseek_eagle and not qwen3_eagle3" mirror: amd: + label: ":amd: (MI300) Spec Decode Eagle 2: Llama 3 + Qwen VL + Other" dind: false device: mi300_1 timeout_in_minutes: 25 @@ -81,7 +84,7 @@ steps: - tests/v1/e2e/spec_decode/ - vllm/platforms/rocm.py -- label: Spec Decode Eagle Nightly (B200) +- label: ":nvidia: (B200) Spec Decode Eagle Nightly" key: spec-decode-eagle-nightly-b200 timeout_in_minutes: 25 device: b200-k8s @@ -93,7 +96,7 @@ steps: commands: - pytest -v -s v1/e2e/spec_decode/eagle/ -- label: Spec Decode Speculators + MTP +- label: ":nvidia: (H200) Spec Decode Speculators + MTP" key: spec-decode-speculators-mtp timeout_in_minutes: 50 device: h200_35gb @@ -108,6 +111,7 @@ steps: - pytest -v -s v1/e2e/spec_decode/mtp/ mirror: amd: + label: ":amd: (MI300) Spec Decode Speculators + MTP" dind: false device: mi300_1 timeout_in_minutes: 75 @@ -123,7 +127,7 @@ steps: - tests/v1/e2e/spec_decode/ - vllm/platforms/rocm.py -- label: Spec Decode Speculators + MTP Nightly (B200) +- label: ":nvidia: (B200) Spec Decode Speculators + MTP Nightly" key: spec-decode-speculators-mtp-nightly-b200 timeout_in_minutes: 30 device: b200-k8s @@ -137,7 +141,7 @@ steps: - pytest -v -s v1/e2e/spec_decode/speculators/ - pytest -v -s v1/e2e/spec_decode/mtp/ -- label: Spec Decode Ngram + Suffix +- label: ":nvidia: (H200) Spec Decode N-Gram + Suffix" key: spec-decode-ngram-suffix timeout_in_minutes: 20 device: h200_35gb @@ -151,6 +155,7 @@ steps: - python3 spec_decode/test_custom_proposer.py mirror: amd: + label: ":amd: (MI300) Spec Decode N-Gram + Suffix" dind: false device: mi300_1 timeout_in_minutes: 35 @@ -165,7 +170,7 @@ steps: - tests/v1/e2e/spec_decode/ - vllm/platforms/rocm.py -- label: Spec Decode Draft Model +- label: ":nvidia: (H200) Spec Decode Draft Model" key: spec-decode-draft-model timeout_in_minutes: 45 device: h200_18gb @@ -177,6 +182,7 @@ steps: - pytest -v -s v1/e2e/spec_decode/draft_model/ mirror: amd: + label: ":amd: (MI300) Spec Decode Draft Model" dind: false device: mi300_1 timeout_in_minutes: 55 @@ -191,7 +197,7 @@ steps: - tests/v1/e2e/spec_decode/ - vllm/platforms/rocm.py -- label: Spec Decode Draft Model Nightly (B200) +- label: ":nvidia: (B200) Spec Decode Draft Model Nightly" key: spec-decode-draft-model-nightly-b200 timeout_in_minutes: 40 device: b200-k8s @@ -203,7 +209,7 @@ steps: commands: - pytest -v -s v1/e2e/spec_decode/draft_model/ -- label: Speculators Correctness Nightly (H100) +- label: ":nvidia: (H100) Speculators Correctness Nightly" key: speculators-correctness timeout_in_minutes: 30 device: h100 @@ -218,7 +224,7 @@ steps: - export VLLM_ALLOW_INSECURE_SERIALIZATION=1 - pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test -- label: Spec Decode DeepSeek MTP Parallel Load (2xB200-2xMI300) +- label: ":nvidia: (B200) Spec Decode DeepSeek MTP Parallel Load" key: spec-decode-deepseek-mtp-parallel-load-2xb200-2xmi300 timeout_in_minutes: 30 device: b200-k8s @@ -235,6 +241,7 @@ steps: - pytest -v -s v1/e2e/spec_decode/test_mtp_parallel_load.py mirror: amd: + label: ":amd: (MI300) Spec Decode DeepSeek MTP Parallel Load" dind: false device: mi300_2 timeout_in_minutes: 45 @@ -254,7 +261,7 @@ steps: - tests/v1/e2e/spec_decode/test_mtp_parallel_load.py - vllm/platforms/rocm.py -- label: Spec Decode AL DFlash Nightly +- label: ":nvidia: (H200) Spec Decode AL DFlash Nightly" key: spec-decode-dflash-nightly timeout_in_minutes: 90 device: h200_35gb @@ -269,6 +276,7 @@ steps: - pytest -v -s v1/e2e/spec_decode/acceptance_rates/dflash/ mirror: amd: + label: ":amd: (MI300) Spec Decode AL DFlash Nightly" dind: false device: mi300_1 timeout_in_minutes: 120 @@ -289,7 +297,7 @@ steps: - tests/v1/e2e/spec_decode/ - vllm/platforms/rocm.py -- label: Spec Decode AL DSpark Nightly +- label: ":nvidia: (H200) Spec Decode AL DSpark Nightly" key: spec-decode-dspark-nightly timeout_in_minutes: 60 device: h200_35gb @@ -304,6 +312,7 @@ steps: - pytest -v -s v1/e2e/spec_decode/acceptance_rates/dspark/ mirror: amd: + label: ":amd: (MI300) Spec Decode AL DSpark Nightly" dind: false device: mi300_1 timeout_in_minutes: 90 @@ -328,7 +337,7 @@ steps: - tests/v1/e2e/spec_decode/ - vllm/platforms/rocm.py -- label: Spec Decode AL MTP + Other Acceptance Nightly +- label: ":nvidia: (H200) Spec Decode AL MTP + Other Acceptance Nightly" key: spec-decode-mtp-other-acceptance-nightly timeout_in_minutes: 60 device: h200_35gb diff --git a/.buildkite/test_areas/torch_abi.yaml b/.buildkite/test_areas/torch_abi.yaml index eaef3551664b..9e8375587adb 100644 --- a/.buildkite/test_areas/torch_abi.yaml +++ b/.buildkite/test_areas/torch_abi.yaml @@ -2,7 +2,7 @@ group: Torch ABI depends_on: - image-build steps: -- label: Torch Stable ABI Audit +- label: ":nvidia: (L4) Torch Stable ABI Audit" key: torch-stable-abi-audit timeout_in_minutes: 5 source_file_dependencies: diff --git a/.buildkite/test_areas/weight_loading.yaml b/.buildkite/test_areas/weight_loading.yaml index f587b6210e5f..0aa25dbad426 100644 --- a/.buildkite/test_areas/weight_loading.yaml +++ b/.buildkite/test_areas/weight_loading.yaml @@ -2,7 +2,7 @@ group: Weight Loading depends_on: - image-build steps: -- label: Weight Loading Multiple GPU # 33min +- label: ":nvidia: (L4) Weight Loading Multi-GPU" key: weight-loading-multiple-gpu timeout_in_minutes: 50 working_dir: "/vllm-workspace/tests" @@ -16,6 +16,7 @@ steps: - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models.txt mirror: amd: + label: ":amd: (MI300) Weight Loading Multi-GPU" dind: false device: mi300_2 timeout_in_minutes: 35 From 6a391a931a92250aa45aa711aaa3f099f5a8d05f Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Tue, 18 Aug 2026 13:04:28 -0700 Subject: [PATCH 104/839] [Rust Frontend][RL] add routed expert prompt offset (#52703) Signed-off-by: Biswa Panda Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- rust/src/engine-core-client/src/protocol/sampling.rs | 4 ++++ rust/src/engine-core-client/src/tests/client.rs | 2 ++ rust/src/engine-core-client/src/tests/python_compat.py | 2 ++ rust/src/text/src/lower.rs | 7 +++++++ 4 files changed, 15 insertions(+) diff --git a/rust/src/engine-core-client/src/protocol/sampling.rs b/rust/src/engine-core-client/src/protocol/sampling.rs index d5c6a852081f..19c579ac33a0 100644 --- a/rust/src/engine-core-client/src/protocol/sampling.rs +++ b/rust/src/engine-core-client/src/protocol/sampling.rs @@ -145,6 +145,9 @@ pub struct EngineCoreSamplingParams { pub skip_reading_prefix_cache: Option, /// Additional request parameters for custom extensions (from `vllm_xargs`). pub extra_args: Option>, + /// Number of prompt tokens to skip from returned routed-expert data. + /// A value of zero returns routing data for the entire prompt. + pub routed_experts_prompt_start: u32, } impl EngineCoreSamplingParams { @@ -175,6 +178,7 @@ impl EngineCoreSamplingParams { logprob_token_ids: None, skip_reading_prefix_cache: None, extra_args: None, + routed_experts_prompt_start: 0, } } } diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index df9eec90b361..fa5d8e0db855 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -159,6 +159,7 @@ fn sample_request_with_id(request_id: &str) -> EngineCoreRequest { stop_token_ids: vec![151643], eos_token_id: Some(151645), all_stop_token_ids: BTreeSet::from([151643, 151645]), + routed_experts_prompt_start: 1, ..EngineCoreSamplingParams::for_test() }), arrival_time: 42.5, @@ -2680,6 +2681,7 @@ fn python_msgpack_fixtures_match_rust_encoding() { logprob_token_ids: None, skip_reading_prefix_cache: None, extra_args: None, + routed_experts_prompt_start: 0, }, ); diff --git a/rust/src/engine-core-client/src/tests/python_compat.py b/rust/src/engine-core-client/src/tests/python_compat.py index ac033914b0d1..0cbd2f4a5973 100755 --- a/rust/src/engine-core-client/src/tests/python_compat.py +++ b/rust/src/engine-core-client/src/tests/python_compat.py @@ -47,6 +47,7 @@ class EngineCoreSamplingParams(msgspec.Struct, dict=True, omit_defaults=True): stop_token_ids: list[int] = [] _eos_token_id: int | None = None _all_stop_token_ids: set[int] = set() + routed_experts_prompt_start: int = 0 output_kind: RequestOutputKind = RequestOutputKind.DELTA @@ -135,6 +136,7 @@ class EngineCoreOutputs( stop_token_ids=[151643], _eos_token_id=151645, _all_stop_token_ids={151643, 151645}, + routed_experts_prompt_start=1, output_kind=RequestOutputKind.FINAL_ONLY, ), pooling_params=None, diff --git a/rust/src/text/src/lower.rs b/rust/src/text/src/lower.rs index 9c0a5905d547..55be99cc4e57 100644 --- a/rust/src/text/src/lower.rs +++ b/rust/src/text/src/lower.rs @@ -188,6 +188,7 @@ pub fn lower_sampling_params( logprob_token_ids, skip_reading_prefix_cache, extra_args: vllm_xargs, + routed_experts_prompt_start: 0, }; validate_resolved_sampling_params(¶ms)?; validate_vocab_range(¶ms, &sampling_limits)?; @@ -645,6 +646,7 @@ mod tests { logprob_token_ids: None, skip_reading_prefix_cache: None, extra_args: None, + routed_experts_prompt_start: 0, } "#]] .assert_debug_eq(¶ms); @@ -694,6 +696,7 @@ mod tests { logprob_token_ids: None, skip_reading_prefix_cache: None, extra_args: None, + routed_experts_prompt_start: 0, } "#]] .assert_debug_eq(¶ms); @@ -859,6 +862,7 @@ mod tests { logprob_token_ids: None, skip_reading_prefix_cache: None, extra_args: None, + routed_experts_prompt_start: 0, } "#]] .assert_debug_eq(¶ms); @@ -926,6 +930,7 @@ mod tests { logprob_token_ids: None, skip_reading_prefix_cache: None, extra_args: None, + routed_experts_prompt_start: 0, } "#]] .assert_debug_eq(¶ms); @@ -986,6 +991,7 @@ mod tests { logprob_token_ids: None, skip_reading_prefix_cache: None, extra_args: None, + routed_experts_prompt_start: 0, } "#]] .assert_debug_eq(¶ms); @@ -1235,6 +1241,7 @@ mod tests { logprob_token_ids: None, skip_reading_prefix_cache: None, extra_args: None, + routed_experts_prompt_start: 0, } "#]] .assert_debug_eq(¶ms); From 7ddb50788ddf0846518feff972028534e39bcb86 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Wed, 19 Aug 2026 04:49:35 +0800 Subject: [PATCH 105/839] [Bugfix][DP] Don't assume the engines started when forwarding a wake (#51481) Signed-off-by: aoshen02 Co-authored-by: Claude Opus 5 (1M context) --- tests/v1/distributed/test_async_llm_dp.py | 137 ++++++++++++++++++++++ vllm/v1/engine/coordinator.py | 6 +- vllm/v1/engine/core.py | 14 +++ 3 files changed, 155 insertions(+), 2 deletions(-) diff --git a/tests/v1/distributed/test_async_llm_dp.py b/tests/v1/distributed/test_async_llm_dp.py index 711daaef7400..f51fde90534d 100644 --- a/tests/v1/distributed/test_async_llm_dp.py +++ b/tests/v1/distributed/test_async_llm_dp.py @@ -299,6 +299,143 @@ async def test_dp_pause_resume_basic(expert_parallel: bool): assert out.finished +async def _consume(generator) -> None: + async for _ in generator: + pass + + +async def _poll_flag(engine: AsyncLLM, want: bool, timeout: float) -> bool: + """Wait for the front-end's view of the DP engines to reach ``want``.""" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if engine.engine_core.dp_engines_running() == want: + return True + await asyncio.sleep(0.05) + return False + + +@pytest.mark.asyncio +async def test_dp_pause_late_request_does_not_block_drain(): + """A request arriving after pause must not leave the coordinator believing + the engines are running. + + Paused engines discard START_DP_WAVE, so nothing can report wave + completion afterwards; if forwarding the wake also marks the engines as + running, drain has no way back and can only end in a timeout. + + MoE only: wave coordination is enabled iff the model is MoE, so a dense + model never reaches the coordinator state this exercises. + """ + with ExitStack() as after: + engine_args = _get_dp_pause_engine_args(expert_parallel=True) + engine = AsyncLLM.from_engine_args(engine_args) + after.callback(engine.shutdown) + + # Run a wave first, so the engines are quiesced by the pause rather + # than by never having started, and drain has something to observe. + long_request = asyncio.create_task( + _consume( + engine.generate( + request_id="warmup", + prompt=DP_PAUSE_PROMPT, + sampling_params=SamplingParams(max_tokens=400, ignore_eos=True), + ) + ) + ) + + # A design that never reports the engines as running would satisfy + # every drain assertion below while destroying the signal, so pin it + # down first. The front-end sets its own copy optimistically when it + # forwards the wake, so sample only after several coordinator + # publishes (every 100ms while stats change) have overwritten it. + await asyncio.sleep(2) + assert not long_request.done(), "the warmup request was too short to sample" + assert engine.engine_core.dp_engines_running(), ( + "the coordinator does not report the engines as running while they are" + ) + + await long_request + assert await _poll_flag(engine, False, timeout=30) + + await engine.pause_generation(mode="abort") + await engine.wait_for_requests_to_drain(drain_timeout=30) + + # Awaiting add_request guarantees the new-request notification has been + # sent to the coordinator - the message that used to latch the flag. + collector = await engine.add_request( + request_id="late", + prompt=DP_PAUSE_PROMPT, + params=SamplingParams(max_tokens=5), + ) + + # The front-end marks the engines running off the back of that + # notification. This is what makes the test non-vacuous: it is the + # path that used to leave the coordinator stuck. + assert await _poll_flag(engine, True, timeout=5), ( + "the late request did not notify the coordinator" + ) + + # It must settle back by itself. Unfixed it never does, because the + # paused engines discard the wake and so never report wave completion. + assert await _poll_flag(engine, False, timeout=60) + + # The late request was held rather than dropped: it completes on resume. + await engine.resume_generation() + while True: + out = await asyncio.wait_for(collector.get(), timeout=60) + if out.finished: + break + + +@pytest.mark.asyncio +async def test_dp_sleep_late_request_does_not_block_drain(): + """The same latch, reached through sleep rather than pause. + + Sleep stops the engines stepping just as pause does, so a request arriving + while they are asleep can leave the coordinator believing they are running + with nothing able to say otherwise. This is worth pinning separately from + the pause case because it is the shape that reaches + `_drain_requests_for_elastic_ep`, which decides from the same signal + whether it is safe to scale. + """ + with ExitStack() as after: + engine_args = _get_dp_pause_engine_args(expert_parallel=True) + engine = AsyncLLM.from_engine_args(engine_args) + after.callback(engine.shutdown) + + async for _ in engine.generate( + request_id="warmup", + prompt=DP_PAUSE_PROMPT, + sampling_params=SamplingParams(max_tokens=5), + ): + pass + assert await _poll_flag(engine, False, timeout=30) + + await engine.sleep(level=1) + assert await engine.is_sleeping() + + collector = await engine.add_request( + request_id="while-asleep", + prompt=DP_PAUSE_PROMPT, + params=SamplingParams(max_tokens=5), + ) + + assert await _poll_flag(engine, True, timeout=5), ( + "the request did not notify the coordinator" + ) + + # Sleeping engines cannot report wave completion, so a coordinator that + # marked them running when it forwarded the wake never hears otherwise. + assert await _poll_flag(engine, False, timeout=60) + + await engine.wake_up() + assert not await engine.is_sleeping() + while True: + out = await asyncio.wait_for(collector.get(), timeout=60) + if out.finished: + break + + @pytest.mark.asyncio @pytest.mark.parametrize("expert_parallel", [False, True]) async def test_dp_pause_abort(expert_parallel: bool): diff --git a/vllm/v1/engine/coordinator.py b/vllm/v1/engine/coordinator.py index d7f05cffc8a6..efadbfea2617 100644 --- a/vllm/v1/engine/coordinator.py +++ b/vllm/v1/engine/coordinator.py @@ -358,8 +358,10 @@ def process_input_socket( # is handled by all the engines. engine_to_exclude = None - engines_running = True - wave_state_changed = True + # engines_running is only set from the engines' + # own notifications; a paused engine discards + # START_DP_WAVE, so sending it is not evidence + # that the engines are running. self._send_start_wave( publish_back, current_wave, engine_to_exclude ) diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 641c5ecc4e01..71a791d72185 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -2156,6 +2156,7 @@ def run_busy_loop(self): # Loop until process is sent a SIGINT or SIGTERM while self._handle_shutdown(): # 1) Poll the input queue until there is work to do. + was_running = self.engines_running self._process_input_queue() # Publish request counts before and after GPU step to ensure freshness. self._maybe_publish_request_counts() @@ -2216,6 +2217,19 @@ def run_busy_loop(self): # Increment wave count and reset step counter. self.current_wave += 1 self.step_counter = 0 + elif ( + not was_running + and self.has_coordinator + and self.dp_rank == 0 + and not self.pending_pause + ): + # Mirror of the wave_complete notification above: the + # coordinator must observe this edge too rather than assume + # that a START_DP_WAVE it sent was acted upon, since a paused + # engine discards it. + self.output_queue.put_nowait( + (-1, EngineCoreOutputs(start_wave=self.current_wave)) + ) raise SystemExit From 12f64b39d29282437e35be9aa5db432fb2a1a6e6 Mon Sep 17 00:00:00 2001 From: Flora Feng <4florafeng@gmail.com> Date: Tue, 18 Aug 2026 17:21:43 -0400 Subject: [PATCH 106/839] [Bugfix][Structured Output] Stop XGrammar token batches at termination (#52805) Signed-off-by: sfeng33 <4florafeng@gmail.com> --- .../spec_decode/test_mtp_structured_output.py | 48 +++++++++++++++++++ vllm/v1/structured_output/backend_xgrammar.py | 19 ++++++-- 2 files changed, 62 insertions(+), 5 deletions(-) diff --git a/tests/v1/spec_decode/test_mtp_structured_output.py b/tests/v1/spec_decode/test_mtp_structured_output.py index 619f3ad6fded..8bd599733a24 100644 --- a/tests/v1/spec_decode/test_mtp_structured_output.py +++ b/tests/v1/spec_decode/test_mtp_structured_output.py @@ -262,6 +262,54 @@ def test_validate_tokens_then_bitmask_round_trip(backend): assert not grammar.is_terminated() +def test_xgrammar_accept_tokens_stops_at_termination(capfd): + """Tokens after a terminating EOS do not reach the matcher.""" + tokenizer, _, request, prompt = _make_manager_and_request("xgrammar") + grammar = request.structured_output_request.grammar + + assert grammar.accept_tokens(request.request_id, prompt) + + eos = tokenizer.eos_token_id + trailing = tokenizer.encode("\n")[0] + processed_before = grammar.num_processed_tokens + + assert grammar.accept_tokens(request.request_id, [eos, trailing]) + assert grammar.is_terminated() + assert grammar.num_processed_tokens == processed_before + 1 + assert "trying to accept new token" not in capfd.readouterr().err + + processed_after_eos = grammar.num_processed_tokens + assert grammar.accept_tokens(request.request_id, [trailing]) + assert grammar.num_processed_tokens == processed_after_eos + assert "trying to accept new token" not in capfd.readouterr().err + + grammar.reset() + assert not grammar.is_terminated() + assert grammar.num_processed_tokens == 0 + + +def test_xgrammar_validate_tokens_stops_at_termination(capfd): + """Validation rolls back after reaching a terminating EOS.""" + tokenizer, _, request, prompt = _make_manager_and_request("xgrammar") + grammar = request.structured_output_request.grammar + + assert grammar.accept_tokens(request.request_id, prompt) + + eos = tokenizer.eos_token_id + trailing = tokenizer.encode("\n")[0] + + assert grammar.validate_tokens([eos, trailing]) == [eos] + assert "trying to accept new token" not in capfd.readouterr().err + # Check matcher state directly to verify validation rolled it back. + assert not grammar.matcher.is_terminated() + + assert grammar.accept_tokens(request.request_id, [eos]) + assert grammar.is_terminated() + + assert grammar.validate_tokens([trailing]) == [] + assert "trying to accept new token" not in capfd.readouterr().err + + class _MarkerReasoner: """Stub reasoner whose reasoning-end marker is a single fixed token.""" diff --git a/vllm/v1/structured_output/backend_xgrammar.py b/vllm/v1/structured_output/backend_xgrammar.py index 258b1dff32f1..5b24a19780aa 100644 --- a/vllm/v1/structured_output/backend_xgrammar.py +++ b/vllm/v1/structured_output/backend_xgrammar.py @@ -157,11 +157,12 @@ class XgrammarGrammar(StructuredOutputGrammar): def accept_tokens(self, request_id: str, tokens: list[int]) -> bool: """Accepts a list of tokens and advances the FSM. - Returns True if the FSM was advanced successfully. - Returns False if the FSM failed to advance. + Returns True if all grammar-constrained tokens were accepted. + Tokens after termination are ignored. Returns False if the FSM + failed to advance. """ if self._is_terminated: - return False + return True for token in tokens: if not self.matcher.accept_token(token): logger.error( @@ -172,7 +173,9 @@ def accept_tokens(self, request_id: str, tokens: list[int]) -> bool: ) return False self.num_processed_tokens += 1 - self._is_terminated = self.matcher.is_terminated() + self._is_terminated = self.matcher.is_terminated() + if self._is_terminated: + break return True def validate_tokens(self, tokens: list[int]) -> list[int]: @@ -181,10 +184,15 @@ def validate_tokens(self, tokens: list[int]) -> list[int]: Returns the prefix list of tokens that are accepted by the FSM. """ + if self._is_terminated: + return [] + accepted_tokens = [] for token in tokens: if self.matcher.accept_token(token): accepted_tokens.append(token) + if self.matcher.is_terminated(): + break else: break if len(accepted_tokens) > 0: @@ -204,8 +212,9 @@ def is_terminated(self) -> bool: return self._is_terminated def reset(self): - self.num_processed_tokens = 0 self.matcher.reset() + self.num_processed_tokens = 0 + self._is_terminated = False # cf https://github.com/mlc-ai/xgrammar/blob/a32ac892676d2eedc0327416105b9b06edfb94b2/cpp/json_schema_converter.cc From 5d8a4cf9761ec890ae9f11502fe0f4a13e9981a5 Mon Sep 17 00:00:00 2001 From: yinfengLiu Date: Wed, 19 Aug 2026 05:37:10 +0800 Subject: [PATCH 107/839] [ROCm] Pad non-aligned AITER MLA heads (#51647) Signed-off-by: Liuyinfeng01 Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../test_rocm_aiter_mla_head_padding.py | 125 +++++++++++++++--- .../test_rocm_aiter_mla_mtp_split.py | 15 ++- .../attention/backends/mla/rocm_aiter_mla.py | 110 ++++++++++----- .../backends/mla/rocm_aiter_mla_sparse.py | 8 +- 4 files changed, 200 insertions(+), 58 deletions(-) diff --git a/tests/kernels/attention/test_rocm_aiter_mla_head_padding.py b/tests/kernels/attention/test_rocm_aiter_mla_head_padding.py index 744f1852baf5..8a284e320aac 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_head_padding.py +++ b/tests/kernels/attention/test_rocm_aiter_mla_head_padding.py @@ -1,13 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Head-count padding + kernel selection for the ROCm AITER MLA backend. - -Kimi-K3 at TP8 puts 12 heads/rank and at TP16 puts 6 heads/rank on the AITER -MLA decode. The asm persistent decode requires exactly 16 heads, so small -head counts are tile-padded to 16 and the padding heads are sliced back off -the output. Divisor counts (1/2/4/8) may keep the Gluon kernel, but only on -gfx950 where that kernel has a build; every small head count on gfx942 (which -has no Gluon build) is routed to the asm persistent decode instead. +"""Head padding and kernel selection for the ROCm AITER MLA backend. + +The asm persistent decode requires a 16-aligned head count, so unaligned +counts through 128 are tile-padded to the next multiple of 16 and sliced back +off the output. Small divisor counts (1/2/4/8) preserve their existing +repeat-interleave path and may keep the Gluon kernel on gfx950. """ import math @@ -19,6 +17,7 @@ from vllm._aiter_ops import is_aiter_found from vllm.platforms import current_platform +from vllm.v1.attention.backends.mla import rocm_aiter_mla from vllm.v1.attention.backends.mla.rocm_aiter_mla import ( AiterMLADecodeMetadata, AiterMLAHelper, @@ -33,12 +32,20 @@ CONTEXT_LEN = 4096 SCALE = 1.0 / math.sqrt(QK_HEAD_DIM) -# Non-divisor counts must go through the tile-and-slice path; divisor counts -# (of 16) keep repeat_interleave. Both must pad to exactly 16 and round-trip. +# Small non-divisor counts use tile-and-slice; divisors of 16 keep +# repeat_interleave. Both pad to exactly 16 and round-trip. NON_DIVISOR_HEADS = [3, 5, 6, 7, 9, 10, 11, 12, 13, 14, 15] DIVISOR_HEADS = [1, 2, 4, 8] +@pytest.fixture(autouse=True) +def _disable_native_h24(monkeypatch): + """Exercise the padding fallback unless a test explicitly enables H24.""" + monkeypatch.setattr( + rocm_aiter_mla, "_aiter_mla_native_h24_supported", lambda: False + ) + + def _rocm_aiter_available() -> bool: return current_platform.is_rocm() and is_aiter_found() and torch.cuda.is_available() @@ -123,6 +130,72 @@ def test_h12_output_discards_padding_heads(): torch.testing.assert_close(unpadded_o, o[:, :12]) +def test_h24_query_is_tile_padded_to_h32(): + q = torch.arange(2 * 24 * 4, dtype=torch.float32).view(2, 24, 4) + + padded_q = AiterMLAHelper.get_mla_padded_q(24, q) + + assert padded_q.shape == (2, 32, 4) + assert padded_q.is_contiguous() + torch.testing.assert_close(padded_q[:, :24], q) + torch.testing.assert_close(padded_q[:, 24:], q[:, :8]) + + +def test_h24_output_discards_h32_padding_heads(): + o = torch.arange(2 * 32 * 4, dtype=torch.float32).view(2, 32, 4) + + unpadded_o = AiterMLAHelper.get_mla_unpadded_o(24, o) + + assert unpadded_o.shape == (2, 24, 4) + torch.testing.assert_close(unpadded_o, o[:, :24]) + + +def test_h24_reducer_without_metadata_still_pads_to_h32(monkeypatch): + monkeypatch.setattr( + rocm_aiter_mla, "_aiter_mla_native_h24_reducer_supported", lambda: True + ) + monkeypatch.setattr( + rocm_aiter_mla, "_aiter_mla_native_h24_metadata_supported", lambda: False + ) + monkeypatch.setattr( + rocm_aiter_mla, + "_aiter_mla_native_h24_supported", + lambda: ( + rocm_aiter_mla._aiter_mla_native_h24_reducer_supported() + and rocm_aiter_mla._aiter_mla_native_h24_metadata_supported() + ), + ) + q = torch.arange(2 * 24 * 4, dtype=torch.float32).view(2, 24, 4) + + assert AiterMLAHelper.get_actual_mla_num_heads(24) == 32 + padded_q = AiterMLAHelper.get_mla_padded_q(24, q) + assert padded_q.shape == (2, 32, 4) + torch.testing.assert_close(padded_q[:, :24], q) + torch.testing.assert_close(padded_q[:, 24:], q[:, :8]) + + +def test_native_h24_requires_reducer_and_metadata(monkeypatch): + monkeypatch.setattr( + rocm_aiter_mla, "_aiter_mla_native_h24_reducer_supported", lambda: True + ) + monkeypatch.setattr( + rocm_aiter_mla, "_aiter_mla_native_h24_metadata_supported", lambda: True + ) + monkeypatch.setattr( + rocm_aiter_mla, + "_aiter_mla_native_h24_supported", + lambda: ( + rocm_aiter_mla._aiter_mla_native_h24_reducer_supported() + and rocm_aiter_mla._aiter_mla_native_h24_metadata_supported() + ), + ) + q = torch.arange(2 * 24 * 4, dtype=torch.float32).view(2, 24, 4) + + assert AiterMLAHelper.get_actual_mla_num_heads(24) == 24 + assert AiterMLAHelper.get_mla_padded_q(24, q) is q + assert AiterMLAHelper.get_mla_unpadded_o(24, q) is q + + def test_existing_divisor_head_mapping_is_unchanged(): q = torch.arange(2 * 8 * 4, dtype=torch.bfloat16).view(2, 8, 4) @@ -134,6 +207,18 @@ def test_existing_divisor_head_mapping_is_unchanged(): torch.testing.assert_close(unpadded_o, q) +@pytest.mark.parametrize("num_heads", [17, 24, 31]) +def test_unaligned_head_counts_round_trip_through_h32(num_heads: int): + q = torch.arange(2 * num_heads * 4, dtype=torch.float32).view(2, num_heads, 4) + + padded_q = AiterMLAHelper.get_mla_padded_q(num_heads, q) + unpadded_o = AiterMLAHelper.get_mla_unpadded_o(num_heads, padded_q) + + assert padded_q.shape == (2, 32, 4) + assert padded_q.is_contiguous() + torch.testing.assert_close(unpadded_o, q) + + @pytest.mark.parametrize("num_heads", NON_DIVISOR_HEADS + DIVISOR_HEADS) def test_all_small_head_counts_pad_to_16_and_round_trip(num_heads: int): q = torch.arange(2 * num_heads * 4, dtype=torch.float32).view(2, num_heads, 4) @@ -147,26 +232,30 @@ def test_all_small_head_counts_pad_to_16_and_round_trip(num_heads: int): torch.testing.assert_close(unpadded_o, q) -def test_num_heads_ge_16_is_passthrough(): - q = torch.arange(2 * 16 * 4, dtype=torch.float32).view(2, 16, 4) - assert AiterMLAHelper.get_mla_padded_q(16, q) is q - assert AiterMLAHelper.get_mla_unpadded_o(16, q) is q +def test_aligned_h32_is_zero_copy(): + q = torch.arange(2 * 32 * 4, dtype=torch.float32).view(2, 32, 4) + assert AiterMLAHelper.get_mla_padded_q(32, q) is q + assert AiterMLAHelper.get_mla_unpadded_o(32, q) is q def test_is_valid_num_heads(): - for n in range(1, 16): + for n in range(1, 129): assert AiterMLAHelper.is_valid_num_heads(n) - assert AiterMLAHelper.is_valid_num_heads(16) - assert AiterMLAHelper.is_valid_num_heads(32) + assert AiterMLAHelper.is_valid_num_heads(24) + assert AiterMLAHelper.is_valid_num_heads(127) + # Aligned counts remain valid above the range where padding is supported. + assert AiterMLAHelper.is_valid_num_heads(144) assert not AiterMLAHelper.is_valid_num_heads(0) + assert not AiterMLAHelper.is_valid_num_heads(129) def test_nondivisor_and_multitoken_never_use_gluon(): # Non-divisor decode always takes the asm path (12 heads/rank at TP8). assert not AiterMLAHelper.use_gluon_decode(12, 1, "auto") assert not AiterMLAHelper.use_gluon_decode(6, 1, "auto") - # >=16 heads never pad, never Gluon. + # >=16 heads never use Gluon, including unaligned counts padded for asm. assert not AiterMLAHelper.use_gluon_decode(16, 1, "auto") + assert not AiterMLAHelper.use_gluon_decode(24, 1, "auto") # Multi-token (verify / qlen>1) is never the single-token Gluon decode. assert not AiterMLAHelper.use_gluon_decode(8, 4, "auto") assert not AiterMLAHelper.use_gluon_decode(12, 4, "auto") diff --git a/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py b/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py index c5eb921a22e2..1894d8aa9cdf 100644 --- a/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py +++ b/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py @@ -15,6 +15,7 @@ from vllm.v1.attention.backends.mla import rocm_aiter_mla # noqa: E402 from vllm.v1.attention.backends.mla.rocm_aiter_mla import ( # noqa: E402 + AiterMLAHelper, AiterMLAMetadataBuilder, ) @@ -83,7 +84,7 @@ def _builder( _uniform_padded_mtp_qo_len=(AiterMLAMetadataBuilder._uniform_padded_mtp_qo_len), _use_persistent_metadata=False, kernel_block_size=kernel_block_size, - _num_attention_heads=16, + _num_attention_heads=AiterMLAHelper.get_actual_mla_num_heads(num_heads), _mla_work_meta_data=torch.empty(1, dtype=torch.int32), _mla_work_info_set=torch.empty(1, dtype=torch.int32), _mla_work_indptr=torch.empty(1, dtype=torch.int32), @@ -109,7 +110,7 @@ def test_backend_declares_uniform_batch_support(): ) -@pytest.mark.parametrize("num_heads", [8, 16, 32, 64, 128]) +@pytest.mark.parametrize("num_heads", [8, 16, 24, 32, 64, 128]) @pytest.mark.parametrize( "spec_method, parallel_drafting", [ @@ -127,8 +128,8 @@ def test_mtp_builder_init_sizes_native_fp8_metadata( ): """Aiter init sizes the metadata for every query length decode can be handed. - Sweeping num_heads asserts the max(16, num_heads) clamp is what sizes the - metadata, covering the fp8 nhead=32 (TP4) fold path. + Sweeping num_heads asserts metadata is sized for the padded decode shape, + covering Kimi-K3 TP4's 24 -> 32 head path and native fp8 nhead=32 folding. """ dtypes = SimpleNamespace(fp8="fp8", fp16="fp16", bf16="bf16") @@ -207,7 +208,7 @@ def init_common_builder(self, *args, **kwargs): { "max_batch_size": config.scheduler_config.max_num_seqs, "max_qo_len": builder.reorder_batch_threshold, - "num_attention_heads": max(16, num_heads), + "num_attention_heads": AiterMLAHelper.get_actual_mla_num_heads(num_heads), "q_dtype": dtypes.fp8, "kv_dtype": dtypes.fp8, "is_sparse": False, @@ -381,6 +382,7 @@ def test_decode_expands_kernel_block_page_indices(monkeypatch): (1, 1, 16, "auto", True), # non-MTP decode (4, 2, 16, "auto", True), # MTP deployment, in-range step (4, 4, 16, "auto", True), # MTP deployment, full-qlen verification step + (1, 1, 24, "auto", True), # unaligned H24 pads to H32 persistent decode (2, 4, 16, "auto", False), # step demand exceeds provisioned K -> fallback (1, 1, 8, "auto", False), # divisor head count -> Gluon decode owns qlen==1 (4, 4, 8, "auto", False), # small head count -> Gluon flatten owns qlen>1 @@ -410,7 +412,8 @@ def test_persistent_metadata_gate( K = _mtp_decode_qlen sizes the metadata buffers at init; a decode step gets the pre-built schedule only when its qlen fits those buffers, otherwise it falls back to the kernel computing its own. qlen==1 (non-MTP) must stay - in-range -- dropping it is the regression this guards. + in-range -- dropping it is the regression this guards. This includes + unaligned H24, whose padded H32 decode uses persistent metadata. Only the Gluon paths ignore the schedule, so the gate follows the routing predicates rather than the raw head count. Reading `num_heads >= 16` instead diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py index 2f832a1357a9..51f487947a2b 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py @@ -3,6 +3,7 @@ import functools from dataclasses import dataclass +from pathlib import Path from typing import ClassVar, Final import torch @@ -77,6 +78,48 @@ def _fp8_mla_prefill_supported() -> bool: return True +@functools.lru_cache(maxsize=1) +def _aiter_mla_native_h24_reducer_supported() -> bool: + """Whether AITER's JIT reducer supports the native H24/512 shape.""" + try: + from aiter.jit.core import AITER_CSRC_DIR + + reduce_source = Path(AITER_CSRC_DIR) / "kernels" / "mla" / "reduce.cu" + source = "".join(reduce_source.read_text(encoding="utf-8").split()) + except (ImportError, OSError): + return False + return "MLA_REDUCE_CASE_EF(NUM_HEAD,24,HEAD_DIM,512," in source + + +@functools.lru_cache(maxsize=1) +def _aiter_mla_native_h24_metadata_supported() -> bool: + """Whether AITER's fast MLA metadata planner accepts native H24. + + The reducer and metadata planner have independent shape dispatch. Checking + only the reducer can route H24 into a planner that rejects it before the + attention kernel launches. Until AITER exposes a capability API, inspect + the shipped JIT source for an explicit native-H24 planner branch. + """ + try: + from aiter.jit.core import AITER_CSRC_DIR + + metadata_source = ( + Path(AITER_CSRC_DIR) / "kernels" / "mla" / "metadata" / "v1_2_device.cuh" + ) + source = "".join(metadata_source.read_text(encoding="utf-8").split()) + except (ImportError, OSError): + return False + return "num_heads==24" in source + + +def _aiter_mla_native_h24_supported() -> bool: + """Whether the complete AITER decode path supports native H24.""" + return ( + _aiter_mla_native_h24_reducer_supported() + and _aiter_mla_native_h24_metadata_supported() + ) + + @functools.lru_cache(maxsize=1) def _gluon_mla_decode_supported() -> bool: """The small-head Gluon MLA decode kernel only has a gfx950 (CDNA4) build. @@ -308,10 +351,11 @@ def __init__( from aiter import dtypes, get_mla_metadata_info_v1 - # For num_attention_heads < 16 (e.g. kimi-k2.5 head=8 with TP8), - # make sure get_mla_metadata_info_v1 / get_mla_metadata_v1 are consistent - # with the actual tensor shape passed to mla_decode_fwd. - self._num_attention_heads = max(16, self.num_heads) + # Keep metadata sizing consistent with the padded tensor shape passed + # to mla_decode_fwd. + self._num_attention_heads = AiterMLAHelper.get_actual_mla_num_heads( + self.num_heads + ) kv_cache_dtype_str = getattr(vllm_config.cache_config, "cache_dtype", "auto") if kv_cache_dtype_str in ("fp8", "fp8_e4m3", "fp8_e5m2"): kv_cache_dtype_str = "fp8" @@ -839,14 +883,16 @@ def _expand_page_indices_kernel( class AiterMLAHelper: """ - AITER MLA persistent (asm) decode requires num_heads >= 16. Head counts - < 16 are padded up to exactly 16: divisors of 16 by repeat_interleave, - other counts (e.g. 12 heads/rank at TP8, 6 at TP16) by tiling the query - heads and slicing to 16. Non-divisor padded decodes take the asm path; - divisors and max_qo_len > 1 small-head verify still use Gluon. + AITER MLA persistent (asm) decode requires a multiple of 16 heads. Unaligned + head counts through 128 are padded to the next multiple of 16 by tiling the + query heads and slicing to the padded size. Native H24 AITER builds bypass + that padding. Small divisors of 16 retain the existing repeat_interleave and + strided-unpad behavior. Native and aligned counts pass through without + copies. """ _AITER_MIN_MLA_HEADS: Final = 16 + _AITER_MAX_PADDED_MLA_HEADS: Final = 128 # Largest qlen the padded gqa=16 asm decode has a bf16 persistent kernel # for. Above it only the non-persistent qseqlen=8 entry exists, and the # fold that reaches a persistent one is gfx950-only. @@ -856,48 +902,50 @@ class AiterMLAHelper: @staticmethod def check_num_heads_validity(num_heads: int): assert AiterMLAHelper.is_valid_num_heads(num_heads), ( - "ROCM AITER MLA requires 1-15 heads (padded to 16 for asm " - "persistent decode; exact divisors of 16 may keep Gluon) or a " - f"multiple of 16 heads, but got {num_heads}.\n" + "ROCM AITER MLA requires a positive multiple of 16 heads, or an " + "unaligned head count up to 128 (padded to the next multiple of " + f"16), but got {num_heads}.\n" f"Try adjusting tensor_parallel_size value." ) @staticmethod def is_valid_num_heads(num_heads: int) -> bool: - return num_heads > 0 and ( - num_heads < AiterMLAHelper._AITER_MIN_MLA_HEADS - or num_heads % AiterMLAHelper._AITER_MIN_MLA_HEADS == 0 + return ( + num_heads > 0 + and num_heads not in AiterMLAHelper._AITER_UNSUPPORTED_HEADS + and ( + num_heads <= AiterMLAHelper._AITER_MAX_PADDED_MLA_HEADS + or num_heads % AiterMLAHelper._AITER_MIN_MLA_HEADS == 0 + ) ) @staticmethod def get_actual_mla_num_heads(num_heads: int) -> int: - return max(num_heads, AiterMLAHelper._AITER_MIN_MLA_HEADS) + if num_heads == 24 and _aiter_mla_native_h24_supported(): + return num_heads + m = AiterMLAHelper._AITER_MIN_MLA_HEADS + return -(-num_heads // m) * m @staticmethod def get_mla_padded_q(num_heads: int, q: torch.Tensor) -> torch.Tensor: - m = AiterMLAHelper._AITER_MIN_MLA_HEADS - if num_heads >= m: + m = AiterMLAHelper.get_actual_mla_num_heads(num_heads) + if num_heads == m: return q if m % num_heads == 0: return q.repeat_interleave(m // num_heads, dim=1) - # Non-divisor head counts (e.g. 12 heads/rank at TP8, 6 at TP16) cannot - # be padded by repeat_interleave. Tile the query heads and slice to - # exactly m; this reaches m for any 0 < num_heads < m (unlike a single - # append, which under-pads when num_heads < m - num_heads). MLA - # attention is independent per query head over the shared KV, so the - # padding heads cannot affect heads [0:num_heads]; they are sliced back - # off in get_mla_unpadded_o. + # Non-divisor head counts cannot be padded by repeat_interleave. Tile + # the query heads and slice to exactly m. MLA attention is independent + # per query head over the shared KV, so padding heads cannot affect + # heads [0:num_heads]; they are sliced back off the output. reps = -(-m // num_heads) # ceil(m / num_heads) - # Slicing a tiled tensor down to m yields a non-contiguous view whenever - # reps * num_heads > m (the common case: TP8 12->24->16, TP16 6->18->16). - # The asm persistent decode reads q as a packed [tokens, m, head_dim] - # buffer, so materialize a contiguous copy. No-op when already contiguous. + # Slicing a tiled tensor yields a non-contiguous view. The asm decode + # reads q as packed [tokens, m, head_dim], so materialize it. return q.repeat(1, reps, 1)[:, :m, :].contiguous() @staticmethod def get_mla_unpadded_o(num_heads: int, o: torch.Tensor) -> torch.Tensor: - m = AiterMLAHelper._AITER_MIN_MLA_HEADS - if num_heads >= m: + m = AiterMLAHelper.get_actual_mla_num_heads(num_heads) + if num_heads == m: return o if m % num_heads == 0: return o[:, :: m // num_heads, :] diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py index 3f10d6105349..6ae880251f87 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py @@ -422,9 +422,11 @@ def __init__( # so the buffers are large enough for any decode shape we might see. from aiter import dtypes, get_mla_metadata_info_v1 - # Aiter sparse MLA also requires num_heads >= 16 (will be padded by - # AiterMLAHelper.get_mla_padded_q in forward). - self._num_attention_heads = max(16, self.num_heads) + # Keep metadata sizing consistent with the padded tensor shape passed + # to the sparse decode kernel. + self._num_attention_heads = AiterMLAHelper.get_actual_mla_num_heads( + self.num_heads + ) q_dtype = self.model_dtype kv_cache_dtype_str = getattr(vllm_config.cache_config, "cache_dtype", "auto") From 6066bb3d502107d27383a64b306d36148421d4b9 Mon Sep 17 00:00:00 2001 From: stefankoncarevic Date: Tue, 18 Aug 2026 23:40:30 +0200 Subject: [PATCH 108/839] [ROCM][CI] Attention test speedup (#52763) Signed-off-by: Stefan Koncarevic --- tests/kernels/attention/test_attention.py | 2 ++ tests/kernels/attention/test_cache.py | 2 ++ tests/kernels/attention/test_cutlass_mla_decode.py | 2 ++ tests/kernels/attention/test_flashmla.py | 2 ++ tests/kernels/attention/test_merge_attn_states.py | 14 ++++++++------ tests/kernels/attention/test_prefix_prefill.py | 2 ++ .../attention/test_triton_decode_attention.py | 2 ++ .../attention/test_triton_unified_attention.py | 2 ++ .../test_triton_unified_attention_diffkv.py | 2 ++ 9 files changed, 24 insertions(+), 6 deletions(-) diff --git a/tests/kernels/attention/test_attention.py b/tests/kernels/attention/test_attention.py index 662e14a9bb87..ef3a860eef89 100644 --- a/tests/kernels/attention/test_attention.py +++ b/tests/kernels/attention/test_attention.py @@ -14,6 +14,8 @@ from vllm.utils.mem_utils import get_max_shared_memory_bytes from vllm.utils.torch_utils import set_random_seed +pytestmark = pytest.mark.skip_global_cleanup + FLOAT32_BYTES = torch.finfo(torch.float).bits // 8 # This will change depending on the compute capability. # - 512 as a buffer diff --git a/tests/kernels/attention/test_cache.py b/tests/kernels/attention/test_cache.py index b6a029ce08a5..7b7317ebaadf 100644 --- a/tests/kernels/attention/test_cache.py +++ b/tests/kernels/attention/test_cache.py @@ -12,6 +12,8 @@ from vllm.platforms import current_platform from vllm.utils.torch_utils import nvfp4_split_data_scale, set_random_seed +pytestmark = pytest.mark.skip_global_cleanup + COPYING_DIRECTION = [("cuda", "cpu"), ("cuda", "cuda"), ("cpu", "cuda")] DTYPES = [torch.bfloat16, torch.float] NUM_TOKENS = [42] # Arbitrary values for testing diff --git a/tests/kernels/attention/test_cutlass_mla_decode.py b/tests/kernels/attention/test_cutlass_mla_decode.py index c0e319a27ad2..067f66ebc495 100644 --- a/tests/kernels/attention/test_cutlass_mla_decode.py +++ b/tests/kernels/attention/test_cutlass_mla_decode.py @@ -11,6 +11,8 @@ from vllm.triton_utils import triton from vllm.utils.platform_utils import num_compute_units +pytestmark = pytest.mark.skip_global_cleanup + def cal_diff( x: torch.Tensor, diff --git a/tests/kernels/attention/test_flashmla.py b/tests/kernels/attention/test_flashmla.py index 84744a81e797..3af59dd86f9a 100644 --- a/tests/kernels/attention/test_flashmla.py +++ b/tests/kernels/attention/test_flashmla.py @@ -16,6 +16,8 @@ is_flashmla_dense_supported, ) +pytestmark = pytest.mark.skip_global_cleanup + def cal_diff( x: torch.Tensor, y: torch.Tensor, name: str, use_fp8: bool = False diff --git a/tests/kernels/attention/test_merge_attn_states.py b/tests/kernels/attention/test_merge_attn_states.py index 1394ea9df405..eeadf0da0bfa 100644 --- a/tests/kernels/attention/test_merge_attn_states.py +++ b/tests/kernels/attention/test_merge_attn_states.py @@ -15,6 +15,14 @@ merge_attn_states as merge_attn_states_triton, ) +pytestmark = [ + pytest.mark.skip_global_cleanup, + pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="merge_attn_states kernels require CUDA or ROCm.", + ), +] + # Naive PyTorch Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005 # can be used to combine partial attention results (in the split-KV case) @@ -153,12 +161,6 @@ def test_merge_attn_states( input_dtype: torch.dtype, use_fp8: bool, ): - if not current_platform.is_cuda(): - pytest.skip( - "Currently only support compare triton merge_attn_states " - "with custom cuda merge_attn_states kernel" - ) - NUM_TOKENS = num_tokens NUM_HEADS = num_query_heads HEAD_SIZE = head_size diff --git a/tests/kernels/attention/test_prefix_prefill.py b/tests/kernels/attention/test_prefix_prefill.py index e2e9c4e8c201..e6d46d6efc4a 100644 --- a/tests/kernels/attention/test_prefix_prefill.py +++ b/tests/kernels/attention/test_prefix_prefill.py @@ -17,6 +17,8 @@ ) from vllm.v1.attention.ops.prefix_prefill import context_attention_fwd +pytestmark = pytest.mark.skip_global_cleanup + NUM_HEADS = [64] NUM_QUERIES_PER_KV = [1, 64] HEAD_SIZES = [24, 128] diff --git a/tests/kernels/attention/test_triton_decode_attention.py b/tests/kernels/attention/test_triton_decode_attention.py index b4b17d9b5ce8..5ac1d6cdcbf8 100644 --- a/tests/kernels/attention/test_triton_decode_attention.py +++ b/tests/kernels/attention/test_triton_decode_attention.py @@ -8,6 +8,8 @@ from vllm.utils.math_utils import cdiv from vllm.v1.attention.ops.triton_decode_attention import decode_attention_fwd +pytestmark = pytest.mark.skip_global_cleanup + DEVICE_TYPE = current_platform.device_type diff --git a/tests/kernels/attention/test_triton_unified_attention.py b/tests/kernels/attention/test_triton_unified_attention.py index d3435ea665db..1c1d4f803ff1 100644 --- a/tests/kernels/attention/test_triton_unified_attention.py +++ b/tests/kernels/attention/test_triton_unified_attention.py @@ -11,6 +11,8 @@ from vllm.v1.attention.ops.triton_unified_attention import unified_attention from vllm.v1.kv_cache_interface import KVQuantMode +pytestmark = pytest.mark.skip_global_cleanup + DEVICE_TYPE = current_platform.device_type NUM_HEADS = [(4, 4), (8, 2), (5, 1)] diff --git a/tests/kernels/attention/test_triton_unified_attention_diffkv.py b/tests/kernels/attention/test_triton_unified_attention_diffkv.py index 1a19cf34379c..513811a09c04 100644 --- a/tests/kernels/attention/test_triton_unified_attention_diffkv.py +++ b/tests/kernels/attention/test_triton_unified_attention_diffkv.py @@ -21,6 +21,8 @@ unified_attention_diffkv, ) +pytestmark = pytest.mark.skip_global_cleanup + DEVICE_TYPE = current_platform.device_type # (num_query_heads, num_kv_heads): MHA, GQA, and the num_kv_heads==1 From d3fafe0c27f9666a06675858738aaeab949da0f5 Mon Sep 17 00:00:00 2001 From: xuebwang-amd Date: Wed, 19 Aug 2026 05:42:45 +0800 Subject: [PATCH 109/839] [Quantization] Remove the dead ocp_mx_scheme branch from moe_kernel_quantize_input (#52603) Signed-off-by: xuebwang-amd Co-authored-by: Claude Opus 5 --- vllm/model_executor/layers/fused_moe/utils.py | 16 ---------------- 1 file changed, 16 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/utils.py b/vllm/model_executor/layers/fused_moe/utils.py index b873d8321d8b..7a4c7fae5abd 100644 --- a/vllm/model_executor/layers/fused_moe/utils.py +++ b/vllm/model_executor/layers/fused_moe/utils.py @@ -365,25 +365,9 @@ def moe_kernel_quantize_input( per_act_token_quant: bool, block_shape: list[int] | None = None, is_scale_swizzled: bool = True, - ocp_mx_scheme: str | None = None, quantization_emulation: bool = False, mx_alignment: int = 0, ) -> tuple[torch.Tensor, torch.Tensor | None]: - # Handle OCP MX scheme that requires QDQ (quantize-dequantize) for emulation - if ocp_mx_scheme is not None: - if ocp_mx_scheme in {"w_mxfp4", "w_mxfp4_a_mxfp4"}: - pass # No QDQ needed for these schemes - elif ocp_mx_scheme.endswith("a_fp8"): - # Perform QDQ (quantize and dequantize) on activation for emulation - # purpose, because there is no native kernel for weight in ocp_mx_scheme - # and activation in FP8. The implementation is based on existing - # non-emulation ops. - # TODO: Remove this `ocp_mx_scheme is not None` block and rely solely - # on `quantization_emulation`. - return _fp8_quantize_dequantize(A, A_scale) - # else: For other schemes (e.g., *_a_mxfp6_e3m2, *_a_mxfp6_e2m3), - # weights are already dequantized, and we proceed with normal - # activation quantization below. if quant_dtype == current_platform.fp8_dtype(): if quantization_emulation: return _fp8_quantize_dequantize(A, A_scale) From 5f7a20b3162ef88d531ada03aa1174643bb97c11 Mon Sep 17 00:00:00 2001 From: Summer Yang Date: Tue, 18 Aug 2026 14:59:08 -0700 Subject: [PATCH 110/839] [nv] add pcp support in dsv3.2 (#52046) Signed-off-by: Summer Yang --- .../test_fused_deepseek_v32_norm_rope.py | 96 +++++++- .../layers/test_mla_short_prefill_indexer.py | 45 ++++ .../layers/sparse_attn_indexer.py | 12 +- vllm/models/deepseek_v32/attention.py | 123 ++++++++-- vllm/models/deepseek_v32/common/kernels.py | 215 +++++++++++------- .../v1/attention/ops/rocm_aiter_mla_sparse.py | 4 +- 6 files changed, 376 insertions(+), 119 deletions(-) diff --git a/tests/kernels/test_fused_deepseek_v32_norm_rope.py b/tests/kernels/test_fused_deepseek_v32_norm_rope.py index d72c1b31c4d7..7a00d0c3d644 100644 --- a/tests/kernels/test_fused_deepseek_v32_norm_rope.py +++ b/tests/kernels/test_fused_deepseek_v32_norm_rope.py @@ -292,6 +292,81 @@ def test_fused_norm_rope_no_indexer(num_tokens: int): assert (topk == 7).all(), "topk buffer should be untouched on shared layer" +@pytest.mark.parametrize("has_indexer", [False, True]) +def test_fused_norm_rope_materializes_pcp_cache_inputs(has_indexer: bool): + """PCP gets local normalized/rotated K rows without direct cache writes.""" + torch.manual_seed(6) + dev = "cuda" + num_tokens = 17 + max_pos = 8192 + pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) + q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16) + kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16) + k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16) + qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16) + kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16) + ik = ( + torch.randn(num_tokens, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16) + if has_indexer + else None + ) + ikw = ( + torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32) + if has_indexer + else None + ) + ikb = ( + torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32) + if has_indexer + else None + ) + mla_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev) + idx_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev) if has_indexer else None + q_out = torch.empty_like(q_c) + kv_out = torch.empty_like(kv_c) + kpe_out = torch.empty_like(k_pe) + ik_out = torch.empty_like(ik) if ik is not None else None + topk = torch.full((num_tokens, 2048), 7, device=dev, dtype=torch.int32) + + actual_q = K.fused_norm_rope( + pos, + q_c, + qw, + EPS, + kv_c, + kvw, + EPS, + k_pe, + mla_cos_sin, + ik, + ikw, + ikb, + EPS, + idx_cos_sin, + topk, + has_indexer=has_indexer, + index_rope_interleave=True, + q_c_out=q_out, + kv_c_out=kv_out, + k_pe_out=kpe_out, + index_k_out=ik_out, + ) + + assert actual_q.data_ptr() == q_out.data_ptr() + assert_bf16(actual_q, rms_norm(q_c, qw), "PCP q norm") + assert_bf16(kv_out, rms_norm(kv_c, kvw), "PCP kv norm") + assert_bf16( + kpe_out, + rope(k_pe.float(), pos, mla_cos_sin, interleave=True), + "PCP k_pe RoPE", + ) + if has_indexer: + assert ik is not None and ikw is not None and ikb is not None + assert ik_out is not None and idx_cos_sin is not None + ik_ref = rope(layer_norm(ik, ikw, ikb), pos, idx_cos_sin, interleave=True) + assert_bf16(ik_out, ik_ref, "PCP indexer-K") + + @pytest.mark.parametrize("num_tokens", [1, 4, 17, 512]) def test_fused_norm_rope_ds_mla(num_tokens: int): """fp8_ds_mla MLA cache layout (FlashMLA sparse, bf16-query path; SM90/SM100). @@ -405,19 +480,29 @@ def test_fused_norm_rope_supports_large_token_count(): # ── fused_q ────────────────────────────────────────────────────────────────── -@pytest.mark.parametrize("num_tokens", [1, 4, 17, 512, 4096]) +@pytest.mark.parametrize( + ("num_tokens", "num_q_heads"), + [ + (1, NUM_HEADS), + (4, NUM_HEADS), + (17, NUM_HEADS), + (512, NUM_HEADS), + (4096, NUM_HEADS), + (17, 64), + ], +) @pytest.mark.parametrize("index_interleave", [True, False]) -def test_fused_q(num_tokens: int, index_interleave: bool): +def test_fused_q(num_tokens: int, num_q_heads: int, index_interleave: bool): torch.manual_seed(2) dev = "cuda" max_pos = 8192 pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos q_pe = torch.randn( - num_tokens, NUM_HEADS, ROPE_DIM, device=dev, dtype=torch.bfloat16 + num_tokens, num_q_heads, ROPE_DIM, device=dev, dtype=torch.bfloat16 ) ql_nope = torch.randn( - num_tokens, NUM_HEADS, KV_LORA, device=dev, dtype=torch.bfloat16 + num_tokens, num_q_heads, KV_LORA, device=dev, dtype=torch.bfloat16 ) index_q = torch.randn( num_tokens, INDEX_HEADS, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16 @@ -449,7 +534,7 @@ def test_fused_q(num_tokens: int, index_interleave: bool): mqa_nope_ref = (ql_nope.float() / s).to(FP8) qpe_ref = rope( q_pe.float(), - pos.unsqueeze(-1).expand(num_tokens, NUM_HEADS), + pos.unsqueeze(-1).expand(num_tokens, num_q_heads), q_cos_sin, interleave=True, ) @@ -484,7 +569,6 @@ def test_fused_q_no_indexer(num_tokens: int): ) q_scale = torch.tensor([0.5], device=dev, dtype=torch.float32) q_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev) - _, _, mqa = K.fused_q( pos, q_pe, diff --git a/tests/model_executor/layers/test_mla_short_prefill_indexer.py b/tests/model_executor/layers/test_mla_short_prefill_indexer.py index 47bfcd0b38d9..e798d8c7f5dc 100644 --- a/tests/model_executor/layers/test_mla_short_prefill_indexer.py +++ b/tests/model_executor/layers/test_mla_short_prefill_indexer.py @@ -177,3 +177,48 @@ def run_indexer(): # K cache is always updated before the scoring decision. torch.testing.assert_close(observed["k"], k[: slot_mapping.numel()]) assert observed["slots"] is slot_mapping + + +def test_skipped_k_cache_insert_accepts_no_k( + monkeypatch: pytest.MonkeyPatch, +) -> None: + indexer_metadata = make_indexer_metadata( + num_prefills=0, + num_prefill_tokens=0, + slot_mapping=torch.empty(0, dtype=torch.long), + ) + monkeypatch.setattr( + sparse_indexer, + "get_forward_context", + lambda: SimpleNamespace( + attn_metadata={INDEXER_LAYER: indexer_metadata}, + cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE, + ), + ) + monkeypatch.setattr( + sparse_indexer.current_platform, "fp8_dtype", lambda: torch.float16 + ) + + topk_indices = torch.full((1, 2048), 17, dtype=torch.int32) + result = sparse_indexer.sparse_attn_indexer( + torch.empty(1, 1), + INDEXER_LAYER, + torch.empty(1), + torch.empty(1, 1), + None, + None, + torch.empty(1, 1), + 128, + "ue8m0", + 2048, + 4, + 4096, + 4096, + topk_indices, + True, + False, + "", + ) + + assert result is topk_indices + assert torch.all(topk_indices == -1) diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index d94b377dbe58..4f0e00cfabff 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -299,7 +299,7 @@ def sparse_attn_indexer( kv_cache: torch.Tensor, q_quant: torch.Tensor, q_scale: torch.Tensor | None, - k: torch.Tensor, + k: torch.Tensor | None, weights: torch.Tensor, quant_block_size: int, scale_fmt: str | None, @@ -695,7 +695,7 @@ def sparse_attn_indexer_fake( kv_cache: torch.Tensor, q_quant: torch.Tensor, q_scale: torch.Tensor | None, - k: torch.Tensor, + k: torch.Tensor | None, weights: torch.Tensor, quant_block_size: int, scale_fmt: str | None, @@ -781,7 +781,7 @@ def forward_native( self, hidden_states: torch.Tensor, q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor], - k: torch.Tensor, + k: torch.Tensor | None, weights: torch.Tensor, ): if current_platform.is_cuda() or current_platform.is_xpu(): @@ -798,7 +798,7 @@ def forward_cuda( self, hidden_states: torch.Tensor, q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor], - k: torch.Tensor, + k: torch.Tensor | None, weights: torch.Tensor, ): # FP8 path: single tensor (per-token scale is folded into `weights`). @@ -835,7 +835,7 @@ def forward_xpu( self, hidden_states: torch.Tensor, q_fp8: torch.Tensor, - k: torch.Tensor, + k: torch.Tensor | None, weights: torch.Tensor, ): return self.forward_cuda(hidden_states, q_fp8, k, weights) @@ -844,7 +844,7 @@ def forward_hip( self, hidden_states: torch.Tensor, q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor], - k: torch.Tensor, + k: torch.Tensor | None, weights: torch.Tensor, ): assert not self.use_fp4_cache, "AMD platform doesn't support fp4 cache yet" diff --git a/vllm/models/deepseek_v32/attention.py b/vllm/models/deepseek_v32/attention.py index cc270619895b..253971b7bd05 100644 --- a/vllm/models/deepseek_v32/attention.py +++ b/vllm/models/deepseek_v32/attention.py @@ -1,5 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import TYPE_CHECKING, cast + import torch import torch.nn as nn from transformers import DeepseekV2Config, DeepseekV3Config @@ -7,8 +9,14 @@ from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.config import CacheConfig, VllmConfig from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.distributed.parallel_state import get_tp_group from vllm.forward_context import get_forward_context from vllm.model_executor.layers.attention import MLAAttention +from vllm.model_executor.layers.attention.attention import get_attention_context +from vllm.model_executor.layers.attention.pcp import ( + finalize_mla_pcp_decode, + maybe_gather_mla_latent_cache_inputs, +) from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm from vllm.model_executor.layers.linear import ( ColumnParallelLinear, @@ -35,6 +43,9 @@ from vllm.platforms import current_platform from vllm.utils.torch_utils import is_quantized_kv_cache +if TYPE_CHECKING: + from vllm.model_executor.layers.attention.mla_attention import MLACommonMetadata + class DeepseekV32Indexer(nn.Module): indexer_cache_cls = DeepseekV32IndexerCache @@ -363,6 +374,7 @@ def forward( # type: ignore[override] attn_metadata = attn_metadata_raw[0].get(self.layer_name) else: attn_metadata = attn_metadata_raw + attn_metadata = cast("MLACommonMetadata | None", attn_metadata) slot_mapping = forward_context.slot_mapping assert isinstance(slot_mapping, dict) @@ -374,7 +386,8 @@ def forward( # type: ignore[override] indexer_k_norm_bias = self.indexer.k_norm.bias indexer_k_norm_eps = self.indexer.k_norm.eps indexer_k_rope_cos_sin_cache = self.indexer_rope_emb.cos_sin_cache - indexer_k_cache = self.indexer.k_cache.kv_cache + indexer_k_cache = None if self.use_pcp else self.indexer.k_cache.kv_cache + index_k_out = torch.empty_like(index_k) if self.use_pcp else None indexer_softmax_scale = self.indexer.softmax_scale indexer_n_head_scale = self.indexer.n_head**-0.5 else: @@ -384,10 +397,11 @@ def forward( # type: ignore[override] indexer_k_norm_eps = 1e-6 indexer_k_rope_cos_sin_cache = None indexer_k_cache = None + index_k_out = None indexer_softmax_scale = 0.0 indexer_n_head_scale = 0.0 - if attn_metadata is None: + if attn_metadata is None or self.use_pcp: mla_kv_cache = None mla_k_scale = None indexer_k_cache = None @@ -396,6 +410,8 @@ def forward( # type: ignore[override] mla_kv_cache = self.kv_cache mla_k_scale = self._k_scale + kv_c_out = torch.empty_like(kv_c) if self.use_pcp else None + k_pe_out = torch.empty_like(k_pe) if self.use_pcp else None q_c = fused_norm_rope( positions, q_c, @@ -419,12 +435,14 @@ def forward( # type: ignore[override] mla_k_scale=mla_k_scale, has_indexer=has_indexer, index_rope_interleave=self._index_rope_interleave, + kv_c_out=kv_c_out, + k_pe_out=k_pe_out, + index_k_out=index_k_out, ) q = self.q_b_proj(q_c)[0].view(-1, self.num_local_heads, self.qk_head_dim) q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) - q_nope = q_nope.transpose(0, 1) - ql_nope = torch.bmm(q_nope, self.W_UK_T).transpose(0, 1) + ql_nope = torch.bmm(q_nope.transpose(0, 1), self.W_UK_T).transpose(0, 1) if self.indexer is not None and not self.skip_topk: index_q = self.indexer.wq_b(q_c)[0] @@ -451,7 +469,10 @@ def forward( # type: ignore[override] self._sparse_indexer_and_attn( q_c, index_q_fp8, + index_k_out, index_weights_out, + kv_c_out, + k_pe_out, ql_nope, mqa_q, output, @@ -463,19 +484,26 @@ def _sparse_indexer_and_attn( self, q_c: torch.Tensor, index_q_fp8: torch.Tensor | None, + index_k: torch.Tensor | None, index_weights_out: torch.Tensor | None, + kv_c: torch.Tensor | None, + k_pe: torch.Tensor | None, ql_nope: torch.Tensor, mqa_q: torch.Tensor, output: torch.Tensor, ) -> None: if self.indexer is not None and not self.skip_topk: + assert index_q_fp8 is not None + assert index_weights_out is not None + if self.use_pcp: + assert index_k is not None sparse_attn_indexer( q_c, self.indexer.k_cache.prefix, self.indexer.k_cache.kv_cache, index_q_fp8, - None, # q_scale folded into weights on the fp8 path - None, # k unused when skip_k_cache_insert=True + None, + index_k, index_weights_out, self.indexer.quant_block_size, self.indexer.scale_fmt, @@ -484,28 +512,56 @@ def _sparse_indexer_and_attn( self.indexer.max_model_len, self.indexer.max_total_seq_len, self.topk_indices_buffer, - skip_k_cache_insert=True, - use_pcp=False, + skip_k_cache_insert=not self.use_pcp, + use_pcp=self.use_pcp, dense_mha_metadata_layer_name=self._dense_mha_metadata_layer_name, - use_fp4_cache=False, - # fused_norm_rope already cleared the topk buffer this forward. + dcp_rank=( + self.dcp_manager.group.rank_in_group + if self.dcp_manager is not None + else 0 + ), + dcp_world_size=( + self._vllm_config.parallel_config.decode_context_parallel_size + ), + cp_kv_cache_interleave_size=( + self._vllm_config.parallel_config.cp_kv_cache_interleave_size + ), skip_topk_buffer_clear=True, ) - attn_metadata_raw = get_forward_context().attn_metadata - if isinstance(attn_metadata_raw, dict): - attn_metadata = attn_metadata_raw.get(self.layer_name) - elif isinstance(attn_metadata_raw, list): - attn_metadata = attn_metadata_raw[0].get(self.layer_name) - else: - attn_metadata = attn_metadata_raw - + attn_metadata, _, kv_cache, layer_slot_mapping = get_attention_context( + self.layer_name + ) if attn_metadata is None: output.zero_() return + attn_metadata = cast("MLACommonMetadata", attn_metadata) + + if self.use_pcp: + assert kv_c is not None and k_pe is not None + kv_for_cache, kpe_for_cache, cache_slot_mapping = ( + maybe_gather_mla_latent_cache_inputs( + kv_c, + k_pe.unsqueeze(1), + layer_slot_mapping, + attn_metadata.num_decode_tokens, + True, + ) + ) + self.impl.do_kv_cache_update( # type: ignore[attr-defined] + kv_for_cache, + kpe_for_cache, + kv_cache, + cache_slot_mapping, + self.kv_cache_dtype, + self._k_scale, + ) num_actual = attn_metadata.num_actual_tokens # type: ignore[attr-defined] - kv_cache = self.kv_cache + if num_actual == 0: + output.zero_() + return + if self._fp8_kv_needs_view: kv_cache = kv_cache.view(torch.float8_e4m3fn) if self._fp8_query: @@ -515,10 +571,35 @@ def _sparse_indexer_and_attn( ] else: mqa_q_arg = (ql_nope[:num_actual], mqa_q[:num_actual]) - attn_out, _ = self.impl.forward_mqa( # type: ignore[attr-defined] + + if self.use_pcp and self.impl.dcp_world_size > self.impl.pcp_world_size: + if isinstance(mqa_q_arg, tuple): + mqa_q_arg = torch.cat(mqa_q_arg, dim=-1) + mqa_q_arg = get_tp_group().all_gather(mqa_q_arg, dim=1) + attn_out, lse = self.impl.forward_mqa( # type: ignore[attr-defined] mqa_q_arg, kv_cache, attn_metadata, self ) + if self.use_pcp and self.impl.dcp_world_size > 1: + assert lse is not None and self.dcp_manager is not None + seq_lens = ( + attn_metadata.decode.seq_lens + if attn_metadata.decode is not None + else cast(torch.Tensor, attn_metadata.seq_lens)[ # type: ignore[attr-defined] + : attn_metadata.num_decodes + ] + ) + query_start_loc = attn_metadata.query_start_loc[ + : attn_metadata.num_decodes + 1 + ] + attn_out = self.dcp_manager.combine( + attn_out, + lse, + seq_lens=seq_lens, + query_start_loc=query_start_loc, + ) + attn_out = finalize_mla_pcp_decode(attn_out, self.num_heads) + # NOTE(woosuk): While the below does not need to be in the eager region, # we put it here to avoid copying the attention output. Move this back to the # captured region once forward_mqa supports `out` argument. @@ -531,3 +612,5 @@ def _sparse_indexer_and_attn( .transpose(0, 1) ) torch.bmm(x, self.W_UV, out=out) + if self.use_pcp and num_actual < output.shape[0]: + output[num_actual:].zero_() diff --git a/vllm/models/deepseek_v32/common/kernels.py b/vllm/models/deepseek_v32/common/kernels.py index ec642409cbdb..c5c6c97d55a9 100644 --- a/vllm/models/deepseek_v32/common/kernels.py +++ b/vllm/models/deepseek_v32/common/kernels.py @@ -103,12 +103,16 @@ def _fused_norm_rope_kernel( kv_stride, kv_rms_norm_w_ptr, kv_rms_eps, + kv_out_ptr, + kv_out_stride, KV_DIM: tl.constexpr, # KV RoPE kpe_ptr, kpe_stride, kpe_rope_cos_sin_cache_ptr, kpe_rope_cos_sin_cache_stride, + kpe_out_ptr, + kpe_out_stride, KPE_HALF_ROT_DIM: tl.constexpr, # Index K layer norm index_k_ptr, @@ -121,6 +125,8 @@ def _fused_norm_rope_kernel( # Index K RoPE index_k_rope_cos_sin_cache_ptr, index_k_rope_cos_sin_cache_stride, + index_k_out_ptr, + index_k_out_stride, INDEX_K_HALF_ROT_DIM: tl.constexpr, # Cache params (shared by indexer K and MLA) slot_mapping_ptr, @@ -174,10 +180,9 @@ def _fused_norm_rope_kernel( return if slot_mapping_ptr is None: - # Memory profiling run. - return - slot_idx = tl.load(slot_mapping_ptr + tok_idx) - if slot_idx < 0: + if kv_out_ptr is None and kpe_out_ptr is None and index_k_out_ptr is None: + return + elif tl.load(slot_mapping_ptr + tok_idx) < 0: # Padding return @@ -217,63 +222,81 @@ def _fused_norm_rope_kernel( r1 = x1 * cos - x2 * sin r2 = x2 * cos + x1 * sin - # MLA concat_and_cache: write [kv_c_normed, k_pe_roped] to cache. - if mla_cache_entry_stride == 0: - return - - mla_block_size = mla_cache_block_stride // mla_cache_entry_stride - mla_block_idx = slot_idx // mla_block_size - mla_block_off = slot_idx % mla_block_size - - if MLA_CACHE_DS_MLA: - # fp8_ds_mla layout (DeepSeek-V3.2, KV_DIM == 512): per-128-element - # tile of the NoPE is dynamically quantized to fp8 with its own - # float32 scale; the RoPE tail is stored unquantized in bf16. - # bytes [0, KV_DIM) : KV_DIM fp8 NoPE values - # bytes [KV_DIM, KV_DIM + 16) : MLA_NUM_TILES float32 scales - # bytes [KV_DIM + 16, ...) : 2 * KPE_HALF_ROT_DIM bf16 RoPE - # mla_cache_block_stride / mla_cache_entry_stride are byte strides - # (mla_cache_ptr is the 1-byte fp8 view of the uint8 cache). - byte_base = ( - mla_block_idx * mla_cache_block_stride - + mla_block_off * mla_cache_entry_stride + # PCP materializes K for the cross-rank gather, then inserts it into cache. + if kv_out_ptr is not None: + tl.store(kv_out_ptr + tok_idx * kv_out_stride + kv_block, kv_c) + if kpe_out_ptr is not None: + tl.store( + kpe_out_ptr + tok_idx * kpe_out_stride + dim_off * 2, + r1.to(kpe_out_ptr.dtype.element_ty), ) - kv_2d = tl.reshape(kv_c, (MLA_NUM_TILES, MLA_TILE_DIM)) - tile_amax = tl.max(tl.abs(kv_2d), axis=1, keep_dims=True) - # scale = amax / 448 (fp8 e4m3 max), matching the reference - # concat_and_cache_ds_mla kernel; floored to FLT_MIN. - tile_scale = tl.maximum(tile_amax * (1.0 / 448.0), 1.1754944e-38) - kv_c_fp8 = tl.reshape((kv_2d / tile_scale).to(tl.float8e4nv), (KV_DIM,)) - tl.store(mla_cache_ptr + byte_base + kv_block, kv_c_fp8) - tile_off = tl.arange(0, MLA_NUM_TILES) tl.store( - mla_cache_ds_scale_ptr + byte_base // 4 + KV_DIM // 4 + tile_off, - tl.reshape(tile_scale, (MLA_NUM_TILES,)), + kpe_out_ptr + tok_idx * kpe_out_stride + dim_off * 2 + 1, + r2.to(kpe_out_ptr.dtype.element_ty), ) - rope_dst = mla_cache_ds_rope_ptr + byte_base // 2 + (KV_DIM // 2 + 8) - tl.store(rope_dst + dim_off * 2, r1.to(tl.bfloat16)) - tl.store(rope_dst + dim_off * 2 + 1, r2.to(tl.bfloat16)) - return - dst = ( - mla_cache_ptr - + mla_block_idx * mla_cache_block_stride - + mla_block_off * mla_cache_entry_stride - ) - # kv_c_normed (KV_DIM elements) - if MLA_CACHE_FP8: - scale = tl.load(mla_cache_scale_ptr) - kv_c_fp8 = (kv_c.to(tl.float32) / scale).to(tl.float8e4nv) - tl.store(dst + kv_block, kv_c_fp8) - else: - tl.store(dst + kv_block, kv_c) - # k_pe_roped (from registers, interleaved layout) - if MLA_CACHE_FP8: - tl.store(dst + KV_DIM + dim_off * 2, (r1 / scale).to(tl.float8e4nv)) - tl.store(dst + KV_DIM + dim_off * 2 + 1, (r2 / scale).to(tl.float8e4nv)) - else: - tl.store(dst + KV_DIM + dim_off * 2, r1) - tl.store(dst + KV_DIM + dim_off * 2 + 1, r2) + if slot_mapping_ptr is not None: + # MLA concat_and_cache: write [kv_c_normed, k_pe_roped] to cache. + if mla_cache_entry_stride == 0: + return + + slot_idx = tl.load(slot_mapping_ptr + tok_idx) + mla_block_size = mla_cache_block_stride // mla_cache_entry_stride + mla_block_idx = slot_idx // mla_block_size + mla_block_off = slot_idx % mla_block_size + + if MLA_CACHE_DS_MLA: + # fp8_ds_mla layout (DeepSeek-V3.2, KV_DIM == 512): per-128-element + # tile of the NoPE is dynamically quantized to fp8 with its own + # float32 scale; the RoPE tail is stored unquantized in bf16. + # bytes [0, KV_DIM) : KV_DIM fp8 NoPE values + # bytes [KV_DIM, KV_DIM + 16) : MLA_NUM_TILES float32 scales + # bytes [KV_DIM + 16, ...) : 2 * KPE_HALF_ROT_DIM bf16 RoPE + # mla_cache_block_stride / mla_cache_entry_stride are byte strides + # (mla_cache_ptr is the 1-byte fp8 view of the uint8 cache). + byte_base = ( + mla_block_idx * mla_cache_block_stride + + mla_block_off * mla_cache_entry_stride + ) + kv_2d = tl.reshape(kv_c, (MLA_NUM_TILES, MLA_TILE_DIM)) + tile_amax = tl.max(tl.abs(kv_2d), axis=1, keep_dims=True) + # scale = amax / 448 (fp8 e4m3 max), matching the reference + # concat_and_cache_ds_mla kernel; floored to FLT_MIN. + tile_scale = tl.maximum(tile_amax * (1.0 / 448.0), 1.1754944e-38) + kv_c_fp8 = tl.reshape((kv_2d / tile_scale).to(tl.float8e4nv), (KV_DIM,)) + tl.store(mla_cache_ptr + byte_base + kv_block, kv_c_fp8) + tile_off = tl.arange(0, MLA_NUM_TILES) + tl.store( + mla_cache_ds_scale_ptr + byte_base // 4 + KV_DIM // 4 + tile_off, + tl.reshape(tile_scale, (MLA_NUM_TILES,)), + ) + rope_dst = mla_cache_ds_rope_ptr + byte_base // 2 + (KV_DIM // 2 + 8) + tl.store(rope_dst + dim_off * 2, r1.to(tl.bfloat16)) + tl.store(rope_dst + dim_off * 2 + 1, r2.to(tl.bfloat16)) + return + + dst = ( + mla_cache_ptr + + mla_block_idx * mla_cache_block_stride + + mla_block_off * mla_cache_entry_stride + ) + # kv_c_normed (KV_DIM elements) + if MLA_CACHE_FP8: + scale = tl.load(mla_cache_scale_ptr) + kv_c_fp8 = (kv_c.to(tl.float32) / scale).to(tl.float8e4nv) + tl.store(dst + kv_block, kv_c_fp8) + else: + tl.store(dst + kv_block, kv_c) + # k_pe_roped (from registers, interleaved layout) + if MLA_CACHE_FP8: + tl.store(dst + KV_DIM + dim_off * 2, (r1 / scale).to(tl.float8e4nv)) + tl.store( + dst + KV_DIM + dim_off * 2 + 1, + (r2 / scale).to(tl.float8e4nv), + ) + else: + tl.store(dst + KV_DIM + dim_off * 2, r1) + tl.store(dst + KV_DIM + dim_off * 2 + 1, r2) elif pid == 0: if not HAS_INDEXER: # Shared layer: no indexer K to process. @@ -354,20 +377,27 @@ def _fused_norm_rope_kernel( roped = normed * cos_full + sign * normed_partner * sin_full result = tl.where(in_rope, roped, normed) - # 3. FP8 quantize + cache write from registers. - # No need to write back to index_k_ptr — the only consumer - # (sparse_attn_indexer) reads from the cache, not index_k. - _fp8_quant_and_cache_write( - result, - index_k_mask, - slot_idx, - indexer_cache_ptr, - indexer_cache_scale_ptr, - indexer_cache_block_size, - indexer_cache_stride, - index_k_block, - INDEX_K_DIM, - ) + if index_k_out_ptr is not None: + tl.store( + index_k_out_ptr + tok_idx * index_k_out_stride + index_k_block, + result.to(index_k_out_ptr.dtype.element_ty), + mask=index_k_mask, + ) + + # PCP inserts index K after gathering; other paths write it directly. + if indexer_cache_ptr is not None and slot_mapping_ptr is not None: + slot_idx = tl.load(slot_mapping_ptr + tok_idx) + _fp8_quant_and_cache_write( + result, + index_k_mask, + slot_idx, + indexer_cache_ptr, + indexer_cache_scale_ptr, + indexer_cache_block_size, + indexer_cache_stride, + index_k_block, + INDEX_K_DIM, + ) def fused_norm_rope( @@ -395,6 +425,9 @@ def fused_norm_rope( has_indexer: bool = True, index_rope_interleave: bool = False, q_c_out: torch.Tensor | None = None, + kv_c_out: torch.Tensor | None = None, + k_pe_out: torch.Tensor | None = None, + index_k_out: torch.Tensor | None = None, ) -> torch.Tensor: assert positions.ndim == 1 assert q_c.ndim == 2 @@ -420,6 +453,10 @@ def fused_norm_rope( assert index_k_rope_cos_sin_cache is not None index_k_dim = index_k.shape[-1] topk = topk_indices_buffer.shape[-1] + if indexer_k_cache is not None or mla_kv_cache is not None: + assert slot_mapping is not None + else: + slot_mapping = None # --- Indexer K cache setup --- if indexer_k_cache is not None: @@ -430,17 +467,9 @@ def fused_norm_rope( if indexer_k_cache.dtype == torch.uint8: indexer_k_cache = indexer_k_cache.view(torch.float8_e4m3fn) else: - # No indexer cache (shared layer / MLA-only fusion). Use dummies but - # KEEP the caller's slot_mapping so the MLA write (pid 1) still runs. - idx_cache_scale_view = torch.empty(0, dtype=torch.float32, device=device) - indexer_k_cache = torch.empty(0, dtype=torch.float8_e4m3fn, device=device) + idx_cache_scale_view = None idx_cache_block_size = 1 - idx_cache_stride = 1 - if mla_kv_cache is None: - # Pure profiling run (no caches at all): skip all per-token writes. - slot_mapping = torch.full( - (num_tokens,), -1, dtype=torch.int64, device=device - ) + idx_cache_stride = 0 # --- MLA KV cache setup --- mla_cache_ds_mla = mla_kv_cache_dtype == "fp8_ds_mla" @@ -469,15 +498,26 @@ def fused_norm_rope( if mla_k_scale is None: mla_k_scale = torch.ones(1, dtype=torch.float32, device=device) else: - # Dummy values — pid 2 will skip the MLA cache write because - # slot_mapping is all -1. + # Dummy cache values; a zero entry stride disables the cache write. mla_kv_cache = torch.empty(0, dtype=torch.bfloat16, device=device) mla_block_stride = 0 mla_entry_stride = 0 - mla_k_scale = torch.ones(1, dtype=torch.float32, device=device) + mla_k_scale = _dummy((1,), torch.float32, device) if q_c_out is None: q_c_out = torch.empty_like(q_c) + kv_c_out_stride = 0 + k_pe_out_stride = 0 + index_k_out_stride = 0 + if kv_c_out is not None: + assert kv_c_out.shape == kv_c.shape + kv_c_out_stride = kv_c_out.stride(0) + if k_pe_out is not None: + assert k_pe_out.shape == k_pe.shape + k_pe_out_stride = k_pe_out.stride(0) + if index_k_out is not None: + assert index_k_out.shape == index_k.shape + index_k_out_stride = index_k_out.stride(0) use_pdl = current_platform.is_arch_support_pdl() _fused_norm_rope_kernel[(num_tokens, 4)]( positions, @@ -495,12 +535,16 @@ def fused_norm_rope( kv_c.stride(0), kv_rms_norm_w, kv_rms_eps, + kv_c_out, + kv_c_out_stride, kv_dim, # KV RoPE k_pe, k_pe.stride(0), k_rope_cos_sin_cache, k_rope_cos_sin_cache.stride(0), + k_pe_out, + k_pe_out_stride, k_rope_cos_sin_cache.shape[-1] // 2, # Index K layer norm + RoPE + FP8 quant index_k, @@ -512,6 +556,8 @@ def fused_norm_rope( triton.next_power_of_2(index_k_dim), index_k_rope_cos_sin_cache, index_k_rope_cos_sin_cache.stride(0), + index_k_out, + index_k_out_stride, index_k_rope_cos_sin_cache.shape[-1] // 2, # Cache params slot_mapping, @@ -793,7 +839,6 @@ def fused_q( assert ql_nope.ndim == 3 assert ql_nope.shape[:2] == q_pe.shape[:2] assert q_scale.dtype == torch.float32 and q_scale.numel() == 1 - num_tokens = positions.shape[0] num_q_heads = q_pe.shape[1] # Grid's 3rd dim must cover the MQA-pack heads (pid 0/2 iterate 2 heads diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 5c9ab386d4e0..0f55556c5dfb 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -682,7 +682,7 @@ def rocm_aiter_sparse_attn_indexer_fake( k_cache_prefix: LayerNameType, kv_cache: torch.Tensor, q_fp8: torch.Tensor, - k: torch.Tensor, + k: torch.Tensor | None, weights: torch.Tensor, quant_block_size: int, scale_fmt: str | None, @@ -702,7 +702,7 @@ def rocm_aiter_sparse_attn_indexer( k_cache_prefix: LayerNameType, kv_cache: torch.Tensor, q_fp8: torch.Tensor, - k: torch.Tensor, + k: torch.Tensor | None, weights: torch.Tensor, quant_block_size: int, scale_fmt: str | None, From 203926c4778698f11904beec88f6a91b94b7122d Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Tue, 18 Aug 2026 17:16:25 -0500 Subject: [PATCH 111/839] [ROCm][CI] Add AMD CI Pull-Request Commands (#52822) Signed-off-by: Andreas Karatzas --- .github/workflows/new_pr_bot.yml | 4 +- .github/workflows/run-ci-command.yml | 8 +- .github/workflows/scripts/run_ci_command.py | 207 +++++++++++++--- .../workflows/scripts/test_run_ci_command.py | 229 +++++++++++++++++- docs/contributing/README.md | 8 +- 5 files changed, 406 insertions(+), 50 deletions(-) diff --git a/.github/workflows/new_pr_bot.yml b/.github/workflows/new_pr_bot.yml index 6e65ccd93a25..bd74a761d6b2 100644 --- a/.github/workflows/new_pr_bot.yml +++ b/.github/workflows/new_pr_bot.yml @@ -80,9 +80,9 @@ jobs: '', '\u{1f4ac} Join our developer Slack at https://slack.vllm.ai to discuss your PR in `#pr-reviews`, coordinate on features in `#feat-` channels, or join special interest groups in `#sig-` channels.', '', - 'PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment `/ci run` whenever CI signals are needed.', + 'PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment `/ci run` for upstream CI or `/amd-ci run` for AMD CI only whenever CI signals are needed.', '', - 'Once the PR is approved or has the `ready` label, the PR author can also use `/ci run`, `/ci retry`, or `/ci cancel`. New commits do not start CI automatically.', + 'Once the PR is approved or has the `ready` label, the PR author can also use the corresponding `/ci run`, `/ci retry`, and `/ci cancel` commands, or their `/amd-ci` variants. New commits do not start upstream CI automatically.', '', 'If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.', '', diff --git a/.github/workflows/run-ci-command.yml b/.github/workflows/run-ci-command.yml index 6d33add7477b..4df8701a1544 100644 --- a/.github/workflows/run-ci-command.yml +++ b/.github/workflows/run-ci-command.yml @@ -21,7 +21,12 @@ jobs: github.event.comment.body == '/ci run all' || github.event.comment.body == '/ci run nightly' || github.event.comment.body == '/ci retry' || - github.event.comment.body == '/ci cancel') + github.event.comment.body == '/ci cancel' || + github.event.comment.body == '/amd-ci run' || + github.event.comment.body == '/amd-ci run all' || + github.event.comment.body == '/amd-ci run nightly' || + github.event.comment.body == '/amd-ci retry' || + github.event.comment.body == '/amd-ci cancel') runs-on: [self-hosted, linux, x64, vllm-runners] steps: - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 @@ -37,6 +42,7 @@ jobs: .github/workflows/scripts/run_ci_command.py env: BUILDKITE_API_TOKEN: ${{ secrets.BUILDKITE_API_TOKEN }} + BUILDKITE_AMD_PIPELINE: amd-ci BUILDKITE_ORGANIZATION: vllm BUILDKITE_PIPELINE: ci CI_TRUSTED_USERS: ${{ vars.CI_TRUSTED_USERS }} diff --git a/.github/workflows/scripts/run_ci_command.py b/.github/workflows/scripts/run_ci_command.py index 0dfa90cf3f48..91c0ae543445 100644 --- a/.github/workflows/scripts/run_ci_command.py +++ b/.github/workflows/scripts/run_ci_command.py @@ -17,11 +17,40 @@ COMMAND_RUN_CI_NIGHTLY = "/ci run nightly" COMMAND_RETRY_FAILED = "/ci retry" COMMAND_CANCEL_CI = "/ci cancel" +COMMAND_RUN_AMD_CI = "/amd-ci run" +COMMAND_RUN_AMD_CI_ALL = "/amd-ci run all" +COMMAND_RUN_AMD_CI_NIGHTLY = "/amd-ci run nightly" +COMMAND_RETRY_AMD_FAILED = "/amd-ci retry" +COMMAND_CANCEL_AMD_CI = "/amd-ci cancel" RUN_CI_COMMAND_ENV = { COMMAND_RUN_CI: {}, COMMAND_RUN_CI_ALL: {"RUN_ALL": "1"}, COMMAND_RUN_CI_NIGHTLY: {"RUN_ALL": "1", "NIGHTLY": "1"}, + COMMAND_RUN_AMD_CI: {}, + COMMAND_RUN_AMD_CI_ALL: {"RUN_ALL": "1"}, + COMMAND_RUN_AMD_CI_NIGHTLY: {"RUN_ALL": "1", "NIGHTLY": "1"}, } +UPSTREAM_CI_COMMANDS = frozenset( + { + COMMAND_RUN_CI, + COMMAND_RUN_CI_ALL, + COMMAND_RUN_CI_NIGHTLY, + COMMAND_RETRY_FAILED, + COMMAND_CANCEL_CI, + } +) +AMD_CI_COMMANDS = frozenset( + { + COMMAND_RUN_AMD_CI, + COMMAND_RUN_AMD_CI_ALL, + COMMAND_RUN_AMD_CI_NIGHTLY, + COMMAND_RETRY_AMD_FAILED, + COMMAND_CANCEL_AMD_CI, + } +) +RETRY_COMMANDS = frozenset({COMMAND_RETRY_FAILED, COMMAND_RETRY_AMD_FAILED}) +CANCEL_COMMANDS = frozenset({COMMAND_CANCEL_CI, COMMAND_CANCEL_AMD_CI}) +ALL_CI_COMMANDS = UPSTREAM_CI_COMMANDS | AMD_CI_COMMANDS CI_AUTHORIZED_COMMENT_MARKER = "" READY_LABELS = {"ready", "ready-run-all-tests"} TRUSTED_PERMISSIONS = {"admin", "maintain", "write"} @@ -408,11 +437,48 @@ def list_failed_jobs(self, build_number: int) -> list[dict[str, Any]]: def parse_command(body: str) -> str | None: - if body in {*RUN_CI_COMMAND_ENV, COMMAND_RETRY_FAILED, COMMAND_CANCEL_CI}: + if body in ALL_CI_COMMANDS: return body return None +def pipeline_for_command( + command: str, + *, + amd_ci_pipeline: str = "amd-ci", + upstream_ci_pipeline: str = "ci", +) -> str: + if command in AMD_CI_COMMANDS: + return amd_ci_pipeline + if command in UPSTREAM_CI_COMMANDS: + return upstream_ci_pipeline + raise ValueError(f"Unsupported CI command: {command}") + + +def ci_name_for_command(command: str) -> str: + if command in AMD_CI_COMMANDS: + return "AMD CI" + if command in UPSTREAM_CI_COMMANDS: + return "CI" + raise ValueError(f"Unsupported CI command: {command}") + + +def run_command_for_command(command: str) -> str: + if command in AMD_CI_COMMANDS: + return COMMAND_RUN_AMD_CI + if command in UPSTREAM_CI_COMMANDS: + return COMMAND_RUN_CI + raise ValueError(f"Unsupported CI command: {command}") + + +def retry_command_for_command(command: str) -> str: + if command in AMD_CI_COMMANDS: + return COMMAND_RETRY_AMD_FAILED + if command in UPSTREAM_CI_COMMANDS: + return COMMAND_RETRY_FAILED + raise ValueError(f"Unsupported CI command: {command}") + + def parse_trusted_users(value: str = "") -> set[str]: return { user.casefold() for item in value.split(",") for user in item.split() if user @@ -432,6 +498,7 @@ def authorize( actor: str, permission: str, pr: Mapping[str, Any], + run_command: str = COMMAND_RUN_CI, trusted_approval: bool = False, trusted_users: set[str] | None = None, ) -> tuple[bool, str]: @@ -454,7 +521,7 @@ def authorize( return True, "approval from a trusted reviewer" return ( False, - "A reviewer with write access must run `/ci run`, approve the PR, " + f"A reviewer with write access must run `{run_command}`, approve the PR, " "or add the `ready` label first.", ) @@ -499,6 +566,21 @@ def is_active_build(build: Mapping[str, Any]) -> bool: return bool(build.get("blocked")) or build.get("state") in ACTIVE_BUILD_STATES +def is_comment_triggered_build(build: Mapping[str, Any]) -> bool: + metadata = build.get("meta_data") or {} + return bool(metadata.get("github-comment-id")) + + +def blocks_new_run(command: str, build: Mapping[str, Any]) -> bool: + if command in AMD_CI_COMMANDS: + return ( + is_comment_triggered_build(build) + and build.get("state") in ACTIVE_BUILD_STATES + and build.get("state") != "blocked" + ) + return is_active_build(build) + + def select_latest_build( builds: Sequence[dict[str, Any]], pr_number: int, @@ -544,17 +626,21 @@ def create_retry_build_payload( *, actor: str, comment_id: int, + command: str = COMMAND_RETRY_FAILED, pr: Mapping[str, Any], source_build: Mapping[str, Any], step_keys: Sequence[str], ) -> dict[str, Any]: + if command not in RETRY_COMMANDS: + raise ValueError(f"Unsupported retry command: {command}") payload = create_build_payload( actor=actor, comment_id=comment_id, + command=run_command_for_command(command), pr=pr, ) source_number = str(source_build["number"]) - payload["message"] = f"PR #{pr['number']} {COMMAND_RETRY_FAILED} by @{actor}" + payload["message"] = f"PR #{pr['number']} {command} by @{actor}" payload["env"]["VLLM_CI_ONLY_STEP_KEYS"] = json.dumps( step_keys, separators=(",", ":") ) @@ -657,13 +743,15 @@ def notify_authorized( pr["number"], ( f"✅ @{author}, CI is now available for this PR.\n\n" - "- `/ci run` starts a CI build.\n" + "- `/ci run` starts upstream CI; `/amd-ci run` starts AMD CI only.\n" "- `/ci retry` retries failed jobs in the CI build for the current " "PR head. If the current head has no CI build, it starts a new CI " "build for the current head containing only jobs that failed in " "the latest earlier CI build for this PR.\n" + "- `/amd-ci retry` retries failed jobs in AMD CI for the current PR " + "head. Use `/amd-ci run` when the current head has no AMD CI build.\n" "- `/ci cancel` cancels scheduled or running CI builds for this PR " - "branch.\n\n" + "branch; `/amd-ci cancel` does the same for AMD CI only.\n\n" f"{CI_AUTHORIZED_COMMENT_MARKER}" ), ) @@ -712,25 +800,30 @@ def handle_run_ci( github: GitHubClient, pr: Mapping[str, Any], ) -> str: + ci_name = ci_name_for_command(command) duplicate_builds = buildkite.list_builds( pr["head"]["sha"], metadata=("github-comment-id", str(comment_id)), ) duplicate = select_latest_build(duplicate_builds, pr["number"]) if duplicate: - return f"CI was already requested by this comment: {duplicate['web_url']}" + return ( + f"{ci_name} was already requested by this comment: {duplicate['web_url']}" + ) current_builds = buildkite.list_builds(pr["head"]["sha"]) active_build = next( ( build for build in current_builds - if is_build_for_pr(build, pr["number"]) and is_active_build(build) + if is_build_for_pr(build, pr["number"]) and blocks_new_run(command, build) ), None, ) if active_build: - return f"CI is already running for this commit: {active_build['web_url']}" + return ( + f"{ci_name} is already running for this commit: {active_build['web_url']}" + ) current_pr = github.get_pr(pr["number"]) if current_pr["state"] != "open" or current_pr["head"]["sha"] != pr["head"]["sha"]: @@ -748,7 +841,7 @@ def handle_run_ci( ) ) return ( - f"Triggered [Buildkite CI #{build['number']}]({build['web_url']}) " + f"Triggered [Buildkite {ci_name} #{build['number']}]({build['web_url']}) " f"for commit `{current_pr['head']['sha'][:12]}`." ) @@ -758,15 +851,21 @@ def handle_retry_failed( actor: str, buildkite: BuildkiteClient, comment_id: int, + command: str, github: GitHubClient, pr: Mapping[str, Any], ) -> str: + ci_name = ci_name_for_command(command) + run_command = run_command_for_command(command) + retry_command = retry_command_for_command(command) builds = buildkite.list_builds(pr["head"]["sha"]) build = select_latest_build(builds, pr["number"]) if build: metadata = build.get("meta_data") or {} if str(metadata.get("github-comment-id")) == str(comment_id): - return f"CI was already requested by this comment: {build['web_url']}" + return ( + f"{ci_name} was already requested by this comment: {build['web_url']}" + ) retried = buildkite.retry_failed_jobs(build["number"], RETRY_STATES) if retried["retried_jobs_count"] == 0: @@ -776,7 +875,13 @@ def handle_retry_failed( ) return ( f"Queued {retried['retried_jobs_count']} failed job(s) for retry in " - f"[Buildkite CI #{build['number']}]({build['web_url']})." + f"[Buildkite {ci_name} #{build['number']}]({build['web_url']})." + ) + + if command == COMMAND_RETRY_AMD_FAILED: + return ( + "No AMD CI build exists for the current PR head. " + f"Use `{run_command}` first." ) previous_builds = buildkite.list_builds( @@ -790,18 +895,23 @@ def handle_retry_failed( ] source_build = select_latest_build(previous_builds, pr["number"]) if not source_build: - return "No earlier CI build exists for this PR. Use `/ci run` first." + return ( + f"No earlier {ci_name} build exists for this PR. Use `{run_command}` first." + ) if not source_build.get("finished_at") or is_active_build(source_build): - return f"The previous CI build is still running: {source_build['web_url']}" + return ( + f"The previous {ci_name} build is still running: {source_build['web_url']}" + ) failed_jobs = buildkite.list_failed_jobs(source_build["number"]) failed_script_jobs = [job for job in failed_jobs if job.get("type") == "script"] missing_step_keys = [job for job in failed_script_jobs if not job.get("step_key")] if missing_step_keys: return ( - f"[Buildkite CI #{source_build['number']}]" + f"[Buildkite {ci_name} #{source_build['number']}]" f"({source_build['web_url']}) has failed jobs without stable step " - "keys, so they cannot be retried on a new commit. Use `/ci run`." + "keys, so they cannot be retried on a new commit. " + f"Use `{run_command}`." ) failed_step_keys = {str(job["step_key"]) for job in failed_script_jobs} @@ -812,16 +922,16 @@ def handle_retry_failed( ) if setup_failures: return ( - f"[Buildkite CI #{source_build['number']}]" + f"[Buildkite {ci_name} #{source_build['number']}]" f"({source_build['web_url']}) failed during CI setup, so its test " - "failure set is incomplete. Use `/ci run` for the new commit." + f"failure set is incomplete. Use `{run_command}` for the new commit." ) step_keys = sorted(failed_step_keys) if not step_keys: return ( "No failed, timed-out, or expired jobs need retrying in " - f"[Buildkite CI #{source_build['number']}]" + f"[Buildkite {ci_name} #{source_build['number']}]" f"({source_build['web_url']})." ) @@ -829,23 +939,24 @@ def handle_retry_failed( if current_pr["state"] != "open" or current_pr["head"]["sha"] != pr["head"]["sha"]: return ( "The PR head changed while processing the command. " - "Comment `/ci retry` again." + f"Comment `{retry_command}` again." ) retry_build = buildkite.create_build( create_retry_build_payload( actor=actor, comment_id=comment_id, + command=command, pr=current_pr, source_build=source_build, step_keys=step_keys, ) ) return ( - f"Triggered [Buildkite CI #{retry_build['number']}]" + f"Triggered [Buildkite {ci_name} #{retry_build['number']}]" f"({retry_build['web_url']}) for commit " f"`{current_pr['head']['sha'][:12]}`, running {len(step_keys)} failed " - f"step(s) from [Buildkite CI #{source_build['number']}]" + f"step(s) from [Buildkite {ci_name} #{source_build['number']}]" f"({source_build['web_url']})." ) @@ -853,23 +964,36 @@ def handle_retry_failed( def handle_cancel_ci( *, buildkite: BuildkiteClient, + command: str = COMMAND_CANCEL_CI, pr: Mapping[str, Any], ) -> str: + ci_name = ci_name_for_command(command) branch = pr["head"]["ref"] - builds = buildkite.list_builds( - None, - branch=branch, - states=CANCELABLE_BUILD_STATES, - ) + branches = [branch] + states = CANCELABLE_BUILD_STATES + if command in AMD_CI_COMMANDS: + head_label = pr["head"].get("label") + if head_label and head_label not in branches: + branches.append(str(head_label)) + + builds_by_number: dict[int, dict[str, Any]] = {} + for candidate_branch in branches: + for build in buildkite.list_builds( + None, + branch=candidate_branch, + states=states, + ): + builds_by_number[build["number"]] = build + cancelable_builds = [ build - for build in builds - if build.get("branch") == branch + for build in builds_by_number.values() + if build.get("branch") in branches and is_build_for_pr(build, pr["number"]) - and build.get("state") in CANCELABLE_BUILD_STATES + and build.get("state") in states ] if not cancelable_builds: - return f"No cancelable CI build is running for branch `{branch}`." + return f"No cancelable {ci_name} build is running for branch `{branch}`." for build in cancelable_builds: buildkite.cancel_build(build["number"]) @@ -879,7 +1003,9 @@ def handle_cancel_ci( ) count = len(cancelable_builds) noun = "build" if count == 1 else "builds" - return f"Requested cancellation of {count} CI {noun} for `{branch}`: {links}." + return ( + f"Requested cancellation of {count} {ci_name} {noun} for `{branch}`: {links}." + ) def run( @@ -929,6 +1055,7 @@ def run( actor=actor, permission=permission, pr=pr, + run_command=run_command_for_command(command), trusted_approval=trusted_approval, trusted_users=trusted_users, ) @@ -949,19 +1076,23 @@ def run( github=github, pr=pr, ) - elif command == COMMAND_RETRY_FAILED: + elif command in RETRY_COMMANDS: message = handle_retry_failed( actor=actor, buildkite=buildkite, comment_id=comment_id, + command=command, github=github, pr=pr, ) - else: + elif command in CANCEL_COMMANDS: message = handle_cancel_ci( buildkite=buildkite, + command=command, pr=pr, ) + else: + raise ValueError(f"Unsupported CI command: {command}") add_reaction_safely(github, comment_id, "rocket") github.add_comment(issue_number, f"✅ {message}") except Exception: @@ -1002,13 +1133,19 @@ def main() -> None: ) return - if not parse_command(event["comment"]["body"]): + command = parse_command(event["comment"]["body"]) + if not command: return + pipeline = pipeline_for_command( + command, + amd_ci_pipeline=os.environ.get("BUILDKITE_AMD_PIPELINE", "amd-ci"), + upstream_ci_pipeline=os.environ.get("BUILDKITE_PIPELINE", "ci"), + ) buildkite = BuildkiteClient( os.environ.get("BUILDKITE_API_TOKEN", ""), os.environ.get("BUILDKITE_ORGANIZATION", "vllm"), - os.environ.get("BUILDKITE_PIPELINE", "ci"), + pipeline, ) run( event, diff --git a/.github/workflows/scripts/test_run_ci_command.py b/.github/workflows/scripts/test_run_ci_command.py index f20b22525fd6..0a3d39eda2ac 100644 --- a/.github/workflows/scripts/test_run_ci_command.py +++ b/.github/workflows/scripts/test_run_ci_command.py @@ -11,8 +11,13 @@ from run_ci_command import ( CANCELABLE_BUILD_STATES, CI_AUTHORIZED_COMMENT_MARKER, + COMMAND_CANCEL_AMD_CI, COMMAND_CANCEL_CI, + COMMAND_RETRY_AMD_FAILED, COMMAND_RETRY_FAILED, + COMMAND_RUN_AMD_CI, + COMMAND_RUN_AMD_CI_ALL, + COMMAND_RUN_AMD_CI_NIGHTLY, COMMAND_RUN_CI, COMMAND_RUN_CI_ALL, COMMAND_RUN_CI_NIGHTLY, @@ -28,6 +33,7 @@ notify_authorized, parse_command, parse_trusted_users, + pipeline_for_command, resolve_workflow_run_pr, run, select_latest_build, @@ -137,6 +143,7 @@ def __init__( self.failed_job_lists = failed_job_lists or [] self.job_list_calls: list[int] = [] self.list_calls: list[tuple[str | None, tuple[str, str] | None]] = [] + self.list_requests: list[dict[str, Any]] = [] self.retry_calls: list[tuple[int, str]] = [] self.cancel_calls: list[int] = [] @@ -149,6 +156,14 @@ def list_builds( states: tuple[str, ...] = (), ) -> list[dict[str, Any]]: self.list_calls.append((commit, metadata)) + self.list_requests.append( + { + "branch": branch, + "commit": commit, + "metadata": metadata, + "states": states, + } + ) return self.build_lists.pop(0) def create_build(self, body: dict[str, Any]) -> dict[str, Any]: @@ -291,21 +306,52 @@ def test_http_transport_does_not_retry_permission_error(self, urlopen: Any) -> N self.assertEqual(urlopen.call_count, 1) def test_only_exact_ci_commands_are_accepted(self) -> None: - self.assertEqual(parse_command(COMMAND_RUN_CI), COMMAND_RUN_CI) - self.assertEqual(parse_command(COMMAND_RUN_CI_ALL), COMMAND_RUN_CI_ALL) - self.assertEqual( - parse_command(COMMAND_RUN_CI_NIGHTLY), + commands = ( + COMMAND_RUN_CI, + COMMAND_RUN_CI_ALL, COMMAND_RUN_CI_NIGHTLY, - ) - self.assertEqual( - parse_command(COMMAND_RETRY_FAILED), COMMAND_RETRY_FAILED, + COMMAND_CANCEL_CI, + COMMAND_RUN_AMD_CI, + COMMAND_RUN_AMD_CI_ALL, + COMMAND_RUN_AMD_CI_NIGHTLY, + COMMAND_RETRY_AMD_FAILED, + COMMAND_CANCEL_AMD_CI, ) - self.assertEqual(parse_command(COMMAND_CANCEL_CI), COMMAND_CANCEL_CI) + for command in commands: + with self.subTest(command=command): + self.assertEqual(parse_command(command), command) + self.assertIsNone(parse_command("/ci run please")) self.assertIsNone(parse_command("/ci run all please")) self.assertIsNone(parse_command("/ci cancel please")) self.assertIsNone(parse_command(" /ci run")) + self.assertIsNone(parse_command("/amd-ci run please")) + self.assertIsNone(parse_command("/amd-ci retry ")) + self.assertIsNone(parse_command("/AMD-CI run")) + self.assertIsNone(parse_command("/amdci run")) + + def test_commands_select_only_their_configured_pipeline(self) -> None: + cases = ( + (COMMAND_RUN_CI, "ci"), + (COMMAND_RUN_CI_ALL, "ci"), + (COMMAND_RUN_CI_NIGHTLY, "ci"), + (COMMAND_RETRY_FAILED, "ci"), + (COMMAND_CANCEL_CI, "ci"), + (COMMAND_RUN_AMD_CI, "amd-ci"), + (COMMAND_RUN_AMD_CI_ALL, "amd-ci"), + (COMMAND_RUN_AMD_CI_NIGHTLY, "amd-ci"), + (COMMAND_RETRY_AMD_FAILED, "amd-ci"), + (COMMAND_CANCEL_AMD_CI, "amd-ci"), + ) + for command, expected_pipeline in cases: + with self.subTest(command=command): + self.assertEqual( + pipeline_for_command(command), + expected_pipeline, + ) + with self.assertRaisesRegex(ValueError, "Unsupported CI command"): + pipeline_for_command("/amd-ci run arbitrary-pipeline") def test_write_access_authorizes_reviewers_and_authors(self) -> None: allowed, _ = authorize( @@ -466,6 +512,45 @@ def test_write_reviewer_runs_ci_without_delegation(self) -> None: self.assertTrue(github.comments[0].startswith("✅ ")) self.assertIn("Buildkite CI #123", github.comments[0]) + def test_amd_run_ignores_blocked_builds(self) -> None: + for metadata in ({}, {"github-comment-id": "98"}): + with self.subTest(metadata=metadata): + github = FakeGitHub() + blocked_build = { + "blocked": True, + "created_at": "2026-08-18T01:00:00Z", + "meta_data": metadata, + "number": 122, + "pull_request": {"id": 42}, + "state": "blocked", + "web_url": "https://buildkite.example/amd-ci/builds/122", + } + buildkite = FakeBuildkite([[], [blocked_build]]) + + run(make_event(COMMAND_RUN_AMD_CI_ALL), github, buildkite) + + self.assertEqual(len(buildkite.created_builds), 1) + self.assertEqual(buildkite.created_builds[0]["env"]["RUN_ALL"], "1") + + def test_amd_run_deduplicates_comment_triggered_active_build(self) -> None: + github = FakeGitHub() + command_build = { + "blocked": False, + "created_at": "2026-08-18T01:00:00Z", + "meta_data": {"github-comment-id": "98"}, + "number": 122, + "pull_request": {"id": 42}, + "source": "api", + "state": "running", + "web_url": "https://buildkite.example/amd-ci/builds/122", + } + buildkite = FakeBuildkite([[], [command_build]]) + + run(make_event(COMMAND_RUN_AMD_CI_ALL), github, buildkite) + + self.assertEqual(buildkite.created_builds, []) + self.assertIn("AMD CI is already running", github.comments[0]) + def test_run_all_sets_buildkite_environment(self) -> None: github = FakeGitHub() buildkite = FakeBuildkite([[], []]) @@ -488,6 +573,32 @@ def test_run_nightly_sets_buildkite_environment(self) -> None: self.assertEqual(payload["env"]["RUN_ALL"], "1") self.assertEqual(payload["env"]["NIGHTLY"], "1") + def test_amd_run_variants_set_buildkite_environment(self) -> None: + cases = ( + (COMMAND_RUN_AMD_CI, {}), + (COMMAND_RUN_AMD_CI_ALL, {"RUN_ALL": "1"}), + ( + COMMAND_RUN_AMD_CI_NIGHTLY, + {"RUN_ALL": "1", "NIGHTLY": "1"}, + ), + ) + for command, expected_env in cases: + with self.subTest(command=command): + github = FakeGitHub() + buildkite = FakeBuildkite([[], []]) + + run(make_event(command), github, buildkite) + + payload = buildkite.created_builds[0] + self.assertEqual(payload["message"], f"PR #42 {command} by @reviewer") + command_env = { + key: value + for key, value in payload["env"].items() + if key in {"RUN_ALL", "NIGHTLY"} + } + self.assertEqual(command_env, expected_env) + self.assertIn("Buildkite AMD CI #123", github.comments[0]) + def test_unapproved_authors_are_denied_without_buildkite(self) -> None: github = FakeGitHub( permission="read", @@ -505,6 +616,20 @@ def test_unapproved_authors_are_denied_without_buildkite(self) -> None: run(make_event(COMMAND_RUN_CI, "author"), github, buildkite) self.assertEqual(len(github.comments), 1) + def test_unapproved_author_gets_amd_specific_guidance(self) -> None: + github = FakeGitHub( + permission="read", + pr=make_pr(), + review_decision="REVIEW_REQUIRED", + ) + buildkite = FakeBuildkite() + + run(make_event(COMMAND_RUN_AMD_CI, "author"), github, buildkite) + + self.assertEqual(buildkite.list_calls, []) + self.assertIn("`/amd-ci run`", github.comments[0]) + self.assertNotIn("`/ci run`", github.comments[0]) + def test_untrusted_approval_cannot_launch_ci(self) -> None: github = FakeGitHub( permission="read", @@ -540,9 +665,12 @@ def test_ready_label_notifies_author_once(self) -> None: self.assertEqual(len(github.comments), 1) comment = github.comments[0] self.assertTrue(comment.startswith("✅ @author")) - self.assertIn("`/ci run` starts a CI build", comment) + self.assertIn("`/ci run` starts upstream CI", comment) self.assertIn("`/ci retry` retries failed jobs", comment) self.assertIn("`/ci cancel` cancels scheduled or running", comment) + self.assertIn("`/amd-ci run` starts AMD CI only", comment) + self.assertIn("`/amd-ci retry` retries failed jobs in AMD CI", comment) + self.assertIn("`/amd-ci cancel` does the same for AMD CI only", comment) self.assertIn("CI build for the current PR head", comment) self.assertIn("only jobs that failed in the latest earlier CI build", comment) self.assertNotIn(COMMAND_RUN_CI_ALL, comment) @@ -723,6 +851,47 @@ def test_ci_retry_retries_failed_jobs_while_build_is_running(self) -> None: self.assertEqual(buildkite.retry_calls, [(123, RETRY_STATES)]) self.assertIn("Queued 3 failed job", github.comments[0]) + def test_amd_ci_retry_retries_only_the_current_head_build(self) -> None: + github = FakeGitHub( + permission="read", + pr=make_pr(labels=[{"name": "ready"}]), + ) + buildkite = FakeBuildkite( + [ + [ + { + "created_at": "2026-07-28T01:00:00Z", + "number": 321, + "pull_request": {"id": 42}, + "state": "failing", + "web_url": "https://buildkite.example/amd-ci/builds/321", + } + ] + ] + ) + + run(make_event(COMMAND_RETRY_AMD_FAILED, "author"), github, buildkite) + + self.assertEqual(buildkite.retry_calls, [(321, RETRY_STATES)]) + self.assertIn("Buildkite AMD CI #321", github.comments[0]) + + def test_amd_ci_retry_requires_a_build_for_the_current_head(self) -> None: + github = FakeGitHub( + permission="read", + pr=make_pr(labels=[{"name": "ready"}]), + ) + buildkite = FakeBuildkite([[]]) + + run(make_event(COMMAND_RETRY_AMD_FAILED, "author"), github, buildkite) + + self.assertEqual(buildkite.list_calls, [("0123456789abcdef", None)]) + self.assertEqual(buildkite.retry_calls, []) + self.assertEqual(buildkite.created_builds, []) + self.assertIn( + "No AMD CI build exists for the current PR head", github.comments[0] + ) + self.assertIn("Use `/amd-ci run`", github.comments[0]) + def test_ci_retry_creates_filtered_build_for_new_head(self) -> None: github = FakeGitHub( permission="read", @@ -926,6 +1095,48 @@ def test_ci_cancel_is_a_noop_without_active_builds(self) -> None: self.assertEqual(buildkite.cancel_calls, []) self.assertIn("No cancelable CI build is running", github.comments[0]) + def test_amd_ci_cancel_handles_command_and_fork_webhook_branches(self) -> None: + pr = make_pr() + pr["head"]["label"] = "contributor:feature" + github = FakeGitHub(pr=pr) + buildkite = FakeBuildkite( + [ + [ + { + "branch": "feature", + "number": 321, + "pull_request": {"id": 42}, + "state": "running", + "web_url": "https://buildkite.example/amd-ci/builds/321", + } + ], + [ + { + "branch": "contributor:feature", + "number": 322, + "pull_request": {"id": 42}, + "state": "failing", + "web_url": "https://buildkite.example/amd-ci/builds/322", + } + ], + ] + ) + + run(make_event(COMMAND_CANCEL_AMD_CI), github, buildkite) + + self.assertEqual(buildkite.cancel_calls, [321, 322]) + self.assertEqual( + [request["branch"] for request in buildkite.list_requests], + ["feature", "contributor:feature"], + ) + self.assertTrue( + all( + request["states"] == CANCELABLE_BUILD_STATES + for request in buildkite.list_requests + ) + ) + self.assertIn("cancellation of 2 AMD CI builds", github.comments[0]) + def test_buildkite_cancel_uses_cancel_build_endpoint(self) -> None: transport = FakeTransport({"number": 123, "state": "canceling"}) client = BuildkiteClient( diff --git a/docs/contributing/README.md b/docs/contributing/README.md index 25e7ecbdf1c7..3206094dee21 100644 --- a/docs/contributing/README.md +++ b/docs/contributing/README.md @@ -303,9 +303,11 @@ review process: clarification or discuss the suggestion. - Note that not all CI checks will be executed due to limited computational resources. Reviewers with write access and configured trusted contributors - can comment `/ci run` when CI signals are needed before a PR is ready. After - the PR is approved or has the `ready` label, the PR author can use `/ci run` - or `/ci retry`. New commits do not start CI automatically. + can comment `/ci run` for upstream CI or `/amd-ci run` for AMD CI only when + CI signals are needed before a PR is ready. After the PR is approved or has + the `ready` label, the PR author can use `/ci run`, `/ci retry`, `/ci cancel`, + or the corresponding `/amd-ci` variants. New commits do not start upstream + CI automatically. ### Pull Request Limits and Escalation From aa6abec49a7ed7e40ee4c2a07586d0d83887fa2a Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Tue, 18 Aug 2026 17:34:50 -0500 Subject: [PATCH 112/839] [CI][ROCm] Prevent Git maintenance races during shallow fetches (#52810) Signed-off-by: Andreas Karatzas Co-authored-by: OpenAI Codex --- .buildkite/scripts/ci-bake-rocm.sh | 6 ++-- .../tools/test_docker_build_metadata_args.py | 35 ++++++++++++++++++- 2 files changed, 38 insertions(+), 3 deletions(-) diff --git a/.buildkite/scripts/ci-bake-rocm.sh b/.buildkite/scripts/ci-bake-rocm.sh index e3681e214a42..614987f6d148 100644 --- a/.buildkite/scripts/ci-bake-rocm.sh +++ b/.buildkite/scripts/ci-bake-rocm.sh @@ -199,11 +199,13 @@ get_buildkite_target_repo_url() { git_fetch_with_timeout() { local timeout_secs="${ROCM_CACHE_GIT_FETCH_TIMEOUT:-60}" + local -a fetch_command=(git fetch --no-auto-maintenance) + # Detached maintenance can race a later shallow fetch on .git/shallow. if command -v timeout >/dev/null 2>&1; then - timeout "${timeout_secs}s" git fetch "$@" + timeout "${timeout_secs}s" "${fetch_command[@]}" "$@" else - git fetch "$@" + "${fetch_command[@]}" "$@" fi } diff --git a/tests/tools/test_docker_build_metadata_args.py b/tests/tools/test_docker_build_metadata_args.py index 925ed4bfc9a2..2c8d8608b19a 100644 --- a/tests/tools/test_docker_build_metadata_args.py +++ b/tests/tools/test_docker_build_metadata_args.py @@ -8,6 +8,7 @@ REPO_ROOT = Path(__file__).resolve().parents[2] HELPER = REPO_ROOT / ".buildkite" / "scripts" / "docker-build-metadata-args.sh" +ROCM_CI_BAKE = REPO_ROOT / ".buildkite" / "scripts" / "ci-bake-rocm.sh" def run_helper( @@ -165,7 +166,7 @@ def test_rocm_ci_base_bake_embeds_content_hash_label() -> None: def test_rocm_ci_base_metadata_inputs_cover_ci_base_files() -> None: - ci_bake = (REPO_ROOT / ".buildkite" / "scripts" / "ci-bake-rocm.sh").read_text() + ci_bake = ROCM_CI_BAKE.read_text() for expected in ( "requirements/common.txt", @@ -174,3 +175,35 @@ def test_rocm_ci_base_metadata_inputs_cover_ci_base_files() -> None: "docker/Dockerfile.rocm", ): assert expected in ci_bake + + +def test_rocm_git_fetch_disables_automatic_maintenance(tmp_path: Path) -> None: + fake_bin = tmp_path / "bin" + fake_bin.mkdir() + git = fake_bin / "git" + git.write_text('#!/bin/sh\nprintf "%s\\n" "$@"\n') + git.chmod(0o755) + + env = os.environ.copy() + env["PATH"] = f"{fake_bin}:{env['PATH']}" + result = subprocess.run( + [ + "bash", + "-c", + 'source "$1"; git_fetch_with_timeout --quiet origin HEAD', + "bash", + str(ROCM_CI_BAKE), + ], + check=True, + env=env, + stdout=subprocess.PIPE, + text=True, + ) + + assert result.stdout.splitlines() == [ + "fetch", + "--no-auto-maintenance", + "--quiet", + "origin", + "HEAD", + ] From 8f4a7f45c53ab52b17023d3ca804e477daa36a23 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Tue, 18 Aug 2026 17:45:47 -0500 Subject: [PATCH 113/839] [ROCm][CI] Gating more ROCm tests (#44969) Signed-off-by: Andreas Karatzas Signed-off-by: Andreas Karatzas Co-authored-by: OpenAI Codex --- .buildkite/test_areas/engine.yaml | 10 ++++++- .buildkite/test_areas/entrypoints.yaml | 24 +++++++++++++++++ .buildkite/test_areas/lora.yaml | 22 +++++++++++++++ .buildkite/test_areas/misc.yaml | 8 ++++++ .buildkite/test_areas/model_runner_v2.yaml | 21 +++++++++++++++ .buildkite/test_areas/models_distributed.yaml | 27 +++++++++++++++++++ .buildkite/test_areas/models_language.yaml | 12 +++++++++ .buildkite/test_areas/models_multimodal.yaml | 24 +++++++++++++++++ .buildkite/test_areas/plugins.yaml | 11 ++++++++ .buildkite/test_areas/pytorch.yaml | 11 ++++++++ .buildkite/test_areas/spec_decode.yaml | 8 ++++++ 11 files changed, 177 insertions(+), 1 deletion(-) diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml index d7bb76c69eff..7cf03c781602 100644 --- a/.buildkite/test_areas/engine.yaml +++ b/.buildkite/test_areas/engine.yaml @@ -133,7 +133,7 @@ steps: depends_on: - image-build-amd -- label: ":nvidia: (H100) V1 E2E" +- label: ":nvidia: (H100) V1 E2E Hybrid Chunked Prefill" key: v1-e2e-4xh100 timeout_in_minutes: 35 device: h100 @@ -145,3 +145,11 @@ steps: - tests/v1/e2e/test_hybrid_chunked_prefill.py commands: - pytest -v -s v1/e2e/test_hybrid_chunked_prefill.py + mirror: + amd: + label: ":amd: (MI300) V1 E2E Hybrid Chunked Prefill" + dind: false + device: mi300_4 + timeout_in_minutes: 35 + depends_on: + - image-build-amd diff --git a/.buildkite/test_areas/entrypoints.yaml b/.buildkite/test_areas/entrypoints.yaml index ed143b881205..56d4e012581e 100644 --- a/.buildkite/test_areas/entrypoints.yaml +++ b/.buildkite/test_areas/entrypoints.yaml @@ -14,6 +14,14 @@ steps: commands: - pytest -v -s entrypoints/unit_tests - pytest -v -s entrypoints/weight_transfer + mirror: + amd: + label: ":amd: (MI355) Entrypoints Unit" + dind: false + device: mi355_1 + timeout_in_minutes: 35 + depends_on: + - image-build-amd - label: ":nvidia: (H200) Entrypoints Integration (LLM)" device: h200_35gb @@ -144,6 +152,14 @@ steps: - tests/entrypoints/openai/responses commands: - pytest -v -s entrypoints/openai/responses + mirror: + amd: + label: ":amd: (MI355) Entrypoints Integration (Responses API)" + dind: false + device: mi355_1 + timeout_in_minutes: 50 + depends_on: + - image-build-amd - label: ":nvidia: (H200) Entrypoints Integration (Speech to Text)" device: h200_35gb @@ -170,6 +186,14 @@ steps: commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/multimodal + mirror: + amd: + label: ":amd: (MI355) Entrypoints Integration (Multimodal)" + dind: false + device: mi355_1 + timeout_in_minutes: 55 + depends_on: + - image-build-amd - label: ":nvidia: (H200) Entrypoints Integration (Pooling)" device: h200_35gb diff --git a/.buildkite/test_areas/lora.yaml b/.buildkite/test_areas/lora.yaml index 5c3c7c0e8900..55a7f54d8290 100644 --- a/.buildkite/test_areas/lora.yaml +++ b/.buildkite/test_areas/lora.yaml @@ -50,3 +50,25 @@ steps: - pytest -v -s -x lora/test_gptoss_tp.py - pytest -v -s -x lora/test_qwen35_densemodel_lora.py - pytest -v -s -x lora/test_gemma4_tp.py + mirror: + amd: + label: ":amd: (MI300) LoRA TP (Distributed)" + dind: false + device: mi300_4 + working_dir: "/vllm-workspace/tests" + timeout_in_minutes: 55 + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/lora + - vllm/model_executor/layers/fused_moe/ + - vllm/platforms/rocm.py + - tests/lora + commands: + - pytest -v -s -x lora/test_chatglm3_tp.py + - pytest -v -s -x lora/test_llama_tp.py + - pytest -v -s -x lora/test_qwen3_with_multi_loras.py + - pytest -v -s -x lora/test_olmoe_tp.py + - pytest -v -s -x lora/test_gptoss_tp.py + - pytest -v -s -x lora/test_qwen35_densemodel_lora.py + - pytest -v -s -x lora/test_gemma4_tp.py diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 750f51283abe..a73ca9bb8192 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -317,6 +317,14 @@ steps: - pytest -v -s detokenizer - pytest -v -s -m 'not cpu_test' multimodal - pytest -v -s utils_ + mirror: + amd: + label: ":amd: (MI355) Async Engine, Inputs, Utils, Worker" + dind: false + device: mi355_1 + timeout_in_minutes: 40 + depends_on: + - image-build-amd - label: ":computer: (CPU) Async Engine, Inputs, Utils, Worker, Config" key: async-engine-inputs-utils-worker-config-cpu diff --git a/.buildkite/test_areas/model_runner_v2.yaml b/.buildkite/test_areas/model_runner_v2.yaml index 1b1bef43aa79..eaec866f58d5 100644 --- a/.buildkite/test_areas/model_runner_v2.yaml +++ b/.buildkite/test_areas/model_runner_v2.yaml @@ -81,6 +81,27 @@ steps: - export NCCL_CUMEM_HOST_ENABLE=0 - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py -k "not ray" - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py + mirror: + amd: + label: ":amd: (MI300) Model Runner V2 Distributed" + dind: false + device: mi300_2 + timeout_in_minutes: 45 + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/v1/worker/gpu/ + - vllm/v1/worker/gpu_worker.py + - tests/basic_correctness/test_basic_correctness.py + - tests/v1/distributed/test_async_llm_dp.py + - tests/v1/distributed/test_eagle_dp.py + - vllm/platforms/rocm.py + commands: + - set -x + - export VLLM_USE_V2_MODEL_RUNNER=1 + - TARGET_TEST_SUITE=MI300 pytest -v -s basic_correctness/test_basic_correctness.py -m 'distributed(num_gpus=2)' -k "not ray and not True" + - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py -k "not ray" + - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py - label: ":nvidia: (L4) Model Runner V2 Pipeline Parallelism" key: model-runner-v2-pipeline-parallelism-4-gpus diff --git a/.buildkite/test_areas/models_distributed.yaml b/.buildkite/test_areas/models_distributed.yaml index 730c51b89b84..e34e539ca294 100644 --- a/.buildkite/test_areas/models_distributed.yaml +++ b/.buildkite/test_areas/models_distributed.yaml @@ -22,3 +22,30 @@ steps: - pytest models/multimodal/generation/test_phi4siglip.py -v -s -m 'distributed(num_gpus=2)' - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_phi4siglip.py - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest models/multimodal/generation/test_whisper.py -v -s -m 'distributed(num_gpus=2)' + mirror: + amd: + label: ":amd: (MI300) Distributed Models" + dind: false + device: mi300_2 + timeout_in_minutes: 75 + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/model_executor/model_loader/sharded_state_loader.py + - vllm/model_executor/models/ + - vllm/model_executor/layers/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + - tests/basic_correctness/ + - tests/model_executor/model_loader/test_sharded_state_loader.py + - tests/models/ + commands: + - TARGET_TEST_SUITE=MI300 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' + - HIP_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)' + - pytest models/transformers/test_backend.py -v -s -m 'distributed(num_gpus=2)' + - pytest models/language -v -s -m 'distributed(num_gpus=2)' + - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_phi4siglip.py + - pytest models/multimodal/generation/test_phi4siglip.py -v -s -m 'distributed(num_gpus=2)' + - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest models/multimodal/generation/test_whisper.py -v -s -m 'distributed(num_gpus=2)' diff --git a/.buildkite/test_areas/models_language.yaml b/.buildkite/test_areas/models_language.yaml index e8675141237e..12c291efe5df 100644 --- a/.buildkite/test_areas/models_language.yaml +++ b/.buildkite/test_areas/models_language.yaml @@ -114,6 +114,18 @@ steps: - MAMBA_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0' - CAUSAL_CONV1D_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)' + mirror: + amd: + label: ":amd: (MI355) Language Models (Extended Generation)" + dind: false + device: mi355_1 + timeout_in_minutes: 70 + depends_on: + - image-build-amd + commands: + - MAMBA_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr' + - CAUSAL_CONV1D_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' + - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)' - label: ":nvidia: (H200) Language Models (PPL)" key: language-models-test-ppl diff --git a/.buildkite/test_areas/models_multimodal.yaml b/.buildkite/test_areas/models_multimodal.yaml index 1d97c3f5e622..62883c341e37 100644 --- a/.buildkite/test_areas/models_multimodal.yaml +++ b/.buildkite/test_areas/models_multimodal.yaml @@ -110,6 +110,14 @@ steps: - tests/models/registry.py commands: - pytest -v -s models/multimodal/processing/test_tensor_schema.py + mirror: + amd: + label: ":amd: (MI355) Multimodal Processor" + dind: false + device: mi355_1 + timeout_in_minutes: 115 + depends_on: + - image-build-amd - label: ":nvidia: (H200) Multimodal Accuracy Eval (Small Models)" device: h200_35gb @@ -187,6 +195,14 @@ steps: - tests/models/multimodal/generation commands: - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB + mirror: + amd: + label: ":amd: (MI355) Multimodal Models (Extended Generation 2) Shard %N" + dind: false + device: mi355_1 + timeout_in_minutes: 100 + depends_on: + - image-build-amd - label: ":nvidia: (H200) Multimodal Models (Extended Generation 3)" device: h200_35gb @@ -198,6 +214,14 @@ steps: - tests/models/multimodal/generation commands: - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model' + mirror: + amd: + label: ":amd: (MI355) Multimodal Models (Extended Generation 3)" + dind: false + device: mi355_1 + timeout_in_minutes: 90 + depends_on: + - image-build-amd - label: ":nvidia: (H200) Multimodal Models (Extended Pooling)" key: multi-modal-models-extended-pooling diff --git a/.buildkite/test_areas/plugins.yaml b/.buildkite/test_areas/plugins.yaml index 2cc4cb8e8bc3..19fa0e4d1e8e 100644 --- a/.buildkite/test_areas/plugins.yaml +++ b/.buildkite/test_areas/plugins.yaml @@ -49,6 +49,17 @@ steps: - pytest -v -s plugins_tests/test_oot_registration_online.py # it needs a clean process - pytest -v -s plugins_tests/test_oot_registration_offline.py # it needs a clean process - pytest -v -s plugins_tests/lora_resolvers # unit tests for in-tree lora resolver plugins + mirror: + amd: + label: ":amd: (MI250) Plugin Integration" + device: mi250_2 + timeout_in_minutes: 50 + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/plugins/ + - tests/plugins/ + - vllm/platforms/rocm.py - label: ":nvidia: (H200) GGUF Plugin" diff --git a/.buildkite/test_areas/pytorch.yaml b/.buildkite/test_areas/pytorch.yaml index 78cebb2e596d..810d335cf062 100644 --- a/.buildkite/test_areas/pytorch.yaml +++ b/.buildkite/test_areas/pytorch.yaml @@ -42,6 +42,17 @@ steps: # However, find does not normally propagate error codes, so we combine it with xargs # (using -0 for proper path handling) - "find compile/ -maxdepth 1 -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'" + mirror: + amd: + label: ":amd: (MI300) PyTorch Compilation" + dind: false + device: mi300_1 + timeout_in_minutes: 90 + depends_on: + - image-build-amd + source_file_dependencies: + - csrc/ + - vllm/platforms/rocm.py - label: ":nvidia: (H200) PyTorch Compilation H100 Cases" key: pytorch-compilation-unit-tests-h100 diff --git a/.buildkite/test_areas/spec_decode.yaml b/.buildkite/test_areas/spec_decode.yaml index 7486a467fbbe..0f3a45838b78 100644 --- a/.buildkite/test_areas/spec_decode.yaml +++ b/.buildkite/test_areas/spec_decode.yaml @@ -223,6 +223,14 @@ steps: commands: - export VLLM_ALLOW_INSECURE_SERIALIZATION=1 - pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test + mirror: + amd: + label: ":amd: (MI300) Speculators Correctness Nightly" + dind: false + device: mi300_1 + timeout_in_minutes: 40 + depends_on: + - image-build-amd - label: ":nvidia: (B200) Spec Decode DeepSeek MTP Parallel Load" key: spec-decode-deepseek-mtp-parallel-load-2xb200-2xmi300 From ef47a897e2ad9a404cce9c9e7df15934deb8ffbe Mon Sep 17 00:00:00 2001 From: Russell Bryant Date: Tue, 18 Aug 2026 19:14:39 -0400 Subject: [PATCH 114/839] [Core] Make prefix-cache NONE_HASH deterministic by default (#51875) Signed-off-by: Russell Bryant Co-authored-by: Claude Opus 4.8 (1M context) --- docs/features/kv_offloading_usage.md | 10 ++- .../mooncake_store_connector_usage.md | 8 +- tests/v1/core/test_kv_cache_utils.py | 68 ++++++++++++++-- .../v1/kv_offload/tiering/p2p/test_manager.py | 51 ++++++++---- vllm/v1/core/kv_cache_utils.py | 77 ++++++++++++++----- vllm/v1/kv_offload/tiering/fs/manager.py | 15 ++-- vllm/v1/kv_offload/tiering/p2p/manager.py | 35 +++++---- .../kv_offload/tiering/p2p/session/client.py | 5 +- .../tiering/p2p/session/protocol.py | 7 +- .../kv_offload/tiering/p2p/session/session.py | 5 +- 10 files changed, 204 insertions(+), 77 deletions(-) diff --git a/docs/features/kv_offloading_usage.md b/docs/features/kv_offloading_usage.md index 83c98c3d019f..e2450c662549 100644 --- a/docs/features/kv_offloading_usage.md +++ b/docs/features/kv_offloading_usage.md @@ -152,10 +152,12 @@ Inside that subdirectory, blocks are sharded across hash-prefix subdirectories t #### Cross-Process Sharing -To enable KV cache sharing between multiple vLLM instances using the same `root_dir` (e.g., via a shared PVC), the `PYTHONHASHSEED` environment variable must be set to the same fixed value (e.g., `"0"`) on every instance. Without this, each process initializes `NONE_HASH` (the chain-hash seed for block content hashes) with random bytes, producing different block filenames for identical token content. +KV cache sharing between multiple vLLM instances using the same `root_dir` (e.g., via a shared PVC) works by default: `NONE_HASH` (the chain-hash seed for block content hashes) is derived from a fixed default seed, so identical token content produces identical block filenames across instances. To use a custom shared seed instead, set the `PYTHONHASHSEED` environment variable to the same value on every instance. + +The exception is the non-cryptographic `xxhash` and `xxhash_cbor` values of `--prefix-caching-hash-algo`, which seed `NONE_HASH` randomly per process so the seed stays unpredictable. Sharing a cache across instances with those algorithms requires setting the same `PYTHONHASHSEED` on every instance. ```bash -PYTHONHASHSEED=0 vllm serve ... +PYTHONHASHSEED= vllm serve ... ``` ### Object Store (OBJ) @@ -182,13 +184,13 @@ The object-store tier (`type: "obj"`) offloads blocks to an S3-compatible object | `region` | no | `""` | Bucket region, if the endpoint requires one. | | `ca_bundle` | no | `""` | CA bundle path for TLS verification. | -Object keys follow the same run-configuration digest scheme as the filesystem tier (see [On-Disk Layout](#on-disk-layout)) and are stored under the optional `prefix`. The [Cross-Process Sharing](#cross-process-sharing) requirement (`PYTHONHASHSEED`) applies to shared buckets as well, so instances sharing a bucket produce identical keys for identical content. At startup the tier probes object store connectivity and fails fast with a configuration error if the bucket is unreachable. +Object keys follow the same run-configuration digest scheme as the filesystem tier (see [On-Disk Layout](#on-disk-layout)) and are stored under the optional `prefix`. The [Cross-Process Sharing](#cross-process-sharing) behavior applies to shared buckets as well, so instances sharing a bucket produce identical keys for identical content; set a shared `PYTHONHASHSEED` if you want a custom seed. At startup the tier probes object store connectivity and fails fast with a configuration error if the bucket is unreachable. ### P2P (Including P/D) The P2P tier (`type: "p2p"`) shares completed KV blocks between vLLM instances over RDMA via NIXL. Each instance binds a control socket on `host:port` and exchanges blocks directly with peers — no shared filesystem required. -The `PYTHONHASHSEED` environment variable must be set to the same fixed value (e.g. `"0"`) on all nodes so that block content hashes match across instances (see [Cross-Process Sharing](#cross-process-sharing)). This is enforced: a P2P instance started without `PYTHONHASHSEED` set fails at startup, and each peer's value is verified during the connect handshake — a peer advertising a different `PYTHONHASHSEED` is rejected. +Block content hashes must match across instances for peers to exchange blocks (see [Cross-Process Sharing](#cross-process-sharing)). This works by default via the deterministic `NONE_HASH` seed, so setting `PYTHONHASHSEED` is optional. If you do set it, it must be the same value on all nodes. Each peer's effective seed is verified during the connect handshake — a peer advertising a different seed is rejected. With the `xxhash`/`xxhash_cbor` algorithms the seed is random per process, so `PYTHONHASHSEED` must be set on every peer or the handshake rejects them. | Key | Required | Default | Notes | | --- | --- | --- | --- | diff --git a/docs/features/mooncake_store_connector_usage.md b/docs/features/mooncake_store_connector_usage.md index e2470bf8fc2b..a07716dea627 100644 --- a/docs/features/mooncake_store_connector_usage.md +++ b/docs/features/mooncake_store_connector_usage.md @@ -235,10 +235,12 @@ Strict isolation requires a Mooncake master started with `--enable_multi_tenants ### Reproducible Block Hashes Across Processes -The `MooncakeStoreConnector` relies on consistent block hashes across all vLLM processes sharing the distributed store. Because Python randomizes its hash seed per process by default, identical prompts can produce different block hashes on different processes — preventing cross-process prefix cache hits. +The `MooncakeStoreConnector` relies on consistent block hashes across all vLLM processes sharing the distributed store. Block hashes chain from `NONE_HASH`, which is derived from a fixed default seed, so identical prompts produce identical block hashes across processes by default — enabling cross-process prefix cache hits without extra configuration. -Set a fixed `PYTHONHASHSEED` on every instance that shares the store (DP ranks, separate prefiller/decoder nodes, and any other vLLM process pointed at the same Mooncake store): +The exception is the non-cryptographic `xxhash`/`xxhash_cbor` values of `--prefix-caching-hash-algo`, which seed `NONE_HASH` randomly per process; sharing a store with those requires `PYTHONHASHSEED`. + +To use a custom shared seed, set the same `PYTHONHASHSEED` on every instance that shares the store (DP ranks, separate prefiller/decoder nodes, and any other vLLM process pointed at the same Mooncake store): ```bash -PYTHONHASHSEED=0 vllm serve ... +PYTHONHASHSEED= vllm serve ... ``` diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index bfe4996b1aa8..829f45608a48 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -20,7 +20,7 @@ PlaceholderRange, ) from vllm.sampling_params import SamplingParams -from vllm.utils.hashing import sha256, sha256_cbor +from vllm.utils.hashing import sha256, sha256_cbor, xxhash, xxhash_cbor from vllm.utils.mem_constants import GiB_bytes from vllm.v1.core.kv_cache_manager import KVCacheManager from vllm.v1.core.kv_cache_utils import ( @@ -199,14 +199,20 @@ def new_mamba_spec( def test_none_hash(monkeypatch, hash_fn): import vllm.v1.core.kv_cache_utils - # case 1: PYTHONHASHSEED is not set, use random + # case 1: PYTHONHASHSEED is not set -> deterministic default seed so that + # independent processes compute identical block hashes for identical + # content (e.g. for KV cache reuse across nodes). with monkeypatch.context() as m: m.delenv("PYTHONHASHSEED", raising=False) reloaded_kv_cache_utils = importlib.reload(vllm.v1.core.kv_cache_utils) reloaded_kv_cache_utils.init_none_hash(hash_fn) - assert reloaded_kv_cache_utils.NONE_HASH is not None - assert isinstance(reloaded_kv_cache_utils.NONE_HASH, bytes) - assert reloaded_kv_cache_utils.NONE_HASH != b"" + none_hash = reloaded_kv_cache_utils.NONE_HASH + assert isinstance(none_hash, bytes) + assert none_hash != b"" + assert none_hash == hash_fn(reloaded_kv_cache_utils.DEFAULT_NONE_HASH_SEED) + # deterministic across re-initialization within the same environment + reloaded_kv_cache_utils.init_none_hash(hash_fn) + assert none_hash == reloaded_kv_cache_utils.NONE_HASH # case 2: PYTHONHASHSEED is set, use the seed and hash_fn with monkeypatch.context() as m: @@ -218,6 +224,58 @@ def test_none_hash(monkeypatch, hash_fn): assert hash_fn("python hash seed") == reloaded_kv_cache_utils.NONE_HASH +@pytest.mark.parametrize("non_crypto_fn", [xxhash, xxhash_cbor]) +def test_none_hash_seed_random_for_non_crypto(monkeypatch, non_crypto_fn): + """Non-cryptographic algorithms keep the per-process random seed. + + A deterministic seed is safe for SHA-256, whose collision resistance does + not depend on a secret, but xxHash is not collision resistant: a known seed + would let an attacker precompute colliding blocks offline. Keep the + unpredictable seed there unless the operator opts into a shared one. + """ + # PYTHONHASHSEED unset -> unpredictable, differs per resolution. + with monkeypatch.context() as m: + m.delenv("PYTHONHASHSEED", raising=False) + seeds = {kv_cache_utils.resolve_none_hash_seed(non_crypto_fn) for _ in range(5)} + assert len(seeds) == 5 + assert kv_cache_utils.DEFAULT_NONE_HASH_SEED not in seeds + + # PYTHONHASHSEED set -> operator opt-in wins, so peers can share a cache. + with monkeypatch.context() as m: + m.setenv("PYTHONHASHSEED", "12345") + assert kv_cache_utils.resolve_none_hash_seed(non_crypto_fn) == "12345" + + +@pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor]) +def test_none_hash_seed_deterministic_for_crypto(monkeypatch, hash_fn): + with monkeypatch.context() as m: + m.delenv("PYTHONHASHSEED", raising=False) + seed = kv_cache_utils.resolve_none_hash_seed(hash_fn) + assert seed == kv_cache_utils.DEFAULT_NONE_HASH_SEED + assert seed == kv_cache_utils.resolve_none_hash_seed(hash_fn) + + +def test_get_none_hash_seed_reports_effective_seed(monkeypatch): + """P2P advertises the seed NONE_HASH was actually derived from. + + The P2P tier is constructed before init_none_hash runs, so it must read the + resolved seed lazily rather than re-deriving it. + """ + import vllm.v1.core.kv_cache_utils + + with monkeypatch.context() as m: + m.delenv("PYTHONHASHSEED", raising=False) + reloaded = importlib.reload(vllm.v1.core.kv_cache_utils) + reloaded.init_none_hash(sha256) + assert reloaded.get_none_hash_seed() == reloaded.DEFAULT_NONE_HASH_SEED + + with monkeypatch.context() as m: + m.setenv("PYTHONHASHSEED", "12345") + reloaded = importlib.reload(vllm.v1.core.kv_cache_utils) + reloaded.init_none_hash(sha256) + assert reloaded.get_none_hash_seed() == "12345" + + def test_kv_cache_block(): # Test KVCacheBlock initialization block = KVCacheBlock(block_id=0) diff --git a/tests/v1/kv_offload/tiering/p2p/test_manager.py b/tests/v1/kv_offload/tiering/p2p/test_manager.py index df7a35219206..2cc9618303af 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_manager.py +++ b/tests/v1/kv_offload/tiering/p2p/test_manager.py @@ -15,6 +15,8 @@ import numpy as np import pytest +from vllm.utils.hashing import sha256 +from vllm.v1.core.kv_cache_utils import DEFAULT_NONE_HASH_SEED, init_none_hash from vllm.v1.kv_offload.base import LookupResult, ReqContext, ScheduleEndContext from vllm.v1.kv_offload.tiering.base import JobResult, TransferJob from vllm.v1.kv_offload.tiering.p2p import manager as manager_module @@ -123,23 +125,12 @@ def _init_offloading_spec() -> SimpleNamespace: # --------------------------------------------------------------------------- -# Tests for __init__ PYTHONHASHSEED assertion +# Tests for __init__ hash seed resolution # --------------------------------------------------------------------------- -class TestInitHashSeedAssertion: - def test_missing_pythonhashseed_raises(self, monkeypatch): - """P2P instance refuses to start when PYTHONHASHSEED is unset.""" - monkeypatch.delenv("PYTHONHASHSEED", raising=False) - with pytest.raises(ValueError, match="PYTHONHASHSEED"): - P2PSecondaryTierManager( - offloading_spec=_init_offloading_spec(), - primary_kv_view=memoryview(bytearray(16)), - ) - - def test_pythonhashseed_set_succeeds(self, monkeypatch): - """With PYTHONHASHSEED set, __init__ records it for the handshake.""" - monkeypatch.setenv("PYTHONHASHSEED", "12345") +class TestInitHashSeed: + def _build(self, monkeypatch) -> P2PSecondaryTierManager: monkeypatch.setattr(manager_module, "NixlTransport", lambda *a, **k: object()) monkeypatch.setattr(manager_module, "ZmqTransport", lambda *a, **k: object()) monkeypatch.setattr( @@ -147,11 +138,39 @@ def test_pythonhashseed_set_succeeds(self, monkeypatch): "from_offloading_spec", lambda **k: SimpleNamespace(get_run_config=lambda: {}), ) - mgr = P2PSecondaryTierManager( + return P2PSecondaryTierManager( offloading_spec=_init_offloading_spec(), primary_kv_view=memoryview(bytearray(16)), ) - assert mgr._hash_seed == "12345" + + def test_missing_pythonhashseed_uses_default(self, monkeypatch): + """P2P falls back to the deterministic default seed when unset.""" + monkeypatch.delenv("PYTHONHASHSEED", raising=False) + mgr = self._build(monkeypatch) + init_none_hash(sha256) + assert mgr._get_hash_seed() == DEFAULT_NONE_HASH_SEED + + def test_pythonhashseed_set_succeeds(self, monkeypatch): + """With PYTHONHASHSEED set, the handshake advertises it.""" + monkeypatch.setenv("PYTHONHASHSEED", "12345") + mgr = self._build(monkeypatch) + init_none_hash(sha256) + assert mgr._get_hash_seed() == "12345" + + def test_seed_resolved_after_init_none_hash(self, monkeypatch): + """The seed is read lazily, not at construction time. + + This tier is built before init_none_hash runs, and a non-cryptographic + hash algorithm seeds NONE_HASH randomly, so resolving in __init__ would + advertise a value that does not match the NONE_HASH actually in use. + """ + monkeypatch.delenv("PYTHONHASHSEED", raising=False) + mgr = self._build(monkeypatch) + assert mgr._hash_seed is None + monkeypatch.setattr( + manager_module, "get_none_hash_seed", lambda: "random-seed-abc" + ) + assert mgr._get_hash_seed() == "random-seed-abc" # --------------------------------------------------------------------------- diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index d6e401d1184e..4ac55dbbf70d 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -15,7 +15,7 @@ from vllm import envs from vllm.config import VllmConfig from vllm.logger import init_logger -from vllm.utils.hashing import sha256_cbor, xxhash_cbor +from vllm.utils.hashing import xxhash, xxhash_cbor from vllm.utils.math_utils import cdiv, round_up from vllm.utils.mem_utils import format_gib from vllm.utils.torch_utils import get_dtype_size @@ -87,32 +87,73 @@ def maybe_convert_block_hash(hash_bytes: BlockHash) -> ExternalBlockHash: # The hash seed for the first block of any prefix block sequence. # -# We use a random value to avoid hash collisions or PYTHONHASHSEED environment -# variable if set such that processes can share the seed if needed. This aligns -# with the behavior of Python's hash() function, which also uses a random seed -# if PYTHONHASHSEED is not set. +# For cryptographic hash algorithms it is derived deterministically from a fixed +# default seed, so independent vLLM processes compute identical block hashes for +# identical content and can share a prefix cache (e.g. KV cache reuse across +# nodes) without extra configuration. This does not weaken collision resistance, +# which for SHA-256 does not depend on keeping the seed secret; ``cache_salt`` +# remains the mechanism for intentional cache isolation. +# +# Non-cryptographic algorithms keep a per-process random seed, because a +# predictable seed would let an attacker precompute colliding blocks offline +# (see #12621). Setting PYTHONHASHSEED overrides the seed in both cases. # # The function `init_none_hash` initializes this variable globally. NONE_HASH: BlockHash -_CBOR_HASH_FUNCTIONS = frozenset({sha256_cbor, xxhash_cbor}) +# Fixed seed used when the PYTHONHASHSEED environment variable is not set and +# the hash algorithm is cryptographic. +DEFAULT_NONE_HASH_SEED = "vllm-none-hash" + +# Algorithms that are not collision resistant, so the seed must stay secret. +_NON_CRYPTO_HASH_FUNCTIONS = frozenset({xxhash, xxhash_cbor}) + +# The seed NONE_HASH was derived from, set by init_none_hash. +_NONE_HASH_SEED: str | None = None -def init_none_hash(hash_fn: Callable[[Any], bytes]): - global NONE_HASH +def resolve_none_hash_seed(hash_fn: Callable[[Any], bytes]) -> str: + """Resolve the seed to derive NONE_HASH from. + + PYTHONHASHSEED wins if set. Otherwise cryptographic algorithms get the + fixed default (shareable across processes) and non-cryptographic ones get + fresh random bytes, keeping the seed unpredictable where collision + resistance depends on it. + """ hash_seed = os.getenv("PYTHONHASHSEED") - if hash_seed is None and hash_fn in _CBOR_HASH_FUNCTIONS: + if hash_seed is not None: + return hash_seed + if hash_fn in _NON_CRYPTO_HASH_FUNCTIONS: + return os.urandom(32).hex() + return DEFAULT_NONE_HASH_SEED + + +def get_none_hash_seed() -> str: + """Return the seed NONE_HASH was derived from. + + Components that must agree on NONE_HASH across processes (the P2P tier + advertises this during its connect handshake) read the resolved seed here + instead of re-deriving it, so they observe the random seed too. Falls back + to the deterministic seed before ``init_none_hash`` has run. + """ + if _NONE_HASH_SEED is None: + return DEFAULT_NONE_HASH_SEED + return _NONE_HASH_SEED + + +def init_none_hash(hash_fn: Callable[[Any], bytes]): + global NONE_HASH, _NONE_HASH_SEED + + _NONE_HASH_SEED = resolve_none_hash_seed(hash_fn) + if hash_fn in _NON_CRYPTO_HASH_FUNCTIONS and os.getenv("PYTHONHASHSEED") is None: logger.warning( - "PYTHONHASHSEED is not set. This will lead to non-reproducible " - "block-hashes when using CBOR-based hash functions such as " - "sha256_cbor or xxhash_cbor. Consider setting PYTHONHASHSEED to a " - "fixed value for reproducibility." + "Using a random per-process NONE_HASH seed because %s is not " + "collision resistant. Block hashes are therefore not reproducible " + "across processes; set PYTHONHASHSEED to a shared value to reuse " + "the prefix cache across instances, or use sha256.", + hash_fn.__name__, ) - - if hash_seed is None: - NONE_HASH = BlockHash(os.urandom(32)) - else: - NONE_HASH = BlockHash(hash_fn(hash_seed)) + NONE_HASH = BlockHash(hash_fn(_NONE_HASH_SEED)) @dataclass(slots=True) diff --git a/vllm/v1/kv_offload/tiering/fs/manager.py b/vllm/v1/kv_offload/tiering/fs/manager.py index dae0d8d5008c..6de2d074605a 100644 --- a/vllm/v1/kv_offload/tiering/fs/manager.py +++ b/vllm/v1/kv_offload/tiering/fs/manager.py @@ -95,13 +95,14 @@ class FileSystemTierManager(SecondaryTierManager): get_finished_jobs() polls job completion and returns completed JobResults. Cross-process sharing: - In order to enable KV cache sharing between multiple vLLM instances - using the same ``root_dir`` (e.g., via a shared PVC) the environment - variable ``PYTHONHASHSEED`` must be set to the same fixed value - (e.g., "0") on all instances. Without this, each process initializes - ``NONE_HASH`` (the chain-hash seed for block content hashes) with - random bytes, producing different block filenames for identical token - content. + KV cache sharing between multiple vLLM instances using the same + ``root_dir`` (e.g., via a shared PVC) works by default: ``NONE_HASH`` + (the chain-hash seed for block content hashes) is derived from a fixed + default seed, so identical token content produces identical block + filenames across instances. Setting the ``PYTHONHASHSEED`` environment + variable to the same value on all instances overrides the default seed, + and is required to share a cache when using a non-cryptographic + prefix-caching hash algorithm, which seeds ``NONE_HASH`` randomly. """ medium: ClassVar[Medium] = Medium.STORAGE diff --git a/vllm/v1/kv_offload/tiering/p2p/manager.py b/vllm/v1/kv_offload/tiering/p2p/manager.py index f76ed59b5c78..8050be0f4f76 100644 --- a/vllm/v1/kv_offload/tiering/p2p/manager.py +++ b/vllm/v1/kv_offload/tiering/p2p/manager.py @@ -8,7 +8,6 @@ from __future__ import annotations -import os import time import uuid from collections.abc import Iterable, Sequence @@ -19,6 +18,7 @@ import vllm.envs as envs from vllm.logger import init_logger +from vllm.v1.core.kv_cache_utils import get_none_hash_seed from vllm.v1.kv_offload.base import ( LookupResult, OffloadKey, @@ -247,20 +247,15 @@ def __init__( **kwargs: Reserved for future tier-specific options. """ super().__init__(offloading_spec, primary_kv_view, tier_type) - # Block hashes chain from NONE_HASH, seeded from PYTHONHASHSEED - # (see init_none_hash in v1/core/kv_cache_utils.py). Peers with - # different seeds compute different hashes for identical content, so - # lookups silently miss and no KV crosses the wire. Require it here so - # a misconfigured P2P instance fails at startup rather than degrading - # silently; the value is also verified against each peer on handshake. - hash_seed = os.getenv("PYTHONHASHSEED") - if hash_seed is None: - raise ValueError( - "PYTHONHASHSEED must be set for P2P KV offload so that block " - "hashes match across instances. Set it to a fixed value (e.g. " - "PYTHONHASHSEED=0) on every P2P peer." - ) - self._hash_seed = hash_seed + # Block hashes chain from NONE_HASH (see v1/core/kv_cache_utils.py). + # Peers whose seeds differ compute different hashes for identical + # content, so lookups silently miss and no KV crosses the wire. The + # seed is advertised and verified against each peer during the + # handshake, so a mismatch is rejected loudly instead of degrading + # silently. Resolved lazily in _get_hash_seed: this tier is built + # before init_none_hash runs, and a non-cryptographic hash algorithm + # seeds NONE_HASH randomly, so the value is only known afterwards. + self._hash_seed: str | None = None if host is None: host = envs.VLLM_P2P_SIDE_CHANNEL_HOST if port is None: @@ -602,6 +597,12 @@ def on_schedule_end(self, context: ScheduleEndContext) -> None: # Internal # ------------------------------------------------------------------ + def _get_hash_seed(self) -> str: + """The seed NONE_HASH was derived from, resolved on first session.""" + if self._hash_seed is None: + self._hash_seed = get_none_hash_seed() + return self._hash_seed + def _get_or_create_session(self, peer_id: str) -> P2PSession: """Return the existing session for peer_id, or open one outbound. @@ -621,7 +622,7 @@ def _get_or_create_session(self, peer_id: str) -> P2PSession: local_id=self._local_id, transport=self._data, local_block_len=self._data.block_len, - local_hash_seed=self._hash_seed, + local_hash_seed=self._get_hash_seed(), conn=conn, ) self._sessions[peer_id] = session @@ -643,7 +644,7 @@ def _accept_new_peers(self, new_connections: Sequence[ControlConnection]) -> Non local_id=self._local_id, transport=self._data, local_block_len=self._data.block_len, - local_hash_seed=self._hash_seed, + local_hash_seed=self._get_hash_seed(), conn=conn, ) logger.info( diff --git a/vllm/v1/kv_offload/tiering/p2p/session/client.py b/vllm/v1/kv_offload/tiering/p2p/session/client.py index 07a4620780eb..66b838047bd6 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/client.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/client.py @@ -312,8 +312,9 @@ def on_abort_ack(self, kv_request_id: str, round_seq: int) -> None: if st is not None and load is not None: logger.warning( "P2PSession %s: load request %s (job_id=%d) timed out; " - "load job completed with failure. If this recurs, ensure " - "PYTHONHASHSEED is set to the same value on all nodes.", + "load job completed with failure. If this recurs, ensure all " + "nodes use the same prefix-cache hash seed (matching " + "PYTHONHASHSEED, if set) and hash algorithm.", self._peer_id, kv_request_id, load.job_id, diff --git a/vllm/v1/kv_offload/tiering/p2p/session/protocol.py b/vllm/v1/kv_offload/tiering/p2p/session/protocol.py index 33ae5bf2f793..ec4153e25b39 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/protocol.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/protocol.py @@ -125,9 +125,10 @@ class ConnectMsg: BLOCK_LEN: Size in bytes of each block (must match between peers). CONFIG_FINGERPRINT: SHA-256 prefix of the model configuration. Peers with different fingerprints are incompatible. - HASH_SEED: The peer's PYTHONHASHSEED. Block hashes chain from a seed - derived from it, so peers with different values compute different - hashes for identical content and must not exchange blocks. + HASH_SEED: The peer's effective prefix-cache hash seed (PYTHONHASHSEED + if set, otherwise the built-in default). Block hashes chain from a + seed derived from it, so peers with different values compute + different hashes for identical content and must not exchange blocks. """ TYPE = "connect" diff --git a/vllm/v1/kv_offload/tiering/p2p/session/session.py b/vllm/v1/kv_offload/tiering/p2p/session/session.py index 4bd780159a47..ba2aca6438c6 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/session.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/session.py @@ -458,9 +458,10 @@ def _on_connect(self, msg: dict) -> None: ) if msg[ConnectMsg.HASH_SEED] != self._local_hash_seed: raise ValueError( - f"PYTHONHASHSEED mismatch from {self.peer_id}: " + f"hash seed mismatch from {self.peer_id}: " f"remote={msg[ConnectMsg.HASH_SEED]!r}, " - f"local={self._local_hash_seed!r}" + f"local={self._local_hash_seed!r}. Ensure PYTHONHASHSEED " + "(if set) matches on all P2P peers." ) self._transport.add_remote_peer( self.peer_id, From a9f4afb66f77ba122d905611fa3f620abda98220 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=84=8D=F0=9D=95=A0=F0=9D=95=9D=F0=9D=95=9D=F0=9D=95=A0?= =?UTF-8?q?=F0=9D=95=A8=20=F0=9D=95=84=F0=9D=95=92=F0=9D=95=9F?= Date: Wed, 19 Aug 2026 03:02:18 +0300 Subject: [PATCH 115/839] [Bugfix] Fix DeepSeek V4 mHC broadcast buffer for dummy load (#51368) Signed-off-by: Hollow Man --- tests/models/test_deepseek_v4_mega_moe.py | 34 +++++++++++++++++++++++ vllm/models/deepseek_v4/nvidia/dspark.py | 5 +++- vllm/models/deepseek_v4/nvidia/model.py | 9 +++--- vllm/models/deepseek_v4/nvidia/mtp.py | 5 +++- 4 files changed, 46 insertions(+), 7 deletions(-) diff --git a/tests/models/test_deepseek_v4_mega_moe.py b/tests/models/test_deepseek_v4_mega_moe.py index 25ea9952c94c..ee1e765db6cf 100644 --- a/tests/models/test_deepseek_v4_mega_moe.py +++ b/tests/models/test_deepseek_v4_mega_moe.py @@ -9,10 +9,13 @@ from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( bind_routed_experts_capturer, ) +from vllm.models.deepseek_v4.nvidia.dspark import DSparkDeepseekV4ForCausalLM from vllm.models.deepseek_v4.nvidia.model import ( + DeepseekV4ForCausalLM, DeepseekV4MegaMoEExperts, make_deepseek_v4_expert_params_mapping, ) +from vllm.models.deepseek_v4.nvidia.mtp import DeepSeekV4MTP from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs from vllm.platforms import current_platform @@ -243,6 +246,37 @@ def test_deepseek_v4_mega_moe_fused_input_staging_is_bitwise_exact(): ) +def test_deepseek_v4_pwal_hook_finalizes_mega_moe_and_mhc_broadcast(): + """The loader invokes the model-level PWAL hook for every load format, + so it must finalize megamoe + mhc broadcast weights to cover dummy + load, which skips load_weights().""" + calls = [] + stub = SimpleNamespace( + model=SimpleNamespace( + finalize_mega_moe_weights=lambda: calls.append("mega_moe"), + finalize_mhc_broadcast_weights=lambda: calls.append("mhc"), + ) + ) + + DeepseekV4ForCausalLM.process_weights_after_loading(stub) + + assert calls == ["mega_moe", "mhc"] + + +def test_deepseek_v4_drafter_pwal_hooks_finalize_mega_moe(): + """MTP/DSpark drafters load as their own top-level models, so each needs + its own PWAL hook now that the megamoe forward no longer finalizes + weights lazily on first use.""" + calls = [] + mtp = SimpleNamespace(finalize_mega_moe_weights=lambda: calls.append("mtp")) + DeepSeekV4MTP.process_weights_after_loading(mtp) + + dspark = SimpleNamespace(_finalize_moe=lambda: calls.append("dspark")) + DSparkDeepseekV4ForCausalLM.process_weights_after_loading(dspark) + + assert calls == ["mtp", "dspark"] + + @pytest.mark.skipif( not torch.cuda.is_available(), reason="DeepSeek V4 MegaMoE fused input staging requires CUDA.", diff --git a/vllm/models/deepseek_v4/nvidia/dspark.py b/vllm/models/deepseek_v4/nvidia/dspark.py index b702d7ec9cbd..1a3c2c780d28 100644 --- a/vllm/models/deepseek_v4/nvidia/dspark.py +++ b/vllm/models/deepseek_v4/nvidia/dspark.py @@ -504,9 +504,9 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: weight_loader(param, loaded_weight) loaded_params.add(name) - self._finalize_moe() if self.model.confidence_head is not None and not loaded_confidence_head: self.model.confidence_head = None + self.process_weights_after_loading() logger.info_once("DSpark draft model loaded: %d params", len(loaded_params)) return loaded_params @@ -514,6 +514,9 @@ def _finalize_moe(self) -> None: for layer in self.model.layers: layer.ffn.finalize_mega_moe_weights() + def process_weights_after_loading(self) -> None: + self._finalize_moe() + def _remap_dspark_name(self, name: str) -> str | None: """Map a checkpoint ``mtp.{i}.*`` name to this model's parameter path. diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 17f29f4a1ae1..bd4e7614563b 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -504,10 +504,6 @@ def forward( is_padding=is_padding, ) - # This method must have been already called during the weight loading phase. - # We call it again here to cover the dummy weight loading case. - self.finalize_weights() - assert self._transformed_l1_weights is not None assert self._transformed_l2_weights is not None deep_gemm.fp8_fp4_mega_moe( @@ -1543,9 +1539,12 @@ def get_mtp_target_hidden_states(self) -> torch.Tensor | None: def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self, skip_substrs=["mtp."]) loaded_params = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + self.process_weights_after_loading() + return loaded_params + + def process_weights_after_loading(self) -> None: self.model.finalize_mega_moe_weights() self.model.finalize_mhc_broadcast_weights() - return loaded_params def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: return self.model.get_expert_mapping() diff --git a/vllm/models/deepseek_v4/nvidia/mtp.py b/vllm/models/deepseek_v4/nvidia/mtp.py index eaf105205c1c..7dd32fe97267 100644 --- a/vllm/models/deepseek_v4/nvidia/mtp.py +++ b/vllm/models/deepseek_v4/nvidia/mtp.py @@ -502,7 +502,7 @@ def _find_mtp_layer_idx(name: str) -> int: f"Use a checkpoint that includes MTP layer weights, " f"or disable speculative decoding." ) - self.finalize_mega_moe_weights() + self.process_weights_after_loading() logger.info_once("MTP draft model loaded: %d params", len(loaded_params)) return loaded_params @@ -510,6 +510,9 @@ def finalize_mega_moe_weights(self) -> None: for layer in self.model.layers.values(): layer.mtp_block.ffn.finalize_mega_moe_weights() + def process_weights_after_loading(self) -> None: + self.finalize_mega_moe_weights() + def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: """ Rewrite the weight name to match the format of the original model. From f1178f3a06fa30a0cc282376924210cedad08c44 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Tue, 18 Aug 2026 17:02:59 -0700 Subject: [PATCH 116/839] Revert DSv4 eager workspace reuse (#52836) Signed-off-by: Woosuk Kwon Co-authored-by: OpenAI Codex --- ...deepseek_v4_qnorm_rope_kv_insert_kernel.cu | 30 +--- csrc/libtorch_stable/ops.h | 8 - csrc/libtorch_stable/torch_bindings.cpp | 7 - tests/kernels/test_compressor_kv_cache.py | 18 --- ..._fused_deepseek_v4_qnorm_rope_kv_insert.py | 14 +- .../test_fused_indexer_q_rope_quant.py | 26 ---- vllm/models/deepseek_v4/attention.py | 37 +---- .../deepseek_v4/common/ops/cache_utils.py | 12 +- .../deepseek_v4/common/ops/fused_indexer_q.py | 42 ++---- vllm/models/deepseek_v4/compressor.py | 11 +- vllm/models/deepseek_v4/eager_scratch.py | 137 ------------------ .../deepseek_v4/nvidia/flashinfer_sparse.py | 4 - vllm/models/deepseek_v4/nvidia/flashmla.py | 3 - vllm/models/deepseek_v4/nvidia/model.py | 20 --- .../ops/sparse_attn_compress_cutedsl.py | 15 +- 15 files changed, 30 insertions(+), 354 deletions(-) delete mode 100644 vllm/models/deepseek_v4/eager_scratch.py diff --git a/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu b/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu index ae93266a4593..7bc435b8e0da 100644 --- a/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu +++ b/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu @@ -942,10 +942,9 @@ static void launchFullCacheKernel( // ──────────────────────────────────────────────────────────────────────────── // Torch op wrapper // ──────────────────────────────────────────────────────────────────────────── -void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out( +torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( torch::stable::Tensor const& q_in, // [N, num_heads_q, 512] bf16 torch::stable::Tensor const& kv, // [N, 512] bf16 (read-only) - torch::stable::Tensor& q_out, // [N, q_head_padded, 512] torch::stable::Tensor& k_cache, // [num_blocks, block_bytes] uint8 torch::stable::Tensor const& slot_mapping, // [N] int64 torch::stable::Tensor const& position_ids, // [N] int64 @@ -971,16 +970,8 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out( STD_TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]"); STD_TORCH_CHECK(q_in.scalar_type() == kv.scalar_type(), "q_in and kv dtype must match"); - STD_TORCH_CHECK(q_out.device() == q_in.device() && q_out.is_contiguous(), - "q_out must be contiguous and on the same device as q_in"); - STD_TORCH_CHECK(q_out.scalar_type() == q_in.scalar_type(), - "q_out dtype must match q_in"); STD_TORCH_CHECK(q_head_padded >= q_in.size(1), "q_head_padded must be >= q_in.size(1) (num_heads_q)"); - STD_TORCH_CHECK(q_out.dim() == 3 && q_out.size(0) == q_in.size(0) && - q_out.size(1) == q_head_padded && - q_out.size(2) == q_in.size(2), - "q_out shape [N, q_head_padded, 512]"); STD_TORCH_CHECK(k_cache.scalar_type() == torch::headeronly::ScalarType::Byte, "k_cache must be uint8"); STD_TORCH_CHECK(cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64, @@ -1008,6 +999,11 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out( q_in.get_device_index()); const cudaStream_t stream = get_current_cuda_stream(q_in.get_device_index()); + // Allocate the padded q output. The kernel writes every element (live + // region gets RMSNorm+RoPE; pad region gets zeros), so `empty` is safe. + auto q_out = torch::stable::new_empty( + q_in, {q_in.size(0), q_head_padded, q_in.size(2)}, q_in.scalar_type()); + VLLM_STABLE_DISPATCH_HALF_TYPES( q_in.scalar_type(), "fused_deepseek_v4_qnorm_rope_kv_insert", [&] { using qkv_scalar_t = scalar_t; @@ -1024,20 +1020,6 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out( num_heads_q_padded, cache_block_size_i, kv_block_stride, stream); }); -} - -torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( - torch::stable::Tensor const& q_in, torch::stable::Tensor const& kv, - torch::stable::Tensor& k_cache, - torch::stable::Tensor const& slot_mapping, - torch::stable::Tensor const& position_ids, - torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded, - double eps, int64_t cache_block_size) { - auto q_out = torch::stable::new_empty( - q_in, {q_in.size(0), q_head_padded, q_in.size(2)}, q_in.scalar_type()); - fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out( - q_in, kv, q_out, k_cache, slot_mapping, position_ids, cos_sin_cache, - q_head_padded, eps, cache_block_size); return q_out; } diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 5c595bfaf49c..682ede771a1e 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -269,14 +269,6 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded, double eps, int64_t cache_block_size); -void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out( - torch::stable::Tensor const& q_in, torch::stable::Tensor const& kv, - torch::stable::Tensor& q_out, torch::stable::Tensor& k_cache, - torch::stable::Tensor const& slot_mapping, - torch::stable::Tensor const& position_ids, - torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded, - double eps, int64_t cache_block_size); - void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert( torch::stable::Tensor& q, torch::stable::Tensor const& kv, torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 28090ce5dcaf..3e92e7f91171 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -433,11 +433,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "Tensor q_in, Tensor kv, Tensor! k_cache, " "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, " "int q_head_padded, float eps, int cache_block_size) -> Tensor"); - ops.def( - "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out(" - "Tensor q_in, Tensor kv, Tensor! q_out, Tensor! k_cache, " - "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, " - "int q_head_padded, float eps, int cache_block_size) -> ()"); // FlashInfer V4 full-cache variants: write Q in place (bf16) or to a separate // FP8 tensor, and KV into a contiguous 512-wide token-strided cache. @@ -773,8 +768,6 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { ops.impl("fused_qk_norm_rope", TORCH_BOX(&fused_qk_norm_rope)); ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert)); - ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out", - TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out)); ops.impl( "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert", TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert)); diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index cfdec96e61d0..b300d0ebeccb 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -20,7 +20,6 @@ from vllm import _custom_ops as ops from vllm.models.deepseek_v4.common.ops import ( - compute_global_topk_indices_and_lens, dequantize_and_gather_k_cache, quantize_and_insert_k_cache, ) @@ -113,23 +112,6 @@ def test_get_dspark_swa_index_width( assert get_dspark_swa_index_width(window_size, num_speculative_tokens) == expected -def test_compute_global_topk_reuses_output_buffers(): - device = "cuda" - topk_indices = torch.tensor( - [[0, 3, -1], [1, 2, -1]], dtype=torch.int32, device=device - ) - token_to_req = torch.tensor([0, 1], dtype=torch.int32, device=device) - block_table = torch.tensor([[5, 7], [11, 13]], dtype=torch.int32, device=device) - is_valid = torch.tensor([True, False], device=device) - args = (topk_indices, token_to_req, block_table, 2, is_valid) - expected = compute_global_topk_indices_and_lens(*args) - outputs = tuple(torch.empty_like(tensor) for tensor in expected) - actual = compute_global_topk_indices_and_lens(*args, output_buffers=outputs) - for result, output, reference in zip(actual, outputs, expected): - assert result.data_ptr() == output.data_ptr() - torch.testing.assert_close(result, reference) - - def _ue8m0_reference(x: torch.Tensor, block_size: int, fp8_max: float): """PyTorch reference for UE8M0 FP8 quantization (per-block, power-of-2 scale). diff --git a/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py b/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py index 5572ee89f210..ed163a0472ad 100644 --- a/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py +++ b/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py @@ -257,18 +257,8 @@ def test_q_path_matches_reference(num_tokens: int, n_heads: int, padded_heads: i num_blocks, bs, HEAD_BYTES, dtype=torch.uint8, device=device ).view(num_blocks, -1) slot_mapping = torch.full((num_tokens,), -1, dtype=torch.int64, device=device) - q_out = torch.empty(num_tokens, padded_heads, HEAD_DIM, dtype=dtype, device=device) - torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out( - q, - kv, - q_out, - k_cache, - slot_mapping, - positions, - cos_sin_cache, - padded_heads, - eps, - bs, + q_out = _call_fused( + q, padded_heads, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs ) torch.testing.assert_close(q_out[:, :n_heads], q_ref, rtol=1e-2, atol=1e-2) diff --git a/tests/kernels/test_fused_indexer_q_rope_quant.py b/tests/kernels/test_fused_indexer_q_rope_quant.py index f8bf944a98f0..6114b7efd6e7 100644 --- a/tests/kernels/test_fused_indexer_q_rope_quant.py +++ b/tests/kernels/test_fused_indexer_q_rope_quant.py @@ -150,23 +150,6 @@ def test_fused_indexer_q_rope_quant_matches_unfused( q_quant_ref, weights_ref = _reference( positions, q, cos_sin_cache, weights, softmax_scale, head_scale, use_fp4 ) - output_buffers: tuple[torch.Tensor, ...] | None = None - OUTPUT_BUFFER_TEST_NUM_TOKENS = 7 - if num_tokens == OUTPUT_BUFFER_TEST_NUM_TOKENS and cache_dtype == torch.float32: - if use_fp4: - q_ref, q_scale_ref = q_quant_ref - output_buffers = ( - torch.empty_like(q_ref), - torch.empty_like(q_scale_ref) - .view(torch.uint8) - .reshape(num_tokens, N_HEAD, -1), - torch.empty_like(weights_ref), - ) - else: - output_buffers = ( - torch.empty_like(q_quant_ref), - torch.empty_like(weights_ref), - ) # use_cutedsl=False: force the triton path even when cutedsl is installed # by patching the dispatcher's has_cutedsl() binding to return False. cutedsl_patch = ( @@ -186,17 +169,8 @@ def test_fused_indexer_q_rope_quant_matches_unfused( softmax_scale, head_scale, use_fp4, - output_buffers=output_buffers, ) - if output_buffers is not None: - if use_fp4: - assert q_quant_fused[0].data_ptr() == output_buffers[0].data_ptr() - assert q_quant_fused[1].data_ptr() == output_buffers[1].data_ptr() - else: - assert q_quant_fused.data_ptr() == output_buffers[0].data_ptr() - assert weights_fused.data_ptr() == output_buffers[-1].data_ptr() - if use_fp4: q_quant_ref, q_scale_ref = q_quant_ref q_quant_fused, q_scale_fused = q_quant_fused diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index 128debf70cb1..ef6f9ea96fdc 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -29,7 +29,6 @@ from vllm.models.deepseek_v4.common.ops.fused_indexer_q import MXFP4_BLOCK_SIZE if TYPE_CHECKING: - from vllm.models.deepseek_v4.eager_scratch import DeepseekV4EagerScratchPool from vllm.v1.attention.backends.mla.sparse_swa import ( DeepseekSparseSWAMetadata, ) @@ -185,7 +184,6 @@ def __init__( prefix: str, topk_indices_buffer: torch.Tensor | None = None, aux_stream_list: list[torch.cuda.Stream] | None = None, - eager_scratch_pool: "DeepseekV4EagerScratchPool | None" = None, ) -> None: super().__init__() config = vllm_config.model_config.hf_config @@ -274,7 +272,6 @@ def __init__( ) self.indexer_rotary_emb = self.rotary_emb self.topk_indices_buffer = topk_indices_buffer - self.eager_scratch_pool = eager_scratch_pool self.indexer = None if self.compress_ratio == 4: @@ -296,7 +293,6 @@ def __init__( compress_ratio=self.compress_ratio, prefix=f"{prefix}.indexer", aux_stream=indexer_aux_stream, - eager_scratch_pool=eager_scratch_pool, ) self._prepare_and_attn_fn = self._prepare_and_attn @@ -355,7 +351,6 @@ def __init__( rotate=True, prefix=f"{prefix}.compressor", k_cache_prefix=self.prefix, - eager_scratch_pool=eager_scratch_pool, ) def forward( @@ -637,24 +632,10 @@ def _fused_qnorm_rope_kv_insert( if cache_dtype == torch.uint8: # fp8_ds_mla UE8M0 paged path. Horizontally fused: # Q side: per-head RMSNorm (no weight) + GPT-J RoPE, zero-filling - # the padding head slots. + # the padding head slots; the kernel allocates and returns + # the padded q tensor. # KV side: GPT-J RoPE + UE8M0 FP8 quant + paged cache insert. swa_kv_cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1) - if self.eager_scratch_pool is not None: - q_out = self.eager_scratch_pool.q_out(q.shape[0]) - torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out( - q, - kv, - q_out, - swa_kv_cache_2d, - swa_metadata.slot_mapping, - positions, - cos_sin_cache, - self.padded_heads, - self.eps, - swa_metadata.block_size, - ) - return q_out return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( q, kv, @@ -703,13 +684,6 @@ def _fused_qnorm_rope_kv_insert( ) return q_fp8 - def _global_topk_output_buffers( - self, topk_indices: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor] | None: - if self.compress_ratio != 4 or self.eager_scratch_pool is None: - return None - return self.eager_scratch_pool.global_topk_outputs(topk_indices) - def get_attn_backend(self) -> type[AttentionBackend]: return self.backend_cls @@ -792,7 +766,6 @@ def __init__( compress_ratio: int = 1, prefix: str = "", aux_stream: torch.cuda.Stream | None = None, - eager_scratch_pool: "DeepseekV4EagerScratchPool | None" = None, ): super().__init__() self.vllm_config = vllm_config @@ -805,7 +778,6 @@ def __init__( self.rope_dim = config.qk_rope_head_dim # 64 self.q_lora_rank = q_lora_rank # 1536 self.compress_ratio = compress_ratio - self.eager_scratch_pool = eager_scratch_pool self.use_fp4_kv = dsa_indexer_uses_fp4(vllm_config) logger.info_once( "Using %s indexer cache for Lightning Indexer.", @@ -869,7 +841,6 @@ def __init__( prefix=f"{prefix}.compressor", k_cache_prefix=self.k_cache.prefix, use_fp4_cache=self.use_fp4_kv, - eager_scratch_pool=eager_scratch_pool, ) self.indexer_op = SparseAttnIndexer( @@ -933,9 +904,6 @@ def wq_b_and_q_quant(): # ReplicatedLinear returns (output, bias); bias is None. q, _ = self.wq_b(qr) q = q.view(-1, self.n_head, self.head_dim) - outputs = None - if self.eager_scratch_pool is not None and self.use_fp4_kv: - outputs = self.eager_scratch_pool.indexer_q_outputs(q.shape[0]) return fused_indexer_q_rope_quant( positions, q, @@ -944,7 +912,6 @@ def wq_b_and_q_quant(): self.softmax_scale, self.n_head**-0.5, use_fp4=self.use_fp4_kv, - output_buffers=outputs, ) # compressor returns None and writes K to the indexer KV cache; the diff --git a/vllm/models/deepseek_v4/common/ops/cache_utils.py b/vllm/models/deepseek_v4/common/ops/cache_utils.py index 1f53c96e9367..8aa5e5d64bec 100644 --- a/vllm/models/deepseek_v4/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v4/common/ops/cache_utils.py @@ -438,7 +438,6 @@ def compute_global_topk_indices_and_lens( block_table: torch.Tensor, block_size: int, is_valid_token: torch.Tensor, - output_buffers: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Map local topk indices to global KV cache slots and count valid entries. @@ -448,15 +447,8 @@ def compute_global_topk_indices_and_lens( 3. Masking padding tokens to length 0 """ num_tokens = topk_indices.shape[0] - if output_buffers is None: - global_topk_indices = torch.empty_like(topk_indices) - topk_lens = torch.empty( - num_tokens, dtype=torch.int32, device=topk_indices.device - ) - else: - global_topk_indices, topk_lens = output_buffers - assert global_topk_indices.shape == topk_indices.shape - assert topk_lens.shape == (num_tokens,) + global_topk_indices = torch.empty_like(topk_indices) + topk_lens = torch.empty(num_tokens, dtype=torch.int32, device=topk_indices.device) _compute_global_topk_indices_and_lens_kernel[(num_tokens,)]( global_topk_indices, global_topk_indices.stride(0), diff --git a/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py b/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py index 5aa00174079a..3ec0c2b29d5d 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py +++ b/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py @@ -295,7 +295,6 @@ def fused_indexer_q_rope_quant( index_weights_softmax_scale: float, index_weights_head_scale: float, use_fp4: bool = False, - output_buffers: tuple[torch.Tensor, ...] | None = None, ) -> tuple[ torch.Tensor | tuple[torch.Tensor, torch.Tensor], torch.Tensor, @@ -333,13 +332,7 @@ def fused_indexer_q_rope_quant( num_index_q_heads = index_q.shape[1] index_q_head_dim = index_q.shape[2] - if output_buffers is None: - index_weights_out = torch.empty_like(index_weights, dtype=torch.float32) - else: - expected_num_buffers = 3 if use_fp4 else 2 - assert len(output_buffers) == expected_num_buffers - index_weights_out = output_buffers[-1] - assert index_weights_out.shape == index_weights.shape + index_weights_out = torch.empty_like(index_weights, dtype=torch.float32) if use_fp4: assert index_q_head_dim % MXFP4_BLOCK_SIZE == 0, ( @@ -347,23 +340,16 @@ def fused_indexer_q_rope_quant( f"size {MXFP4_BLOCK_SIZE}" ) num_scale_blocks = index_q_head_dim // MXFP4_BLOCK_SIZE - packed_shape = (num_tokens, num_index_q_heads, index_q_head_dim // 2) - scale_shape = (num_tokens, num_index_q_heads, num_scale_blocks) - if output_buffers is None: - index_q_packed = torch.empty( - packed_shape, - dtype=torch.uint8, - device=index_q.device, - ) - index_q_scale = torch.empty( - scale_shape, - dtype=torch.uint8, - device=index_q.device, - ) - else: - index_q_packed, index_q_scale, _ = output_buffers - assert index_q_packed.shape == packed_shape - assert index_q_scale.shape == scale_shape + index_q_packed = torch.empty( + (num_tokens, num_index_q_heads, index_q_head_dim // 2), + dtype=torch.uint8, + device=index_q.device, + ) + index_q_scale = torch.empty( + (num_tokens, num_index_q_heads, num_scale_blocks), + dtype=torch.uint8, + device=index_q.device, + ) if has_cutedsl(): # lazily import, otherwise some tests fail due to CUDA driver init failure. from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import ( @@ -432,11 +418,7 @@ def fused_indexer_q_rope_quant( fp8_dtype = current_platform.fp8_dtype() use_fnuz = fp8_dtype == torch.float8_e4m3fnuz fp8_max = 224.0 if use_fnuz else 448.0 - if output_buffers is None: - index_q_fp8 = torch.empty_like(index_q, dtype=fp8_dtype) - else: - index_q_fp8, _ = output_buffers - assert index_q_fp8.shape == index_q.shape + index_q_fp8 = torch.empty_like(index_q, dtype=fp8_dtype) if has_cutedsl(): # lazily import, otherwise some tests fail due to CUDA driver init failure. from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import ( diff --git a/vllm/models/deepseek_v4/compressor.py b/vllm/models/deepseek_v4/compressor.py index 93ae1e89c919..590e748145a4 100644 --- a/vllm/models/deepseek_v4/compressor.py +++ b/vllm/models/deepseek_v4/compressor.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, ClassVar, cast +from typing import Any, ClassVar, cast import torch from torch import nn @@ -35,9 +35,6 @@ SlidingWindowMLASpec, ) -if TYPE_CHECKING: - from vllm.models.deepseek_v4.eager_scratch import DeepseekV4EagerScratchPool - def _prefer_two_stage_compressor() -> bool: # Platforms that favor the triton variant of two-stage compressor split. @@ -229,7 +226,6 @@ def __init__( prefix: str = "", k_cache_prefix="", use_fp4_cache: bool = False, - eager_scratch_pool: "DeepseekV4EagerScratchPool | None" = None, ): super().__init__() self.compress_ratio = compress_ratio @@ -239,7 +235,6 @@ def __init__( self.prefix = prefix self.k_cache_prefix = k_cache_prefix self.use_fp4_cache = use_fp4_cache - self.eager_scratch_pool = eager_scratch_pool config = vllm_config.model_config.hf_config self.rope_head_dim = config.qk_rope_head_dim @@ -433,10 +428,6 @@ def forward( store_full_fp8=store_full_fp8, fp8_scale=fp8_scale, ) - if not self.overlap and self.eager_scratch_pool is not None: - extra_kwargs["compress_scratch"] = ( - self.eager_scratch_pool.compressor_scratch(num_actual) - ) elif self._use_two_stage_fused_compressor: # head=512 cr>=128 (no overlap): two-pass split compressor on the # prefill suffix, single-pass on the decode prefix. diff --git a/vllm/models/deepseek_v4/eager_scratch.py b/vllm/models/deepseek_v4/eager_scratch.py deleted file mode 100644 index bc46239bfe34..000000000000 --- a/vllm/models/deepseek_v4/eager_scratch.py +++ /dev/null @@ -1,137 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from math import prod - -import torch - -from vllm.models.deepseek_v4.common.ops.fused_indexer_q import MXFP4_BLOCK_SIZE -from vllm.utils.math_utils import round_up - - -class DeepseekV4EagerScratchPool: - """Model-wide outputs and scratch used inside the attention eager break.""" - - _ALIGNMENT = 256 - - def __init__( - self, - max_num_tokens: int, - padded_q_heads: int, - q_head_dim: int, - index_q_heads: int, - index_q_head_dim: int, - index_topk: int, - device: torch.device | str, - ) -> None: - self.max_num_tokens = max_num_tokens - self.index_topk = index_topk - self._q = torch.empty( - (max_num_tokens, padded_q_heads, q_head_dim), - dtype=torch.bfloat16, - device=device, - ) - - fp4_specs = ( - ((max_num_tokens, index_q_heads, index_q_head_dim // 2), torch.uint8), - ( - ( - max_num_tokens, - index_q_heads, - index_q_head_dim // MXFP4_BLOCK_SIZE, - ), - torch.uint8, - ), - ((max_num_tokens, index_q_heads), torch.float32), - ) - global_specs = ( - ((max_num_tokens, index_topk), torch.int32), - ((max_num_tokens,), torch.int32), - ) - compressor_specs = (((max_num_tokens, q_head_dim), torch.float32),) - # FP4 indexer is C4 only, global mapping after FP4 indexer - # compressor scratch is C128 only - # so here we use max instead of sum - aux_bytes = max( - self._packed_size(specs) - for specs in (fp4_specs, global_specs, compressor_specs) - ) - storage = torch.empty(aux_bytes, dtype=torch.uint8, device=device) - - self._q_outputs: dict[int, torch.Tensor] = {} - fp4_values, fp4_scales, fp4_weights = self._views(storage, fp4_specs) - self._fp4_template = (fp4_values, fp4_scales, fp4_weights) - self._fp4_outputs: dict[ - int, tuple[torch.Tensor, torch.Tensor, torch.Tensor] - ] = {} - global_indices, global_lens = self._views(storage, global_specs) - self._global_template = (global_indices, global_lens) - self._global_outputs: dict[int, tuple[torch.Tensor, torch.Tensor]] = {} - self._compressor_template = self._views(storage, compressor_specs)[0] - self._compressor_outputs: dict[int, torch.Tensor] = {} - self._storage = storage - - @classmethod - def _packed_size( - cls, specs: tuple[tuple[tuple[int, ...], torch.dtype], ...] - ) -> int: - offset = 0 - for shape, dtype in specs: - offset = round_up(offset, cls._ALIGNMENT) + prod(shape) * dtype.itemsize - return round_up(offset, cls._ALIGNMENT) - - @classmethod - def _views( - cls, - storage: torch.Tensor, - specs: tuple[tuple[tuple[int, ...], torch.dtype], ...], - ) -> list[torch.Tensor]: - offset = 0 - views = [] - for shape, dtype in specs: - offset = round_up(offset, cls._ALIGNMENT) - num_bytes = prod(shape) * dtype.itemsize - views.append(storage[offset : offset + num_bytes].view(dtype).view(shape)) - offset += num_bytes - return views - - def q_out(self, num_tokens: int) -> torch.Tensor: - output = self._q_outputs.get(num_tokens) - if output is None: - output = self._q[:num_tokens] - self._q_outputs[num_tokens] = output - return output - - def compressor_scratch(self, num_tokens: int) -> torch.Tensor: - output = self._compressor_outputs.get(num_tokens) - if output is None: - output = self._compressor_template[:num_tokens] - self._compressor_outputs[num_tokens] = output - return output - - def indexer_q_outputs( - self, - num_tokens: int, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - output = self._fp4_outputs.get(num_tokens) - if output is None: - values, scales, weights = self._fp4_template - output = ( - values[:num_tokens], - scales[:num_tokens], - weights[:num_tokens], - ) - self._fp4_outputs[num_tokens] = output - return output - - def global_topk_outputs( - self, topk_indices: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor]: - num_tokens, topk = topk_indices.shape - assert topk == self.index_topk - output = self._global_outputs.get(num_tokens) - if output is None: - indices, lens = self._global_template - output = (indices[:num_tokens], lens[:num_tokens]) - self._global_outputs[num_tokens] = output - return output diff --git a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py index e73156c1c121..1250934b8c65 100644 --- a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py +++ b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py @@ -763,9 +763,6 @@ def _forward_decode( attn_metadata.block_table[:num_decodes], block_size, is_valid, - output_buffers=self._global_topk_output_buffers( - self.topk_indices_buffer[:num_decode_tokens] - ), ) ) extra_sparse_indices = global_indices.view(num_decode_tokens, 1, -1) @@ -855,7 +852,6 @@ def _forward_prefill( attn_metadata.block_table, block_size, swa_metadata.is_valid_token[prefill_token_slice], - output_buffers=self._global_topk_output_buffers(local_topk_indices), ) ) diff --git a/vllm/models/deepseek_v4/nvidia/flashmla.py b/vllm/models/deepseek_v4/nvidia/flashmla.py index a6a24ae0cc71..b8eed8aefe5b 100644 --- a/vllm/models/deepseek_v4/nvidia/flashmla.py +++ b/vllm/models/deepseek_v4/nvidia/flashmla.py @@ -197,9 +197,6 @@ def _forward_decode( attn_metadata.block_table[:num_decodes], block_size, is_valid, - output_buffers=self._global_topk_output_buffers( - self.topk_indices_buffer[:num_decode_tokens] - ), ) topk_indices = global_indices.view(num_decode_tokens, 1, -1) else: diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index bd4e7614563b..922d1c449871 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -72,7 +72,6 @@ sp_shard, ) from vllm.models.deepseek_v4.attention import DeepseekV4Attention -from vllm.models.deepseek_v4.eager_scratch import DeepseekV4EagerScratchPool from vllm.models.deepseek_v4.nvidia.flashinfer_sparse import ( DeepseekV4FlashInferMLAAttention, DeepseekV4FlashInferSM120Attention, @@ -815,7 +814,6 @@ def __init__( prefix, topk_indices_buffer: torch.Tensor | None = None, aux_stream_list: list[torch.cuda.Stream] | None = None, - eager_scratch_pool: DeepseekV4EagerScratchPool | None = None, ): super().__init__() @@ -829,7 +827,6 @@ def __init__( prefix=f"{prefix}.attn", topk_indices_buffer=topk_indices_buffer, aux_stream_list=aux_stream_list, - eager_scratch_pool=eager_scratch_pool, ) if self.use_sequence_parallel: self.attn.wo_b.reduce_results = False @@ -1017,22 +1014,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): # (compressor kv_score, indexer.weights_proj, indexer.compressor # kv_score). fused_wqa_wkv stays on the default stream. aux_stream_list = [torch.cuda.Stream() for _ in range(3)] - padded_heads = _select_dsv4_attn_cls(vllm_config).get_padded_num_q_heads( - config.num_attention_heads // get_tensor_model_parallel_world_size() - ) - self.eager_scratch_pool: DeepseekV4EagerScratchPool | None = None - if not vllm_config.parallel_config.use_ubatching: - # TODO: support dbo if needed - # this requires the buffer to have ubatch dim - self.eager_scratch_pool = DeepseekV4EagerScratchPool( - vllm_config.scheduler_config.max_num_batched_tokens, - padded_heads, - config.head_dim, - config.index_n_heads, - config.index_head_dim, - config.index_topk, - current_platform.device_type, - ) # Reserved topk indices buffer for all Indexer layers to reuse. self.topk_indices_buffer = torch.empty( @@ -1058,7 +1039,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=prefix, topk_indices_buffer=self.topk_indices_buffer, aux_stream_list=aux_stream_list, - eager_scratch_pool=self.eager_scratch_pool, ), prefix=f"{prefix}.layers", ) diff --git a/vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py b/vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py index 960c516c4522..4ff4b232d10f 100644 --- a/vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py +++ b/vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py @@ -2097,7 +2097,6 @@ def compress_norm_rope_store_cutedsl( store_full_kv: bool = False, store_full_fp8: bool = False, fp8_scale: torch.Tensor | None = None, - compress_scratch: torch.Tensor | None = None, ) -> None: if compress_ratio == 4: # For C4A, the single fused kernel is faster than the two-kernel version. @@ -2130,15 +2129,11 @@ def compress_norm_rope_store_cutedsl( ) else: # For C128, the two-kernel version is faster than the single fused kernel. - if compress_scratch is None: - compressed_kv = torch.empty( - (num_actual, head_dim), - dtype=torch.float32, - device=state_cache.device, - ) - else: - assert compress_scratch.shape == (num_actual, head_dim) - compressed_kv = compress_scratch + compressed_kv = torch.empty( + (num_actual, head_dim), + dtype=torch.float32, + device=state_cache.device, + ) split_kv_compress_norm_rope_insert_sparse_attn_cutedsl( state_cache, token_to_req_indices, From b05ae5dc008850a620dec6de66635dec2b5913fd Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Tue, 18 Aug 2026 18:04:10 -0700 Subject: [PATCH 117/839] [CI][Bugfix] Complete DeepSeek-V4 FSE test fixture contract (#52842) Signed-off-by: Thomas Parnell Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Co-authored-by: Thomas Parnell Co-authored-by: Claude --- tests/model_executor/layers/test_fused_shared_expert.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/model_executor/layers/test_fused_shared_expert.py b/tests/model_executor/layers/test_fused_shared_expert.py index 7080c66f128c..c667f326f6da 100644 --- a/tests/model_executor/layers/test_fused_shared_expert.py +++ b/tests/model_executor/layers/test_fused_shared_expert.py @@ -479,6 +479,8 @@ def test_models_fse_init( is_mm_prefix_lm=False, multimodal_config=None, quantization_config=None, + runner_type="generate", + is_moe=True, ) vllm_config.parallel_config.enable_expert_parallel = False if model_type == "deepseek_v4": @@ -535,7 +537,8 @@ def test_models_fse_init( ) vllm_config.speculative_config = SimpleNamespace( - draft_model_config=SimpleNamespace(hf_config=config) + draft_model_config=SimpleNamespace(hf_config=config), + method="mtp", ) mtp = DeepSeekV4MTP(vllm_config=vllm_config) assert model.is_fused_shared_expert_enabled is (fse_enabled and not exclude) From d36cc4254eb6f745427e9bd328aab5ac34e5cb7e Mon Sep 17 00:00:00 2001 From: Itay Etelis <92247226+Etelis@users.noreply.github.com> Date: Wed, 19 Aug 2026 04:37:29 +0300 Subject: [PATCH 118/839] [Bugfix][Elastic EP] Reject scale below the minimum data parallel size (#52702) Signed-off-by: Itay Etelis Co-authored-by: Itay Etelis --- vllm/entrypoints/serve/elastic_ep/api_router.py | 2 ++ vllm/v1/engine/core_client.py | 16 ++++++++++++++-- 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/vllm/entrypoints/serve/elastic_ep/api_router.py b/vllm/entrypoints/serve/elastic_ep/api_router.py index 02a242509050..52c68d332174 100644 --- a/vllm/entrypoints/serve/elastic_ep/api_router.py +++ b/vllm/entrypoints/serve/elastic_ep/api_router.py @@ -75,6 +75,8 @@ async def scale_elastic_ep(raw_request: Request): detail="Scale failed due to request drain timeout " f"after {drain_timeout} seconds", ) from e + except ValueError as e: + raise HTTPException(status_code=400, detail=str(e)) from e except Exception as e: logger.error("Scale failed: %s", e) raise HTTPException(status_code=500, detail="Scale failed") from e diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index d692a13d658a..a626af89e1f3 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -1637,9 +1637,21 @@ async def prepare_elastic_ep(self, new_data_parallel_size: int) -> None: ) parallel_config = self.vllm_config.parallel_config num_experts = self.vllm_config.model_config.get_num_experts() - num_redundant_experts = ( + num_physical_experts = ( num_experts + parallel_config.eplb_config.num_redundant_experts - ) * new_data_parallel_size // cur_data_parallel_size - num_experts + ) + num_redundant_experts = ( + num_physical_experts * new_data_parallel_size // cur_data_parallel_size + - num_experts + ) + if num_redundant_experts < 0: + # Scaling keeps physical experts per engine fixed, so below this + # size the logical experts no longer fit. + raise ValueError( + f"Cannot scale to data_parallel_size {new_data_parallel_size}, " + f"minimum is " + f"{-(-num_experts * cur_data_parallel_size // num_physical_experts)}" + ) if new_data_parallel_size < cur_data_parallel_size: await self._prepare_scale_down_elastic_ep(new_data_parallel_size) else: From a2257f95b79287933defef5be40b74ce93053da3 Mon Sep 17 00:00:00 2001 From: Itay Alroy <75032521+itayalroy@users.noreply.github.com> Date: Wed, 19 Aug 2026 05:13:06 +0300 Subject: [PATCH 119/839] [Elastic EP] Reduce eager-mode reconfiguration downtime (#51885) Signed-off-by: Itay Alroy --- tests/distributed/test_eplb_execute.py | 122 ------------- .../worker/test_gpu_model_runner_v2_eplb.py | 3 +- vllm/config/parallel.py | 2 +- .../distributed/elastic_ep/elastic_execute.py | 162 ++++++++++++------ vllm/distributed/elastic_ep/elastic_state.py | 49 ++---- vllm/distributed/eplb/eplb_communicator.py | 48 ++---- vllm/distributed/eplb/eplb_state.py | 71 +++++--- vllm/distributed/parallel_state.py | 14 +- vllm/v1/worker/gpu/eplb_utils.py | 2 - vllm/v1/worker/gpu/model_runner.py | 2 - vllm/v1/worker/gpu_model_runner.py | 20 +-- vllm/v1/worker/gpu_worker.py | 2 + 12 files changed, 194 insertions(+), 303 deletions(-) diff --git a/tests/distributed/test_eplb_execute.py b/tests/distributed/test_eplb_execute.py index 4c9b98b62cf5..21fa057fd201 100644 --- a/tests/distributed/test_eplb_execute.py +++ b/tests/distributed/test_eplb_execute.py @@ -782,125 +782,3 @@ def test_rearrange_expert_weights_profile_mode(world_size): _test_rearrange_expert_weights_profile_mode, world_size, ) - - -def _test_nixl_deferred_init_worker( - env, - world_size: int, - num_layers: int, - num_local_experts: int, - num_logical_experts: int, -) -> None: - """Exercise NixlEplbCommunicator with defer_remote_setup=True (elastic EP path).""" - from vllm.distributed.eplb.eplb_communicator import NixlEplbCommunicator - - set_env_vars_and_device(env) - - vllm_config = VllmConfig() - vllm_config.parallel_config.tensor_parallel_size = world_size - - with set_current_vllm_config(vllm_config): - ensure_model_parallel_initialized( - tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1 - ) - - ep_group_coordinator = get_tp_group() - ep_group = ep_group_coordinator.cpu_group - ep_rank = torch.distributed.get_rank() - device = torch.device(f"cuda:{ep_rank}") - - total_physical_experts = world_size * num_local_experts - hidden_sizes = [32, 64] - - redundancy_config = create_redundancy_config( - num_logical_experts, total_physical_experts - ) - old_indices = create_expert_indices_with_redundancy( - num_layers, - num_logical_experts, - total_physical_experts, - redundancy_config, - ) - - new_redundancy_config = create_redundancy_config( - num_logical_experts, total_physical_experts - ) - new_indices = create_expert_indices_with_redundancy( - num_layers, - num_logical_experts, - total_physical_experts, - new_redundancy_config, - ) - - expert_weights = create_expert_weights( - num_layers, num_local_experts, hidden_sizes, ep_rank, device, old_indices - ) - - expert_buffer = [torch.empty_like(w) for w in expert_weights[0]] - - communicator = NixlEplbCommunicator( - cpu_group=ep_group_coordinator.cpu_group, - all_expert_weights=expert_weights, - expert_buffer=expert_buffer, - defer_remote_setup=True, - ) - assert not communicator._remote_state_initialized - - rearrange_expert_weights_inplace( - old_indices, - new_indices, - expert_weights, - expert_buffer, - ep_group, - communicator, - ) - - assert communicator._remote_state_initialized - - local_ok = verify_expert_weights_after_shuffle( - expert_weights, - new_indices, - hidden_sizes, - ep_rank, - num_local_experts, - ) - - local_ok = ( - verify_redundant_experts_have_same_weights( - expert_weights, - new_indices, - hidden_sizes, - ep_rank, - world_size, - num_local_experts, - ) - and local_ok - ) - assert_verification_synced( - local_ok, - "Deferred NIXL init verification failed on at least one rank.", - ) - - -@pytest.mark.skipif(not has_nixl(), reason="NIXL is not available") -@pytest.mark.parametrize( - "world_size,num_layers,num_local_experts,num_logical_experts", - [(2, 2, 3, 4)], -) -def test_nixl_deferred_init( - world_size, - num_layers, - num_local_experts, - num_logical_experts, -): - """Test NixlEplbCommunicator with defer_remote_setup=True (elastic EP path).""" - - if torch.accelerator.device_count() < world_size: - pytest.skip(f"Need at least {world_size} GPUs to run the test") - distributed_run( - _test_nixl_deferred_init_worker, - world_size, - num_layers, - num_local_experts, - num_logical_experts, - ) diff --git a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py index ebb4beb2a5a9..637d078c7323 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py +++ b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py @@ -159,13 +159,12 @@ def test_v2_setup_eplb_from_mapping_rebuilds_state(monkeypatch): runner = _make_runner(model=SimpleNamespace(is_moe=True)) mapping = torch.tensor([[0, 1, 2, 3]], dtype=torch.int64) - mrv2.GPUModelRunner.setup_eplb_from_mapping(runner, mapping, 2) + mrv2.GPUModelRunner.setup_eplb_from_mapping(runner, mapping) assert runner.eplb_state is not None assert runner.eplb_state.built_from_mapping is True assert FakeEplbState.from_mapping_kwargs is not None assert FakeEplbState.from_mapping_kwargs["expanded_physical_to_logical"] is mapping - assert FakeEplbState.from_mapping_kwargs["num_valid_physical_experts"] == 2 def test_v2_sample_tokens_runs_eplb_on_non_last_pp_rank(monkeypatch): diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index eb346ed4a34a..d0118c6503b6 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -973,7 +973,7 @@ def __post_init__(self) -> None: if self.enable_eplb and self.eplb_config.communicator is None: # Prefer NIXL when available: zero-copy RDMA reads, compatible - # with both async EPLB and elastic EP (deferred remote setup). + # with both async EPLB and elastic EP. # Fallbacks: pynccl for elastic EP (stateless groups need it), # torch_gloo for static EP. torch_nccl is avoided because NCCL # is incompatible with async EPLB (multi-stream conflicts) and diff --git a/vllm/distributed/elastic_ep/elastic_execute.py b/vllm/distributed/elastic_ep/elastic_execute.py index 99325a1706ab..6124d22b44d8 100644 --- a/vllm/distributed/elastic_ep/elastic_execute.py +++ b/vllm/distributed/elastic_ep/elastic_execute.py @@ -29,12 +29,15 @@ create_standby_groups, get_standby_dp_group, get_standby_ep_group, + get_standby_eplb_group, pop_standby_groups, ) -from vllm.distributed.eplb.eplb_communicator import create_eplb_communicator +from vllm.distributed.eplb.eplb_communicator import ( + EplbCommunicator, +) from vllm.distributed.parallel_state import ( + GroupCoordinator, _replace_active_groups, - get_eplb_group, ) from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator from vllm.logger import init_logger @@ -149,10 +152,12 @@ def __init__(self, worker): self.worker_ref = weakref.ref(worker) self.reconfig_request = None self._staged_moe_quant_methods: dict[nn.Module, FusedMoEMethodBase] = {} + self._prepared_eplb_communicator: EplbCommunicator | None = None self._async_executor = ThreadPoolExecutor( max_workers=1, thread_name_prefix="ElasticEPAsync" ) self._async_future: Future[None] | None = None + self._group_cleanup_future: Future[None] | None = None @property def worker(self): @@ -207,12 +212,40 @@ def clear_async(self) -> None: self._async_future = None future.result() + def _destroy_retired_groups( + self, groups: tuple[GroupCoordinator | None, ...] + ) -> None: + from vllm.platforms import current_platform + + current_platform.set_device(self.worker.device) + for group in groups: + if group is not None: + group.destroy() + + def _start_group_cleanup(self, groups: tuple[GroupCoordinator | None, ...]) -> None: + assert self._group_cleanup_future is None + self._group_cleanup_future = self._async_executor.submit( + self._destroy_retired_groups, groups + ) + + def _wait_for_group_cleanup(self) -> None: + if (future := self._group_cleanup_future) is not None: + self._group_cleanup_future = None + future.result() + + def shutdown(self) -> None: + try: + self._wait_for_group_cleanup() + finally: + self._async_executor.shutdown() + def load_model(self) -> None: self.worker.load_model(load_dummy_weights=True) - def create_standby_groups( + def prepare_reconfiguration( self, reconfig_request: ReconfigureDistributedRequest, use_all2all: bool ) -> None: + self._wait_for_group_cleanup() self.reconfig_request = reconfig_request new_dp_size = reconfig_request.new_data_parallel_size old_dp_size = get_dp_group().world_size @@ -227,8 +260,20 @@ def create_standby_groups( use_all2all=use_all2all, enable_eplb=parallel_config.enable_eplb, ) - if new_dp_size < old_dp_size: - self.stage_standby_moe_quant_methods() + self.stage_standby_moe_quant_methods() + self._prepare_eplb_communicator(get_standby_eplb_group()) + if new_dp_size > old_dp_size: + self.transfer_weights(old_dp_size, new_dp_size) + self._warm_target_groups(get_standby_dp_group(), get_standby_ep_group()) + + def _prepare_eplb_communicator(self, eplb_group) -> None: + assert eplb_group is not None + model_runner = self.worker.model_runner + eplb_state = model_runner.eplb_state + assert eplb_state is not None + self._prepared_eplb_communicator = eplb_state.create_communicator( + model_runner.model_config, eplb_group + ) def transfer_weights(self, old_dp_size: int, new_dp_size: int) -> None: standby_dp_group = get_standby_dp_group() @@ -276,6 +321,15 @@ def transfer_weights(self, old_dp_size: int, new_dp_size: int) -> None: ) torch.accelerator.synchronize() + def _warm_target_groups(self, dp_group, ep_group) -> None: + assert dp_group is not None and ep_group is not None + stream = torch.Stream(device=dp_group.device) + with stream: + tensor = torch.zeros(1, dtype=torch.int32, device=dp_group.device) + for group in (dp_group, ep_group): + torch.distributed.all_reduce(tensor, group=group.device_group) + stream.synchronize() + def broadcast_expert_mapping(self) -> None: standby_dp_group = get_standby_dp_group() assert standby_dp_group is not None @@ -362,15 +416,22 @@ def _release_cuda_graphs(self) -> None: torch.accelerator.empty_cache() def switch_and_remove(self) -> None: + # Removing ranks skipped preparation, so wait for prior cleanup here. + self._wait_for_group_cleanup() self._release_cuda_graphs() - _replace_active_groups(world=None, dp=None, ep=None, eplb=None, node_count=None) + retired_groups = _replace_active_groups( + world=None, dp=None, ep=None, eplb=None, node_count=None + ) + self._start_group_cleanup(retired_groups) + # Finish collective cleanup before this worker is shut down. + self._wait_for_group_cleanup() - def switch_and_prepare(self) -> None: + def switch_and_prepare(self) -> tuple[GroupCoordinator | None, ...]: old_dp_size = get_dp_group().world_size old_ep_size = get_ep_group().world_size self._release_cuda_graphs() - _replace_active_groups(**pop_standby_groups()) + retired_groups = _replace_active_groups(**pop_standby_groups()) parallel_config = self.worker.vllm_config.parallel_config reconfig_request = self.reconfig_request @@ -453,7 +514,6 @@ def switch_and_prepare(self) -> None: ) eplb_model_state.expert_load_pass = expanded_expert_load_pass eplb_model_state.expert_load_window = expanded_expert_load_window - eplb_state.num_valid_physical_experts = old_num_physical_experts else: assert pad_size < 0 eplb_model_state.expert_load_pass = eplb_model_state.expert_load_pass[ @@ -462,7 +522,6 @@ def switch_and_prepare(self) -> None: eplb_model_state.expert_load_window = eplb_model_state.expert_load_window[ :, :, :num_physical_experts ] - eplb_state.num_valid_physical_experts = num_physical_experts model = self.worker.model_runner.get_model() model.expert_weights = [] @@ -485,18 +544,9 @@ def switch_and_prepare(self) -> None: if getattr(module._quant_method, "wraps_legacy_quant_method", False): module._replace_quant_method(module._quant_method.old_quant_method) - eplb_model_state.expert_buffer = [ - torch.empty_like(w) for w in model.expert_weights[0] - ] - assert parallel_config.eplb_config.communicator is not None, ( - "EPLB communicator backend must be set by ParallelConfig" - ) - eplb_model_state.communicator = create_eplb_communicator( - group_coordinator=get_eplb_group(), - backend=parallel_config.eplb_config.communicator, - expert_weights=model.expert_weights, - expert_buffer=eplb_model_state.expert_buffer, - ) + assert self._prepared_eplb_communicator is not None + eplb_state.update_communicator(model_config, self._prepared_eplb_communicator) + self._prepared_eplb_communicator = None if ( self.worker.vllm_config.compilation_config.mode @@ -513,8 +563,12 @@ def switch_and_prepare(self) -> None: compilation_counter.stock_torch_compile_count += 1 self.worker.model_runner.model.compile(fullgraph=True, backend=backend) + return retired_groups + def _perform_eplb_reshuffle( - self, rank_mapping: dict[int, int] | None = None + self, + rank_mapping: dict[int, int] | None = None, + async_op: bool = False, ) -> None: if get_ep_group().rank == 0: logger.info("[Elastic EP] Starting expert resharding...") @@ -522,21 +576,19 @@ def _perform_eplb_reshuffle( eplb_state = self.worker.model_runner.eplb_state assert eplb_state is not None - model_config = self.worker.model_runner.model_config - eplb_model_state = eplb_state.model_states[model_config.compute_hash()] is_async_enabled = eplb_state.is_async - eplb_state.is_async = False + run_async = async_op and is_async_enabled + eplb_state.is_async = run_async if rank_mapping is None: eplb_state.rearrange() else: eplb_state.rearrange(rank_mapping=rank_mapping) - # NOTE(yongji): check whether we need to synchronize here - torch.accelerator.synchronize() + if not run_async: + # Wait for non-blocking expert resharding copies before continuing + # the Elastic EP reconfiguration. + torch.accelerator.synchronize() # reset expert_rearrangement_step to ensure all ranks are synchronized eplb_state.expert_rearrangement_step = 0 - eplb_state.num_valid_physical_experts = ( - eplb_model_state.physical_to_logical_map.shape[1] - ) eplb_state.is_async = is_async_enabled # Start the async worker thread if it doesn't exist yet (idempotent). # This is needed for new workers after scale-up: they create EplbState @@ -544,25 +596,31 @@ def _perform_eplb_reshuffle( # groups aren't ready yet. eplb_state.start_async_loop() if get_ep_group().rank == 0: - logger.info("[Elastic EP] Expert resharding completed") + logger.info( + "[Elastic EP] Expert resharding %s", + "scheduled" if run_async else "completed", + ) def commit_scale_up(self, is_existing_worker: bool) -> None: if is_existing_worker: self.broadcast_expert_mapping() - self.switch_and_prepare() + retired_groups = self.switch_and_prepare() else: - mapping, _, num_valid_experts = self.receive_expert_mapping() - self.worker.model_runner.setup_eplb_from_mapping(mapping, num_valid_experts) - self._perform_eplb_reshuffle() + mapping = self.receive_expert_mapping() + self.worker.model_runner.setup_eplb_from_mapping(mapping) self.warm_and_capture() + self._perform_eplb_reshuffle(async_op=True) + if is_existing_worker: + self._start_group_cleanup(retired_groups) def commit_scale_down(self, new_dp_size: int, removing: bool) -> None: self.perform_scale_down_eplb_reshuffle(new_dp_size) if removing: self.switch_and_remove() else: - self.switch_and_prepare() + retired_groups = self.switch_and_prepare() self.warm_and_capture() + self._start_group_cleanup(retired_groups) def perform_scale_down_eplb_reshuffle(self, new_dp_size: int) -> None: eplb_state = self.worker.model_runner.eplb_state @@ -578,7 +636,7 @@ def perform_scale_down_eplb_reshuffle(self, new_dp_size: int) -> None: } self._perform_eplb_reshuffle(rank_mapping=rank_mapping) - def receive_weights(self) -> None: + def prepare_new_worker(self) -> None: dp_group = get_dp_group() assert isinstance(dp_group, StatelessGroupCoordinator) new_dp_size = dp_group.world_size @@ -618,19 +676,18 @@ def receive_weights(self) -> None: expert_weights=expert_weights, ) torch.accelerator.synchronize() + self._warm_target_groups(get_dp_group(), get_ep_group()) - def receive_expert_mapping(self) -> tuple[torch.Tensor, int, int]: + def receive_expert_mapping(self) -> torch.Tensor: dp_group = get_dp_group() assert isinstance(dp_group, StatelessGroupCoordinator) - physical_to_logical, num_local_physical_experts, num_logical_experts = ( - broadcast_expert_mapping( - physical_to_logical=None, - num_local_physical_experts=None, - num_logical_experts=None, - dp_group=dp_group, - src_rank=0, - device=self.worker.device, - ) + physical_to_logical, num_local_physical_experts, _ = broadcast_expert_mapping( + physical_to_logical=None, + num_local_physical_experts=None, + num_logical_experts=None, + dp_group=dp_group, + src_rank=0, + device=self.worker.device, ) num_moe_layers = physical_to_logical.shape[0] new_dp_size = get_dp_group().world_size @@ -642,13 +699,10 @@ def receive_expert_mapping(self) -> tuple[torch.Tensor, int, int]: dtype=physical_to_logical.dtype, device=physical_to_logical.device, ) - old_num_physical_experts = physical_to_logical.shape[1] - expanded_physical_to_logical[:, :old_num_physical_experts] = physical_to_logical - return ( - expanded_physical_to_logical, - num_logical_experts, - old_num_physical_experts, + expanded_physical_to_logical[:, : physical_to_logical.shape[1]] = ( + physical_to_logical ) + return expanded_physical_to_logical def warmup_local_kernels(self) -> None: with set_current_vllm_config(self.worker.vllm_config): diff --git a/vllm/distributed/elastic_ep/elastic_state.py b/vllm/distributed/elastic_ep/elastic_state.py index 1b15ba8dc710..6718cd8ae873 100644 --- a/vllm/distributed/elastic_ep/elastic_state.py +++ b/vllm/distributed/elastic_ep/elastic_state.py @@ -30,12 +30,10 @@ class ScaleUpExistingEngineState(enum.IntEnum): - CREATE_STANDBY_GROUPS = 0 - STAGE_QUANT_METHODS = 1 - TRANSFER_WEIGHTS = 2 - SYNC_KV_CACHE_MEMORY_SIZE = 3 - COMMIT_SCALE_UP = 4 # Blocks forward passes. - COMPLETE = 5 + PREPARE = 0 + SYNC_KV_CACHE_MEMORY_SIZE = 1 + COMMIT_SCALE_UP = 2 # Blocks forward passes. + COMPLETE = 3 class ScaleUpNewEngineState(enum.IntEnum): @@ -95,7 +93,7 @@ def __init__( self.state = ( ScaleUpNewEngineState.PRE_KV_INIT if worker_type == "new" - else ScaleUpExistingEngineState.CREATE_STANDBY_GROUPS + else ScaleUpExistingEngineState.PREPARE ) else: self.state = ( @@ -166,20 +164,8 @@ def _progress_existing_engine(self) -> bool: state = self.state assert self.old_dp_group is not None - if state == ScaleUpExistingEngineState.CREATE_STANDBY_GROUPS: - if not self._create_standby_groups(): - return False - self.state = ScaleUpExistingEngineState.STAGE_QUANT_METHODS - return True - - elif state == ScaleUpExistingEngineState.STAGE_QUANT_METHODS: - if not self._execute_async("stage_standby_moe_quant_methods"): - return False - self.state = ScaleUpExistingEngineState.TRANSFER_WEIGHTS - return True - - elif state == ScaleUpExistingEngineState.TRANSFER_WEIGHTS: - if not self._transfer_weights(): + if state == ScaleUpExistingEngineState.PREPARE: + if not self._prepare_workers(): return False self.state = ScaleUpExistingEngineState.SYNC_KV_CACHE_MEMORY_SIZE return True @@ -210,7 +196,7 @@ def _progress_new_engine(self) -> bool: assert self.new_dp_group is not None and self.new_dp_store is not None if state == ScaleUpNewEngineState.PRE_KV_INIT: - self._collective_rpc("elastic_ep_execute", args=("receive_weights",)) + self._collective_rpc("elastic_ep_execute", args=("prepare_new_worker",)) self.engine_core.available_gpu_memory_for_kv_cache = ( ParallelConfig.sync_kv_cache_memory_size(self.new_dp_group, -1) ) @@ -243,7 +229,7 @@ def _progress_remaining_engine(self) -> bool: assert self.old_dp_group is not None if state == ScaleDownRemainingEngineState.PREPARE: - if self._create_standby_groups(): + if self._prepare_workers(): self.state = ScaleDownRemainingEngineState.COMMIT_SCALE_DOWN self._mark_ready_for_switch() return True @@ -328,29 +314,18 @@ def _ensure_new_dp_group(self) -> bool: self._prepare_future = None return True - def _create_standby_groups(self) -> bool: + def _prepare_workers(self) -> bool: assert self.old_dp_group is not None if not self._ensure_new_dp_group(): return False if not self._execute_async( - "create_standby_groups", + "prepare_reconfiguration", self.reconfig_request, self.new_parallel_config.use_all2all, ): return False if self.old_dp_group.rank() == 0: - logger.info("[Elastic EP] Created standby communication groups") - return True - - def _transfer_weights(self) -> bool: - assert self.reconfig_request is not None and self.old_dp_group is not None - old_dp_size = self.old_dp_group.size() - new_dp_size = self.reconfig_request.new_data_parallel_size - - if not self._execute_async("transfer_weights", old_dp_size, new_dp_size): - return False - if self.old_dp_group.rank() == 0: - logger.info("[Elastic EP] Transferred weights to new workers") + logger.info("[Elastic EP] Prepared reconfiguration") return True def _sync_kv_cache_memory_size(self) -> bool: diff --git a/vllm/distributed/eplb/eplb_communicator.py b/vllm/distributed/eplb/eplb_communicator.py index 3f04b7bc91ed..14efd16177cc 100644 --- a/vllm/distributed/eplb/eplb_communicator.py +++ b/vllm/distributed/eplb/eplb_communicator.py @@ -248,7 +248,6 @@ def __init__( cpu_group: ProcessGroup, all_expert_weights: Sequence[Sequence[torch.Tensor]], expert_buffer: Sequence[torch.Tensor], - defer_remote_setup: bool = False, ) -> None: """Create a NIXL-backed EPLB communicator. @@ -256,11 +255,6 @@ def __init__( cpu_group: CPU process group for metadata exchange. all_expert_weights: Expert weight tensors for all MoE layers. expert_buffer: Pre-allocated receive buffer tensors. - defer_remote_setup: If True, postpone the collective - all-gather of NIXL agent metadata until the first - ``set_transfer_context`` call. Required for elastic EP - where ranks join asynchronously and cannot participate - in collectives at construction time. """ assert all_expert_weights, ( "NixlEplbCommunicator requires non-empty all_expert_weights." @@ -317,29 +311,17 @@ def __init__( ] = {} self._cuda_device_id = int(self._device.index or 0) - self._remote_state_initialized = False self._init_step("buffers", self._init_registered_buffers) - if defer_remote_setup: - logger.info_once("NIXL EPLB: deferring remote agent setup (elastic EP).") - else: - self._init_remote_state() + self._init_remote_state() self._log_initialized() def _init_remote_state(self) -> None: """Exchange NIXL agent metadata and RDMA pointer info with all peers. This is a collective operation (uses ``all_gather_object`` twice). - Under elastic EP the call is deferred to the first - ``set_transfer_context`` invocation, where all ranks are - guaranteed to be synchronized. """ self._init_step("agents", self._init_remote_agents) self._init_step("send meta", self._exchange_remote_send_meta) - self._remote_state_initialized = True - - def _ensure_remote_state(self) -> None: - if not self._remote_state_initialized: - self._init_remote_state() @property def needs_profile_buffer_reservation(self) -> bool: @@ -373,7 +355,6 @@ def add_send( pass def set_transfer_context(self, old_indices: np.ndarray, layer_idx: int) -> None: - self._ensure_remote_state() assert not self._xfer_entries, ( f"set_transfer_context() called with {len(self._xfer_entries)} " f"pending transfers from layer {self._layer_idx}; " @@ -673,9 +654,8 @@ def create_eplb_communicator( Falls back to ``"torch_nccl"`` when *None*. Stateless (elastic EP) groups support ``"torch_nccl"``, ``"pynccl"``, and ``"nixl"``; ``"torch_nccl"`` is silently - promoted to ``"pynccl"``. ``"nixl"`` uses deferred remote - agent setup to avoid collective deadlocks during elastic - scaling. When tensors reside on CPU, ``"torch_gloo"`` or + promoted to ``"pynccl"``. When tensors reside on CPU, + ``"torch_gloo"`` or ``"torch_nccl"`` are used via the CPU process group. expert_weights: Expert weight tensors for *all* MoE layers. Shape ``(num_layers)(num_tensors_per_layer)``. @@ -729,22 +709,19 @@ def _create_pynccl() -> EplbCommunicator: ) from exc is_stateless = isinstance(group_coordinator, StatelessGroupCoordinator) - if is_stateless: - if backend == "nixl": - pass # handled below with defer_remote_setup=True - elif backend not in ("torch_nccl", "pynccl"): + if is_stateless and backend != "nixl": + if backend not in ("torch_nccl", "pynccl"): raise ValueError( f"Elastic EP requires 'torch_nccl', 'pynccl', or 'nixl' " f"EPLB communicator (got '{backend}')." ) - else: - if backend == "torch_nccl": - logger.warning( - "Stateless elastic EP requires PyNCCL backend. " - "Forcing EPLB communicator to 'pynccl'." - ) - backend = "pynccl" - return _create_pynccl() + if backend == "torch_nccl": + logger.warning( + "Stateless elastic EP requires PyNCCL backend. " + "Forcing EPLB communicator to 'pynccl'." + ) + backend = "pynccl" + return _create_pynccl() if backend == "nixl": if not has_nixl(): @@ -761,7 +738,6 @@ def _create_pynccl() -> EplbCommunicator: cpu_group=group_coordinator.cpu_group, all_expert_weights=expert_weights, expert_buffer=expert_buffer, - defer_remote_setup=is_stateless, ) except Exception as exc: raise RuntimeError( diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py index 5b2bb05ed0aa..0e3e38b25369 100644 --- a/vllm/distributed/eplb/eplb_state.py +++ b/vllm/distributed/eplb/eplb_state.py @@ -37,6 +37,7 @@ from vllm.config import ModelConfig, ParallelConfig from vllm.config.utils import compute_hash_cached from vllm.distributed.parallel_state import ( + GroupCoordinator, get_ep_group, get_eplb_group, get_node_count, @@ -289,14 +290,6 @@ def __init__(self, parallel_config: ParallelConfig, device: torch.device): """ CUDA device index for the async EPLB worker thread. """ - self.num_valid_physical_experts: int = 0 - """ - Number of valid physical experts. - This is the number of physical experts that are - actually mapped to logical experts. In elastic EP, - newly started EP ranks may not have physical experts - mapped yet. - """ if self.device.type == "cuda": self.cuda_device_index = self.device.index if self.cuda_device_index is None and torch.cuda.is_available(): @@ -499,7 +492,6 @@ def add_model( num_unpadded_tokens_tensors=num_unpadded_tokens_tensors, ) self.model_states[model_config.compute_hash()] = model_state - self.num_valid_physical_experts = model.num_physical_experts def prepare_forward( self, @@ -766,28 +758,28 @@ def rearrange( # Map the physical expert load to global logical experts global_expert_load_windows = [] for eplb_model_state in self.model_states.values(): - expert_load_window = eplb_model_state.expert_load_window[ - :, :, : self.num_valid_physical_experts - ] + expert_load_window = eplb_model_state.expert_load_window + physical_to_logical = eplb_model_state.physical_to_logical_map + invalid_idx = eplb_model_state.model.num_logical_experts logical_expert_load_window = torch.zeros( self.expert_load_window_size, eplb_model_state.model.num_moe_layers, - eplb_model_state.model.num_logical_experts, + invalid_idx + 1, dtype=eplb_model_state.expert_load_window.dtype, device=eplb_model_state.expert_load_window.device, ) logical_expert_load_window.scatter_add_( dim=-1, - index=eplb_model_state.physical_to_logical_map[ - :, : self.num_valid_physical_experts - ] + index=physical_to_logical.masked_fill( + physical_to_logical < 0, invalid_idx + ) .unsqueeze(0) .expand_as(expert_load_window) .long(), src=expert_load_window, ) - global_expert_load_window = logical_expert_load_window.sum(dim=0) + global_expert_load_window = logical_expert_load_window[..., :-1].sum(dim=0) global_expert_load_windows.append(global_expert_load_window) # Perform all-reduce to get the expert load across all ranks for each model global_expert_load_windows = self._allreduce_list(global_expert_load_windows) @@ -960,7 +952,7 @@ def drain_async(self) -> None: Each pending result is acknowledged (consumed_event recorded) so the async worker can proceed, but the transferred weights are intentionally - NOT applied — a full synchronous rearrange is expected to follow. + NOT applied — a full rearrange is expected to follow. Ranks are kept in lockstep via _all_ranks_result_ready (all_reduce on the EP CPU group). The async worker's coordinated-stop collectives @@ -1060,7 +1052,6 @@ def from_mapping( device: torch.device, parallel_config: ParallelConfig, expanded_physical_to_logical: torch.Tensor, - num_valid_physical_experts: int, ) -> "EplbState": eplb_state = cls( parallel_config=parallel_config, @@ -1070,12 +1061,24 @@ def from_mapping( model=model, model_config=model_config, ) - eplb_state.num_valid_physical_experts = num_valid_physical_experts - eplb_model_state = eplb_state.model_states[model_config.compute_hash()] + eplb_state.update_mapping( + model_config, + expanded_physical_to_logical, + ) + + return eplb_state + + def update_mapping( + self, + model_config: ModelConfig, + expanded_physical_to_logical: torch.Tensor, + ) -> None: + eplb_model_state = self.model_states[model_config.compute_hash()] eplb_model_state.physical_to_logical_map.copy_(expanded_physical_to_logical) (logical_to_physical_map_cpu, logical_replica_count_cpu) = compute_logical_maps( - expanded_physical_to_logical.cpu(), model.num_logical_experts + expanded_physical_to_logical.cpu(), + eplb_model_state.model.num_logical_experts, ) max_num_replicas = eplb_model_state.logical_to_physical_map.shape[-1] @@ -1087,13 +1090,31 @@ def from_mapping( max_num_replicas - num_replicas, ), value=-1, - ).to(device) - logical_replica_count = logical_replica_count_cpu.to(device) + ).to(self.device) + logical_replica_count = logical_replica_count_cpu.to(self.device) eplb_model_state.logical_to_physical_map.copy_(logical_to_physical_map) eplb_model_state.logical_replica_count.copy_(logical_replica_count) - return eplb_state + def create_communicator( + self, model_config: ModelConfig, group_coordinator: GroupCoordinator + ) -> EplbCommunicator: + model_state = self.model_states[model_config.compute_hash()] + backend = self.parallel_config.eplb_config.communicator + assert backend is not None + return create_eplb_communicator( + group_coordinator, + backend, + model_state.model.expert_weights, + model_state.expert_buffer, + ) + + def update_communicator( + self, + model_config: ModelConfig, + communicator: EplbCommunicator, + ) -> None: + self.model_states[model_config.compute_hash()].communicator = communicator @dataclass diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index 128db903d4d6..96a95a69fd54 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -1366,21 +1366,21 @@ def _replace_active_groups( ep: GroupCoordinator | None, eplb: GroupCoordinator | None, node_count: int | None, -) -> None: - """Destroy the current DP/EP/WORLD/EPLB groups and replace them. +) -> tuple[GroupCoordinator | None, ...]: + """Replace the active groups and return the groups they replaced. - Destruction is collective — all ranks in the old groups must call this - function together. Pass all-``None`` to tear down without replacement. + The caller must destroy the returned DP, EP, WORLD, and EPLB groups + collectively and in that order. Pass all-``None`` to remove the active + groups without replacement. """ global _WORLD, _DP, _EP, _EPLB, _NODE_COUNT - for group in (_DP, _EP, _WORLD, _EPLB): - if group is not None: - group.destroy() + old_groups = _DP, _EP, _WORLD, _EPLB _WORLD = world _DP = dp _EP = ep _EPLB = eplb _NODE_COUNT = node_count + return old_groups _TP: GroupCoordinator | None = None diff --git a/vllm/v1/worker/gpu/eplb_utils.py b/vllm/v1/worker/gpu/eplb_utils.py index 57a4e6db2c23..24c0f0c4e1c1 100644 --- a/vllm/v1/worker/gpu/eplb_utils.py +++ b/vllm/v1/worker/gpu/eplb_utils.py @@ -144,7 +144,6 @@ def setup_from_mapping( model: nn.Module, model_config: Any, expanded_physical_to_logical: torch.Tensor, - old_num_physical_experts: int, ) -> None: moe_model = get_mixture_of_experts_model(model) assert moe_model is not None @@ -155,6 +154,5 @@ def setup_from_mapping( device=self.device, parallel_config=self.parallel_config, expanded_physical_to_logical=expanded_physical_to_logical, - num_valid_physical_experts=old_num_physical_experts, ) self._has_registered_models = True diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 4f65769c9580..3861bb52649d 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1962,13 +1962,11 @@ def eep_eplb_suppressed(self, suppressed: bool) -> None: def setup_eplb_from_mapping( self, expanded_physical_to_logical: torch.Tensor, - old_num_physical_experts: int, ) -> None: self.eplb.setup_from_mapping( self.model, self.model_config, expanded_physical_to_logical, - old_num_physical_experts, ) ########### EPLB methods end ########### diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index a8d08a7bf293..ae11b2b602bb 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -3517,17 +3517,11 @@ def eplb_step(self, is_dummy: bool = False, is_profile: bool = False) -> None: def setup_eplb_from_mapping( self, expanded_physical_to_logical: torch.Tensor, - old_num_physical_experts: int, ) -> None: - assert self._moe_model is not None - - self.eplb_state = EplbState.from_mapping( - model=self._moe_model, - model_config=self.model_config, - device=self.device, - parallel_config=self.parallel_config, - expanded_physical_to_logical=expanded_physical_to_logical, - num_valid_physical_experts=old_num_physical_experts, + assert self.eplb_state is not None + self.eplb_state.update_mapping( + self.model_config, + expanded_physical_to_logical, ) def _pool( @@ -5488,11 +5482,7 @@ def load_model(self, load_dummy_weights: bool = False) -> None: # MixtureOfExperts themselves. self._moe_model = get_mixture_of_experts_model(self.model) - if ( - self.parallel_config.enable_eplb - and not load_dummy_weights - and self._moe_model is not None - ): + if self._moe_model is not None and self.parallel_config.enable_eplb: logger.info_once( "EPLB is enabled for MoE part of model %s.", self.model_config.model, diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index a3b00aaad2a2..958972f3c54e 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -1376,6 +1376,8 @@ def shutdown(self) -> None: if weight_transfer_engine := getattr(self, "weight_transfer_engine", None): weight_transfer_engine.shutdown() + self.elastic_ep_executor.shutdown() + # Release GPU resources held by the model runner so that memory # can be reclaimed when running in-process if model_runner := getattr(self, "model_runner", None): From b1d9337e92692da328bdf921286e59d8f8f36725 Mon Sep 17 00:00:00 2001 From: liuzhenwei Date: Wed, 19 Aug 2026 10:16:08 +0800 Subject: [PATCH 120/839] [EPD] Allow KV consumers to omit MM embeddings (#52697) Signed-off-by: zhenwei-intel --- .../integration/run_epd_correctness_test.sh | 3 +++ vllm/config/multimodal.py | 12 ++++----- vllm/config/vllm.py | 26 +++++++++---------- vllm/model_executor/models/colqwen3.py | 2 +- vllm/model_executor/models/colqwen3_5.py | 2 +- vllm/model_executor/models/hunyuan_vision.py | 2 +- vllm/model_executor/models/keye.py | 2 +- vllm/model_executor/models/keye_vl1_5.py | 2 +- .../model_executor/models/llava_onevision2.py | 2 +- vllm/model_executor/models/minicpmv.py | 2 +- vllm/model_executor/models/opencua.py | 2 +- .../models/qwen2_5_omni_thinker.py | 2 +- vllm/model_executor/models/qwen2_vl.py | 2 +- vllm/model_executor/models/qwen3_vl.py | 2 +- vllm/multimodal/parse.py | 13 +++++----- vllm/multimodal/processing/context.py | 13 +++------- 16 files changed, 43 insertions(+), 46 deletions(-) diff --git a/tests/v1/ec_connector/integration/run_epd_correctness_test.sh b/tests/v1/ec_connector/integration/run_epd_correctness_test.sh index c58df0c076a5..b48949f92210 100644 --- a/tests/v1/ec_connector/integration/run_epd_correctness_test.sh +++ b/tests/v1/ec_connector/integration/run_epd_correctness_test.sh @@ -183,6 +183,7 @@ run_epd_1e_1pd() { --port "$PREFILL_DECODE_PORT" \ --max-model-len "$MAX_MODEL_LEN" \ --enforce-eager \ + --enable-mm-embeds \ --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \ --enable-request-id-headers \ --max-num-seqs "$MAX_NUM_SEQS" \ @@ -389,6 +390,7 @@ run_epd_1e_1p_1d() { --port "$PREFILL_PORT" \ --max-model-len "$MAX_MODEL_LEN" \ --enforce-eager \ + --enable-mm-embeds \ --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \ --enable-request-id-headers \ --max-num-seqs "$MAX_NUM_SEQS" \ @@ -415,6 +417,7 @@ run_epd_1e_1p_1d() { --port "$DECODE_PORT" \ --max-model-len "$MAX_MODEL_LEN" \ --enforce-eager \ + --enable-mm-embeds \ --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \ --enable-request-id-headers \ --max-num-seqs "$MAX_NUM_SEQS" \ diff --git a/vllm/config/multimodal.py b/vllm/config/multimodal.py index e1343bc3524f..04e94735a1cd 100644 --- a/vllm/config/multimodal.py +++ b/vllm/config/multimodal.py @@ -236,17 +236,17 @@ class MultiModalConfig: - "direct_rpc": Use msgspec serialization via RPC - "torch_shm": Use torch.multiprocessing shared memory for zero-copy IPC Defaults to "direct_rpc". """ - mm_embeds_from_ec_connector: bool = False + allow_missing_mm_embeddings: bool = False """Whether a pre-computed-embedding input may omit the `*_embeds` tensor. In an encode/prefill/decode (EPD) deployment the encoder instance publishes - embeddings through the EC connector, so the request that reaches the - prefill/decode instance only needs to carry the grid/size metadata that - sizes the placeholder range — the embeddings themselves come from the - connector, keyed by `mm_hash`. + embeddings through the EC connector. An EC consumer loads those embeddings + from the connector, while a KV consumer receives the resulting prompt KV + cache. Their requests only need the grid/size metadata that sizes the + placeholder range. Derived, not user-settable: `VllmConfig.__post_init__` sets this to True - exactly on EC consumers. Everywhere else it stays False so that a request + on EC and KV consumers. Everywhere else it stays False so that a request which forgets its embeddings still fails fast in the frontend, with a clear error, rather than deep inside the model.""" diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 5faf09fc347e..ff93bed86fb6 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1577,7 +1577,7 @@ def has_blocked_weights(): ) current_platform.check_and_update_config(self) - self._resolve_mm_embeds_from_ec_connector() + self._resolve_allow_missing_mm_embeddings() self._resolve_mm_processor_device() self._validate_mm_processor_device() @@ -2254,14 +2254,13 @@ def __str__(self): f"kernel_config={self.kernel_config!r}" ) - def _resolve_mm_embeds_from_ec_connector(self) -> None: - """Allow `*_embeds` to be omitted only where the connector supplies them. + def _resolve_allow_missing_mm_embeddings(self) -> None: + """Allow `*_embeds` tensors to be omitted on disaggregated consumers. - That is exactly an EC consumer: the encoder instance publishes the - embeddings through the EC connector, so the request only has to carry - the grid metadata that sizes the placeholder range. On every other - deployment a missing `*_embeds` is a client error and must keep failing - fast in the frontend. + An EC consumer loads embeddings from its connector. A KV consumer + receives the prompt KV produced from those embeddings, so it does not + need the tensors either. On every other deployment a missing tensor is + a client error and must keep failing fast in the frontend. """ model_config = self.model_config if model_config is None: @@ -2271,15 +2270,16 @@ def _resolve_mm_embeds_from_ec_connector(self) -> None: return ec_config = self.ec_transfer_config + kv_config = self.kv_transfer_config # Derived, so overwrite unconditionally rather than honouring a value # that was set by hand. - mm_config.mm_embeds_from_ec_connector = ( + mm_config.allow_missing_mm_embeddings = ( ec_config is not None and ec_config.is_ec_consumer - ) - if mm_config.mm_embeds_from_ec_connector: + ) or (kv_config is not None and kv_config.is_kv_consumer) + if mm_config.allow_missing_mm_embeddings: logger.info_once( - "EC consumer: pre-computed-embedding inputs may omit the " - "embedding tensor; embeddings are loaded from the EC connector." + "EC/KV consumer: pre-computed-embedding inputs may " + "omit the embedding tensor." ) def _resolve_mm_processor_device(self) -> None: diff --git a/vllm/model_executor/models/colqwen3.py b/vllm/model_executor/models/colqwen3.py index dee975ae2045..a2012fffd2ad 100644 --- a/vllm/model_executor/models/colqwen3.py +++ b/vllm/model_executor/models/colqwen3.py @@ -104,7 +104,7 @@ def get_data_parser(self): spatial_merge_size, video_needs_metadata=self._supports_video, expected_hidden_size=self._get_expected_hidden_size(), - embeds_from_ec_connector=self.embeds_from_ec_connector, + allow_missing_mm_embeddings=self.allow_missing_mm_embeddings, ) diff --git a/vllm/model_executor/models/colqwen3_5.py b/vllm/model_executor/models/colqwen3_5.py index 309da6d68688..760a54f485f7 100644 --- a/vllm/model_executor/models/colqwen3_5.py +++ b/vllm/model_executor/models/colqwen3_5.py @@ -101,7 +101,7 @@ def get_data_parser(self): spatial_merge_size, video_needs_metadata=self._supports_video, expected_hidden_size=self._get_expected_hidden_size(), - embeds_from_ec_connector=self.embeds_from_ec_connector, + allow_missing_mm_embeddings=self.allow_missing_mm_embeddings, ) diff --git a/vllm/model_executor/models/hunyuan_vision.py b/vllm/model_executor/models/hunyuan_vision.py index d1c19f07f52f..74e4c1a6cec4 100644 --- a/vllm/model_executor/models/hunyuan_vision.py +++ b/vllm/model_executor/models/hunyuan_vision.py @@ -599,7 +599,7 @@ def get_image_processor( def get_data_parser(self): return HunYuanVLMultiModalDataParser( expected_hidden_size=self._get_expected_hidden_size(), - embeds_from_ec_connector=self.embeds_from_ec_connector, + allow_missing_mm_embeddings=self.allow_missing_mm_embeddings, ) def get_supported_mm_limits(self) -> Mapping[str, int | None]: diff --git a/vllm/model_executor/models/keye.py b/vllm/model_executor/models/keye.py index 28ccddd47e9f..2fee80b271c3 100644 --- a/vllm/model_executor/models/keye.py +++ b/vllm/model_executor/models/keye.py @@ -936,7 +936,7 @@ def get_image_processor(self, **kwargs: object): def get_data_parser(self): return KeyeMultiModalDataParser( expected_hidden_size=self._get_expected_hidden_size(), - embeds_from_ec_connector=self.embeds_from_ec_connector, + allow_missing_mm_embeddings=self.allow_missing_mm_embeddings, ) def get_supported_mm_limits( diff --git a/vllm/model_executor/models/keye_vl1_5.py b/vllm/model_executor/models/keye_vl1_5.py index 26870ebc26ac..c7ecc68db1f2 100644 --- a/vllm/model_executor/models/keye_vl1_5.py +++ b/vllm/model_executor/models/keye_vl1_5.py @@ -363,7 +363,7 @@ class KeyeVL1_5ProcessingInfo(KeyeProcessingInfo): def get_data_parser(self): return KeyeVL1_5MultiModalDataParser( expected_hidden_size=self._get_expected_hidden_size(), - embeds_from_ec_connector=self.embeds_from_ec_connector, + allow_missing_mm_embeddings=self.allow_missing_mm_embeddings, ) def get_max_frame_per_video(self) -> int: diff --git a/vllm/model_executor/models/llava_onevision2.py b/vllm/model_executor/models/llava_onevision2.py index daae740c3851..01f073a0d92c 100644 --- a/vllm/model_executor/models/llava_onevision2.py +++ b/vllm/model_executor/models/llava_onevision2.py @@ -1279,7 +1279,7 @@ def get_data_parser(self) -> MultiModalDataParser: return LlavaOnevision2MultiModalDataParser( self.get_hf_config().vision_config.spatial_merge_size, video_needs_metadata=True, - embeds_from_ec_connector=self.embeds_from_ec_connector, + allow_missing_mm_embeddings=self.allow_missing_mm_embeddings, ) def get_hf_processor(self, **kwargs: object): diff --git a/vllm/model_executor/models/minicpmv.py b/vllm/model_executor/models/minicpmv.py index d9358fed425f..e85ac81ee583 100644 --- a/vllm/model_executor/models/minicpmv.py +++ b/vllm/model_executor/models/minicpmv.py @@ -653,7 +653,7 @@ def get_image_processor(self, **kwargs: object): def get_data_parser(self): return MiniCPMVMultiModalDataParser( expected_hidden_size=self._get_expected_hidden_size(), - embeds_from_ec_connector=self.embeds_from_ec_connector, + allow_missing_mm_embeddings=self.allow_missing_mm_embeddings, ) def get_model_version(self): diff --git a/vllm/model_executor/models/opencua.py b/vllm/model_executor/models/opencua.py index 9fb0b526b39c..a34475539b42 100644 --- a/vllm/model_executor/models/opencua.py +++ b/vllm/model_executor/models/opencua.py @@ -57,7 +57,7 @@ def get_data_parser(self): return Qwen2VLMultiModalDataParser( self.get_hf_config().vision_config.spatial_merge_size, expected_hidden_size=self._get_expected_hidden_size(), - embeds_from_ec_connector=self.embeds_from_ec_connector, + allow_missing_mm_embeddings=self.allow_missing_mm_embeddings, ) def get_hf_config(self): diff --git a/vllm/model_executor/models/qwen2_5_omni_thinker.py b/vllm/model_executor/models/qwen2_5_omni_thinker.py index 96a58df1ddaa..a512be850dde 100644 --- a/vllm/model_executor/models/qwen2_5_omni_thinker.py +++ b/vllm/model_executor/models/qwen2_5_omni_thinker.py @@ -373,7 +373,7 @@ def get_data_parser(self): target_sr=feature_extractor.sampling_rate, target_channels=self.get_target_channels(), expected_hidden_size=self._get_expected_hidden_size(), - embeds_from_ec_connector=self.embeds_from_ec_connector, + allow_missing_mm_embeddings=self.allow_missing_mm_embeddings, ) def get_target_channels(self) -> int: diff --git a/vllm/model_executor/models/qwen2_vl.py b/vllm/model_executor/models/qwen2_vl.py index e4a56eadfc7e..49e026534ca1 100644 --- a/vllm/model_executor/models/qwen2_vl.py +++ b/vllm/model_executor/models/qwen2_vl.py @@ -857,7 +857,7 @@ def get_data_parser(self): return Qwen2VLMultiModalDataParser( self.get_hf_config().vision_config.spatial_merge_size, expected_hidden_size=self._get_expected_hidden_size(), - embeds_from_ec_connector=self.embeds_from_ec_connector, + allow_missing_mm_embeddings=self.allow_missing_mm_embeddings, ) def get_supported_mm_limits(self) -> Mapping[str, int | None]: diff --git a/vllm/model_executor/models/qwen3_vl.py b/vllm/model_executor/models/qwen3_vl.py index de341d302a35..3f383b9bd8a1 100644 --- a/vllm/model_executor/models/qwen3_vl.py +++ b/vllm/model_executor/models/qwen3_vl.py @@ -910,7 +910,7 @@ def get_data_parser(self): self.get_hf_config().vision_config.spatial_merge_size, video_needs_metadata=True, expected_hidden_size=self._get_expected_hidden_size(), - embeds_from_ec_connector=self.embeds_from_ec_connector, + allow_missing_mm_embeddings=self.allow_missing_mm_embeddings, ) def _get_vision_info( diff --git a/vllm/multimodal/parse.py b/vllm/multimodal/parse.py index 3053670a647b..6306d6a1362d 100644 --- a/vllm/multimodal/parse.py +++ b/vllm/multimodal/parse.py @@ -554,10 +554,9 @@ class MultiModalDataParser: embedding inputs. If provided, validates that user-supplied embeddings have the correct hidden size to prevent crashes during model inference. - embeds_from_ec_connector (bool): Whether pre-computed embeddings may be - absent from the request because an encode/prefill/decode encoder - instance publishes them through an EC connector instead. Derived by - `BaseProcessingInfo.embeds_from_ec_connector`. + allow_missing_mm_embeddings (bool): Whether pre-computed embedding + tensors may be absent from the request on a disaggregated consumer. + Derived by `BaseProcessingInfo.allow_missing_mm_embeddings`. """ embedding_fields: Mapping[str, Mapping[str, EmbeddingFieldRole]] = {} @@ -594,7 +593,7 @@ def embedding_field_sets(self, modality: str) -> tuple[set[str], set[str]]: """ metadata = self.placeholder_metadata_fields(modality) values = set(self.embedding_fields.get(modality, {})) - metadata - if self.embeds_from_ec_connector: + if self.allow_missing_mm_embeddings: return metadata, values return metadata | values, set() @@ -606,11 +605,11 @@ def __init__( audio_resample_method: Literal["pyav", "scipy", "soxr"] = "pyav", video_needs_metadata: bool = False, expected_hidden_size: int | None = None, - embeds_from_ec_connector: bool = False, + allow_missing_mm_embeddings: bool = False, ) -> None: super().__init__() - self.embeds_from_ec_connector = embeds_from_ec_connector + self.allow_missing_mm_embeddings = allow_missing_mm_embeddings self.audio_resampler = AudioResampler( target_sr=target_sr, diff --git a/vllm/multimodal/processing/context.py b/vllm/multimodal/processing/context.py index a1d495867415..45acdf8f4d7b 100644 --- a/vllm/multimodal/processing/context.py +++ b/vllm/multimodal/processing/context.py @@ -367,15 +367,10 @@ def _get_expected_hidden_size(self) -> int | None: return None @property - def embeds_from_ec_connector(self) -> bool: - """Whether pre-computed embeddings may arrive outside the request. - - True only on an EC consumer, where an encode/prefill/decode encoder - instance publishes them through the connector instead, so the request - carries only the metadata that sizes the placeholder range. - """ + def allow_missing_mm_embeddings(self) -> bool: + """Whether pre-computed embedding tensors may be omitted.""" mm_config = self.ctx.model_config.multimodal_config - return mm_config is not None and mm_config.mm_embeds_from_ec_connector + return mm_config is not None and mm_config.allow_missing_mm_embeddings def get_data_parser(self) -> MultiModalDataParser: """ @@ -389,7 +384,7 @@ def get_data_parser(self) -> MultiModalDataParser: """ return MultiModalDataParser( expected_hidden_size=self._get_expected_hidden_size(), - embeds_from_ec_connector=self.embeds_from_ec_connector, + allow_missing_mm_embeddings=self.allow_missing_mm_embeddings, ) @cached_property From f485081e8b76135091398818ea63040ffe7e1c5b Mon Sep 17 00:00:00 2001 From: Chaojun Zhang Date: Wed, 19 Aug 2026 10:25:23 +0800 Subject: [PATCH 121/839] [XPU] Support EC connector KV Offloading on XPU (#49532) Signed-off-by: Chaojun Zhang --- .buildkite/intel_jobs/misc_intel.yaml | 6 +- .../unit/cpu/worker/test_worker.py | 79 ++++++++++--------- .../unit/test_ec_cpu_connector.py | 11 ++- .../ec_connector/cpu/worker/__init__.py | 16 ++-- .../cpu/worker/descriptor_buffers.py | 39 +++++++-- 5 files changed, 93 insertions(+), 58 deletions(-) diff --git a/.buildkite/intel_jobs/misc_intel.yaml b/.buildkite/intel_jobs/misc_intel.yaml index 679e626b3ed6..897ae7d22814 100644 --- a/.buildkite/intel_jobs/misc_intel.yaml +++ b/.buildkite/intel_jobs/misc_intel.yaml @@ -8,7 +8,7 @@ steps: agent_tags: label: production gpu: 1+ - mem: 16+ + mem: 24+ no_plugin: true working_dir: "." env: @@ -22,6 +22,7 @@ steps: - tests/v1/kv_offload - tests/v1/worker - tests/v1/kv_connector/unit + - tests/v1/ec_connector/unit - tests/v1/metrics - tests/entrypoints/openai/correctness/test_lmeval.py commands: @@ -30,7 +31,8 @@ steps: 'pip install -r requirements/kv_connectors.txt && export VLLM_WORKER_MULTIPROC_METHOD=spawn && cd tests && - pytest -v -s v1/executor' + pytest -v -s v1/executor && + VLLM_BATCH_INVARIANT=1 pytest -v -s -m 'not cpu_test' v1/ec_connector/unit' - label: V1 Sample + Logits timeout_in_minutes: 90 diff --git a/tests/v1/ec_connector/unit/cpu/worker/test_worker.py b/tests/v1/ec_connector/unit/cpu/worker/test_worker.py index 4668abf25127..3daf4ee7de1d 100644 --- a/tests/v1/ec_connector/unit/cpu/worker/test_worker.py +++ b/tests/v1/ec_connector/unit/cpu/worker/test_worker.py @@ -2,10 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Tests for ECCPUWorker. -The byte-level tests exercise real CUDA stream/event coordination -against a real ``ECSharedRegion`` mmap and are skipped on hosts without -CUDA. The lifecycle tests don't fire any ``torch.cuda.*`` primitives and -run anywhere. +The byte-level tests exercise real accelerator (CUDA/XPU/...) stream/event +coordination against a real ``ECSharedRegion`` mmap and are skipped on hosts +without an accelerator. The lifecycle tests don't fire any +``current_platform.*`` device primitives and run anywhere. Mocking policy -------------- @@ -37,6 +37,7 @@ ECSharedRegion, ) from vllm.distributed.ec_transfer.ec_connector.cpu.worker import ECCPUWorker +from vllm.platforms import current_platform # ── shape constants ────────────────────────────────────────────────────────── @@ -47,9 +48,11 @@ _BLOCK_SIZE_BYTES = _HIDDEN_DIM * _DTYPE.itemsize _NUM_BLOCKS = 8 -_requires_cuda = pytest.mark.skipif( - not torch.cuda.is_available(), - reason="exercises real CUDA stream/event coordination in ECCPUWorker", +DEVICE_TYPE = current_platform.device_type + +_requires_accelerator = pytest.mark.skipif( + not (current_platform.is_cuda() or current_platform.is_xpu()), + reason="exercises real accelerator stream/event coordination in ECCPUWorker", ) @@ -136,7 +139,7 @@ def factory( # ── save_caches ────────────────────────────────────────────────────────────── -@_requires_cuda +@_requires_accelerator @pytest.mark.parametrize( "n_elements,n_blocks", [ @@ -158,7 +161,7 @@ def test_save_caches_writes_to_assigned_blocks(make_worker, n_elements, n_blocks sentinel = 0x5A worker._region.blocks.fill_(sentinel) - src = torch.arange(n_elements, dtype=_DTYPE, device="cuda") + src = torch.arange(n_elements, dtype=_DTYPE, device=DEVICE_TYPE) expected_bytes = src.cpu().reshape(-1).view(torch.uint8) total_bytes = n_elements * _DTYPE.itemsize @@ -228,15 +231,15 @@ def test_save_caches_raises_when_allocated_blocks_too_small(make_worker): worker.save_caches({"h": src}, "h", _meta(saves={"h": [0, 1]})) -@_requires_cuda +@_requires_accelerator def test_save_caches_batches_multiple_hashes(make_worker): """Multiple save_caches calls are batched into a single flush.""" worker = make_worker() sentinel = 0x5A worker._region.blocks.fill_(sentinel) - src_a = torch.arange(_HIDDEN_DIM, dtype=_DTYPE, device="cuda") - src_b = torch.arange(_HIDDEN_DIM, 2 * _HIDDEN_DIM, dtype=_DTYPE, device="cuda") + src_a = torch.arange(_HIDDEN_DIM, dtype=_DTYPE, device=DEVICE_TYPE) + src_b = torch.arange(_HIDDEN_DIM, 2 * _HIDDEN_DIM, dtype=_DTYPE, device=DEVICE_TYPE) cache = {"a": src_a, "b": src_b} worker.save_caches(cache, "a", _meta(saves={"a": [1], "b": [3]})) @@ -258,7 +261,7 @@ def test_save_caches_batches_multiple_hashes(make_worker): # ── start_load_caches ──────────────────────────────────────────────────────── -@_requires_cuda +@_requires_accelerator def test_start_load_caches_copies_with_correct_shape_dtype_and_bytes(make_worker): """Single batched load across all hashes with correct byte→dtype→shape.""" worker = make_worker() @@ -276,20 +279,22 @@ def test_start_load_caches_copies_with_correct_shape_dtype_and_bytes(make_worker worker.start_load_caches(encoder_cache, _meta(loads={"h": block_ids})) out = encoder_cache["h"] - assert out.is_cuda, "consumer worker must place the tensor on the GPU" + assert out.device.type == DEVICE_TYPE, ( + "consumer worker must place the tensor on the accelerator" + ) assert out.shape == (n_blocks, _HIDDEN_DIM) assert out.dtype == _DTYPE assert torch.equal(out.cpu(), src_orig) -@_requires_cuda +@_requires_accelerator def test_start_load_caches_preserves_existing_encoder_cache_entry(make_worker): """If ``encoder_cache`` already holds the ``mm_hash``, the worker must not overwrite it.""" worker = make_worker() worker._region.blocks[0].fill_(0x42) - sentinel = torch.full((_HIDDEN_DIM,), 7.0, dtype=_DTYPE, device="cuda") + sentinel = torch.full((_HIDDEN_DIM,), 7.0, dtype=_DTYPE, device=DEVICE_TYPE) encoder_cache = {"h": sentinel} worker.start_load_caches(encoder_cache, _meta(loads={"h": [0]})) @@ -298,7 +303,7 @@ def test_start_load_caches_preserves_existing_encoder_cache_entry(make_worker): ) -@_requires_cuda +@_requires_accelerator def test_start_load_caches_noop_when_loads_is_empty(make_worker): """When ``meta.loads`` is empty the early-return must fire.""" worker = make_worker() @@ -308,7 +313,7 @@ def test_start_load_caches_noop_when_loads_is_empty(make_worker): assert encoder_cache == {} -@_requires_cuda +@_requires_accelerator def test_start_load_caches_skips_cached_and_loads_new_in_same_step(make_worker): """Cached entries are preserved while new ones are loaded.""" worker = make_worker() @@ -321,7 +326,7 @@ def test_start_load_caches_skips_cached_and_loads_new_in_same_step(make_worker): for i, idx in enumerate(new_block_ids): worker._region.blocks[idx].copy_(src_int8[i]) - cached_tensor = torch.full((1, _HIDDEN_DIM), 99.0, dtype=_DTYPE, device="cuda") + cached_tensor = torch.full((1, _HIDDEN_DIM), 99.0, dtype=_DTYPE, device=DEVICE_TYPE) encoder_cache: dict[str, torch.Tensor] = {"cached_h": cached_tensor} worker.start_load_caches( encoder_cache, @@ -336,7 +341,7 @@ def test_start_load_caches_skips_cached_and_loads_new_in_same_step(make_worker): assert torch.equal(out.cpu(), src_orig) -@_requires_cuda +@_requires_accelerator @pytest.mark.parametrize( "tp_rank,pcp_rank", [(0, 0), (1, 0), (0, 1), (1, 1)], @@ -358,23 +363,23 @@ def test_start_load_caches_works_on_all_ranks(make_worker, tp_rank, pcp_rank): worker.start_load_caches(encoder_cache, _meta(loads={"h": block_ids})) out = encoder_cache["h"] - assert out.is_cuda + assert out.device.type == DEVICE_TYPE assert torch.equal(out.cpu(), src_orig) # ── round-trip ─────────────────────────────────────────────────────────────── -@_requires_cuda +@_requires_accelerator def test_save_then_load_round_trips_bytes(make_worker): """Full producer→mmap→consumer byte path in one shot.""" worker = make_worker() n_blocks = 3 block_ids = [5, 1, 6] - src = torch.arange(n_blocks * _HIDDEN_DIM, dtype=_DTYPE, device="cuda").reshape( - n_blocks, _HIDDEN_DIM - ) + src = torch.arange( + n_blocks * _HIDDEN_DIM, dtype=_DTYPE, device=DEVICE_TYPE + ).reshape(n_blocks, _HIDDEN_DIM) worker.save_caches({"h": src}, "h", _meta(saves={"h": block_ids})) worker.flush_saves() @@ -390,12 +395,12 @@ def test_save_then_load_round_trips_bytes(make_worker): # ── buffer recycling ──────────────────────────────────────────────────────── -@_requires_cuda +@_requires_accelerator def test_buffer_pool_is_reused_across_save_steps(make_worker): """After flush_saves, descriptor buffers are returned to the pool and reused on the next flush — no reallocation.""" worker = make_worker() - src = torch.arange(_HIDDEN_DIM, dtype=_DTYPE, device="cuda") + src = torch.arange(_HIDDEN_DIM, dtype=_DTYPE, device=DEVICE_TYPE) worker.save_caches({"h": src}, "h", _meta(saves={"h": [0]})) worker.flush_saves() @@ -411,7 +416,7 @@ def test_buffer_pool_is_reused_across_save_steps(make_worker): assert id(worker._buf_pool._pool[0].src_ptrs) == buf_id -@_requires_cuda +@_requires_accelerator def test_buffer_pool_is_reused_across_load_steps(make_worker): """After start_load_caches, descriptor buffers are returned to the pool and reused on the next call.""" @@ -435,11 +440,11 @@ def test_buffer_pool_is_reused_across_load_steps(make_worker): # ── stream management ──────────────────────────────────────────────────────── -@_requires_cuda +@_requires_accelerator def test_stream_initialized_at_construction(make_worker): - """``_load_stream`` must be a fully initialized CUDA stream.""" + """``_load_stream`` must be a fully initialized accelerator stream.""" worker = make_worker() - assert isinstance(worker._load_stream, torch.cuda.Stream) + assert isinstance(worker._load_stream, current_platform.Stream) # ── lifecycle ──────────────────────────────────────────────────────────────── @@ -518,14 +523,14 @@ def test_shutdown_calls_region_cleanup_and_swallows_errors(caplog_vllm): # ── e2e: scheduler + worker pipeline ──────────────────────────────────────── -@_requires_cuda +@_requires_accelerator def test_e2e_scheduler_worker_save_then_load(make_worker, monkeypatch): """Full pipeline: scheduler allocates blocks, worker saves GPU tensor to mmap via flush_saves, scheduler marks ready after step delay, worker loads from mmap back to GPU, and the result matches the original. Exercises the real scheduler + worker cooperation through a shared - ECSharedRegion, with real CUDA transfers and stream coordination. + ECSharedRegion, with real accelerator transfers and stream coordination. """ import vllm.distributed.ec_transfer.ec_connector.cpu.scheduler as sched_mod from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler import ( @@ -551,9 +556,9 @@ class _Cfg: # -- Step 1: scheduler allocates, worker saves -- n_blocks = 3 - src = torch.arange(n_blocks * _HIDDEN_DIM, dtype=_DTYPE, device="cuda").reshape( - n_blocks, _HIDDEN_DIM - ) + src = torch.arange( + n_blocks * _HIDDEN_DIM, dtype=_DTYPE, device=DEVICE_TYPE + ).reshape(n_blocks, _HIDDEN_DIM) class _Pos: offset = 0 @@ -590,7 +595,7 @@ class _Request: worker.start_load_caches(load_cache, meta_load) out = load_cache["img_001"] - assert out.is_cuda + assert out.device.type == DEVICE_TYPE assert out.shape == src.shape assert out.dtype == src.dtype assert torch.equal(out.cpu(), src.cpu()) diff --git a/tests/v1/ec_connector/unit/test_ec_cpu_connector.py b/tests/v1/ec_connector/unit/test_ec_cpu_connector.py index 3a2600860aa7..6a1ccf857154 100644 --- a/tests/v1/ec_connector/unit/test_ec_cpu_connector.py +++ b/tests/v1/ec_connector/unit/test_ec_cpu_connector.py @@ -7,7 +7,7 @@ - Accuracy: outputs from EC CPU cache match fresh encoder computation. - Latency: loading from EC CPU cache is faster than a cold encoder run. -Requires a CUDA GPU and the Qwen2-VL-2B-Instruct model. +Requires a CUDA or XPU GPU and the Qwen2-VL-2B-Instruct model. """ import time @@ -57,8 +57,8 @@ def _wait_for_ec_ready(llm: LLM) -> None: def _latency_test(llm: LLM) -> None: """Verify EC CPU cache hit is faster than cold encoder computation.""" - if not current_platform.is_cuda(): - pytest.skip("Latency test requires CUDA") + if not (current_platform.is_cuda() or current_platform.is_xpu()): + pytest.skip("Latency test requires an accelerator (CUDA or XPU)") sampling_params = SamplingParams(max_tokens=1, temperature=0) base_image = ImageAsset("stop_sign").pil_image @@ -134,7 +134,10 @@ def _accuracy_test(llm: LLM) -> None: ) -@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") +@pytest.mark.skipif( + not (current_platform.is_cuda() or current_platform.is_xpu()), + reason="Requires an accelerator (CUDA or XPU)", +) def test_ec_cpu_offloading() -> None: """Tests ECCPUConnector accuracy and latency with a VLM model.""" ec_transfer_config = ECTransferConfig( diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/worker/__init__.py b/vllm/distributed/ec_transfer/ec_connector/cpu/worker/__init__.py index b923a2218df1..5a957c4966bc 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/worker/__init__.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/worker/__init__.py @@ -102,16 +102,15 @@ def save_caches( assert self._save_count + len(block_ids) <= self._save_bufs.src_ptrs.numel() - src_ptrs, dst_ptrs, sizes = self._save_bufs + bufs = self._save_bufs src_base = src.view(-1).view(torch.uint8).data_ptr() dst_base = self._region.blocks.data_ptr() idx = self._save_count for i, block_idx in enumerate(block_ids): start = i * block_size - src_ptrs[idx] = src_base + start - dst_ptrs[idx] = dst_base + block_idx * block_size - sizes[idx] = min(block_size, total_bytes - start) + bufs.set_ptrs(idx, src_base + start, dst_base + block_idx * block_size) + bufs.sizes[idx] = min(block_size, total_bytes - start) idx += 1 self._save_count = idx @@ -123,7 +122,7 @@ def flush_saves(self) -> None: bufs = self._save_bufs assert bufs is not None - src_ptrs, dst_ptrs, sizes = bufs + src_ptrs, dst_ptrs, sizes = bufs.src_ptrs, bufs.dst_ptrs, bufs.sizes n = self._save_count swap_blocks_batch(src_ptrs[:n], dst_ptrs[:n], sizes[:n]) @@ -173,8 +172,11 @@ def start_load_caches( op_idx = 0 for block_ids in load_items.values(): for block_idx in block_ids: - src_ptrs[op_idx] = src_base + block_idx * block_size - dst_ptrs[op_idx] = dst_buf_base + op_idx * block_size + bufs.set_ptrs( + op_idx, + src_base + block_idx * block_size, + dst_buf_base + op_idx * block_size, + ) op_idx += 1 swap_blocks_batch(src_ptrs, dst_ptrs, sizes, is_src_access_order_any=True) diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/worker/descriptor_buffers.py b/vllm/distributed/ec_transfer/ec_connector/cpu/worker/descriptor_buffers.py index a75d99ee8181..a42d6ebda4de 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/worker/descriptor_buffers.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/worker/descriptor_buffers.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Reusable pool of (src_ptrs, dst_ptrs, sizes) int64 tensor triples. +"""Reusable pool of (src_ptrs, dst_ptrs, sizes) tensor triples. Used by ECCPUWorker to batch swap_blocks_batch descriptors without per-step allocation overhead. @@ -8,20 +8,46 @@ from typing import NamedTuple +import numpy as np import torch +from vllm.platforms import current_platform + +# CUDA/ROCm cache_kernels.cu requires int64 pointers; the XPU DMA engine +# requires uint64 (see vllm._custom_ops.swap_blocks_batch). +_PTR_DTYPE = torch.uint64 if current_platform.is_xpu() else torch.int64 + class DescriptorBuffers(NamedTuple): src_ptrs: torch.Tensor dst_ptrs: torch.Tensor sizes: torch.Tensor + # Numpy aliases of src_ptrs/dst_ptrs, written via set_ptrs(). + src_np: np.ndarray + dst_np: np.ndarray + + def set_ptrs(self, idx: int, src: int, dst: int) -> None: + """Record the source and destination address of descriptor *idx*. + + TODO(torch>=2.14): drop this indirection and assign the tensors + directly once the minimum supported torch is 2.14. The numpy detour + exists only because torch's setitem unpacks the value as a signed + long long before 2.14 (pytorch#191458), rejecting XPU USM pointers + >= 2**63, while the two's-complement rewrite those versions accept is + in turn rejected by uint64 from 2.14 on. Numpy casts against the + array dtype and so works on either. `sizes` holds byte counts and + needs no such care. + """ + self.src_np[idx] = src + self.dst_np[idx] = dst class DescriptorBufferPool: """Pool of descriptor buffer triples for swap_blocks_batch. - Each buffer is a `DescriptorBuffers` namedtuple of three 1-D int64 - tensors of equal length. Buffers are recycled across steps; if a + Each buffer is a `DescriptorBuffers` namedtuple of three 1-D tensors + (dtype `_PTR_DTYPE`, platform-dependent) of equal length, paired with + numpy aliases used to fill them. Buffers are recycled across steps; if a returned buffer is too small it is discarded and a fresh one allocated. """ @@ -35,11 +61,8 @@ def acquire(self, n: int) -> DescriptorBuffers: bufs = self._pool.pop() if bufs.src_ptrs.numel() >= n: return bufs - return DescriptorBuffers( - torch.empty(n, dtype=torch.int64), - torch.empty(n, dtype=torch.int64), - torch.empty(n, dtype=torch.int64), - ) + src, dst, sizes = (torch.empty(n, dtype=_PTR_DTYPE) for _ in range(3)) + return DescriptorBuffers(src, dst, sizes, src.numpy(), dst.numpy()) def release(self, bufs: DescriptorBuffers) -> None: """Return a buffer triple to the pool for reuse.""" From e4d61d0d222ad21437a1507bb34ca7e916aac880 Mon Sep 17 00:00:00 2001 From: "Li, Jiang" Date: Wed, 19 Aug 2026 10:28:52 +0800 Subject: [PATCH 122/839] [CPU] Add AMX-only high-performance MLA backend for DeepSeek V2/V3/R1 (#52616) Signed-off-by: jiang1.li --- .buildkite/hardware_tests/cpu.yaml | 4 + cmake/cpu_extension.cmake | 6 +- csrc/cpu/sgl-kernels/bmm.cpp | 208 ++ csrc/cpu/sgl-kernels/decode.cpp | 1821 +++++++++++++++++ csrc/cpu/sgl-kernels/extend.cpp | 568 +++++ csrc/cpu/sgl-kernels/flash_attn.h | 250 +++ csrc/cpu/sgl-kernels/mla_cache.cpp | 111 + csrc/cpu/torch_bindings.cpp | 62 + tests/kernels/attention/test_amx_mla.py | 471 +++++ tests/v1/attention/test_mla_backends.py | 2 + .../v1/attention/test_mla_prefill_selector.py | 1 + vllm/_custom_ops.py | 103 + .../layers/attention/mla_attention.py | 41 + vllm/model_executor/layers/linear.py | 6 + vllm/platforms/cpu.py | 36 +- vllm/v1/attention/backends/mla/amx_mla.py | 405 ++++ .../backends/mla/prefill/cpu_native.py | 61 + .../backends/mla/prefill/registry.py | 3 + .../backends/mla/prefill/selector.py | 2 + vllm/v1/attention/backends/registry.py | 1 + 20 files changed, 4156 insertions(+), 6 deletions(-) create mode 100644 csrc/cpu/sgl-kernels/bmm.cpp create mode 100644 csrc/cpu/sgl-kernels/decode.cpp create mode 100644 csrc/cpu/sgl-kernels/extend.cpp create mode 100644 csrc/cpu/sgl-kernels/flash_attn.h create mode 100644 csrc/cpu/sgl-kernels/mla_cache.cpp create mode 100644 tests/kernels/attention/test_amx_mla.py create mode 100644 vllm/v1/attention/backends/mla/amx_mla.py create mode 100644 vllm/v1/attention/backends/mla/prefill/cpu_native.py diff --git a/.buildkite/hardware_tests/cpu.yaml b/.buildkite/hardware_tests/cpu.yaml index efd3e30f6f89..77ebfb99bedc 100644 --- a/.buildkite/hardware_tests/cpu.yaml +++ b/.buildkite/hardware_tests/cpu.yaml @@ -12,6 +12,9 @@ steps: - vllm/_custom_ops.py - tests/kernels/attention/test_cpu_attn.py - tests/v1/attention/test_group_head_counts.py + - tests/kernels/attention/test_amx_mla.py + - vllm/v1/attention/backends/mla/amx_mla.py + - vllm/model_executor/layers/attention/mla_attention.py - tests/kernels/moe/test_cpu_fused_moe.py - tests/kernels/moe/test_cpu_quant_fused_moe.py - tests/kernels/test_onednn.py @@ -29,6 +32,7 @@ steps: bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m " pytest -x -v -s tests/kernels/attention/test_cpu_attn.py pytest -x -v -s tests/v1/attention/test_group_head_counts.py + pytest -x -v -s tests/kernels/attention/test_amx_mla.py pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py pytest -x -v -s tests/kernels/moe/test_cpu_quant_fused_moe.py pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index 46ed96c600c3..92c3daae64bc 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -519,7 +519,11 @@ if (ENABLE_X86_ISA) "csrc/cpu/sgl-kernels/moe.cpp" "csrc/cpu/sgl-kernels/moe_int8.cpp" "csrc/cpu/sgl-kernels/moe_int4.cpp" - "csrc/cpu/sgl-kernels/moe_fp8.cpp") + "csrc/cpu/sgl-kernels/moe_fp8.cpp" + "csrc/cpu/sgl-kernels/bmm.cpp" + "csrc/cpu/sgl-kernels/decode.cpp" + "csrc/cpu/sgl-kernels/extend.cpp" + "csrc/cpu/sgl-kernels/mla_cache.cpp") set(VLLM_EXT_SRC_AVX512 "csrc/cpu/sgl-kernels/fla.cpp" diff --git a/csrc/cpu/sgl-kernels/bmm.cpp b/csrc/cpu/sgl-kernels/bmm.cpp new file mode 100644 index 000000000000..7f1ef9fb8bc6 --- /dev/null +++ b/csrc/cpu/sgl-kernels/bmm.cpp @@ -0,0 +1,208 @@ +// Adapted from +// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu + +// clang-format off + +#include "common.h" +#include "gemm.h" +#include "vec.h" + +namespace { + +template +void bmm_kernel_impl( + scalar_t* __restrict__ out, + const scalar_t* __restrict__ mat1, + const packed_t* __restrict__ mat2, + int64_t B, + int64_t M, + int64_t N, + int64_t K, + int64_t mat1_strideB, + int64_t mat1_strideM, + int64_t out_strideB, + int64_t out_strideM, + float scale = 0.f) { + constexpr int64_t BLOCK_M = block_size_m(); + constexpr int64_t BLOCK_N = block_size_n(); + const int64_t MB = div_up(M, BLOCK_M); + const int64_t NB = div_up(N, BLOCK_N); + + // mat2 contiguous in [B, N, K] + int64_t mat2_strideB = N * K; + int64_t mat2_strideN = K; + + const bool use_brgemm = can_use_brgemm(M); + + // parallel on [B, MB, NB] + at::parallel_for(0, B * MB * NB, 0, [&](int64_t begin, int64_t end) { + int64_t bs{0}, mb{0}, nb{0}; + data_index_init(begin, bs, B, mb, MB, nb, NB); + + // for brgemm, use float32 for accumulate + alignas(64) float Ctmp[BLOCK_M * BLOCK_N]; + + for (int i = begin; i < end; ++i) { + UNUSED(i); + int mb_start = mb * BLOCK_M; + int mb_size = std::min(M - mb_start, BLOCK_M); + int nb_start = nb * BLOCK_N; + int nb_size = std::min(N - nb_start, BLOCK_N); + + tinygemm_kernel( + /* A */ mat1 + bs * mat1_strideB + mb_start * mat1_strideM, + /* B */ mat2 + bs * mat2_strideB + nb_start * mat2_strideN /* nb * BLOCK_N * K */, + /* C */ out + bs * out_strideB + mb_start * out_strideM + nb_start, + /* Ctmp*/ Ctmp, + /* M */ mb_size, + /* N */ nb_size, + /* K */ K, + /* lda */ mat1_strideM, + /* ldb */ nb_size, + /* ldc */ out_strideM, + /* brg */ use_brgemm); + + // move to the next index + data_index_step(bs, B, mb, MB, nb, NB); + } + + if (use_brgemm) { + at::native::cpublas::brgemm_release(); + } + }); +} + +template <> +void bmm_kernel_impl( + at::BFloat16* __restrict__ out, + const at::BFloat16* __restrict__ mat1, + const at::Float8_e4m3fn* __restrict__ mat2, + int64_t B, + int64_t M, + int64_t N, + int64_t K, + int64_t mat1_strideB, + int64_t mat1_strideM, + int64_t out_strideB, + int64_t out_strideM, + float scale) { + constexpr int64_t BLOCK_M = block_size_m(); + constexpr int64_t BLOCK_N = block_size_n(); + const int64_t MB = div_up(M, BLOCK_M); + const int64_t NB = div_up(N, BLOCK_N); + + // mat2 contiguous in [B, N, K] + int64_t mat2_strideB = N * K; + int64_t mat2_strideN = K; + + const bool use_brgemm = can_use_brgemm(M); + + // parallel on [B, MB, NB] + parallel_2d(B * MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) { + // for brgemm, use float32 for accumulate + alignas(64) float Ctmp[BLOCK_M * BLOCK_N]; + // for brgemm when mat2 is float8_e4m3 + alignas(64) at::BFloat16 Btmp[BLOCK_N * BLOCK_K]; + + loop_2d(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) { + int64_t bs = mb / MB; + int64_t mb_start = (mb % MB) * BLOCK_M; + int64_t mb_size = std::min(M - mb_start, BLOCK_M); + int64_t nb_start = nb * BLOCK_N; + int64_t nb_size = std::min(N - nb_start, BLOCK_N); + + tinygemm_kernel( + /* A */ mat1 + bs * mat1_strideB + mb_start * mat1_strideM, + /* B */ mat2 + bs * mat2_strideB + nb_start * mat2_strideN /* nb * BLOCK_N * K */, + /* C */ out + bs * out_strideB + mb_start * out_strideM + nb_start, + /* Btmp*/ Btmp, + /* Ctmp*/ Ctmp, + /*scale*/ scale, + /* M */ mb_size, + /* N */ nb_size, + /* K */ K, + /* lda */ mat1_strideM, + /* ldb */ nb_size, + /* ldc */ out_strideM, + /* brg */ use_brgemm); + }); + + if (use_brgemm) { + at::native::cpublas::brgemm_release(); + } + }); +} + +} // anonymous namespace + +// mat1 : [B, M, K] +// mat2 : [B, N, K] or [B, OC, IC] +// out : [B, M, N] +// scale: [] 0-dim tensor for per tensor quant +// +void bmm_cpu( + at::Tensor& out, at::Tensor& mat1, at::Tensor& mat2, bool is_vnni, const std::optional& scale) { + auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2); + + // input and out could be non-contiguous + // weight needs to be contiguous in [OC, IC] order + CHECK_LAST_DIM_CONTIGUOUS_INPUT(mat1); + CHECK_LAST_DIM_CONTIGUOUS_INPUT(out); + CHECK_INPUT(mat2); + CHECK_DIM(3, out); + CHECK_DIM(3, mat1); + CHECK_DIM(3, mat2); + + int64_t B = mat1.size(0); + int64_t M = mat1.size(1); + int64_t N = mat2.size(1); + int64_t K = mat1.size(2); + + const bool use_fp8_w8a16 = scale.has_value(); + TORCH_CHECK(N % 32 == 0, "tinygemm requires N to be 32x."); + + int64_t mat1_strideB = mat1.stride(0); + int64_t mat1_strideM = mat1.stride(1); + int64_t out_strideB = out.stride(0); + int64_t out_strideM = out.stride(1); + + // check shapes + TORCH_CHECK(mat2.size(0) == B && mat2.size(2) == K, "bmm: mat2 shape mismatch!"); + TORCH_CHECK(out.size(0) == B && out.size(1) == M, "bmm: out shape mismatch!"); + if (!use_fp8_w8a16) { + AT_DISPATCH_REDUCED_FLOATING_TYPES(mat1.scalar_type(), "bmm_kernel_impl", [&] { + bmm_kernel_impl( + out.data_ptr(), + mat1.data_ptr(), + packed_w.data_ptr(), + B, + M, + N, + K, + mat1_strideB, + mat1_strideM, + out_strideB, + out_strideM); + }); + } else { // fp8 bmm + float scale_val = 0.f; + + auto scale_tensor = scale.value(); + TORCH_CHECK(scale_tensor.ndimension() == 0, "bmm: expect scale to be 0-dim tensor."); + scale_val = scale_tensor.item(); + + bmm_kernel_impl( + out.data_ptr(), + mat1.data_ptr(), + packed_w.data_ptr(), + B, + M, + N, + K, + mat1_strideB, + mat1_strideM, + out_strideB, + out_strideM, + scale_val); + } +} diff --git a/csrc/cpu/sgl-kernels/decode.cpp b/csrc/cpu/sgl-kernels/decode.cpp new file mode 100644 index 000000000000..edc6e3818cb4 --- /dev/null +++ b/csrc/cpu/sgl-kernels/decode.cpp @@ -0,0 +1,1821 @@ +// Adapted from +// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu + +// clang-format off + +#include "common.h" +#include "gemm.h" +#include "vec.h" + +namespace { + +// [NOTE] TODO list for this kernel: +// 1. tune the value for BLOCK_N +// 2. planning for {batches, num_heads, num_kv_splits} +// and use actual num_kv_splits for small seq length +// 3. try fast impl of `.tanh()` +// 4. provide amx kernel for index_gemm_kernel_nn when M = 16 +// + +#if defined(CPU_CAPABILITY_AVX512) +// key: from [N, 32] to [32/2, N, 2] +// val: from [N, 32] to [N/2, 32, 2] +template +inline void pack_vnni_Nx32( + scalar_t* __restrict__ dst0, + scalar_t* __restrict__ dst1, + const scalar_t* __restrict__ src, + const index_t* __restrict__ ind, + int N, + int ld_src, + int ld_dst0, + int ld_dst1, + bool convert_v) { + __m512i vinputs[16]; + int n = 0; + for (; n < N; ++n) { + vinputs[n] = _mm512_loadu_si512(src + ind[n] * ld_src); + } + // padding with zero to avoid uninitialized vectors + for (; n < 16; ++n) { + vinputs[n] = _mm512_set1_epi32(0); + } + + // pack value, skip 64 elems for deepseek + // handle 2 vectors at a time from [2, 32] to [32, 2] + if (convert_v) { + for (int n = 0; n < 16; n += 2) { + __m512i d0, d1; + std::tie(d0, d1) = transpose_2x32_16bit(vinputs[n], vinputs[n + 1]); + _mm512_storeu_si512(dst1 + (n >> 1) * ld_dst1 * 2, d0); + _mm512_storeu_si512(dst1 + (n >> 1) * ld_dst1 * 2 + 32, d1); + } + } + + // pack key + transpose_16x16_32bit(vinputs); + + const __mmask16 vmask = (1 << N) - 1; + for (int k = 0; k < 16; ++k) { + _mm512_mask_storeu_epi32(dst0 + k * ld_dst0 * 2, vmask, vinputs[k]); + } +} +#endif + +// [NOTE]: MLA vnni format conversion +// +// here we apply same strategy as `FlashMLA`: +// each kv_cache is loaded once and packed twice (L2 cache hit) +// +// * for key: from [N, K/2, 2] to [K/2, N, 2] +// * for value: from [N/2, 2, Kv] to [N/2, Kv, 2] +// +template +void pack_vnni( + scalar_t* __restrict__ dst0, + scalar_t* __restrict__ dst1, + const scalar_t* __restrict__ src, + const index_t* __restrict__ ind, + int N, + int K, + int Kv, + int ld_src, + int ld_dst0, + int ld_dst1) { +#if defined(CPU_CAPABILITY_AVX512) + const int NB = div_up(N, 16); + const int KB = K / 32; // no remainder + const int KBv = Kv / 32; // no remainder + + for (int nb = 0; nb < NB; ++nb) { + for (int kb = 0; kb < KB; ++kb) { + // handle 16x512bits each block + int nb_size = std::min(N - nb * 16, 16); + pack_vnni_Nx32( + /* dst0 */ dst0 + ((kb * 32) >> 1) * ld_dst0 * 2 + nb * 16 * 2, + /* dst1 */ dst1 + ((nb * 16) >> 1) * ld_dst1 * 2 + kb * 32 * 2, + /* src */ src + kb * 32, + /* ind */ ind + nb * 16, + /* N */ nb_size, + /* ld_src */ ld_src, + /* ld_dst0 */ ld_dst0, + /* ld_dst1 */ ld_dst1, + /* cvt_v */ kb < KBv); + } + } +#else + for (int n = 0; n < N; ++n) { + index_t index = ind[n]; + for (int k = 0; k < K / 2; ++k) { + for (int d = 0; d < 2; ++d) { + dst0[k * ld_dst0 * 2 + n * 2 + d] = src[index * ld_src + k * 2 + d]; + } + } + } + // from [N/2, 2, K] to [N/2, K, 2] + for (int n = 0; n < (N >> 1) * 2; n += 2) { + index_t index0 = ind[n + 0]; + index_t index1 = ind[n + 1]; + for (int k = 0; k < Kv; ++k) { + dst1[(n >> 1) * ld_dst1 * 2 + k * 2 + 0] = src[index0 * ld_src + k]; + dst1[(n >> 1) * ld_dst1 * 2 + k * 2 + 1] = src[index1 * ld_src + k]; + } + } + if (N % 2 != 0) { + index_t index = ind[N - 1]; + for (int k = 0; k < Kv; ++k) { + dst1[(N >> 1) * ld_dst1 * 2 + k * 2 + 0] = src[index * ld_src + k]; + dst1[(N >> 1) * ld_dst1 * 2 + k * 2 + 1] = 0; + } + } +#endif +} + +template +inline void fill_stub(scalar_t* __restrict__ out, float val, int64_t size) { + using Vec = at::vec::Vectorized; + constexpr int kVecSize = Vec::size(); + const Vec data_vec = Vec(static_cast(val)); + int64_t d = 0; +#pragma GCC unroll 4 + for (; d <= size - kVecSize; d += kVecSize) { + data_vec.store(out + d); + } + if (size - d > 0) { + data_vec.store(out + d, size - d); + } +} + +template +inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ acc, float s, int64_t size) { + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + constexpr int kVecSize = bVec::size(); + const fVec s_fvec = fVec(s); + int64_t d = 0; +#pragma GCC unroll 4 + for (; d <= size - kVecSize; d += kVecSize) { + auto [a_fvec0, a_fvec1] = load_float_vec2(acc + d); + a_fvec0 = a_fvec0 * s_fvec; + a_fvec1 = a_fvec1 * s_fvec; + bVec out_bvec = convert_from_float_ext(a_fvec0, a_fvec1); + out_bvec.store(out + d); + } + for (; d < size; ++d) { + out[d] = static_cast(acc[d] * s); + } +} + +template +inline void copy_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ src, int64_t size) { + using bVec = at::vec::Vectorized; + constexpr int kVecSize = bVec::size(); + int64_t d = 0; +#pragma GCC unroll 4 + for (; d <= size - kVecSize; d += kVecSize) { + bVec out_bvec = bVec::loadu(src + d); + out_bvec.store(out + d); + } + for (; d < size; ++d) { + out[d] = src[d]; + } +} + +template +inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ input) { + static_assert(BLOCK_N % 32 == 0); + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + + constexpr int COLS = BLOCK_N / 16; + auto store = [&](auto i) { + constexpr int col = i % COLS; + // for COLS = 2, 4 use 512bit store + if constexpr (col % 2 == 0) { + auto [a_fvec0, a_fvec1] = load_float_vec2(input + col * 16); + bVec out_bvec = convert_from_float_ext(a_fvec0, a_fvec1); + out_bvec.store(out + col * 16); + } + }; + Unroll{}(store); +} + +// GEMM handles query @ key (indexed) x scale +// A : [M, K] +// B : [N, K] indexed +// C : [M, N] +// +template +struct tinygemm_kernel_nt { + static inline void apply( + const scalar_t* __restrict__ A, + const scalar_t* __restrict__ B, + float* __restrict__ C, + const index_t* __restrict__ indices, + float scale, + int64_t lda, + int64_t ldb, + int64_t ldc, + int64_t K, + int64_t max_tokens) { + for (int64_t m = 0; m < BLOCK_M; ++m) { + for (int64_t n = 0; n < BLOCK_N; ++n) { + float sum = 0.f; + int64_t b_idx = indices[n]; + TORCH_CHECK(b_idx < max_tokens, "token index out of scope!"); + for (int64_t k = 0; k < K; ++k) { + sum += scale * static_cast(A[m * lda + k]) * static_cast(B[b_idx * ldb + k]); + } + C[m * ldc + n] = sum; + } + } + } +}; + +#if defined(CPU_CAPABILITY_AVX512) +template +struct tinygemm_kernel_nt { + static inline void apply( + const at::BFloat16* __restrict__ A, + const at::BFloat16* __restrict__ B, + float* __restrict__ C, + const index_t* __restrict__ indices, + float scale, + int64_t lda, + int64_t ldb, + int64_t ldc, + int64_t K, + int64_t max_tokens) { + constexpr int ROWS = BLOCK_M; + constexpr int COLS = BLOCK_N; + + __m512bh va; + __m512bh vb[COLS]; + __m512 vc[ROWS * COLS]; + __m512 vscale = _mm512_set1_ps(scale); + + auto loadc = [&](auto i) { vc[i] = _mm512_setzero_ps(); }; + Unroll{}(loadc); + + // for main loop + auto compute = [&](auto i, int64_t k) { + constexpr int row = i / COLS; + constexpr int col = i % COLS; + + if constexpr (col == 0) { + va = (__m512bh)(_mm512_loadu_si512(A + row * lda + k)); + } + if constexpr (row == 0) { + if constexpr (col + 1 < COLS) { + int64_t b_idx_prefetch = indices[col + 1]; + _mm_prefetch(B + b_idx_prefetch * ldb + k, _MM_HINT_T0); + } + int64_t b_idx = indices[col]; + TORCH_CHECK(b_idx < max_tokens, "token index out of scope!"); + vb[col] = (__m512bh)(_mm512_loadu_si512(B + b_idx * ldb + k)); + } + vc[i] = _mm512_dpbf16_ps(vc[i], va, vb[col]); + }; + + // for remainder + auto compute2 = [&](auto i, int64_t k, __mmask32 mask) { + constexpr int row = i / COLS; + constexpr int col = i % COLS; + + if constexpr (col == 0) { + va = (__m512bh)(_mm512_maskz_loadu_epi16(mask, A + row * lda + k)); + } + if constexpr (row == 0) { + int64_t b_idx = indices[col]; + TORCH_CHECK(b_idx < max_tokens, "token index out of scope!"); + vb[col] = (__m512bh)(_mm512_maskz_loadu_epi16(mask, B + b_idx * ldb + k)); + } + vc[i] = _mm512_dpbf16_ps(vc[i], va, vb[col]); + }; + + int64_t k = 0; + for (; k <= K - 32; k += 32) { + Unroll{}(compute, k); + } + int64_t count = K - k; + if (count > 0) { + __mmask32 mask = (1ULL << count) - 1; + Unroll{}(compute2, k, mask); + } + + auto storec = [&](auto i) { + constexpr int row = i / COLS; + constexpr int col = i % COLS; + C[row * ldc + col] = _mm512_reduce_add_ps(_mm512_mul_ps(vc[i], vscale)); + }; + Unroll{}(storec); + } +}; +#endif + +#if defined(CPU_CAPABILITY_AVX512) +template +struct tinygemm_kernel_nt { + static inline void apply( + const at::Half* __restrict__ A, + const at::Half* __restrict__ B, + float* __restrict__ C, + const index_t* __restrict__ indices, + float scale, + int64_t lda, + int64_t ldb, + int64_t ldc, + int64_t K, + int64_t max_tokens) { + constexpr int ROWS = BLOCK_M; + constexpr int COLS = BLOCK_N; + + __m512 va0, va1; + __m512 vb0[COLS], vb1[COLS]; + __m512 vc[ROWS * COLS]; + __m512 vscale = _mm512_set1_ps(scale); + + auto loadc = [&](auto i) { vc[i] = _mm512_setzero_ps(); }; + Unroll{}(loadc); + + auto compute = [&](auto i, int64_t k) { + constexpr int row = i / COLS; + constexpr int col = i % COLS; + + if constexpr (col == 0) { + __m512i a16 = _mm512_loadu_si512((__m512i const*)(A + row * lda + k)); + va0 = CVT_FP16_TO_FP32(_mm512_extracti32x8_epi32(a16, 0)); + va1 = CVT_FP16_TO_FP32(_mm512_extracti32x8_epi32(a16, 1)); + } + + if constexpr (row == 0) { + int64_t b_idx = indices[col]; + TORCH_CHECK(b_idx < max_tokens, "token index out of scope!"); + __m512i b16 = _mm512_loadu_si512((__m512i const*)(B + b_idx * ldb + k)); + vb0[col] = CVT_FP16_TO_FP32(_mm512_extracti32x8_epi32(b16, 0)); + vb1[col] = CVT_FP16_TO_FP32(_mm512_extracti32x8_epi32(b16, 1)); + } + + vc[i] = _mm512_fmadd_ps(va0, vb0[col], _mm512_fmadd_ps(va1, vb1[col], vc[i])); + }; + + auto compute2 = [&](auto i, int64_t k, __mmask32 mask) { + constexpr int row = i / COLS; + constexpr int col = i % COLS; + + if constexpr (col == 0) { + __m512i a16 = _mm512_maskz_loadu_epi16(mask, (const void*)(A + row * lda + k)); + va0 = CVT_FP16_TO_FP32(_mm512_extracti32x8_epi32(a16, 0)); + va1 = CVT_FP16_TO_FP32(_mm512_extracti32x8_epi32(a16, 1)); + } + + if constexpr (row == 0) { + int64_t b_idx = indices[col]; + TORCH_CHECK(b_idx < max_tokens, "token index out of scope!"); + __m512i b16 = _mm512_maskz_loadu_epi16(mask, (const void*)(B + b_idx * ldb + k)); + vb0[col] = CVT_FP16_TO_FP32(_mm512_extracti32x8_epi32(b16, 0)); + vb1[col] = CVT_FP16_TO_FP32(_mm512_extracti32x8_epi32(b16, 1)); + } + + vc[i] = _mm512_fmadd_ps(va0, vb0[col], _mm512_fmadd_ps(va1, vb1[col], vc[i])); + }; + + int64_t k = 0; + for (; k <= K - 32; k += 32) { + Unroll{}(compute, k); + } + int64_t count = K - k; + if (count > 0) { + __mmask32 mask = (1ULL << count) - 1; + Unroll{}(compute2, k, mask); + } + + auto storec = [&](auto i) { + constexpr int row = i / COLS; + constexpr int col = i % COLS; + C[row * ldc + col] = _mm512_reduce_add_ps(_mm512_mul_ps(vc[i], vscale)); + }; + Unroll{}(storec); + } +}; +#endif + +#define LAUNCH_TINYGEMM_KERNEL_NT(MB_SIZE, NB_SIZE) \ + tinygemm_kernel_nt::apply( \ + A + mb_start * lda, B, C + mb_start * ldc + nb_start, indices + nb_start, scale, lda, ldb, ldc, K, max_tokens); + +// this is used when N isn't multiple of 16, +// N corresponds to `head_size_v` which should be 16x +template +inline void tinygemm_kernel_nn_scalar( + const float* __restrict__ A, + const scalar_t* __restrict__ B, + float* __restrict__ C, + const index_t* __restrict__ indices, + const float* __restrict__ scale, + int64_t M, + int64_t N, + int64_t K, + int64_t lda, + int64_t ldb, + int64_t ldc, + int64_t max_tokens) { + for (int64_t m = 0; m < M; ++m) { + for (int64_t n = 0; n < N; ++n) { + C[m * ldc + n] *= scale[m]; + for (int64_t k = 0; k < K; ++k) { + int64_t b_idx = indices[k]; + TORCH_CHECK(b_idx < max_tokens, "token index out of scope!"); + C[m * ldc + n] += A[m * lda + k] * static_cast(B[b_idx * ldb + n]); + } + } + } +} + +// GEMM handles v' * scale + attn @ value (indexed) +// A : [M, K] +// B : [K, N] indexed +// C ：[M, N] +// +template +struct tinygemm_kernel_nn { + static inline void apply( + const float* __restrict__ A, + const scalar_t* __restrict__ B, + float* __restrict__ C, + const index_t* __restrict__ indices, + const float* __restrict__ scale, + int64_t lda, + int64_t ldb, + int64_t ldc, + int64_t K, + int64_t max_tokens) { + tinygemm_kernel_nn_scalar(A, B, C, indices, scale, BLOCK_M, BLOCK_N, K, lda, ldb, ldc, max_tokens); + } +}; + +#if defined(CPU_CAPABILITY_AVX512) +template +struct tinygemm_kernel_nn { + static inline void apply( + const float* __restrict__ A, + const at::BFloat16* __restrict__ B, + float* __restrict__ C, + const index_t* __restrict__ indices, + const float* __restrict__ scale, + int64_t lda, + int64_t ldb, + int64_t ldc, + int64_t K, + int64_t max_tokens) { + constexpr int ROWS = BLOCK_M; + constexpr int COLS = BLOCK_N / 16; + + __m512 va; + __m512 vb[COLS]; + __m512 vc[ROWS * COLS]; + __m512 vscale; + + auto loadc = [&](auto i) { + constexpr int row = i / COLS; + constexpr int col = i % COLS; +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Warray-bounds" + if constexpr (col == 0) { + vscale = _mm512_set1_ps(scale[row]); + } +#pragma GCC diagnostic pop + vc[i] = _mm512_loadu_ps(C + row * ldc + col * 16); + vc[i] = _mm512_mul_ps(vc[i], vscale); + }; + Unroll{}(loadc); + + auto compute = [&](auto i, int64_t k) { + constexpr int row = i / COLS; + constexpr int col = i % COLS; + + if constexpr (col == 0) { + va = _mm512_set1_ps(A[row * lda + k]); + } + if constexpr (row == 0) { + if (k + 1 < K) { + int64_t b_idx_prefetch = indices[k + 1]; + _mm_prefetch(B + b_idx_prefetch * ldb + col * 16, _MM_HINT_T0); + } + int64_t b_idx = indices[k]; + TORCH_CHECK(b_idx < max_tokens, "token index out of scope!"); + + // for COLS = 2, 4, 6, 8 use 512 bit load + // for COLS = 1, 3, 5, 7 use 256 bit load + if constexpr (COLS % 2 == 0) { + if constexpr (col % 2 == 0) { + __m512i b16 = _mm512_loadu_si512(reinterpret_cast(B + b_idx * ldb + col * 16)); + vb[col + 0] = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(b16, 0)); + vb[col + 1] = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(b16, 1)); + } + } else { + __m256i b16 = _mm256_loadu_si256(reinterpret_cast(B + b_idx * ldb + col * 16)); + vb[col] = CVT_BF16_TO_FP32(b16); + } + } + vc[i] = _mm512_fmadd_ps(va, vb[col], vc[i]); + }; + + for (int64_t k = 0; k < K; ++k) { + Unroll{}(compute, k); + } + + auto storec = [&](auto i) { + constexpr int row = i / COLS; + constexpr int col = i % COLS; + _mm512_storeu_ps(C + row * ldc + col * 16, vc[i]); + }; + Unroll{}(storec); + } +}; +#endif + +#if defined(CPU_CAPABILITY_AVX512) +template +struct tinygemm_kernel_nn { + static inline void apply( + const float* __restrict__ A, + const at::Half* __restrict__ B, + float* __restrict__ C, + const index_t* __restrict__ indices, + const float* __restrict__ scale, + int64_t lda, + int64_t ldb, + int64_t ldc, + int64_t K, + int64_t max_tokens) { + constexpr int ROWS = BLOCK_M; + constexpr int COLS = BLOCK_N / 16; + + __m512 va; + __m512 vb[COLS]; + __m512 vc[ROWS * COLS]; + __m512 vscale; + + auto loadc = [&](auto i) { + constexpr int row = i / COLS; + constexpr int col = i % COLS; +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Warray-bounds" + if constexpr (col == 0) { + vscale = _mm512_set1_ps(scale[row]); + } +#pragma GCC diagnostic pop + vc[i] = _mm512_loadu_ps(C + row * ldc + col * 16); + vc[i] = _mm512_mul_ps(vc[i], vscale); + }; + Unroll{}(loadc); + + auto compute = [&](auto i, int64_t k) { + constexpr int row = i / COLS; + constexpr int col = i % COLS; + + if constexpr (col == 0) { + va = _mm512_set1_ps(A[row * lda + k]); + } + if constexpr (row == 0) { + if (k + 1 < K) { + int64_t b_idx_prefetch = indices[k + 1]; + _mm_prefetch(B + b_idx_prefetch * ldb + col * 16, _MM_HINT_T0); + } + int64_t b_idx = indices[k]; + TORCH_CHECK(b_idx < max_tokens, "token index out of scope!"); + + // for COLS = 2, 4, 6, 8 use 512 bit load + // for COLS = 1, 3, 5, 7 use 256 bit load + if constexpr (COLS % 2 == 0) { + if constexpr (col % 2 == 0) { + __m512i b16 = _mm512_loadu_si512(reinterpret_cast(B + b_idx * ldb + col * 16)); + vb[col + 0] = CVT_FP16_TO_FP32(_mm512_extracti32x8_epi32(b16, 0)); + vb[col + 1] = CVT_FP16_TO_FP32(_mm512_extracti32x8_epi32(b16, 1)); + } + } else { + __m256i b16 = _mm256_loadu_si256(reinterpret_cast(B + b_idx * ldb + col * 16)); + vb[col] = CVT_FP16_TO_FP32(b16); + } + } + vc[i] = _mm512_fmadd_ps(va, vb[col], vc[i]); + }; + + for (int64_t k = 0; k < K; ++k) { + Unroll{}(compute, k); + } + + auto storec = [&](auto i) { + constexpr int row = i / COLS; + constexpr int col = i % COLS; + _mm512_storeu_ps(C + row * ldc + col * 16, vc[i]); + }; + Unroll{}(storec); + } +}; +#endif + +#define LAUNCH_TINYGEMM_KERNEL_NN(MB_SIZE, NB_SIZE) \ + tinygemm_kernel_nn::apply( \ + A + mb_start * lda, \ + B + nb_start, \ + C + mb_start * ldc + nb_start, \ + indices, \ + scale + mb_start, \ + lda, \ + ldb, \ + ldc, \ + K, \ + max_tokens); + +template +void index_gemm_kernel_nt( + const scalar_t* __restrict__ A, + const scalar_t* __restrict__ B, + float* __restrict__ C, + const index_t* __restrict__ indices, + float scale, + int64_t M, + int64_t N, + int64_t K, + int64_t lda, + int64_t ldb, + int64_t ldc, + int64_t max_tokens) { + // pattern: 1-8-8 + if (M == 1) { + constexpr int64_t BLOCK_N = 8; + const int64_t NB = div_up(N, BLOCK_N); + int64_t mb_start = 0, lda = 1, ldc = 1; + + for (int64_t nb = 0; nb < NB; ++nb) { + int64_t nb_start = nb * BLOCK_N; + int64_t nb_size = std::min(BLOCK_N, N - nb_start); + + switch (nb_size) { + case 1: + LAUNCH_TINYGEMM_KERNEL_NT(1, 1); + break; + case 2: + LAUNCH_TINYGEMM_KERNEL_NT(1, 2); + break; + case 3: + LAUNCH_TINYGEMM_KERNEL_NT(1, 3); + break; + case 4: + LAUNCH_TINYGEMM_KERNEL_NT(1, 4); + break; + case 5: + LAUNCH_TINYGEMM_KERNEL_NT(1, 5); + break; + case 6: + LAUNCH_TINYGEMM_KERNEL_NT(1, 6); + break; + case 7: + LAUNCH_TINYGEMM_KERNEL_NT(1, 7); + break; + case 8: + LAUNCH_TINYGEMM_KERNEL_NT(1, 8); + break; + default: + TORCH_CHECK(false, "Unexpected block size, 1x", "nb_size"); + } + } + return; + } + + // default pattern: 1-6-24 + // FP16 pattern: 2-8-16 + constexpr int64_t BLOCK_M = 4; + constexpr int64_t BLOCK_N = std::is_same_v ? 4 : 6; + const int64_t MB = div_up(M, BLOCK_M); + const int64_t NB = div_up(N, BLOCK_N); + + for (int64_t mb = 0; mb < MB; ++mb) { + int64_t mb_start = mb * BLOCK_M; + int64_t mb_size = std::min(BLOCK_M, M - mb_start); + for (int64_t nb = 0; nb < NB; ++nb) { + int64_t nb_start = nb * BLOCK_N; + int64_t nb_size = std::min(BLOCK_N, N - nb_start); + + switch (mb_size << 4 | nb_size) { + // mb_size = 1 + case 0x11: + LAUNCH_TINYGEMM_KERNEL_NT(1, 1); + break; + case 0x12: + LAUNCH_TINYGEMM_KERNEL_NT(1, 2); + break; + case 0x13: + LAUNCH_TINYGEMM_KERNEL_NT(1, 3); + break; + case 0x14: + LAUNCH_TINYGEMM_KERNEL_NT(1, 4); + break; + case 0x15: + LAUNCH_TINYGEMM_KERNEL_NT(1, 5); + break; + case 0x16: + LAUNCH_TINYGEMM_KERNEL_NT(1, 6); + break; + // mb_size = 2 + case 0x21: + LAUNCH_TINYGEMM_KERNEL_NT(2, 1); + break; + case 0x22: + LAUNCH_TINYGEMM_KERNEL_NT(2, 2); + break; + case 0x23: + LAUNCH_TINYGEMM_KERNEL_NT(2, 3); + break; + case 0x24: + LAUNCH_TINYGEMM_KERNEL_NT(2, 4); + break; + case 0x25: + LAUNCH_TINYGEMM_KERNEL_NT(2, 5); + break; + case 0x26: + LAUNCH_TINYGEMM_KERNEL_NT(2, 6); + break; + // mb_size = 3 + case 0x31: + LAUNCH_TINYGEMM_KERNEL_NT(3, 1); + break; + case 0x32: + LAUNCH_TINYGEMM_KERNEL_NT(3, 2); + break; + case 0x33: + LAUNCH_TINYGEMM_KERNEL_NT(3, 3); + break; + case 0x34: + LAUNCH_TINYGEMM_KERNEL_NT(3, 4); + break; + case 0x35: + LAUNCH_TINYGEMM_KERNEL_NT(3, 5); + break; + case 0x36: + LAUNCH_TINYGEMM_KERNEL_NT(3, 6); + break; + // mb_size = 4 + case 0x41: + LAUNCH_TINYGEMM_KERNEL_NT(4, 1); + break; + case 0x42: + LAUNCH_TINYGEMM_KERNEL_NT(4, 2); + break; + case 0x43: + LAUNCH_TINYGEMM_KERNEL_NT(4, 3); + break; + case 0x44: + LAUNCH_TINYGEMM_KERNEL_NT(4, 4); + break; + case 0x45: + LAUNCH_TINYGEMM_KERNEL_NT(4, 5); + break; + case 0x46: + LAUNCH_TINYGEMM_KERNEL_NT(4, 6); + break; + default: + TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", "nb_size"); + } + } + } +} + +template +void index_gemm_kernel_nn( + const float* __restrict__ A, + const scalar_t* __restrict__ B, + float* __restrict__ C, + const index_t* __restrict__ indices, + float* __restrict__ scale, + int64_t M, + int64_t N, + int64_t K, + int64_t lda, + int64_t ldb, + int64_t ldc, + int64_t max_tokens) { + constexpr int kVecSize = 16; + if ((N & (kVecSize - 1)) != 0) { + tinygemm_kernel_nn_scalar(A, B, C, indices, scale, M, N, K, lda, ldb, ldc, max_tokens); + return; + } + + // pattern: 1-8-8 + if (M == 1) { + constexpr int64_t BLOCK_N = 8 * kVecSize; + const int64_t NB = div_up(N, BLOCK_N); + int64_t mb_start = 0, lda = 1, ldc = 1; + + for (int64_t nb = 0; nb < NB; ++nb) { + int64_t nb_start = nb * BLOCK_N; + int64_t nb_size = std::min(BLOCK_N, N - nb_start); + + switch (nb_size >> 4) { + case 1: + LAUNCH_TINYGEMM_KERNEL_NN(1, 16); + break; + case 2: + LAUNCH_TINYGEMM_KERNEL_NN(1, 32); + break; + case 3: + LAUNCH_TINYGEMM_KERNEL_NN(1, 48); + break; + case 4: + LAUNCH_TINYGEMM_KERNEL_NN(1, 64); + break; + case 5: + LAUNCH_TINYGEMM_KERNEL_NN(1, 80); + break; + case 6: + LAUNCH_TINYGEMM_KERNEL_NN(1, 96); + break; + case 7: + LAUNCH_TINYGEMM_KERNEL_NN(1, 112); + break; + case 8: + LAUNCH_TINYGEMM_KERNEL_NN(1, 128); + break; + default: + TORCH_CHECK(false, "Unexpected block size, 1x", "nb_size"); + } + } + return; + } + + constexpr int64_t BLOCK_M = 4; + constexpr int64_t BLOCK_N = 6 * kVecSize; + const int64_t MB = div_up(M, BLOCK_M); + const int64_t NB = div_up(N, BLOCK_N); + + for (int64_t mb = 0; mb < MB; ++mb) { + int64_t mb_start = mb * BLOCK_M; + int64_t mb_size = std::min(BLOCK_M, M - mb_start); + for (int64_t nb = 0; nb < NB; ++nb) { + int64_t nb_start = nb * BLOCK_N; + int64_t nb_size = std::min(BLOCK_N, N - nb_start); + + switch (mb_size << 4 | nb_size >> 4) { + // mb_size = 1 + case 0x11: + LAUNCH_TINYGEMM_KERNEL_NN(1, 16); + break; + case 0x12: + LAUNCH_TINYGEMM_KERNEL_NN(1, 32); + break; + case 0x13: + LAUNCH_TINYGEMM_KERNEL_NN(1, 48); + break; + case 0x14: + LAUNCH_TINYGEMM_KERNEL_NN(1, 64); + break; + case 0x15: + LAUNCH_TINYGEMM_KERNEL_NN(1, 80); + break; + case 0x16: + LAUNCH_TINYGEMM_KERNEL_NN(1, 96); + break; + // mb_size = 2 + case 0x21: + LAUNCH_TINYGEMM_KERNEL_NN(2, 16); + break; + case 0x22: + LAUNCH_TINYGEMM_KERNEL_NN(2, 32); + break; + case 0x23: + LAUNCH_TINYGEMM_KERNEL_NN(2, 48); + break; + case 0x24: + LAUNCH_TINYGEMM_KERNEL_NN(2, 64); + break; + case 0x25: + LAUNCH_TINYGEMM_KERNEL_NN(2, 80); + break; + case 0x26: + LAUNCH_TINYGEMM_KERNEL_NN(2, 96); + break; + // mb_size = 3 + case 0x31: + LAUNCH_TINYGEMM_KERNEL_NN(3, 16); + break; + case 0x32: + LAUNCH_TINYGEMM_KERNEL_NN(3, 32); + break; + case 0x33: + LAUNCH_TINYGEMM_KERNEL_NN(3, 48); + break; + case 0x34: + LAUNCH_TINYGEMM_KERNEL_NN(3, 64); + break; + case 0x35: + LAUNCH_TINYGEMM_KERNEL_NN(3, 80); + break; + case 0x36: + LAUNCH_TINYGEMM_KERNEL_NN(3, 96); + break; + // mb_size = 4 + case 0x41: + LAUNCH_TINYGEMM_KERNEL_NN(4, 16); + break; + case 0x42: + LAUNCH_TINYGEMM_KERNEL_NN(4, 32); + break; + case 0x43: + LAUNCH_TINYGEMM_KERNEL_NN(4, 48); + break; + case 0x44: + LAUNCH_TINYGEMM_KERNEL_NN(4, 64); + break; + case 0x45: + LAUNCH_TINYGEMM_KERNEL_NN(4, 80); + break; + case 0x46: + LAUNCH_TINYGEMM_KERNEL_NN(4, 96); + break; + default: + TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", "nb_size"); + } + } + } +} + +template +void decode_set_kv_buffer( + scalar_t* __restrict__ k_buffer, + scalar_t* __restrict__ v_buffer, + const scalar_t* __restrict__ key, + const scalar_t* __restrict__ value, + const int64_t* __restrict__ loc, + int64_t batches, + int64_t num_heads_kv, + int64_t head_size, + int64_t head_size_v, + int64_t k_strideN, + int64_t k_strideH, + int64_t v_strideN, + int64_t v_strideH, + int64_t nk_strideN, + int64_t nk_strideH, + int64_t nv_strideN, + int64_t nv_strideH, + bool is_mla) { + at::parallel_for(0, batches * num_heads_kv, 0, [&](int64_t begin, int64_t end) { + int64_t bs{0}, head_kv_id{0}; + data_index_init(begin, bs, batches, head_kv_id, num_heads_kv); + + for (int64_t i = begin; i < end; i++) { + int64_t loc_val = loc[bs]; + scalar_t* k_buffer_ptr = k_buffer + loc_val * k_strideN + head_kv_id * k_strideH; + const scalar_t* new_key_ptr = key + bs * nk_strideN + head_kv_id * nk_strideH; + copy_stub(k_buffer_ptr, new_key_ptr, head_size); + if (!is_mla) { + scalar_t* v_buffer_ptr = v_buffer + loc_val * v_strideN + head_kv_id * v_strideH; + const scalar_t* new_value_ptr = value + bs * nv_strideN + head_kv_id * nv_strideH; + copy_stub(v_buffer_ptr, new_value_ptr, head_size_v); + } + + // move to the next index + data_index_step(bs, batches, head_kv_id, num_heads_kv); + } + }); +} + +template +void decode_accumulate_kv_splits( + scalar_t* __restrict__ output, + float* __restrict__ attn_logits, + const scalar_t* __restrict__ sinks_ptr, + int64_t batches, + int64_t num_heads, + int64_t head_size_v, + int64_t num_kv_splits, + int64_t l_stride1, + int64_t l_stride2, + bool has_sink) { + using Vec = at::vec::Vectorized; + + // parallel on [batches, num_heads] + at::parallel_for(0, batches * num_heads, 0, [&](int64_t begin, int64_t end) { + int64_t bi{0}, ni{0}; + data_index_init(begin, bi, batches, ni, num_heads); + // NB: here we use logits[b][h][0] as acc, since + // for the first kv split (kv_id == 0): + // m_delta = std::exp(-inf) = 0 + // e_logic = std::exp(0) = 1 + // acc = acc * m_delta + tv * e_logic = tv + for (int64_t i = begin; i < end; ++i) { + float* __restrict__ acc = attn_logits + i * l_stride1; + + float s_prime = 0.f; + float m_prime = -std::numeric_limits::infinity(); + + // update acc with from each kv_split + for (int64_t kv_id = 0; kv_id < num_kv_splits; ++kv_id) { + float* __restrict__ tv = acc + kv_id * l_stride2; + const float tlogic = (acc + kv_id * l_stride2)[head_size_v]; + + float m_i = std::max(tlogic, m_prime); + float m_delta = std::exp(m_prime - m_i); + float e_logic = std::exp(tlogic - m_i); + if (kv_id != 0) { + at::vec::map2( + [m_delta, e_logic](Vec x, Vec y) { return x * Vec(m_delta) + y * Vec(e_logic); }, + acc, + acc, + tv, + head_size_v); + } + + s_prime = s_prime * m_delta + e_logic; + m_prime = m_i; + } + if (has_sink) { + s_prime += std::exp(sinks_ptr[ni] - m_prime); + } + copy_stub(output + i * head_size_v, acc, 1 / s_prime, head_size_v); + // move to the next index + data_index_step(bi, batches, ni, num_heads); + } + }); +} + +template +void decode_attention_kernel_impl( + scalar_t* __restrict__ output, + float* __restrict__ attn_logits, + const scalar_t* __restrict__ query, + const scalar_t* __restrict__ k_buffer, + const scalar_t* __restrict__ v_buffer, + const index_t* __restrict__ req_to_token, + const int64_t* __restrict__ req_pool_indices, + const int64_t* __restrict__ seq_lens, + const int64_t* __restrict__ encoder_lens, + const scalar_t* __restrict__ sinks, + int64_t batches, + int64_t num_heads, + int64_t head_size, + int64_t head_size_v, + int64_t num_kv_splits, + int64_t q_strideM, + int64_t q_strideH, + int64_t k_strideN, + int64_t k_strideH, + int64_t v_strideN, + int64_t v_strideH, + float sm_scale, + float logit_cap, + int64_t max_num_reqs, + int64_t max_context_len, + int64_t max_total_num_tokens, + int64_t sliding_window_size, + bool is_cross_attn, + bool has_encoder_lens, + bool has_sink) { + using Vec = at::vec::Vectorized; + + // strides + const int64_t l_stride1 = num_kv_splits * (head_size_v + 1); + const int64_t l_stride2 = head_size_v + 1; + + const bool has_logit_cap = logit_cap > 0; + float rlogit_cap = has_logit_cap ? 1 / logit_cap : 0.f; + + // parallel on [batches, num_heads, num_kv_splits] + at::parallel_for(0, batches * num_heads * num_kv_splits, 0, [&](int64_t begin, int64_t end) { + int64_t bs{0}, head_id{0}, kv_id{0}; + data_index_init(begin, bs, batches, head_id, num_heads, kv_id, num_kv_splits); + + // s_prime and s_delta + alignas(64) float s_i[BLOCK_N]; + float* __restrict__ s_delta = s_i; + + for (int64_t i = begin; i < end; ++i) { + // get query + const scalar_t* __restrict__ q_ptr = query + bs * q_strideM + head_id * q_strideH; + + // get key/value + int64_t seq_len_kv = is_cross_attn ? encoder_lens[bs] : seq_lens[bs]; + int64_t req_pool_id = req_pool_indices[bs]; + int64_t kv_offset = (has_encoder_lens && (!is_cross_attn)) ? encoder_lens[bs] : 0; + if (sliding_window_size > 0 && seq_len_kv > sliding_window_size) { + kv_offset = seq_len_kv - sliding_window_size; + seq_len_kv = sliding_window_size; + } + TORCH_CHECK(seq_len_kv <= max_context_len, "seq_len_kv out of scope!"); + TORCH_CHECK(req_pool_id < max_num_reqs, "req_pool_id out of scope!"); + + const int64_t SPLIT_SIZE = div_up(seq_len_kv, num_kv_splits); + const int64_t kv_start = kv_id * SPLIT_SIZE; + const int64_t kv_end = std::min(kv_start + SPLIT_SIZE, seq_len_kv); + + float m_prime = -std::numeric_limits::infinity(); + float s_prime = 0.f; + + // get v_prime, and init to zero + float* __restrict__ v_prime = attn_logits + i * (head_size_v + 1); + fill_stub(v_prime, 0.f, head_size_v); + + // loop over K and V sequence with BLOCK_N + for (int64_t n = kv_start; n < kv_end; n += BLOCK_N) { + int64_t n_size = std::min(BLOCK_N, kv_end - n); + + // calculate s_i <- scale * Q @ K + index_gemm_kernel_nt( + /* A */ q_ptr, + /* B */ k_buffer + head_id * k_strideH, + /* C */ s_i, + /* ind */ req_to_token + req_pool_id * max_context_len + n + kv_offset, + /* scl */ sm_scale, + /* M */ 1, + /* N */ n_size, + /* K */ head_size, + /* lda */ 1, + /* ldb */ k_strideN, + /* ldc */ 1, + /* mtt */ max_total_num_tokens); + + // TODO: `tanh` from torch uses sleef u10, going to be slow + if (has_logit_cap) { + at::vec::map( + [logit_cap, rlogit_cap](Vec x) { return Vec(logit_cap) * (x * Vec(rlogit_cap)).tanh(); }, + s_i, + s_i, + n_size); + } + + // m_i: max value per row + float m_i = at::vec::reduce_all([](Vec& x, Vec& y) { return at::vec::maximum(x, y); }, s_i, n_size); + m_i = std::max(m_i, m_prime); + + // m_delta <- exp(m' - m_i) + float m_delta = std::exp(m_prime - m_i); + + // s_delta <- exp(s_i - m_i) + at::vec::map([m_i](Vec x) { return (x - Vec(m_i)).exp_u20(); }, s_delta, s_i, n_size); + + // s' <- s' * m_delta + sum(s_delta) + s_prime *= m_delta; + s_prime += at::vec::reduce_all([](Vec& x, Vec& y) { return x + y; }, s_delta, n_size); + + m_prime = m_i; + + // calculate V' <- s_delta @ V + V' * m_delta + index_gemm_kernel_nn( + /* A */ s_delta, + /* B */ v_buffer + head_id * v_strideH, + /* C */ v_prime, + /* ind */ req_to_token + req_pool_id * max_context_len + n + kv_offset, + /* scl */ &m_delta, + /* M */ 1, + /* N */ head_size_v, + /* K */ n_size, + /* lda */ 1, + /* ldb */ v_strideN, + /* ldc */ 1, + /* mtt */ max_total_num_tokens); + } // loop with KV blocks + + // only update v' when kv_split_size > 0 + if (kv_end > kv_start) { + float s = 1 / s_prime; + at::vec::map([s](Vec out) { return out * Vec(s); }, v_prime, v_prime, head_size_v); + + v_prime[head_size_v] = m_prime + std::log(s_prime); + } else { + v_prime[head_size_v] = -std::numeric_limits::infinity(); + } + + // move to the next index + data_index_step(bs, batches, head_id, num_heads, kv_id, num_kv_splits); + } + }); + + decode_accumulate_kv_splits( + output, attn_logits, sinks, batches, num_heads, head_size_v, num_kv_splits, l_stride1, l_stride2, has_sink); +} // MHA + +template +void decode_attention_mla_kernel_impl( + scalar_t* __restrict__ output, + float* __restrict__ attn_logits, + const scalar_t* __restrict__ query, + const scalar_t* __restrict__ k_buffer, + const scalar_t* __restrict__ v_buffer, + const index_t* __restrict__ req_to_token, + const int64_t* __restrict__ req_pool_indices, + const int64_t* __restrict__ seq_lens, + scalar_t* __restrict__ buffer, + const scalar_t* __restrict__ sinks, + int64_t batches, + int64_t num_heads, + int64_t head_size, + int64_t head_size_v, + int64_t num_kv_splits, + int64_t q_strideM, + int64_t q_strideH, + int64_t k_strideN, + int64_t k_strideH, + int64_t v_strideN, + int64_t v_strideH, + float sm_scale, + float logit_cap, + int64_t max_num_reqs, + int64_t max_context_len, + int64_t max_total_num_tokens, + int64_t buffer_size_per_thread, + bool has_sink) { + using Vec = at::vec::Vectorized; + + // block length for heads + const int64_t BLOCK_H = batches == 1 ? 6 : (batches > 16 ? 22 : 11); + + // strides + const int64_t l_stride0 = num_heads * num_kv_splits * (head_size_v + 1); + const int64_t l_stride1 = num_kv_splits * (head_size_v + 1); + const int64_t l_stride2 = head_size_v + 1; + + TORCH_CHECK(logit_cap == 0.f, "decode MLA: expect no logit_cap."); + + // partition the heads into blocks for parallel + const int64_t num_blocks = div_up(num_heads, BLOCK_H); + + // parallel on [batches, num_blocks, num_kv_splits] + at::parallel_for(0, batches * num_blocks * num_kv_splits, 0, [&](int64_t begin, int64_t end) { + int64_t bs{0}, block_id{0}, kv_id{0}; + data_index_init(begin, bs, batches, block_id, num_blocks, kv_id, num_kv_splits); + + int tid = at::get_thread_num(); + scalar_t* __restrict__ Btmp0 = buffer + tid * buffer_size_per_thread; + scalar_t* __restrict__ Btmp1 = Btmp0 + BLOCK_N * head_size; + + // init Btmp1 just once for each thread to prevent NaN + // Btmp0 is not needed as it computes full K every single time + fill_stub(Btmp1, 0.f, BLOCK_N * head_size_v); + + alignas(64) float s_i[BLOCK_H * BLOCK_N]; + float* __restrict__ s_delta = s_i; + alignas(64) scalar_t s_delta2[BLOCK_H * BLOCK_N]; + + alignas(64) float s_prime[BLOCK_H]; + alignas(64) float m_prime[BLOCK_H]; + alignas(64) float m_delta[BLOCK_H]; + + for (int64_t i = begin; i < end; ++i) { + const int64_t h_start = block_id * BLOCK_H; + const int64_t h_end = std::min(block_id * BLOCK_H + BLOCK_H, num_heads); + const int64_t h_size = h_end - h_start; + + // get query + const scalar_t* __restrict__ q_ptr = query + bs * q_strideM + h_start * q_strideH; + + int64_t seq_len_kv = seq_lens[bs]; + int64_t req_pool_id = req_pool_indices[bs]; + TORCH_CHECK(seq_len_kv <= max_context_len, "seq_len_kv out of scope!"); + TORCH_CHECK(req_pool_id < max_num_reqs, "req_pool_id out of scope!"); + + const int64_t SPLIT_SIZE = div_up(seq_len_kv, num_kv_splits); + const int64_t kv_start = kv_id * SPLIT_SIZE; + const int64_t kv_end = std::min(kv_start + SPLIT_SIZE, seq_len_kv); + + fill_stub(s_prime, 0.f, BLOCK_H); + fill_stub(m_prime, -std::numeric_limits::infinity(), BLOCK_H); + + // get v_prime, and init to zero + float* __restrict__ v_prime = attn_logits + bs * l_stride0 + h_start * l_stride1 + kv_id * l_stride2; + for (int64_t h = 0; h < h_size; ++h) { + fill_stub(v_prime + h * l_stride1, 0.f, head_size_v); + } + + // loop over K and V sequence with BLOCK_N + for (int64_t n = kv_start; n < kv_end; n += BLOCK_N) { + int64_t n_size = std::min(BLOCK_N, kv_end - n); + const int64_t padded_n_size = div_up(int(n_size), TILE_K) * TILE_K; + + // get key and pack + pack_vnni( + /* dst0 */ Btmp0, + /* dst1 */ Btmp1, + /* src */ k_buffer + /* head_kv_id */ 0 * k_strideH, + /* ind */ req_to_token + req_pool_id * max_context_len + n, + /* N */ n_size, + /* K */ head_size, + /* Kv */ head_size_v, + /* ld_src */ k_strideN, + /* ld_dst0 */ BLOCK_N, + /* ld_dst1 */ head_size_v); + + // calculate s_i <- Q @ K + at::native::cpublas::brgemm( + /* M */ h_size, + /* N */ n_size, + /* K */ head_size, + /* lda */ q_strideH, + /* ldb */ BLOCK_N, + /* ldc */ BLOCK_N, + /* add_C */ false, + /* A */ q_ptr, + /* B */ Btmp0, + /* C */ s_i); + + const Vec scale_vec = Vec(sm_scale); + for (int64_t h = 0; h < h_size; ++h) { + // s_i <- s_i * scale + at::vec::map( + [scale_vec](Vec x) { return x * scale_vec; }, s_i + h * BLOCK_N, s_i + h * BLOCK_N, n_size); + + // m_i: max value per row + float m_i = at::vec::reduce_all( + [](Vec& x, Vec& y) { return at::vec::maximum(x, y); }, s_i + h * BLOCK_N, n_size); + m_i = std::max(m_i, m_prime[h]); + + // m_delta <- exp(m' - m_i) + m_delta[h] = std::exp(m_prime[h] - m_i); + + // s_delta <- exp(s_i - m_i) + at::vec::map( + [m_i](Vec x) { return (x - Vec(m_i)).exp_u20(); }, s_delta + h * BLOCK_N, s_i + h * BLOCK_N, n_size); + + // s' <- s' * m_delta + sum(s_delta) + s_prime[h] *= m_delta[h]; + s_prime[h] += at::vec::reduce_all([](Vec& x, Vec& y) { return x + y; }, s_delta + h * BLOCK_N, n_size); + + m_prime[h] = m_i; + + // v' <- v' * m_delta + float scale_m = m_delta[h]; + at::vec::map( + [scale_m](Vec x) { return x * Vec(scale_m); }, + v_prime + h * l_stride1, + v_prime + h * l_stride1, + head_size_v); + + // pad s_delta with 0 first and then convert to scalar_t + fill_stub(s_delta + h * BLOCK_N + n_size, 0.f, padded_n_size - n_size); + copy_stub(s_delta2 + h * BLOCK_N, s_delta + h * BLOCK_N); + } + + // calculate V' <- s_delta @ V + V' + at::native::cpublas::brgemm( + /* M */ h_size, + /* N */ head_size_v, + /* K */ padded_n_size, // n_size + /* lda */ BLOCK_N, + /* ldb */ head_size_v, + /* ldc */ l_stride1, + /* add_C */ true, + /* A */ s_delta2, + /* B */ Btmp1, + /* C */ v_prime); + } // loop with KV blocks + + // only update v' when kv_split_size > 0 + if (kv_end > kv_start) { + for (int64_t h = 0; h < h_size; ++h) { + float s = 1 / s_prime[h]; + at::vec::map( + [s](Vec out) { return out * Vec(s); }, v_prime + h * l_stride1, v_prime + h * l_stride1, head_size_v); + (v_prime + h * l_stride1)[head_size_v] = m_prime[h] + std::log(s_prime[h]); + } + } else { + for (int64_t h = 0; h < h_size; ++h) { + (v_prime + h * l_stride1)[head_size_v] = -std::numeric_limits::infinity(); + } + } + + // move to the next index + data_index_step(bs, batches, block_id, num_blocks, kv_id, num_kv_splits); + } + at::native::cpublas::brgemm_release(); + }); + + decode_accumulate_kv_splits( + output, attn_logits, sinks, batches, num_heads, head_size_v, num_kv_splits, l_stride1, l_stride2, has_sink); +} // MLA + +template +void decode_attention_grouped_kernel_impl( + scalar_t* __restrict__ output, + float* __restrict__ attn_logits, + const scalar_t* __restrict__ query, + const scalar_t* __restrict__ k_buffer, + const scalar_t* __restrict__ v_buffer, + const index_t* __restrict__ req_to_token, + const int64_t* __restrict__ req_pool_indices, + const int64_t* __restrict__ seq_lens, + const int64_t* __restrict__ encoder_lens, + const scalar_t* __restrict__ sinks, + int64_t batches, + int64_t num_heads, + int64_t num_heads_kv, + int64_t head_size, + int64_t head_size_v, + int64_t num_kv_splits, + int64_t q_strideM, + int64_t q_strideH, + int64_t k_strideN, + int64_t k_strideH, + int64_t v_strideN, + int64_t v_strideH, + float sm_scale, + float logit_cap, + int64_t max_num_reqs, + int64_t max_context_len, + int64_t max_total_num_tokens, + int64_t sliding_window_size, + bool is_cross_attn, + bool has_encoder_lens, + bool has_sink) { + using Vec = at::vec::Vectorized; + + // block length for heads + // we parallel on [batches, divup(num_heads, BLOCK_H), num_kv_splits] + // use smaller BLOCK_H when batches is small to utilize all cores + constexpr int64_t kBLOCK_H = 16; + const int64_t BLOCK_H = std::min(4 * batches, kBLOCK_H); + + // strides + const int64_t l_stride0 = num_heads * num_kv_splits * (head_size_v + 1); + const int64_t l_stride1 = num_kv_splits * (head_size_v + 1); + const int64_t l_stride2 = head_size_v + 1; + + const bool has_logit_cap = logit_cap > 0; + float rlogit_cap = has_logit_cap ? 1 / logit_cap : 0.f; + + // partition the heads into blocks for parallel + const int64_t num_groups = num_heads / num_heads_kv; + const int64_t num_blocks = div_up(num_groups, BLOCK_H); + + // parallel on [batches, num_heads_kv, num_blocks, num_kv_splits] + at::parallel_for(0, batches * num_heads_kv * num_blocks * num_kv_splits, 0, [&](int64_t begin, int64_t end) { + int64_t bs{0}, head_kv_id{0}, block_id{0}, kv_id{0}; + data_index_init(begin, bs, batches, head_kv_id, num_heads_kv, block_id, num_blocks, kv_id, num_kv_splits); + + alignas(64) float s_i[BLOCK_H * BLOCK_N]; + float* __restrict__ s_delta = s_i; + + alignas(64) float s_prime[BLOCK_H]; + alignas(64) float m_prime[BLOCK_H]; + alignas(64) float m_delta[BLOCK_H]; + + for (int64_t i = begin; i < end; ++i) { + const int64_t h_start = head_kv_id * num_groups + block_id * BLOCK_H; + const int64_t h_end = head_kv_id * num_groups + std::min(block_id * BLOCK_H + BLOCK_H, num_groups); + const int64_t h_size = h_end - h_start; + + // get query + const scalar_t* __restrict__ q_ptr = query + bs * q_strideM + h_start * q_strideH; + + int64_t seq_len_kv = is_cross_attn ? encoder_lens[bs] : seq_lens[bs]; + int64_t req_pool_id = req_pool_indices[bs]; + int64_t kv_offset = (has_encoder_lens && (!is_cross_attn)) ? encoder_lens[bs] : 0; + TORCH_CHECK(seq_len_kv <= max_context_len, "seq_len_kv out of scope!"); + TORCH_CHECK(req_pool_id < max_num_reqs, "req_pool_id out of scope!"); + if (sliding_window_size > 0 && seq_len_kv > sliding_window_size) { + kv_offset = seq_len_kv - sliding_window_size; + seq_len_kv = sliding_window_size; + } + const int64_t SPLIT_SIZE = div_up(seq_len_kv, num_kv_splits); + const int64_t kv_start = kv_id * SPLIT_SIZE; + const int64_t kv_end = std::min(kv_start + SPLIT_SIZE, seq_len_kv); + + fill_stub(s_prime, 0.f, BLOCK_H); + fill_stub(m_prime, -std::numeric_limits::infinity(), BLOCK_H); + + // get v_prime, and init to zero + float* __restrict__ v_prime = attn_logits + bs * l_stride0 + h_start * l_stride1 + kv_id * l_stride2; + for (int64_t h = 0; h < h_size; ++h) { + fill_stub(v_prime + h * l_stride1, 0.f, head_size_v); + } + + // loop over K and V sequence with BLOCK_N + for (int64_t n = kv_start; n < kv_end; n += BLOCK_N) { + int64_t n_size = std::min(BLOCK_N, kv_end - n); + + // calculate Q @ K + index_gemm_kernel_nt( + /* A */ q_ptr, + /* B */ k_buffer + head_kv_id * k_strideH, + /* C */ s_i, + /* ind */ req_to_token + req_pool_id * max_context_len + n + kv_offset, + /* scl */ sm_scale, + /* M */ h_size, + /* N */ n_size, + /* K */ head_size, + /* lda */ q_strideH, + /* ldb */ k_strideN, + /* ldc */ BLOCK_N, + /* mtt */ max_total_num_tokens); + + if (has_logit_cap) { + at::vec::map( + [logit_cap, rlogit_cap](Vec x) { return Vec(logit_cap) * (x * Vec(rlogit_cap)).tanh(); }, + s_i, + s_i, + BLOCK_H * BLOCK_N); + } + + // update the sm_scale coefficients + for (int64_t h = 0; h < h_size; ++h) { + // m_i: max value per row + float m_i = at::vec::reduce_all( + [](Vec& x, Vec& y) { return at::vec::maximum(x, y); }, s_i + h * BLOCK_N, n_size); + m_i = std::max(m_i, m_prime[h]); + + // m_delta <- exp(m' - m_i) + m_delta[h] = std::exp(m_prime[h] - m_i); + + // s_delta <- exp(s_i - m_i) + at::vec::map( + [m_i](Vec x) { return (x - Vec(m_i)).exp_u20(); }, s_delta + h * BLOCK_N, s_i + h * BLOCK_N, n_size); + + // s' <- s' * m_delta + sum(s_delta) + s_prime[h] *= m_delta[h]; + s_prime[h] += at::vec::reduce_all([](Vec& x, Vec& y) { return x + y; }, s_delta + h * BLOCK_N, n_size); + + m_prime[h] = m_i; + } + + // calculate V' <- s_delta @ V + V' * m_delta + index_gemm_kernel_nn( + /* A */ s_delta, + /* B */ v_buffer + head_kv_id * v_strideH, + /* C */ v_prime, + /* ind */ req_to_token + req_pool_id * max_context_len + n + kv_offset, + /* scl */ m_delta, + /* M */ h_size, + /* N */ head_size_v, + /* K */ n_size, + /* lda */ BLOCK_N, + /* ldb */ v_strideN, + /* ldc */ l_stride1, + /* mtt */ max_total_num_tokens); + } // loop with KV blocks + + // only update v' when kv_split_size > 0 + if (kv_end > kv_start) { + for (int64_t h = 0; h < h_size; ++h) { + float s = 1 / s_prime[h]; + at::vec::map( + [s](Vec out) { return out * Vec(s); }, v_prime + h * l_stride1, v_prime + h * l_stride1, head_size_v); + (v_prime + h * l_stride1)[head_size_v] = m_prime[h] + std::log(s_prime[h]); + } + } else { + for (int64_t h = 0; h < h_size; ++h) { + (v_prime + h * l_stride1)[head_size_v] = -std::numeric_limits::infinity(); + } + } + + // move to the next index + data_index_step(bs, batches, head_kv_id, num_heads_kv, block_id, num_blocks, kv_id, num_kv_splits); + } + }); + + decode_accumulate_kv_splits( + output, attn_logits, sinks, batches, num_heads, head_size_v, num_kv_splits, l_stride1, l_stride2, has_sink); +} // GQA/MQA + +} // anonymous namespace + +// query: [num_tokens, num_heads, head_size] +// output: [num_tokens, num_heads, head_size] +// k_buffer: [max_total_num_tokens, num_heads, head_size] +// v_buffer: [max_total_num_tokens, num_heads, head_size_v] +// attn_logits: [num_seqs, num_heads, num_kv_splits, head_size_v + 1] +// req_to_token: [max_num_reqs, max_context_len] int32 or int64 +// req_pool_indices: [num_seqs] int64 +// seq_lens: [num_seqs] int64 +// encoder_lens: [num_seqs] int64 or None +// sinks: [num_heads] or None +void decode_attention_cpu( + at::Tensor& query, + at::Tensor& k_buffer, + at::Tensor& v_buffer, + at::Tensor& output, + const std::optional& key, + const std::optional& value, + const std::optional& loc, + at::Tensor& attn_logits, + at::Tensor& req_to_token, + at::Tensor& req_pool_indices, + at::Tensor& seq_lens, + double sm_scale, + double logit_cap, + bool is_cross_attn, + int64_t sliding_window_size, + std::optional encoder_lens, + std::optional sinks) { + CHECK_LAST_DIM_CONTIGUOUS_INPUT(query); + CHECK_LAST_DIM_CONTIGUOUS_INPUT(k_buffer); + CHECK_LAST_DIM_CONTIGUOUS_INPUT(v_buffer); + CHECK_DIM(3, query); + CHECK_DIM(3, k_buffer); + CHECK_DIM(3, v_buffer); + + int64_t num_seqs = seq_lens.size(0); + int64_t max_num_reqs = req_to_token.size(0); + int64_t max_context_len = req_to_token.size(1); + int64_t max_total_num_tokens = k_buffer.size(0); + + int64_t num_heads = query.size(1); + int64_t num_heads_kv = k_buffer.size(1); + int64_t head_size = query.size(2); + int64_t head_size_v = v_buffer.size(2); + + int64_t num_kv_splits = attn_logits.size(2); + + CHECK_EQ(attn_logits.size(0), num_seqs); + CHECK_EQ(attn_logits.size(1), num_heads); + CHECK_EQ(attn_logits.size(3), head_size_v + 1); + CHECK_EQ(attn_logits.scalar_type(), at::kFloat); + + // strides for query + int64_t q_strideM = query.stride(0); + int64_t q_strideH = query.stride(1); + + // strides for k_buffer and v_buffer + int64_t k_strideN = k_buffer.stride(0); + int64_t k_strideH = k_buffer.stride(1); + int64_t v_strideN = v_buffer.stride(0); + int64_t v_strideH = v_buffer.stride(1); + + // check index data types + const auto index_dtype = req_to_token.scalar_type(); + TORCH_CHECK( + index_dtype == at::kInt || index_dtype == at::kLong, + "decode: expect req_to_token to be int32 or int64, got ", + index_dtype); + TORCH_CHECK(seq_lens.scalar_type() == at::kLong, "decode: expect req_lens to be int64, got ", seq_lens.scalar_type()); + TORCH_CHECK( + req_pool_indices.scalar_type() == at::kLong, + "decode: expect req_pool_indices to be int64, got ", + req_pool_indices.scalar_type()); + + // check if we have MLA here + void* k_buffer_data = k_buffer.data_ptr(); + void* v_buffer_data = v_buffer.data_ptr(); + const bool is_mla = (k_buffer_data == v_buffer_data) && (num_heads_kv == 1) && (head_size == head_size_v + 64); + + // block length for k_buffer and v_buffer + constexpr int BLOCK_N = 256; + + // buffer for packing k_cache and v_cache + int num_threads = at::get_num_threads(); + int64_t size_per_thread = is_mla ? BLOCK_N * head_size + BLOCK_N * head_size_v : 0; + auto buffer = at::empty({num_threads, size_per_thread}, k_buffer.options()); + bool has_encoder_lens = encoder_lens.has_value(); + // Since encoder_lens is not used when it is None, encoder_lens_t can be initialized as any tensor of int64_t dtype. + at::Tensor encoder_lens_t = seq_lens; + if (has_encoder_lens) { + encoder_lens_t = encoder_lens.value(); + CHECK_EQ(encoder_lens_t.size(0), num_seqs); + } + bool has_sink = sinks.has_value(); + at::Tensor sinks_tensor = has_sink ? sinks.value() : at::empty({num_heads}, query.options()); + CHECK_DIM(1, sinks_tensor); + CHECK_EQ(sinks_tensor.size(0), num_heads); + AT_DISPATCH_REDUCED_FLOATING_TYPES(query.scalar_type(), "decode_attention_kernel", [&] { + AT_DISPATCH_INDEX_TYPES(index_dtype, "decode_attention_indices", [&] { + if (key.has_value()) { + TORCH_CHECK(value.has_value(), "key and value should have values at the same time") + TORCH_CHECK(loc.has_value(), "loc must be given when key/value are given") + auto loc_tensor = loc.value(); + CHECK_DIM(1, loc_tensor); + CHECK_EQ(loc_tensor.numel(), num_seqs); + auto key_tensor = key.value(); + auto value_tensor = value.value(); + // for MLA, key and value shares the same storage and value could be non-contiguous + CHECK_LAST_DIM_CONTIGUOUS_INPUT(key_tensor); + CHECK_LAST_DIM_CONTIGUOUS_INPUT(value_tensor); + CHECK_DIM(3, key_tensor); + CHECK_DIM(3, value_tensor); + // strides for new key and value + int64_t nk_strideN = key_tensor.stride(0); + int64_t nk_strideH = key_tensor.stride(1); + int64_t nv_strideN = value_tensor.stride(0); + int64_t nv_strideH = value_tensor.stride(1); + // update the kv buffer + decode_set_kv_buffer( + (scalar_t*)k_buffer_data, + (scalar_t*)v_buffer_data, + key_tensor.data_ptr(), + value_tensor.data_ptr(), + loc_tensor.data_ptr(), + num_seqs, + num_heads_kv, + head_size, + head_size_v, + k_strideN, + k_strideH, + v_strideN, + v_strideH, + nk_strideN, + nk_strideH, + nv_strideN, + nv_strideH, + is_mla); + } + + if (num_heads == num_heads_kv) { + // MHA + decode_attention_kernel_impl( + output.data_ptr(), + attn_logits.data_ptr(), + query.data_ptr(), + (const scalar_t*)k_buffer_data, + (const scalar_t*)v_buffer_data, + req_to_token.data_ptr(), + req_pool_indices.data_ptr(), + seq_lens.data_ptr(), + encoder_lens_t.data_ptr(), + sinks_tensor.data_ptr(), + num_seqs, + num_heads, + head_size, + head_size_v, + num_kv_splits, + q_strideM, + q_strideH, + k_strideN, + k_strideH, + v_strideN, + v_strideH, + sm_scale, + logit_cap, + max_num_reqs, + max_context_len, + max_total_num_tokens, + sliding_window_size, + is_cross_attn, + has_encoder_lens, + has_sink); + } else if (is_mla) { + // MLA + decode_attention_mla_kernel_impl( + output.data_ptr(), + attn_logits.data_ptr(), + query.data_ptr(), + (const scalar_t*)k_buffer_data, + (const scalar_t*)v_buffer_data, + req_to_token.data_ptr(), + req_pool_indices.data_ptr(), + seq_lens.data_ptr(), + buffer.data_ptr(), + sinks_tensor.data_ptr(), + num_seqs, + num_heads, + head_size, + head_size_v, + num_kv_splits, + q_strideM, + q_strideH, + k_strideN, + k_strideH, + v_strideN, + v_strideH, + sm_scale, + logit_cap, + max_num_reqs, + max_context_len, + max_total_num_tokens, + size_per_thread, + has_sink); + } else { + // GQA/MQA + decode_attention_grouped_kernel_impl( + output.data_ptr(), + attn_logits.data_ptr(), + query.data_ptr(), + (const scalar_t*)k_buffer_data, + (const scalar_t*)v_buffer_data, + req_to_token.data_ptr(), + req_pool_indices.data_ptr(), + seq_lens.data_ptr(), + encoder_lens_t.data_ptr(), + sinks_tensor.data_ptr(), + num_seqs, + num_heads, + num_heads_kv, + head_size, + head_size_v, + num_kv_splits, + q_strideM, + q_strideH, + k_strideN, + k_strideH, + v_strideN, + v_strideH, + sm_scale, + logit_cap, + max_num_reqs, + max_context_len, + max_total_num_tokens, + sliding_window_size, + is_cross_attn, + has_encoder_lens, + has_sink); + } + }); + }); +} diff --git a/csrc/cpu/sgl-kernels/extend.cpp b/csrc/cpu/sgl-kernels/extend.cpp new file mode 100644 index 000000000000..3c0a9ebe5460 --- /dev/null +++ b/csrc/cpu/sgl-kernels/extend.cpp @@ -0,0 +1,568 @@ +// Adapted from +// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu + +// clang-format off + +#include "common.h" +#include "flash_attn.h" +#include "gemm.h" + +namespace { + +// [NOTE]: extend attention for CPU +// 1. BLOCK_M and BLOCK_N tuned for various seq lengths +// 2. can handle non-contiguous k_extend and v_extend +// 3. computes attention for prefix and extend separately +// 4. TODO: apply head dimension blocking to optimize GQA +// 5. optional tree mask for speculative decoding TARGET_VERIFY (EAGLE topk > 1): +// `tree_mask` is a flat [batches * qlen * qlen] bool tensor in +// TreeMaskMode::QLEN_ONLY layout, where qlen == extend_seq_lens[bs] == +// max_len_extend (uniform across the batch, equal to draft_token_num). +// Row i = query draft token, column j = key draft token; true means query i +// may attend key j (each row marks self + ancestors + root). The committed +// prefix (stage 1) is implicitly fully visible to every draft token, which +// is why the mask only covers the qlen x qlen new-token block; the GPU +// FULL_MASK layout carries the prefix columns explicitly but they are +// all-true for EAGLE. When tree_mask is absent, stage 2 falls back to the +// plain causal mask (correct for non-spec extend and topk == 1 chains). +// + +template +void extend_attention_kernel_impl( + scalar_t* __restrict__ o_extend, + const scalar_t* __restrict__ q_extend, + const scalar_t* __restrict__ k_extend, + const scalar_t* __restrict__ v_extend, + const scalar_t* __restrict__ k_buffer, + const scalar_t* __restrict__ v_buffer, + const index_t* __restrict__ req_to_token, + const int64_t* __restrict__ req_pool_indices, + const int64_t* __restrict__ seq_lens, + const int64_t* __restrict__ encoder_lens, + const index_t* __restrict__ extend_seq_lens, + const index_t* __restrict__ extend_start_loc, + const void* __restrict__ buffer, + const scalar_t* __restrict__ sinks, + const bool* __restrict__ tree_mask, + int batches, + int num_heads, + int num_heads_kv, + int head_size, + int head_size_v, + int q_strideM, + int q_strideH, + int ke_strideN, + int ke_strideH, + int ve_strideN, + int ve_strideH, + int k_strideN, + int k_strideH, + int v_strideN, + int v_strideH, + float sm_scale, + int max_num_reqs, + int max_context_len, + int max_total_num_tokens, + int max_len_extend, + int buffer_size_per_thread, + int64_t sliding_window_size, + bool is_prefix_skipped, + bool is_cross_attn, + bool has_encoder_lens, + bool has_sink) { + // strides + const int o_strideM = num_heads * head_size_v; + const int o_strideH = head_size_v; + + // we use same buffer for packed key and value + const int ldb_tmp = std::max(head_size, head_size_v); + + const int num_groups = num_heads / num_heads_kv; + TORCH_CHECK(num_groups * num_heads_kv == num_heads); + + // number of blocks along M + int MB = div_up(max_len_extend, BLOCK_M); + + // parallel on [batches, num_heads, BM] + at::parallel_for(0, batches * num_heads * MB, 0, [&](int begin, int end) { + int bs{0}, head_id{0}, mb{0}; + data_index_init(begin, bs, batches, head_id, num_heads, mb, MB); + + int tid = at::get_thread_num(); + // s_i: [BLOCK_M, BLOCK_N] + float* __restrict__ s_i = reinterpret_cast((char*)(buffer) + tid * buffer_size_per_thread); + + // v_prime: [BLOCK_M, head_size_v] + float* __restrict__ v_prime = s_i + BLOCK_M * BLOCK_N; + + // s_delta: [BLOCK_M, BLOCK_N] + scalar_t* __restrict__ s_delta = reinterpret_cast(v_prime + BLOCK_M * head_size_v); + + // Btmp: [BLOCK_N, max(head_size, head_size_v)] + scalar_t* __restrict__ Btmp = reinterpret_cast(s_delta + BLOCK_M * BLOCK_N); + + // init Btmp just once for each thread to prevent NaN + fill_stub(Btmp, 0.f, BLOCK_N * ldb_tmp); + fill_stub(s_delta, 0.f, BLOCK_M * BLOCK_N); + + alignas(64) float s_prime[BLOCK_M]; + alignas(64) float m_prime[BLOCK_M]; + + for (int i = begin; i < end; ++i) { + // seq_len = prefix + extend + int head_kv_id = head_id / num_groups; + int seq_len = seq_lens[bs]; + int seq_len_extend = extend_seq_lens[bs]; + int seq_len_prefix = seq_len - seq_len_extend; + int seq_extend_start_loc = extend_start_loc[bs]; + + int req_pool_id = req_pool_indices[bs]; + int kv_offset = (has_encoder_lens && (!is_cross_attn)) ? encoder_lens[bs] : 0; + TORCH_CHECK(seq_len_prefix >= 0, "prefix len < 0!"); + TORCH_CHECK(seq_len <= max_context_len, "seq_len out of scope!"); + TORCH_CHECK(req_pool_id < max_num_reqs, "req_pool_id out of scope!"); + + if (is_prefix_skipped) { + TORCH_CHECK(seq_len_prefix == 0, "extend attention: expect seq_len_prefix to be 0, got ", seq_len_prefix); + } + + if (tree_mask != nullptr) { + // QLEN_ONLY layout assumes a uniform qlen across the batch (TARGET_VERIFY) + TORCH_CHECK( + seq_len_extend == max_len_extend, + "extend attention: tree_mask requires uniform extend_seq_lens, got ", + seq_len_extend, + " vs ", + max_len_extend); + } + + // offset and size in MB + int m = mb * BLOCK_M; + int m_size = std::min(BLOCK_M, seq_len_extend - m); + + if (m_size <= 0) { + data_index_step(bs, batches, head_id, num_heads, mb, MB); + continue; + } + + // get query + const scalar_t* __restrict__ q_ptr = q_extend + (seq_extend_start_loc + m) * q_strideM + head_id * q_strideH; + + // init v', s' and m' + fill_stub(v_prime, 0.f, m_size * head_size_v); + fill_stub(s_prime, 0.f, m_size); + fill_stub(m_prime, -std::numeric_limits::infinity(), m_size); + // stage 1: compute scores with prefix + int kv_start = 0; + int kv_end = is_cross_attn ? encoder_lens[bs] : seq_len_prefix; + for (int n = kv_start; n < kv_end; n += BLOCK_N) { + int n_size = std::min(BLOCK_N, kv_end - n); + + // `n_size` is K in 2nd gemm, pad to TILE_K; + const int padded_n_size = div_up(n_size, TILE_K) * TILE_K; + + // get key and pack + pack_vnni( + /* dst */ Btmp, + /* src */ k_buffer + head_kv_id * k_strideH, + /* ind */ req_to_token + req_pool_id * max_context_len + n + kv_offset, + /* N */ n_size, + /* K */ head_size, + /* ld_src */ k_strideN, + /* ld_dst */ BLOCK_N); + + // calculate s_i <- Q @ K + at::native::cpublas::brgemm( + /* M */ m_size, + /* N */ n_size, + /* K */ head_size, + /* lda */ q_strideM, + /* ldb */ BLOCK_N, + /* ldc */ BLOCK_N, + /* add_C */ false, + /* A */ q_ptr, + /* B */ Btmp, + /* C */ s_i); + + for (int row = 0; row < m_size; ++row) { + if (sliding_window_size > 0) { + int last_col = seq_len_prefix + row + m - sliding_window_size + 1; + if (last_col >= n + n_size) { + continue; + } + fill_stub(s_i + row * BLOCK_N, -std::numeric_limits::infinity(), last_col - n); + } + flash_attn_softmax::apply( + s_i, s_delta, v_prime, s_prime, m_prime, m_size, n_size, padded_n_size, head_size_v, sm_scale, row); + } + + // get value and pack + pack_vnni2( + /* dst */ Btmp, + /* src */ v_buffer + head_kv_id * v_strideH, + /* ind */ req_to_token + req_pool_id * max_context_len + n + kv_offset, + /* K */ n_size, + /* N */ head_size_v, + /* ld_src */ v_strideN, + /* ld_dst */ head_size_v); + + // calculate V' <- s_delta @ V + V' + at::native::cpublas::brgemm( + /* M */ m_size, + /* N */ head_size_v, + /* K */ padded_n_size, // n_size + /* lda */ BLOCK_N, + /* ldb */ head_size_v, + /* ldc */ head_size_v, + /* add_C */ true, + /* A */ s_delta, + /* B */ Btmp, + /* C */ v_prime); + } // loop with seq_len_prefix + if (!is_cross_attn) { + // stage 2: compute the triangle part + int num_keys = std::min(seq_len_extend, m + BLOCK_M); + for (int n = 0; n < num_keys; n += BLOCK_N) { + int n_size = std::min(BLOCK_N, num_keys - n); + + // `n_size` is K in 2nd gemm, pad to TILE_K; + const int padded_n_size = div_up(n_size, TILE_K) * TILE_K; + + // get key and pack + pack_vnni( + /* dst */ Btmp, + /* src */ k_extend + (seq_extend_start_loc + n) * ke_strideN + head_kv_id * ke_strideH, + /* N */ n_size, + /* K */ head_size, + /* ld_src */ ke_strideN, + /* ld_dst */ BLOCK_N); + + // calculate s_i <- Q @ K + at::native::cpublas::brgemm( + /* M */ m_size, + /* N */ n_size, + /* K */ head_size, + /* lda */ q_strideM, + /* ldb */ BLOCK_N, + /* ldc */ BLOCK_N, + /* add_C */ false, + /* A */ q_ptr, + /* B */ Btmp, + /* C */ s_i); + + // apply tree mask (speculative TARGET_VERIFY) or causal mask + if (tree_mask != nullptr) { + // [Note] tree mask for EAGLE topk > 1 (TreeMaskMode::QLEN_ONLY). + // mask[bs][m + row][n + col] == false -> query draft token (m + row) + // may not attend key draft token (n + col); set the score to -inf + // before softmax. The tree mask subsumes the causal constraint: + // ancestors always precede descendants in the draft token ordering, + // so permitted keys satisfy j <= i and the causal `num_keys` bound + // above remains valid. + const bool* __restrict__ mask_base = + tree_mask + (static_cast(bs) * seq_len_extend + m) * seq_len_extend + n; + for (int row = 0; row < m_size; ++row) { + float* __restrict__ row_ptr = s_i + row * BLOCK_N; + const bool* __restrict__ mask_ptr = mask_base + static_cast(row) * seq_len_extend; + for (int col = 0; col < n_size; ++col) { + if (!mask_ptr[col]) { + row_ptr[col] = -std::numeric_limits::infinity(); + } + } + } + } else if (n + n_size - 1 > m) { + // apply causal mask + // [Note] condition to apply causal mask. + // Mask any block whose last key (n + n_size - 1) is strictly after the first query position (m), i.e. n + + // n_size - 1 > m. The original condition was `num_keys - n <= BLOCK_N` (last n-block only). That was + // correct when BLOCK_M <= BLOCK_N/2 because earlier n-blocks were guaranteed to contain only past keys. + // With BLOCK_M=512, BLOCK_N=768: + // BLOCK_M > BLOCK_N/2, so the first n-block can contain future keys. + // Example: m=512 (mb=1), num_keys=1024, first n-block covers keys [0, 768). + // Query row=0 is at position 512, so keys 513..767 are future and must be + // masked — but `num_keys - 0 = 1024 > BLOCK_N` skips masking entirely, + // producing wrong (non-causal) attention for rows 0..254 of this m-block. + for (int row = 0; row < m_size; ++row) { + int last_col = m + row - n; + // [Note] mask the entire row if last_col < 0. + // Clamp to -1: when n > m + row every key in this block is a future + // key, so the entire row should be masked. Without this clamp, + // last_col+1 <= 0 and fill_stub would write before row_ptr. + last_col = std::max(last_col, -1); + // fill [last_col + 1, n_size) to -inf + float* row_ptr = s_i + row * BLOCK_N; + fill_stub(row_ptr + last_col + 1, -std::numeric_limits::infinity(), n_size - last_col - 1); + } + } + + for (int row = 0; row < m_size; ++row) { + if (sliding_window_size > 0 && row + m + 1 >= n + sliding_window_size - 1 && + row + m + 1 < n + sliding_window_size + n_size) { + fill_stub( + s_i + row * BLOCK_N, -std::numeric_limits::infinity(), row + m - n - sliding_window_size + 1); + } else if (sliding_window_size > 0 && row + m + 1 >= n + sliding_window_size) { + continue; + } + flash_attn_softmax::apply( + s_i, s_delta, v_prime, s_prime, m_prime, m_size, n_size, padded_n_size, head_size_v, sm_scale, row); + } + + // get value and pack + pack_vnni2( + /* dst */ Btmp, + /* src */ v_extend + (seq_extend_start_loc + n) * ve_strideN + head_kv_id * ve_strideH, + /* K */ n_size, + /* N */ head_size_v, + /* ld_src */ ve_strideN, + /* ld_dst */ head_size_v); + + // calculate V' <- s_delta @ V + V' + at::native::cpublas::brgemm( + /* M */ m_size, + /* N */ head_size_v, + /* K */ padded_n_size, // n_size + /* lda */ BLOCK_N, + /* ldb */ head_size_v, + /* ldc */ head_size_v, + /* add_C */ true, + /* A */ s_delta, + /* B */ Btmp, + /* C */ v_prime); + } // loop with seq_len_extend + } + scalar_t* __restrict__ out_ptr = o_extend + (seq_extend_start_loc + m) * o_strideM + head_id * o_strideH; + for (int row = 0; row < m_size; ++row) { + if (has_sink) { + s_prime[row] += std::exp(sinks[head_id] - m_prime[row]); + } + float s = 1 / s_prime[row]; + copy_stub(out_ptr + row * o_strideM, v_prime + row * head_size_v, s, head_size_v); + } + + // move to the next index + data_index_step(bs, batches, head_id, num_heads, mb, MB); + } + at::native::cpublas::brgemm_release(); + }); +} + +} // anonymous namespace + +template +inline int resize_buffer(at::Tensor& buffer, int num_threads, int head_size, int head_size_v) { + static_assert(BLOCK_M <= BLOCK_N, "Make sure BLOCK_M <= BLOCK_N to prevent buffer overflows during causal masking"); + const int size_per_thread = + /* s_i */ BLOCK_M * BLOCK_N * sizeof(float) + + /* v_prime */ BLOCK_M * head_size_v * sizeof(float) + + /* s_delta */ BLOCK_M * BLOCK_N * sizeof(uint16_t) + + /* Btmp */ BLOCK_N * std::max(head_size, head_size_v) * sizeof(uint16_t); + + buffer.resize_({num_threads, size_per_thread}); + return size_per_thread; +} + +#define LAUNCH_EXTEND_ATTENTION_KERNEL(BLOCK_M, BLOCK_N) \ + do { \ + int sz = resize_buffer(buffer, num_threads, head_size, head_size_v); \ + \ + extend_attention_kernel_impl( \ + o_extend.data_ptr(), \ + q_extend.data_ptr(), \ + k_extend.data_ptr(), \ + v_extend.data_ptr(), \ + k_buffer.data_ptr(), \ + v_buffer.data_ptr(), \ + req_to_token.data_ptr(), \ + req_pool_indices.data_ptr(), \ + seq_lens.data_ptr(), \ + encoder_lens_t.data_ptr(), \ + extend_seq_lens.data_ptr(), \ + extend_start_loc.data_ptr(), \ + buffer.data_ptr(), \ + sinks_tensor.data_ptr(), \ + tree_mask_ptr, \ + num_seqs, \ + num_heads, \ + num_heads_kv, \ + head_size, \ + head_size_v, \ + q_strideM, \ + q_strideH, \ + ke_strideN, \ + ke_strideH, \ + ve_strideN, \ + ve_strideH, \ + k_strideN, \ + k_strideH, \ + v_strideN, \ + v_strideH, \ + sm_scale, \ + max_num_reqs, \ + max_context_len, \ + max_total_num_tokens, \ + max_len_extend, \ + sz, \ + sliding_window_size, \ + is_prefix_skipped, \ + is_cross_attn, \ + has_encoder_lens, \ + has_sink); \ + } while (0) + +// q_extend, k_extend, v_extend, o_extend: contiguous tensors +// k_buffer, v_buffer: (prefix + extend) tensors in mem_manager +// +// q_extend: [num_tokens, num_heads, head_size] +// k_extend: [num_extend_tokens, num_heads, head_size] +// v_extend: [num_extend_tokens, num_heads, head_size] +// o_extend: [num_tokens, num_heads, head_size] +// k_buffer: [max_total_num_tokens, num_heads, head_size] +// v_buffer: [max_total_num_tokens, num_heads, head_size] +// req_to_token: [max_num_reqs, max_context_len] int32 or int64 +// req_pool_indices: [num_seqs] int64 +// seq_lens: [num_seqs] int64 +// extend_seq_lens: [num_seqs] +// extend_start_loc: [num_seqs] +// encoder_lens: [num_seqs] int64 or None +// sinks: [num_heads] or None +// tree_mask: [num_seqs * max_len_extend * max_len_extend] bool or None +// TreeMaskMode::QLEN_ONLY tree mask for speculative TARGET_VERIFY; see [NOTE] 5 above. +void extend_attention_cpu( + at::Tensor& q_extend, + const std::optional& k_extend_opt, + const std::optional& v_extend_opt, + at::Tensor& o_extend, + at::Tensor& k_buffer, + at::Tensor& v_buffer, + at::Tensor& req_to_token, + at::Tensor& req_pool_indices, + at::Tensor& seq_lens, + at::Tensor& extend_seq_lens, + at::Tensor& extend_start_loc, + int64_t max_len_extend, + double sm_scale, + double logit_cap, + bool is_cross_attn, + int64_t sliding_window_size, + std::optional encoder_lens, + std::optional sinks, + std::optional tree_mask) { + if (!is_cross_attn) { + TORCH_CHECK( + k_extend_opt.has_value() && v_extend_opt.has_value(), + "k_extend and v_extend are required for non-cross attention"); + } + // Since k_extend and v_extend are not used for cross attention, they can be initialized as k_buffer and v_buffer + // here. + auto k_extend = k_extend_opt.has_value() ? k_extend_opt.value() : k_buffer; + auto v_extend = v_extend_opt.has_value() ? v_extend_opt.value() : v_buffer; + + CHECK_LAST_DIM_CONTIGUOUS_INPUT(q_extend); + CHECK_INPUT(o_extend); + CHECK_LAST_DIM_CONTIGUOUS_INPUT(k_extend); + CHECK_LAST_DIM_CONTIGUOUS_INPUT(v_extend); + CHECK_LAST_DIM_CONTIGUOUS_INPUT(k_buffer); + CHECK_LAST_DIM_CONTIGUOUS_INPUT(v_buffer); + + int num_seqs = seq_lens.size(0); + int max_num_reqs = req_to_token.size(0); + int max_context_len = req_to_token.size(1); + int max_total_num_tokens = k_buffer.size(0); + + int num_heads = q_extend.size(1); + int num_heads_kv = k_extend.size(1); + int head_size = q_extend.size(2); + int head_size_v = v_extend.size(2); + + // strides for q_extend, k_extend and v_extend + int q_strideM = q_extend.stride(0); + int q_strideH = q_extend.stride(1); + int ke_strideN = k_extend.stride(0); + int ke_strideH = k_extend.stride(1); + int ve_strideN = v_extend.stride(0); + int ve_strideH = v_extend.stride(1); + + // strides for k_buffer and v_buffer + int k_strideN = k_buffer.stride(0); + int k_strideH = k_buffer.stride(1); + int v_strideN = v_buffer.stride(0); + int v_strideH = v_buffer.stride(1); + + // check sizes + CHECK_EQ(req_pool_indices.size(0), num_seqs); + CHECK_EQ(extend_seq_lens.size(0), num_seqs); + CHECK_EQ(extend_start_loc.size(0), num_seqs); + CHECK_EQ(v_extend.size(1), num_heads_kv); + CHECK_EQ(k_buffer.size(1), v_buffer.size(1)); + + // MLA will skip prefix part + const bool is_prefix_skipped = k_buffer.size(1) != num_heads_kv; + + // check index data types + const auto index_dtype = req_to_token.scalar_type(); + TORCH_CHECK( + index_dtype == at::kInt || index_dtype == at::kLong, + "extend: expect req_to_token to be int32 or int64, got ", + index_dtype); + TORCH_CHECK(seq_lens.scalar_type() == at::kLong, "extend: expect req_lens to be int64, got ", seq_lens.scalar_type()); + TORCH_CHECK( + req_pool_indices.scalar_type() == at::kLong, + "extend: expect req_pool_indices to be int64, got ", + req_pool_indices.scalar_type()); + TORCH_CHECK( + extend_seq_lens.scalar_type() == index_dtype && extend_start_loc.scalar_type() == index_dtype, + "extend: expect extend_seq_lens and extend_start_loc to have same dtype as req_to_token."); + + // D and DV need to be 32x as we transpose by 512-bit + TORCH_CHECK(head_size % 32 == 0, "invalid head_size ", head_size); + TORCH_CHECK(head_size_v % 32 == 0, "invalid head_size_v ", head_size_v); + + int num_threads = at::get_num_threads(); + auto buffer = at::empty({}, q_extend.options().dtype(at::kChar)); + + bool has_encoder_lens = encoder_lens.has_value(); + // Since encoder_lens is not used when it is None, encoder_lens_t can be initialized as any tensor of int64_t dtype. + at::Tensor encoder_lens_t = seq_lens; + if (has_encoder_lens) { + encoder_lens_t = encoder_lens.value(); + CHECK_EQ(encoder_lens_t.size(0), num_seqs); + } + bool has_sink = sinks.has_value(); + at::Tensor sinks_tensor = has_sink ? sinks.value() : at::empty({num_heads}, q_extend.options()); + CHECK_DIM(1, sinks_tensor); + CHECK_EQ(sinks_tensor.size(0), num_heads); + + const bool* tree_mask_ptr = nullptr; + if (tree_mask.has_value()) { + const at::Tensor& tree_mask_t = tree_mask.value(); + CHECK_INPUT(tree_mask_t); + TORCH_CHECK( + tree_mask_t.scalar_type() == at::kBool, "extend: expect tree_mask to be bool, got ", tree_mask_t.scalar_type()); + TORCH_CHECK( + tree_mask_t.numel() == static_cast(num_seqs) * max_len_extend * max_len_extend, + "extend: expect tree_mask numel to be num_seqs * max_len_extend^2 = ", + static_cast(num_seqs) * max_len_extend * max_len_extend, + ", got ", + tree_mask_t.numel()); + TORCH_CHECK(!is_cross_attn, "extend: tree_mask is not supported for cross attention"); + // The window mask derives query positions from the row index + // (seq_len_prefix + m + row), but tree-mask rows sit at their tree depth, + // which is <= the row index; combining the two would over-mask the prefix. + TORCH_CHECK(sliding_window_size <= 0, "extend: tree_mask is not supported with sliding window attention"); + tree_mask_ptr = tree_mask_t.data_ptr(); + } + + AT_DISPATCH_REDUCED_FLOATING_TYPES(q_extend.scalar_type(), "extend_attention_kernel", [&] { + AT_DISPATCH_INDEX_TYPES(index_dtype, "extend_attention_indices", [&] { + if (max_len_extend <= 256) { + LAUNCH_EXTEND_ATTENTION_KERNEL(32, 64); + } else if (max_len_extend <= 1024) { + LAUNCH_EXTEND_ATTENTION_KERNEL(128, 256); + } else if (max_len_extend <= 4096) { + LAUNCH_EXTEND_ATTENTION_KERNEL(256, 768); + } else { // max_len_extend > 4096 + LAUNCH_EXTEND_ATTENTION_KERNEL(512, 768); + } + }); + }); +} diff --git a/csrc/cpu/sgl-kernels/flash_attn.h b/csrc/cpu/sgl-kernels/flash_attn.h new file mode 100644 index 000000000000..24ceec64b5dd --- /dev/null +++ b/csrc/cpu/sgl-kernels/flash_attn.h @@ -0,0 +1,250 @@ +// Adapted from +// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu + +// clang-format off + +#pragma once +#include "common.h" +#include "vec.h" +#include "vec_pack.h" + +template +inline void fill_stub(scalar_t* __restrict__ out, float val, int size) { + using Vec = at::vec::Vectorized; + constexpr int kVecSize = Vec::size(); + const Vec data_vec = Vec(static_cast(val)); + int d = 0; +#pragma GCC unroll 4 + for (; d <= size - kVecSize; d += kVecSize) { + data_vec.store(out + d); + } + if (size - d > 0) { + data_vec.store(out + d, size - d); + } +} + +template +inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ input) { + static_assert(BLOCK_N % 32 == 0); + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + + constexpr int COLS = BLOCK_N / 16; + auto store = [&](auto i) { + constexpr int col = i % COLS; + // for COLS = 2, 4 use 512bit store + if constexpr (col % 2 == 0) { + auto [a_fvec0, a_fvec1] = load_float_vec2(input + col * 16); + bVec out_bvec = convert_from_float_ext(a_fvec0, a_fvec1); + out_bvec.store(out + col * 16); + } + }; + Unroll{}(store); +} + +template +inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ acc, float s, int size) { + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + constexpr int kVecSize = bVec::size(); + const fVec s_fvec = fVec(s); + int d = 0; +#pragma GCC unroll 4 + for (; d <= size - kVecSize; d += kVecSize) { + auto [a_fvec0, a_fvec1] = load_float_vec2(acc + d); + a_fvec0 = a_fvec0 * s_fvec; + a_fvec1 = a_fvec1 * s_fvec; + bVec out_bvec = convert_from_float_ext(a_fvec0, a_fvec1); + out_bvec.store(out + d); + } + for (; d < size; ++d) { + out[d] = static_cast(acc[d] * s); + } +} + +#if defined(CPU_CAPABILITY_AVX512) +template <> +inline void copy_stub(at::BFloat16* __restrict__ out, const float* __restrict__ acc, float s, int size) { + const __m512 vscale = _mm512_set1_ps(s); + int d = 0; +#pragma GCC unroll 4 + for (; d <= size - 32; d += 32) { + __m512 va0 = _mm512_mul_ps(_mm512_loadu_ps(acc + d), vscale); + __m512 va1 = _mm512_mul_ps(_mm512_loadu_ps(acc + d + 16), vscale); + __m512i vb = (__m512i)(_mm512_cvtne2ps_pbh(va1, va0)); + _mm512_storeu_si512(out + d, vb); + } + int remainder = size - d; + if (remainder > 0) { + if (remainder <= 16) { + const __mmask16 vmask = (1ULL << remainder) - 1; + __m512 va = _mm512_mul_ps(_mm512_maskz_loadu_ps(vmask, acc + d), vscale); + __m256i vb = (__m256i)(_mm512_cvtneps_pbh(va)); + _mm256_mask_storeu_epi16(reinterpret_cast<__m256i*>(out + d), vmask, vb); + } else { // remainder > 16 + const __mmask16 vmask = (1ULL << (remainder - 16)) - 1; + __m512 va0 = _mm512_mul_ps(_mm512_loadu_ps(acc + d), vscale); + __m512 va1 = _mm512_mul_ps(_mm512_maskz_loadu_ps(vmask, acc + d + 16), vscale); + __m512i vb = (__m512i)(_mm512_cvtne2ps_pbh(va1, va0)); + const __mmask32 vmask2 = (1ULL << remainder) - 1; + _mm512_mask_storeu_epi16(reinterpret_cast<__m512i*>(out + d), vmask2, vb); + } + } +} +#endif + +template +struct flash_attn_softmax { + static inline void apply( + float* __restrict__ s_i, + scalar_t* __restrict__ s_delta2, + float* __restrict__ v_prime, + float* __restrict__ s_prime, + float* __restrict__ m_prime, + int m_size, + int n_size, + int padded_n_size, + int head_size_v, + const float sm_scale, + int row) { + using Vec = at::vec::Vectorized; + const Vec scale_vec = Vec(sm_scale); + float* s_delta = s_i; + // s_i <- s_i * scale + at::vec::map([scale_vec](Vec x) { return x * scale_vec; }, s_i + row * BLOCK_N, s_i + row * BLOCK_N, n_size); + + // m_i: max value per row + float m_i = + at::vec::reduce_all([](Vec& x, Vec& y) { return at::vec::maximum(x, y); }, s_i + row * BLOCK_N, n_size); + m_i = std::max(m_i, m_prime[row]); + + // m_delta <- exp(m' - m_i) + float m_delta = std::exp(m_prime[row] - m_i); + + // s_delta <- exp(s_i - m_i) + at::vec::map( + [m_i](Vec x) { return (x - Vec(m_i)).fexp_u20(); }, s_delta + row * BLOCK_N, s_i + row * BLOCK_N, n_size); + + // s' <- s' * m_delta + sum(s_delta) + s_prime[row] *= m_delta; + s_prime[row] += at::vec::reduce_all([](Vec& x, Vec& y) { return x + y; }, s_delta + row * BLOCK_N, n_size); + + m_prime[row] = m_i; + + // v' <- v' * m_delta + at::vec::map( + [m_delta](Vec x) { return x * Vec(m_delta); }, + v_prime + row * head_size_v, + v_prime + row * head_size_v, + head_size_v); + + // Keep s_delta row-major for the following brgemm(P @ V), and only + // convert the columns that brgemm will consume. + fill_stub(s_delta + row * BLOCK_N + n_size, 0.f, padded_n_size - n_size); + copy_stub(s_delta2 + row * BLOCK_N, s_delta + row * BLOCK_N, 1.f, padded_n_size); + } +}; + +#if defined(CPU_CAPABILITY_AVX512) +template +struct flash_attn_softmax { + static inline void apply( + float* __restrict__ s_i, + at::BFloat16* __restrict__ s_delta2, + float* __restrict__ v_prime, + float* __restrict__ s_prime, + float* __restrict__ m_prime, + int m_size, + int n_size, + int padded_n_size, + int head_size_v, + const float sm_scale, + int row) { + float* s_delta = s_i; + const __m512 vscale = _mm512_set1_ps(sm_scale); + + int n_remainder = n_size & 15; // 0xF + const __mmask16 vmask = (1ULL << n_remainder) - 1; + + int v_remainder = head_size_v & 15; // 0xF + const __mmask16 vmask1 = (1ULL << v_remainder) - 1; + + constexpr float NEG_INF = -std::numeric_limits::infinity(); + + __m512 va; + __m256i vb; + __m512 vmax; + __m512 vsum; + __m512 vmdelta; + + const __m512 vneg_inf = _mm512_set1_ps(NEG_INF); + + int m = row; + vmax = vneg_inf; + + // s_i <- s_i * scale + int n = 0; + for (; n <= n_size - 16; n += 16) { + va = _mm512_mul_ps(_mm512_loadu_ps(s_i + m * BLOCK_N + n), vscale); + vmax = _mm512_max_ps(va, vmax); + } + if (n_remainder > 0) { + va = _mm512_mul_ps(_mm512_mask_loadu_ps(vneg_inf, vmask, s_i + m * BLOCK_N + n), vscale); + vmax = _mm512_max_ps(va, vmax); + } + + // m_i: max value per row + float m_i = _mm512_reduce_max_ps(vmax); + m_i = std::max(m_i, m_prime[m]); + vmax = _mm512_set1_ps(m_i); + + // m_delta <- exp(m' - m_i) + float m_delta = std::exp(m_prime[m] - m_i); + + // s_delta <- exp(s_i - m_i) + vsum = _mm512_setzero_ps(); + for (n = 0; n <= n_size - 16; n += 16) { + va = _mm512_mul_ps(_mm512_loadu_ps(s_i + m * BLOCK_N + n), vscale); + va = _mm512_fexp_u20_ps(_mm512_sub_ps(va, vmax)); + vsum = _mm512_add_ps(vsum, va); + + vb = (__m256i)(_mm512_cvtneps_pbh(va)); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(s_delta2 + m * BLOCK_N + n), vb); + } + if (n_remainder > 0) { + va = _mm512_mul_ps(_mm512_mask_loadu_ps(vneg_inf, vmask, s_i + m * BLOCK_N + n), vscale); + va = _mm512_fexp_u20_ps(_mm512_sub_ps(va, vmax)); + vsum = _mm512_add_ps(vsum, va); + + vb = (__m256i)(_mm512_cvtneps_pbh(va)); + _mm256_mask_storeu_epi16(reinterpret_cast<__m256i*>(s_delta2 + m * BLOCK_N + n), vmask, vb); + } + + // s' <- s' * m_delta + sum(s_delta) + s_prime[m] *= m_delta; + s_prime[m] += _mm512_reduce_add_ps(vsum); + + m_prime[m] = m_i; + + // pad s_delta with 0, pad_size range from [0, 32) + int pad_size = padded_n_size - n_size; + if (pad_size > 0) { + const __m512i vzero = _mm512_setzero_si512(); + __mmask32 vmask2 = (1ULL << pad_size) - 1; + _mm512_mask_storeu_epi16(reinterpret_cast<__m512i*>(s_delta2 + m * BLOCK_N + n_size), vmask2, vzero); + } + + // v' <- v' * m_delta + vmdelta = _mm512_set1_ps(m_delta); + int k = 0; + for (; k <= head_size_v - 16; k += 16) { + va = _mm512_mul_ps(_mm512_loadu_ps(v_prime + m * head_size_v + k), vmdelta); + _mm512_storeu_ps(reinterpret_cast<__m512*>(v_prime + m * head_size_v + k), va); + } + if (v_remainder > 0) { + va = _mm512_mul_ps(_mm512_maskz_loadu_ps(vmask1, v_prime + m * head_size_v + k), vmdelta); + _mm512_mask_storeu_ps(reinterpret_cast<__m512*>(v_prime + m * head_size_v + k), vmask1, va); + } + } +}; +#endif diff --git a/csrc/cpu/sgl-kernels/mla_cache.cpp b/csrc/cpu/sgl-kernels/mla_cache.cpp new file mode 100644 index 000000000000..e4629f84a21b --- /dev/null +++ b/csrc/cpu/sgl-kernels/mla_cache.cpp @@ -0,0 +1,111 @@ +// vLLM-native CPU cache-write op for MLA's single-latent-buffer KV cache. +// +// `concat_and_cache_mla` (the generic MLA cache-write op used by every GPU +// backend) is registered CUDA-only. This is the CPU counterpart, adapted in +// spirit from SGLang's `store_cache_cpu` (csrc/cpu/kvcache.cpp) but +// generalized to write two source tensors (`kv_c_normed`, `k_pe`) into two +// different column-offset ranges of the SAME destination row -- SGLang's +// version assumes k/v land in two independent, equal-row-width cache +// tensors, which doesn't hold here since MLA's cache is one 576-wide buffer +// and the two column ranges (512-wide, 64-wide) don't match the buffer's +// true per-token stride, so the write can't reuse `store_cache_cpu` as-is. + +#include "common.h" +#include "vec.h" + +namespace { + +template +inline void copy_stub(scalar_t* __restrict__ dst, + const scalar_t* __restrict__ src, int64_t size) { + int64_t d = 0; +#if defined(CPU_CAPABILITY_AVX512) + using Vec = at::vec::Vectorized; + constexpr int64_t kVecSize = Vec::size(); + for (; d <= size - kVecSize; d += kVecSize) { + Vec data = Vec::loadu(src + d); + data.store(dst + d); + } +#endif + for (; d < size; ++d) { + dst[d] = src[d]; + } +} + +template +void concat_and_cache_mla_kernel_impl( + const scalar_t* __restrict__ kv_c_normed, // [num_tokens, kv_lora_rank] + const scalar_t* __restrict__ k_pe, // [num_tokens, qk_rope_head_dim] + scalar_t* __restrict__ kv_cache, // [.., kv_lora_rank + qk_rope_head_dim] + const index_t* __restrict__ slot_mapping, // [num_tokens] + int64_t num_tokens, int64_t kv_lora_rank, int64_t qk_rope_head_dim, + int64_t kv_c_stride, int64_t k_pe_stride, int64_t cache_stride) { + at::parallel_for(0, num_tokens, 0, [&](int64_t begin, int64_t end) { + for (int64_t i = begin; i < end; ++i) { + int64_t slot = static_cast(slot_mapping[i]); + if (slot < 0) { + // padded/invalid token, matches `concat_and_cache_mla`'s semantics. + continue; + } + scalar_t* __restrict__ cache_row = kv_cache + slot * cache_stride; + copy_stub(cache_row, kv_c_normed + i * kv_c_stride, kv_lora_rank); + copy_stub(cache_row + kv_lora_rank, k_pe + i * k_pe_stride, + qk_rope_head_dim); + } + }); +} + +} // namespace + +// kv_c_normed : [num_tokens, kv_lora_rank] +// k_pe : [num_tokens, qk_rope_head_dim] or [num_tokens, 1, +// qk_rope_head_dim] kv_cache : [num_blocks, block_size, kv_lora_rank + +// qk_rope_head_dim] slot_mapping: [num_tokens] int32/int64, absolute physical +// row index +// (block_id * block_size + block_offset); negative entries are +// skipped (padded tokens), matching `concat_and_cache_mla`. +void concat_and_cache_mla_cpu(const at::Tensor& kv_c_normed, + const at::Tensor& k_pe, at::Tensor& kv_cache, + const at::Tensor& slot_mapping) { + TORCH_CHECK(kv_c_normed.dim() == 2, + "kv_c_normed must be 2D [num_tokens, kv_lora_rank]"); + TORCH_CHECK(k_pe.dim() == 2 || k_pe.dim() == 3, "k_pe must be 2D or 3D"); + TORCH_CHECK(kv_cache.dim() == 3, + "kv_cache must be 3D [num_blocks, block_size, head_size]"); + TORCH_CHECK(kv_c_normed.stride(-1) == 1, + "kv_c_normed innermost dim must be contiguous"); + TORCH_CHECK(k_pe.stride(-1) == 1, "k_pe innermost dim must be contiguous"); + TORCH_CHECK(kv_cache.stride(-1) == 1, + "kv_cache innermost dim must be contiguous"); + + int64_t num_tokens = kv_c_normed.size(0); + int64_t kv_lora_rank = kv_c_normed.size(1); + int64_t qk_rope_head_dim = k_pe.size(-1); + int64_t head_size = kv_cache.size(-1); + TORCH_CHECK(head_size == kv_lora_rank + qk_rope_head_dim, + "kv_cache head_size must equal kv_lora_rank + qk_rope_head_dim"); + TORCH_CHECK(slot_mapping.size(0) == num_tokens, "slot_mapping size mismatch"); + + // Real physical per-token stride of the cache, read from the tensor's own + // strides rather than assumed to equal head_size (paged/pooled buffers are + // not guaranteed to have zero inter-row padding in general, even though in + // practice a freshly-allocated MLA cache is fully contiguous). + int64_t cache_stride = kv_cache.stride(1); + + const auto dtype = kv_cache.scalar_type(); + TORCH_CHECK(dtype == kv_c_normed.scalar_type() && dtype == k_pe.scalar_type(), + "concat_and_cache_mla_cpu: dtype mismatch"); + const auto index_dtype = slot_mapping.scalar_type(); + TORCH_CHECK(index_dtype == at::kLong || index_dtype == at::kInt, + "slot_mapping must be int32 or int64"); + + AT_DISPATCH_REDUCED_FLOATING_TYPES(dtype, "concat_and_cache_mla_cpu", [&] { + AT_DISPATCH_INDEX_TYPES(index_dtype, "concat_and_cache_mla_cpu_index", [&] { + concat_and_cache_mla_kernel_impl( + kv_c_normed.data_ptr(), k_pe.data_ptr(), + kv_cache.data_ptr(), slot_mapping.data_ptr(), + num_tokens, kv_lora_rank, qk_rope_head_dim, kv_c_normed.stride(0), + k_pe.stride(0), cache_stride); + }); + }); +} diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index 4bb0edf2b22a..c393000617c7 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -91,6 +91,39 @@ at::Tensor fp8_scaled_mm_cpu(at::Tensor& mat1, at::Tensor& mat2, const std::optional& bias, at::ScalarType out_dtype, bool is_vnni); +// Adapted from sglang: MLA CPU kernels (AMX-only) +void decode_attention_cpu(at::Tensor& query, at::Tensor& k_buffer, + at::Tensor& v_buffer, at::Tensor& output, + const std::optional& key, + const std::optional& value, + const std::optional& loc, + at::Tensor& attn_logits, at::Tensor& req_to_token, + at::Tensor& req_pool_indices, at::Tensor& seq_lens, + double sm_scale, double logit_cap, bool is_cross_attn, + int64_t sliding_window_size, + std::optional encoder_lens, + std::optional sinks); + +void extend_attention_cpu( + at::Tensor& q_extend, const std::optional& k_extend, + const std::optional& v_extend, at::Tensor& o_extend, + at::Tensor& k_buffer, at::Tensor& v_buffer, at::Tensor& req_to_token, + at::Tensor& req_pool_indices, at::Tensor& seq_lens, + at::Tensor& extend_seq_lens, at::Tensor& extend_start_loc, + int64_t max_len_extend, double sm_scale, double logit_cap, + bool is_cross_attn, int64_t sliding_window_size, + std::optional encoder_lens, std::optional sinks, + std::optional tree_mask); + +void bmm_cpu(at::Tensor& out, at::Tensor& mat1, at::Tensor& mat2, bool is_vnni, + const std::optional& scale); + +// vLLM-native: CPU cache-write op for MLA's single-latent-buffer KV cache +// (the CUDA-only `concat_and_cache_mla` has no CPU dispatch). +void concat_and_cache_mla_cpu(const at::Tensor& kv_c_normed, + const at::Tensor& k_pe, at::Tensor& kv_cache, + const at::Tensor& slot_mapping); + // Adapted from sglang: INT4 W4A8 kernels std::tuple convert_weight_packed_scale_zp( at::Tensor qweight, // awq: (*, K, N / 8) || gptq: (*, K / 8, N) , int32 @@ -521,6 +554,35 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "pad_slot_id, " "bool is_vnni) -> Tensor"); ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu); + + // Adapted from sglang: MLA CPU kernels (AMX-only, DeepSeek V2/V3/R1) + ops.def( + "decode_attention_cpu(Tensor query, Tensor k_buffer, Tensor v_buffer, " + "Tensor(a!) output, Tensor? key, Tensor? value, Tensor? loc, Tensor " + "attn_logits, Tensor req_to_token, Tensor req_pool_indices, Tensor " + "seq_lens, float sm_scale, float logit_cap, bool is_cross_attn, int " + "sliding_window_size, Tensor? encoder_lens, Tensor? sinks) -> ()"); + ops.impl("decode_attention_cpu", torch::kCPU, &decode_attention_cpu); + + ops.def( + "extend_attention_cpu(Tensor q_extend, Tensor? k_extend, Tensor? " + "v_extend, Tensor(a!) o_extend, Tensor k_buffer, Tensor v_buffer, " + "Tensor req_to_token, Tensor req_pool_indices, Tensor seq_lens, Tensor " + "extend_seq_lens, Tensor extend_start_loc, int max_len_extend, float " + "sm_scale, float logit_cap, bool is_cross_attn, int " + "sliding_window_size, Tensor? encoder_lens, Tensor? sinks, Tensor? " + "tree_mask=None) -> ()"); + ops.impl("extend_attention_cpu", torch::kCPU, &extend_attention_cpu); + + ops.def( + "bmm_cpu(Tensor(a!) out, Tensor mat1, Tensor mat2, bool is_vnni, " + "Tensor? scale) -> ()"); + ops.impl("bmm_cpu", torch::kCPU, &bmm_cpu); + + ops.def( + "concat_and_cache_mla_cpu(Tensor kv_c_normed, Tensor k_pe, " + "Tensor(a!) kv_cache, Tensor slot_mapping) -> ()"); + ops.impl("concat_and_cache_mla_cpu", torch::kCPU, &concat_and_cache_mla_cpu); #endif #if (defined(__AVX512BF16__) && defined(__AVX512F__) && \ diff --git a/tests/kernels/attention/test_amx_mla.py b/tests/kernels/attention/test_amx_mla.py new file mode 100644 index 000000000000..f5fa2fad16c2 --- /dev/null +++ b/tests/kernels/attention/test_amx_mla.py @@ -0,0 +1,471 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Correctness tests for the AMX-only CPU MLA backend: the vendored +decode/extend/bmm kernels, the KV cache write, and the ``AMXMLAImpl`` +backend built on top of them. +""" + +import pytest +import torch + +from vllm.platforms import current_platform + +if not current_platform.is_cpu(): + pytest.skip("skipping CPU-only tests", allow_module_level=True) + +if not torch.cpu._is_amx_tile_supported(): + pytest.skip("AMX MLA requires an AMX-capable host", allow_module_level=True) + +torch.cpu._init_amx() + +from vllm import _custom_ops as ops # noqa: E402 +from vllm.utils.torch_utils import set_random_seed # noqa: E402 +from vllm.v1.attention.backends.mla.amx_mla import ( # noqa: E402 + AMXMLAImpl, + _compute_num_kv_splits, + _expand_block_table, +) + +KV_LORA_RANK = 64 +QK_ROPE_HEAD_DIM = 32 +HEAD_SIZE = KV_LORA_RANK + QK_ROPE_HEAD_DIM # 96, kv-cache row width +# The extend/decode kernels transpose by 512-bit lanes: both the cache row +# width and the (kv_lora_rank-wide) value width must be 32-element aligned. +DTYPE = torch.bfloat16 +ATOL = 2e-2 +RTOL = 2e-2 +# Backend-level tests chain absorb -> attend -> de-absorb through bf16, +# accumulating more rounding error than a single kernel op. +_IMPL_ATOL = 1.5e-1 +_IMPL_RTOL = 1.5e-1 + + +def _flatten_cache(kv_cache: torch.Tensor) -> torch.Tensor: + """(num_blocks, block_size, head_size) -> (num_slots, head_size).""" + return kv_cache.view(-1, kv_cache.size(-1)) + + +def _random_paged_cache(num_slots: int, head_size: int) -> torch.Tensor: + return torch.randn(num_slots, head_size, dtype=DTYPE) + + +def _ref_latent_attn( + q: torch.Tensor, + keys: torch.Tensor, + scale: float, + kv_lora_rank: int, +) -> torch.Tensor: + """Plain-PyTorch causal-free latent-space MQA attention for one request. + + q: (num_heads, head_size); keys: (ctx_len, head_size). Values are the + first ``kv_lora_rank`` columns of the same latent rows (MLA's K/V + aliasing), matching what decode_attention_cpu/extend_attention_cpu + compute internally for the MLA-shaped case. + """ + values = keys[:, :kv_lora_rank] + scores = (q.float() @ keys.float().T) * scale + probs = torch.softmax(scores, dim=-1) + return (probs @ values.float()).to(q.dtype) + + +def test_bmm_cpu_matches_torch_bmm(): + set_random_seed(0) + # bmm_cpu requires the output's last dim (mat2's out-features) to be a + # multiple of 32 (tinygemm's tile width); the contraction dim is free. + n, b, p, l = 8, 32, 32, 32 # noqa: E741 + mat1 = torch.randn(n, b, p, dtype=DTYPE) + mat2 = torch.randn(n, l, p, dtype=DTYPE) # (N, OUT, IN) Linear convention + + ref = torch.bmm(mat1.float(), mat2.float().transpose(1, 2)).to(DTYPE) + + out = torch.empty(n, b, l, dtype=DTYPE) + ops.bmm_cpu(out, mat1, mat2, False, None) + torch.testing.assert_close(out, ref, atol=ATOL, rtol=RTOL) + + packed = torch.ops._C.convert_weight_packed(mat2.contiguous()) + out_vnni = torch.empty(n, b, l, dtype=DTYPE) + ops.bmm_cpu(out_vnni, mat1, packed, True, None) + torch.testing.assert_close(out_vnni, ref, atol=ATOL, rtol=RTOL) + + +def test_concat_and_cache_mla_round_trip(): + set_random_seed(0) + num_tokens = 37 + num_blocks, block_size = 8, 16 + kv_cache = torch.zeros(num_blocks, block_size, HEAD_SIZE, dtype=DTYPE) + + kv_c_normed = torch.randn(num_tokens, KV_LORA_RANK, dtype=DTYPE) + k_pe = torch.randn(num_tokens, QK_ROPE_HEAD_DIM, dtype=DTYPE) + slot_mapping = torch.randperm(num_blocks * block_size)[:num_tokens].to(torch.int64) + + ops.amx_mla_concat_and_cache(kv_c_normed, k_pe, kv_cache, slot_mapping) + + flat = _flatten_cache(kv_cache) + read_back = flat[slot_mapping] + expected = torch.cat([kv_c_normed, k_pe], dim=-1) + torch.testing.assert_close(read_back, expected) + + untouched = torch.ones(num_blocks * block_size, dtype=torch.bool) + untouched[slot_mapping] = False + assert (flat[untouched] == 0).all() + + +@pytest.mark.parametrize("seq_lens", [[5], [1, 300, 33], [513, 1, 129, 7]]) +def test_decode_attention_cpu_matches_reference(seq_lens): + set_random_seed(0) + num_seqs = len(seq_lens) + num_heads = 4 + block_size = 32 + max_len = max(seq_lens) + max_blocks = (max_len + block_size - 1) // block_size + num_blocks = num_seqs * max_blocks + + kv_cache = _random_paged_cache(num_blocks * block_size, HEAD_SIZE) + block_table = torch.randperm(num_blocks).view(num_seqs, max_blocks).to(torch.int64) + req_to_token = _expand_block_table(block_table, block_size) + req_pool_indices = torch.arange(num_seqs, dtype=torch.int64) + seq_lens_t = torch.tensor(seq_lens, dtype=torch.int64) + scale = HEAD_SIZE**-0.5 + + q = torch.randn(num_seqs, num_heads, HEAD_SIZE, dtype=DTYPE) + kv_cache_flat = kv_cache.view(-1, 1, HEAD_SIZE) + v_buffer = kv_cache_flat[..., :KV_LORA_RANK] + + o = torch.zeros(num_seqs, num_heads, KV_LORA_RANK, dtype=DTYPE) + attn_logits = torch.zeros( + num_seqs, num_heads, 4, KV_LORA_RANK + 1, dtype=torch.float32 + ) + ops.cpu_mla_decode( + q, + kv_cache_flat, + v_buffer, + o, + None, + None, + None, + attn_logits, + req_to_token, + req_pool_indices, + seq_lens_t, + scale, + 0.0, + False, + 0, + None, + None, + ) + + flat_cache = _flatten_cache(kv_cache) + for i, seq_len in enumerate(seq_lens): + token_ids = req_to_token[i, :seq_len] + keys = flat_cache[token_ids] + ref = _ref_latent_attn(q[i], keys, scale, KV_LORA_RANK) + torch.testing.assert_close(o[i], ref, atol=ATOL, rtol=RTOL) + + +@pytest.mark.parametrize( + "prefix_lens,extend_lens", + [ + ([0], [9]), # fresh prefill, no cached prefix + ([64], [17]), # continuation of a cached prefix + ([0, 40, 128], [23, 5, 31]), # mixed batch + ], +) +def test_extend_attention_cpu_matches_reference(prefix_lens, extend_lens): + set_random_seed(0) + num_seqs = len(prefix_lens) + num_heads = 4 + block_size = 32 + seq_lens = [p + e for p, e in zip(prefix_lens, extend_lens)] + max_len = max(seq_lens) + max_blocks = (max_len + block_size - 1) // block_size + num_blocks = num_seqs * max_blocks + + kv_cache = _random_paged_cache(num_blocks * block_size, HEAD_SIZE) + block_table = torch.randperm(num_blocks).view(num_seqs, max_blocks).to(torch.int64) + req_to_token = _expand_block_table(block_table, block_size) + req_pool_indices = torch.arange(num_seqs, dtype=torch.int64) + seq_lens_t = torch.tensor(seq_lens, dtype=torch.int64) + extend_seq_lens_t = torch.tensor(extend_lens, dtype=torch.int64) + extend_start_loc = torch.cumsum( + torch.tensor([0, *extend_lens[:-1]], dtype=torch.int64), dim=0 + ) + scale = HEAD_SIZE**-0.5 + total_new_tokens = sum(extend_lens) + + # The new tokens' latent K/V, already written into the cache (mirrors + # do_kv_cache_update running before forward_mha in real usage). + k_extend_flat = torch.randn(total_new_tokens, HEAD_SIZE, dtype=DTYPE) + flat_cache = _flatten_cache(kv_cache) + for i in range(num_seqs): + prefix_len = prefix_lens[i] + ext_len = extend_lens[i] + new_token_ids = req_to_token[i, prefix_len : prefix_len + ext_len] + start = extend_start_loc[i].item() + flat_cache[new_token_ids] = k_extend_flat[start : start + ext_len] + + q_extend = torch.randn(total_new_tokens, num_heads, HEAD_SIZE, dtype=DTYPE) + k_extend = k_extend_flat.unsqueeze(1) + v_extend = k_extend[..., :KV_LORA_RANK] + kv_cache_flat = kv_cache.view(-1, 1, HEAD_SIZE) + v_buffer = kv_cache_flat[..., :KV_LORA_RANK] + + o_extend = torch.empty(total_new_tokens, num_heads, KV_LORA_RANK, dtype=DTYPE) + ops.cpu_mla_extend( + q_extend, + k_extend, + v_extend, + o_extend, + kv_cache_flat, + v_buffer, + req_to_token, + req_pool_indices, + seq_lens_t, + extend_seq_lens_t, + extend_start_loc, + max(extend_lens), + scale, + 0.0, + False, + 0, + None, + None, + None, + ) + + for i in range(num_seqs): + prefix_len = prefix_lens[i] + ext_len = extend_lens[i] + start = extend_start_loc[i].item() + for t in range(ext_len): + visible = prefix_len + t + 1 + token_ids = req_to_token[i, :visible] + keys = flat_cache[token_ids] + ref = _ref_latent_attn(q_extend[start + t], keys, scale, KV_LORA_RANK) + torch.testing.assert_close(o_extend[start + t], ref, atol=ATOL, rtol=RTOL) + + +class _FakeLinear: + """Minimal duck-typed stand-in for kv_b_proj: only .weight and + .quant_method are read by get_and_maybe_dequant_weights().""" + + def __init__(self, weight: torch.Tensor): + self.weight = weight + self.quant_method = None + + +def _make_amx_mla_impl(num_heads, qk_nope_head_dim, v_head_dim, kv_lora_rank): + weight = torch.randn( + num_heads * (qk_nope_head_dim + v_head_dim), kv_lora_rank, dtype=DTYPE + ) + kv_b_proj = _FakeLinear(weight) + impl = AMXMLAImpl( + num_heads=num_heads, + head_size=kv_lora_rank + QK_ROPE_HEAD_DIM, + scale=(qk_nope_head_dim + QK_ROPE_HEAD_DIM) ** -0.5, + num_kv_heads=1, + alibi_slopes=None, + sliding_window=None, + kv_cache_dtype="auto", + logits_soft_cap=None, + attn_type="decoder", + kv_sharing_target_layer_name=None, + q_lora_rank=None, + kv_lora_rank=kv_lora_rank, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=QK_ROPE_HEAD_DIM, + qk_head_dim=qk_nope_head_dim + QK_ROPE_HEAD_DIM, + v_head_dim=v_head_dim, + kv_b_proj=kv_b_proj, + ) + impl.process_weights_after_loading(DTYPE) + return impl, weight + + +class _FakePrefillMetadata: + def __init__( + self, block_table, cpu_seq_lens, query_start_loc, req_to_token, req_pool_indices + ): + self.block_table = block_table + self.cpu_seq_lens = cpu_seq_lens + query_start_loc_i64 = query_start_loc.to(torch.int64) + extend_seq_lens = query_start_loc_i64[1:] - query_start_loc_i64[:-1] + self.extend_seq_lens = extend_seq_lens + self.extend_start_loc = query_start_loc_i64[:-1] + self.max_len_extend = int(extend_seq_lens.max().item()) + self.req_to_token = req_to_token + self.req_pool_indices = req_pool_indices + + +class _FakeDecodeMetadata: + def __init__(self, block_table, seq_lens, req_to_token, req_pool_indices): + self.block_table = block_table + self.seq_lens_i64 = seq_lens.to(torch.int64) + self.req_to_token = req_to_token + self.req_pool_indices = req_pool_indices + self.num_kv_splits = _compute_num_kv_splits( + int(seq_lens.max().item()), current_platform.num_compute_units() + ) + + +class _FakeAttnMetadata: + def __init__(self, decode=None, prefill=None, max_seq_len=0): + self.decode = decode + self.prefill = prefill + self.max_seq_len = max_seq_len + + +def test_amx_mla_impl_forward_mqa_matches_reference(default_vllm_config): + """forward_mqa receives already-absorbed Q (as MLAAttention.forward_impl + would produce via layer.W_UK_T) and must reproduce plain per-head MLA + decode attention through W_UK/W_UV.""" + set_random_seed(1) + num_heads, qk_nope_head_dim, v_head_dim, kv_lora_rank = 4, 32, 32, 64 + impl, kv_b_weight = _make_amx_mla_impl( + num_heads, qk_nope_head_dim, v_head_dim, kv_lora_rank + ) + w_uk, w_uv = kv_b_weight.T.view( + kv_lora_rank, num_heads, qk_nope_head_dim + v_head_dim + ).split([qk_nope_head_dim, v_head_dim], dim=-1) + + seq_lens = [1, 40, 5] + num_seqs = len(seq_lens) + block_size = 32 + max_blocks = (max(seq_lens) + block_size - 1) // block_size + num_blocks = num_seqs * max_blocks + head_size = kv_lora_rank + QK_ROPE_HEAD_DIM + + kv_cache = torch.randn(num_blocks, block_size, head_size, dtype=DTYPE) + block_table = torch.randperm(num_blocks).view(num_seqs, max_blocks).to(torch.int64) + req_to_token = _expand_block_table(block_table, block_size) + + q_nope = torch.randn(num_seqs, num_heads, qk_nope_head_dim, dtype=DTYPE) + q_pe = torch.randn(num_seqs, num_heads, QK_ROPE_HEAD_DIM, dtype=DTYPE) + # Pre-absorb, mirroring MLAAttention.forward_impl's own bmm against W_UK_T. + ql_nope = torch.einsum("bnp,lnp->bnl", q_nope.float(), w_uk.float()).to(DTYPE) + + attn_metadata = _FakeAttnMetadata( + decode=_FakeDecodeMetadata( + block_table=block_table, + seq_lens=torch.tensor(seq_lens, dtype=torch.int64), + req_to_token=req_to_token, + req_pool_indices=torch.arange(num_seqs, dtype=torch.int64), + ), + max_seq_len=max(seq_lens), + ) + o, lse = impl.forward_mqa((ql_nope, q_pe), kv_cache, attn_metadata, layer=None) + assert lse is None + # De-absorb, mirroring MLAAttention.forward_impl's own bmm against + # layer._v_up_proj, to get back into real per-head V space for + # comparison against a from-scratch reference. + out_real = torch.einsum("bnl,lnv->bnv", o.float(), w_uv.float()) + + flat_cache = _flatten_cache(kv_cache) + scale = impl.scale + for i, seq_len in enumerate(seq_lens): + token_ids = req_to_token[i, :seq_len] + keys_latent = flat_cache[token_ids] # (L, head_size) + k_nope_latent, k_pe = keys_latent.split( + [kv_lora_rank, QK_ROPE_HEAD_DIM], dim=-1 + ) + k_real = torch.einsum("lp,pnd->nld", k_nope_latent.float(), w_uk.float()) + k_pe_b = k_pe.float().unsqueeze(0).expand(num_heads, -1, -1) + k_full = torch.cat([k_real, k_pe_b], dim=-1) # (num_heads, L, qk_head_dim) + q_full = torch.cat( + [q_nope[i], q_pe[i]], dim=-1 + ).float() # (num_heads, qk_head_dim) + scores = torch.einsum("nd,nld->nl", q_full, k_full) * scale + probs = torch.softmax(scores, dim=-1) + v_real = torch.einsum("lp,pnv->nlv", k_nope_latent.float(), w_uv.float()) + ref = torch.einsum("nl,nlv->nv", probs, v_real) + torch.testing.assert_close(out_real[i], ref, atol=_IMPL_ATOL, rtol=_IMPL_RTOL) + + +def test_amx_mla_impl_forward_mha_matches_reference(default_vllm_config): + """forward_mha receives raw unabsorbed Q and must attend correctly for a + prefill batch with a mix of fresh and continued (cached-prefix) + sequences.""" + set_random_seed(2) + num_heads, qk_nope_head_dim, v_head_dim, kv_lora_rank = 4, 32, 32, 64 + impl, kv_b_weight = _make_amx_mla_impl( + num_heads, qk_nope_head_dim, v_head_dim, kv_lora_rank + ) + w_uk, w_uv = kv_b_weight.T.view( + kv_lora_rank, num_heads, qk_nope_head_dim + v_head_dim + ).split([qk_nope_head_dim, v_head_dim], dim=-1) + + prefix_lens = [0, 20] + extend_lens = [13, 7] + seq_lens = [p + e for p, e in zip(prefix_lens, extend_lens)] + num_seqs = len(prefix_lens) + block_size = 32 + max_blocks = (max(seq_lens) + block_size - 1) // block_size + num_blocks = num_seqs * max_blocks + head_size = kv_lora_rank + QK_ROPE_HEAD_DIM + + kv_cache = torch.randn(num_blocks, block_size, head_size, dtype=DTYPE) + block_table = torch.randperm(num_blocks).view(num_seqs, max_blocks).to(torch.int64) + req_to_token = _expand_block_table(block_table, block_size) + flat_cache = _flatten_cache(kv_cache) + + total_new_tokens = sum(extend_lens) + query_start_loc = torch.tensor( + [0, *torch.cumsum(torch.tensor(extend_lens), dim=0).tolist()], + dtype=torch.int64, + ) + + q_nope = torch.randn(total_new_tokens, num_heads, qk_nope_head_dim, dtype=DTYPE) + q_pe = torch.randn(total_new_tokens, num_heads, QK_ROPE_HEAD_DIM, dtype=DTYPE) + q = torch.cat([q_nope, q_pe], dim=-1) + + kv_c_normed = torch.randn(total_new_tokens, kv_lora_rank, dtype=DTYPE) + k_pe = torch.randn(total_new_tokens, 1, QK_ROPE_HEAD_DIM, dtype=DTYPE) + + # Write the new tokens' latent K/V into the cache ahead of time, mirroring + # do_kv_cache_update running before forward_mha in real usage. + for i in range(num_seqs): + prefix_len = prefix_lens[i] + ext_len = extend_lens[i] + new_token_ids = req_to_token[i, prefix_len : prefix_len + ext_len] + start = query_start_loc[i].item() + flat_cache[new_token_ids, :kv_lora_rank] = kv_c_normed[start : start + ext_len] + flat_cache[new_token_ids, kv_lora_rank:] = k_pe[start : start + ext_len, 0] + + attn_metadata = _FakeAttnMetadata( + prefill=_FakePrefillMetadata( + block_table=block_table, + cpu_seq_lens=torch.tensor(seq_lens, dtype=torch.int64), + query_start_loc=query_start_loc, + req_to_token=req_to_token, + req_pool_indices=torch.arange(num_seqs, dtype=torch.int64), + ) + ) + output = torch.empty(total_new_tokens, num_heads * v_head_dim, dtype=DTYPE) + impl.forward_mha( + q, kv_c_normed, k_pe, kv_cache, attn_metadata, k_scale=None, output=output + ) + out_view = output.view(total_new_tokens, num_heads, v_head_dim).float() + + scale = impl.scale + for i in range(num_seqs): + prefix_len = prefix_lens[i] + ext_len = extend_lens[i] + start = query_start_loc[i].item() + for t in range(ext_len): + visible = prefix_len + t + 1 + token_ids = req_to_token[i, :visible] + keys_latent = flat_cache[token_ids] + k_nope_latent, k_pe_ref = keys_latent.split( + [kv_lora_rank, QK_ROPE_HEAD_DIM], dim=-1 + ) + k_real = torch.einsum("lp,pnd->nld", k_nope_latent.float(), w_uk.float()) + k_pe_b = k_pe_ref.float().unsqueeze(0).expand(num_heads, -1, -1) + k_full = torch.cat([k_real, k_pe_b], dim=-1) + q_full = torch.cat([q_nope[start + t], q_pe[start + t]], dim=-1).float() + scores = torch.einsum("nd,nld->nl", q_full, k_full) * scale + probs = torch.softmax(scores, dim=-1) + v_real = torch.einsum("lp,pnv->nlv", k_nope_latent.float(), w_uv.float()) + ref = torch.einsum("nl,nlv->nv", probs, v_real) + torch.testing.assert_close( + out_view[start + t], ref, atol=_IMPL_ATOL, rtol=_IMPL_RTOL + ) diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index cd164250a2fc..950de764108e 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -172,9 +172,11 @@ def test_mla_post_load_preserves_runtime_weight_addresses(monkeypatch): layer.kv_b_proj.quant_method = None layer.is_aiter_triton_fp4_bmm_enabled = False layer.is_aiter_triton_fp8_bmm_enabled = False + layer.is_amx_bmm_enabled = False layer.dcp_q_replicate = False layer.quant_config = None layer.layer_name = "test" + layer.impl = SimpleNamespace(process_weights_after_loading=lambda act_dtype: None) monkeypatch.setattr( mla_attention_module, "set_default_quant_scales", lambda *_, **__: None diff --git a/tests/v1/attention/test_mla_prefill_selector.py b/tests/v1/attention/test_mla_prefill_selector.py index e6d9f939ea58..bb03aa496a86 100644 --- a/tests/v1/attention/test_mla_prefill_selector.py +++ b/tests/v1/attention/test_mla_prefill_selector.py @@ -63,6 +63,7 @@ def test_no_device_capability_returns_flash_attn(self): with patch("vllm.platforms.current_platform") as mock_platform: mock_platform.get_device_capability.return_value = None + mock_platform.is_cpu.return_value = False backend = get_mla_prefill_backend(vllm_config) assert backend.get_name() == "FLASH_ATTN" diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 5d63ffa0de37..550bf94c43e3 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -3937,6 +3937,109 @@ def cpu_attention_with_kv_cache( ) +def cpu_mla_decode( + query: torch.Tensor, + k_buffer: torch.Tensor, + v_buffer: torch.Tensor, + output: torch.Tensor, + key: torch.Tensor | None, + value: torch.Tensor | None, + loc: torch.Tensor | None, + attn_logits: torch.Tensor, + req_to_token: torch.Tensor, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + logit_cap: float, + is_cross_attn: bool, + sliding_window_size: int, + encoder_lens: torch.Tensor | None, + sinks: torch.Tensor | None, +) -> None: + torch.ops._C.decode_attention_cpu( + query, + k_buffer, + v_buffer, + output, + key, + value, + loc, + attn_logits, + req_to_token, + req_pool_indices, + seq_lens, + sm_scale, + logit_cap, + is_cross_attn, + sliding_window_size, + encoder_lens, + sinks, + ) + + +def cpu_mla_extend( + q_extend: torch.Tensor, + k_extend: torch.Tensor | None, + v_extend: torch.Tensor | None, + o_extend: torch.Tensor, + k_buffer: torch.Tensor, + v_buffer: torch.Tensor, + req_to_token: torch.Tensor, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + extend_seq_lens: torch.Tensor, + extend_start_loc: torch.Tensor, + max_len_extend: int, + sm_scale: float, + logit_cap: float, + is_cross_attn: bool, + sliding_window_size: int, + encoder_lens: torch.Tensor | None, + sinks: torch.Tensor | None, + tree_mask: torch.Tensor | None = None, +) -> None: + torch.ops._C.extend_attention_cpu( + q_extend, + k_extend, + v_extend, + o_extend, + k_buffer, + v_buffer, + req_to_token, + req_pool_indices, + seq_lens, + extend_seq_lens, + extend_start_loc, + max_len_extend, + sm_scale, + logit_cap, + is_cross_attn, + sliding_window_size, + encoder_lens, + sinks, + tree_mask, + ) + + +def bmm_cpu( + out: torch.Tensor, + mat1: torch.Tensor, + mat2: torch.Tensor, + is_vnni: bool, + scale: torch.Tensor | None = None, +) -> None: + torch.ops._C.bmm_cpu(out, mat1, mat2, is_vnni, scale) + + +def amx_mla_concat_and_cache( + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + kv_cache: torch.Tensor, + slot_mapping: torch.Tensor, +) -> None: + torch.ops._C.concat_and_cache_mla_cpu(kv_c_normed, k_pe, kv_cache, slot_mapping) + + def cpu_gemm_wna16( input: torch.Tensor, q_weight: torch.Tensor, diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 27ccd8ca0ccb..37dc2c95ff49 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -540,6 +540,11 @@ def __init__( **extra_impl_args, ) self.q_pad_num_heads = getattr(self.impl, "q_pad_num_heads", None) + self.is_amx_bmm_enabled = getattr(self.impl, "uses_amx_bmm", False) + # AMX reads kv_b_proj's weight directly and never calls it live; the + # reference CPU MLA backend calls it but isn't perf-critical. Skip + # the packed-kernel dispatch either way. + kv_b_proj._cpu_skip_gemm_dispatch = True self.use_direct_call = not current_platform.opaque_attention_op() vllm_config = get_current_vllm_config() @@ -908,6 +913,20 @@ def forward_impl( group_size=128, transpose_bm=True, ) + elif self.is_amx_bmm_enabled: + # bmm_cpu computes out[n] = mat1[n] @ mat2[n]^T against + # AMXMLAImpl's own (N, L, P) packed W_UK -- same as prefill. + N, B, P = mqa_q_nope.shape + L = self.kv_lora_rank + mqa_ql_nope = mqa_q_nope.new_empty((N, B, L)) + ops.bmm_cpu( + mqa_ql_nope, + mqa_q_nope, + self.impl._w_uk_packed, # type: ignore[attr-defined] + True, + None, + ) + mqa_ql_nope = mqa_ql_nope.transpose(0, 1) else: # Pads the head_dim if necessary (for the underlying kernel) N, B, P = mqa_q_nope.shape @@ -1026,6 +1045,18 @@ def forward_impl( return output_padded def process_weights_after_loading(self, act_dtype: torch.dtype): + # Let per-backend impls do their own weight packing first (no-op + # unless overridden), mirroring Attention.process_weights_after_loading. + self.impl.process_weights_after_loading(act_dtype) + + if self.is_amx_bmm_enabled: + # AMXMLAImpl already packed its own W_UK/W_UV above, for both + # prefill and decode. Release the now-unused raw weight. + self.kv_b_proj.weight = torch.nn.Parameter( + torch.empty(0), requires_grad=False + ) + return + # we currently do not have quantized bmm's which are needed for # `W_UV` and `W_UK_T`, we just store fp16/bf16 copies and perform # the bmm's in 16-bit, the extra memory overhead of this is fairly low @@ -1192,6 +1223,16 @@ def _v_up_proj(self, x: torch.Tensor, out: torch.Tensor): x = rocm_aiter_ops.triton_fp8_bmm( x, self.W_V, self.W_V_scale, group_size=128, transpose_bm=True, YQ=out ) + elif self.is_amx_bmm_enabled: + # bmm_cpu computes out[n] = mat1[n] @ mat2[n]^T against + # AMXMLAImpl's own (N, V, L) packed W_UV -- same as prefill. + ops.bmm_cpu( + out.transpose(0, 1), + x, + self.impl._w_uv_packed, # type: ignore[attr-defined] + True, + None, + ) else: # Multiply + Transpose (N, B, L) x (N, L, V)->(N, B, V)->(B, N, V) torch.bmm(x, self.W_UV, out=out.transpose(0, 1)) diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index 050e7a0acd8f..1e31a64d685f 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -197,6 +197,12 @@ def create_weights( def process_weights_after_loading(self, layer: torch.nn.Module) -> None: if current_platform.is_cpu(): + # MLA's kv_b_proj (see `_cpu_skip_gemm_dispatch`): not + # perf-critical, so skip packing and use a plain fallback. + if getattr(layer, "_cpu_skip_gemm_dispatch", False): + layer.cpu_linear = torch.nn.functional.linear + return + from vllm.model_executor.layers.utils import dispatch_cpu_unquantized_gemm dispatch_cpu_unquantized_gemm(layer, remove_weight=True) diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index dd18363e0d26..02ca99ff33a9 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -89,13 +89,30 @@ def get_attn_backend_cls( if attn_selector_config.use_sparse: raise NotImplementedError("Sparse Attention is not supported on CPU.") if attn_selector_config.use_mla: + amx_available = ( + cls.get_cpu_architecture() == CpuArchEnum.X86 + and torch.cpu._is_amx_tile_supported() + ) + if amx_available and selected_backend != AttentionBackendEnum.CPU_MLA: + # Prefer AMX when available, unless CPU_MLA was explicitly requested. + if ( + selected_backend + and selected_backend != AttentionBackendEnum.AMX_MLA + ): + logger.info("Cannot use %s backend on CPU.", selected_backend) + logger.info_once("Using %s backend.", AttentionBackendEnum.AMX_MLA.name) + return AttentionBackendEnum.AMX_MLA.get_path() # Reference MLA implementation on CPU. Performance is not the # goal here; the backend simply wires the CPU decode kernel # (`mla_decode_kvcache`) and an SDPA-based prefill together with # the shared MLA scaffolding so that DeepSeek-style models can # execute on CPU. - if selected_backend and selected_backend != AttentionBackendEnum.CPU_MLA: + if selected_backend and selected_backend not in ( + AttentionBackendEnum.CPU_MLA, + AttentionBackendEnum.AMX_MLA, + ): logger.info("Cannot use %s backend on CPU.", selected_backend) + logger.info_once("Using %s backend.", AttentionBackendEnum.CPU_MLA.name) return AttentionBackendEnum.CPU_MLA.get_path() if selected_backend and selected_backend != AttentionBackendEnum.CPU_ATTN: logger.info("Cannot use %s backend on CPU.", selected_backend) @@ -134,11 +151,20 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None: # The CPU MLA decode kernel only compiles with block_size=16 today # (see csrc/cpu/mla_decode.cpp). If the model uses MLA we override # the default block size regardless of user preference to avoid a - # runtime kernel dispatch failure. + # runtime kernel dispatch failure. AMX MLA has no such constraint + # (same AMX-available condition as get_attn_backend_cls), so it's + # excluded from this override. cpu_mla_enabled = model_config is not None and getattr( model_config, "use_mla", False ) - if cpu_mla_enabled: + amx_mla_enabled = ( + cpu_mla_enabled + and cls.get_cpu_architecture() == CpuArchEnum.X86 + and torch.cpu._is_amx_tile_supported() + and vllm_config.attention_config.backend != AttentionBackendEnum.CPU_MLA + ) + reference_cpu_mla_enabled = cpu_mla_enabled and not amx_mla_enabled + if reference_cpu_mla_enabled: if cache_config.user_specified_block_size and cache_config.block_size != 16: logger.warning( "CPU MLA backend requires block_size=16, overriding " @@ -149,7 +175,7 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None: elif not cache_config.user_specified_block_size: cache_config.block_size = 128 - if not cpu_mla_enabled and cache_config.block_size % 32 != 0: + if not reference_cpu_mla_enabled and cache_config.block_size % 32 != 0: logger.warning( "CPU backend prefers block_size is multiples of 32, " "otherwise the performance is not optimized." @@ -342,7 +368,7 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None: vllm_config.parallel_config.tensor_parallel_size ) - if model_config is not None and model_config.use_mla: + if model_config is not None and model_config.use_mla and not amx_mla_enabled: logger.info_once( "MLA is enabled on a non-GPU platform; forcing chunked " "prefill and prefix caching to be disabled." diff --git a/vllm/v1/attention/backends/mla/amx_mla.py b/vllm/v1/attention/backends/mla/amx_mla.py new file mode 100644 index 000000000000..5f9ac7f17adf --- /dev/null +++ b/vllm/v1/attention/backends/mla/amx_mla.py @@ -0,0 +1,405 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""AMX-only, high-performance MLA backend for DeepSeek V2/V3/R1 on CPU. + +Built on the AMX decode/extend/bmm kernels vendored under +``csrc/cpu/sgl-kernels/``, plugged into vLLM's ``MLACommonBackend``/ +``MLACommonImpl`` abstraction the same way every other concrete MLA backend +(``TritonMLAImpl``, etc.) does. + +This is a separate backend from the reference ``CPUMLABackend`` +(``vllm/v1/attention/backends/mla/cpu_mla.py``): that one targets every CPU +(any dtype, block_size=16, ``mla_decode_kvcache`` decode kernel + inherited +SDPA-style prefill) as a functional/CI reference, not performance. This +backend instead requires AMX (bf16 only, block_size a multiple of 32) and is +selected by the platform layer in preference to the reference backend +whenever the host supports it -- see ``CpuPlatform.get_attn_backend_cls``. + +Two points where this backend differs structurally from the GPU backends, +both explained in the CPU MLA design plan: + +- ``forward_mha`` is fully overridden (not inherited from + ``MLACommonBaseImpl``): the GPU "compute-friendly" prefill path depends on + a pluggable ``MLAPrefillBackend`` and CUDA-only chunked-context gather ops, + neither of which exist on CPU. Instead, this attends directly in + latent-MQA space via the ``extend_attention_cpu`` kernel, which handles + cached-prefix continuation and fresh prefill in one causal pass (the KV + cache already contains the new tokens by the time this runs, since + ``do_kv_cache_update`` executes first). +- Weight absorption for the prefill path is done with this impl's own + VNNI-packed copies of ``W_UK``/``W_UV`` (computed in + ``process_weights_after_loading``, using ``bmm_cpu``), rather than reading + the generic ``layer.W_UK_T``/``layer.W_UV``/``layer._v_up_proj`` the way + GPU backends do, since ``forward_mha``'s abstract signature has no + ``layer`` parameter (only ``forward_mqa`` does) and that signature is left + untouched. The decode path needs no such packing: ``MLAAttention.forward_impl`` + already absorbs/de-absorbs Q around the ``forward_mqa`` call using the + generic (unpacked) ``layer.W_UK_T``/``layer._v_up_proj``, so + ``forward_mqa`` here only has to invoke the decode kernel. +""" + +from __future__ import annotations + +from typing import ClassVar + +import torch + +from vllm import _custom_ops as ops +from vllm.model_executor.layers.attention.mla_attention import ( + MLACommonBackend, + MLACommonImpl, + MLACommonMetadata, + MLACommonMetadataBuilder, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + get_and_maybe_dequant_weights, +) +from vllm.platforms import current_platform +from vllm.v1.attention.backend import AttentionLayer, AttentionType, MultipleOf + +_MIN_WORK_PER_SPLIT = 512 +_SPLIT_OCCUPANCY_MULTIPLIER = 2 + + +def _compute_num_kv_splits(max_seq_len: int, num_threads: int) -> int: + """Mirrors TritonMLAImpl's _compute_num_kv_splits, using the CPU thread + count in place of SM count.""" + ideal_splits = 1 + while ideal_splits < max(1, max_seq_len // _MIN_WORK_PER_SPLIT): + ideal_splits *= 2 + max_splits = num_threads * _SPLIT_OCCUPANCY_MULTIPLIER + return min(ideal_splits, max_splits) + + +def _expand_block_table(block_table: torch.Tensor, block_size: int) -> torch.Tensor: + """Adapter: vLLM's block-table paging -> the flat per-(request, position) + physical row index the decode/extend kernels expect. Called once per step + from ``AMXMLAMetadataBuilder.build``, not per layer. + """ + offsets = torch.arange( + block_size, device=block_table.device, dtype=block_table.dtype + ) + flat = block_table.unsqueeze(-1) * block_size + offsets + return flat.reshape(block_table.size(0), -1).contiguous() + + +class AMXMLABackend(MLACommonBackend): + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16] + + @classmethod + def get_supported_head_sizes(cls) -> list[int]: + # kv_lora_rank(512) + qk_rope_head_dim(64), DeepSeek V2/V3/R1/V3.2. + return [576] + + @staticmethod + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: + return [MultipleOf(32)] + + @classmethod + def supports_block_size(cls, block_size: int | None) -> bool: + if block_size is None: + return True + return block_size % 32 == 0 + + @staticmethod + def get_name() -> str: + return "AMX_MLA" + + @staticmethod + def get_impl_cls() -> type[AMXMLAImpl]: + return AMXMLAImpl + + @staticmethod + def get_builder_cls() -> type[AMXMLAMetadataBuilder]: + return AMXMLAMetadataBuilder + + +class AMXMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]): + def build(self, common_prefix_len, common_attn_metadata, fast_build: bool = False): + attn_metadata = super().build( + common_prefix_len, common_attn_metadata, fast_build + ) + block_size = self.kv_cache_spec.block_size + if attn_metadata.decode is not None: + # Built once per step and reused by every layer, instead of + # recomputing the same tensor arithmetic in every forward_mqa call. + decode = attn_metadata.decode + decode.req_to_token = _expand_block_table( # type: ignore[attr-defined] + decode.block_table, block_size + ) + decode.seq_lens_i64 = decode.seq_lens.to(torch.int64) # type: ignore[attr-defined] + decode.req_pool_indices = torch.arange( # type: ignore[attr-defined] + decode.block_table.size(0), + dtype=torch.int64, + device=decode.block_table.device, + ) + decode.num_kv_splits = _compute_num_kv_splits( # type: ignore[attr-defined] + attn_metadata.max_seq_len, current_platform.num_compute_units() + ) + if attn_metadata.prefill is not None: + # cpu_seq_lens: total per-request context length (prefix + new + # tokens), the one thing the extend kernel needs that isn't + # already on the shared MLACommonPrefillMetadata. + prefill = attn_metadata.prefill + num_decodes = attn_metadata.num_decodes + num_prefills = attn_metadata.num_prefills + prefill.cpu_seq_lens = common_attn_metadata.seq_lens[ # type: ignore[attr-defined] + num_decodes : num_decodes + num_prefills + ].to(torch.int64) + prefill.req_to_token = _expand_block_table( # type: ignore[attr-defined] + prefill.block_table, block_size + ).to(torch.int64) + prefill.req_pool_indices = torch.arange( # type: ignore[attr-defined] + prefill.block_table.size(0), + dtype=torch.int64, + device=prefill.block_table.device, + ) + query_start_loc_i64 = prefill.query_start_loc.to(torch.int64) + extend_seq_lens = query_start_loc_i64[1:] - query_start_loc_i64[:-1] + prefill.extend_seq_lens = extend_seq_lens # type: ignore[attr-defined] + prefill.extend_start_loc = query_start_loc_i64[:-1] # type: ignore[attr-defined] + prefill.max_len_extend = int(extend_seq_lens.max().item()) # type: ignore[attr-defined] + return attn_metadata + + +class AMXMLAImpl(MLACommonImpl[MLACommonMetadata]): + # Tells MLAAttention (mla_attention.py) to use this impl's own packed + # W_UK/W_UV + bmm_cpu for the shared decode absorb/de-absorb, mirroring + # is_aiter_triton_fp8_bmm_enabled/is_aiter_triton_fp4_bmm_enabled. + uses_amx_bmm = True + + def __init__( + self, + num_heads: int, + head_size: int, + scale: float, + num_kv_heads: int, + alibi_slopes: list[float] | None, + sliding_window: int | None, + kv_cache_dtype: str, + logits_soft_cap: float | None, + attn_type: str, + kv_sharing_target_layer_name: str | None, + **mla_args, + ) -> None: + super().__init__( + num_heads, + head_size, + scale, + num_kv_heads, + alibi_slopes, + sliding_window, + kv_cache_dtype, + logits_soft_cap, + attn_type, + kv_sharing_target_layer_name, + **mla_args, + ) + unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap] + if any(unsupported_features): + raise NotImplementedError( + "AMXMLAImpl does not support one of the following: " + "alibi_slopes, sliding_window, logits_soft_cap" + ) + if attn_type != AttentionType.DECODER: + raise NotImplementedError( + "Encoder self-attention and encoder/decoder cross-attention " + "are not implemented for AMXMLAImpl" + ) + self._w_uk_packed: torch.Tensor | None = None + self._w_uv_packed: torch.Tensor | None = None + + def process_weights_after_loading(self, act_dtype: torch.dtype) -> None: + # forward_mha (unlike forward_mqa) receives no `layer` argument, so it + # cannot read the generic layer.W_UK_T/W_UV/._v_up_proj that + # MLAAttention.process_weights_after_loading already computes. This + # impl derives its own copies from self.kv_b_proj (already available + # via the constructor) and VNNI-packs them once, ahead of time, for + # use with bmm_cpu in forward_mha. + kv_b_proj_weight = get_and_maybe_dequant_weights( + self.kv_b_proj, out_dtype=act_dtype + ).T + kv_b_proj_weight = kv_b_proj_weight.view( + self.kv_lora_rank, self.num_heads, self.qk_nope_head_dim + self.v_head_dim + ) + w_uk, w_uv = kv_b_proj_weight.split( + [self.qk_nope_head_dim, self.v_head_dim], dim=-1 + ) + # bmm_cpu computes out[b] = mat1[b] @ mat2[b]^T (Linear-weight + # convention: mat2[b] is (OUT, IN)). + # absorb: ql_nope(N,B,L) = q_nope(N,B,P) @ w_uk_for_bmm(N,L,P)^T + # de-absorb: out(N,B,V) = attn_out(N,B,L) @ w_uv_for_bmm(N,V,L)^T + w_uk_for_bmm = w_uk.permute(1, 0, 2).contiguous() # (N, L, P) + w_uv_for_bmm = w_uv.permute(1, 2, 0).contiguous() # (N, V, L) + self._w_uk_packed = torch.ops._C.convert_weight_packed(w_uk_for_bmm) + self._w_uv_packed = torch.ops._C.convert_weight_packed(w_uv_for_bmm) + + def do_kv_cache_update( + self, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + kv_cache: torch.Tensor, + slot_mapping: torch.Tensor, + kv_cache_dtype: str, + k_scale: torch.Tensor, + ) -> None: + # Overrides the default (CUDA-only concat_and_cache_mla). + if kv_cache.numel() == 0: + return + assert kv_cache_dtype == "auto", ( + "AMXMLAImpl only supports an unquantized (bf16) KV cache; fp8 " + "KV cache for CPU MLA is not yet implemented." + ) + ops.amx_mla_concat_and_cache( + kv_c_normed, k_pe.squeeze(1), kv_cache, slot_mapping.flatten() + ) + + def forward_mqa( + self, + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + kv_c_and_k_pe_cache: torch.Tensor, + attn_metadata: MLACommonMetadata, + layer: AttentionLayer, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + # q arrives already absorbed (tensor or (ql_nope, q_pe) tuple) -- + # MLAAttention.forward_impl does the Q-absorption bmm generically + # before calling forward_mqa, using layer.W_UK_T. De-absorption is + # likewise done generically afterwards via layer._v_up_proj. Nothing + # left for this impl to do beyond invoking the decode kernel. + if isinstance(q, tuple): + q = torch.cat(q, dim=-1) + assert attn_metadata.decode is not None + + num_tokens = q.shape[0] + kv_cache_flat = kv_c_and_k_pe_cache.view(-1, 1, self.head_size) + v_buffer = kv_cache_flat[..., : self.kv_lora_rank] + + seq_lens = attn_metadata.decode.seq_lens_i64 # type: ignore[attr-defined] + req_to_token = attn_metadata.decode.req_to_token # type: ignore[attr-defined] + req_pool_indices = attn_metadata.decode.req_pool_indices # type: ignore[attr-defined] + + num_kv_splits = attn_metadata.decode.num_kv_splits # type: ignore[attr-defined] + o = torch.zeros( + num_tokens, + self.num_heads, + self.kv_lora_rank, + dtype=q.dtype, + device=q.device, + ) + attn_logits = torch.zeros( + num_tokens, + self.num_heads, + num_kv_splits, + self.kv_lora_rank + 1, + dtype=torch.float32, + device=q.device, + ) + + ops.cpu_mla_decode( + q, + kv_cache_flat, + v_buffer, + o, + None, + None, + None, # loc: only meaningful when key/value are given; we + # already wrote the new K/V via do_kv_cache_update before this. + attn_logits, + req_to_token, + req_pool_indices, + seq_lens, + self.scale, + 0.0, + False, + 0, + None, + None, + ) + return o, None + + def forward_mha( + self, + q: torch.Tensor, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + kv_c_and_k_pe_cache: torch.Tensor, + attn_metadata: MLACommonMetadata, + k_scale: torch.Tensor, + output: torch.Tensor, + output_scale: torch.Tensor | None = None, + ) -> None: + assert output_scale is None, ( + "AMXMLAImpl.forward_mha does not support fused output quantization" + ) + prefill = attn_metadata.prefill + assert prefill is not None + assert self._w_uk_packed is not None and self._w_uv_packed is not None + + num_tokens = q.shape[0] + num_heads = self.num_heads + + q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + + # Absorb: (N, B, P) x (N, L, P)^T -> (N, B, L) + # bmm_cpu tolerates non-contiguous mat1/out as long as their last dim + # is contiguous, so q_nope's transpose and ql_nope's transposed view + # can be passed directly -- no copy needed for either. + q_nope_t = q_nope.transpose(0, 1) + ql_nope = torch.empty( + num_tokens, num_heads, self.kv_lora_rank, dtype=q.dtype, device=q.device + ) + ops.bmm_cpu(ql_nope.transpose(0, 1), q_nope_t, self._w_uk_packed, True, None) + mqa_q = torch.cat([ql_nope, q_pe], dim=-1) + + kv_cache_flat = kv_c_and_k_pe_cache.view(-1, 1, self.head_size) + v_buffer = kv_cache_flat[..., : self.kv_lora_rank] + + req_to_token = prefill.req_to_token # type: ignore[attr-defined] + req_pool_indices = prefill.req_pool_indices # type: ignore[attr-defined] + + seq_lens = prefill.cpu_seq_lens # type: ignore[attr-defined] + extend_seq_lens = prefill.extend_seq_lens # type: ignore[attr-defined] + extend_start_loc = prefill.extend_start_loc # type: ignore[attr-defined] + max_len_extend = prefill.max_len_extend # type: ignore[attr-defined] + + # k_extend/v_extend: the new tokens' own latent K/V, aliased in one + # 576-wide buffer (v_extend is a view of k_extend's first + # kv_lora_rank columns) so the kernel can gather both from one read, + # mirroring the cache's own layout. + k_extend = torch.cat( + [kv_c_normed, k_pe.reshape(num_tokens, self.qk_rope_head_dim)], dim=-1 + ).unsqueeze(1) + v_extend = k_extend[..., : self.kv_lora_rank] + + attn_out = torch.empty( + num_tokens, num_heads, self.kv_lora_rank, dtype=q.dtype, device=q.device + ) + ops.cpu_mla_extend( + mqa_q, + k_extend, + v_extend, + attn_out, + kv_cache_flat, + v_buffer, + req_to_token, + req_pool_indices, + seq_lens, + extend_seq_lens, + extend_start_loc, + max_len_extend, + self.scale, + 0.0, + False, + 0, + None, + None, + None, + ) + + # De-absorb directly into `output` (view, no extra copy on either side): + # (N, B, L) x (N, V, L)^T -> (N, B, V) + attn_out_t = attn_out.transpose(0, 1) + output_view = output.view(num_tokens, num_heads, self.v_head_dim).transpose( + 0, 1 + ) + ops.bmm_cpu(output_view, attn_out_t, self._w_uv_packed, True, None) diff --git a/vllm/v1/attention/backends/mla/prefill/cpu_native.py b/vllm/v1/attention/backends/mla/prefill/cpu_native.py new file mode 100644 index 000000000000..52ba8b7e2422 --- /dev/null +++ b/vllm/v1/attention/backends/mla/prefill/cpu_native.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inert MLA prefill backend for the CPU platform. + +CPU's MLA attention impl (``CPUMLAImpl``) fully overrides ``forward_mha`` with a +kernel that attends directly against the paged latent KV cache (covering both +fresh prefill and cached-prefix continuation in one pass), so it never calls +into a pluggable ``MLAPrefillBackend``'s ``run_prefill_new_tokens``/ +``run_prefill_context_chunk``. This class exists only to satisfy +``MLACommonMetadataBuilder``'s structural requirement that +``MLAAttention.prefill_backend`` be a real, cloneable object. +""" + +from typing import TYPE_CHECKING + +import torch + +from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend + +if TYPE_CHECKING: + from vllm.model_executor.layers.attention.mla_attention import ( + MLACommonPrefillMetadata, + ) + + +class CPUNativeMLAPrefillBackend(MLAPrefillBackend): + """Placeholder MLA prefill backend for CPU; never actually invoked.""" + + @staticmethod + def get_name() -> str: + return "CPU_NATIVE" + + def run_prefill_new_tokens( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + return_softmax_lse: bool, + out: torch.Tensor | None = None, + output_scale: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + raise AssertionError( + "CPUNativeMLAPrefillBackend.run_prefill_new_tokens is unreachable: " + "CPUMLAImpl.forward_mha fully overrides the dense-MHA prefill path " + "and never calls the pluggable prefill backend." + ) + + def run_prefill_context_chunk( + self, + chunk: "MLACommonPrefillMetadata.ContextChunk", + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + raise AssertionError( + "CPUNativeMLAPrefillBackend.run_prefill_context_chunk is " + "unreachable: CPUMLAImpl.forward_mha fully overrides the " + "dense-MHA prefill path and never calls the pluggable prefill " + "backend." + ) diff --git a/vllm/v1/attention/backends/mla/prefill/registry.py b/vllm/v1/attention/backends/mla/prefill/registry.py index 0d818a084ba6..b521f4ff1ec2 100644 --- a/vllm/v1/attention/backends/mla/prefill/registry.py +++ b/vllm/v1/attention/backends/mla/prefill/registry.py @@ -52,6 +52,9 @@ class MLAPrefillBackendEnum(Enum, metaclass=_MLAPrefillBackendEnumMeta): "vllm.v1.attention.backends.mla.prefill.aiter_flash_attn." "AiterFlashAttnPrefillBackend" ) + CPU_NATIVE = ( + "vllm.v1.attention.backends.mla.prefill.cpu_native.CPUNativeMLAPrefillBackend" + ) # Placeholder for third-party/custom backends - must be registered before use # set to None to avoid alias with other backend, whose value is an empty string CUSTOM = None diff --git a/vllm/v1/attention/backends/mla/prefill/selector.py b/vllm/v1/attention/backends/mla/prefill/selector.py index c313c9017723..5cdfe3c6cae1 100644 --- a/vllm/v1/attention/backends/mla/prefill/selector.py +++ b/vllm/v1/attention/backends/mla/prefill/selector.py @@ -109,6 +109,8 @@ def get_mla_prefill_backend( device_capability = current_platform.get_device_capability() if device_capability is None: + if current_platform.is_cpu(): + return MLAPrefillBackendEnum.CPU_NATIVE.get_class() logger.info_once( "Device capability not available, using FlashAttention MLA prefill backend." ) diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index 646227cc4f9c..6fc9dd1e1fde 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -123,6 +123,7 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): ) CPU_ATTN = "vllm.v1.attention.backends.cpu_attn.CPUAttentionBackend" CPU_MLA = "vllm.v1.attention.backends.mla.cpu_mla.CPUMLABackend" + AMX_MLA = "vllm.v1.attention.backends.mla.amx_mla.AMXMLABackend" TURBOQUANT = "vllm.v1.attention.backends.turboquant_attn.TurboQuantAttentionBackend" # Placeholder for third-party/custom backends - must be registered before use # set to None to avoid alias with other backend, whose value is an empty string From f936a267f9a774b7ac09f987226d8e5a591139da Mon Sep 17 00:00:00 2001 From: Willow Lopez <100782273+Oxygen56@users.noreply.github.com> Date: Wed, 19 Aug 2026 10:45:41 +0800 Subject: [PATCH 123/839] [Bugfix][XPU] Fix Mamba state pointer overflow (#48109) Signed-off-by: Oxygen56 --- tests/v1/worker/test_mamba_utils.py | 89 +++++++++++++++++++++++++++++ vllm/v1/worker/mamba_utils.py | 11 +++- 2 files changed, 98 insertions(+), 2 deletions(-) diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index 0534795f9e84..c31116042faa 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -18,6 +18,7 @@ from vllm.v1.worker.mamba_utils import ( MambaCopyBuffers, MambaSpecDecodeGPUContext, + _reinterpret_u64_as_i64, batch_memcpy, collect_mamba_copy_meta, do_mamba_copy_block, @@ -189,6 +190,94 @@ def copy_to_gpu(self, n: int | None = None) -> torch.Tensor: return self.gpu[:n].copy_(self.cpu[:n], non_blocking=True) +class _FakeDataPtrTensor: + """Tensor wrapper that exposes a controlled data_ptr for metadata tests.""" + + def __init__(self, tensor: torch.Tensor, data_ptr: int): + self._tensor = tensor + self._data_ptr = data_ptr + self.shape = tensor.shape + + def data_ptr(self) -> int: + return self._data_ptr + + def dim(self) -> int: + return self._tensor.dim() + + def stride(self, *args): + return self._tensor.stride(*args) + + def numel(self) -> int: + return self._tensor.numel() + + def element_size(self) -> int: + return self._tensor.element_size() + + def size(self, *args): + return self._tensor.size(*args) + + def __getitem__(self, item): + return self._tensor[item] + + +def test_reinterpret_u64_as_i64_preserves_pointer_bits(): + ptrs = [ + 0, + 1, + (1 << 63) - 1, + 1 << 63, + (1 << 63) + 1234, + (1 << 64) - 1, + ] + ptr_tensor = torch.zeros(len(ptrs), dtype=torch.int64) + + for idx, ptr in enumerate(ptrs): + ptr_tensor[idx] = _reinterpret_u64_as_i64(ptr) + + assert ptr_tensor.numpy().view(np.uint64).tolist() == ptrs + + +def test_gpu_context_reinterprets_high_data_ptrs_for_int64_metadata(): + cfg = _TestConfig(num_layers=1) + device = torch.device("cpu") + kv_cache_config = _make_kv_cache_config(cfg, ["layer_0"]) + gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device) + conv_ptr = 1 << 63 + temporal_ptr = (1 << 64) - 8 + block_table_ptr = (1 << 63) + 42 + + conv_state = _FakeDataPtrTensor( + torch.empty( + cfg.num_blocks, + cfg.conv_width, + cfg.conv_inner_dim, + dtype=cfg.dtype, + ), + conv_ptr, + ) + temporal_state = _FakeDataPtrTensor( + torch.empty(cfg.num_blocks, cfg.temporal_state_dim, dtype=cfg.dtype), + temporal_ptr, + ) + block_table = _FakeDataPtrTensor( + torch.empty(1, 4, dtype=torch.int32), + block_table_ptr, + ) + forward_context = {"layer_0": _make_mock_attention(conv_state, temporal_state)} + + gpu_ctx.initialize_from_forward_context( + kv_cache_config, forward_context, _COPY_FUNCS, [block_table] + ) + + assert gpu_ctx.state_base_addrs.tolist() == [ + _reinterpret_u64_as_i64(conv_ptr), + _reinterpret_u64_as_i64(temporal_ptr), + ] + assert gpu_ctx.block_table_ptrs.tolist() == [ + _reinterpret_u64_as_i64(block_table_ptr) + ] + + def _make_postprocess_scheduler_output( req_ids: list[str], num_scheduled_tokens: dict[str, int], diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index a0b89303dd86..5dfa6f474d3b 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -114,6 +114,11 @@ def _memcpy_u64_tiled( tl.store(dst_u8 + i + offsets, data, mask=mask) +def _reinterpret_u64_as_i64(value: int) -> int: + """Preserve a uint64 pointer bit pattern in a torch.int64 tensor.""" + return value if value < (1 << 63) else value - (1 << 64) + + @triton.jit def _copy_mamba_state_block( state_idx, @@ -836,7 +841,9 @@ def _populate_metadata( for state_type_idx, copy_func in enumerate(mamba_state_copy_funcs): state = kv_caches[state_type_idx] # Base address - self.state_base_addrs[idx] = state.data_ptr() + self.state_base_addrs[idx] = _reinterpret_u64_as_i64( + state.data_ptr() + ) # Block stride (bytes between consecutive blocks) # state shape: [num_blocks, ...], stride(0) = elements per block @@ -923,7 +930,7 @@ def _populate_metadata( ) self.block_table_stride_req = int(next(iter(strides))) for i, bt in enumerate(block_tables): - self.block_table_ptrs[i] = bt.data_ptr() + self.block_table_ptrs[i] = _reinterpret_u64_as_i64(bt.data_ptr()) self.is_initialized = True From 03b87dc42f1f3750ddc37da359fe2b670d78e925 Mon Sep 17 00:00:00 2001 From: Qiming Zhang Date: Tue, 18 Aug 2026 19:47:37 -0700 Subject: [PATCH 124/839] [XPU] upgrade requirements/test/xpu.txt (#52672) Signed-off-by: mayuyuace Signed-off-by: Qiming Zhang --- requirements/test/xpu.txt | 352 ++++++++++++++++++++------------------ 1 file changed, 183 insertions(+), 169 deletions(-) diff --git a/requirements/test/xpu.txt b/requirements/test/xpu.txt index 6e8ac715acc3..c4376e00b056 100644 --- a/requirements/test/xpu.txt +++ b/requirements/test/xpu.txt @@ -1,14 +1,14 @@ # This file was autogenerated by uv via the following command: # uv pip compile requirements/test/xpu.in -c requirements/xpu.txt -o requirements/test/xpu.txt --index-strategy unsafe-best-match --python-platform x86_64-manylinux_2_39 --python-version 3.12 -absl-py==2.4.0 +absl-py==2.5.0 # via # -r requirements/test/xpu.in # rouge-score -accelerate==1.13.0 +accelerate==1.14.0 # via -r requirements/test/xpu.in -aiohappyeyeballs==2.6.1 +aiohappyeyeballs==2.7.1 # via aiohttp -aiohttp==3.13.4 +aiohttp==3.14.3 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -17,30 +17,33 @@ aiohttp==3.13.4 # lm-eval aiosignal==1.4.0 # via aiohttp -albumentations==1.4.6 +albucore==0.0.24 + # via albumentations +albumentations==2.0.8 # via -r requirements/test/xpu.in -annotated-doc==0.0.4 +annotated-doc==0.0.5 # via # fastapi # typer -annotated-types==0.7.0 +annotated-types==0.8.0 # via pydantic -anthropic==0.112.0 +anthropic==0.122.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt -anyio==4.13.0 +anyio==4.14.2 # via # anthropic # httpx + # httpx2 # mcp # openai # sse-starlette # starlette # watchfiles -apache-tvm-ffi==0.1.12 +apache-tvm-ffi==0.1.13.post3 # via xgrammar -arctic-inference==0.1.1 +arctic-inference==0.2.0 # via -r requirements/test/xpu.in astor==0.8.1 # via depyf @@ -49,37 +52,35 @@ attrs==26.1.0 # aiohttp # jsonschema # referencing -audioread==3.0.1 - # via - # -r requirements/test/xpu.in - # librosa +audioread==3.1.0 + # via -r requirements/test/xpu.in blake3==1.0.9 # via -r requirements/test/../common.txt -blobfile==3.0.0 +blobfile==3.2.0 # via -r requirements/test/xpu.in -bm25s==0.2.13 +bm25s==0.3.10 # via # -r requirements/test/xpu.in # mteb bounded-pool-executor==0.0.3 # via pqdm -cachetools==7.1.4 +cachetools==7.1.7 # via -r requirements/test/../common.txt -cbor2==6.1.2 +cbor2==6.1.4 # via -r requirements/test/../common.txt -certifi==2026.2.25 +certifi==2026.7.22 # via # httpcore # httpx # requests # sentry-sdk -cffi==2.0.0 +cffi==2.1.1 # via # cryptography # soundfile -chardet==5.2.0 +chardet==6.0.0.post1 # via mbstrdecoder -charset-normalizer==3.4.6 +charset-normalizer==3.5.1 # via requests chz==0.4.0 # via gpt-oss @@ -90,7 +91,6 @@ click==8.4.2 # nltk # rich-toolkit # schemathesis - # typer # uvicorn cloudpickle==3.1.2 # via -r requirements/test/../common.txt @@ -100,21 +100,23 @@ compressed-tensors==0.17.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt -coverage==7.13.5 +coverage==7.15.4 # via pytest-cov -cryptography==49.0.0 +cryptography==50.0.0 # via pyjwt -dataproperty==1.1.0 +dataproperty==1.1.1 # via # pytablewriter # tabledata -datasets==4.8.4 +datasets==5.0.1 # via # evaluate # lm-eval # mteb -decorator==5.2.1 +decorator==5.3.1 # via librosa +defusedxml==0.7.1 + # via nltk depyf==0.20.0 # via # -c requirements/common.txt @@ -134,7 +136,7 @@ distro==1.9.0 # openai dnspython==2.8.0 # via email-validator -docker==7.1.0 +docker==7.2.0 # via gpt-oss docopt==0.6.2 # via num2words @@ -156,19 +158,21 @@ email-validator==2.3.0 # pydantic evaluate==0.4.6 # via lm-eval -fastapi==0.135.2 +fastapi==0.136.3 # via # -c requirements/common.txt # -r requirements/test/../common.txt # gpt-oss # model-hosting-container-standards -fastapi-cli==0.0.27 +fastapi-cli==0.0.32 # via fastapi -fastapi-cloud-cli==0.21.0 +fastapi-cloud-cli==0.23.0 # via fastapi-cli fastar==0.11.0 - # via fastapi-cloud-cli -filelock==3.25.2 + # via + # fastapi + # fastapi-cloud-cli +filelock==3.32.3 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -181,27 +185,28 @@ frozenlist==1.8.0 # via # aiohttp # aiosignal -fsspec==2026.2.0 +fsspec==2026.6.0 # via # datasets # evaluate # huggingface-hub # torch -googleapis-common-protos==1.75.0 +googleapis-common-protos==1.75.1 # via # opentelemetry-exporter-otlp-proto-grpc # opentelemetry-exporter-otlp-proto-http -gpt-oss==0.0.8 +gpt-oss==0.0.9 # via -r requirements/test/xpu.in -graphql-core==3.2.8 +graphql-core==3.2.11 # via hypothesis-graphql -grpcio==1.81.1 +grpcio==1.83.0 # via opentelemetry-exporter-otlp-proto-grpc h11==0.16.0 # via # httpcore + # httpcore2 # uvicorn -harfile==0.4.0 +harfile==0.5.0 # via schemathesis hf-xet==1.6.0 # via huggingface-hub @@ -209,6 +214,8 @@ html2text==2025.4.15 # via gpt-oss httpcore==1.0.9 # via httpx +httpcore2==2.10.0 + # via httpx2 httptools==0.8.0 # via uvicorn httpx==0.28.1 @@ -218,12 +225,11 @@ httpx==0.28.1 # fastapi # fastapi-cloud-cli # huggingface-hub - # mcp # model-hosting-container-standards +httpx2==2.10.0 + # via + # mcp # openai - # schemathesis -httpx-sse==0.4.3 - # via mcp huggingface-hub==1.28.0 # via # -c requirements/common.txt @@ -235,26 +241,25 @@ huggingface-hub==1.28.0 # timm # tokenizers # transformers -hypothesis==6.151.10 +hypothesis==6.165.10 # via # hypothesis-graphql # hypothesis-jsonschema # schemathesis -hypothesis-graphql==0.12.0 +hypothesis-graphql==0.13.1 # via schemathesis hypothesis-jsonschema==0.23.1 # via schemathesis -idna==3.11 +idna==3.18 # via # anyio # email-validator # httpx + # httpx2 # requests # yarl -ijson==3.5.0 +ijson==3.5.1 # via -r requirements/test/../common.txt -imageio==2.37.3 - # via scikit-image impi-rt==2021.18.0 # via # oneccl @@ -304,7 +309,7 @@ jinja2==3.1.6 # fastapi # lm-eval # torch -jiter==0.15.0 +jiter==0.16.0 # via # anthropic # openai @@ -324,22 +329,21 @@ jsonschema==4.26.0 # hypothesis-jsonschema # mcp # mistral-common - # schemathesis -jsonschema-rs==0.45.0 +jsonschema-rs==0.49.9 # via schemathesis jsonschema-specifications==2025.9.1 # via jsonschema -junit-xml==1.9 - # via schemathesis +jupyter-client==8.9.1 + # via gpt-oss +jupyter-core==5.9.1 + # via jupyter-client lark==1.2.2 # via # -c requirements/common.txt # -r requirements/test/../common.txt lazy-loader==0.5 - # via - # librosa - # scikit-image -librosa==0.10.2.post1 + # via librosa +librosa==1.0.0 # via -r requirements/test/xpu.in llguidance==1.7.6 # via @@ -355,27 +359,29 @@ lm-format-enforcer==0.11.3 # -r requirements/test/../common.txt loguru==0.7.3 # via compressed-tensors -lxml==6.0.2 +lxml==6.1.1 # via # blobfile # gpt-oss # sacrebleu -markdown-it-py==4.0.0 +markdown-it-py==4.2.0 # via rich markupsafe==3.0.3 # via # jinja2 # werkzeug -mbstrdecoder==1.1.4 +mbstrdecoder==1.1.5 # via # dataproperty # pytablewriter # typepy -mcp==1.28.1 +mcp==2.0.0 # via -r requirements/test/../common.txt +mcp-types==2.0.0 + # via mcp mdurl==0.1.2 # via markdown-it-py -mistral-common==1.11.6 +mistral-common==1.11.7 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -392,17 +398,17 @@ model-hosting-container-standards==0.1.16 # via # -c requirements/common.txt # -r requirements/test/../common.txt -modelscope==1.35.3 +modelscope==1.37.1 # via -r requirements/test/xpu.in -more-itertools==10.8.0 +more-itertools==11.1.0 # via lm-eval mpmath==1.3.0 # via sympy -msgpack==1.1.2 +msgpack==1.2.1 # via librosa msgspec==0.21.1 # via -r requirements/test/../common.txt -mteb==2.12.7 +mteb==2.19.4 # via -r requirements/test/xpu.in multidict==6.7.1 # via @@ -412,13 +418,13 @@ multiprocess==0.70.19 # via # datasets # evaluate +narwhals==2.24.0 + # via scikit-learn networkx==3.6.1 - # via - # scikit-image - # torch + # via torch ninja==1.13.0 # via -r requirements/test/../common.txt -nltk==3.9.4 +nltk==3.10.3 # via rouge-score num2words==0.5.14 # via -r requirements/test/xpu.in @@ -426,15 +432,15 @@ numba==0.65.0 # via # -c requirements/xpu.txt # librosa -numpy==2.2.6 +numpy==2.3.5 # via # -r requirements/test/../common.txt # accelerate + # albucore # albumentations # bm25s # datasets # evaluate - # imageio # librosa # lm-eval # mistral-common @@ -445,13 +451,11 @@ numpy==2.2.6 # pytrec-eval-terrier # rouge-score # sacrebleu - # scikit-image # scikit-learn # scipy # sentence-transformers # soundfile # soxr - # tifffile # torchvision # transformers # xgrammar @@ -478,7 +482,7 @@ onemkl-sycl-rng==2026.0.0 # via torch onemkl-sycl-sparse==2026.0.0 # via torch -openai==2.44.0 +openai==3.2.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -487,45 +491,47 @@ openai-harmony==0.0.8 # -c requirements/common.txt # -r requirements/test/../common.txt # gpt-oss -opencv-python-headless==4.13.0.92 +opencv-python-headless==5.0.0.93 # via # -c requirements/common.txt # -r requirements/test/../common.txt + # albucore # albumentations # mistral-common -opentelemetry-api==1.43.0 +opentelemetry-api==1.44.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt + # mcp # opentelemetry-exporter-otlp-proto-grpc # opentelemetry-exporter-otlp-proto-http # opentelemetry-sdk # opentelemetry-semantic-conventions -opentelemetry-exporter-otlp==1.43.0 +opentelemetry-exporter-otlp==1.44.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt -opentelemetry-exporter-otlp-proto-common==1.43.0 +opentelemetry-exporter-otlp-proto-common==1.44.0 # via # opentelemetry-exporter-otlp-proto-grpc # opentelemetry-exporter-otlp-proto-http -opentelemetry-exporter-otlp-proto-grpc==1.43.0 +opentelemetry-exporter-otlp-proto-grpc==1.44.0 # via opentelemetry-exporter-otlp -opentelemetry-exporter-otlp-proto-http==1.43.0 +opentelemetry-exporter-otlp-proto-http==1.44.0 # via opentelemetry-exporter-otlp -opentelemetry-proto==1.43.0 +opentelemetry-proto==1.44.0 # via # opentelemetry-exporter-otlp-proto-common # opentelemetry-exporter-otlp-proto-grpc # opentelemetry-exporter-otlp-proto-http -opentelemetry-sdk==1.43.0 +opentelemetry-sdk==1.44.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt # opentelemetry-exporter-otlp-proto-grpc # opentelemetry-exporter-otlp-proto-http # opentelemetry-semantic-conventions-ai -opentelemetry-semantic-conventions==0.64b0 +opentelemetry-semantic-conventions==0.65b0 # via # opentelemetry-sdk # opentelemetry-semantic-conventions-ai @@ -537,7 +543,7 @@ outlines-core==0.2.14 # via # -c requirements/common.txt # -r requirements/test/../common.txt -packaging==26.0 +packaging==26.3 # via # -c requirements/xpu.txt # accelerate @@ -550,10 +556,9 @@ packaging==26.0 # pooch # pytest # pytest-rerunfailures - # scikit-image # transformers # typepy -pandas==3.0.1 +pandas==3.0.5 # via # datasets # evaluate @@ -561,41 +566,41 @@ partial-json-parser==0.2.1.1.post7 # via -r requirements/test/../common.txt pathvalidate==3.3.1 # via pytablewriter -pillow==12.1.1 +pillow==12.3.0 # via # -r requirements/test/../common.txt - # imageio # mistral-common - # scikit-image # torchvision -platformdirs==4.9.4 - # via pooch +platformdirs==4.11.3 + # via + # jupyter-core + # pooch pluggy==1.6.0 # via # pytest # pytest-cov -polars==1.39.3 +polars==1.43.2 # via mteb -polars-runtime-32==1.39.3 +polars-runtime-32==1.43.2 # via polars -pooch==1.8.2 +pooch==1.9.0 # via # -r requirements/test/xpu.in # librosa -portalocker==3.2.0 +portalocker==4.1.0 # via sacrebleu pqdm==0.2.0 # via -r requirements/test/xpu.in -prometheus-client==0.25.0 +prometheus-client==0.26.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt # prometheus-fastapi-instrumentator -prometheus-fastapi-instrumentator==8.0.2 +prometheus-fastapi-instrumentator==8.1.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt -propcache==0.4.1 +propcache==0.5.2 # via # aiohttp # yarl @@ -609,13 +614,11 @@ psutil==7.2.2 # via # -r requirements/test/../common.txt # accelerate -py==1.11.0 - # via pytest-forked py-cpuinfo==9.0.0 # via -r requirements/test/../common.txt -pyarrow==23.0.1 +pyarrow==25.0.1 # via datasets -pybase64==1.4.3 +pybase64==1.5.0 # via -r requirements/test/../common.txt pycountry==26.2.16 # via pydantic-extra-types @@ -623,7 +626,7 @@ pycparser==3.0 # via cffi pycryptodomex==3.23.0 # via blobfile -pydantic==2.12.5 +pydantic==2.13.4 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -635,6 +638,7 @@ pydantic==2.12.5 # gpt-oss # lm-format-enforcer # mcp + # mcp-types # mistral-common # model-hosting-container-standards # mteb @@ -643,33 +647,31 @@ pydantic==2.12.5 # pydantic-extra-types # pydantic-settings # xgrammar -pydantic-core==2.41.5 +pydantic-core==2.46.4 # via pydantic pydantic-extra-types==2.11.1 # via # fastapi # mistral-common -pydantic-settings==2.14.2 - # via - # fastapi - # mcp -pyelftools==0.32 +pydantic-settings==2.15.0 + # via fastapi +pyelftools==0.33 # via triton-xpu -pygments==2.20.0 +pygments==2.21.0 # via # pytest # rich pyjwt==2.13.0 # via mcp -pyrate-limiter==4.1.0 +pyrate-limiter==4.4.0 # via schemathesis -pystemmer==3.0.0 +pystemmer==3.1.0 # via # -r requirements/test/xpu.in # mteb pytablewriter==1.2.1 # via lm-eval -pytest==9.0.2 +pytest==9.1.1 # via # -r requirements/test/xpu.in # pytest-asyncio @@ -679,27 +681,28 @@ pytest==9.0.2 # pytest-shard # pytest-timeout # schemathesis -pytest-asyncio==1.3.0 +pytest-asyncio==1.4.0 # via -r requirements/test/xpu.in -pytest-cov==6.3.0 +pytest-cov==7.1.0 # via -r requirements/test/xpu.in -pytest-forked==1.6.0 +pytest-forked==1.7.5 # via -r requirements/test/xpu.in -pytest-rerunfailures==14.0 +pytest-rerunfailures==16.6 # via -r requirements/test/xpu.in pytest-shard==0.1.2 # via -r requirements/test/xpu.in -pytest-timeout==2.3.1 +pytest-timeout==2.4.0 # via -r requirements/test/xpu.in python-dateutil==2.9.0.post0 # via + # jupyter-client # pandas # typepy -python-dotenv==1.2.2 +python-dotenv==1.2.3 # via # pydantic-settings # uvicorn -python-json-logger==4.1.0 +python-json-logger==4.2.0 # via -r requirements/test/../common.txt python-multipart==0.0.32 # via @@ -707,7 +710,7 @@ python-multipart==0.0.32 # mcp pytrec-eval-terrier==0.5.10 # via mteb -pytz==2026.1.post1 +pytz==2026.3.post1 # via typepy pyyaml==6.0.3 # via @@ -727,7 +730,8 @@ pyzmq==27.1.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt -rapidfuzz==3.12.1 + # jupyter-client +rapidfuzz==3.14.5 # via # -r requirements/test/xpu.in # jiwer @@ -735,14 +739,14 @@ referencing==0.37.0 # via # jsonschema # jsonschema-specifications -regex==2026.3.32 +regex==2026.7.19 # via # -r requirements/test/../common.txt # nltk # sacrebleu # tiktoken # transformers -requests==2.33.1 +requests==2.34.2 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -759,21 +763,21 @@ requests==2.33.1 # schemathesis # starlette-testclient # tiktoken -rich==14.3.3 +rich==15.0.0 # via # mteb # rich-toolkit # schemathesis # typer -rich-toolkit==0.20.1 +rich-toolkit==0.20.3 # via # fastapi-cli # fastapi-cloud-cli -rignore==0.7.6 +rignore==0.8.1 # via fastapi-cloud-cli rouge-score==0.1.2 # via lm-eval -rpds-py==0.30.0 +rpds-py==2026.6.3 # via # jsonschema # referencing @@ -786,32 +790,27 @@ safetensors==0.8.0 # accelerate # timm # transformers -schemathesis==4.14.2 +schemathesis==4.24.3 # via -r requirements/test/xpu.in -scikit-image==0.26.0 - # via albumentations -scikit-learn==1.8.0 +scikit-learn==1.9.0 # via - # albumentations # librosa # lm-eval # mteb # sentence-transformers -scipy==1.17.1 +scipy==1.18.0 # via # albumentations - # bm25s # librosa # mteb # pytrec-eval-terrier - # scikit-image # scikit-learn # sentence-transformers -sentence-transformers==5.3.0 +sentence-transformers==5.7.0 # via mteb -sentencepiece==0.2.1 +sentencepiece==0.2.2 # via -r requirements/test/../common.txt -sentry-sdk==2.63.0 +sentry-sdk==2.68.0 # via fastapi-cloud-cli setproctitle==1.3.7 # via -r requirements/test/../common.txt @@ -826,11 +825,12 @@ setuptools==80.10.2 # torch shellingham==1.5.4 # via typer +simsimd==6.5.16 + # via albucore six==1.17.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt - # junit-xml # python-dateutil # rouge-score sniffio==1.3.1 @@ -839,21 +839,21 @@ sniffio==1.3.1 # openai sortedcontainers==2.4.0 # via hypothesis -soundfile==0.13.1 +soundfile==0.14.0 # via # -r requirements/test/xpu.in # librosa # mistral-common -soxr==0.5.0.post1 +soxr==1.1.0 # via # -r requirements/test/xpu.in # librosa # mistral-common sqlitedict==2.1.0 # via lm-eval -sse-starlette==3.4.5 +sse-starlette==3.4.8 # via mcp -starlette==1.3.1 +starlette==1.6.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -865,13 +865,15 @@ starlette==1.3.1 # starlette-testclient starlette-testclient==0.4.1 # via schemathesis -structlog==25.5.0 +stringzilla==5.1.2 + # via albucore +structlog==26.1.0 # via gpt-oss supervisor==4.3.0 # via model-hosting-container-standards sympy==1.14.0 # via torch -tabledata==1.3.4 +tabledata==1.3.5 # via pytablewriter tabulate==0.10.0 # via sacrebleu @@ -880,7 +882,7 @@ tbb==2023.0.0 # intel-opencl-rt # mkl # torch -tblib==3.1.0 +tblib==3.2.2 # via -r requirements/test/xpu.in tcmlib==1.5.0 # via @@ -893,26 +895,24 @@ tenacity==9.1.4 # via # gpt-oss # lm-eval - # schemathesis termcolor==3.3.0 # via gpt-oss threadpoolctl==3.6.0 # via scikit-learn -tifffile==2026.3.3 - # via scikit-image -tiktoken==0.12.0 +tiktoken==0.14.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt # gpt-oss # lm-eval # mistral-common -timm==1.0.17 +timm==1.0.28 # via -r requirements/test/xpu.in tokenizers==0.22.2 # via # -c requirements/common.txt # -r requirements/test/../common.txt + # sentence-transformers # transformers torch==2.13.0+xpu # via @@ -926,7 +926,9 @@ torch==2.13.0+xpu # xgrammar torchvision==0.28.0+xpu # via timm -tqdm==4.67.3 +tornado==6.5.8 + # via jupyter-client +tqdm==4.70.0 # via # -r requirements/test/../common.txt # datasets @@ -940,12 +942,17 @@ tqdm==4.67.3 # pqdm # sentence-transformers # transformers +traitlets==5.16.1 + # via + # jupyter-client + # jupyter-core transformers==5.15.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt # -r requirements/test/xpu.in # compressed-tensors + # mteb # sentence-transformers # xgrammar triton==3.7.2+xpu @@ -956,32 +963,38 @@ triton-xpu==3.7.2 # via # torch # triton -typepy==1.3.4 +truststore==0.10.4 + # via + # httpcore2 + # httpx2 +typepy==1.3.5 # via # dataproperty # pytablewriter # tabledata -typer==0.24.1 +typer==0.27.1 # via # fastapi-cli # fastapi-cloud-cli # transformers -typing-extensions==4.15.0 +typing-extensions==4.16.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt + # aiohttp # aiosignal - # albumentations # anthropic # anyio # apache-tvm-ffi # chz # fastapi # grpcio + # httpx2 # huggingface-hub - # librosa + # jupyter-client # lm-eval # mcp + # mcp-types # mistral-common # mteb # openai @@ -999,11 +1012,12 @@ typing-extensions==4.15.0 # rich-toolkit # schemathesis # sentence-transformers + # soundfile # starlette # torch # typing-inspection # xgrammar -typing-inspection==0.4.2 +typing-inspection==0.4.4 # via # fastapi # mcp @@ -1013,14 +1027,14 @@ umf==1.1.0 # via # intel-cmplr-lib-ur # torch -urllib3==2.6.3 +urllib3==2.7.0 # via # blobfile # docker # modelscope # requests # sentry-sdk -uvicorn==0.42.0 +uvicorn==0.52.3 # via # fastapi # fastapi-cli @@ -1033,9 +1047,9 @@ watchfiles==1.2.0 # via # -r requirements/test/../common.txt # uvicorn -websockets==16.0 +websockets==17.0.1 # via uvicorn -werkzeug==3.1.7 +werkzeug==3.1.8 # via schemathesis word2number==1.1 # via lm-eval @@ -1043,9 +1057,9 @@ xgrammar==0.2.3 # via # -c requirements/common.txt # -r requirements/test/../common.txt -xxhash==3.6.0 +xxhash==4.0.1 # via # datasets # evaluate -yarl==1.23.0 +yarl==1.24.5 # via aiohttp From aabc1a0a0b55c768ab1f8962157f5fa086932b88 Mon Sep 17 00:00:00 2001 From: Tyler Michael Smith Date: Tue, 18 Aug 2026 22:52:25 -0400 Subject: [PATCH 125/839] [Bugfix][Benchmark] Check readiness before tokenizer init in rust vllm-bench (#51863) Signed-off-by: Tyler Michael Smith Co-authored-by: OpenAI Codex --- rust/src/bench/src/benchmark.rs | 51 +++--------- rust/src/bench/src/multi_turn.rs | 50 ++++------- rust/src/bench/src/ready_checker.rs | 124 +++++++++++++++++++++++++++- rust/src/bench/src/tokenizer.rs | 15 ++-- 4 files changed, 159 insertions(+), 81 deletions(-) diff --git a/rust/src/bench/src/benchmark.rs b/rust/src/bench/src/benchmark.rs index cc4acc4f0614..d2a7bf84802d 100644 --- a/rust/src/bench/src/benchmark.rs +++ b/rust/src/bench/src/benchmark.rs @@ -18,7 +18,7 @@ use crate::metrics::steady_state; use crate::output::console::print_results; use crate::output::json::{append_result, build_result_json, compute_result_filename, save_result}; use crate::rate_control::compute_schedule; -use crate::ready_checker::wait_for_endpoint; +use crate::ready_checker::{get_first_model, wait_for_endpoint}; /// Pre-resolve the hostname in `base_url` and pin all resolved IPs on the /// client builder via [`reqwest::ClientBuilder::resolve_to_addrs`]. This @@ -281,40 +281,6 @@ pub(crate) fn assign_lora_modules( Some(out) } -/// Fetch the first model from the server's /v1/models endpoint. -async fn get_first_model_from_server( - base_url: &str, - client: &reqwest::Client, - extra_headers: &Option>, -) -> Result<(String, String)> { - let url = format!("{base_url}/v1/models"); - let mut request = client.get(&url); - if let Some(headers) = extra_headers { - for (k, v) in headers { - request = request.header(k, v); - } - } - // Add API key from environment - if let Ok(api_key) = std::env::var("OPENAI_API_KEY") { - request = request.header("Authorization", format!("Bearer {api_key}")); - } - - let response = request.send().await?; - let data: serde_json::Value = response.json().await?; - - if let Some(models) = data.get("data").and_then(|d| d.as_array()) - && let Some(first) = models.first() - { - let id = first.get("id").and_then(|v| v.as_str()).unwrap_or_default().to_string(); - let root = first.get("root").and_then(|v| v.as_str()).unwrap_or(&id).to_string(); - return Ok((id, root)); - } - - Err(BenchError::Config(format!( - "No models found on the server at {base_url}" - ))) -} - /// Run the complete benchmark. /// /// This is the core orchestrator matching Python's benchmark() + main_async(). @@ -348,8 +314,13 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { (m.clone(), config.model_name.clone()) } else { tracing::info!(base_url = %config.base_url, "fetching first model from server"); - let (name, id) = - get_first_model_from_server(&config.base_url, &client, &config.extra_headers).await?; + let (name, id) = get_first_model( + &config.base_url, + &client, + &config.extra_headers, + config.ready_check_timeout_sec, + ) + .await?; tracing::info!( model_name = name, model_id = id, @@ -364,7 +335,11 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { } else { let tid = config.tokenizer_id.as_deref().unwrap_or(&model_id); tracing::info!(tokenizer = tid, "loading tokenizer"); - let server_info = Some((config.base_url.as_str(), model_id.as_str())); + let server_info = Some(( + config.base_url.as_str(), + model_id.as_str(), + config.ready_check_timeout_sec, + )); let t = crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info).await?; Some(t) diff --git a/rust/src/bench/src/multi_turn.rs b/rust/src/bench/src/multi_turn.rs index 4b962b879c01..e0a31381705f 100644 --- a/rust/src/bench/src/multi_turn.rs +++ b/rust/src/bench/src/multi_turn.rs @@ -23,6 +23,7 @@ use crate::output::console::print_multi_turn_results; use crate::output::json::{ append_result, build_multi_turn_result_json, compute_result_filename, save_result, }; +use crate::ready_checker::{get_first_model, wait_for_endpoint}; /// Output from a single turn within a conversation. #[derive(Debug, Clone)] @@ -74,7 +75,13 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result Result Result>, -) -> Result<(String, String)> { - let url = format!("{base_url}/v1/models"); - let mut request = client.get(&url); - if let Some(headers) = extra_headers { - for (k, v) in headers { - request = request.header(k, v); - } - } - if let Ok(api_key) = std::env::var("OPENAI_API_KEY") { - request = request.header("Authorization", format!("Bearer {api_key}")); - } - - let response = request.send().await?; - let data: serde_json::Value = response.json().await?; - - if let Some(models) = data.get("data").and_then(|d| d.as_array()) - && let Some(first) = models.first() - { - let id = first.get("id").and_then(|v| v.as_str()).unwrap_or_default().to_string(); - let root = first.get("root").and_then(|v| v.as_str()).unwrap_or(&id).to_string(); - return Ok((id, root)); - } - - Err(BenchError::Config(format!( - "No models found on the server at {base_url}" - ))) -} - #[cfg(test)] mod tests { use std::sync::Arc; diff --git a/rust/src/bench/src/ready_checker.rs b/rust/src/bench/src/ready_checker.rs index 144a1ebd99ea..77d35fb278e8 100644 --- a/rust/src/bench/src/ready_checker.rs +++ b/rust/src/bench/src/ready_checker.rs @@ -1,14 +1,96 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project +use std::future::Future; use std::time::Instant; use indicatif::{ProgressBar, ProgressStyle}; +use thiserror_ext::AsReport as _; use crate::backends::{RequestFuncInput, RequestFuncOutput, get_backend}; use crate::cli::BackendKind; use crate::error::{BenchError, Result}; +/// Retry an operation until it succeeds or the readiness timeout expires. +pub(crate) async fn retry_with_timeout( + timeout_seconds: u64, + retry_interval: u64, + mut operation: F, +) -> Result +where + F: FnMut() -> Fut, + Fut: Future>, +{ + if timeout_seconds == 0 { + return operation().await; + } + + let deadline = Instant::now() + std::time::Duration::from_secs(timeout_seconds); + + loop { + let error = match operation().await { + Ok(value) => return Ok(value), + Err(error) => error, + }; + tracing::warn!(error = %error.as_report(), "endpoint is not ready"); + + let remaining = deadline.saturating_duration_since(Instant::now()); + if remaining.is_zero() { + return Err(BenchError::EndpointTimeout( + timeout_seconds, + format!("{}", error.as_report()), + )); + } + + let sleep_duration = + std::cmp::min(std::time::Duration::from_secs(retry_interval), remaining); + if !sleep_duration.is_zero() { + tokio::time::sleep(sleep_duration).await; + } + } +} + +/// Fetch the first model from the server, retrying while it starts. +pub(crate) async fn get_first_model( + base_url: &str, + client: &reqwest::Client, + extra_headers: &Option>, + timeout_seconds: u64, +) -> Result<(String, String)> { + let url = format!("{base_url}/v1/models"); + retry_with_timeout(timeout_seconds, 5, || async { + let mut request = client.get(&url); + if let Some(headers) = extra_headers { + for (key, value) in headers { + request = request.header(key, value); + } + } + // Add API key from environment + if let Ok(api_key) = std::env::var("OPENAI_API_KEY") { + request = request.header("Authorization", format!("Bearer {api_key}")); + } + + let response = request.send().await?.error_for_status()?; + let data: serde_json::Value = response.json().await?; + if let Some(model) = data + .get("data") + .and_then(|value| value.as_array()) + .and_then(|models| models.first()) + { + let id = + model.get("id").and_then(|value| value.as_str()).unwrap_or_default().to_string(); + let root = + model.get("root").and_then(|value| value.as_str()).unwrap_or(&id).to_string(); + return Ok((id, root)); + } + + Err(BenchError::Config(format!( + "No models found on the server at {base_url}" + ))) + }) + .await +} + /// Wait for the serving endpoint to become available. /// /// Sends test requests with retry until success or timeout. @@ -63,7 +145,7 @@ pub async fn wait_for_endpoint( last_error = err; } Err(e) => { - last_error = e.to_string(); + last_error = format!("{}", e.as_report()); } } @@ -76,3 +158,43 @@ pub async fn wait_for_endpoint( Err(BenchError::EndpointTimeout(timeout_seconds, last_error)) } + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use super::*; + + #[tokio::test] + async fn retry_succeeds_after_transient_failures() { + let attempts = AtomicUsize::new(0); + + let value = retry_with_timeout(1, 0, || async { + if attempts.fetch_add(1, Ordering::Relaxed) < 2 { + Err(BenchError::Backend("not ready".into())) + } else { + Ok(42) + } + }) + .await + .unwrap(); + + assert_eq!(value, 42); + assert_eq!(attempts.load(Ordering::Relaxed), 3); + } + + #[tokio::test] + async fn zero_timeout_does_not_retry() { + let attempts = AtomicUsize::new(0); + + let error = retry_with_timeout(0, 0, || async { + attempts.fetch_add(1, Ordering::Relaxed); + Err::<(), _>(BenchError::Backend("not ready".into())) + }) + .await + .unwrap_err(); + + assert!(matches!(error, BenchError::Backend(_))); + assert_eq!(attempts.load(Ordering::Relaxed), 1); + } +} diff --git a/rust/src/bench/src/tokenizer.rs b/rust/src/bench/src/tokenizer.rs index add890c39166..9854386e9a9e 100644 --- a/rust/src/bench/src/tokenizer.rs +++ b/rust/src/bench/src/tokenizer.rs @@ -30,7 +30,7 @@ pub struct ServerTokenizer { impl ServerTokenizer { /// Create a new server tokenizer and verify connectivity. - pub async fn new(base_url: &str, model: &str) -> Result { + pub async fn new(base_url: &str, model: &str, timeout_seconds: u64) -> Result { let client = reqwest::Client::builder() .timeout(std::time::Duration::from_secs(30)) .build() @@ -49,7 +49,10 @@ impl ServerTokenizer { }; // Probe the endpoint to verify it works and discover vocab size - let test_tokens = st.encode_async("test").await?; + let test_tokens = crate::ready_checker::retry_with_timeout(timeout_seconds, 5, || { + st.encode_async("test") + }) + .await?; let max_id = test_tokens.iter().copied().max().unwrap_or(0); let estimated_vocab = (max_id * 2).max(131072); @@ -225,11 +228,11 @@ impl TokenizerKind { /// 2. Tiktoken model file (for Kimi, Qwen, etc.) /// 3. Server-side /tokenize + /detokenize endpoints /// -/// `server_info` is `Some((base_url, model))` to enable server-side fallback. +/// `server_info` is `Some((base_url, model, timeout_seconds))` to enable server-side fallback. pub async fn load_tokenizer( model_id: &str, _trust_remote_code: bool, - server_info: Option<(&str, &str)>, + server_info: Option<(&str, &str, u64)>, ) -> Result { // 0. Check for built-in tiktoken encoding names (no HF download needed). These are useful for // consistent cross-model token counting (e.g. Artificial Analysis). @@ -275,13 +278,13 @@ pub async fn load_tokenizer( } Err(tiktoken_err) => { // 3. Try server-side fallback - if let Some((base_url, model)) = server_info { + if let Some((base_url, model, timeout_seconds)) = server_info { tracing::info!( model = model_id, error = %tiktoken_err.as_report(), "tiktoken unavailable; trying server-side tokenization" ); - match ServerTokenizer::new(base_url, model).await { + match ServerTokenizer::new(base_url, model, timeout_seconds).await { Ok(srv) => { tracing::info!( model = model_id, From deeeae75d0e42199aaeb627ad66e38b28337d2bf Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Tue, 18 Aug 2026 19:56:33 -0700 Subject: [PATCH 126/839] [MM] Keep more metadata tensors on CPU (#52827) Signed-off-by: Nick Hill --- vllm/model_executor/models/cohere2_vision.py | 2 +- vllm/model_executor/models/deepseek_ocr.py | 4 +++- vllm/model_executor/models/deepseek_ocr2.py | 4 +++- vllm/model_executor/models/deepseek_vl2.py | 4 +++- vllm/model_executor/models/ernie45_vl.py | 4 ++-- vllm/model_executor/models/fireredasr2.py | 2 +- vllm/model_executor/models/fireredlid.py | 4 ++-- vllm/model_executor/models/funasr.py | 2 +- vllm/model_executor/models/funaudiochat.py | 4 +++- vllm/model_executor/models/glmasr.py | 4 ++-- vllm/model_executor/models/granite4_vision.py | 2 +- .../models/hyperclovax_vision.py | 10 +++++++--- vllm/model_executor/models/isaac.py | 2 +- vllm/model_executor/models/kanana_v.py | 4 ++-- vllm/model_executor/models/keye_vl1_5.py | 8 +++++--- vllm/model_executor/models/kimi_audio.py | 2 +- vllm/model_executor/models/kimi_vl.py | 2 +- .../model_executor/models/llava_onevision2.py | 8 ++++---- vllm/model_executor/models/mimo_v2_omni.py | 14 ++++++++----- vllm/model_executor/models/minicpmv.py | 8 ++++---- vllm/model_executor/models/minicpmv4_6.py | 10 ++++++---- vllm/model_executor/models/mllama4.py | 4 ++-- vllm/model_executor/models/molmo.py | 6 ++++-- vllm/model_executor/models/molmo2.py | 20 +++++++++++-------- .../models/moss_transcribe_diarize.py | 7 +++++-- vllm/model_executor/models/muse_glimmer.py | 8 ++++++-- .../model_executor/models/nano_nemotron_vl.py | 14 +++++++------ vllm/model_executor/models/ovis.py | 2 +- vllm/model_executor/models/paddleocr_vl.py | 2 +- vllm/model_executor/models/phi4mm.py | 2 +- .../models/qwen2_5_omni_thinker.py | 18 +++++++++++++---- vllm/model_executor/models/qwen2_audio.py | 2 +- vllm/model_executor/models/qwen3_asr.py | 10 +++++++--- .../models/qwen3_omni_moe_thinker.py | 6 +++++- vllm/model_executor/models/step3_vl.py | 2 +- .../models/transformers/multimodal.py | 18 ++++++++++++++--- vllm/model_executor/models/ultravox.py | 2 +- vllm/models/inkling/common/mm_preprocess.py | 4 ++-- 38 files changed, 148 insertions(+), 83 deletions(-) diff --git a/vllm/model_executor/models/cohere2_vision.py b/vllm/model_executor/models/cohere2_vision.py index b58df12d0d2c..b6def10aab3b 100644 --- a/vllm/model_executor/models/cohere2_vision.py +++ b/vllm/model_executor/models/cohere2_vision.py @@ -268,7 +268,7 @@ def _get_mm_fields_config( num_patches = hf_inputs.get("num_patches", torch.empty(0)) return dict( pixel_values=MultiModalFieldConfig.flat_from_sizes("image", num_patches), - num_patches=MultiModalFieldConfig.batched("image"), + num_patches=MultiModalFieldConfig.batched("image", keep_on_cpu=True), image_embeds=MultiModalFieldConfig.batched("image"), ) diff --git a/vllm/model_executor/models/deepseek_ocr.py b/vllm/model_executor/models/deepseek_ocr.py index ead643a7f819..4bb5598db3a6 100644 --- a/vllm/model_executor/models/deepseek_ocr.py +++ b/vllm/model_executor/models/deepseek_ocr.py @@ -320,7 +320,9 @@ def _get_mm_fields_config( patches_per_image = torch.where(is_tiled, images_spatial_crop.prod(dim=-1), 0) return dict( pixel_values=MultiModalFieldConfig.batched("image"), - images_spatial_crop=MultiModalFieldConfig.batched("image"), + images_spatial_crop=MultiModalFieldConfig.batched( + "image", keep_on_cpu=True + ), images_crop=MultiModalFieldConfig.flat_from_sizes( "image", patches_per_image ), diff --git a/vllm/model_executor/models/deepseek_ocr2.py b/vllm/model_executor/models/deepseek_ocr2.py index f96fb9c14610..aa6e9b1aec43 100644 --- a/vllm/model_executor/models/deepseek_ocr2.py +++ b/vllm/model_executor/models/deepseek_ocr2.py @@ -191,7 +191,9 @@ def _get_mm_fields_config( patches_per_image = torch.where(is_tiled, images_spatial_crop.prod(dim=-1), 0) return dict( pixel_values=MultiModalFieldConfig.batched("image"), - images_spatial_crop=MultiModalFieldConfig.batched("image"), + images_spatial_crop=MultiModalFieldConfig.batched( + "image", keep_on_cpu=True + ), images_crop=MultiModalFieldConfig.flat_from_sizes( "image", patches_per_image ), diff --git a/vllm/model_executor/models/deepseek_vl2.py b/vllm/model_executor/models/deepseek_vl2.py index a514bef57608..9ca05b6c77d1 100644 --- a/vllm/model_executor/models/deepseek_vl2.py +++ b/vllm/model_executor/models/deepseek_vl2.py @@ -270,7 +270,9 @@ def _get_mm_fields_config( return dict( pixel_values=MultiModalFieldConfig.flat_from_sizes("image", num_patches), - images_spatial_crop=MultiModalFieldConfig.batched("image"), + images_spatial_crop=MultiModalFieldConfig.batched( + "image", keep_on_cpu=True + ), image_embeds=MultiModalFieldConfig.batched("image"), ) diff --git a/vllm/model_executor/models/ernie45_vl.py b/vllm/model_executor/models/ernie45_vl.py index 0c1781b04fa8..9fa04470020d 100644 --- a/vllm/model_executor/models/ernie45_vl.py +++ b/vllm/model_executor/models/ernie45_vl.py @@ -1246,11 +1246,11 @@ def _get_mm_fields_config( pixel_values=MultiModalFieldConfig.flat_from_sizes( "image", image_grid_sizes ), - image_grid_thw=MultiModalFieldConfig.batched("image"), + image_grid_thw=MultiModalFieldConfig.batched("image", keep_on_cpu=True), pixel_values_videos=MultiModalFieldConfig.flat_from_sizes( "video", video_grid_sizes ), - video_grid_thw=MultiModalFieldConfig.batched("video"), + video_grid_thw=MultiModalFieldConfig.batched("video", keep_on_cpu=True), ) diff --git a/vllm/model_executor/models/fireredasr2.py b/vllm/model_executor/models/fireredasr2.py index eea0c7d8897e..231417d83367 100644 --- a/vllm/model_executor/models/fireredasr2.py +++ b/vllm/model_executor/models/fireredasr2.py @@ -260,7 +260,7 @@ def _get_mm_fields_config( return dict( input_features=MultiModalFieldConfig.batched("audio"), speech_lengths=MultiModalFieldConfig.batched("audio"), - fake_token_lengths=MultiModalFieldConfig.batched("audio"), + fake_token_lengths=MultiModalFieldConfig.batched("audio", keep_on_cpu=True), ) def _get_prompt_updates( diff --git a/vllm/model_executor/models/fireredlid.py b/vllm/model_executor/models/fireredlid.py index 804ed2bc9fd9..69fe995b6b23 100644 --- a/vllm/model_executor/models/fireredlid.py +++ b/vllm/model_executor/models/fireredlid.py @@ -482,8 +482,8 @@ def _get_mm_fields_config( ) -> Mapping[str, MultiModalFieldConfig]: return dict( input_features=MultiModalFieldConfig.batched("audio"), - speech_lengths=MultiModalFieldConfig.batched("audio"), - fake_token_lengths=MultiModalFieldConfig.batched("audio"), + speech_lengths=MultiModalFieldConfig.batched("audio", keep_on_cpu=True), + fake_token_lengths=MultiModalFieldConfig.batched("audio", keep_on_cpu=True), ) def _get_prompt_updates( diff --git a/vllm/model_executor/models/funasr.py b/vllm/model_executor/models/funasr.py index 56b1b68f4f7e..69f63d770baa 100644 --- a/vllm/model_executor/models/funasr.py +++ b/vllm/model_executor/models/funasr.py @@ -779,7 +779,7 @@ def _get_mm_fields_config( return dict( input_features=MultiModalFieldConfig.batched("audio"), speech_lengths=MultiModalFieldConfig.batched("audio"), - fake_token_lengths=MultiModalFieldConfig.batched("audio"), + fake_token_lengths=MultiModalFieldConfig.batched("audio", keep_on_cpu=True), ) def _get_prompt_updates( diff --git a/vllm/model_executor/models/funaudiochat.py b/vllm/model_executor/models/funaudiochat.py index 0a9e286783c4..b336b9801e2b 100644 --- a/vllm/model_executor/models/funaudiochat.py +++ b/vllm/model_executor/models/funaudiochat.py @@ -695,7 +695,9 @@ def _get_mm_fields_config( ) -> Mapping[str, MultiModalFieldConfig]: return { "speech_ids": MultiModalFieldConfig.batched("audio"), - "speech_attention_mask": MultiModalFieldConfig.batched("audio"), + "speech_attention_mask": MultiModalFieldConfig.batched( + "audio", keep_on_cpu=True + ), "input_features": MultiModalFieldConfig.batched("audio"), "feature_attention_mask": MultiModalFieldConfig.batched("audio"), "feature_exist_mask": MultiModalFieldConfig.batched("audio"), diff --git a/vllm/model_executor/models/glmasr.py b/vllm/model_executor/models/glmasr.py index 1ff7224d7e4e..454687f469f7 100644 --- a/vllm/model_executor/models/glmasr.py +++ b/vllm/model_executor/models/glmasr.py @@ -625,13 +625,13 @@ def _glmasr_field_config( feature_attention_mask=MultiModalFieldConfig.flat_from_sizes( "audio", chunk_counts, dim=0 ), - chunk_counts=MultiModalFieldConfig.batched("audio"), + chunk_counts=MultiModalFieldConfig.batched("audio", keep_on_cpu=True), ) return dict( audio_embeds=MultiModalFieldConfig.batched("audio"), input_features=MultiModalFieldConfig.batched("audio"), feature_attention_mask=MultiModalFieldConfig.batched("audio"), - chunk_counts=MultiModalFieldConfig.batched("audio"), + chunk_counts=MultiModalFieldConfig.batched("audio", keep_on_cpu=True), ) diff --git a/vllm/model_executor/models/granite4_vision.py b/vllm/model_executor/models/granite4_vision.py index 5c56717aa631..e82a42bf87c8 100644 --- a/vllm/model_executor/models/granite4_vision.py +++ b/vllm/model_executor/models/granite4_vision.py @@ -398,7 +398,7 @@ def _get_mm_fields_config( ) -> Mapping[str, MultiModalFieldConfig]: return dict( pixel_values=MultiModalFieldConfig.batched("image"), - image_sizes=MultiModalFieldConfig.batched("image"), + image_sizes=MultiModalFieldConfig.batched("image", keep_on_cpu=True), ) diff --git a/vllm/model_executor/models/hyperclovax_vision.py b/vllm/model_executor/models/hyperclovax_vision.py index 593be5d88f6b..647103aa44e1 100644 --- a/vllm/model_executor/models/hyperclovax_vision.py +++ b/vllm/model_executor/models/hyperclovax_vision.py @@ -324,10 +324,14 @@ def _get_mm_fields_config( ) -> Mapping[str, MultiModalFieldConfig]: fields = dict( pixel_values_images=MultiModalFieldConfig.batched("image"), - image_sizes_images=MultiModalFieldConfig.batched("image"), - vision_query_lengths_images=MultiModalFieldConfig.batched("image"), + image_sizes_images=MultiModalFieldConfig.batched("image", keep_on_cpu=True), + vision_query_lengths_images=MultiModalFieldConfig.batched( + "image", keep_on_cpu=True + ), pixel_values_videos=MultiModalFieldConfig.batched("video"), - vision_query_lengths_videos=MultiModalFieldConfig.batched("video"), + vision_query_lengths_videos=MultiModalFieldConfig.batched( + "video", keep_on_cpu=True + ), ) return fields diff --git a/vllm/model_executor/models/isaac.py b/vllm/model_executor/models/isaac.py index 87932b50328e..4979971b02db 100644 --- a/vllm/model_executor/models/isaac.py +++ b/vllm/model_executor/models/isaac.py @@ -440,7 +440,7 @@ def _get_mm_fields_config( "pixel_values": MultiModalFieldConfig.flat_from_sizes( "image", image_grid_sizes ), - "image_grid_thw": MultiModalFieldConfig.batched("image"), + "image_grid_thw": MultiModalFieldConfig.batched("image", keep_on_cpu=True), } def _get_prompt_updates( diff --git a/vllm/model_executor/models/kanana_v.py b/vllm/model_executor/models/kanana_v.py index b1a5f78b1b33..7a1421b515d7 100644 --- a/vllm/model_executor/models/kanana_v.py +++ b/vllm/model_executor/models/kanana_v.py @@ -561,8 +561,8 @@ def _get_mm_fields_config( mm_fields_config = dict( pixel_values=MultiModalFieldConfig.flat_from_sizes("image", pixel_sizes), - vision_grid_thw=MultiModalFieldConfig.batched("image"), - image_token_thw=MultiModalFieldConfig.batched("image"), + vision_grid_thw=MultiModalFieldConfig.batched("image", keep_on_cpu=True), + image_token_thw=MultiModalFieldConfig.batched("image", keep_on_cpu=True), ) return mm_fields_config diff --git a/vllm/model_executor/models/keye_vl1_5.py b/vllm/model_executor/models/keye_vl1_5.py index c7ecc68db1f2..a4dcd729db3b 100644 --- a/vllm/model_executor/models/keye_vl1_5.py +++ b/vllm/model_executor/models/keye_vl1_5.py @@ -309,13 +309,15 @@ def _keye_field_config( return dict( pixel_values=MultiModalFieldConfig.flat_from_sizes("image", image_grid_sizes), image_embeds=MultiModalFieldConfig.flat_from_sizes("image", image_grid_sizes), - image_grid_thw=MultiModalFieldConfig.batched("image"), + image_grid_thw=MultiModalFieldConfig.batched("image", keep_on_cpu=True), pixel_values_videos=MultiModalFieldConfig.flat_from_sizes( "video", video_num_patches ), video_embeds=MultiModalFieldConfig.flat_from_sizes("video", video_num_patches), - video_grid_thw=MultiModalFieldConfig.flat_from_sizes("video", video_num_grids), - num_frames=MultiModalFieldConfig.batched("video"), + video_grid_thw=MultiModalFieldConfig.flat_from_sizes( + "video", video_num_grids, keep_on_cpu=True + ), + num_frames=MultiModalFieldConfig.batched("video", keep_on_cpu=True), ) diff --git a/vllm/model_executor/models/kimi_audio.py b/vllm/model_executor/models/kimi_audio.py index cb3c83e70891..900894612ce2 100644 --- a/vllm/model_executor/models/kimi_audio.py +++ b/vllm/model_executor/models/kimi_audio.py @@ -201,7 +201,7 @@ def get_dummy_processor_inputs( # Field config for Kimi-Audio multimodal data _KIMIAUDIO_FIELD_CONFIG = { "whisper_input_features": MultiModalFieldConfig.batched("audio"), - "feature_attention_mask": MultiModalFieldConfig.batched("audio"), + "feature_attention_mask": MultiModalFieldConfig.batched("audio", keep_on_cpu=True), } diff --git a/vllm/model_executor/models/kimi_vl.py b/vllm/model_executor/models/kimi_vl.py index 2b08fc6c1fd2..e5fa22de7e1a 100644 --- a/vllm/model_executor/models/kimi_vl.py +++ b/vllm/model_executor/models/kimi_vl.py @@ -251,7 +251,7 @@ def _get_mm_fields_config( pixel_values=MultiModalFieldConfig.flat_from_sizes( "image", image_grid_sizes ), - image_grid_hws=MultiModalFieldConfig.batched("image"), + image_grid_hws=MultiModalFieldConfig.batched("image", keep_on_cpu=True), ) def _get_prompt_updates( diff --git a/vllm/model_executor/models/llava_onevision2.py b/vllm/model_executor/models/llava_onevision2.py index 01f073a0d92c..6c13f5b14b61 100644 --- a/vllm/model_executor/models/llava_onevision2.py +++ b/vllm/model_executor/models/llava_onevision2.py @@ -431,20 +431,20 @@ def _field_config(hf_inputs: Mapping[str, torch.Tensor]): image_embeds=MultiModalFieldConfig.flat_from_sizes( "image", image_embed_grid_sizes ), - image_grid_thw=MultiModalFieldConfig.batched("image"), + image_grid_thw=MultiModalFieldConfig.batched("image", keep_on_cpu=True), # OV2 first-class MM kwarg: per-patch (t,h,w) # positions required by the 3-D vision RoPE. patch_positions=MultiModalFieldConfig.flat_from_sizes( - "image", image_pixel_grid_sizes + "image", image_pixel_grid_sizes, keep_on_cpu=True ), pixel_values_videos=MultiModalFieldConfig.flat_from_sizes( "video", video_patch_sizes ), video_grid_thw=MultiModalFieldConfig.flat_from_sizes( - "video", video_num_frames + "video", video_num_frames, keep_on_cpu=True ), patch_positions_videos=MultiModalFieldConfig.flat_from_sizes( - "video", video_patch_sizes + "video", video_patch_sizes, keep_on_cpu=True ), video_num_frames=MultiModalFieldConfig.batched("video", keep_on_cpu=True), frame_timestamps=MultiModalFieldConfig.batched("video", keep_on_cpu=True), diff --git a/vllm/model_executor/models/mimo_v2_omni.py b/vllm/model_executor/models/mimo_v2_omni.py index 589ca4f61eab..bd2bdbf9018b 100644 --- a/vllm/model_executor/models/mimo_v2_omni.py +++ b/vllm/model_executor/models/mimo_v2_omni.py @@ -893,17 +893,21 @@ def _get_mm_fields_config( merge_size = self.info.get_hf_config().vision_config.spatial_merge_size fields: dict[str, MultiModalFieldConfig] = dict( **_create_qwen2vl_field_factory(merge_size)(hf_inputs), - second_per_grid_ts=MultiModalFieldConfig.batched("video"), - video_start_times=MultiModalFieldConfig.batched("video"), + second_per_grid_ts=MultiModalFieldConfig.batched("video", keep_on_cpu=True), + video_start_times=MultiModalFieldConfig.batched("video", keep_on_cpu=True), audio_features=MultiModalFieldConfig.batched("audio"), - audio_token_lens=MultiModalFieldConfig.batched("audio"), + audio_token_lens=MultiModalFieldConfig.batched("audio", keep_on_cpu=True), ) # video_audio fields: only present when video_audio content was processed if "video_audio_n_segs" in hf_inputs: - fields["video_audio_n_segs"] = MultiModalFieldConfig.batched("video") + fields["video_audio_n_segs"] = MultiModalFieldConfig.batched( + "video", keep_on_cpu=True + ) # video_audio_seg_lens: list of per-video 1D tensors, batched("video") if "video_audio_seg_lens" in hf_inputs: - fields["video_audio_seg_lens"] = MultiModalFieldConfig.batched("video") + fields["video_audio_seg_lens"] = MultiModalFieldConfig.batched( + "video", keep_on_cpu=True + ) if "va_audio_features" in hf_inputs: fields["va_audio_features"] = MultiModalFieldConfig.batched("va_audio") return fields diff --git a/vllm/model_executor/models/minicpmv.py b/vllm/model_executor/models/minicpmv.py index e85ac81ee583..5d8eecfaaaa9 100644 --- a/vllm/model_executor/models/minicpmv.py +++ b/vllm/model_executor/models/minicpmv.py @@ -472,12 +472,12 @@ def get_version_by_config(config: PretrainedConfig) -> tuple[int, ...]: def _minicpmv_field_config(hf_inputs: Mapping[str, torch.Tensor]): return dict( pixel_values=MultiModalFieldConfig.batched("image"), - image_sizes=MultiModalFieldConfig.batched("image"), - tgt_sizes=MultiModalFieldConfig.batched("image"), + image_sizes=MultiModalFieldConfig.batched("image", keep_on_cpu=True), + tgt_sizes=MultiModalFieldConfig.batched("image", keep_on_cpu=True), image_embeds=MultiModalFieldConfig.batched("image"), video_pixel_values=MultiModalFieldConfig.batched("video"), - video_image_sizes=MultiModalFieldConfig.batched("video"), - video_tgt_sizes=MultiModalFieldConfig.batched("video"), + video_image_sizes=MultiModalFieldConfig.batched("video", keep_on_cpu=True), + video_tgt_sizes=MultiModalFieldConfig.batched("video", keep_on_cpu=True), video_embeds=MultiModalFieldConfig.batched("video"), ) diff --git a/vllm/model_executor/models/minicpmv4_6.py b/vllm/model_executor/models/minicpmv4_6.py index ac9efabaa3ab..a7e8d7c75f65 100644 --- a/vllm/model_executor/models/minicpmv4_6.py +++ b/vllm/model_executor/models/minicpmv4_6.py @@ -75,15 +75,17 @@ def _minicpmv4_6_field_config(hf_inputs: Mapping[str, torch.Tensor]): fields = dict( pixel_values=MultiModalFieldConfig.batched("image"), - tgt_sizes=MultiModalFieldConfig.batched("image"), + tgt_sizes=MultiModalFieldConfig.batched("image", keep_on_cpu=True), image_embeds=MultiModalFieldConfig.batched("image"), video_pixel_values=MultiModalFieldConfig.batched("video"), - video_image_sizes=MultiModalFieldConfig.batched("video"), - video_tgt_sizes=MultiModalFieldConfig.batched("video"), + video_image_sizes=MultiModalFieldConfig.batched("video", keep_on_cpu=True), + video_tgt_sizes=MultiModalFieldConfig.batched("video", keep_on_cpu=True), video_embeds=MultiModalFieldConfig.batched("video"), ) if "use_vit_merger" in hf_inputs: - fields["use_vit_merger"] = MultiModalFieldConfig.batched("image") + fields["use_vit_merger"] = MultiModalFieldConfig.batched( + "image", keep_on_cpu=True + ) return fields diff --git a/vllm/model_executor/models/mllama4.py b/vllm/model_executor/models/mllama4.py index 178dae506c6b..fe5ff17a3b46 100644 --- a/vllm/model_executor/models/mllama4.py +++ b/vllm/model_executor/models/mllama4.py @@ -653,8 +653,8 @@ def _get_mm_fields_config( pixel_values=MultiModalFieldConfig.flat_from_sizes( "image", patches_per_image ), - patches_per_image=MultiModalFieldConfig.batched("image"), - aspect_ratios=MultiModalFieldConfig.batched("image"), + patches_per_image=MultiModalFieldConfig.batched("image", keep_on_cpu=True), + aspect_ratios=MultiModalFieldConfig.batched("image", keep_on_cpu=True), ) def _get_prompt_updates( diff --git a/vllm/model_executor/models/molmo.py b/vllm/model_executor/models/molmo.py index bdc46d83b934..3fccb7a86cd2 100644 --- a/vllm/model_executor/models/molmo.py +++ b/vllm/model_executor/models/molmo.py @@ -1208,8 +1208,10 @@ def _get_mm_fields_config( images=MultiModalFieldConfig.flat_from_sizes("image", num_crops), image_masks=MultiModalFieldConfig.flat_from_sizes("image", num_crops), image_input_idx=MultiModalFieldConfig.flat_from_sizes("image", num_crops), - num_crops=MultiModalFieldConfig.batched("image"), - img_patch_id=MultiModalFieldConfig.shared("image", num_images), + num_crops=MultiModalFieldConfig.batched("image", keep_on_cpu=True), + img_patch_id=MultiModalFieldConfig.shared( + "image", num_images, keep_on_cpu=True + ), ) def _get_prompt_updates( diff --git a/vllm/model_executor/models/molmo2.py b/vllm/model_executor/models/molmo2.py index d8de4f5ed17c..42f5ebfdc5f6 100644 --- a/vllm/model_executor/models/molmo2.py +++ b/vllm/model_executor/models/molmo2.py @@ -2127,26 +2127,30 @@ def _get_mm_fields_config( image_token_pooling=MultiModalFieldConfig.flat_from_sizes( "image", image_num_pooled_patches ), - image_num_crops=MultiModalFieldConfig.batched("image"), - image_num_pooled_patches=MultiModalFieldConfig.batched("image"), - image_num_patches=MultiModalFieldConfig.batched("image"), + image_num_crops=MultiModalFieldConfig.batched("image", keep_on_cpu=True), + image_num_pooled_patches=MultiModalFieldConfig.batched( + "image", keep_on_cpu=True + ), + image_num_patches=MultiModalFieldConfig.batched("image", keep_on_cpu=True), image_tokens=MultiModalFieldConfig.flat_from_sizes( "image", num_image_tokens ), - num_image_tokens=MultiModalFieldConfig.batched("image"), + num_image_tokens=MultiModalFieldConfig.batched("image", keep_on_cpu=True), pixel_values_videos=MultiModalFieldConfig.flat_from_sizes( "video", video_num_crops ), video_token_pooling=MultiModalFieldConfig.flat_from_sizes( "video", video_num_pooled_patches ), - video_num_crops=MultiModalFieldConfig.batched("video"), - video_num_pooled_patches=MultiModalFieldConfig.batched("video"), - video_num_patches=MultiModalFieldConfig.batched("video"), + video_num_crops=MultiModalFieldConfig.batched("video", keep_on_cpu=True), + video_num_pooled_patches=MultiModalFieldConfig.batched( + "video", keep_on_cpu=True + ), + video_num_patches=MultiModalFieldConfig.batched("video", keep_on_cpu=True), video_tokens=MultiModalFieldConfig.flat_from_sizes( "video", num_video_tokens ), - num_video_tokens=MultiModalFieldConfig.batched("video"), + num_video_tokens=MultiModalFieldConfig.batched("video", keep_on_cpu=True), ) def _get_prompt_updates( diff --git a/vllm/model_executor/models/moss_transcribe_diarize.py b/vllm/model_executor/models/moss_transcribe_diarize.py index 819881e0f213..347c61cb5410 100644 --- a/vllm/model_executor/models/moss_transcribe_diarize.py +++ b/vllm/model_executor/models/moss_transcribe_diarize.py @@ -352,9 +352,12 @@ def _mtd_field_config( audio_feature_lengths=MultiModalFieldConfig.flat_from_sizes( "audio", audio_chunk_counts, + keep_on_cpu=True, + ), + audio_chunk_counts=MultiModalFieldConfig.batched("audio", keep_on_cpu=True), + audio_token_lengths=MultiModalFieldConfig.batched( + "audio", keep_on_cpu=True ), - audio_chunk_counts=MultiModalFieldConfig.batched("audio"), - audio_token_lengths=MultiModalFieldConfig.batched("audio"), ) return fields diff --git a/vllm/model_executor/models/muse_glimmer.py b/vllm/model_executor/models/muse_glimmer.py index 74e000131822..83ec17a75d56 100644 --- a/vllm/model_executor/models/muse_glimmer.py +++ b/vllm/model_executor/models/muse_glimmer.py @@ -378,12 +378,16 @@ def _get_mm_fields_config( if "image_pixel_values" in hf_inputs: fields.update( image_pixel_values=MultiModalFieldConfig.batched("image"), - image_feature_sizes=MultiModalFieldConfig.batched("image"), + image_feature_sizes=MultiModalFieldConfig.batched( + "image", keep_on_cpu=True + ), ) if "video_pixel_values" in hf_inputs: fields.update( video_pixel_values=MultiModalFieldConfig.batched("video"), - video_feature_sizes=MultiModalFieldConfig.batched("video"), + video_feature_sizes=MultiModalFieldConfig.batched( + "video", keep_on_cpu=True + ), ) return fields diff --git a/vllm/model_executor/models/nano_nemotron_vl.py b/vllm/model_executor/models/nano_nemotron_vl.py index c339b0e435b0..7ecf8677bb33 100644 --- a/vllm/model_executor/models/nano_nemotron_vl.py +++ b/vllm/model_executor/models/nano_nemotron_vl.py @@ -386,10 +386,12 @@ def _get_image_fields_config(self, hf_inputs: BatchFeature): return dict( pixel_values_flat=pixel_values_flat, - image_num_patches=MultiModalFieldConfig.batched("image"), + image_num_patches=MultiModalFieldConfig.batched("image", keep_on_cpu=True), image_embeds=MultiModalFieldConfig.batched("image"), - num_tokens_per_image=MultiModalFieldConfig.batched("image"), - imgs_sizes=MultiModalFieldConfig.batched("image"), + num_tokens_per_image=MultiModalFieldConfig.batched( + "image", keep_on_cpu=True + ), + imgs_sizes=MultiModalFieldConfig.batched("image", keep_on_cpu=True), ) def _get_video_fields_config(self, hf_inputs: BatchFeature): @@ -399,9 +401,9 @@ def _get_video_fields_config(self, hf_inputs: BatchFeature): pixel_values_flat_video=MultiModalFieldConfig.flat_from_sizes( "video", video_num_patches ), - video_num_patches=MultiModalFieldConfig.batched("video"), - frames_indices=MultiModalFieldConfig.batched("video"), - frame_duration_ms=MultiModalFieldConfig.batched("video"), + video_num_patches=MultiModalFieldConfig.batched("video", keep_on_cpu=True), + frames_indices=MultiModalFieldConfig.batched("video", keep_on_cpu=True), + frame_duration_ms=MultiModalFieldConfig.batched("video", keep_on_cpu=True), ) def _get_audio_fields_config(self, hf_inputs: BatchFeature): diff --git a/vllm/model_executor/models/ovis.py b/vllm/model_executor/models/ovis.py index f25585fd7643..5ec5f8cd1825 100644 --- a/vllm/model_executor/models/ovis.py +++ b/vllm/model_executor/models/ovis.py @@ -386,7 +386,7 @@ def _get_mm_fields_config( ) -> Mapping[str, MultiModalFieldConfig]: return dict( pixel_values=MultiModalFieldConfig.batched("image"), - grids=MultiModalFieldConfig.batched("image"), + grids=MultiModalFieldConfig.batched("image", keep_on_cpu=True), indicator_tokens=MultiModalFieldConfig.batched("image"), ) diff --git a/vllm/model_executor/models/paddleocr_vl.py b/vllm/model_executor/models/paddleocr_vl.py index bd635aaada2d..5f5302e7b78d 100644 --- a/vllm/model_executor/models/paddleocr_vl.py +++ b/vllm/model_executor/models/paddleocr_vl.py @@ -279,7 +279,7 @@ def _get_mm_fields_config( ) -> Mapping[str, MultiModalFieldConfig]: return dict( pixel_values=MultiModalFieldConfig.batched("image"), - image_grid_thw=MultiModalFieldConfig.batched("image"), + image_grid_thw=MultiModalFieldConfig.batched("image", keep_on_cpu=True), ) def _get_prompt_updates( diff --git a/vllm/model_executor/models/phi4mm.py b/vllm/model_executor/models/phi4mm.py index 9a8da5ed3215..359cc1613bea 100644 --- a/vllm/model_executor/models/phi4mm.py +++ b/vllm/model_executor/models/phi4mm.py @@ -918,7 +918,7 @@ def _get_mm_fields_config( input_image_embeds=MultiModalFieldConfig.batched("image"), image_attention_mask=MultiModalFieldConfig.batched("image"), image_sizes=MultiModalFieldConfig.batched("image", keep_on_cpu=True), - num_img_tokens=MultiModalFieldConfig.batched("image"), + num_img_tokens=MultiModalFieldConfig.batched("image", keep_on_cpu=True), input_audio_embeds=MultiModalFieldConfig.batched("audio"), ) diff --git a/vllm/model_executor/models/qwen2_5_omni_thinker.py b/vllm/model_executor/models/qwen2_5_omni_thinker.py index a512be850dde..149ab7bc15da 100644 --- a/vllm/model_executor/models/qwen2_5_omni_thinker.py +++ b/vllm/model_executor/models/qwen2_5_omni_thinker.py @@ -301,8 +301,12 @@ def _qwen2_5_omni_thinker_field_config(hf_inputs: Mapping[str, torch.Tensor]): input_audio_features=MultiModalFieldConfig.flat_from_sizes( "audio", audio_feature_lengths, dim=1 ), - feature_attention_mask=MultiModalFieldConfig.batched("audio"), - audio_feature_lengths=MultiModalFieldConfig.batched("audio"), + feature_attention_mask=MultiModalFieldConfig.batched( + "audio", keep_on_cpu=True + ), + audio_feature_lengths=MultiModalFieldConfig.batched( + "audio", keep_on_cpu=True + ), pixel_values=MultiModalFieldConfig.flat_from_sizes( "image", image_pixel_grid_sizes ), @@ -318,7 +322,9 @@ def _qwen2_5_omni_thinker_field_config(hf_inputs: Mapping[str, torch.Tensor]): ), video_grid_thw=MultiModalFieldConfig.batched("video", keep_on_cpu=True), second_per_grid_ts=MultiModalFieldConfig.batched("video", keep_on_cpu=True), - use_audio_in_video=MultiModalFieldConfig.shared("video", num_videos), + use_audio_in_video=MultiModalFieldConfig.shared( + "video", num_videos, keep_on_cpu=True + ), ) return _qwen2_5_omni_thinker_field_config @@ -1031,7 +1037,11 @@ def _process_audio_input( audio_input: Qwen2_5OmniAudioFeatureInputs, ) -> torch.Tensor: input_features = audio_input["input_features"] - audio_feature_lengths = audio_input["audio_feature_lengths"] + # audio_feature_lengths is keep_on_cpu; the audio tower derives + # device placement from feature_lens, so move it explicitly. + audio_feature_lengths = audio_input["audio_feature_lengths"].to( + input_features.device, non_blocking=True + ) audio_feat_lengths, audio_output_lengths = ( self.audio_tower._get_feat_extract_output_lengths(audio_feature_lengths) diff --git a/vllm/model_executor/models/qwen2_audio.py b/vllm/model_executor/models/qwen2_audio.py index 115a7f7f79b6..c989ec85b493 100644 --- a/vllm/model_executor/models/qwen2_audio.py +++ b/vllm/model_executor/models/qwen2_audio.py @@ -130,7 +130,7 @@ def _qwen2audio_field_config(hf_inputs: Mapping[str, torch.Tensor]): return dict( audio_embeds=MultiModalFieldConfig.batched("audio"), input_features=MultiModalFieldConfig.batched("audio"), - feature_attention_mask=MultiModalFieldConfig.batched("audio"), + feature_attention_mask=MultiModalFieldConfig.batched("audio", keep_on_cpu=True), ) diff --git a/vllm/model_executor/models/qwen3_asr.py b/vllm/model_executor/models/qwen3_asr.py index f5ec1a452629..328fcdbdac27 100644 --- a/vllm/model_executor/models/qwen3_asr.py +++ b/vllm/model_executor/models/qwen3_asr.py @@ -254,8 +254,8 @@ def _qwen3asr_field_config(hf_inputs: Mapping[str, torch.Tensor]): input_audio_features=MultiModalFieldConfig.flat_from_sizes( "audio", audio_feature_lengths, dim=1 ), - feature_attention_mask=MultiModalFieldConfig.batched("audio"), - audio_feature_lengths=MultiModalFieldConfig.batched("audio"), + feature_attention_mask=MultiModalFieldConfig.batched("audio", keep_on_cpu=True), + audio_feature_lengths=MultiModalFieldConfig.batched("audio", keep_on_cpu=True), ) @@ -459,7 +459,11 @@ def _process_audio_input( audio_input: Qwen2_5OmniAudioFeatureInputs, ) -> torch.Tensor: input_features = audio_input["input_features"] - audio_feature_lengths = audio_input["audio_feature_lengths"] + # audio_feature_lengths is keep_on_cpu; the audio tower derives + # device placement from feature_lens, so move it explicitly. + audio_feature_lengths = audio_input["audio_feature_lengths"].to( + input_features.device, non_blocking=True + ) audio_output_lengths = _get_feat_extract_output_lengths(audio_feature_lengths) diff --git a/vllm/model_executor/models/qwen3_omni_moe_thinker.py b/vllm/model_executor/models/qwen3_omni_moe_thinker.py index 48a16e97d1b4..0df29a7e713e 100755 --- a/vllm/model_executor/models/qwen3_omni_moe_thinker.py +++ b/vllm/model_executor/models/qwen3_omni_moe_thinker.py @@ -1553,7 +1553,11 @@ def _process_audio_input( audio_input: Qwen2_5OmniAudioFeatureInputs, ) -> tuple[torch.Tensor, ...]: input_features = audio_input["input_features"] - audio_feature_lengths = audio_input["audio_feature_lengths"] + # audio_feature_lengths is keep_on_cpu; the audio tower derives + # device placement from feature_lens, so move it explicitly. + audio_feature_lengths = audio_input["audio_feature_lengths"].to( + input_features.device, non_blocking=True + ) audio_output_lengths = _get_feat_extract_output_lengths(audio_feature_lengths) diff --git a/vllm/model_executor/models/step3_vl.py b/vllm/model_executor/models/step3_vl.py index f9b22705ed08..a06c221968cd 100644 --- a/vllm/model_executor/models/step3_vl.py +++ b/vllm/model_executor/models/step3_vl.py @@ -202,7 +202,7 @@ def _get_mm_fields_config( ), num_patches=MultiModalFieldConfig.batched("image", keep_on_cpu=True), patch_newline_mask=MultiModalFieldConfig.flat_from_sizes( - "image", num_patches + "image", num_patches, keep_on_cpu=True ), ) diff --git a/vllm/model_executor/models/transformers/multimodal.py b/vllm/model_executor/models/transformers/multimodal.py index f328a18a3ce0..e94490232c08 100644 --- a/vllm/model_executor/models/transformers/multimodal.py +++ b/vllm/model_executor/models/transformers/multimodal.py @@ -300,11 +300,17 @@ def _get_mm_fields_config( # Keep these as batched, as they always have batch size as first dim if "audio" in modalities: - mm_fields["num_audio_tokens"] = MultiModalFieldConfig.batched("audio") + mm_fields["num_audio_tokens"] = MultiModalFieldConfig.batched( + "audio", keep_on_cpu=True + ) if "image" in modalities: - mm_fields["image_grid_thw"] = MultiModalFieldConfig.batched("image") + mm_fields["image_grid_thw"] = MultiModalFieldConfig.batched( + "image", keep_on_cpu=True + ) # TODO: route to "video" once the video modality is supported - mm_fields["video_grid_thw"] = MultiModalFieldConfig.batched("image") + mm_fields["video_grid_thw"] = MultiModalFieldConfig.batched( + "image", keep_on_cpu=True + ) mm_fields["num_image_patches"] = MultiModalFieldConfig.batched( "image", keep_on_cpu=True ) @@ -755,6 +761,12 @@ def _process_image_input(self, **kwargs) -> list[torch.Tensor] | None: num_image_patches = kwargs.pop("num_image_patches") + # grid_thw fields are registered keep_on_cpu; restore the on-device + # placement that HF get_image_features implementations expect. + for key, value in kwargs.items(): + if isinstance(value, torch.Tensor): + kwargs[key] = value.to(pixel_values.device, non_blocking=True) + # The underlying HuggingFace `get_image_features` implementations # contain model-internal syncs (e.g. Idefics3 filters all-zero # padding images via boolean-mask indexing, LlavaOnevision diff --git a/vllm/model_executor/models/ultravox.py b/vllm/model_executor/models/ultravox.py index 551c2c5694b2..6bd41f300366 100644 --- a/vllm/model_executor/models/ultravox.py +++ b/vllm/model_executor/models/ultravox.py @@ -251,7 +251,7 @@ def _get_mm_fields_config( "audio", num_chunks, keep_on_cpu=True ), # num_chunks can convert audio_chunked to audio batch dimension - audio_num_chunks=MultiModalFieldConfig.batched("audio"), + audio_num_chunks=MultiModalFieldConfig.batched("audio", keep_on_cpu=True), audio_embeds=MultiModalFieldConfig.batched("audio"), ) diff --git a/vllm/models/inkling/common/mm_preprocess.py b/vllm/models/inkling/common/mm_preprocess.py index eb3c7b591ff5..6688d37cd82b 100644 --- a/vllm/models/inkling/common/mm_preprocess.py +++ b/vllm/models/inkling/common/mm_preprocess.py @@ -307,12 +307,12 @@ def _get_mm_fields_config( return dict( # Ragged per-image patches, grouped by num_patches. pixel_values=MultiModalFieldConfig.flat_from_sizes("image", num_patches), - num_patches=MultiModalFieldConfig.batched("image"), + num_patches=MultiModalFieldConfig.batched("image", keep_on_cpu=True), # Ragged per-audio frames, grouped by num_audio_tokens. input_audio_features=MultiModalFieldConfig.flat_from_sizes( "audio", num_audio_tokens ), - num_audio_tokens=MultiModalFieldConfig.batched("audio"), + num_audio_tokens=MultiModalFieldConfig.batched("audio", keep_on_cpu=True), ) def _get_prompt_updates( From e575b5f1a967b7ebb44eb65c4a3203e269cfc54d Mon Sep 17 00:00:00 2001 From: Qiming Zhang Date: Tue, 18 Aug 2026 20:30:47 -0700 Subject: [PATCH 127/839] [XPU][CI] fix hf runner (#52730) Signed-off-by: mayuyuace --- tests/conftest.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index bc6d9cb281e3..53d448b9b530 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -893,12 +893,15 @@ def predict(self, prompts: list[list[str]], *args, **kwargs) -> torch.Tensor: return self.model.predict(prompts, *args, convert_to_tensor=True, **kwargs) def __enter__(self): - if current_platform.is_rocm(): - # Record starting memory usage stats on ROCm so that we can wait for - # memory to roughly settle back below these levels on shutdown. This is - # helpful in cases where the HfRunner is initialized after significant GPU - # memory is already occupied, e.g. in + if current_platform.is_rocm() or current_platform.is_xpu(): + # Record starting memory usage stats on ROCm/XPU so that we can + # wait for memory to roughly settle back below these levels on + # shutdown. This is helpful in cases where the HfRunner is + # initialized after significant GPU memory is already occupied, + # e.g. in # tests/basic_correctness/test_basic_correctness.py::test_models_distributed + # where vllm worker processes are still alive and holding GPU + # memory when hf_runner.__exit__ is called. from tests.utils import ( get_physical_device_indices, record_gpu_memory_usage_stats, @@ -918,8 +921,8 @@ def __exit__(self, exc_type, exc_value, traceback): del self.model cleanup_dist_env_and_memory() - # ROCm frees VRAM lazily; wait so a runner started right after this HF - # model exits does not OOM on its startup memory guard. + # ROCm/XPU free VRAM lazily; wait so a runner started right after this + # HF model exits does not OOM on its startup memory guard. wait_for_memory_to_settle( threshold_ratio=getattr(self, "threshold_ratios", None) ) From 8e46accab22bc864f3ddbe8276d7c7cd2eb71239 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Tue, 18 Aug 2026 23:32:41 -0400 Subject: [PATCH 128/839] [Attention] Vectorize sparse MLA mask loads (#52217) Signed-off-by: Matthew Bonanni Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/v1/attention/test_sparse_mla_mask.py | 22 ++++++++++++++++ .../layers/attention/sparse_mla_attention.py | 9 ++++--- .../layers/attention/sparse_mla_mask.py | 26 ++++++++++++++----- 3 files changed, 46 insertions(+), 11 deletions(-) diff --git a/tests/v1/attention/test_sparse_mla_mask.py b/tests/v1/attention/test_sparse_mla_mask.py index 3ecc179be4b4..20c7631e1d31 100644 --- a/tests/v1/attention/test_sparse_mla_mask.py +++ b/tests/v1/attention/test_sparse_mla_mask.py @@ -8,6 +8,7 @@ from vllm.model_executor.layers.attention.sparse_mla_attention import ( _build_topk_mask, + _topk_mask_shape, ) @@ -27,3 +28,24 @@ def test_build_topk_mask_single_request_matches_generic_path() -> None: torch.testing.assert_close(single_req[0, 0], generic[0, 0]) torch.testing.assert_close(single_req[0, 1], generic[1, 0]) + + +def test_topk_mask_rows_are_aligned_for_vectorized_loads() -> None: + assert _topk_mask_shape(2, 129, 129) == (2, 256, 8) + assert _topk_mask_shape(2, 129, 129, reserve_key_starts_word=True) == ( + 2, + 256, + 8, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_build_topk_mask_preserves_aligned_row_storage() -> None: + shape = _topk_mask_shape(1, 1, 129) + out = torch.zeros(shape, dtype=torch.int32, device="cuda") + topk = torch.tensor([[0, 128]], dtype=torch.int32, device="cuda") + + mask = _build_topk_mask([topk], [1], 1, 129, out) + + assert mask.shape == (1, 1, 8) + assert mask[0, 0].tolist() == [1, 0, 0, 0, 1, 0, 0, 0] diff --git a/vllm/model_executor/layers/attention/sparse_mla_attention.py b/vllm/model_executor/layers/attention/sparse_mla_attention.py index a619b7403540..85e09cd5d49b 100644 --- a/vllm/model_executor/layers/attention/sparse_mla_attention.py +++ b/vllm/model_executor/layers/attention/sparse_mla_attention.py @@ -58,6 +58,7 @@ def _topk_mask_shape( tile_m = 128 if max_query_len <= 128 else 256 padded_q_len = triton.cdiv(max_query_len, tile_m) * tile_m num_words = triton.cdiv(max_key_len, 32) + int(reserve_key_starts_word) + num_words = triton.cdiv(num_words, 4) * 4 return batch_size, padded_q_len, num_words @@ -432,12 +433,12 @@ def _build_topk_mask( max_seq_len: int, out: torch.Tensor, ) -> torch.Tensor: - """Build a bit-packed top-k mask into ``out[:B, :max_Q, :num_words]``.""" + """Build a bit-packed top-k mask while preserving padded row storage.""" batch_size = len(q_lens) num_words = (max_seq_len + 31) // 32 total_rows = batch_size * max_q_len if total_rows == 0: - return out[:batch_size, :max_q_len, :num_words] + return out[:batch_size, :max_q_len] total_q = sum(q_lens) mask_row_stride = out.stride(-2) @@ -457,7 +458,7 @@ def _build_topk_mask( BLOCK_TOPK=triton.next_power_of_2(num_topk), BLOCK_WORDS=block_words, ) - return out[:1, :max_q_len, :num_words] + return out[:1, :max_q_len] topk_packed = torch.cat(topk_indices_per_req, dim=0) num_topk = topk_packed.shape[1] @@ -478,7 +479,7 @@ def _build_topk_mask( BLOCK_TOPK=triton.next_power_of_2(num_topk), BLOCK_WORDS=block_words, ) - return out[:batch_size, :max_q_len, :num_words] + return out[:batch_size, :max_q_len] class SparseMLACommonImpl(MLACommonBaseImpl[T], Generic[T]): diff --git a/vllm/model_executor/layers/attention/sparse_mla_mask.py b/vllm/model_executor/layers/attention/sparse_mla_mask.py index 1cd151034c2e..8659786bd15b 100644 --- a/vllm/model_executor/layers/attention/sparse_mla_mask.py +++ b/vllm/model_executor/layers/attention/sparse_mla_mask.py @@ -20,15 +20,27 @@ def dense_mask_mod( batch_idx = utils.ssa_to_scalar(batch) q_idx = utils.ssa_to_scalar(q_idx) kv_idx = utils.ssa_to_scalar(kv_idx) - word_idx = kv_idx >> 5 - bit_idx = cutlass.Uint32(kv_idx & 31) - word = dense_mask[batch_idx, q_idx, word_idx] - result = cute.make_rmem_tensor(1, dtype=cutlass.Uint32) - result[0] = utils.shr_u32(cutlass.Uint32(word), bit_idx) - return result.load() + batch_stride, query_stride, _ = dense_mask.stride + aligned_mask = cute.make_tensor( + dense_mask.iterator, + cute.make_layout( + dense_mask.shape, + stride=( + cute.assume(batch_stride, divby=4), + cute.assume(query_stride, divby=4), + 1, + ), + ), + ) + mask_row = aligned_mask[batch_idx, q_idx, None] + mask_chunks = cute.flat_divide(mask_row, (4,)) + mask_chunk = mask_chunks[None, (kv_idx >> 5) >> 2] + loaded = cute.make_rmem_tensor_like(mask_chunk) + cute.autovec_copy(mask_chunk, loaded) + return cute.recast_tensor(loaded, cutlass.Uint32).load() -dense_mask_mod.__vec_size__ = 32 +dense_mask_mod.__vec_size__ = 128 @cute.jit From 31305736307ecdfe3398b643240770651b523a84 Mon Sep 17 00:00:00 2001 From: Davis Wertheimer Date: Tue, 18 Aug 2026 23:32:46 -0400 Subject: [PATCH 129/839] [Model] Add GraniteSWA and GraniteMoeSWA via existing Granite (#52706) Signed-off-by: Davis Wertheimer Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- docs/models/supported_models.md | 2 + .../language/generation/test_granite.py | 49 ++++- tests/models/registry.py | 6 + vllm/model_executor/models/granite.py | 76 ++++++- vllm/model_executor/models/granitemoe.py | 207 +++++------------- .../model_executor/models/granitemoeshared.py | 51 +---- vllm/model_executor/models/registry.py | 2 + 7 files changed, 186 insertions(+), 207 deletions(-) diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 0ed82796290b..f9c70adec763 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -405,6 +405,8 @@ th { | `GraniteMoeForCausalLM` | Granite 3.0 MoE, PowerMoE | `ibm-granite/granite-3.0-1b-a400m-base`, `ibm-granite/granite-3.0-3b-a800m-instruct`, `ibm/PowerMoE-3b`, etc. | ✅︎ | ✅︎ | | `GraniteMoeHybridForCausalLM` | Granite 4.0 MoE Hybrid | `ibm-granite/granite-4.0-tiny-preview`, etc. | ✅︎ | ✅︎ | | `GraniteMoeSharedForCausalLM` | Granite MoE Shared | `ibm-research/moe-7b-1b-active-shared-experts` (test model) | ✅︎ | ✅︎ | +| `GraniteMoeSWAForCausalLM` | Granite MoE SWA | `ibm-granite/granite-swash-3b-a600m` | ✅︎ | ✅︎ | +| `GraniteSWAForCausalLM` | Granite SWA | `ibm-granite/granite-swash-2b` | ✅︎ | ✅︎ | | `GritLM` | GritLM | `parasail-ai/GritLM-7B-vllm`. | ✅︎ | ✅︎ | | `HrmTextForCausalLM` | HRM-Text | `sapientinc/HRM-Text-1B`, etc. | | | | `HunYuanDenseV1ForCausalLM` | Hunyuan Dense | `tencent/Hunyuan-7B-Instruct` | ✅︎ | ✅︎ | diff --git a/tests/models/language/generation/test_granite.py b/tests/models/language/generation/test_granite.py index c0498b2f7de1..207e798fef82 100644 --- a/tests/models/language/generation/test_granite.py +++ b/tests/models/language/generation/test_granite.py @@ -1,14 +1,20 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest +from transformers import GraniteConfig -from ...utils import check_logprobs_close +from vllm.model_executor.models.granite import granite_layer_attn_params -MODELS = [ +from ...utils import check_logprobs_close, check_transformers_version + +# model -> minimum transformers version, or None if unconstrained +MODELS = { # TODO(sang): Sliding window should be tested separately. - "ibm/PowerLM-3b", - "ibm/PowerMoE-3b", -] + "ibm/PowerLM-3b": None, + "ibm/PowerMoE-3b": None, + "ibm-granite/granite-swash-2b": "5.15.1", + "ibm-granite/granite-swash-3b-a600m": "5.15.1", +} @pytest.mark.parametrize("model", MODELS) @@ -25,6 +31,8 @@ def test_models( max_tokens: int, num_logprobs: int, ) -> None: + check_transformers_version(model, min_transformers_version=MODELS[model]) + with hf_runner(model, dtype=dtype) as hf_model: hf_outputs = hf_model.generate_greedy_logprobs_limit( example_prompts, max_tokens, num_logprobs @@ -40,3 +48,34 @@ def test_models( name_0="hf", name_1="vllm", ) + + +def test_granite_swa_features_are_off_without_swa_config(): + """A plain Granite config must not pick up sliding windows or sinks.""" + config = GraniteConfig(num_hidden_layers=4) + theta = config.rope_parameters["rope_theta"] + + assert [granite_layer_attn_params(config, i) for i in range(4)] == [ + (None, theta, False) + ] * 4 + + +@pytest.mark.parametrize( + "attention_sinks, expected_sink", [(None, True), (False, False)] +) +def test_granite_swa_features_resolve_per_layer(attention_sinks, expected_sink): + """`layer_types`/`layer_rope_theta` apply per layer; theta 0 means NoPE.""" + kwargs = {} if attention_sinks is None else {"attention_sinks": attention_sinks} + config = GraniteConfig( + num_hidden_layers=3, + sliding_window=128, + layer_types=["full_attention", "sliding_attention", "sliding_attention"], + layer_rope_theta=[10000.0, 0.0, 1000000.0], + **kwargs, + ) + + assert [granite_layer_attn_params(config, i) for i in range(3)] == [ + (None, 10000.0, expected_sink), + (128, 0.0, expected_sink), + (128, 1000000.0, expected_sink), + ] diff --git a/tests/models/registry.py b/tests/models/registry.py index 16604ec7ff83..64efae5261be 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -320,6 +320,12 @@ def check_available_online( "GraniteMoeSharedForCausalLM": _HfExamplesInfo( "ibm-research/moe-7b-1b-active-shared-experts" ), + "GraniteMoeSWAForCausalLM": _HfExamplesInfo( + "ibm-granite/granite-swash-3b-a600m", min_transformers_version="5.15.1" + ), + "GraniteSWAForCausalLM": _HfExamplesInfo( + "ibm-granite/granite-swash-2b", min_transformers_version="5.15.1" + ), "HrmTextForCausalLM": _HfExamplesInfo( "sapientinc/HRM-Text-1B", min_transformers_version="5.9.0", diff --git a/vllm/model_executor/models/granite.py b/vllm/model_executor/models/granite.py index a0bc8350b43a..df813fd1e918 100644 --- a/vllm/model_executor/models/granite.py +++ b/vllm/model_executor/models/granite.py @@ -22,7 +22,13 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Inference-only IBM Granite model compatible with HuggingFace weights.""" +"""Inference-only IBM Granite model compatible with HuggingFace weights. + +Also serves the `granite_swa` checkpoints (`GraniteSWAForCausalLM`), supporting +three additional features: per-layer sliding window attention (`layer_types`), a +learnable per-head attention sink (`self_attn.sinks`), and a per-layer RoPE base +(`layer_rope_theta`, with 0 for NoPE). +""" from collections.abc import Iterable from itertools import islice @@ -30,6 +36,7 @@ import torch from torch import nn from transformers import GraniteConfig +from transformers.configuration_utils import PretrainedConfig from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig @@ -49,6 +56,8 @@ ParallelLMHead, VocabParallelEmbedding, ) +from vllm.model_executor.model_loader.weight_utils import sharded_weight_loader +from vllm.model_executor.utils import set_weight_attrs from vllm.sequence import IntermediateTensors from .interfaces import SupportsLoRA, SupportsPP, SupportsQuant @@ -56,11 +65,44 @@ AutoWeightsLoader, PPMissingLayer, WeightsMapper, + extract_layer_index, make_layers, maybe_prefix, ) +def granite_layer_attn_params( + config: PretrainedConfig, layer_idx: int +) -> tuple[int | None, float, bool]: + """Resolve one layer's sliding window, RoPE base and sink usage. + + Plain Granite configs carry no SWA fields and fall back to full + attention, global RoPE base and no sink. HF SWA checkpoints use + sinks without a dedicated flag, so assume true when `layer_types` + is used, and allow `attention_sinks` to override that decision. + + Returns: + Sliding window size (`None` for full attention), RoPE base theta (`0` + for NoPE), and attention sink presence/absence. + """ + layer_types = getattr(config, "layer_types", None) + sliding_window = ( + config.sliding_window + if layer_types is not None and layer_types[layer_idx] == "sliding_attention" + else None + ) + + layer_rope_theta = getattr(config, "layer_rope_theta", None) + rope_theta = ( + layer_rope_theta[layer_idx] + if layer_rope_theta is not None + else config.rope_parameters["rope_theta"] + ) + + has_sink = getattr(config, "attention_sinks", layer_types is not None) + return sliding_window, rope_theta, has_sink + + class GraniteMLP(nn.Module): def __init__( self, @@ -154,11 +196,25 @@ def __init__( prefix=f"{prefix}.o_proj", ) - self.rotary_emb = get_rope( - self.head_dim, - max_position=max_position_embeddings, - rope_parameters=config.rope_parameters, + sliding_window, rope_theta, has_sink = granite_layer_attn_params( + config, extract_layer_index(prefix) ) + + self.use_rope = rope_theta != 0 + if self.use_rope: + self.rotary_emb = get_rope( + self.head_dim, + max_position=max_position_embeddings, + rope_parameters={**config.rope_parameters, "rope_theta": rope_theta}, + ) + + # Per-head sink, applied in backend as extra logit in the softmax denom + if has_sink: + self.sinks = nn.Parameter(torch.empty(self.num_heads), requires_grad=False) + set_weight_attrs(self.sinks, {"weight_loader": sharded_weight_loader(0)}) + else: + self.sinks = None + self.attn = Attention( self.num_heads, self.head_dim, @@ -166,7 +222,9 @@ def __init__( num_kv_heads=self.num_kv_heads, cache_config=cache_config, quant_config=quant_config, + per_layer_sliding_window=sliding_window, prefix=f"{prefix}.attn", + sinks=self.sinks, ) def forward( @@ -176,7 +234,8 @@ def forward( ) -> torch.Tensor: qkv, _ = self.qkv_proj(hidden_states) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - q, k = self.rotary_emb(positions, q, k) + if self.use_rope: + q, k = self.rotary_emb(positions, q, k) attn_output = self.attn(q, k, v) output, _ = self.o_proj(attn_output) return output @@ -413,8 +472,5 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # processed with quantization, LoRA, fine-tuning, etc. skip_prefixes = ["lm_head."] if self.config.tie_word_embeddings else None - loader = AutoWeightsLoader( - self, - skip_prefixes=skip_prefixes, - ) + loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes) return loader.load_weights(weights) diff --git a/vllm/model_executor/models/granitemoe.py b/vllm/model_executor/models/granitemoe.py index e1219e3337dd..d775de50ca4c 100644 --- a/vllm/model_executor/models/granitemoe.py +++ b/vllm/model_executor/models/granitemoe.py @@ -26,10 +26,10 @@ from collections.abc import Iterable from itertools import islice -from typing import Any import torch from torch import nn +from transformers.configuration_utils import PretrainedConfig from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig @@ -39,10 +39,7 @@ tensor_model_parallel_all_gather, ) from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.fused_moe import ( - FusedMoEFactory, - fused_moe_make_expert_params_mapping, -) +from vllm.model_executor.layers.fused_moe import FusedMoEFactory from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( QKVParallelLinear, @@ -56,15 +53,20 @@ ParallelLMHead, VocabParallelEmbedding, ) -from vllm.model_executor.model_loader.weight_utils import ( - default_weight_loader, - maybe_remap_kv_scale_name, -) +from vllm.model_executor.model_loader.weight_utils import sharded_weight_loader from vllm.model_executor.models.utils import sequence_parallel_chunk +from vllm.model_executor.utils import set_weight_attrs from vllm.sequence import IntermediateTensors +from .granite import granite_layer_attn_params from .interfaces import SupportsLoRA, SupportsPP -from .utils import AutoWeightsLoader, is_pp_missing_parameter, make_layers, maybe_prefix +from .utils import ( + AutoWeightsLoader, + WeightsMapper, + extract_layer_index, + make_layers, + maybe_prefix, +) class GraniteMoeMoE(nn.Module): @@ -139,11 +141,11 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: class GraniteMoeAttention(nn.Module): def __init__( self, + config: PretrainedConfig, hidden_size: int, num_heads: int, num_kv_heads: int, max_position: int = 4096 * 32, - rope_parameters: dict[str, Any] | None = None, cache_config: CacheConfig | None = None, quant_config: QuantizationConfig | None = None, attention_multiplier: float | None = None, @@ -190,12 +192,25 @@ def __init__( quant_config=quant_config, prefix=f"{prefix}.o_proj", ) - self.rotary_emb = get_rope( - self.head_dim, - max_position=max_position, - rope_parameters=rope_parameters, - is_neox_style=True, + sliding_window, rope_theta, has_sink = granite_layer_attn_params( + config, extract_layer_index(prefix) ) + + self.use_rope = rope_theta != 0 + if self.use_rope: + self.rotary_emb = get_rope( + self.head_dim, + max_position=max_position, + rope_parameters={**config.rope_parameters, "rope_theta": rope_theta}, + is_neox_style=True, + ) + + if has_sink: + self.sinks = nn.Parameter(torch.empty(self.num_heads), requires_grad=False) + set_weight_attrs(self.sinks, {"weight_loader": sharded_weight_loader(0)}) + else: + self.sinks = None + self.attn = Attention( self.num_heads, self.head_dim, @@ -203,7 +218,9 @@ def __init__( num_kv_heads=self.num_kv_heads, cache_config=cache_config, quant_config=quant_config, + per_layer_sliding_window=sliding_window, prefix=f"{prefix}.attn", + sinks=self.sinks, ) def forward( @@ -213,7 +230,8 @@ def forward( ) -> torch.Tensor: qkv, _ = self.qkv_proj(hidden_states) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - q, k = self.rotary_emb(positions, q, k) + if self.use_rope: + q, k = self.rotary_emb(positions, q, k) attn_output = self.attn(q, k, v) output, _ = self.o_proj(attn_output) return output @@ -234,11 +252,11 @@ def __init__( self.hidden_size = config.hidden_size self.self_attn = GraniteMoeAttention( + config=config, hidden_size=self.hidden_size, num_heads=config.num_attention_heads, max_position=config.max_position_embeddings, num_kv_heads=config.num_key_value_heads, - rope_parameters=config.rope_parameters, cache_config=cache_config, quant_config=quant_config, prefix=f"{prefix}.self_attn", @@ -284,6 +302,23 @@ def forward( @support_torch_compile class GraniteMoeModel(nn.Module): + hf_to_vllm_mapper: WeightsMapper = WeightsMapper( + orig_to_new_suffix={ + # Legacy names to new names + "moe.input_linear.weight": "moe.experts.gate_up_proj", + "moe.output_linear.weight": "moe.experts.down_proj", + ".router.layer.weight": ".gate.weight", + # Checkpoint name to vLLM name + ".router.weight": ".gate.weight", + }, + orig_to_new_stacked={ + # weight_name: (param_name, shard_id) + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + }, + ) + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() @@ -339,135 +374,9 @@ def forward( hidden_states = self.norm(hidden_states) return hidden_states - def _load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - """ - This function is copied from `MixtralModel.load_weights`, mainly to - decouple from mixtral, avoiding impact on support like BNB - quantization. - """ - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - ("qkv_proj", "q_proj", "q"), - ("qkv_proj", "k_proj", "k"), - ("qkv_proj", "v_proj", "v"), - ] - - # Params for weights, fp8 weight scales, fp8 activation scales - # (param_name, weight_name, expert_id, shard_id) - expert_params_mapping = fused_moe_make_expert_params_mapping( - self, - ckpt_gate_proj_name="w1", - ckpt_down_proj_name="w2", - ckpt_up_proj_name="w3", - num_experts=self.config.num_local_experts, - ) - - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - # Skip loading extra bias for GPTQ models. - if ( - name.endswith(".bias") or name.endswith("_bias") - ) and name not in params_dict: - continue - # Skip layers on other devices. - if is_pp_missing_parameter(name, self): - continue - if name.endswith("scale"): - # Remapping the name of FP8 kv-scale. - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - break - else: - for mapping in expert_params_mapping: - param_name, weight_name, expert_id, shard_id = mapping - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - # Skip layers on other devices. - if is_pp_missing_parameter(name, self): - continue - if ( - name.endswith(".bias") or name.endswith("_bias") - ) and name not in params_dict: - continue - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader( - param, - loaded_weight, - name, - shard_id=shard_id, - expert_id=expert_id, - ) - break - else: - # Skip loading extra bias for GPTQ models. - if ( - name.endswith(".bias") or name.endswith("_bias") - ) and name not in params_dict: - continue - # Skip layers on other devices. - if is_pp_missing_parameter(name, self): - continue - # Remapping the name of FP8 kv-scale. - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue - - param = params_dict[name] - weight_loader = getattr( - param, "weight_loader", default_weight_loader - ) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - new_weights = {} - for n, p in weights: - if n.endswith(".block_sparse_moe.input_linear.weight"): - for e in range(p.size(0)): - w1_name = n.replace( - ".block_sparse_moe.input_linear.weight", - f".block_sparse_moe.experts.{e}.w1.weight", - ) - w3_name = n.replace( - ".block_sparse_moe.input_linear.weight", - f".block_sparse_moe.experts.{e}.w3.weight", - ) - w1_param, w3_param = p[e].chunk(2, dim=0) - assert w1_name not in new_weights - assert w3_name not in new_weights - new_weights[w1_name] = w1_param - new_weights[w3_name] = w3_param - elif n.endswith(".block_sparse_moe.output_linear.weight"): - for e in range(p.size(0)): - w2_name = n.replace( - ".block_sparse_moe.output_linear.weight", - f".block_sparse_moe.experts.{e}.w2.weight", - ) - w2_param = p[e] - assert w2_name not in new_weights - new_weights[w2_name] = w2_param - elif n.endswith(".block_sparse_moe.router.layer.weight"): - gate_name = n.replace( - ".block_sparse_moe.router.layer.weight", - ".block_sparse_moe.gate.weight", - ) - assert gate_name not in new_weights - new_weights[gate_name] = p - else: - new_weights[n] = p - return self._load_weights(new_weights.items()) + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) class GraniteMoeForCausalLM(nn.Module, SupportsLoRA, SupportsPP): @@ -543,8 +452,6 @@ def make_empty_intermediate_tensors( ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + skip_prefixes = ["lm_head."] if self.config.tie_word_embeddings else None + loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes) return loader.load_weights(weights) diff --git a/vllm/model_executor/models/granitemoeshared.py b/vllm/model_executor/models/granitemoeshared.py index 603afc60d683..ecd942da4042 100644 --- a/vllm/model_executor/models/granitemoeshared.py +++ b/vllm/model_executor/models/granitemoeshared.py @@ -4,6 +4,10 @@ The architecture is the same as granitemoe but with the addition of shared experts. + +Also serves the `granitemoe_swa` checkpoints (`GraniteMoeSWAForCausalLM`), which +add the same per-layer sliding window, attention sink and per-layer RoPE support +as `granite_swa` (see `granite.py`). """ from collections.abc import Iterable @@ -85,11 +89,11 @@ def __init__( super().__init__() self.hidden_size = config.hidden_size self.self_attn = GraniteMoeAttention( + config=config, hidden_size=self.hidden_size, num_heads=config.num_attention_heads, max_position=config.max_position_embeddings, num_kv_heads=config.num_key_value_heads, - rope_parameters=config.rope_parameters, cache_config=cache_config, quant_config=quant_config, prefix=f"{prefix}.self_attn", @@ -207,43 +211,8 @@ def forward( hidden_states = self.norm(hidden_states) return hidden_states - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - new_weights = {} - for n, p in weights: - if n.endswith(".block_sparse_moe.input_linear.weight"): - for e in range(p.size(0)): - w1_name = n.replace( - ".block_sparse_moe.input_linear.weight", - f".block_sparse_moe.experts.{e}.w1.weight", - ) - w3_name = n.replace( - ".block_sparse_moe.input_linear.weight", - f".block_sparse_moe.experts.{e}.w3.weight", - ) - w1_param, w3_param = p[e].chunk(2, dim=0) - assert w1_name not in new_weights - assert w3_name not in new_weights - new_weights[w1_name] = w1_param - new_weights[w3_name] = w3_param - elif n.endswith(".block_sparse_moe.output_linear.weight"): - for e in range(p.size(0)): - w2_name = n.replace( - ".block_sparse_moe.output_linear.weight", - f".block_sparse_moe.experts.{e}.w2.weight", - ) - w2_param = p[e] - assert w2_name not in new_weights - new_weights[w2_name] = w2_param - elif n.endswith(".block_sparse_moe.router.layer.weight"): - gate_name = n.replace( - ".block_sparse_moe.router.layer.weight", - ".block_sparse_moe.gate.weight", - ) - assert gate_name not in new_weights - new_weights[gate_name] = p - else: - new_weights[n] = p - return GraniteMoeModel._load_weights(self, new_weights.items()) + hf_to_vllm_mapper = GraniteMoeModel.hf_to_vllm_mapper + load_weights = GraniteMoeModel.load_weights class GraniteMoeSharedForCausalLM(nn.Module, SupportsLoRA, SupportsPP): @@ -320,8 +289,6 @@ def make_empty_intermediate_tensors( ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + skip_prefixes = ["lm_head."] if self.config.tie_word_embeddings else None + loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes) return loader.load_weights(weights) diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 1e4bf6731c96..95877c319d1f 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -124,6 +124,8 @@ "GraniteMoeForCausalLM": ("granitemoe", "GraniteMoeForCausalLM"), "GraniteMoeHybridForCausalLM": ("granitemoehybrid", "GraniteMoeHybridForCausalLM"), "GraniteMoeSharedForCausalLM": ("granitemoeshared", "GraniteMoeSharedForCausalLM"), + "GraniteMoeSWAForCausalLM": ("granitemoeshared", "GraniteMoeSharedForCausalLM"), + "GraniteSWAForCausalLM": ("granite", "GraniteForCausalLM"), "GritLM": ("gritlm", "GritLM"), "HrmTextForCausalLM": ("hrm_text", "HrmTextForCausalLM"), "HunYuanMoEV1ForCausalLM": ("hunyuan_v1", "HunYuanMoEV1ForCausalLM"), From 63ff748f657af0e2a95a712fd2cc994f4299ff8c Mon Sep 17 00:00:00 2001 From: drakosha Date: Wed, 19 Aug 2026 07:02:09 +0300 Subject: [PATCH 130/839] [Attention][MLA] FlashMLA sparse: DCP on the fp8_ds_mla mixed-batch path + MTP (#46514) Signed-off-by: Mikhail Kostryukov Co-authored-by: Claude Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../v1/attention/test_sparse_mla_backends.py | 154 +++++++++++++++++ vllm/v1/attention/backends/mla/flashmla.py | 12 ++ .../attention/backends/mla/flashmla_sparse.py | 163 ++++++++++++++---- 3 files changed, 295 insertions(+), 34 deletions(-) diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index dec43fd87438..ac9a87d08c74 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -1492,6 +1492,8 @@ def run_kernel(**kwargs): topk_indices_buffer=topk_indices, num_heads=2, kv_lora_rank=1, + dcp_world_size=1, + need_to_return_lse_for_decode=False, _fp8_flash_mla_kernel=run_kernel, ) impl._forward_fp8_kv_mixed_batch = MethodType( @@ -1512,3 +1514,155 @@ def run_kernel(**kwargs): assert kernel_q_shapes == [(1, num_decode_tokens, 2, 3)] assert output.shape == (num_decode_tokens, 2, 1) assert lse is None + + +def _build_sparse_dcp_vllm_config( + local_heads: int, + dcp_world_size: int, + comm_backend: str = "ag_rs", +): + """Minimal sparse-MLA VllmConfig for the FlashMLASparse DCP head-envelope + guard. TP is simulated by mocking ``get_num_attention_heads`` to return the + per-rank head count, as the decode-correctness test above does. + """ + kv_lora_rank = 512 + qk_nope_head_dim = 128 + qk_rope_head_dim = 64 + v_head_dim = 128 + head_size = kv_lora_rank + qk_rope_head_dim + topk_tokens = 128 + + vllm_config = create_vllm_config( + model_name="deepseek-ai/DeepSeek-V2-Lite-Chat", + tensor_parallel_size=1, + max_model_len=4096, + block_size=64, + hf_config_override={ + "index_topk": topk_tokens, + "attn_module_list_cfg": [{"topk_tokens": topk_tokens}], + }, + ) + model_config = vllm_config.model_config + model_config.dtype = torch.bfloat16 + model_config.hf_text_config = SimpleNamespace( + q_lora_rank=None, + kv_lora_rank=kv_lora_rank, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=v_head_dim, + model_type="deepseek_v2", + ) + model_config.get_num_attention_heads = MethodType( + lambda self, parallel_config: local_heads, model_config + ) + model_config.get_num_kv_heads = MethodType( + lambda self, parallel_config: 1, model_config + ) + model_config.get_head_size = MethodType(lambda self: head_size, model_config) + model_config.get_sliding_window = MethodType(lambda self: None, model_config) + + vllm_config.cache_config.cache_dtype = "fp8_ds_mla" + vllm_config.parallel_config.decode_context_parallel_size = dcp_world_size + vllm_config.parallel_config.dcp_comm_backend = comm_backend + # The base builder clones the layer's dense-MHA prefill backend from + # static_forward_context; the guard tests never run prefill. + vllm_config.compilation_config.static_forward_context["placeholder"] = ( + SimpleNamespace(prefill_backend=None) + ) + return vllm_config + + +@pytest.mark.skipif( + torch.cuda.get_device_capability() < (9, 0), + reason="FlashMLASparseBackend requires CUDA 9.0 or higher", +) +@pytest.mark.parametrize( + "local_heads,dcp_world_size,should_raise", + [ + (16, 8, True), + (24, 4, True), + (16, 4, False), + (16, 1, False), + ], +) +def test_fp8_dcp_head_envelope_guard(local_heads, dcp_world_size, should_raise): + """The fp8 decode envelope (head padding + tile-scheduler metadata) is + sized from the local head count while the kernel runs on the DCP-gathered + heads, so the builder must reject configs where the two pad differently. + """ + device = torch.device(DEVICE_TYPE) + vllm_config = _build_sparse_dcp_vllm_config(local_heads, dcp_world_size) + kv_cache_spec = create_standard_kv_cache_spec(vllm_config) + builder_cls = FlashMLASparseBackend.get_builder_cls() + + if should_raise: + with pytest.raises(NotImplementedError, match="envelope"): + builder_cls(kv_cache_spec, ["placeholder"], vllm_config, device) + else: + builder = builder_cls(kv_cache_spec, ["placeholder"], vllm_config, device) + gathered_heads = local_heads * dcp_world_size + local_pad = 64 if local_heads <= 64 else 128 + gathered_pad = 64 if gathered_heads <= 64 else 128 + assert builder.fp8_decode_padded_heads == local_pad + assert local_pad == gathered_pad + + +def test_fp8_mixed_batch_dcp_neutralizes_empty_rows(monkeypatch): + """A decode row whose top-k shard holds no local candidates (all -1) has + undefined kernel out/lse; it must come back as (0, -inf), the identity of + the cross-rank LSE merge, or a NaN would survive the merge even at zero + weight (0 * NaN = NaN).""" + num_tokens, num_heads, head_dim = 3, 2, 3 + q = torch.empty(num_tokens, num_heads, head_dim, device=DEVICE_TYPE) + local_indices = torch.tensor( + [[0, 1, -1, -1], [-1, -1, -1, -1], [2, -1, 3, -1]], + dtype=torch.int32, + device=DEVICE_TYPE, + ) + + monkeypatch.setattr( + "vllm.v1.attention.backends.mla.flashmla_sparse." + "triton_filter_and_convert_dcp_index", + lambda *args, **kwargs: local_indices, + ) + + def run_kernel(**kwargs): + out = torch.full( + (1, num_tokens, num_heads, 1), float("nan"), device=DEVICE_TYPE + ) + lse = torch.full((1, num_heads, num_tokens), float("nan"), device=DEVICE_TYPE) + for token_id in (0, 2): # rows with local candidates get real values + out[0, token_id] = float(token_id + 1) + lse[0, :, token_id] = float(token_id + 1) + return out, lse + + metadata = SimpleNamespace( + fp8_extra_metadata=FlashMLASparseMetadata.FP8KernelMetadata( + scheduler_metadata=object(), # type: ignore[arg-type] + dummy_block_table=torch.empty(1, 1, dtype=torch.int32, device=DEVICE_TYPE), + cache_lens=torch.empty(1, dtype=torch.int32, device=DEVICE_TYPE), + ), + req_id_per_token=torch.empty(num_tokens, dtype=torch.int32, device=DEVICE_TYPE), + block_table=torch.empty(1, 1, dtype=torch.int32, device=DEVICE_TYPE), + block_size=64, + cp_kv_cache_interleave_size=1, + ) + impl = SimpleNamespace( + dcp_world_size=2, + dcp_rank=0, + need_to_return_lse_for_decode=True, + _fp8_flash_mla_kernel=run_kernel, + ) + + out, lse = FlashMLASparseImpl._forward_fp8_kv_mixed_batch( + impl, q, torch.empty(0, device=DEVICE_TYPE), local_indices, metadata + ) + + assert torch.equal(out[1], torch.zeros_like(out[1])) + assert torch.isneginf(lse[1]).all() + for token_id in (0, 2): + assert torch.equal(out[token_id], torch.full_like(out[token_id], token_id + 1)) + assert torch.equal(lse[token_id], torch.full_like(lse[token_id], token_id + 1)) + assert out.is_contiguous() + assert not out.isnan().any() + assert not lse.isnan().any() diff --git a/vllm/v1/attention/backends/mla/flashmla.py b/vllm/v1/attention/backends/mla/flashmla.py index bb6efe59c8c1..a8c7b2555d4e 100644 --- a/vllm/v1/attention/backends/mla/flashmla.py +++ b/vllm/v1/attention/backends/mla/flashmla.py @@ -343,4 +343,16 @@ def forward_mqa( o = reshape_attn_output_for_spec_decode(o) + if self.need_to_return_lse_for_decode: + # FlashMLA returns LSE as [batch, heads, seq_len]; the DCP reducer + # consumes [tokens, heads]. Flattening matters under spec-decode, + # where seq_len > 1. Only DCP consumes lse, so skip the copy + # otherwise. + num_decodes, q_num_heads, seq_len = lse.shape + lse = ( + lse.permute(0, 2, 1) + .reshape(num_decodes * seq_len, q_num_heads) + .contiguous() + ) + return o, lse diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index e8be7ad5bb39..7b0be952ca83 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -28,6 +28,7 @@ ) from vllm.v1.attention.backends.mla.sparse_utils import ( triton_convert_req_index_to_global_index, + triton_filter_and_convert_dcp_index, ) from vllm.v1.attention.backends.utils import ( reshape_attn_output_for_spec_decode, @@ -254,9 +255,12 @@ def __init__( ) else: threshold = {16: 128, 32: 128, 64: 256, 128: 256}.get(num_q_heads, 256) + # Varlen decodes are safe under DCP: causality comes from the + # indexer's top-k indices, not from the kernel metadata. self._init_reorder_batch_threshold( threshold, supports_spec_as_decode=True, + supports_dcp_with_varlen=(parallel_config.cp_kv_cache_interleave_size == 1), ) sm_count = num_compute_units(device.index) @@ -313,6 +317,42 @@ def __init__( device=device, ) + self.fp8_use_mixed_batch = self.num_heads < MIN_HEADS_FOR_BF16_PREFILL + + if parallel_config.decode_context_parallel_size > 1: + if parallel_config.dcp_comm_backend != "ag_rs": + raise NotImplementedError( + "DCP for FlashMLA sparse is only validated with the " + "default 'ag_rs' DCP comm backend; got " + f"'{parallel_config.dcp_comm_backend}'" + ) + if not self.fp8_use_mixed_batch: + raise NotImplementedError( + "DCP for FlashMLA sparse is only supported on the " + "mixed-batch fp8 path (num_heads < " + f"{MIN_HEADS_FOR_BF16_PREFILL}); the separate " + "prefill/decode path returns the LSE for decode tokens " + "only, while the DCP merge needs it for every token" + ) + # Head padding (and the tile-scheduler metadata sized from it) is + # computed from the local head count, but the kernel runs on the + # DCP-gathered heads. + gathered_num_heads = ( + self.num_heads * parallel_config.decode_context_parallel_size + ) + gathered_padded_heads = FlashMLASparseImpl._compute_fp8_decode_padded_heads( + gathered_num_heads + ) + if self.fp8_decode_padded_heads != gathered_padded_heads: + raise NotImplementedError( + "DCP for FlashMLA sparse requires the local and " + "DCP-gathered head counts to pad to the same fp8 decode " + f"kernel envelope; got {self.num_heads} local heads " + f"(pad to {self.fp8_decode_padded_heads}) vs " + f"{gathered_num_heads} gathered heads (pad to " + f"{gathered_padded_heads})" + ) + def _build_fp8_mixed_decode_prefill( self, common_attn_metadata: CommonAttentionMetadata, @@ -496,10 +536,9 @@ def build( ) -> FlashMLASparseMetadata: metadata = super().build(common_prefix_len, common_attn_metadata, fast_build) - fp8_use_mixed_batch = self.num_heads < MIN_HEADS_FOR_BF16_PREFILL - metadata.fp8_use_mixed_batch = fp8_use_mixed_batch + metadata.fp8_use_mixed_batch = self.fp8_use_mixed_batch if self.use_fp8_kv_cache: - if fp8_use_mixed_batch: + if self.fp8_use_mixed_batch: metadata.fp8_extra_metadata = self._build_fp8_mixed_decode_prefill( common_attn_metadata ) @@ -512,6 +551,8 @@ def build( class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]): + can_return_lse_for_decode: bool = True + @staticmethod def _compute_fp8_decode_padded_heads(num_heads: int) -> int: # FP8 decode kernel only supports h_q = 64 or 128 @@ -566,6 +607,14 @@ def __init__( "fp8_ds_mla kv-cache dtype" ) + if self.need_to_return_lse_for_decode and not is_quantized_kv_cache( + kv_cache_dtype + ): + raise NotImplementedError( + "DCP for FlashMLA sparse requires an fp8_ds_mla kv-cache; " + "the bf16 sparse path is not supported under DCP." + ) + if kv_cache_dtype == "fp8_ds_mla": # Reserve workspace during initialization assert vllm_config is not None and vllm_config.model_config is not None @@ -590,7 +639,7 @@ def _forward_bf16_kv( kv_c_and_k_pe_cache: torch.Tensor, topk_indices: torch.Tensor, attn_metadata: FlashMLASparseMetadata, - ) -> torch.Tensor: + ) -> tuple[torch.Tensor, torch.Tensor]: # Convert per-request indices to global slots (decode) or workspace # offsets (prefill). req_id_per_token covers the whole batch; slice it # to the MQA tokens (q may exclude prefill tokens routed to dense MHA). @@ -712,7 +761,7 @@ def _fp8_decode( chunk_topk_indices_workspace = topk_indices[chunk.tokens_slice] chunk_topk_length = topk_length[chunk.tokens_slice] - attn_out[chunk.tokens_slice] = self._bf16_flash_mla_kernel( + attn_out[chunk.tokens_slice], _ = self._bf16_flash_mla_kernel( chunk_q, chunk_workspace, chunk_topk_indices_workspace, @@ -727,22 +776,42 @@ def _forward_fp8_kv_mixed_batch( kv_c_and_k_pe_cache: torch.Tensor, topk_indices: torch.Tensor, attn_metadata: FlashMLASparseMetadata, - ) -> torch.Tensor: + ) -> tuple[torch.Tensor, torch.Tensor | None]: """Mixed batch FP8 forward path that treats all tokens as one batch. This is equivalent to main branch's approach and avoids the BF16 prefill kernel which has head padding overhead when num_heads is small. Used when use_mixed_batch is True. + + The lse is only returned when DCP needs it, otherwise None. """ - # Convert per-request indices to global slots (decode) or workspace - # offsets (prefill). - topk_indices = triton_convert_req_index_to_global_index( - attn_metadata.req_id_per_token[: topk_indices.shape[0]], - attn_metadata.block_table, - topk_indices, - BLOCK_SIZE=attn_metadata.block_size, - NUM_TOPK_TOKENS=topk_indices.shape[1], - ) + if self.dcp_world_size > 1: + # The indexer emits global token ids; keep this rank's shard and + # convert to local slots. compact_valid_to_front=False keeps the + # scattered -1s, which the fp8 kernel masks natively and the + # empty-row neutralization below relies on. req_id is sliced to + # topk_indices rows (the converter grids from req_id). + topk_indices = triton_filter_and_convert_dcp_index( + attn_metadata.req_id_per_token[: topk_indices.shape[0]], + attn_metadata.block_table, + topk_indices, + dcp_size=self.dcp_world_size, + dcp_rank=self.dcp_rank, + cp_kv_cache_interleave_size=attn_metadata.cp_kv_cache_interleave_size, + BLOCK_SIZE=attn_metadata.block_size, + NUM_TOPK_TOKENS=topk_indices.shape[1], + compact_valid_to_front=False, + ) + else: + # Convert per-request indices to global slots (decode) or workspace + # offsets (prefill). + topk_indices = triton_convert_req_index_to_global_index( + attn_metadata.req_id_per_token[: topk_indices.shape[0]], + attn_metadata.block_table, + topk_indices, + BLOCK_SIZE=attn_metadata.block_size, + NUM_TOPK_TOKENS=topk_indices.shape[1], + ) assert attn_metadata.fp8_extra_metadata is not None assert isinstance( @@ -750,15 +819,29 @@ def _forward_fp8_kv_mixed_batch( ) fp8_metadata = attn_metadata.fp8_extra_metadata - _attn_out, _ = self._fp8_flash_mla_kernel( + _attn_out, _lse = self._fp8_flash_mla_kernel( q=q.unsqueeze(0), # unsqueeze to add batch_dim: (T, H, D) -> (1, T, H, D) kv_c_and_k_pe_cache=kv_c_and_k_pe_cache, topk_indices=topk_indices.unsqueeze(0), # (T, topk) -> (1, T, topk) kernel_metadata=fp8_metadata, ) - # Output is (1, T, H, D_v), squeeze back to (T, H, D_v) - return _attn_out.squeeze(0) + out = _attn_out.squeeze(0) + + if not self.need_to_return_lse_for_decode: + return out, None + + # Kernel LSE is (1, H, T); the DCP merge consumes (T, H). + lse = _lse.squeeze(0).transpose(0, 1) + # Rows where this rank owns none of the selected tokens (all indices + # -1) have undefined out/lse; (0, -inf) is the identity element of the + # cross-rank LSE merge, so it drops this rank from those rows. + empty_rows = (topk_indices == -1).all(dim=-1) + out.masked_fill_(empty_rows.view(-1, 1, 1), 0.0) + lse.masked_fill_(empty_rows.view(-1, 1), float("-inf")) + # The head-padding slice above can leave `out` non-contiguous, and the + # merge feeds it to reduce_scatter. + return out.contiguous(), lse def _fp8_flash_mla_kernel( self, @@ -793,9 +876,10 @@ def _fp8_flash_mla_kernel( softmax_scale=self.softmax_scale, ) - # Slice output back to actual head count if we padded + # Slice output and lse back to actual head count if we padded if actual_num_heads < padded_num_heads: out = out[:, :, :actual_num_heads, :] + lse = lse[:, :actual_num_heads, :] return out, lse @@ -805,35 +889,42 @@ def _bf16_flash_mla_kernel( kv_c_and_k_pe_cache: torch.Tensor, topk_indices: torch.Tensor, topk_length: torch.Tensor | None = None, - ) -> torch.Tensor: + ) -> tuple[torch.Tensor, torch.Tensor]: num_tokens = q.shape[0] kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.view( -1, 1, kv_c_and_k_pe_cache.shape[-1] ) # NOTE(Chen): kernel requires num_local_head to be a multiple of - # 64 on hopper and 128 on blackwell - if self.num_heads % self.prefill_padding != 0: - assert self.prefill_padding % self.num_heads == 0 + # 64 on hopper and 128 on blackwell. Pad from q's head count, not + # self.num_heads: under DCP the heads are all-gathered before this. + actual_num_heads = q.shape[1] + padded_num_heads = ( + (actual_num_heads + self.prefill_padding - 1) + // self.prefill_padding + * self.prefill_padding + ) + if actual_num_heads < padded_num_heads: logger.warning_once( - f"Padding num_heads from {self.num_heads} to " - f"{self.prefill_padding} for BF16 sparse prefill kernel" + f"Padding num_heads from {actual_num_heads} to " + f"{padded_num_heads} for BF16 sparse prefill kernel" ) - q_padded = q.new_empty((q.shape[0], self.prefill_padding, q.shape[2])) - q_padded[:, : self.num_heads, :] = q + q_padded = q.new_empty((q.shape[0], padded_num_heads, q.shape[2])) + q_padded[:, :actual_num_heads, :] = q q = q_padded topk_indices = topk_indices.view(num_tokens, 1, -1) - output = flash_mla_sparse_fwd( + output, _, lse = flash_mla_sparse_fwd( q, kv_c_and_k_pe_cache, topk_indices, self.softmax_scale, topk_length=topk_length, - )[0] + ) - output = output[:, : self.num_heads, :] - return output + output = output[:, :actual_num_heads, :] + lse = lse[:, :actual_num_heads] + return output, lse def forward_mqa( self, @@ -859,12 +950,16 @@ def forward_mqa( use_fp8_cache = self.kv_cache_dtype == "fp8_ds_mla" + lse: torch.Tensor | None = None + if not use_fp8_cache: - attn_out = self._forward_bf16_kv( + attn_out, bf16_lse = self._forward_bf16_kv( q, kv_c_and_k_pe_cache, topk_indices, attn_metadata ) + if self.need_to_return_lse_for_decode: + lse = bf16_lse elif attn_metadata.fp8_use_mixed_batch: - attn_out = self._forward_fp8_kv_mixed_batch( + attn_out, lse = self._forward_fp8_kv_mixed_batch( q, kv_c_and_k_pe_cache, topk_indices, attn_metadata ) else: @@ -872,4 +967,4 @@ def forward_mqa( q, kv_c_and_k_pe_cache, topk_indices, attn_metadata ) - return attn_out, None + return attn_out, lse From 08afae27863e884d68937f76e5bb87bba171c5de Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Tue, 18 Aug 2026 23:36:05 -0500 Subject: [PATCH 131/839] [ModelRunner v2] Enable MRV2 for pooling models by default (#48290) Signed-off-by: Taneem Ibrahim --- tests/models/language/pooling/test_colbert.py | 4 +-- tests/test_config.py | 25 ++++++++++++++++++- vllm/config/vllm.py | 3 --- 3 files changed, 26 insertions(+), 6 deletions(-) diff --git a/tests/models/language/pooling/test_colbert.py b/tests/models/language/pooling/test_colbert.py index dd4e1dea9772..75323d54bf8d 100644 --- a/tests/models/language/pooling/test_colbert.py +++ b/tests/models/language/pooling/test_colbert.py @@ -323,8 +323,8 @@ def test_colbert_embed_not_supported(self, colbert_model): [ pytest.param("bert", True, id="bert-v2"), pytest.param("modernbert", True, id="modernbert-v2"), - pytest.param("jina", False, id="jina-v1"), - pytest.param("lfm2", False, id="lfm2-v1"), + pytest.param("jina", True, id="jina-v2"), + pytest.param("lfm2", True, id="lfm2-v2"), ], ) def test_colbert_hf_comparison(vllm_runner, monkeypatch, backend, use_v2): diff --git a/tests/test_config.py b/tests/test_config.py index a2797e52126a..8560e4f334cd 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -342,15 +342,38 @@ def test_resolve_cudagraph_mode_adjusts_spec_decode_sizes_only_for_v1( ), False, ), + ( + SimpleNamespace( + model="sentence-transformers/all-MiniLM-L6-v2", + architectures=["BertModel"], + runner_type="pooling", + is_multimodal_model=False, + is_moe=False, + is_quantized=False, + ), + True, + ), ( SimpleNamespace( model="Qwen/Qwen3-Embedding-0.6B", architectures=["Qwen3ForCausalLM"], runner_type="pooling", + is_multimodal_model=False, is_moe=False, is_quantized=False, ), - False, + True, + ), + ( + SimpleNamespace( + model="TomoroAI/tomoro-colqwen3-embed-4b", + architectures=["ColQwen3"], + runner_type="pooling", + is_multimodal_model=True, + is_moe=False, + is_quantized=False, + ), + True, ), ], ) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index ff93bed86fb6..cee60e99b635 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -676,9 +676,6 @@ def _is_default_v2_model_runner_model(self) -> bool: if model_config is None: return False - if model_config.runner_type != "generate": - return False - architectures = getattr(model_config, "architectures", []) default_architectures = default_v2_model_runner_architectures() is_default_v2_architecture = any( From 86a89a93486d3bd4bd1c94149066c992cb7e2f09 Mon Sep 17 00:00:00 2001 From: Vineeth Sai Varikuntla Date: Tue, 18 Aug 2026 21:48:50 -0700 Subject: [PATCH 132/839] [Bugfix][Frontend] Run the serve arg checks for `vllm launch` too (#52825) Signed-off-by: Vineeth Sai --- tests/entrypoints/openai/test_cli_args.py | 44 +++++++++++++++++++++++ vllm/entrypoints/openai/cli_args.py | 4 ++- 2 files changed, 47 insertions(+), 1 deletion(-) diff --git a/tests/entrypoints/openai/test_cli_args.py b/tests/entrypoints/openai/test_cli_args.py index 1f764202e55e..c8c5fd0f55a4 100644 --- a/tests/entrypoints/openai/test_cli_args.py +++ b/tests/entrypoints/openai/test_cli_args.py @@ -214,6 +214,50 @@ def test_per_request_metrics_requires_log_stats(serve_parser): validate_parsed_serve_args(args) +def _build_launch_render_parser(): + """Mirror `vllm launch render`. + + `vllm/entrypoints/cli/main.py` parses subcommands with ``dest="subparser"``, + and `LaunchSubcommandBase.add_cli_args` builds the component parser with + `make_arg_parser`, so a launch component carries serve args under + ``args.subparser == "launch"``. + """ + vllm_parser = FlexibleArgumentParser() + subparsers = vllm_parser.add_subparsers(required=False, dest="subparser") + launch_parser = subparsers.add_parser("launch") + launch_subparsers = launch_parser.add_subparsers( + required=True, dest="launch_component" + ) + render_parser = launch_subparsers.add_parser("render") + make_arg_parser(render_parser) + return vllm_parser + + +@pytest.fixture +def launch_render_parser(): + return _build_launch_render_parser() + + +def test_launch_render_validates_serve_args(launch_render_parser): + """`vllm launch render` reuses the serve parser, so it gets the serve checks""" + args = launch_render_parser.parse_args( + args=["launch", "render", "--enable-auto-tool-choice"] + ) + assert args.subparser == "launch" + with pytest.raises(TypeError): + validate_parsed_serve_args(args) + + +def test_subcommand_without_serve_args_skips_validation(): + """A subcommand that does not use the serve parser is still skipped""" + vllm_parser = FlexibleArgumentParser() + subparsers = vllm_parser.add_subparsers(required=False, dest="subparser") + subparsers.add_parser("chat") + args = vllm_parser.parse_args(args=["chat"]) + + validate_parsed_serve_args(args) + + @pytest.mark.parametrize( "cli_args, expected_middleware", [ diff --git a/vllm/entrypoints/openai/cli_args.py b/vllm/entrypoints/openai/cli_args.py index b387fb63573a..e9284c5dd69b 100644 --- a/vllm/entrypoints/openai/cli_args.py +++ b/vllm/entrypoints/openai/cli_args.py @@ -412,7 +412,9 @@ def make_arg_parser(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: def validate_parsed_serve_args(args: argparse.Namespace): """Quick checks for model serve args that raise prior to loading.""" - if hasattr(args, "subparser") and args.subparser != "serve": + # `vllm launch ` builds its parser with make_arg_parser too (see + # LaunchSubcommandBase.add_cli_args), so its args are serve args as well. + if hasattr(args, "subparser") and args.subparser not in ("serve", "launch"): return # Ensure that the chat template is valid; raises if it likely isn't From 5a4c8d99242e9e069b604d0e9b969e77f7dd501d Mon Sep 17 00:00:00 2001 From: Agoni-02 <114491222+Agoni-02@users.noreply.github.com> Date: Wed, 19 Aug 2026 12:55:20 +0800 Subject: [PATCH 133/839] [Bugfix][LoRA] Add embedding_modules for Qwen3.5 CausalLM (#48850) Signed-off-by: m0_68169237 Co-authored-by: m0_68169237 Co-authored-by: Cursor Composer Co-authored-by: Jee Jee Li --- vllm/model_executor/models/qwen3_5.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/vllm/model_executor/models/qwen3_5.py b/vllm/model_executor/models/qwen3_5.py index 3fb3b7a782c8..e8cc041ce423 100644 --- a/vllm/model_executor/models/qwen3_5.py +++ b/vllm/model_executor/models/qwen3_5.py @@ -309,6 +309,11 @@ class Qwen3_5ForCausalLMBase( "in_proj_qkvz": ["in_proj_qkv", "in_proj_z"], "in_proj_ba": ["in_proj_b", "in_proj_a"], } + # Maps PEFT embed/lm_head LoRA targets onto vLLM embedding wrappers. + embedding_modules = { + "embed_tokens": "input_embeddings", + "lm_head": "output_embeddings", + } # Some community text-only checkpoints keep the extraneous # `model.language_model.` prefix inherited from the VL training stack. From ee117307512f7ab0fe5ce70f8a2963bb83dfc442 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Wed, 19 Aug 2026 00:07:40 -0700 Subject: [PATCH 134/839] [BugFix] Revert incorrect MM keep_on_cpu=True changes (#52881) Signed-off-by: Nick Hill --- vllm/model_executor/models/ernie45_vl.py | 12 ++++++++---- vllm/model_executor/models/fireredlid.py | 2 +- vllm/model_executor/models/isaac.py | 4 +++- vllm/model_executor/models/llava_onevision2.py | 4 ++-- vllm/model_executor/models/minicpmv4_6.py | 4 ++-- .../model_executor/models/moss_transcribe_diarize.py | 1 - vllm/model_executor/models/qwen2_audio.py | 2 +- 7 files changed, 17 insertions(+), 12 deletions(-) diff --git a/vllm/model_executor/models/ernie45_vl.py b/vllm/model_executor/models/ernie45_vl.py index 9fa04470020d..8c4b150e30d5 100644 --- a/vllm/model_executor/models/ernie45_vl.py +++ b/vllm/model_executor/models/ernie45_vl.py @@ -1726,7 +1726,9 @@ def encoder_eager_forward( # Eager fallback: run the full pipeline (ViT + resampler). The result # is scattered directly, so it must be the post-merge embeddings. pixel_values = mm_kwargs["pixel_values"].type(self.vision_model.dtype) - grid_thw = mm_kwargs["image_grid_thw"].to(self.vision_model.device) + grid_thw = mm_kwargs["image_grid_thw"].to( + self.vision_model.device, non_blocking=True + ) image_features = self.vision_model(pixel_values, grid_thw) return self.resampler_model(image_features, grid_thw) @@ -1744,11 +1746,13 @@ def postprocess_encoder_output( # the actual batch grid_thw, then scatter the post-merge embeddings. # Ernie only uses the single "default" encoder path. output = outputs["default"] - grid_thw = batch_mm_kwargs["image_grid_thw"].to(output.device) + grid_thw_cpu = batch_mm_kwargs["image_grid_thw"] + grid_thw = grid_thw_cpu.to(output.device, non_blocking=True) # The valid token count slices the graph output for the eager # resampler call, so it has to come back to the host. - with gpu_sync_allowed(): - num_valid = int((grid_thw[:, 0] * grid_thw[:, 1] * grid_thw[:, 2]).sum()) + num_valid = int( + (grid_thw_cpu[:, 0] * grid_thw_cpu[:, 1] * grid_thw_cpu[:, 2]).sum() + ) image_embeds = self.resampler_model(output[:num_valid], grid_thw) scatter_output_slices(image_embeds, indices, per_item_out_tokens, dest, clone) diff --git a/vllm/model_executor/models/fireredlid.py b/vllm/model_executor/models/fireredlid.py index 69fe995b6b23..09faf75da1ff 100644 --- a/vllm/model_executor/models/fireredlid.py +++ b/vllm/model_executor/models/fireredlid.py @@ -482,7 +482,7 @@ def _get_mm_fields_config( ) -> Mapping[str, MultiModalFieldConfig]: return dict( input_features=MultiModalFieldConfig.batched("audio"), - speech_lengths=MultiModalFieldConfig.batched("audio", keep_on_cpu=True), + speech_lengths=MultiModalFieldConfig.batched("audio"), fake_token_lengths=MultiModalFieldConfig.batched("audio", keep_on_cpu=True), ) diff --git a/vllm/model_executor/models/isaac.py b/vllm/model_executor/models/isaac.py index 4979971b02db..789c2403b905 100644 --- a/vllm/model_executor/models/isaac.py +++ b/vllm/model_executor/models/isaac.py @@ -961,7 +961,9 @@ def _process_image_input( device = next(self.language_model.parameters()).device dtype = self.vision_embedding.linear_fc1.weight.dtype pixel_values = pixel_values.to(device=device, dtype=dtype) - spatial_grids = image_grid_thw[:, 1:3].to(device, dtype=torch.int32) + spatial_grids = image_grid_thw[:, 1:3].to( + device, dtype=torch.int32, non_blocking=True + ) vision_embeddings = self.vision_embedding((pixel_values, spatial_grids)) merge_size = self.config.vision_config.pixel_shuffle_scale_factor diff --git a/vllm/model_executor/models/llava_onevision2.py b/vllm/model_executor/models/llava_onevision2.py index 6c13f5b14b61..c09b477f5333 100644 --- a/vllm/model_executor/models/llava_onevision2.py +++ b/vllm/model_executor/models/llava_onevision2.py @@ -435,7 +435,7 @@ def _field_config(hf_inputs: Mapping[str, torch.Tensor]): # OV2 first-class MM kwarg: per-patch (t,h,w) # positions required by the 3-D vision RoPE. patch_positions=MultiModalFieldConfig.flat_from_sizes( - "image", image_pixel_grid_sizes, keep_on_cpu=True + "image", image_pixel_grid_sizes ), pixel_values_videos=MultiModalFieldConfig.flat_from_sizes( "video", video_patch_sizes @@ -444,7 +444,7 @@ def _field_config(hf_inputs: Mapping[str, torch.Tensor]): "video", video_num_frames, keep_on_cpu=True ), patch_positions_videos=MultiModalFieldConfig.flat_from_sizes( - "video", video_patch_sizes, keep_on_cpu=True + "video", video_patch_sizes ), video_num_frames=MultiModalFieldConfig.batched("video", keep_on_cpu=True), frame_timestamps=MultiModalFieldConfig.batched("video", keep_on_cpu=True), diff --git a/vllm/model_executor/models/minicpmv4_6.py b/vllm/model_executor/models/minicpmv4_6.py index a7e8d7c75f65..69124a061c73 100644 --- a/vllm/model_executor/models/minicpmv4_6.py +++ b/vllm/model_executor/models/minicpmv4_6.py @@ -75,11 +75,11 @@ def _minicpmv4_6_field_config(hf_inputs: Mapping[str, torch.Tensor]): fields = dict( pixel_values=MultiModalFieldConfig.batched("image"), - tgt_sizes=MultiModalFieldConfig.batched("image", keep_on_cpu=True), + tgt_sizes=MultiModalFieldConfig.batched("image"), image_embeds=MultiModalFieldConfig.batched("image"), video_pixel_values=MultiModalFieldConfig.batched("video"), video_image_sizes=MultiModalFieldConfig.batched("video", keep_on_cpu=True), - video_tgt_sizes=MultiModalFieldConfig.batched("video", keep_on_cpu=True), + video_tgt_sizes=MultiModalFieldConfig.batched("video"), video_embeds=MultiModalFieldConfig.batched("video"), ) if "use_vit_merger" in hf_inputs: diff --git a/vllm/model_executor/models/moss_transcribe_diarize.py b/vllm/model_executor/models/moss_transcribe_diarize.py index 347c61cb5410..a862c27e6f25 100644 --- a/vllm/model_executor/models/moss_transcribe_diarize.py +++ b/vllm/model_executor/models/moss_transcribe_diarize.py @@ -352,7 +352,6 @@ def _mtd_field_config( audio_feature_lengths=MultiModalFieldConfig.flat_from_sizes( "audio", audio_chunk_counts, - keep_on_cpu=True, ), audio_chunk_counts=MultiModalFieldConfig.batched("audio", keep_on_cpu=True), audio_token_lengths=MultiModalFieldConfig.batched( diff --git a/vllm/model_executor/models/qwen2_audio.py b/vllm/model_executor/models/qwen2_audio.py index c989ec85b493..115a7f7f79b6 100644 --- a/vllm/model_executor/models/qwen2_audio.py +++ b/vllm/model_executor/models/qwen2_audio.py @@ -130,7 +130,7 @@ def _qwen2audio_field_config(hf_inputs: Mapping[str, torch.Tensor]): return dict( audio_embeds=MultiModalFieldConfig.batched("audio"), input_features=MultiModalFieldConfig.batched("audio"), - feature_attention_mask=MultiModalFieldConfig.batched("audio", keep_on_cpu=True), + feature_attention_mask=MultiModalFieldConfig.batched("audio"), ) From b09bd69b5bf14911abef9a0e8e493b83c8a38fa6 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Wed, 19 Aug 2026 00:13:52 -0700 Subject: [PATCH 135/839] [Model][NVIDIA] Route DSA models to the CUDA non-compiled path (#52861) --- tests/compile/fusions_e2e/conftest.py | 21 +-- tests/compile/fusions_e2e/models.py | 24 --- tests/compile/fusions_e2e/test_tp1_quant.py | 6 +- tests/compile/fusions_e2e/test_tp2_ar_rms.py | 6 +- tests/compile/h100/test_startup.py | 13 -- .../test_fused_deepseek_v32_norm_rope.py | 7 +- tests/models/registry.py | 5 + tests/test_config.py | 144 +++++++++++++++++- vllm/config/speculative.py | 12 +- vllm/config/vllm.py | 106 ++++++++----- vllm/model_executor/models/registry.py | 8 +- vllm/models/deepseek_v32/__init__.py | 25 ++- vllm/models/deepseek_v32/amd/rocm.py | 1 - vllm/models/deepseek_v32/attention.py | 19 +-- vllm/models/deepseek_v32/common/kernels.py | 3 +- vllm/v1/spec_decode/llm_base_proposer.py | 6 +- 16 files changed, 260 insertions(+), 146 deletions(-) diff --git a/tests/compile/fusions_e2e/conftest.py b/tests/compile/fusions_e2e/conftest.py index 929b27c11c12..f3df3c358ea2 100644 --- a/tests/compile/fusions_e2e/conftest.py +++ b/tests/compile/fusions_e2e/conftest.py @@ -90,11 +90,8 @@ def run( backend_name = attn_backend.backend.name.lower() requires_mla = "deepseek" in model_name.lower() is_mla = "mla" in backend_name - # DeepSeek V3.2 uses sparse MLA - requires_sparse = "v3.2" in model_name.lower() - is_sparse = "sparse" in backend_name - if requires_mla != is_mla or requires_sparse != is_sparse: + if requires_mla != is_mla: pytest.skip( f"Incompatible model '{model_name}' and " f"attention backend '{attn_backend.backend.name}'" @@ -128,22 +125,6 @@ def run( # models (e.g. Llama-4-Scout-FP8) at 16384 tokens may trigger OOM. model_kwargs.setdefault("max_num_batched_tokens", 8192) - # Sparse MLA models (DSv3.2) hit an over-strict inductor assertion in - # decompose_auto_functionalized when +rotary_embedding is forced into - # the compile graph. Disable qk_norm+rope fusion (which auto-enables - # +rotary_embedding) for this combo to avoid the known torch bug. - # TODO: remove once upstream torch fix lands. - if requires_sparse: - if "pass_config" in compilation_config: - compilation_config["pass_config"].enable_qk_norm_rope_fusion = False - matches_check = [m for m in matches_check if m != "norm_rope_fusion"] - # DSv3.2 sparse indexer uses persistent_topk with k=config.index_topk - # (2048 for the default config). max_model_len must be >= index_topk - # or the topk kernel raises "k out of range" at runtime. - model_kwargs["max_model_len"] = max( - model_kwargs.get("max_model_len", 0), 2048 - ) - # Always compile the full graph instead of piecewise if not compilation_config["use_inductor_graph_partition"]: compilation_config["splitting_ops"] = [] diff --git a/tests/compile/fusions_e2e/models.py b/tests/compile/fusions_e2e/models.py index 2d407fbc3025..7879cbe65096 100644 --- a/tests/compile/fusions_e2e/models.py +++ b/tests/compile/fusions_e2e/models.py @@ -58,18 +58,6 @@ id="TRITON_MLA", ) -FLASHMLA_SPARSE_ATTN = pytest.param( - AttentionBackendCase( - backend=AttentionBackendEnum.FLASHMLA_SPARSE, - model_kwargs=dict(kv_cache_dtype="fp8_ds_mla"), - ), - id="FLASHMLA_SPARSE", - marks=pytest.mark.skipif( - not is_blackwell(), - reason="FlashMLA Sparse requires Blackwell", - ), -) - # Models llama3_8b = ModelFusionInfo( model_name="meta-llama/Llama-3.1-8B-Instruct", @@ -197,18 +185,6 @@ ), ) -deepseek_v32_fp4 = ModelFusionInfo( - model_name="nvidia/DeepSeek-V3.2-NVFP4", - matches=lambda n_layers: Matches( - rms_quant_fusion=0, - # silu+quant on dense layers only; MoE hides the act+quant site - act_quant_fusion=min(3, n_layers), - # MLA attn + NVFP4 output quant fuses on sparse MLA output path - attn_quant_fusion=n_layers, - ar_rms_fusion=n_layers * 2 + 1, - ), -) - gpt_oss_20b = ModelFusionInfo( model_name="openai/gpt-oss-20b", matches=lambda n_layers: Matches( diff --git a/tests/compile/fusions_e2e/test_tp1_quant.py b/tests/compile/fusions_e2e/test_tp1_quant.py index fbb382b4458d..3fab133f963e 100644 --- a/tests/compile/fusions_e2e/test_tp1_quant.py +++ b/tests/compile/fusions_e2e/test_tp1_quant.py @@ -18,7 +18,6 @@ from .models import ( FLASHINFER_ATTN, FLASHINFER_MLA_ATTN, - FLASHMLA_SPARSE_ATTN, ROCM_AITER_UNIFIED_ATTN, ROCM_ATTN, TRITON_ATTN, @@ -26,7 +25,6 @@ deepseek_coder_v2_lite_fp8, deepseek_r1_fp4, deepseek_v3_fp8, - deepseek_v32_fp4, llama3_8b_fp4, llama3_8b_fp8, llama4_scout_fp4, @@ -149,11 +147,11 @@ def test_tp1_fp8_fusions( @pytest.mark.parametrize( "model_name, matches_fn, model_kwargs, hf_overrides", - [llama3_8b_fp4, llama4_scout_fp4, deepseek_r1_fp4, deepseek_v32_fp4], + [llama3_8b_fp4, llama4_scout_fp4, deepseek_r1_fp4], ) @pytest.mark.parametrize( "attn_backend", - [FLASHINFER_ATTN, FLASHINFER_MLA_ATTN, FLASHMLA_SPARSE_ATTN], + [FLASHINFER_ATTN, FLASHINFER_MLA_ATTN], ) @pytest.mark.parametrize("n_layers", [6]) @pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm")) diff --git a/tests/compile/fusions_e2e/test_tp2_ar_rms.py b/tests/compile/fusions_e2e/test_tp2_ar_rms.py index b6ad4e2e6e85..c88d47cd314f 100644 --- a/tests/compile/fusions_e2e/test_tp2_ar_rms.py +++ b/tests/compile/fusions_e2e/test_tp2_ar_rms.py @@ -18,14 +18,12 @@ from .models import ( FLASHINFER_ATTN, FLASHINFER_MLA_ATTN, - FLASHMLA_SPARSE_ATTN, ROCM_AITER_UNIFIED_ATTN, ROCM_ATTN, TRITON_ATTN, deepseek_coder_v2_lite_fp8, deepseek_r1_fp4, deepseek_v3_fp8, - deepseek_v32_fp4, gpt_oss_20b, llama3_8b, llama3_8b_fp4, @@ -120,11 +118,11 @@ def test_tp2_ar_rms_fp8_fusions( @multi_gpu_test(num_gpus=2) @pytest.mark.parametrize( "model_name, matches_fn, model_kwargs, hf_overrides", - [llama3_8b_fp4, llama4_scout_fp4, deepseek_r1_fp4, deepseek_v32_fp4], + [llama3_8b_fp4, llama4_scout_fp4, deepseek_r1_fp4], ) @pytest.mark.parametrize( "attn_backend", - [FLASHINFER_ATTN, FLASHINFER_MLA_ATTN, FLASHMLA_SPARSE_ATTN], + [FLASHINFER_ATTN, FLASHINFER_MLA_ATTN], ) @pytest.mark.parametrize("n_layers", [4]) @pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm")) diff --git a/tests/compile/h100/test_startup.py b/tests/compile/h100/test_startup.py index 075fc8e24972..e57cf85b4a77 100644 --- a/tests/compile/h100/test_startup.py +++ b/tests/compile/h100/test_startup.py @@ -132,19 +132,6 @@ class ModelStartupSpec(NamedTuple): ), id="gpt_oss_120b", ), - # NOTE: DeepSeek-V3.2 requires sparse MLA (index_topk) which needs - # Hopper+ GPUs. This test must run on H100 (see pytorch.yaml). - pytest.param( - ModelStartupSpec( - model="deepseek-ai/DeepSeek-V3.2", - hf_overrides=_SMALL_MOE_OVERRIDES, - cold_artifacts_saved=9, - # https://github.com/vllm-project/vllm/issues/38051 - warm_artifacts_saved=0 if is_torch_equal_or_newer("2.12.0") else 9, - warm_artifacts_loaded=9 if is_torch_equal_or_newer("2.12.0") else 0, - ), - id="deepseek_v3.2", - ), pytest.param( ModelStartupSpec( model="moonshotai/Kimi-K2.5", diff --git a/tests/kernels/test_fused_deepseek_v32_norm_rope.py b/tests/kernels/test_fused_deepseek_v32_norm_rope.py index 7a00d0c3d644..d3d5f3b0dc75 100644 --- a/tests/kernels/test_fused_deepseek_v32_norm_rope.py +++ b/tests/kernels/test_fused_deepseek_v32_norm_rope.py @@ -146,8 +146,8 @@ def assert_fp8(got: torch.Tensor, ref: torch.Tensor, msg: str): @pytest.mark.parametrize("num_tokens", [1, 4, 17, 512, 4096]) @pytest.mark.parametrize("index_interleave", [True, False]) -@pytest.mark.parametrize("mla_fp8", [False, True]) -def test_fused_norm_rope(num_tokens: int, index_interleave: bool, mla_fp8: bool): +@pytest.mark.parametrize("mla_dtype", ["auto", "bfloat16", "fp8"]) +def test_fused_norm_rope(num_tokens: int, index_interleave: bool, mla_dtype: str): torch.manual_seed(0) dev = "cuda" max_pos = 8192 @@ -167,13 +167,12 @@ def test_fused_norm_rope(num_tokens: int, index_interleave: bool, mla_fp8: bool) bs = max_pos # single block covering all tokens mla_dim = KV_LORA + ROPE_DIM + mla_fp8 = mla_dtype == "fp8" if mla_fp8: mla_cache = torch.zeros(1, bs, mla_dim, device=dev, dtype=torch.uint8) - mla_dtype = "fp8" mla_k_scale = torch.tensor([0.3], device=dev, dtype=torch.float32) else: mla_cache = torch.zeros(1, bs, mla_dim, device=dev, dtype=torch.bfloat16) - mla_dtype = "auto" mla_k_scale = None idx_row = INDEX_HEAD_DIM + INDEX_HEAD_DIM // 128 * 4 # 132 idx_cache = torch.zeros(1, bs, idx_row, device=dev, dtype=torch.uint8) diff --git a/tests/models/registry.py b/tests/models/registry.py index 64efae5261be..0411d64f9489 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -1679,6 +1679,11 @@ def check_available_online( speculative_model="luccafong/deepseek_mtp_draft_random", trust_remote_code=True, ), + "DeepseekV32MTPModel": _HfExamplesInfo( + "nvidia/DeepSeek-V3.2-NVFP4", + speculative_model="nvidia/DeepSeek-V3.2-NVFP4", + is_available_online=False, + ), "DeepSeekV4MTPModel": _HfExamplesInfo( "deepseek-ai/DeepSeek-V4-Flash", speculative_model="deepseek-ai/DeepSeek-V4-Flash", diff --git a/tests/test_config.py b/tests/test_config.py index 8560e4f334cd..2d9c3b83f760 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -113,18 +113,154 @@ def test_v2_model_runner_env_tri_state(monkeypatch, env_value, expected): assert envs.VLLM_USE_V2_MODEL_RUNNER is expected -def test_rocm_defaults_deepseek_v4_to_mrv1(monkeypatch): - """ROCm keeps DeepSeek V4 on MRV1, which is still faster there.""" - from vllm.config.vllm import default_v2_model_runner_architectures +def test_rocm_keeps_compiled_deepseek_defaults(monkeypatch): + """ROCm keeps DeepSeek V3.2 and V4 on their compiled MRV1 paths.""" + from vllm.config.vllm import ( + default_breakable_cudagraph_architectures, + default_v2_model_runner_architectures, + ) from vllm.platforms import current_platform monkeypatch.setattr(current_platform, "is_rocm", lambda: True) # The lookup is lru_cached against a fixed platform. default_v2_model_runner_architectures.cache_clear() + default_breakable_cudagraph_architectures.cache_clear() + try: + v2_architectures = default_v2_model_runner_architectures() + breakable_architectures = default_breakable_cudagraph_architectures() + + assert "DeepseekV32ForCausalLM" not in v2_architectures + assert "DeepseekV4ForCausalLM" not in v2_architectures + assert "DeepseekV32ForCausalLM" not in breakable_architectures + assert "DeepseekV32MTPModel" not in breakable_architectures + finally: + default_v2_model_runner_architectures.cache_clear() + default_breakable_cudagraph_architectures.cache_clear() + + +@pytest.mark.parametrize( + ("model", "architecture"), + [ + ("nvidia/GLM-5.2-NVFP4", "GlmMoeDsaForCausalLM"), + ("zai-org/GLM-5.2-FP8", "GlmMoeDsaForCausalLM"), + ("nvidia/DeepSeek-V3.2-NVFP4", "DeepseekV32ForCausalLM"), + ], +) +@pytest.mark.parametrize("with_mtp", [False, True], ids=["no-mtp", "mtp"]) +def test_dsa_models_default_to_mrv2_and_breakable_cudagraph( + monkeypatch, model, architecture, with_mtp +): + from vllm.compilation.breakable_cudagraph import ( + is_breakable_cudagraph_enabled, + ) + from vllm.config.vllm import ( + default_breakable_cudagraph_architectures, + default_v2_model_runner_architectures, + ) + from vllm.platforms import current_platform + + monkeypatch.delenv("VLLM_USE_BREAKABLE_CUDAGRAPH", raising=False) + monkeypatch.delenv("VLLM_USE_V2_MODEL_RUNNER", raising=False) + monkeypatch.setattr(vllm_config_module, "HAS_TRITON", True) + monkeypatch.setattr(current_platform, "is_rocm", lambda: False) + default_v2_model_runner_architectures.cache_clear() + default_breakable_cudagraph_architectures.cache_clear() + + model_config = SimpleNamespace( + model=model, + architectures=[architecture], + runner_type="generate", + is_moe=True, + is_hybrid=False, + is_attention_free=False, + is_diffusion=False, + ) + config = SimpleNamespace( + model_config=model_config, + speculative_config=SimpleNamespace(method="mtp") if with_mtp else None, + parallel_config=SimpleNamespace(prefill_context_parallel_size=1), + compilation_config=CompilationConfig( + cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE + ), + ) + config._dflash_needs_multi_kv_group = lambda: False + config._is_default_v2_model_runner_model = lambda: ( + VllmConfig._is_default_v2_model_runner_model(config) + ) + config._get_v2_model_runner_unsupported_features = lambda: [] + config._uses_breakable_cudagraph_by_default = lambda: ( + VllmConfig._uses_breakable_cudagraph_by_default(config) + ) + try: - assert "DeepseekV4ForCausalLM" not in default_v2_model_runner_architectures() + assert VllmConfig.use_v2_model_runner.fget(config) + assert VllmConfig._maybe_enable_breakable_cudagraph(config) + assert is_breakable_cudagraph_enabled() + assert config.compilation_config.mode == CompilationMode.NONE + assert config.compilation_config.cudagraph_mode.has_piecewise_cudagraphs() finally: + os.environ.pop("VLLM_USE_BREAKABLE_CUDAGRAPH", None) default_v2_model_runner_architectures.cache_clear() + default_breakable_cudagraph_architectures.cache_clear() + + +@pytest.mark.parametrize( + ("architecture", "is_rocm", "expected"), + [ + ("DeepseekV32ForCausalLM", False, True), + ("DeepseekV32ForCausalLM", True, False), + ("DeepseekV32MTPModel", False, True), + ("DeepseekV32MTPModel", True, False), + ("GlmMoeDsaForCausalLM", False, True), + ("GlmMoeDsaForCausalLM", True, True), + ], +) +def test_dsa_breakable_cudagraph_platform_default( + monkeypatch, architecture, is_rocm, expected +): + from vllm.config.vllm import default_breakable_cudagraph_architectures + from vllm.platforms import current_platform + + monkeypatch.delenv("VLLM_USE_BREAKABLE_CUDAGRAPH", raising=False) + monkeypatch.setattr(current_platform, "is_rocm", lambda: is_rocm) + default_breakable_cudagraph_architectures.cache_clear() + config = SimpleNamespace( + model_config=SimpleNamespace(architectures=[architecture]), + compilation_config=CompilationConfig(), + ) + config._uses_breakable_cudagraph_by_default = lambda: ( + VllmConfig._uses_breakable_cudagraph_by_default(config) + ) + + try: + assert VllmConfig._maybe_enable_breakable_cudagraph(config) is expected + if expected: + assert config.compilation_config.mode == CompilationMode.NONE + finally: + os.environ.pop("VLLM_USE_BREAKABLE_CUDAGRAPH", None) + default_breakable_cudagraph_architectures.cache_clear() + + +@pytest.mark.parametrize( + ("model_type", "expected_architecture"), + [ + ("deepseek_v32", "DeepseekV32MTPModel"), + ("glm_moe_dsa", "DeepseekV32MTPModel"), + ("deepseek_v3", "DeepSeekMTPModel"), + ], +) +def test_dsa_models_select_matching_mtp(model_type, expected_architecture): + from transformers import PretrainedConfig + + hf_config = PretrainedConfig( + architectures=["DeepseekV32ForCausalLM"], + num_nextn_predict_layers=1, + ) + hf_config.model_type = model_type + + SpeculativeConfig.hf_config_override(hf_config) + + assert hf_config.architectures == [expected_architecture] @pytest.mark.parametrize( diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index c3e0866b453d..6cdba55263c3 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -341,6 +341,7 @@ def compute_hash(self) -> str: @staticmethod def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig: initial_architecture = hf_config.architectures[0] + use_v32_mtp = hf_config.model_type in ("deepseek_v32", "glm_moe_dsa") if hf_config.model_type == "dots3_note": n_predict = getattr(hf_config, "num_nextn_predict_layers", 1) mtp_layer_types = getattr(hf_config, "mtp_layer_types", None) @@ -365,7 +366,12 @@ def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig: if hf_config.model_type == "deepseek_mtp": n_predict = getattr(hf_config, "num_nextn_predict_layers", None) hf_config.update( - {"n_predict": n_predict, "architectures": ["DeepSeekMTPModel"]} + { + "n_predict": n_predict, + "architectures": [ + "DeepseekV32MTPModel" if use_v32_mtp else "DeepSeekMTPModel" + ], + } ) if hf_config.model_type == "deepseek_v4": hf_config.model_type = "deepseek_mtp" @@ -759,10 +765,6 @@ def __post_init__(self): if self.method == "mtp": if self.target_model_config is None: raise ValueError("target_model_config must be present for mtp") - if self.target_model_config.hf_text_config.model_type == "deepseek_v32": - # FIXME(luccafong): cudagraph with v32 MTP is not supported, - # remove this when the issue is fixed. - self.enforce_eager = True # use the draft model from the same model: self.model = self.target_model_config.model # Align the quantization of draft model for cases such as diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index cee60e99b635..ade8a666cf27 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -69,7 +69,9 @@ DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset( { "DeepseekV2ForCausalLM", + "DeepseekV32ForCausalLM", "DeepseekV4ForCausalLM", + "GlmMoeDsaForCausalLM", "GraniteMoeForCausalLM", "InklingForCausalLM", "InklingForConditionalGeneration", @@ -79,6 +81,23 @@ } ) +DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES = frozenset( + { + "DeepseekV32MTPModel", + "DeepseekV32ForCausalLM", + "DeepseekV4ForCausalLM", + "DeepSeekV4MTPModel", + "GlmMoeDsaForCausalLM", + "InklingForCausalLM", + "InklingForConditionalGeneration", + "KimiK3ForConditionalGeneration", + "KimiK3MTPModel", + "KimiLinearForCausalLM", + "MiniMaxM3SparseForCausalLM", + "MiniMaxM3SparseForConditionalGeneration", + } +) + @lru_cache def default_v2_model_runner_architectures() -> frozenset[str]: @@ -86,13 +105,28 @@ def default_v2_model_runner_architectures() -> frozenset[str]: from vllm.platforms import current_platform if current_platform.is_rocm(): - # TODO(rocm): DeepSeek V4 is still faster on MRV1 on ROCm. The - # attention layer picks the eager cudagraph region MRV1 needs, so - # this is a perf default only; drop it once MRV2 catches up. - return DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES - {"DeepseekV4ForCausalLM"} + # TODO(rocm): These models are either unsupported by MRV2 or slower with + # MRV2 on AMD GPUs. + return DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES - { + "DeepseekV32ForCausalLM", + "DeepseekV4ForCausalLM", + } return DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES +@lru_cache +def default_breakable_cudagraph_architectures() -> frozenset[str]: + """Architectures defaulting to breakable CUDA graphs on this platform.""" + from vllm.platforms import current_platform + + if current_platform.is_rocm(): + return DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES - { + "DeepseekV32ForCausalLM", + "DeepseekV32MTPModel", + } + return DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES + + class OptimizationLevel(IntEnum): """Optimization level enum.""" @@ -691,6 +725,34 @@ def _is_default_v2_model_runner_model(self) -> bool: return False return is_default_v2_architecture or not model_config.is_moe + def _uses_breakable_cudagraph_by_default(self) -> bool: + model_config = self.model_config + if model_config is None: + return False + + architectures = set(model_config.architectures) + return bool(architectures & default_breakable_cudagraph_architectures()) + + def _maybe_enable_breakable_cudagraph(self) -> bool: + if ( + "VLLM_USE_BREAKABLE_CUDAGRAPH" not in os.environ + and self._uses_breakable_cudagraph_by_default() + ): + os.environ["VLLM_USE_BREAKABLE_CUDAGRAPH"] = "1" + logger.info_once( + "Auto-enabling VLLM_USE_BREAKABLE_CUDAGRAPH=1. " + "Set VLLM_USE_BREAKABLE_CUDAGRAPH=0 to opt out." + ) + + from vllm.compilation.breakable_cudagraph import ( + is_breakable_cudagraph_enabled, + ) + + enabled = is_breakable_cudagraph_enabled() + if enabled: + self.compilation_config.mode = CompilationMode.NONE + return enabled + @property def needs_dp_coordinator(self) -> bool: """ @@ -1288,41 +1350,7 @@ def __post_init__(self): ) self.compilation_config.mode = CompilationMode.NONE - # For model classes don't carry @support_torch_compile — - # the breakable cudagraph is the supported PIECEWISE path. Auto-enable - # it unless the user has explicitly opted out via the env var. - if ( - self.model_config is not None - and "VLLM_USE_BREAKABLE_CUDAGRAPH" not in os.environ - and any( - a - in ( - "DeepseekV4ForCausalLM", - "DeepSeekV4MTPModel", - "InklingForCausalLM", - "InklingForConditionalGeneration", - "KimiK3ForConditionalGeneration", - "KimiK3MTPModel", - "KimiLinearForCausalLM", - "MiniMaxM3SparseForCausalLM", - "MiniMaxM3SparseForConditionalGeneration", - ) - for a in self.model_config.architectures - ) - ): - os.environ["VLLM_USE_BREAKABLE_CUDAGRAPH"] = "1" - logger.info_once( - "Auto-enabling VLLM_USE_BREAKABLE_CUDAGRAPH=1. " - "Set VLLM_USE_BREAKABLE_CUDAGRAPH=0 to opt out." - ) - - from vllm.compilation.breakable_cudagraph import ( - is_breakable_cudagraph_enabled, - ) - - breakable_cudagraph_enabled = is_breakable_cudagraph_enabled() - if breakable_cudagraph_enabled: - self.compilation_config.mode = CompilationMode.NONE + breakable_cudagraph_enabled = self._maybe_enable_breakable_cudagraph() if not breakable_cudagraph_enabled and ( self.compilation_config.backend == "eager" diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 95877c319d1f..027dc10e6d53 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -92,7 +92,10 @@ "DeepseekForCausalLM": ("deepseek_v2", "DeepseekForCausalLM"), "DeepseekV2ForCausalLM": ("deepseek_v2", "DeepseekV2ForCausalLM"), "DeepseekV3ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"), - "DeepseekV32ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"), + "DeepseekV32ForCausalLM": ( + "vllm.models.deepseek_v32", + "DeepseekV32ForCausalLM", + ), "DeepseekV4ForCausalLM": ("vllm.models.deepseek_v4", "DeepseekV4ForCausalLM"), "Ernie4_5ForCausalLM": ("ernie45", "Ernie4_5ForCausalLM"), "Ernie4_5_MoeForCausalLM": ("ernie45_moe", "Ernie4_5_MoeForCausalLM"), @@ -115,7 +118,7 @@ "Glm4ForCausalLM": ("glm4", "Glm4ForCausalLM"), "Glm4MoeForCausalLM": ("glm4_moe", "Glm4MoeForCausalLM"), "Glm4MoeLiteForCausalLM": ("glm4_moe_lite", "Glm4MoeLiteForCausalLM"), - "GlmMoeDsaForCausalLM": ("deepseek_v2", "GlmMoeDsaForCausalLM"), + "GlmMoeDsaForCausalLM": ("vllm.models.deepseek_v32", "GlmMoeDsaForCausalLM"), "GptOssForCausalLM": ("gpt_oss", "GptOssForCausalLM"), "GPT2LMHeadModel": ("gpt2", "GPT2LMHeadModel"), "GPTJForCausalLM": ("gpt_j", "GPTJForCausalLM"), @@ -661,6 +664,7 @@ "Eagle3DeepseekV3ForCausalLM": ("deepseek_eagle3", "Eagle3DeepseekV2ForCausalLM"), "EagleDeepSeekMTPModel": ("deepseek_eagle", "EagleDeepseekV3ForCausalLM"), "DeepSeekMTPModel": ("deepseek_mtp", "DeepSeekMTP"), + "DeepseekV32MTPModel": ("vllm.models.deepseek_v32", "DeepseekV32MTP"), "Dots3NoteMTPModel": ("vllm.models.dots3_note", "Dots3NoteMTP"), "DeepSeekV4MTPModel": ("vllm.models.deepseek_v4", "DeepSeekV4MTP"), "BailingMoeV3MTPModel": ("bailing_moe_v3_mtp", "BailingMoeV3MTPModel"), diff --git a/vllm/models/deepseek_v32/__init__.py b/vllm/models/deepseek_v32/__init__.py index f3eba73142b4..0ac154505f2e 100644 --- a/vllm/models/deepseek_v32/__init__.py +++ b/vllm/models/deepseek_v32/__init__.py @@ -1,26 +1,35 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""DeepSeek V3.2 (``deepseek_v32``) model — hardware-isolated entry point. +"""DeepSeek V3.2 (``deepseek_v32``) platform entry point. DeepSeek V3.2 introduced the DeepSeek Sparse Attention (DSA) architecture: MLA + a "lightning indexer" that selects the top-k tokens for a sparse MLA attend. The same model code serves any DSA checkpoint, including GLM-5.2 (``glm_moe_dsa``), which reuses this architecture. + +The CUDA implementation selects capability-specific kernels internally and +falls back when an optimization is unavailable. Other platforms use the +generic implementation by default. """ from vllm.platforms import current_platform -if current_platform.is_rocm(): - from .amd.model import DeepseekV32ForCausalLM - from .amd.mtp import DeepseekV32MTP -elif current_platform.is_xpu(): - raise NotImplementedError("deepseek_v32 does not yet support XPU.") -else: - # Covers Blackwell (sm100) and all other CUDA devices. +if current_platform.is_cuda(): + # GLM-5.2 (glm_moe_dsa) reuses the CUDA DSA module. Individual optimized + # kernels remain gated on the device capabilities they support. from .nvidia.model import DeepseekV32ForCausalLM + from .nvidia.model import DeepseekV32ForCausalLM as GlmMoeDsaForCausalLM from .nvidia.mtp import DeepseekV32MTP +else: + # ROCm, XPU, and CPU keep the generic implementation. + from vllm.model_executor.models.deepseek_mtp import DeepSeekMTP as DeepseekV32MTP + from vllm.model_executor.models.deepseek_v2 import ( + DeepseekV3ForCausalLM as DeepseekV32ForCausalLM, + ) + from vllm.model_executor.models.deepseek_v2 import GlmMoeDsaForCausalLM __all__ = [ "DeepseekV32ForCausalLM", "DeepseekV32MTP", + "GlmMoeDsaForCausalLM", ] diff --git a/vllm/models/deepseek_v32/amd/rocm.py b/vllm/models/deepseek_v32/amd/rocm.py index 0297fc353863..6d8e7ea855b8 100644 --- a/vllm/models/deepseek_v32/amd/rocm.py +++ b/vllm/models/deepseek_v32/amd/rocm.py @@ -38,7 +38,6 @@ class DeepseekV32ROCmIndexer(DeepseekV32Indexer): class DeepseekV32MLAAttention(DeepseekV32Attention): - require_fp8_kv_cache: bool = False indexer_cls = DeepseekV32ROCmIndexer def __init__(self, vllm_config, config, prefix, topk_indices_buffer=None): diff --git a/vllm/models/deepseek_v32/attention.py b/vllm/models/deepseek_v32/attention.py index 253971b7bd05..02582810ca02 100644 --- a/vllm/models/deepseek_v32/attention.py +++ b/vllm/models/deepseek_v32/attention.py @@ -170,7 +170,6 @@ class DeepseekV32Attention(MLAAttention): indexer: "DeepseekV32Indexer | None" indexer_cls: "type[DeepseekV32Indexer]" = DeepseekV32Indexer - require_fp8_kv_cache: bool = True supports_dense_mha_prefill = False def __init__( @@ -283,21 +282,9 @@ def __init__( self.layer_name if enable_short_prefill_scoring_skip else "" ) - if self.require_fp8_kv_cache: - assert is_quantized_kv_cache(self.kv_cache_dtype), ( - "deepseek_v32 (nvidia) requires an fp8 KV cache served by a sparse " - "MLA backend. Launch with --kv-cache-dtype fp8 (FlashInfer sparse) " - "or --kv-cache-dtype fp8_ds_mla (FlashMLA sparse)." - ) - self._fp8_query = self.impl.supports_quant_query_input - if not self._fp8_query: - assert self.kv_cache_dtype == "fp8_ds_mla", ( - "deepseek_v32 (nvidia) on a bf16-query sparse MLA backend " - "(FlashMLA sparse) requires the fp8_ds_mla KV cache layout. " - "Launch with --kv-cache-dtype fp8_ds_mla." - ) - - self._fp8_kv_needs_view = self.kv_cache_dtype != "fp8_ds_mla" + fp8_attention = is_quantized_kv_cache(self.kv_cache_dtype) + self._fp8_query = fp8_attention and self.impl.supports_quant_query_input + self._fp8_kv_needs_view = fp8_attention and self.kv_cache_dtype != "fp8_ds_mla" self._index_rope_interleave = getattr(config, "indexer_rope_interleave", False) diff --git a/vllm/models/deepseek_v32/common/kernels.py b/vllm/models/deepseek_v32/common/kernels.py index c5c6c97d55a9..31315172482e 100644 --- a/vllm/models/deepseek_v32/common/kernels.py +++ b/vllm/models/deepseek_v32/common/kernels.py @@ -6,6 +6,7 @@ from vllm.platforms import current_platform from vllm.triton_utils import tl, triton +from vllm.utils.torch_utils import is_quantized_kv_cache # Cache of tiny 1-element dummy tensors (per device, dtype) reused by the # has_indexer=False path so the indexer args don't allocate every call. @@ -473,7 +474,7 @@ def fused_norm_rope( # --- MLA KV cache setup --- mla_cache_ds_mla = mla_kv_cache_dtype == "fp8_ds_mla" - mla_cache_fp8 = mla_kv_cache_dtype not in ("auto", "fp8_ds_mla") + mla_cache_fp8 = is_quantized_kv_cache(mla_kv_cache_dtype) and not mla_cache_ds_mla mla_num_tiles = 1 mla_ds_scale_view = torch.empty(0, dtype=torch.float32, device=device) mla_ds_rope_view = torch.empty(0, dtype=torch.bfloat16, device=device) diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index 29a95caa88eb..7a67a00084ab 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -1018,7 +1018,11 @@ def model_returns_tuple(self) -> bool: # feedback into the next draft step. architectures = self.draft_model_config.hf_config.architectures or [] return bool( - {"DeepSeekMTPModel", "KimiK3MTPModel"}.intersection(architectures) + { + "DeepSeekMTPModel", + "DeepseekV32MTPModel", + "KimiK3MTPModel", + }.intersection(architectures) ) return self.method not in ("mtp", "draft_model", "dflash") From 842dd8fd96650063e1ad32e6075742d457d39773 Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Wed, 19 Aug 2026 02:21:08 -0500 Subject: [PATCH 136/839] [Pooling] Use semantic task validation errors (#52867) Signed-off-by: Taneem Ibrahim --- vllm/entrypoints/pooling/pooling/serving.py | 20 +++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/vllm/entrypoints/pooling/pooling/serving.py b/vllm/entrypoints/pooling/pooling/serving.py index e049d1ba18f5..218052d19e7b 100644 --- a/vllm/entrypoints/pooling/pooling/serving.py +++ b/vllm/entrypoints/pooling/pooling/serving.py @@ -5,6 +5,7 @@ from fastapi.responses import JSONResponse, Response, StreamingResponse from typing_extensions import assert_never +from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger from vllm.outputs import PoolingRequestOutput from vllm.tasks import SupportedTask @@ -60,7 +61,9 @@ def get_io_processor(self, request: AnyPoolingRequest) -> PoolingIOProcessor: def _verify_pooling_task(self, request: PoolingRequest) -> str: if getattr(request, "dimensions", None) is not None: - raise ValueError("dimensions is currently not supported") + raise VLLMValidationError( + "dimensions is currently not supported", parameter="dimensions" + ) if request.task is None: request.task = self.pooling_task @@ -74,22 +77,25 @@ def _verify_pooling_task(self, request: PoolingRequest) -> str: # plugin task uses io_processor.parse_request to verify inputs if pooling_task != "plugin" and pooling_task != self.pooling_task: if pooling_task not in self.supported_tasks: - raise ValueError( + raise VLLMValidationError( f"Unsupported task: {pooling_task!r} " - f"Supported tasks: {self.supported_tasks}" + f"Supported tasks: {self.supported_tasks}", + parameter="task", ) else: - raise ValueError( + raise VLLMValidationError( "Try switching the model's pooling_task " - f"via --pooler-config.task {request.task}." + f"via --pooler-config.task {request.task}.", + parameter="task", ) if pooling_task == "plugin" and "plugin" not in self.io_processors: - raise ValueError( + raise VLLMValidationError( "No IOProcessor plugin installed. Please refer " "to the documentation and to the " "'prithvi_geospatial_mae_io_processor' " - "offline inference example for more details." + "offline inference example for more details.", + parameter="task", ) return pooling_task From 9a9aa2b017b468082ca538b6f7a725a60d7f8b47 Mon Sep 17 00:00:00 2001 From: LuckyStep <67696304+Andy365-365@users.noreply.github.com> Date: Wed, 19 Aug 2026 15:52:53 +0800 Subject: [PATCH 137/839] [Bugfix] Redact api_key in non-default args log (#52523) Signed-off-by: LuckyStep <67696304+Andy365-365@users.noreply.github.com> --- .../entrypoints/serve/utils/test_api_utils.py | 37 +++++++++++++++++++ tests/test_envs.py | 6 +++ vllm/entrypoints/serve/utils/api_utils.py | 16 +++++++- vllm/envs.py | 3 ++ 4 files changed, 61 insertions(+), 1 deletion(-) diff --git a/tests/entrypoints/serve/utils/test_api_utils.py b/tests/entrypoints/serve/utils/test_api_utils.py index c23d0e9e6828..669bff18d333 100644 --- a/tests/entrypoints/serve/utils/test_api_utils.py +++ b/tests/entrypoints/serve/utils/test_api_utils.py @@ -1,10 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from argparse import Namespace + import pytest from vllm.entrypoints.openai.engine.protocol import StreamOptions +from vllm.entrypoints.serve.utils import api_utils from vllm.entrypoints.serve.utils.api_utils import ( + _redact_sensitive_args, get_max_tokens, should_include_usage, ) @@ -110,3 +114,36 @@ def test_input_length_exceeds_max_model_len(self): input_length=150, default_sampling_params={"max_tokens": 2048}, ) + + +class TestRedactSensitiveArgs: + API_KEY = "sk-test-secret-12345" + + def test_redact_replaces_sensitive_values_only(self): + args = {"api_key": self.API_KEY, "other": "visible"} + redacted = _redact_sensitive_args(args) + assert redacted == {"api_key": "***", "other": "visible"} + # original dict must not be mutated + assert args == {"api_key": self.API_KEY, "other": "visible"} + + def test_no_sensitive_fields_returns_original(self): + args = {"model_tag": "org/model", "other": "visible"} + assert _redact_sensitive_args(args) is args + + def test_api_key_not_in_log(self, monkeypatch, caplog): + non_default = { + "model_tag": "org/model", + "default_chat_template_kwargs": {"enable_thinking": False}, + "api_key": self.API_KEY, + "enable_auto_tool_choice": True, + "tool_call_parser": "qwen3_coder", + } + monkeypatch.setattr(api_utils, "get_non_default_args", lambda args: non_default) + with caplog.at_level("INFO", logger="vllm.entrypoints.serve.utils.api_utils"): + api_utils.log_non_default_args(args=Namespace()) + message = caplog.text + assert self.API_KEY not in message + assert "'api_key': '***'" in message + # non-sensitive args are still logged + assert "org/model" in message + assert "qwen3_coder" in message diff --git a/tests/test_envs.py b/tests/test_envs.py index 3d214fdbff81..1ef40cde0cc0 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -37,6 +37,12 @@ def test_nixl_side_channel_host_is_not_compile_factor( assert "VLLM_NIXL_SIDE_CHANNEL_HOST" not in envs.compile_factors() +def test_api_key_is_not_compile_factor(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("VLLM_API_KEY", "sk-super-secret") + + assert "VLLM_API_KEY" not in envs.compile_factors() + + def test_p2p_side_channel_defaults_and_override(monkeypatch: pytest.MonkeyPatch): monkeypatch.delenv("VLLM_P2P_SIDE_CHANNEL_HOST", raising=False) monkeypatch.delenv("VLLM_P2P_SIDE_CHANNEL_PORT", raising=False) diff --git a/vllm/entrypoints/serve/utils/api_utils.py b/vllm/entrypoints/serve/utils/api_utils.py index c75cd16e54f4..f186ec862b97 100644 --- a/vllm/entrypoints/serve/utils/api_utils.py +++ b/vllm/entrypoints/serve/utils/api_utils.py @@ -267,9 +267,23 @@ def jsonify_non_default_args( return {key: _jsonify_arg_value(value) for key, value in non_default_args.items()} +# Fields whose values must never be logged verbatim. +_SENSITIVE_ARG_FIELDS = frozenset({"api_key"}) + + +def _redact_sensitive_args(args: dict[str, Any]) -> dict[str, Any]: + """Return a copy of `args` with sensitive values redacted for logging.""" + if not any(key in _SENSITIVE_ARG_FIELDS for key in args): + return args + return { + key: ("***" if key in _SENSITIVE_ARG_FIELDS else value) + for key, value in args.items() + } + + def log_non_default_args(args: Namespace | EngineArgs): non_default_args = get_non_default_args(args) - logger.info("non-default args: %s", non_default_args) + logger.info("non-default args: %s", _redact_sensitive_args(non_default_args)) def should_include_usage( diff --git a/vllm/envs.py b/vllm/envs.py index f80c51207fd6..b49841d9fff0 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -2243,6 +2243,9 @@ def compile_factors() -> dict[str, object]: "S3_ACCESS_KEY_ID", "S3_SECRET_ACCESS_KEY", "S3_ENDPOINT_URL", + # Credential; never affects compiled artifacts and must not be + # persisted in cache_key_factors.json. + "VLLM_API_KEY", "VLLM_USAGE_STATS_SERVER", "VLLM_NO_USAGE_STATS", "VLLM_DO_NOT_TRACK", From 93eea4f665a0a6a188fd52b37f4a18bc37315f62 Mon Sep 17 00:00:00 2001 From: Qiming Zhang Date: Wed, 19 Aug 2026 02:01:12 -0700 Subject: [PATCH 138/839] [XPU][CI] downgrade sentencepiece (#52904) Signed-off-by: mayuyuace --- requirements/test/xpu.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements/test/xpu.txt b/requirements/test/xpu.txt index c4376e00b056..60ffa0ae61e9 100644 --- a/requirements/test/xpu.txt +++ b/requirements/test/xpu.txt @@ -808,7 +808,7 @@ scipy==1.18.0 # sentence-transformers sentence-transformers==5.7.0 # via mteb -sentencepiece==0.2.2 +sentencepiece==0.2.1 # via -r requirements/test/../common.txt sentry-sdk==2.68.0 # via fastapi-cloud-cli From 340b7e4909f14ff33943ab6ac7955191e1b8b4ff Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Wed, 19 Aug 2026 17:29:32 +0800 Subject: [PATCH 139/839] fix: reject string schemas that mix pattern/format with length bounds (#49996) Signed-off-by: Yufeng He <40085740+he-yufeng@users.noreply.github.com> --- tests/v1/structured_output/test_utils.py | 6 ++++++ vllm/v1/structured_output/backend_xgrammar.py | 14 ++++++++++++++ 2 files changed, 20 insertions(+) diff --git a/tests/v1/structured_output/test_utils.py b/tests/v1/structured_output/test_utils.py index c026ab0e4e78..11b3e5c01b38 100644 --- a/tests/v1/structured_output/test_utils.py +++ b/tests/v1/structured_output/test_utils.py @@ -14,6 +14,12 @@ def unsupported_string_schemas(): return [ {"type": "string", "format": "non_existing_format"}, + # pattern/format is compiled but length bounds are silently dropped, + # so the combination must be rejected instead of producing quietly + # wrong output + {"type": "string", "pattern": "^a+$", "maxLength": 2}, + {"type": "string", "pattern": "^a+$", "minLength": 3}, + {"type": "string", "format": "email", "maxLength": 10}, ] diff --git a/vllm/v1/structured_output/backend_xgrammar.py b/vllm/v1/structured_output/backend_xgrammar.py index 5b24a19780aa..3809ba2d91b0 100644 --- a/vllm/v1/structured_output/backend_xgrammar.py +++ b/vllm/v1/structured_output/backend_xgrammar.py @@ -262,6 +262,20 @@ def check_object(obj: dict[str, Any]) -> bool: ): return True + # A string mixing a generative constraint (pattern or format) with + # explicit length bounds. xgrammar compiles the pattern/format side + # and silently drops minLength/maxLength from the grammar, so output + # can violate the bound without any error surfacing. Verified against + # the compiled EBNF: pattern/format grammars come out byte-identical + # with and without the length keywords, while maxLength alone lowers + # to {0, N} correctly. + if ( + obj.get("type") == "string" + and ("pattern" in obj or "format" in obj) + and ("minLength" in obj or "maxLength" in obj) + ): + return True + # Unsupported keywords for objects if obj.get("type") == "object" and any( key in obj for key in ("patternProperties", "propertyNames") From eac636a7fa476983cdae34b45a984e9852aad375 Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Wed, 19 Aug 2026 17:43:46 +0800 Subject: [PATCH 140/839] [Frontend] Move api_server.py out openai folder (#52131) Signed-off-by: wang.yuqi --- .buildkite/test-amd.yaml | 10 +- .buildkite/test_areas/entrypoints.yaml | 2 + .buildkite/test_areas/rust_frontend.yaml | 8 +- docs/deployment/integrations/kthena.md | 2 +- docs/design/arch_overview.md | 16 +- docs/design/endpoint_plugins.md | 2 +- examples/applications/api_server/server.py | 5 +- tests/entrypoints/launchers/__init__.py | 0 .../launchers/api_server/__init__.py | 0 .../api_server}/_api_server_spawn_workers.py | 0 .../test_api_server_process_manager.py | 2 +- .../api_server}/test_multi_api_servers.py | 0 .../test_launch_cli.py | 0 .../completion => launchers}/test_shutdown.py | 2 +- .../test_ssl_cert_refresher.py | 2 +- .../test_http_status_metrics.py | 2 +- .../instrumentator}/test_uds.py | 3 +- tests/lora/test_add_lora.py | 2 +- tests/lora/test_lora_functions.py | 2 +- tests/plugins_tests/test_endpoint_plugins.py | 20 +- vllm/benchmarks/throughput.py | 2 +- vllm/entrypoints/cli/launch.py | 44 +- vllm/entrypoints/cli/serve.py | 6 +- .../launchers/api_server/app_state.py | 149 ++++ .../entrypoints/launchers/api_server/entry.py | 229 ++++++ vllm/entrypoints/launchers/app.py | 56 ++ vllm/entrypoints/{ => launchers}/launcher.py | 98 ++- vllm/entrypoints/launchers/render/__init__.py | 0 .../entrypoints/launchers/render/app_state.py | 104 +++ vllm/entrypoints/launchers/render/entry.py | 120 +++ vllm/entrypoints/launchers/utils/__init__.py | 0 .../{serve => launchers}/utils/constants.py | 5 - .../utils/server_utils.py | 2 +- .../{serve => launchers}/utils/ssl.py | 0 vllm/entrypoints/openai/api_server.py | 703 ++---------------- vllm/entrypoints/openai/cli_args.py | 6 +- vllm/entrypoints/openai/dp_supervisor.py | 4 +- vllm/entrypoints/openai/responses/context.py | 2 +- vllm/entrypoints/openai/run_batch.py | 7 +- .../exception_handling/handlers/vllm_error.py | 2 +- vllm/plugins/endpoint_plugins/interface.py | 40 +- vllm/tasks.py | 1 + vllm/v1/utils.py | 2 +- 43 files changed, 887 insertions(+), 775 deletions(-) create mode 100644 tests/entrypoints/launchers/__init__.py create mode 100644 tests/entrypoints/launchers/api_server/__init__.py rename tests/entrypoints/{unit_tests => launchers/api_server}/_api_server_spawn_workers.py (100%) rename tests/entrypoints/{unit_tests => launchers/api_server}/test_api_server_process_manager.py (99%) rename tests/entrypoints/{openai => launchers/api_server}/test_multi_api_servers.py (100%) rename tests/entrypoints/{unit_tests => launchers}/test_launch_cli.py (100%) rename tests/entrypoints/{openai/completion => launchers}/test_shutdown.py (99%) rename tests/entrypoints/{serve/utils => launchers}/test_ssl_cert_refresher.py (97%) rename tests/entrypoints/{openai => serve/instrumentator}/test_uds.py (96%) create mode 100644 vllm/entrypoints/launchers/api_server/app_state.py create mode 100644 vllm/entrypoints/launchers/api_server/entry.py create mode 100644 vllm/entrypoints/launchers/app.py rename vllm/entrypoints/{ => launchers}/launcher.py (64%) create mode 100644 vllm/entrypoints/launchers/render/__init__.py create mode 100644 vllm/entrypoints/launchers/render/app_state.py create mode 100644 vllm/entrypoints/launchers/render/entry.py create mode 100644 vllm/entrypoints/launchers/utils/__init__.py rename vllm/entrypoints/{serve => launchers}/utils/constants.py (80%) rename vllm/entrypoints/{serve => launchers}/utils/server_utils.py (98%) rename vllm/entrypoints/{serve => launchers}/utils/ssl.py (100%) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 2f13ff83e9f2..bb454f7f6adc 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -990,10 +990,12 @@ steps: - vllm/entrypoints - tests/entrypoints/unit_tests - tests/entrypoints/weight_transfer + - tests/entrypoints/launchers - vllm/platforms/rocm.py commands: - pytest -v -s entrypoints/unit_tests - pytest -v -s entrypoints/weight_transfer + - pytest -v -s entrypoints/launchers - label: Entrypoints Integration (LLM) # TBD timeout_in_minutes: 180 @@ -2320,9 +2322,9 @@ steps: - tests/benchmarks/test_serve_cli.py - tests/entrypoints/openai/chat_completion/test_chat_completion.py - tests/entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py - - tests/entrypoints/openai/completion/test_shutdown.py + - tests/entrypoints/launchers/test_shutdown.py - tests/entrypoints/openai/test_return_token_ids.py - - tests/entrypoints/openai/test_uds.py + - tests/entrypoints/serve/instrumentator/test_uds.py - tests/v1/sample/test_logprobs_e2e.py - vllm/platforms/rocm.py commands: @@ -2331,9 +2333,9 @@ steps: - pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)" - pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py -k "not test_invalid_json_schema and not test_invalid_regex and not test_kv_transfer_prompt_token_ids_round_trip and not test_kv_transfer_prompt_token_ids_streaming" - pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not multiple" - - pytest -v -s entrypoints/openai/completion/test_shutdown.py -k "not engine_failure and not test_abort_timeout_exits_quickly" + - pytest -v -s entrypoints/launchers/test_shutdown.py -k "not engine_failure and not test_abort_timeout_exits_quickly" - pytest -v -s entrypoints/openai/test_return_token_ids.py -k "not test_comparison" - - pytest -v -s entrypoints/openai/test_uds.py + - pytest -v -s entrypoints/serve/instrumentator/test_uds.py - pytest -v -s v1/sample/test_logprobs_e2e.py -k "test_prompt_logprobs_e2e_server" - label: Rust Frontend Serve Admin Coverage # TBD diff --git a/.buildkite/test_areas/entrypoints.yaml b/.buildkite/test_areas/entrypoints.yaml index 56d4e012581e..b75a21ea1195 100644 --- a/.buildkite/test_areas/entrypoints.yaml +++ b/.buildkite/test_areas/entrypoints.yaml @@ -11,9 +11,11 @@ steps: - vllm/entrypoints - tests/entrypoints/unit_tests - tests/entrypoints/weight_transfer + - tests/entrypoints/launchers commands: - pytest -v -s entrypoints/unit_tests - pytest -v -s entrypoints/weight_transfer + - pytest -v -s entrypoints/launchers mirror: amd: label: ":amd: (MI355) Entrypoints Unit" diff --git a/.buildkite/test_areas/rust_frontend.yaml b/.buildkite/test_areas/rust_frontend.yaml index 36b7d375d4ff..21fcd127dbdb 100644 --- a/.buildkite/test_areas/rust_frontend.yaml +++ b/.buildkite/test_areas/rust_frontend.yaml @@ -18,9 +18,9 @@ steps: - tests/entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py # - tests/entrypoints/openai/completion/test_prompt_validation.py - - tests/entrypoints/openai/completion/test_shutdown.py + - tests/entrypoints/launchers/test_shutdown.py - tests/entrypoints/openai/test_return_token_ids.py - - tests/entrypoints/openai/test_uds.py + - tests/entrypoints/serve/instrumentator/test_uds.py - tests/v1/sample/test_logprobs_e2e.py commands: - export VLLM_USE_RUST_FRONTEND=1 @@ -30,11 +30,11 @@ steps: - pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not multiple" # - pytest -v -s entrypoints/openai/completion/test_prompt_validation.py -k "not prompt_embeds" - - pytest -v -s entrypoints/openai/completion/test_shutdown.py -k "not engine_failure and not test_abort_timeout_exits_quickly" + - pytest -v -s entrypoints/launchers/test_shutdown.py -k "not engine_failure and not test_abort_timeout_exits_quickly" # test_comparison streams differently: Rust emits a separate first (prompt_token_ids) chunk and # finish chunk without logprobs, while the test reads `logprobs.tokens` on every chunk. - pytest -v -s entrypoints/openai/test_return_token_ids.py -k "not test_comparison" - - pytest -v -s entrypoints/openai/test_uds.py + - pytest -v -s entrypoints/serve/instrumentator/test_uds.py - pytest -v -s v1/sample/test_logprobs_e2e.py -k "test_prompt_logprobs_e2e_server" - label: ":nvidia: (H200) Rust Frontend Serve/Admin Coverage" diff --git a/docs/deployment/integrations/kthena.md b/docs/deployment/integrations/kthena.md index 7cc3f14a71eb..0285b733c6b8 100644 --- a/docs/deployment/integrations/kthena.md +++ b/docs/deployment/integrations/kthena.md @@ -117,7 +117,7 @@ Commands: - > bash /vllm-workspace/examples/ray_serving/multi-node-serving.sh leader --ray_cluster_size=2; python3 -m - vllm.entrypoints.openai.api_server --port 8080 --model + vllm.entrypoints.launchers.api_server.entry --port 8080 --model meta-llama/Llama-3.1-405B-Instruct --tensor-parallel-size 8 --pipeline-parallel-size 2 ``` diff --git a/docs/design/arch_overview.md b/docs/design/arch_overview.md index c19ea49d96db..711e9a2668eb 100644 --- a/docs/design/arch_overview.md +++ b/docs/design/arch_overview.md @@ -62,20 +62,6 @@ vllm serve The code for the `vllm` CLI can be found in [vllm/entrypoints/cli/main.py](../../vllm/entrypoints/cli/main.py). -Sometimes you may see the API server entrypoint used directly instead of via the -`vllm` CLI command. For example: - -```bash -python -m vllm.entrypoints.openai.api_server --model -``` - -!!! warning - - `python -m vllm.entrypoints.openai.api_server` is deprecated - and may become unsupported in a future release. - -That code can be found in [vllm/entrypoints/openai/api_server.py](../../vllm/entrypoints/openai/api_server.py). - More details on the API server can be found in the [Online Serving](../serving/online_serving/README.md) document. ## V1 Process Architecture @@ -88,7 +74,7 @@ The API server process handles HTTP requests (e.g., the OpenAI-compatible API), By default, there is **1 API server process**, but when data parallelism is used, the API server count automatically scales to match the data parallel size. This can also be manually configured with the `--api-server-count` flag. Each API server connects to **all** engine cores via ZMQ in a many-to-many topology, enabling any API server to route requests to any engine core. Each API server process uses multiple CPU threads for media loading (controlled by `VLLM_MEDIA_LOADING_THREAD_COUNT`, default 8). -The code can be found in [vllm/entrypoints/openai/api_server.py](../../vllm/entrypoints/openai/api_server.py) and [vllm/v1/utils.py](../../vllm/v1/utils.py). +The code can be found in [vllm/entrypoints/launchers/api_server](../../vllm/entrypoints/launchers/api_server) and [vllm/v1/utils.py](../../vllm/v1/utils.py). ### Engine Core Process diff --git a/docs/design/endpoint_plugins.md b/docs/design/endpoint_plugins.md index 9f38fe5da188..98b755788e2e 100644 --- a/docs/design/endpoint_plugins.md +++ b/docs/design/endpoint_plugins.md @@ -1,6 +1,6 @@ # Endpoint Plugins -Endpoint plugins let out-of-tree packages add HTTP routes to the OpenAI compatible API server without editing `vllm/entrypoints/openai/api_server.py`. Their scope is +Endpoint plugins let out-of-tree packages add HTTP routes to the OpenAI compatible API server. Their scope is the **HTTP surface only** registering routes and optionally per app state used by those routes. A plugin reaches the engine the same way an in-tree serving handler does, through the `EngineClient` it is handed at startup (e.g. `engine_client.collective_rpc(...)`). No new engine access path is introduced. !!! warning "Security" diff --git a/examples/applications/api_server/server.py b/examples/applications/api_server/server.py index adac4133210e..e56a7931f666 100644 --- a/examples/applications/api_server/server.py +++ b/examples/applications/api_server/server.py @@ -4,8 +4,7 @@ NOTE: This API server is used only for demonstrating usage of AsyncEngine and simple performance benchmarks. It is not intended for production use. For production use, we recommend using our OpenAI compatible server. -We are also not going to accept PRs modifying this file, please -change `vllm/entrypoints/openai/api_server.py` instead. +We are also not going to accept PRs modifying this file. """ import asyncio @@ -21,7 +20,7 @@ import vllm.envs as envs from vllm.engine.arg_utils import AsyncEngineArgs from vllm.engine.async_llm_engine import AsyncLLMEngine -from vllm.entrypoints.launcher import serve_http +from vllm.entrypoints.launchers.launcher import serve_http from vllm.entrypoints.serve.utils.api_utils import with_cancellation from vllm.logger import init_logger from vllm.sampling_params import SamplingParams diff --git a/tests/entrypoints/launchers/__init__.py b/tests/entrypoints/launchers/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/entrypoints/launchers/api_server/__init__.py b/tests/entrypoints/launchers/api_server/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/entrypoints/unit_tests/_api_server_spawn_workers.py b/tests/entrypoints/launchers/api_server/_api_server_spawn_workers.py similarity index 100% rename from tests/entrypoints/unit_tests/_api_server_spawn_workers.py rename to tests/entrypoints/launchers/api_server/_api_server_spawn_workers.py diff --git a/tests/entrypoints/unit_tests/test_api_server_process_manager.py b/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py similarity index 99% rename from tests/entrypoints/unit_tests/test_api_server_process_manager.py rename to tests/entrypoints/launchers/api_server/test_api_server_process_manager.py index 902d0ce2d744..93d0c4970060 100644 --- a/tests/entrypoints/unit_tests/test_api_server_process_manager.py +++ b/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py @@ -10,7 +10,7 @@ import pytest import zmq -from tests.entrypoints.unit_tests._api_server_spawn_workers import ( +from tests.entrypoints.launchers.api_server._api_server_spawn_workers import ( exit_before_report_worker, ) from vllm.utils.network_utils import make_zmq_socket, split_zmq_path diff --git a/tests/entrypoints/openai/test_multi_api_servers.py b/tests/entrypoints/launchers/api_server/test_multi_api_servers.py similarity index 100% rename from tests/entrypoints/openai/test_multi_api_servers.py rename to tests/entrypoints/launchers/api_server/test_multi_api_servers.py diff --git a/tests/entrypoints/unit_tests/test_launch_cli.py b/tests/entrypoints/launchers/test_launch_cli.py similarity index 100% rename from tests/entrypoints/unit_tests/test_launch_cli.py rename to tests/entrypoints/launchers/test_launch_cli.py diff --git a/tests/entrypoints/openai/completion/test_shutdown.py b/tests/entrypoints/launchers/test_shutdown.py similarity index 99% rename from tests/entrypoints/openai/completion/test_shutdown.py rename to tests/entrypoints/launchers/test_shutdown.py index 4b62ce182264..e3be2919f5ef 100644 --- a/tests/entrypoints/openai/completion/test_shutdown.py +++ b/tests/entrypoints/launchers/test_shutdown.py @@ -147,7 +147,7 @@ async def test_shutdown_on_engine_failure(tmp_path: Path): # dtype, max-len etc set so that this can run in CI sys.executable, "-m", - "vllm.entrypoints.openai.api_server", + "vllm.entrypoints.launchers.api_server.entry", "--model", MODEL_NAME, "--dtype", diff --git a/tests/entrypoints/serve/utils/test_ssl_cert_refresher.py b/tests/entrypoints/launchers/test_ssl_cert_refresher.py similarity index 97% rename from tests/entrypoints/serve/utils/test_ssl_cert_refresher.py rename to tests/entrypoints/launchers/test_ssl_cert_refresher.py index 8f5251374a6b..77d2aa3f3594 100644 --- a/tests/entrypoints/serve/utils/test_ssl_cert_refresher.py +++ b/tests/entrypoints/launchers/test_ssl_cert_refresher.py @@ -7,7 +7,7 @@ import pytest -from vllm.entrypoints.serve.utils.ssl import SSLCertRefresher +from vllm.entrypoints.launchers.utils.ssl import SSLCertRefresher class MockSSLContext(SSLContext): diff --git a/tests/entrypoints/serve/exception_handling/test_http_status_metrics.py b/tests/entrypoints/serve/exception_handling/test_http_status_metrics.py index 060b05bbaca8..9fb1c34b8790 100644 --- a/tests/entrypoints/serve/exception_handling/test_http_status_metrics.py +++ b/tests/entrypoints/serve/exception_handling/test_http_status_metrics.py @@ -14,7 +14,7 @@ from fastapi import HTTPException from prometheus_client import CollectorRegistry -from vllm.entrypoints.openai.api_server import build_app +from vllm.entrypoints.launchers.api_server.entry import build_app from vllm.exceptions import ( VLLMNotFoundError, VLLMServerError, diff --git a/tests/entrypoints/openai/test_uds.py b/tests/entrypoints/serve/instrumentator/test_uds.py similarity index 96% rename from tests/entrypoints/openai/test_uds.py rename to tests/entrypoints/serve/instrumentator/test_uds.py index 7c8557141087..7dfb1f2cd12e 100644 --- a/tests/entrypoints/openai/test_uds.py +++ b/tests/entrypoints/serve/instrumentator/test_uds.py @@ -6,11 +6,10 @@ import httpx import pytest +from tests.utils import RemoteOpenAIServer from vllm import envs from vllm.version import __version__ as VLLM_VERSION -from ...utils import RemoteOpenAIServer - MODEL_NAME = "Qwen/Qwen3-0.6B" diff --git a/tests/lora/test_add_lora.py b/tests/lora/test_add_lora.py index 9a82ab99ea9c..628f5bc28eb5 100644 --- a/tests/lora/test_add_lora.py +++ b/tests/lora/test_add_lora.py @@ -6,7 +6,7 @@ import pytest from vllm.engine.arg_utils import AsyncEngineArgs -from vllm.entrypoints.openai.api_server import ( +from vllm.entrypoints.launchers.api_server.entry import ( build_async_engine_client_from_engine_args, ) from vllm.inputs import TextPrompt diff --git a/tests/lora/test_lora_functions.py b/tests/lora/test_lora_functions.py index 1c692630284d..ecf2ad9ee654 100644 --- a/tests/lora/test_lora_functions.py +++ b/tests/lora/test_lora_functions.py @@ -7,7 +7,7 @@ import pytest from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs -from vllm.entrypoints.openai.api_server import ( +from vllm.entrypoints.launchers.api_server.entry import ( build_async_engine_client_from_engine_args, ) from vllm.lora.request import LoRARequest diff --git a/tests/plugins_tests/test_endpoint_plugins.py b/tests/plugins_tests/test_endpoint_plugins.py index 9385d98f2e34..5733333e5754 100644 --- a/tests/plugins_tests/test_endpoint_plugins.py +++ b/tests/plugins_tests/test_endpoint_plugins.py @@ -17,14 +17,14 @@ from fastapi import FastAPI from vllm_add_dummy_endpoint_plugin import DummyAdminEndpointPlugin -from vllm.entrypoints.openai.api_server import ( - _attach_endpoint_plugins, - _init_endpoint_plugins_state, - build_app, -) +from vllm.entrypoints.launchers.app import build_app from vllm.entrypoints.openai.cli_args import make_arg_parser from vllm.plugins import load_endpoint_plugins -from vllm.plugins.endpoint_plugins.interface import EndpointPlugin +from vllm.plugins.endpoint_plugins.interface import ( + EndpointPlugin, + attach_endpoint_plugins, + init_endpoint_plugins_state, +) from vllm.utils.argparse_utils import FlexibleArgumentParser @@ -157,7 +157,7 @@ def test_attach_is_noop_when_nothing_discovered(monkeypatch: pytest.MonkeyPatch) monkeypatch.delenv("VLLM_PLUGINS", raising=False) app = FastAPI() - _attach_endpoint_plugins(app, ("generate",)) + attach_endpoint_plugins(app, ("generate",)) assert app.state.endpoint_plugins == [] @@ -172,7 +172,7 @@ async def test_init_state_is_noop_without_phase_a(monkeypatch: pytest.MonkeyPatc monkeypatch.setenv("VLLM_PLUGINS", "dummy_admin_endpoint_plugin") state = State() - await _init_endpoint_plugins_state(_FakeEngineClient(), state, _build_args()) + await init_endpoint_plugins_state(_FakeEngineClient(), state, _build_args()) assert not hasattr(state, "dummy_engine_client") @@ -196,7 +196,7 @@ async def test_render_server_attaches_endpoint_plugins_with_no_engine_client( for route in app.routes ) - await _init_endpoint_plugins_state(None, app.state, args) + await init_endpoint_plugins_state(None, app.state, args) assert app.state.dummy_engine_client is None @@ -223,7 +223,7 @@ async def test_endpoint_plugin_end_to_end(monkeypatch: pytest.MonkeyPatch): ) fake_engine_client = _FakeEngineClient(rpc_result=["cfg-a", "cfg-b"]) - await _init_endpoint_plugins_state(fake_engine_client, app.state, args) + await init_endpoint_plugins_state(fake_engine_client, app.state, args) assert app.state.dummy_engine_client is fake_engine_client diff --git a/vllm/benchmarks/throughput.py b/vllm/benchmarks/throughput.py index e8a0ccb3a0bc..5b757d43d1cf 100644 --- a/vllm/benchmarks/throughput.py +++ b/vllm/benchmarks/throughput.py @@ -305,7 +305,7 @@ async def run_vllm_async( disable_detokenize: bool = False, warmup_requests: list[SampleRequest] | None = None, ) -> float: - from vllm.entrypoints.openai.api_server import ( + from vllm.entrypoints.launchers.api_server.entry import ( build_async_engine_client_from_engine_args, ) diff --git a/vllm/entrypoints/cli/launch.py b/vllm/entrypoints/cli/launch.py index 91d4bf094b7d..38df03fac6ba 100644 --- a/vllm/entrypoints/cli/launch.py +++ b/vllm/entrypoints/cli/launch.py @@ -3,18 +3,11 @@ import argparse import inspect -import signal import uvloop -from vllm import envs -from vllm.config import VllmConfig -from vllm.engine.arg_utils import AsyncEngineArgs from vllm.entrypoints.cli.types import CLISubcommand -from vllm.entrypoints.openai.api_server import ( - build_and_serve_renderer, - setup_server, -) +from vllm.entrypoints.launchers.render.entry import run_launch_fastapi from vllm.entrypoints.openai.cli_args import ( make_arg_parser, validate_parsed_serve_args, @@ -120,38 +113,3 @@ def subparser_init( def cmd_init() -> list[CLISubcommand]: return [LaunchSubcommand()] - - -async def run_launch_fastapi(args: argparse.Namespace) -> None: - """Run the online serving layer with FastAPI (no GPU inference).""" - - # Interrupt initialization if SIGTERM arrives before uvicorn installs - # its own signal handlers. Once uvicorn is running it replaces this. - def _interrupt_init(*_) -> None: - raise KeyboardInterrupt("terminated") - - signal.signal(signal.SIGTERM, _interrupt_init) - - # 1. Socket binding - listen_address, sock = setup_server(args, reuse_port=False) - - # 2. Build and serve the API server - engine_args = AsyncEngineArgs.from_cli_args(args) - model_config = engine_args.create_model_config() - - # Render servers preprocess data only — no inference, no quantized kernels. - # Clear quantization so VllmConfig skips quant dtype/capability validation. - model_config.quantization = None - - # Render servers never allocate KV cache; suppress the spurious CPU KV - # cache space warning from CpuPlatform.check_and_update_config. - envs.VLLM_CPU_KVCACHE_SPACE = 0 - - vllm_config = VllmConfig(model_config=model_config) - shutdown_task = await build_and_serve_renderer( - vllm_config, listen_address, sock, args - ) - try: - await shutdown_task - finally: - sock.close() diff --git a/vllm/entrypoints/cli/serve.py b/vllm/entrypoints/cli/serve.py index 9f29cd545f32..1d289712abec 100644 --- a/vllm/entrypoints/cli/serve.py +++ b/vllm/entrypoints/cli/serve.py @@ -10,11 +10,9 @@ import vllm import vllm.envs as envs from vllm.entrypoints.cli.types import CLISubcommand -from vllm.entrypoints.openai.api_server import run_server, setup_server +from vllm.entrypoints.launchers.api_server.entry import run_server, setup_server from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args -from vllm.entrypoints.openai.dp_supervisor import ( - run_dp_supervisor, -) +from vllm.entrypoints.openai.dp_supervisor import run_dp_supervisor from vllm.entrypoints.serve.utils.api_utils import VLLM_SUBCMD_PARSER_EPILOG from vllm.logger import init_logger from vllm.usage.usage_lib import UsageContext diff --git a/vllm/entrypoints/launchers/api_server/app_state.py b/vllm/entrypoints/launchers/api_server/app_state.py new file mode 100644 index 000000000000..c896db43920b --- /dev/null +++ b/vllm/entrypoints/launchers/api_server/app_state.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import warnings +from argparse import Namespace +from typing import cast + +from starlette.datastructures import State + +from vllm.engine.protocol import EngineClient +from vllm.entrypoints.chat_utils import load_chat_template +from vllm.entrypoints.openai.models.protocol import BaseModelPath +from vllm.entrypoints.openai.models.serving import OpenAIServingModels +from vllm.entrypoints.serve.tokenize.serving import ServingTokenization +from vllm.entrypoints.serve.utils.api_utils import process_lora_modules +from vllm.entrypoints.serve.utils.request_logger import RequestLogger +from vllm.plugins.endpoint_plugins.interface import init_endpoint_plugins_state +from vllm.renderers.online_derenderer import OnlineDerenderer +from vllm.renderers.online_renderer import OnlineRenderer +from vllm.tasks import FALLBACK_SUPPORTED_TASKS, POOLING_TASKS, SupportedTask + + +async def init_app_state( + engine_client: EngineClient, + state: State, + args: Namespace, + supported_tasks: tuple["SupportedTask", ...] | None = None, +) -> None: + vllm_config = engine_client.vllm_config + + if args.tool_call_parser is not None: + from vllm.parser.metrics import init_parser_metrics + + init_parser_metrics( + model_name=cast(str, vllm_config.model_config.served_model_name) + ) + + if supported_tasks is None: + warnings.warn( + "The 'supported_tasks' parameter was not provided to " + "init_app_state and will be required in a future version. " + "Please pass 'supported_tasks' explicitly.", + DeprecationWarning, + stacklevel=2, + ) + supported_tasks = FALLBACK_SUPPORTED_TASKS + + if args.served_model_name is not None: + served_model_names = args.served_model_name + else: + served_model_names = [args.model] + + if args.enable_log_requests: + request_logger = RequestLogger(max_log_len=args.max_log_len) + else: + request_logger = None + + base_model_paths = [ + BaseModelPath(name=name, model_path=args.model) for name in served_model_names + ] + + state.engine_client = engine_client + state.log_stats = not args.disable_log_stats + state.vllm_config = vllm_config + state.args = args + resolved_chat_template = load_chat_template(args.chat_template) + + # Merge default_mm_loras into the static lora_modules + default_mm_loras = ( + vllm_config.lora_config.default_mm_loras + if vllm_config.lora_config is not None + else {} + ) + lora_modules = process_lora_modules(args.lora_modules, default_mm_loras) + + state.openai_serving_models = OpenAIServingModels( + engine_client=engine_client, + base_model_paths=base_model_paths, + lora_modules=lora_modules, + ) + await state.openai_serving_models.init_static_loras() + + state.online_renderer = OnlineRenderer( + model_config=engine_client.model_config, + renderer=engine_client.renderer, + request_logger=request_logger, + chat_template=resolved_chat_template, + chat_template_content_format=args.chat_template_content_format, + trust_request_chat_template=args.trust_request_chat_template, + enable_auto_tools=args.enable_auto_tool_choice, + exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none, + tool_parser=args.tool_call_parser, + reasoning_parser=args.structured_outputs_config.reasoning_parser, + default_chat_template_kwargs=args.default_chat_template_kwargs, + log_error_stack=args.log_error_stack, + ) + state.online_renderer.warmup() + + state.online_derenderer = OnlineDerenderer( + model_config=engine_client.model_config, + renderer=engine_client.renderer, + request_logger=request_logger, + chat_template=resolved_chat_template, + chat_template_content_format=args.chat_template_content_format, + trust_request_chat_template=args.trust_request_chat_template, + enable_auto_tools=args.enable_auto_tool_choice, + exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none, + tool_parser=args.tool_call_parser, + reasoning_parser=args.structured_outputs_config.reasoning_parser, + default_chat_template_kwargs=args.default_chat_template_kwargs, + log_error_stack=args.log_error_stack, + ) + + state.serving_tokenization = ServingTokenization( + state.openai_serving_models, + state.online_renderer, + request_logger=request_logger, + chat_template=resolved_chat_template, + chat_template_content_format=args.chat_template_content_format, + default_chat_template_kwargs=args.default_chat_template_kwargs, + trust_request_chat_template=args.trust_request_chat_template, + ) + + if "generate" in supported_tasks: + from vllm.entrypoints.generate.api_router import init_generate_state + + await init_generate_state( + engine_client, state, args, request_logger, supported_tasks + ) + + from vllm.entrypoints.scale_out.factories import init_scale_out_state + + init_scale_out_state(state, args, engine_client, request_logger) + + if "transcription" in supported_tasks or "realtime" in supported_tasks: + from vllm.entrypoints.speech_to_text.factories import init_speech_to_text_state + + init_speech_to_text_state( + engine_client, state, args, request_logger, supported_tasks + ) + + if any(task in POOLING_TASKS for task in supported_tasks): + from vllm.entrypoints.pooling.factories import init_pooling_state + + init_pooling_state(engine_client, state, args, request_logger, supported_tasks) + + await init_endpoint_plugins_state(engine_client, state, args) + + state.enable_server_load_tracking = args.enable_server_load_tracking + state.server_load_metrics = 0 diff --git a/vllm/entrypoints/launchers/api_server/entry.py b/vllm/entrypoints/launchers/api_server/entry.py new file mode 100644 index 000000000000..dadc91ec6e58 --- /dev/null +++ b/vllm/entrypoints/launchers/api_server/entry.py @@ -0,0 +1,229 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import asyncio +import multiprocessing +import multiprocessing.forkserver as forkserver +import os +import signal +import socket +import tempfile +from argparse import Namespace +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from typing import Any + +import vllm.envs as envs +from vllm.engine.arg_utils import AsyncEngineArgs +from vllm.engine.protocol import EngineClient +from vllm.logger import init_logger +from vllm.reasoning import ReasoningParserManager +from vllm.tool_parsers import ToolParserManager +from vllm.usage.usage_lib import UsageContext +from vllm.utils.system_utils import decorate_logs + +from ..app import build_app +from ..launcher import serve_http, setup_server +from ..utils.server_utils import get_uvicorn_log_config +from .app_state import init_app_state + +prometheus_multiproc_dir: tempfile.TemporaryDirectory + +# Cannot use __name__ (https://github.com/vllm-project/vllm/pull/4765) +logger = init_logger("vllm.entrypoints.launchers.api_server.entry") + + +@asynccontextmanager +async def build_async_engine_client( + args: Namespace, + *, + usage_context: UsageContext = UsageContext.OPENAI_API_SERVER, + client_config: dict[str, Any] | None = None, +) -> AsyncIterator[EngineClient]: + if os.getenv("VLLM_WORKER_MULTIPROC_METHOD") == "forkserver": + # The executor is expected to be mp. + # Pre-import heavy modules in the forkserver process + logger.debug("Setup forkserver with pre-imports") + multiprocessing.set_start_method("forkserver") + multiprocessing.set_forkserver_preload(["vllm.v1.engine.async_llm"]) + forkserver.ensure_running() + logger.debug("Forkserver setup complete!") + + # Context manager to handle engine_client lifecycle + # Ensures everything is shutdown and cleaned up on error/exit + engine_args = AsyncEngineArgs.from_cli_args(args) + if client_config: + engine_args._api_process_count = client_config.get("client_count", 1) + engine_args._api_process_rank = client_config.get("client_index", 0) + + async with build_async_engine_client_from_engine_args( + engine_args, + usage_context=usage_context, + client_config=client_config, + ) as engine: + yield engine + + +@asynccontextmanager +async def build_async_engine_client_from_engine_args( + engine_args: AsyncEngineArgs, + *, + usage_context: UsageContext = UsageContext.OPENAI_API_SERVER, + client_config: dict[str, Any] | None = None, +) -> AsyncIterator[EngineClient]: + """ + Create EngineClient, either: + - in-process using the AsyncLLMEngine Directly + - multiprocess using AsyncLLMEngine RPC + + Returns the Client or None if the creation failed. + """ + + # Create the EngineConfig (determines if we can use V1). + vllm_config = engine_args.create_engine_config(usage_context=usage_context) + + from vllm.v1.engine.async_llm import AsyncLLM + + async_llm: AsyncLLM | None = None + + # Don't mutate the input client_config + client_config = dict(client_config) if client_config else {} + client_count = client_config.pop("client_count", 1) + client_index = client_config.pop("client_index", 0) + + try: + async_llm = AsyncLLM.from_vllm_config( + vllm_config=vllm_config, + usage_context=usage_context, + enable_log_requests=engine_args.enable_log_requests, + aggregate_engine_logging=engine_args.aggregate_engine_logging, + disable_log_stats=engine_args.disable_log_stats, + client_addresses=client_config, + client_count=client_count, + client_index=client_index, + ) + + # Don't keep the dummy data in memory + assert async_llm is not None + await async_llm.reset_mm_cache() + + yield async_llm + finally: + if async_llm: + async_llm.shutdown(timeout=vllm_config.shutdown_timeout) + + +async def build_and_serve( + engine_client: EngineClient, + listen_address: str, + sock: socket.socket, + args: Namespace, + **uvicorn_kwargs, +) -> asyncio.Task: + """Build FastAPI app, initialize state, and start serving. + + Returns the shutdown task for the caller to await. + """ + + # Get uvicorn log config (from file or with endpoint filter) + log_config = get_uvicorn_log_config(args) + if log_config is not None: + uvicorn_kwargs["log_config"] = log_config + + supported_tasks = await engine_client.get_supported_tasks() + model_config = engine_client.model_config + + logger.info("Supported tasks: %s", supported_tasks) + app = build_app(args, supported_tasks, model_config) + await init_app_state(engine_client, app.state, args, supported_tasks) + + logger.info("Starting vLLM server on %s", listen_address) + + return await serve_http( + app, + sock=sock, + enable_ssl_refresh=args.enable_ssl_refresh, + host=args.host, + port=args.port, + log_level=args.uvicorn_log_level, + # NOTE: When the 'disable_uvicorn_access_log' value is True, + # no access log will be output. + access_log=not args.disable_uvicorn_access_log, + timeout_keep_alive=envs.VLLM_HTTP_TIMEOUT_KEEP_ALIVE, + ssl_keyfile=args.ssl_keyfile, + ssl_certfile=args.ssl_certfile, + ssl_ca_certs=args.ssl_ca_certs, + ssl_cert_reqs=args.ssl_cert_reqs, + ssl_ciphers=args.ssl_ciphers, + h11_max_incomplete_event_size=args.h11_max_incomplete_event_size, + h11_max_header_count=args.h11_max_header_count, + **uvicorn_kwargs, + ) + + +async def run_server(args, **uvicorn_kwargs) -> None: + """Run a single-worker API server.""" + + decorate_logs("APIServer", skip_if_decorated=True) + + # Interrupt initialization if SIGTERM arrives before uvicorn installs its + # own signal handlers. Once uvicorn is running it replaces this. + def _interrupt_init(*_) -> None: + raise KeyboardInterrupt("terminated") + + signal.signal(signal.SIGTERM, _interrupt_init) + + listen_address, sock = setup_server(args, reuse_port=False) + await run_server_worker(listen_address, sock, args, **uvicorn_kwargs) + + +async def run_server_worker( + listen_address, sock, args, client_config=None, **uvicorn_kwargs +) -> None: + """Run a single API server worker.""" + + if args.tool_parser_plugin and len(args.tool_parser_plugin) > 3: + ToolParserManager.import_tool_parser(args.tool_parser_plugin) + + if args.reasoning_parser_plugin and len(args.reasoning_parser_plugin) > 3: + ReasoningParserManager.import_reasoning_parser(args.reasoning_parser_plugin) + + async with build_async_engine_client( + args, + client_config=client_config, + ) as engine_client: + shutdown_task = await build_and_serve( + engine_client, listen_address, sock, args, **uvicorn_kwargs + ) + # NB: Await server shutdown only after the backend context is exited + try: + await shutdown_task + finally: + sock.close() + + +def main(): + import uvloop + + from vllm.entrypoints.openai.cli_args import ( + make_arg_parser, + validate_parsed_serve_args, + ) + from vllm.entrypoints.serve.utils.api_utils import cli_env_setup + from vllm.utils.argparse_utils import FlexibleArgumentParser + + # NOTE(simon): + # This section should be in sync with vllm/entrypoints/cli/main.py for CLI + # entrypoints. + cli_env_setup() + parser = FlexibleArgumentParser( + description="vLLM OpenAI-Compatible RESTful API server." + ) + parser = make_arg_parser(parser) + args = parser.parse_args() + validate_parsed_serve_args(args) + + uvloop.run(run_server(args)) + + +if __name__ == "__main__": + main() diff --git a/vllm/entrypoints/launchers/app.py b/vllm/entrypoints/launchers/app.py new file mode 100644 index 000000000000..ddea70781051 --- /dev/null +++ b/vllm/entrypoints/launchers/app.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import warnings +from argparse import Namespace + +from fastapi import FastAPI + +from vllm.config import ModelConfig +from vllm.entrypoints.serve.exception_handling.register import init_exception_handler +from vllm.entrypoints.serve.middleware.register import init_entrypoints_middleware +from vllm.entrypoints.serve.sagemaker.api_router import sagemaker_standards_bootstrap +from vllm.plugins.endpoint_plugins.interface import attach_endpoint_plugins +from vllm.tasks import FALLBACK_SUPPORTED_TASKS, SupportedTask + +from .api_server.routers import register_api_routers +from .utils.server_utils import lifespan + + +def build_app( + args: Namespace, + supported_tasks: tuple["SupportedTask", ...] | None = None, + model_config: ModelConfig | None = None, +) -> FastAPI: + if supported_tasks is None: + warnings.warn( + "The 'supported_tasks' parameter was not provided to " + "build_app and will be required in a future version. " + "Defaulting to ('generate',).", + DeprecationWarning, + stacklevel=2, + ) + supported_tasks = FALLBACK_SUPPORTED_TASKS + + if args.disable_fastapi_docs: + app = FastAPI( + openapi_url=None, docs_url=None, redoc_url=None, lifespan=lifespan + ) + elif args.enable_offline_docs: + app = FastAPI(docs_url=None, redoc_url=None, lifespan=lifespan) + else: + app = FastAPI(lifespan=lifespan) + app.state.args = args + app.root_path = args.root_path + + register_api_routers(args, app, supported_tasks, model_config) + + # Endpoint plugins are attached last so their routes are registered after all core + # routers. This runs even for the CPU only render server. A plugin eligible for + # the `render` task still gets its routes registered. It receives + # `engine_client=None` at Phase B (see `_init_endpoint_plugins_state`). + attach_endpoint_plugins(app, supported_tasks) + + init_exception_handler(app) + init_entrypoints_middleware(args, app, supported_tasks) + app = sagemaker_standards_bootstrap(app) + return app diff --git a/vllm/entrypoints/launcher.py b/vllm/entrypoints/launchers/launcher.py similarity index 64% rename from vllm/entrypoints/launcher.py rename to vllm/entrypoints/launchers/launcher.py index 80b210516333..aeacb9f8c22c 100644 --- a/vllm/entrypoints/launcher.py +++ b/vllm/entrypoints/launchers/launcher.py @@ -14,13 +14,23 @@ from vllm import envs from vllm.engine.protocol import EngineClient -from vllm.entrypoints.serve.utils.constants import ( +from vllm.entrypoints.launchers.utils.ssl import SSLCertRefresher +from vllm.entrypoints.serve.utils.api_utils import ( + log_non_default_args, + log_version_and_model, +) +from vllm.logger import init_logger +from vllm.reasoning import ReasoningParserManager +from vllm.tool_parsers import ToolParserManager +from vllm.tracing import instrument +from vllm.utils.network_utils import find_process_using_port, is_valid_ipv6_address +from vllm.utils.system_utils import set_ulimit +from vllm.version import __version__ as VLLM_VERSION + +from .utils.constants import ( H11_MAX_HEADER_COUNT_DEFAULT, H11_MAX_INCOMPLETE_EVENT_SIZE_DEFAULT, ) -from vllm.entrypoints.serve.utils.ssl import SSLCertRefresher -from vllm.logger import init_logger -from vllm.utils.network_utils import find_process_using_port logger = init_logger(__name__) @@ -190,3 +200,83 @@ def terminate_if_errored(server: uvicorn.Server, engine: EngineClient): engine_errored = engine.errored and not engine.is_running if not envs.VLLM_KEEP_ALIVE_ON_ENGINE_DEATH and engine_errored: server.should_exit = True + + +def create_server_socket( + addr: tuple[str, int], + *, + reuse_port: bool, +) -> socket.socket: + family = socket.AF_INET + if is_valid_ipv6_address(addr[0]): + family = socket.AF_INET6 + + sock = socket.socket(family=family, type=socket.SOCK_STREAM) + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + if reuse_port: + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1) + sock.bind(addr) + + return sock + + +def create_server_unix_socket(path: str) -> socket.socket: + sock = socket.socket(family=socket.AF_UNIX, type=socket.SOCK_STREAM) + sock.bind(path) + return sock + + +def validate_api_server_args(args): + valid_tool_parses = ToolParserManager.list_registered() + if args.enable_auto_tool_choice and args.tool_call_parser not in valid_tool_parses: + raise KeyError( + f"invalid tool call parser: {args.tool_call_parser} " + f"(chose from {{ {','.join(valid_tool_parses)} }})" + ) + + valid_reasoning_parsers = ReasoningParserManager.list_registered() + if ( + reasoning_parser := args.structured_outputs_config.reasoning_parser + ) and reasoning_parser not in valid_reasoning_parsers: + raise KeyError( + f"invalid reasoning parser: {reasoning_parser} " + f"(chose from {{ {','.join(valid_reasoning_parsers)} }})" + ) + + +@instrument(span_name="API server setup") +def setup_server(args, *, reuse_port: bool): + """Validate API server args and create the server socket.""" + + log_version_and_model(logger, VLLM_VERSION, args.model) + log_non_default_args(args) + + if args.tool_parser_plugin and len(args.tool_parser_plugin) > 3: + ToolParserManager.import_tool_parser(args.tool_parser_plugin) + + if args.reasoning_parser_plugin and len(args.reasoning_parser_plugin) > 3: + ReasoningParserManager.import_reasoning_parser(args.reasoning_parser_plugin) + + validate_api_server_args(args) + + # workaround to make sure that we bind the port before the engine is set up. + # This avoids race conditions with ray. + # see https://github.com/vllm-project/vllm/issues/8204 + if args.uds: + sock = create_server_unix_socket(args.uds) + else: + sock_addr = (args.host or "", args.port) + sock = create_server_socket(sock_addr, reuse_port=reuse_port) + + # workaround to avoid footguns where uvicorn drops requests with too + # many concurrent requests active + set_ulimit() + + if args.uds: + listen_address = f"unix:{args.uds}" + else: + addr, port = sock_addr + is_ssl = args.ssl_keyfile and args.ssl_certfile + host_part = f"[{addr}]" if is_valid_ipv6_address(addr) else addr or "0.0.0.0" + listen_address = f"http{'s' if is_ssl else ''}://{host_part}:{port}" + return listen_address, sock diff --git a/vllm/entrypoints/launchers/render/__init__.py b/vllm/entrypoints/launchers/render/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/vllm/entrypoints/launchers/render/app_state.py b/vllm/entrypoints/launchers/render/app_state.py new file mode 100644 index 000000000000..2b0b030b4eda --- /dev/null +++ b/vllm/entrypoints/launchers/render/app_state.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from argparse import Namespace + +from starlette.datastructures import State + +from vllm.config import VllmConfig +from vllm.entrypoints.chat_utils import load_chat_template +from vllm.entrypoints.openai.models.protocol import BaseModelPath +from vllm.entrypoints.openai.models.serving import OpenAIModelRegistry +from vllm.entrypoints.scale_out.factories import init_render_state +from vllm.entrypoints.serve.tokenize.serving import ServingTokenization +from vllm.entrypoints.serve.utils.request_logger import RequestLogger +from vllm.plugins.endpoint_plugins.interface import init_endpoint_plugins_state +from vllm.renderers import renderer_from_config +from vllm.renderers.online_derenderer import OnlineDerenderer +from vllm.renderers.online_renderer import OnlineRenderer + + +async def init_render_app_state( + vllm_config: VllmConfig, + state: State, + args: Namespace, +) -> None: + """Initialise FastAPI app state for a CPU-only render server. + + Unlike :func:`init_app_state` this function does not require an + :class:`~vllm.engine.protocol.EngineClient`; it bootstraps the + preprocessing pipeline (renderer, input_processor) + directly from the :class:`~vllm.config.VllmConfig`. + """ + + served_model_names = args.served_model_name or [args.model] + model_registry = OpenAIModelRegistry( + model_config=vllm_config.model_config, + base_model_paths=[ + BaseModelPath(name=name, model_path=args.model) + for name in served_model_names + ], + ) + + if args.enable_log_requests: + request_logger = RequestLogger(max_log_len=args.max_log_len) + else: + request_logger = None + + renderer = renderer_from_config(vllm_config) + resolved_chat_template = load_chat_template(args.chat_template) + + state.online_renderer = OnlineRenderer( + model_config=vllm_config.model_config, + renderer=renderer, + request_logger=request_logger, + chat_template=resolved_chat_template, + chat_template_content_format=args.chat_template_content_format, + trust_request_chat_template=args.trust_request_chat_template, + enable_auto_tools=args.enable_auto_tool_choice, + exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none, + tool_parser=args.tool_call_parser, + reasoning_parser=args.reasoning_parser, + default_chat_template_kwargs=args.default_chat_template_kwargs, + log_error_stack=args.log_error_stack, + ) + state.online_renderer.warmup() + + state.online_derenderer = OnlineDerenderer( + model_config=vllm_config.model_config, + renderer=renderer, + request_logger=request_logger, + chat_template=resolved_chat_template, + chat_template_content_format=args.chat_template_content_format, + trust_request_chat_template=args.trust_request_chat_template, + enable_auto_tools=args.enable_auto_tool_choice, + exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none, + tool_parser=args.tool_call_parser, + reasoning_parser=args.reasoning_parser, + default_chat_template_kwargs=args.default_chat_template_kwargs, + log_error_stack=args.log_error_stack, + ) + + state.openai_serving_models = model_registry + state.serving_tokenization = ServingTokenization( + model_registry, + state.online_renderer, + request_logger=request_logger, + chat_template=resolved_chat_template, + chat_template_content_format=args.chat_template_content_format, + default_chat_template_kwargs=args.default_chat_template_kwargs, + trust_request_chat_template=args.trust_request_chat_template, + ) + + init_render_state(state, request_logger) + + state.vllm_config = vllm_config + # Disable stats logging — there is no engine to poll. + state.log_stats = False + state.engine_client = None + state.args = args + state.enable_server_load_tracking = False + state.server_load_metrics = 0 + + # No `EngineClient` exists for the render server, so plugins get `None` and + # must handle it themselves (see `EndpointPlugin.init_state`). + await init_endpoint_plugins_state(None, state, args) diff --git a/vllm/entrypoints/launchers/render/entry.py b/vllm/entrypoints/launchers/render/entry.py new file mode 100644 index 000000000000..2c7c4b54259a --- /dev/null +++ b/vllm/entrypoints/launchers/render/entry.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import argparse +import asyncio +import signal +import socket +from argparse import Namespace + +from vllm import AsyncEngineArgs, envs +from vllm.config import VllmConfig +from vllm.logger import init_logger + +from ..app import build_app +from ..launcher import serve_http, setup_server +from ..utils.server_utils import get_uvicorn_log_config +from .app_state import init_render_app_state + +logger = init_logger("vllm.entrypoints.launchers.render.entry") + + +async def build_and_serve_renderer( + vllm_config: VllmConfig, + listen_address: str, + sock: socket.socket, + args: Namespace, + **uvicorn_kwargs, +) -> asyncio.Task: + """Build FastAPI app for a CPU-only render server, initialize state, and + start serving. + + Returns the shutdown task for the caller to await. + """ + + # Get uvicorn log config (from file or with endpoint filter) + log_config = get_uvicorn_log_config(args) + if log_config is not None: + uvicorn_kwargs["log_config"] = log_config + + app = build_app(args, ("render",)) + await init_render_app_state(vllm_config, app.state, args) + + logger.info("Starting vLLM server on %s", listen_address) + + return await serve_http( + app, + sock=sock, + enable_ssl_refresh=args.enable_ssl_refresh, + host=args.host, + port=args.port, + log_level=args.uvicorn_log_level, + # NOTE: When the 'disable_uvicorn_access_log' value is True, + # no access log will be output. + access_log=not args.disable_uvicorn_access_log, + timeout_keep_alive=envs.VLLM_HTTP_TIMEOUT_KEEP_ALIVE, + ssl_keyfile=args.ssl_keyfile, + ssl_certfile=args.ssl_certfile, + ssl_ca_certs=args.ssl_ca_certs, + ssl_cert_reqs=args.ssl_cert_reqs, + ssl_ciphers=args.ssl_ciphers, + h11_max_incomplete_event_size=args.h11_max_incomplete_event_size, + h11_max_header_count=args.h11_max_header_count, + **uvicorn_kwargs, + ) + + +async def run_launch_fastapi(args: argparse.Namespace) -> None: + """Run the online serving layer with FastAPI (no GPU inference).""" + + # Interrupt initialization if SIGTERM arrives before uvicorn installs + # its own signal handlers. Once uvicorn is running it replaces this. + def _interrupt_init(*_) -> None: + raise KeyboardInterrupt("terminated") + + signal.signal(signal.SIGTERM, _interrupt_init) + + # 1. Socket binding + listen_address, sock = setup_server(args, reuse_port=False) + + # 2. Build and serve the API server + engine_args = AsyncEngineArgs.from_cli_args(args) + model_config = engine_args.create_model_config() + + # Render servers preprocess data only — no inference, no quantized kernels. + # Clear quantization so VllmConfig skips quant dtype/capability validation. + model_config.quantization = None + + # Render servers never allocate KV cache; suppress the spurious CPU KV + # cache space warning from CpuPlatform.check_and_update_config. + envs.VLLM_CPU_KVCACHE_SPACE = 0 + + vllm_config = VllmConfig(model_config=model_config) + shutdown_task = await build_and_serve_renderer( + vllm_config, listen_address, sock, args + ) + try: + await shutdown_task + finally: + sock.close() + + +if __name__ == "__main__": + import uvloop + + from vllm.entrypoints.openai.cli_args import ( + make_arg_parser, + validate_parsed_serve_args, + ) + from vllm.entrypoints.serve.utils.api_utils import cli_env_setup + from vllm.utils.argparse_utils import FlexibleArgumentParser + + cli_env_setup() + parser = FlexibleArgumentParser( + description="Starts a GPU-less rendering server " + "for preprocessing and postprocessing only" + ) + parser = make_arg_parser(parser) + args = parser.parse_args() + validate_parsed_serve_args(args) + + uvloop.run(run_launch_fastapi(args)) diff --git a/vllm/entrypoints/launchers/utils/__init__.py b/vllm/entrypoints/launchers/utils/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/vllm/entrypoints/serve/utils/constants.py b/vllm/entrypoints/launchers/utils/constants.py similarity index 80% rename from vllm/entrypoints/serve/utils/constants.py rename to vllm/entrypoints/launchers/utils/constants.py index 5726ee0735d4..99750c31749c 100644 --- a/vllm/entrypoints/serve/utils/constants.py +++ b/vllm/entrypoints/launchers/utils/constants.py @@ -1,12 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Shared constants for vLLM entrypoints. -""" # HTTP header limits for h11 parser # These constants help mitigate header abuse attacks H11_MAX_INCOMPLETE_EVENT_SIZE_DEFAULT = 4194304 # 4 MB H11_MAX_HEADER_COUNT_DEFAULT = 256 - -MCP_PREFIX = "mcp_" diff --git a/vllm/entrypoints/serve/utils/server_utils.py b/vllm/entrypoints/launchers/utils/server_utils.py similarity index 98% rename from vllm/entrypoints/serve/utils/server_utils.py rename to vllm/entrypoints/launchers/utils/server_utils.py index 97910d019ec7..082582672381 100644 --- a/vllm/entrypoints/serve/utils/server_utils.py +++ b/vllm/entrypoints/launchers/utils/server_utils.py @@ -12,7 +12,7 @@ from vllm.logger import init_logger from vllm.utils.gc_utils import freeze_gc_heap -logger = init_logger("vllm.entrypoints.openai.server_utils") +logger = init_logger(__name__) def load_log_config(log_config_file: str | None) -> dict | None: diff --git a/vllm/entrypoints/serve/utils/ssl.py b/vllm/entrypoints/launchers/utils/ssl.py similarity index 100% rename from vllm/entrypoints/serve/utils/ssl.py rename to vllm/entrypoints/launchers/utils/ssl.py diff --git a/vllm/entrypoints/openai/api_server.py b/vllm/entrypoints/openai/api_server.py index c541f03f445e..a8e8b58c1877 100644 --- a/vllm/entrypoints/openai/api_server.py +++ b/vllm/entrypoints/openai/api_server.py @@ -1,670 +1,59 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import asyncio -import multiprocessing -import multiprocessing.forkserver as forkserver -import os -import signal -import socket -import tempfile import warnings -from argparse import Namespace -from collections.abc import AsyncIterator -from contextlib import asynccontextmanager -from typing import Any, cast -import uvloop -from fastapi import FastAPI -from starlette.datastructures import State - -import vllm.envs as envs -from vllm.config import ModelConfig, VllmConfig -from vllm.engine.arg_utils import AsyncEngineArgs -from vllm.engine.protocol import EngineClient -from vllm.entrypoints.chat_utils import load_chat_template -from vllm.entrypoints.launcher import serve_http +from vllm.entrypoints.launchers.api_server.app_state import init_app_state +from vllm.entrypoints.launchers.api_server.entry import ( + build_and_serve, + build_async_engine_client, + build_async_engine_client_from_engine_args, + run_server, + run_server_worker, +) from vllm.entrypoints.launchers.api_server.routers import register_api_routers -from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args -from vllm.entrypoints.openai.models.protocol import BaseModelPath -from vllm.entrypoints.openai.models.serving import OpenAIServingModels -from vllm.entrypoints.serve.exception_handling.register import init_exception_handler -from vllm.entrypoints.serve.middleware.register import init_entrypoints_middleware -from vllm.entrypoints.serve.sagemaker.api_router import sagemaker_standards_bootstrap -from vllm.entrypoints.serve.tokenize.serving import ServingTokenization -from vllm.entrypoints.serve.utils.api_utils import ( - cli_env_setup, - log_non_default_args, - log_version_and_model, - process_lora_modules, +from vllm.entrypoints.launchers.app import build_app +from vllm.entrypoints.launchers.launcher import ( + create_server_socket, + create_server_unix_socket, + setup_server, + validate_api_server_args, ) -from vllm.entrypoints.serve.utils.request_logger import RequestLogger -from vllm.entrypoints.serve.utils.server_utils import ( - get_uvicorn_log_config, - lifespan, +from vllm.entrypoints.launchers.render.app_state import init_render_app_state +from vllm.entrypoints.launchers.render.entry import build_and_serve_renderer + +warnings.warn( + "`vllm.entrypoints.openai.api_server` is deprecated and will likely be" + "unsupported in a future version. Use the corresponding function from " + "`vllm.entrypoints.launchers` instead.", + DeprecationWarning, + stacklevel=1, ) -from vllm.logger import init_logger -from vllm.reasoning import ReasoningParserManager -from vllm.renderers.online_derenderer import OnlineDerenderer -from vllm.renderers.online_renderer import OnlineRenderer -from vllm.tasks import POOLING_TASKS, SupportedTask -from vllm.tool_parsers import ToolParserManager -from vllm.tracing import instrument -from vllm.usage.usage_lib import UsageContext -from vllm.utils.argparse_utils import FlexibleArgumentParser -from vllm.utils.network_utils import is_valid_ipv6_address -from vllm.utils.system_utils import decorate_logs, set_ulimit -from vllm.version import __version__ as VLLM_VERSION - -prometheus_multiproc_dir: tempfile.TemporaryDirectory - -# Cannot use __name__ (https://github.com/vllm-project/vllm/pull/4765) -logger = init_logger("vllm.entrypoints.openai.api_server") - -_FALLBACK_SUPPORTED_TASKS: tuple[SupportedTask, ...] = ("generate",) - - -def _attach_endpoint_plugins( - app: FastAPI, supported_tasks: tuple["SupportedTask", ...] -) -> None: - """Phase A of endpoint plugin wiring: discover, gate and attach routes. - - Attached last after all core routers. This is so endpoint plugin routes can - shadow core routes with the same path (see `EndpointPlugin.attach_router` - docstring). No-ops when no plugins are discovered/allowlisted. - """ - from vllm.plugins import load_endpoint_plugins - - endpoint_plugins = load_endpoint_plugins(supported_tasks) - for plugin in endpoint_plugins: - plugin.attach_router(app) - app.state.endpoint_plugins = endpoint_plugins - - -async def _init_endpoint_plugins_state( - engine_client: EngineClient | None, state: State, args: Namespace -) -> None: - """Phase B of endpoint plugin wiring: initialize per app plugin state. - - `state.endpoint_plugins` is set by `_attach_endpoint_plugins` (Phase A) - in `build_app`. Some `init_app_state` callers (e.g. `run_batch.py`) - build their own bare `State` without going through `build_app`. As a result - `endpoint_plugins` may be absent and are treated that the same as "none attached". - - `engine_client` is `None` for the CPU only render server which has no - engine (see `init_render_app_state`). Plugins must handle a `None` - `engine_client` themselves (see `EndpointPlugin.init_state`). - """ - for plugin in getattr(state, "endpoint_plugins", []): - await plugin.init_state(engine_client, state, args) - - -@asynccontextmanager -async def build_async_engine_client( - args: Namespace, - *, - usage_context: UsageContext = UsageContext.OPENAI_API_SERVER, - client_config: dict[str, Any] | None = None, -) -> AsyncIterator[EngineClient]: - if os.getenv("VLLM_WORKER_MULTIPROC_METHOD") == "forkserver": - # The executor is expected to be mp. - # Pre-import heavy modules in the forkserver process - logger.debug("Setup forkserver with pre-imports") - multiprocessing.set_start_method("forkserver") - multiprocessing.set_forkserver_preload(["vllm.v1.engine.async_llm"]) - forkserver.ensure_running() - logger.debug("Forkserver setup complete!") - - # Context manager to handle engine_client lifecycle - # Ensures everything is shutdown and cleaned up on error/exit - engine_args = AsyncEngineArgs.from_cli_args(args) - if client_config: - engine_args._api_process_count = client_config.get("client_count", 1) - engine_args._api_process_rank = client_config.get("client_index", 0) - - async with build_async_engine_client_from_engine_args( - engine_args, - usage_context=usage_context, - client_config=client_config, - ) as engine: - yield engine - - -@asynccontextmanager -async def build_async_engine_client_from_engine_args( - engine_args: AsyncEngineArgs, - *, - usage_context: UsageContext = UsageContext.OPENAI_API_SERVER, - client_config: dict[str, Any] | None = None, -) -> AsyncIterator[EngineClient]: - """ - Create EngineClient, either: - - in-process using the AsyncLLMEngine Directly - - multiprocess using AsyncLLMEngine RPC - - Returns the Client or None if the creation failed. - """ - - # Create the EngineConfig (determines if we can use V1). - vllm_config = engine_args.create_engine_config(usage_context=usage_context) - - from vllm.v1.engine.async_llm import AsyncLLM - - async_llm: AsyncLLM | None = None - - # Don't mutate the input client_config - client_config = dict(client_config) if client_config else {} - client_count = client_config.pop("client_count", 1) - client_index = client_config.pop("client_index", 0) - - try: - async_llm = AsyncLLM.from_vllm_config( - vllm_config=vllm_config, - usage_context=usage_context, - enable_log_requests=engine_args.enable_log_requests, - aggregate_engine_logging=engine_args.aggregate_engine_logging, - disable_log_stats=engine_args.disable_log_stats, - client_addresses=client_config, - client_count=client_count, - client_index=client_index, - ) - - # Don't keep the dummy data in memory - assert async_llm is not None - await async_llm.reset_mm_cache() - - yield async_llm - finally: - if async_llm: - async_llm.shutdown(timeout=vllm_config.shutdown_timeout) - - -def build_app( - args: Namespace, - supported_tasks: tuple["SupportedTask", ...] | None = None, - model_config: ModelConfig | None = None, -) -> FastAPI: - if supported_tasks is None: - warnings.warn( - "The 'supported_tasks' parameter was not provided to " - "build_app and will be required in a future version. " - "Defaulting to ('generate',).", - DeprecationWarning, - stacklevel=2, - ) - supported_tasks = _FALLBACK_SUPPORTED_TASKS - - if args.disable_fastapi_docs: - app = FastAPI( - openapi_url=None, docs_url=None, redoc_url=None, lifespan=lifespan - ) - elif args.enable_offline_docs: - app = FastAPI(docs_url=None, redoc_url=None, lifespan=lifespan) - else: - app = FastAPI(lifespan=lifespan) - app.state.args = args - app.root_path = args.root_path - - register_api_routers(args, app, supported_tasks, model_config) - - # Endpoint plugins are attached last so their routes are registered after all core - # routers. This runs even for the CPU only render server. A plugin eligible for - # the `render` task still gets its routes registered. It receives - # `engine_client=None` at Phase B (see `_init_endpoint_plugins_state`). - _attach_endpoint_plugins(app, supported_tasks) - - init_exception_handler(app) - init_entrypoints_middleware(args, app, supported_tasks) - app = sagemaker_standards_bootstrap(app) - return app - - -async def init_app_state( - engine_client: EngineClient, - state: State, - args: Namespace, - supported_tasks: tuple["SupportedTask", ...] | None = None, -) -> None: - vllm_config = engine_client.vllm_config - - if args.tool_call_parser is not None: - from vllm.parser.metrics import init_parser_metrics - - init_parser_metrics( - model_name=cast(str, vllm_config.model_config.served_model_name) - ) - - if supported_tasks is None: - warnings.warn( - "The 'supported_tasks' parameter was not provided to " - "init_app_state and will be required in a future version. " - "Please pass 'supported_tasks' explicitly.", - DeprecationWarning, - stacklevel=2, - ) - supported_tasks = _FALLBACK_SUPPORTED_TASKS - - if args.served_model_name is not None: - served_model_names = args.served_model_name - else: - served_model_names = [args.model] - - if args.enable_log_requests: - request_logger = RequestLogger(max_log_len=args.max_log_len) - else: - request_logger = None - - base_model_paths = [ - BaseModelPath(name=name, model_path=args.model) for name in served_model_names - ] - - state.engine_client = engine_client - state.log_stats = not args.disable_log_stats - state.vllm_config = vllm_config - state.args = args - resolved_chat_template = load_chat_template(args.chat_template) - - # Merge default_mm_loras into the static lora_modules - default_mm_loras = ( - vllm_config.lora_config.default_mm_loras - if vllm_config.lora_config is not None - else {} - ) - lora_modules = process_lora_modules(args.lora_modules, default_mm_loras) - - state.openai_serving_models = OpenAIServingModels( - engine_client=engine_client, - base_model_paths=base_model_paths, - lora_modules=lora_modules, - ) - await state.openai_serving_models.init_static_loras() - - state.online_renderer = OnlineRenderer( - model_config=engine_client.model_config, - renderer=engine_client.renderer, - request_logger=request_logger, - chat_template=resolved_chat_template, - chat_template_content_format=args.chat_template_content_format, - trust_request_chat_template=args.trust_request_chat_template, - enable_auto_tools=args.enable_auto_tool_choice, - exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none, - tool_parser=args.tool_call_parser, - reasoning_parser=args.structured_outputs_config.reasoning_parser, - default_chat_template_kwargs=args.default_chat_template_kwargs, - log_error_stack=args.log_error_stack, - ) - state.online_renderer.warmup() - - state.online_derenderer = OnlineDerenderer( - model_config=engine_client.model_config, - renderer=engine_client.renderer, - request_logger=request_logger, - chat_template=resolved_chat_template, - chat_template_content_format=args.chat_template_content_format, - trust_request_chat_template=args.trust_request_chat_template, - enable_auto_tools=args.enable_auto_tool_choice, - exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none, - tool_parser=args.tool_call_parser, - reasoning_parser=args.structured_outputs_config.reasoning_parser, - default_chat_template_kwargs=args.default_chat_template_kwargs, - log_error_stack=args.log_error_stack, - ) - - state.serving_tokenization = ServingTokenization( - state.openai_serving_models, - state.online_renderer, - request_logger=request_logger, - chat_template=resolved_chat_template, - chat_template_content_format=args.chat_template_content_format, - default_chat_template_kwargs=args.default_chat_template_kwargs, - trust_request_chat_template=args.trust_request_chat_template, - ) - - if "generate" in supported_tasks: - from vllm.entrypoints.generate.api_router import init_generate_state - - await init_generate_state( - engine_client, state, args, request_logger, supported_tasks - ) - - from vllm.entrypoints.scale_out.factories import init_scale_out_state - - init_scale_out_state(state, args, engine_client, request_logger) - - if "transcription" in supported_tasks or "realtime" in supported_tasks: - from vllm.entrypoints.speech_to_text.factories import init_speech_to_text_state - - init_speech_to_text_state( - engine_client, state, args, request_logger, supported_tasks - ) - - if any(task in POOLING_TASKS for task in supported_tasks): - from vllm.entrypoints.pooling.factories import init_pooling_state - - init_pooling_state(engine_client, state, args, request_logger, supported_tasks) - - await _init_endpoint_plugins_state(engine_client, state, args) - - state.enable_server_load_tracking = args.enable_server_load_tracking - state.server_load_metrics = 0 - - -async def init_render_app_state( - vllm_config: VllmConfig, - state: State, - args: Namespace, -) -> None: - """Initialise FastAPI app state for a CPU-only render server. - - Unlike :func:`init_app_state` this function does not require an - :class:`~vllm.engine.protocol.EngineClient`; it bootstraps the - preprocessing pipeline (renderer, input_processor) - directly from the :class:`~vllm.config.VllmConfig`. - """ - from vllm.entrypoints.chat_utils import load_chat_template - from vllm.entrypoints.openai.models.serving import OpenAIModelRegistry - from vllm.renderers import renderer_from_config - from vllm.renderers.online_renderer import OnlineRenderer - - served_model_names = args.served_model_name or [args.model] - model_registry = OpenAIModelRegistry( - model_config=vllm_config.model_config, - base_model_paths=[ - BaseModelPath(name=name, model_path=args.model) - for name in served_model_names - ], - ) - - if args.enable_log_requests: - request_logger = RequestLogger(max_log_len=args.max_log_len) - else: - request_logger = None - - renderer = renderer_from_config(vllm_config) - resolved_chat_template = load_chat_template(args.chat_template) - - state.online_renderer = OnlineRenderer( - model_config=vllm_config.model_config, - renderer=renderer, - request_logger=request_logger, - chat_template=resolved_chat_template, - chat_template_content_format=args.chat_template_content_format, - trust_request_chat_template=args.trust_request_chat_template, - enable_auto_tools=args.enable_auto_tool_choice, - exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none, - tool_parser=args.tool_call_parser, - reasoning_parser=args.reasoning_parser, - default_chat_template_kwargs=args.default_chat_template_kwargs, - log_error_stack=args.log_error_stack, - ) - state.online_renderer.warmup() - - state.online_derenderer = OnlineDerenderer( - model_config=vllm_config.model_config, - renderer=renderer, - request_logger=request_logger, - chat_template=resolved_chat_template, - chat_template_content_format=args.chat_template_content_format, - trust_request_chat_template=args.trust_request_chat_template, - enable_auto_tools=args.enable_auto_tool_choice, - exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none, - tool_parser=args.tool_call_parser, - reasoning_parser=args.reasoning_parser, - default_chat_template_kwargs=args.default_chat_template_kwargs, - log_error_stack=args.log_error_stack, - ) - - state.openai_serving_models = model_registry - state.serving_tokenization = ServingTokenization( - model_registry, - state.online_renderer, - request_logger=request_logger, - chat_template=resolved_chat_template, - chat_template_content_format=args.chat_template_content_format, - default_chat_template_kwargs=args.default_chat_template_kwargs, - trust_request_chat_template=args.trust_request_chat_template, - ) - - from vllm.entrypoints.scale_out.factories import init_render_state - - init_render_state(state, request_logger) - - state.vllm_config = vllm_config - # Disable stats logging — there is no engine to poll. - state.log_stats = False - state.engine_client = None - state.args = args - state.enable_server_load_tracking = False - state.server_load_metrics = 0 - - # No `EngineClient` exists for the render server, so plugins get `None` and - # must handle it themselves (see `EndpointPlugin.init_state`). - await _init_endpoint_plugins_state(None, state, args) - - -def create_server_socket( - addr: tuple[str, int], - *, - reuse_port: bool, -) -> socket.socket: - family = socket.AF_INET - if is_valid_ipv6_address(addr[0]): - family = socket.AF_INET6 - - sock = socket.socket(family=family, type=socket.SOCK_STREAM) - sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - if reuse_port: - sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1) - sock.bind(addr) - - return sock - - -def create_server_unix_socket(path: str) -> socket.socket: - sock = socket.socket(family=socket.AF_UNIX, type=socket.SOCK_STREAM) - sock.bind(path) - return sock - - -def validate_api_server_args(args): - valid_tool_parses = ToolParserManager.list_registered() - if args.enable_auto_tool_choice and args.tool_call_parser not in valid_tool_parses: - raise KeyError( - f"invalid tool call parser: {args.tool_call_parser} " - f"(chose from {{ {','.join(valid_tool_parses)} }})" - ) - - valid_reasoning_parsers = ReasoningParserManager.list_registered() - if ( - reasoning_parser := args.structured_outputs_config.reasoning_parser - ) and reasoning_parser not in valid_reasoning_parsers: - raise KeyError( - f"invalid reasoning parser: {reasoning_parser} " - f"(chose from {{ {','.join(valid_reasoning_parsers)} }})" - ) - - -@instrument(span_name="API server setup") -def setup_server(args, *, reuse_port: bool): - """Validate API server args and create the server socket.""" - - log_version_and_model(logger, VLLM_VERSION, args.model) - log_non_default_args(args) - - if args.tool_parser_plugin and len(args.tool_parser_plugin) > 3: - ToolParserManager.import_tool_parser(args.tool_parser_plugin) - - if args.reasoning_parser_plugin and len(args.reasoning_parser_plugin) > 3: - ReasoningParserManager.import_reasoning_parser(args.reasoning_parser_plugin) - - validate_api_server_args(args) - - # workaround to make sure that we bind the port before the engine is set up. - # This avoids race conditions with ray. - # see https://github.com/vllm-project/vllm/issues/8204 - if args.uds: - sock = create_server_unix_socket(args.uds) - else: - sock_addr = (args.host or "", args.port) - sock = create_server_socket(sock_addr, reuse_port=reuse_port) - - # workaround to avoid footguns where uvicorn drops requests with too - # many concurrent requests active - set_ulimit() - - if args.uds: - listen_address = f"unix:{args.uds}" - else: - addr, port = sock_addr - is_ssl = args.ssl_keyfile and args.ssl_certfile - host_part = f"[{addr}]" if is_valid_ipv6_address(addr) else addr or "0.0.0.0" - listen_address = f"http{'s' if is_ssl else ''}://{host_part}:{port}" - return listen_address, sock - - -async def build_and_serve( - engine_client: EngineClient, - listen_address: str, - sock: socket.socket, - args: Namespace, - **uvicorn_kwargs, -) -> asyncio.Task: - """Build FastAPI app, initialize state, and start serving. - - Returns the shutdown task for the caller to await. - """ - - # Get uvicorn log config (from file or with endpoint filter) - log_config = get_uvicorn_log_config(args) - if log_config is not None: - uvicorn_kwargs["log_config"] = log_config - - supported_tasks = await engine_client.get_supported_tasks() - model_config = engine_client.model_config - - logger.info("Supported tasks: %s", supported_tasks) - app = build_app(args, supported_tasks, model_config) - await init_app_state(engine_client, app.state, args, supported_tasks) - - logger.info("Starting vLLM server on %s", listen_address) - - return await serve_http( - app, - sock=sock, - enable_ssl_refresh=args.enable_ssl_refresh, - host=args.host, - port=args.port, - log_level=args.uvicorn_log_level, - # NOTE: When the 'disable_uvicorn_access_log' value is True, - # no access log will be output. - access_log=not args.disable_uvicorn_access_log, - timeout_keep_alive=envs.VLLM_HTTP_TIMEOUT_KEEP_ALIVE, - ssl_keyfile=args.ssl_keyfile, - ssl_certfile=args.ssl_certfile, - ssl_ca_certs=args.ssl_ca_certs, - ssl_cert_reqs=args.ssl_cert_reqs, - ssl_ciphers=args.ssl_ciphers, - h11_max_incomplete_event_size=args.h11_max_incomplete_event_size, - h11_max_header_count=args.h11_max_header_count, - **uvicorn_kwargs, - ) - - -async def build_and_serve_renderer( - vllm_config: VllmConfig, - listen_address: str, - sock: socket.socket, - args: Namespace, - **uvicorn_kwargs, -) -> asyncio.Task: - """Build FastAPI app for a CPU-only render server, initialize state, and - start serving. - - Returns the shutdown task for the caller to await. - """ - - # Get uvicorn log config (from file or with endpoint filter) - log_config = get_uvicorn_log_config(args) - if log_config is not None: - uvicorn_kwargs["log_config"] = log_config - - app = build_app(args, ("render",)) - await init_render_app_state(vllm_config, app.state, args) - - logger.info("Starting vLLM server on %s", listen_address) - - return await serve_http( - app, - sock=sock, - enable_ssl_refresh=args.enable_ssl_refresh, - host=args.host, - port=args.port, - log_level=args.uvicorn_log_level, - # NOTE: When the 'disable_uvicorn_access_log' value is True, - # no access log will be output. - access_log=not args.disable_uvicorn_access_log, - timeout_keep_alive=envs.VLLM_HTTP_TIMEOUT_KEEP_ALIVE, - ssl_keyfile=args.ssl_keyfile, - ssl_certfile=args.ssl_certfile, - ssl_ca_certs=args.ssl_ca_certs, - ssl_cert_reqs=args.ssl_cert_reqs, - ssl_ciphers=args.ssl_ciphers, - h11_max_incomplete_event_size=args.h11_max_incomplete_event_size, - h11_max_header_count=args.h11_max_header_count, - **uvicorn_kwargs, - ) - - -async def run_server(args, **uvicorn_kwargs) -> None: - """Run a single-worker API server.""" - - decorate_logs("APIServer", skip_if_decorated=True) - - # Interrupt initialization if SIGTERM arrives before uvicorn installs its - # own signal handlers. Once uvicorn is running it replaces this. - def _interrupt_init(*_) -> None: - raise KeyboardInterrupt("terminated") - - signal.signal(signal.SIGTERM, _interrupt_init) - - listen_address, sock = setup_server(args, reuse_port=False) - await run_server_worker(listen_address, sock, args, **uvicorn_kwargs) - - -async def run_server_worker( - listen_address, sock, args, client_config=None, **uvicorn_kwargs -) -> None: - """Run a single API server worker.""" - - if args.tool_parser_plugin and len(args.tool_parser_plugin) > 3: - ToolParserManager.import_tool_parser(args.tool_parser_plugin) - - if args.reasoning_parser_plugin and len(args.reasoning_parser_plugin) > 3: - ReasoningParserManager.import_reasoning_parser(args.reasoning_parser_plugin) - async with build_async_engine_client( - args, - client_config=client_config, - ) as engine_client: - shutdown_task = await build_and_serve( - engine_client, listen_address, sock, args, **uvicorn_kwargs - ) - # NB: Await server shutdown only after the backend context is exited - try: - await shutdown_task - finally: - sock.close() +__all__ = [ + "build_async_engine_client", + "build_async_engine_client_from_engine_args", + "build_app", + "init_app_state", + "init_render_app_state", + "create_server_socket", + "create_server_unix_socket", + "validate_api_server_args", + "setup_server", + "build_and_serve", + "build_and_serve_renderer", + "run_server", + "run_server_worker", + "register_api_routers", +] if __name__ == "__main__": - # NOTE(simon): - # This section should be in sync with vllm/entrypoints/cli/main.py for CLI - # entrypoints. - cli_env_setup() - parser = FlexibleArgumentParser( - description="vLLM OpenAI-Compatible RESTful API server." + warnings.warn( + "The `python -m vllm.entrypoints.openai.api_server` command is deprecated " + "and may be removed in a future release. Please use `vllm server` instead.", + DeprecationWarning, + stacklevel=1, ) - parser = make_arg_parser(parser) - args = parser.parse_args() - validate_parsed_serve_args(args) + from vllm.entrypoints.launchers.api_server.entry import main - uvloop.run(run_server(args)) + main() diff --git a/vllm/entrypoints/openai/cli_args.py b/vllm/entrypoints/openai/cli_args.py index e9284c5dd69b..49730d34b81f 100644 --- a/vllm/entrypoints/openai/cli_args.py +++ b/vllm/entrypoints/openai/cli_args.py @@ -20,11 +20,11 @@ ChatTemplateContentFormatOption, validate_chat_template, ) -from vllm.entrypoints.openai.models.protocol import LoRAModulePath -from vllm.entrypoints.serve.utils.constants import ( +from vllm.entrypoints.launchers.utils.constants import ( H11_MAX_HEADER_COUNT_DEFAULT, H11_MAX_INCOMPLETE_EVENT_SIZE_DEFAULT, ) +from vllm.entrypoints.openai.models.protocol import LoRAModulePath from vllm.tool_parsers import ToolParserManager from vllm.utils.argparse_utils import FlexibleArgumentParser @@ -444,6 +444,6 @@ def validate_parsed_serve_args(args: argparse.Namespace): def create_parser_for_docs() -> FlexibleArgumentParser: parser_for_docs = FlexibleArgumentParser( - prog="-m vllm.entrypoints.openai.api_server" + prog="-m vllm.entrypoints.launchers.api_server.entry" ) return make_arg_parser(parser_for_docs) diff --git a/vllm/entrypoints/openai/dp_supervisor.py b/vllm/entrypoints/openai/dp_supervisor.py index 4076e1ed3245..c5efe916c025 100644 --- a/vllm/entrypoints/openai/dp_supervisor.py +++ b/vllm/entrypoints/openai/dp_supervisor.py @@ -23,7 +23,7 @@ from fastapi import FastAPI, Response import vllm.envs as envs -from vllm.entrypoints.launcher import NoSignalServer +from vllm.entrypoints.launchers.launcher import NoSignalServer from vllm.logger import init_logger from vllm.utils.system_utils import ( decorate_logs, @@ -236,7 +236,7 @@ async def ready() -> Response: def _run_python_vllm_dp_server(child_args: argparse.Namespace) -> None: - from vllm.entrypoints.openai.api_server import run_server + from vllm.entrypoints.launchers.api_server.entry import run_server uvloop.run(run_server(child_args)) diff --git a/vllm/entrypoints/openai/responses/context.py b/vllm/entrypoints/openai/responses/context.py index 9bc3161f88d1..72b3ea92b8bd 100644 --- a/vllm/entrypoints/openai/responses/context.py +++ b/vllm/entrypoints/openai/responses/context.py @@ -39,7 +39,6 @@ build_response_output_items, construct_tool_dicts, ) -from vllm.entrypoints.serve.utils.constants import MCP_PREFIX from vllm.outputs import RequestOutput from vllm.parser.abstract_parser import Parser from vllm.tokenizers import TokenizerLike @@ -58,6 +57,7 @@ "python": "code_interpreter", "container": "container", } +MCP_PREFIX = "mcp_" def _map_tool_name_to_tool_type(tool_name: str) -> str: diff --git a/vllm/entrypoints/openai/run_batch.py b/vllm/entrypoints/openai/run_batch.py index e9b7c582d6de..f4b7b35bf0e2 100644 --- a/vllm/entrypoints/openai/run_batch.py +++ b/vllm/entrypoints/openai/run_batch.py @@ -31,7 +31,7 @@ from vllm.connections import global_http_connection from vllm.engine.arg_utils import AsyncEngineArgs from vllm.engine.protocol import EngineClient -from vllm.entrypoints.openai.api_server import init_app_state +from vllm.entrypoints.launchers.api_server.entry import init_app_state from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ChatCompletionResponse, @@ -836,8 +836,7 @@ async def run_batch( error_msg=f"URL {request.url} was used. " "Supported endpoints: /v1/chat/completions, /v1/embeddings," " /v1/audio/transcriptions, /v1/audio/translations, /score, " - " /rerank. See vllm/entrypoints/openai/api_server.py " - "for supported score/rerank versions.", + " /rerank.", ) ) @@ -848,7 +847,7 @@ async def run_batch( async def main(args: Namespace): - from vllm.entrypoints.openai.api_server import build_async_engine_client + from vllm.entrypoints.launchers.api_server.entry import build_async_engine_client from vllm.usage.usage_lib import UsageContext validate_run_batch_args(args) diff --git a/vllm/entrypoints/serve/exception_handling/handlers/vllm_error.py b/vllm/entrypoints/serve/exception_handling/handlers/vllm_error.py index 8ac9756f3d7f..213c451167cf 100644 --- a/vllm/entrypoints/serve/exception_handling/handlers/vllm_error.py +++ b/vllm/entrypoints/serve/exception_handling/handlers/vllm_error.py @@ -4,7 +4,7 @@ from fastapi import Request from starlette.responses import JSONResponse -from vllm.entrypoints.launcher import terminate_if_errored +from vllm.entrypoints.launchers.launcher import terminate_if_errored from vllm.entrypoints.openai.engine.protocol import GenerationError from vllm.exceptions import VLLMError from vllm.logger import init_logger diff --git a/vllm/plugins/endpoint_plugins/interface.py b/vllm/plugins/endpoint_plugins/interface.py index 99487f57b68e..00451217d2f2 100644 --- a/vllm/plugins/endpoint_plugins/interface.py +++ b/vllm/plugins/endpoint_plugins/interface.py @@ -20,7 +20,7 @@ posture of exposing plugin defined routes. The CPU only render server (see `build_and_serve_renderer` in -`vllm/entrypoints/openai/api_server.py`) has no `EngineClient`. A plugin +`vllm/entrypoints/launchers/render`) has no `EngineClient`. A plugin eligible for the `render` task (`required_tasks` is `None` or includes `"render"`) still gets `attach_router` called but `init_state` receives `engine_client=None`. Plugins that cannot function without an engine should @@ -34,8 +34,9 @@ from fastapi import FastAPI from starlette.datastructures import State +from vllm.engine.protocol import EngineClient + if TYPE_CHECKING: - from vllm.engine.protocol import EngineClient from vllm.tasks import SupportedTask @@ -85,3 +86,38 @@ async def init_state( cannot function without an engine. """ ... + + +def attach_endpoint_plugins( + app: FastAPI, supported_tasks: tuple["SupportedTask", ...] +) -> None: + """Phase A of endpoint plugin wiring: discover, gate and attach routes. + + Attached last after all core routers. This is so endpoint plugin routes can + shadow core routes with the same path (see `EndpointPlugin.attach_router` + docstring). No-ops when no plugins are discovered/allowlisted. + """ + from vllm.plugins import load_endpoint_plugins + + endpoint_plugins = load_endpoint_plugins(supported_tasks) + for plugin in endpoint_plugins: + plugin.attach_router(app) + app.state.endpoint_plugins = endpoint_plugins + + +async def init_endpoint_plugins_state( + engine_client: EngineClient | None, state: State, args: Namespace +) -> None: + """Phase B of endpoint plugin wiring: initialize per app plugin state. + + `state.endpoint_plugins` is set by `_attach_endpoint_plugins` (Phase A) + in `build_app`. Some `init_app_state` callers (e.g. `run_batch.py`) + build their own bare `State` without going through `build_app`. As a result + `endpoint_plugins` may be absent and are treated that the same as "none attached". + + `engine_client` is `None` for the CPU only render server which has no + engine (see `init_render_app_state`). Plugins must handle a `None` + `engine_client` themselves (see `EndpointPlugin.init_state`). + """ + for plugin in getattr(state, "endpoint_plugins", []): + await plugin.init_state(engine_client, state, args) diff --git a/vllm/tasks.py b/vllm/tasks.py index 366ef0c672f0..d1de4249a107 100644 --- a/vllm/tasks.py +++ b/vllm/tasks.py @@ -41,3 +41,4 @@ def check_removed_pooling_task(task: object) -> None: FRONTEND_TASKS: tuple[FrontendTask, ...] = get_args(FrontendTask) SupportedTask = Literal[GenerationTask, PoolingTask, FrontendTask] +FALLBACK_SUPPORTED_TASKS: tuple[SupportedTask, ...] = ("generate",) diff --git a/vllm/v1/utils.py b/vllm/v1/utils.py index 8d482f58fdf9..0501c2ae3274 100644 --- a/vllm/v1/utils.py +++ b/vllm/v1/utils.py @@ -503,7 +503,7 @@ def run_api_server_worker_proc( ) -> None: """Entrypoint for individual API server worker processes.""" - from vllm.entrypoints.openai.api_server import run_server_worker + from vllm.entrypoints.launchers.api_server.entry import run_server_worker client_config = client_config or {} server_index = client_config.get("client_index", 0) From cba06764d7a9da41e6f535d6355c13f725574f07 Mon Sep 17 00:00:00 2001 From: wangxiyuan Date: Wed, 19 Aug 2026 18:31:19 +0800 Subject: [PATCH 141/839] [Platform] Fill in the missing backend parameter for torch.compile (#51781) Signed-off-by: wangxiyuan --- vllm/model_executor/models/diffusion_gemma.py | 6 +++--- vllm/model_executor/models/kimi_k25_vit.py | 6 +++++- vllm/model_executor/models/parakeet.py | 7 ++++--- vllm/transformers_utils/processors/nano_nemotron_vl.py | 3 ++- vllm/v1/sample/ops/topk_topp_sampler.py | 2 +- 5 files changed, 15 insertions(+), 9 deletions(-) diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 52a3824860ea..3a51e5aad678 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -126,7 +126,7 @@ def get_mm_max_tokens_per_item( return super().get_mm_max_tokens_per_item(seq_len, mm_counts) -@torch.compile(dynamic=True) +@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend) def _softcap_logits(logits: torch.Tensor, cap: float) -> torch.Tensor: # fp32 before tanh for numerical stability (matches HF DiffusionGemma). # Compiling fuses the cast/div/tanh/mul into one elementwise kernel over @@ -454,7 +454,7 @@ def get_placeholder_str(cls, modality: str, i: int) -> str | None: raise ValueError(f"Unsupported modality: {modality}") -@torch.compile(dynamic=True) +@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend) def _compute_num_rejected( num_logits: torch.Tensor, num_sampled: torch.Tensor, @@ -466,7 +466,7 @@ def _compute_num_rejected( return torch.where(is_denoise, query_lens, num_rejected) -@torch.compile(dynamic=True) +@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend) def _compiled_sample_step( # Logits from the model [num_decode * CL, vocab] logits: torch.Tensor, diff --git a/vllm/model_executor/models/kimi_k25_vit.py b/vllm/model_executor/models/kimi_k25_vit.py index dca62d596b42..20f5b105b966 100644 --- a/vllm/model_executor/models/kimi_k25_vit.py +++ b/vllm/model_executor/models/kimi_k25_vit.py @@ -61,7 +61,11 @@ def wrapper(org, interpolation_mode, shape): @get_rope_shape_decorate -@torch.compile(dynamic=True, disable=current_platform.simple_compile_backend == "tpu") +@torch.compile( + dynamic=True, + backend=current_platform.simple_compile_backend, + disable=current_platform.simple_compile_backend == "tpu", +) def get_rope_shape(org, interpolation_mode, shape): return ( F.interpolate( diff --git a/vllm/model_executor/models/parakeet.py b/vllm/model_executor/models/parakeet.py index 5671ba05c56d..ba7b96b8244a 100644 --- a/vllm/model_executor/models/parakeet.py +++ b/vllm/model_executor/models/parakeet.py @@ -19,6 +19,7 @@ from vllm.model_executor.layers.activation import ReLUSquaredActivation from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.platforms import current_platform from vllm.transformers_utils.configs.parakeet import ExtractorConfig, ParakeetConfig logger = init_logger(__name__) @@ -186,7 +187,7 @@ def _torch_extract_fbank_features(self, waveform: torch.Tensor, device: str): ) return self._apply_mel_filters(stft, mel_filters) - @torch.compile(dynamic=True) + @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend) def _apply_mel_filters( self, stft_output: torch.Tensor, mel_filters: torch.Tensor ) -> torch.Tensor: @@ -195,7 +196,7 @@ def _apply_mel_filters( mel_spec = torch.log(mel_spec + LOG_ZERO_GUARD_VALUE) return mel_spec.permute(0, 2, 1) - @torch.compile(dynamic=True) + @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend) def _apply_preemphasis( self, input_features: torch.Tensor, audio_lengths: torch.Tensor ) -> torch.Tensor: @@ -213,7 +214,7 @@ def _apply_preemphasis( input_features = input_features.masked_fill(~timemask, 0.0) return input_features - @torch.compile(dynamic=True) + @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend) def _normalize_mel_features( self, mel_features: torch.Tensor, audio_lengths: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: diff --git a/vllm/transformers_utils/processors/nano_nemotron_vl.py b/vllm/transformers_utils/processors/nano_nemotron_vl.py index 028d207a25ac..022eb28f3f1e 100644 --- a/vllm/transformers_utils/processors/nano_nemotron_vl.py +++ b/vllm/transformers_utils/processors/nano_nemotron_vl.py @@ -26,6 +26,7 @@ from vllm.multimodal.inputs import AudioItem from vllm.multimodal.processing.processor import PromptUpdateDetails from vllm.multimodal.video_prune.evs import compute_retained_tokens_count +from vllm.platforms import current_platform from vllm.tokenizers.hf import HfTokenizer from .internvl import calculate_internvl_targets, get_internvl_target_ratios @@ -56,7 +57,7 @@ def calculate_timestamps( return timestamps -@torch.compile(dynamic=True) +@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend) def _bicubic_resize_and_normalize( tensor: torch.Tensor, size: tuple[int, int] | None = None, diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py index c83ac7c4b90e..e6e0e054c06c 100644 --- a/vllm/v1/sample/ops/topk_topp_sampler.py +++ b/vllm/v1/sample/ops/topk_topp_sampler.py @@ -338,7 +338,7 @@ def forward_xpu( # Note: this is a workaround for # https://github.com/pytorch/pytorch/pull/151218 -@torch.compile(dynamic=True) +@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend) def compiled_random_sample(logits: torch.Tensor) -> torch.Tensor: probs = logits.softmax(dim=-1, dtype=torch.float32) q = torch.empty_like(probs) From b160cab156106b212ee67beb16c9725d366b8857 Mon Sep 17 00:00:00 2001 From: hy123 Date: Wed, 19 Aug 2026 19:18:10 +0800 Subject: [PATCH 142/839] [Bugfix] Restore model info caching for package backends (#52690) Signed-off-by: haoyang.qian Co-authored-by: OpenAI Codex --- tests/models/test_registry.py | 30 ++++++++++++++++++++++++++ vllm/model_executor/models/registry.py | 3 ++- 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/tests/models/test_registry.py b/tests/models/test_registry.py index 70b8b18f76f6..81d6b312152c 100644 --- a/tests/models/test_registry.py +++ b/tests/models/test_registry.py @@ -173,6 +173,36 @@ def test_lazy_modelinfo_package_hash_includes_submodules(tmp_path): assert first_hash != second_hash +def test_lazy_modelinfo_package_attempts_cache_load(monkeypatch): + cached_model_info = object() + loaded_hashes = [] + + def fake_load_cache(self, module_hash): + loaded_hashes.append(module_hash) + return cached_model_info + + monkeypatch.setattr( + _LazyRegisteredModel, + "_load_modelinfo_from_cache", + fake_load_cache, + ) + monkeypatch.setattr( + "vllm.model_executor.models.registry._run_in_subprocess", + lambda _: pytest.fail("Package-backed model should use the cache path"), + ) + + registered_model = _LazyRegisteredModel( + module_name="vllm.model_executor.models.transformers", + class_name="TransformersForCausalLM", + ) + + result = registered_model.inspect_model_cls() + + assert result is cached_model_info + assert len(loaded_hashes) == 1 + assert loaded_hashes[0] + + def test_hf_registry_coverage(): untested_archs = ( ModelRegistry.get_supported_archs() - HF_EXAMPLE_MODELS.get_supported_archs() diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 027dc10e6d53..78e1dae55662 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -1004,9 +1004,10 @@ def inspect_model_cls(self) -> _ModelInfo: # hardware-isolated ``vllm.models.`` layout) live outside # ``vllm/model_executor/models``. Resolve the module spec directly # so the file-hash cache stays warm for them. + model_path: Path | None = None if self.module_name.startswith("vllm.model_executor.models."): model_path = Path(__file__).parent / f"{self.module_name.split('.')[-1]}.py" - else: + if model_path is None or not model_path.exists(): try: spec = importlib.util.find_spec(self.module_name) except (ImportError, ValueError): From 58302b459198eec60fad49af3c506ec10d6af821 Mon Sep 17 00:00:00 2001 From: QWERQWERQWE86 Date: Wed, 19 Aug 2026 19:38:34 +0800 Subject: [PATCH 143/839] [Doc] Fix group numbering in Case 3 of hybrid_kv_cache_manager.md (#52160) Signed-off-by: qwerqwerqwe8688-jpg --- docs/design/hybrid_kv_cache_manager.md | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/docs/design/hybrid_kv_cache_manager.md b/docs/design/hybrid_kv_cache_manager.md index 82d54e9b5c1e..9891c82f97c9 100644 --- a/docs/design/hybrid_kv_cache_manager.md +++ b/docs/design/hybrid_kv_cache_manager.md @@ -115,9 +115,10 @@ Unfortunately, not all models have such a beautiful ratio, and approach in Case - Group 0: 10 full attention layers (full.0 - full.9) - Group 1: 10 sliding window attention layers (sw.0 - sw.9) - Group 2: 10 sliding window attention layers (sw.10 - sw.19) -- ... -- Group 6: 10 sliding window attention layers (sw.40 - sw.49) -- Group 7: 2 sliding window attention layers (sw.50 - sw.51) and 8 padding layers +- Group 3: 10 sliding window attention layers (sw.20 - sw.29) +- Group 4: 10 sliding window attention layers (sw.30 - sw.39) +- Group 5: 10 sliding window attention layers (sw.40 - sw.49) +- Group 6: 2 sliding window attention layers (sw.50 - sw.51) and 8 padding layers We will update this algorithm if this heuristic leads to a bad result when a new model comes out (e.g., 20 full + 30 sw, the group size should be 10 instead of 20). From 2f54100a593e34dd473912c756346abee55a6246 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Wed, 19 Aug 2026 13:25:46 +0100 Subject: [PATCH 144/839] [CI] Fix docs build (#52937) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- docs/mkdocs/gen_files/generate_argparse.py | 44 +++++++++++++--------- vllm/multimodal/video.py | 2 +- 2 files changed, 28 insertions(+), 18 deletions(-) diff --git a/docs/mkdocs/gen_files/generate_argparse.py b/docs/mkdocs/gen_files/generate_argparse.py index c0fa707a1cda..8ba28f60db31 100644 --- a/docs/mkdocs/gen_files/generate_argparse.py +++ b/docs/mkdocs/gen_files/generate_argparse.py @@ -27,16 +27,30 @@ from generated_content import fill_markers # noqa: E402 +HAS_TORCH = importlib.util.find_spec("torch") is not None + def mock_if_no_torch(mock_module: str, mock: MagicMock): - if not importlib.util.find_spec("torch"): + if not HAS_TORCH: sys.modules[mock_module] = mock +class PydanticMagicMock(MagicMock): + """`MagicMock` that's able to generate pydantic-core schemas.""" + + def __init__(self, *args, **kwargs): + name = kwargs.get("name") + super().__init__(*args, **kwargs) + self.__spec__ = ModuleSpec(name, None) + + def __get_pydantic_core_schema__(self, source_type, handler): + return core_schema.any_schema() + + # Mock custom op code class MockCustomOp: @staticmethod - def register(name): + def register(*args, **kwargs): def decorator(cls): return cls @@ -45,7 +59,7 @@ def decorator(cls): class MockPluggableLayer: @staticmethod - def register(name): + def register(*args, **kwargs): def decorator(cls): return cls @@ -69,8 +83,16 @@ def decorator(cls): importlib.metadata.version = lambda name: VERSIONS.get(name) or "0.0.0" -# Make torch.nn.Parameter safe to inherit from -mock_if_no_torch("torch.nn", MagicMock(Parameter=object)) +# Real class because a `MagicMock` base conflicts with `ABCMeta` +class MockModule: + pass + + +# Make torch.nn.Parameter and torch.nn.Module safe to inherit from. +# `import torch.nn` resolves the attribute on `torch`, so mock both. +mock_nn = MagicMock(Parameter=object, Module=MockModule) +mock_if_no_torch("torch.nn", mock_nn) +mock_if_no_torch("torch", PydanticMagicMock(name="torch", nn=mock_nn)) # Mock torch.library.infer_schema for vllm.ir.ops.IrOpInplaceOverload.__init__ @@ -98,18 +120,6 @@ def get_outputs(native_fn: Callable) -> str: ) -class PydanticMagicMock(MagicMock): - """`MagicMock` that's able to generate pydantic-core schemas.""" - - def __init__(self, *args, **kwargs): - name = kwargs.get("name") - super().__init__(*args, **kwargs) - self.__spec__ = ModuleSpec(name, None) - - def __get_pydantic_core_schema__(self, source_type, handler): - return core_schema.any_schema() - - def auto_mock(module_name: str, attr: str, max_mocks: int = 100): """Function that automatically mocks missing modules during imports.""" logger.info("Importing %s from %s", attr, module_name) diff --git a/vllm/multimodal/video.py b/vllm/multimodal/video.py index 4f750c5cfa1b..f1a605311a16 100644 --- a/vllm/multimodal/video.py +++ b/vllm/multimodal/video.py @@ -220,7 +220,7 @@ def load_bytes( frame_recovery: bool = False, *, backend: VideoDecoderBackend = "opencv", - **kwargs, + **kwargs: Any, ) -> tuple[npt.NDArray, dict[str, Any]]: """Load sampled frames from raw video bytes. From 58de6cbdc2fac255ca6fe6855f90c65f80a76d57 Mon Sep 17 00:00:00 2001 From: Nave Assaf <55059536+Naveassaf@users.noreply.github.com> Date: Wed, 19 Aug 2026 15:34:50 +0300 Subject: [PATCH 145/839] Add NemotronH_Omni_Reasoning_V3 as a supported Nemotron architecture (#52929) Signed-off-by: Nave Assaf --- tests/models/registry.py | 4 ++++ vllm/config/speculative.py | 5 ++++- vllm/model_executor/models/registry.py | 1 + 3 files changed, 9 insertions(+), 1 deletion(-) diff --git a/tests/models/registry.py b/tests/models/registry.py index 0411d64f9489..c862e767cd46 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -1250,6 +1250,10 @@ def check_available_online( "NemotronH_Super_Omni_Reasoning_V3": _HfExamplesInfo( "nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16", is_available_online=False ), + # TODO: Change repo id once pertinent archs are public. + "NemotronH_Omni_Reasoning_V3": _HfExamplesInfo( + "nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16", is_available_online=False + ), "OpenCUAForConditionalGeneration": _HfExamplesInfo( "xlangai/OpenCUA-7B", trust_remote_code=True, diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 6cdba55263c3..e378906535ca 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -490,7 +490,10 @@ def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig: {"n_predict": n_predict, "architectures": ["ErnieMTPModel"]} ) - if hf_config.architectures[0] == "NemotronH_Super_Omni_Reasoning_V3": + if hf_config.architectures[0] in ( + "NemotronH_Super_Omni_Reasoning_V3", + "NemotronH_Omni_Reasoning_V3", + ): # Promote VLM's text_config so MTP detection below fires correctly hf_config = hf_config.text_config diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 78e1dae55662..75ef6756ddc9 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -535,6 +535,7 @@ "NemotronH_Nano_VL_V2": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"), "NemotronH_Nano_Omni_Reasoning_V3": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"), "NemotronH_Super_Omni_Reasoning_V3": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"), + "NemotronH_Omni_Reasoning_V3": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"), "NVLM_D": ("nvlm_d", "NVLM_D_Model"), "MuseGlimmerForConditionalGeneration": ("muse_glimmer", "MuseGlimmerForCausalLM"), "OpenCUAForConditionalGeneration": ("opencua", "OpenCUAForConditionalGeneration"), From 92bdee05cb4ea5e94c4cce3eb0544f51d6eece8d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <45557362+qgallouedec@users.noreply.github.com> Date: Wed, 19 Aug 2026 08:48:45 -0400 Subject: [PATCH 146/839] [Bugfix][Frontend] Return all choices from /inference/v1/generate when n > 1 (#52399) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Quentin Gallouédec --- .../token_in_token_out/test_serving_tokens.py | 24 +++++++++++++++++++ .../scale_out/token_in_token_out/serving.py | 5 ++-- 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py b/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py index a1491c1c3e1b..be33a182bccd 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py @@ -183,6 +183,30 @@ async def test_generate_defaults_max_tokens_when_omitted(client): ) +@pytest.mark.asyncio +@pytest.mark.skipif( + envs.VLLM_USE_RUST_FRONTEND, + reason="parallel sampling (n > 1) is not supported by the Rust frontend", +) +async def test_generate_returns_all_choices_when_n_greater_than_one(client): + """Regression: ``n > 1`` must return ``n`` choices. + + Non-streaming requests kept ``SamplingParams``' default output kind, + ``CUMULATIVE``, so only the sequences updated during the last engine step + reached the response and the others were silently dropped. + """ + payload = { + "model": MODEL_NAME, + "token_ids": [1, 2, 3], + "sampling_params": {"max_tokens": 5, "temperature": 1.0, "n": 4}, + "stream": False, + } + resp = await client.post(GEN_ENDPOINT, json=payload) + resp.raise_for_status() + data = resp.json() + assert sorted(choice["index"] for choice in data["choices"]) == [0, 1, 2, 3] + + @pytest.mark.asyncio async def test_generate_stream(client): payload = { diff --git a/vllm/entrypoints/scale_out/token_in_token_out/serving.py b/vllm/entrypoints/scale_out/token_in_token_out/serving.py index 9e9ace877a7c..fbbb8052bed3 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/serving.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/serving.py @@ -218,8 +218,9 @@ async def serve_tokens( if self.force_no_detokenize: sampling_params.detokenize = False - if request.stream: - sampling_params.output_kind = RequestOutputKind.DELTA + sampling_params.output_kind = ( + RequestOutputKind.DELTA if request.stream else RequestOutputKind.FINAL_ONLY + ) self._log_inputs( request_id, From c2e7242ab7287cc9e2bdbea5901c560fc2d17bff Mon Sep 17 00:00:00 2001 From: Eilam <155648591+eilamc14@users.noreply.github.com> Date: Wed, 19 Aug 2026 16:24:26 +0300 Subject: [PATCH 147/839] [Bugfix][LoRA] Guard None group members in expand_packed_lora (partial LoRA on Qwen3.5/3.6 GatedDeltaNet) (#47640) Signed-off-by: Eilam C Co-authored-by: Jee Jee Li --- vllm/lora/layers/column_parallel_linear.py | 31 +++++++++++++++++----- 1 file changed, 25 insertions(+), 6 deletions(-) diff --git a/vllm/lora/layers/column_parallel_linear.py b/vllm/lora/layers/column_parallel_linear.py index 12151699ac48..ea776ec2ea4f 100644 --- a/vllm/lora/layers/column_parallel_linear.py +++ b/vllm/lora/layers/column_parallel_linear.py @@ -265,17 +265,36 @@ def slice_lora_b( def expand_packed_lora( self, - lora_a: list[torch.Tensor], - lora_b: list[torch.Tensor], - ) -> tuple[list[torch.Tensor], list[torch.Tensor]]: + lora_a: list[torch.Tensor | None], + lora_b: list[torch.Tensor | None], + ) -> tuple[list[torch.Tensor | None], list[torch.Tensor | None]]: """ Expand packed adapter groups when they don't match n_slices. - E.g. in_proj_qkv (covers Q+K+V) + in_proj_z + E.g. in_proj_qkv (covers Q+K+V) + in_proj_z. + + A None group member means that member was not adapted; the slice(s) + it covers are emitted as None placeholders so subsequent groups stay + aligned and those slices are left at base weights. This matches the + None-tolerance already present in slice_lora_b() and the set_lora() + stacking loop. """ - expanded_a: list[torch.Tensor] = [] - expanded_b: list[torch.Tensor] = [] + expanded_a: list[torch.Tensor | None] = [] + expanded_b: list[torch.Tensor | None] = [] start_idx = 0 for a_i, b_i in zip(lora_a, lora_b): + if b_i is None: + # Unadapted group member: its row count is unknown (the tensor + # is missing), so infer its coverage as the remaining slices. + # This is exact for the only layout that reaches this path, + # the fused GDN in_proj_qkvz group, whose sole multi-slice + # member (in_proj_qkv, Q+K+V) leads and whose only optional + # member (in_proj_z) trails. + covered = self.n_slices - start_idx + for _ in range(covered): + expanded_a.append(None) + expanded_b.append(None) + start_idx += covered + continue # Determine which output slices this b_i covers. b_rows, cu_rows, covered = b_i.shape[0], 0, 0 for i in range(start_idx, self.n_slices): From 2b7fcbf52782f8729fd6ce6c9ab803617d72897b Mon Sep 17 00:00:00 2001 From: Gabriel Wu <13583761+lucifer1004@users.noreply.github.com> Date: Wed, 19 Aug 2026 21:25:29 +0800 Subject: [PATCH 148/839] [Kernel] SM120: stop routing misaligned-M blockwise FP8 GEMMs to the small-M swapAB config (#52775) Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> --- .../w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh index a9008ce44240..8795326250bf 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh +++ b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh @@ -250,7 +250,7 @@ void cutlass_gemm_blockwise_sm120_fp8_dispatch(torch::stable::Tensor& out, torch::stable::Tensor const& b_scales) { int M = a.size(0); // more heuristic tuning can be done here by checking N/K dimensions as well - bool swap_ab = (M <= 64) || (M % 4 != 0); + bool swap_ab = (M <= 64); if (!swap_ab) { if (M <= 256) { From db92053e97b5630a6a36118693b1dffe9b03be36 Mon Sep 17 00:00:00 2001 From: "Xiaoan (Sean) Liu" <95677580+sseanliu@users.noreply.github.com> Date: Wed, 19 Aug 2026 07:16:39 -0700 Subject: [PATCH 149/839] [Core] Skip broadcasting mm tensor data to workers for prefix-cache-covered items (#52041) Signed-off-by: Xiaoan (Sean) Liu Signed-off-by: Nick Hill Co-authored-by: Nick Hill --- tests/v1/core/test_output.py | 87 +++++++++++++++++++++++++++++++++ vllm/multimodal/utils.py | 31 ++++++++++++ vllm/v1/core/sched/output.py | 8 ++- vllm/v1/core/sched/scheduler.py | 6 ++- 4 files changed, 130 insertions(+), 2 deletions(-) diff --git a/tests/v1/core/test_output.py b/tests/v1/core/test_output.py index 9dea19320e61..bac225f5f0cc 100644 --- a/tests/v1/core/test_output.py +++ b/tests/v1/core/test_output.py @@ -2,6 +2,14 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch +from vllm.multimodal.inputs import ( + MultiModalBatchedField, + MultiModalFeatureSpec, + MultiModalFieldElem, + MultiModalKwargsItem, + PlaceholderRange, +) +from vllm.multimodal.utils import strip_covered_mm_data from vllm.v1.core.sched.output import NewRequestData @@ -34,3 +42,82 @@ def test_repr_with_multi_element_tensor() -> None: assert "prompt_embeds_shape=torch.Size([10, 768])" in repr(new_requests_data) assert "prompt_embeds_shape=torch.Size([10, 768])" in new_requests_data.anon_repr() + + +def _mm_feature(offset: int, length: int) -> MultiModalFeatureSpec: + return MultiModalFeatureSpec( + data=MultiModalKwargsItem.dummy(), + mm_position=PlaceholderRange(offset=offset, length=length), + identifier=f"hash_{offset}", + modality="image", + ) + + +def test_strip_covered_mm_data() -> None: + """Items fully inside the computed prefix lose their data; items touching + the uncomputed region keep it; already-None data stays None.""" + from dataclasses import replace + + covered = _mm_feature(offset=0, length=100) + boundary = _mm_feature(offset=150, length=100) # ends exactly at 250 + uncovered = _mm_feature(offset=300, length=100) + already_none = replace(_mm_feature(offset=100, length=50), data=None) + + stripped = strip_covered_mm_data( + [covered, boundary, uncovered, already_none], num_computed_tokens=250 + ) + + assert stripped[0].data is None # fully covered -> stripped + assert stripped[1].data is None # span end == computed -> covered -> stripped + assert stripped[2].data is not None # extends past prefix -> kept + assert stripped[3].data is None # was already None + # non-data fields are preserved + assert stripped[0].identifier == covered.identifier + assert stripped[0].mm_position == covered.mm_position + # original list is not mutated + assert covered.data is not None + + +def test_strip_covered_mm_data_zero_computed() -> None: + """With no prefix hit nothing is stripped.""" + features = [_mm_feature(offset=0, length=100)] + stripped = strip_covered_mm_data(features, num_computed_tokens=0) + assert stripped[0].data is not None + + +def _mm_feature_mixed(offset: int, length: int) -> MultiModalFeatureSpec: + data = MultiModalKwargsItem( + { + "pixel_values": MultiModalFieldElem( + data=torch.empty(4), field=MultiModalBatchedField() + ), + "image_grid_thw": MultiModalFieldElem( + data=torch.ones(1, 3, dtype=torch.long), + field=MultiModalBatchedField(keep_on_cpu=True), + ), + } + ) + return MultiModalFeatureSpec( + data=data, + mm_position=PlaceholderRange(offset=offset, length=length), + identifier=f"hash_{offset}", + modality="image", + ) + + +def test_strip_covered_mm_data_mrope() -> None: + """For M-RoPE models, covered items keep their keep_on_cpu metadata fields + (the worker needs them to compute positions); payload fields are dropped.""" + covered = _mm_feature_mixed(offset=0, length=100) + uncovered = _mm_feature_mixed(offset=300, length=100) + + stripped = strip_covered_mm_data( + [covered, uncovered], num_computed_tokens=250, uses_mrope=True + ) + + assert stripped[0].data is not None + assert list(stripped[0].data.keys()) == ["image_grid_thw"] + assert stripped[1].data is not None + assert set(stripped[1].data.keys()) == {"pixel_values", "image_grid_thw"} + # original list is not mutated + assert set(covered.data.keys()) == {"pixel_values", "image_grid_thw"} diff --git a/vllm/multimodal/utils.py b/vllm/multimodal/utils.py index fc400710f096..ac9b028b46d2 100644 --- a/vllm/multimodal/utils.py +++ b/vllm/multimodal/utils.py @@ -5,6 +5,7 @@ import mimetypes from collections import defaultdict from collections.abc import Generator, Sequence +from dataclasses import replace from itertools import groupby from typing import TYPE_CHECKING, Any @@ -212,6 +213,36 @@ def _batch_mm_items( } +def strip_covered_mm_data( + mm_features: list[MultiModalFeatureSpec], + num_computed_tokens: int, + uses_mrope: bool = False, +) -> list[MultiModalFeatureSpec]: + """Drop the tensor data of mm items whose placeholder span is fully inside + a prefix-cache-covered region: no encoder run can be scheduled for them, + so the workers never consume the payload fields. M-RoPE models are the + exception: the worker computes positions for the whole prompt from the + CPU-side metadata fields (e.g. grid dims), so those are kept. The + scheduler-side ``Request`` keeps the full features.""" + if not mm_features or num_computed_tokens == 0: + return mm_features + + def maybe_strip(f: MultiModalFeatureSpec) -> MultiModalFeatureSpec: + if f.data is None or ( + f.mm_position.offset + f.mm_position.length > num_computed_tokens + ): + return f + + data = None + if uses_mrope: + data = MultiModalKwargsItem( + {k: elem for k, elem in f.data.items() if elem.field.keep_on_cpu} + ) + return replace(f, data=data) + + return [maybe_strip(f) for f in mm_features] + + def group_and_batch_mm_items( items: Sequence[MultiModalKwargsItem], *, diff --git a/vllm/v1/core/sched/output.py b/vllm/v1/core/sched/output.py index 4d11be021ed4..b26a729264c9 100644 --- a/vllm/v1/core/sched/output.py +++ b/vllm/v1/core/sched/output.py @@ -6,6 +6,7 @@ from typing import TYPE_CHECKING from vllm.config.ec_manager_config import EncoderCacheManagerMetadata +from vllm.multimodal.utils import strip_covered_mm_data if TYPE_CHECKING: import numpy as np @@ -53,11 +54,16 @@ def from_request( request: Request, block_ids: tuple[list[int], ...], prefill_token_ids: list[int] | None = None, + uses_mrope: bool = False, ) -> "NewRequestData": return cls( req_id=request.request_id, prompt_token_ids=request.prompt_token_ids, - mm_features=request.mm_features, + mm_features=strip_covered_mm_data( + request.mm_features, + request.num_computed_tokens, + uses_mrope=uses_mrope, + ), sampling_params=request.sampling_params, pooling_params=request.pooling_params, block_ids=block_ids, diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index e4a21328660a..2115ca673496 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -82,6 +82,7 @@ def __init__( self.scheduler_config = vllm_config.scheduler_config self.cache_config = vllm_config.cache_config self.lora_config = vllm_config.lora_config + self.model_uses_mrope = vllm_config.model_config.uses_mrope self.kv_cache_config = kv_cache_config self.kv_events_config = vllm_config.kv_events_config self.parallel_config = vllm_config.parallel_config @@ -1199,13 +1200,16 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: req, req_to_new_blocks[req.request_id].get_block_ids(), req._all_token_ids, + uses_mrope=self.model_uses_mrope, ) for req in scheduled_new_reqs ] else: new_reqs_data = [ NewRequestData.from_request( - req, req_to_new_blocks[req.request_id].get_block_ids() + req, + req_to_new_blocks[req.request_id].get_block_ids(), + uses_mrope=self.model_uses_mrope, ) for req in scheduled_new_reqs ] From 2d7f42b4f3b2b69dd9b2610d287083d147c997f6 Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Wed, 19 Aug 2026 10:55:09 -0400 Subject: [PATCH 150/839] [Build] Add InstantTensor to CUDA dependencies (#52801) Signed-off-by: mgoin Co-authored-by: OpenAI Codex --- requirements/cuda.txt | 1 + requirements/test/cuda.txt | 4 +++- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/requirements/cuda.txt b/requirements/cuda.txt index 530c223bea53..a2c0364d1f4f 100644 --- a/requirements/cuda.txt +++ b/requirements/cuda.txt @@ -23,6 +23,7 @@ nvidia-cudnn-frontend>=1.19.1 nvtx==0.2.15 # Required for faster safetensors model loading fastsafetensors >= 0.3.3 +instanttensor >= 0.1.9 # QuACK and Cutlass DSL for FA4 (cute-DSL implementation) nvidia-cutlass-dsl[cu13]==4.6.2 diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index 2af564acb806..f28f733cd47b 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -406,7 +406,9 @@ inflect==5.6.2 iniconfig==2.0.0 # via pytest instanttensor==0.1.9 - # via -r requirements/test/cuda.in + # via + # -c requirements/cuda.txt + # -r requirements/test/cuda.in interegular==0.3.3 # via lm-format-enforcer isodate==0.7.2 From be06873198cd544f963b4e2f0f0877b23ab5861e Mon Sep 17 00:00:00 2001 From: Yohann Prigent Date: Wed, 19 Aug 2026 16:55:59 +0200 Subject: [PATCH 151/839] [Bugfix] compressed-tensors: restore int8 grouped WNA16 MoE support (#52002) Signed-off-by: Yohann Prigent Co-authored-by: Misha Goin --- .../compressed_tensors_moe/compressed_tensors_moe_wna16.py | 1 - 1 file changed, 1 deletion(-) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py index bd7ad4234a5f..35894b322482 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py @@ -82,7 +82,6 @@ def __init__( else: scale = kInt4StaticGroupScale elif self.num_bits == 8: - assert self.group_size == -1 scale = kInt8StaticGroupScale else: raise ValueError( From 160f7f0840e1773b2418199800c7ac66f1b4f074 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Wed, 19 Aug 2026 10:19:42 -0500 Subject: [PATCH 152/839] [ROCm][CI] Extended Fused MoE and FP8 MoE test support (#41100) Signed-off-by: Andreas Karatzas Signed-off-by: Andreas Karatzas --- .buildkite/test-amd.yaml | 59 +++++++- .buildkite/test_areas/kernels.yaml | 53 +++++++ .../moe/test_modular_oai_triton_moe.py | 29 +++- tests/kernels/moe/test_moe.py | 41 ++++++ tests/kernels/moe/test_moe_layer.py | 132 ++++++++++++++---- tests/kernels/moe/utils.py | 27 +++- tests/quantization/test_fp8.py | 13 ++ .../layers/fused_moe/fused_moe.py | 6 + .../utils/nvfp4_emulation_utils.py | 2 +- 9 files changed, 331 insertions(+), 31 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index bb454f7f6adc..fdeaf788e0b8 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -3500,9 +3500,44 @@ steps: - vllm/_aiter_ops.py - vllm/platforms/rocm.py commands: - - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - pytest -v -s kernels/moe + --ignore=kernels/moe/test_modular_oai_triton_moe.py + --ignore=kernels/moe/test_gpt_oss_triton_kernels.py + --ignore=kernels/moe/test_moe.py + --ignore=kernels/moe/test_block_int8.py + --ignore=kernels/moe/test_triton_moe_no_act_mul.py + --ignore=kernels/moe/test_triton_moe_ptpc_fp8.py + --ignore=kernels/moe/test_deepep_moe.py + --ignore=kernels/moe/test_moe_layer.py + --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT +- label: Kernels FusedMoE Layer Test (2xB200-2xMI355) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_2 + num_gpus: 2 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - csrc/moe/ + - csrc/rocm/ + - tests/kernels/moe + - vllm/model_executor/layers/fused_moe/ + - vllm/model_executor/layers/quantization/ + - vllm/distributed/ + - vllm/config/ + - vllm/forward_context.py + - vllm/v1/worker/workspace.py + - vllm/utils/import_utils.py + - vllm/utils/math_utils.py + - vllm/utils/torch_utils.py + - vllm/platforms/ + - vllm/_aiter_ops.py + commands: + - pytest -v -s kernels/moe/test_moe_layer.py + - label: Kernels Quantization Test %N # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] @@ -3526,12 +3561,34 @@ steps: commands: - pytest -v -s kernels/quantization +- label: Kernels FP8 MoE Test (2xH100-1xMI355) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_1 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - csrc/moe/ + - vllm/model_executor/layers/fused_moe/ + - tests/kernels/moe/test_deepep_moe.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + - vllm/envs.py + commands: + - pytest -v -s kernels/moe/test_gpt_oss_triton_kernels.py + - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py + - pytest -v -s kernels/moe/test_moe.py + - pytest -v -s kernels/moe/test_block_int8.py + - pytest -v -s kernels/moe/test_triton_moe_no_act_mul.py + - pytest -v -s kernels/moe/test_triton_moe_ptpc_fp8.py + - label: Kernels FP8 MoE Test (2xH100-2xMI355) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false agent_pool: mi355_2 num_gpus: 2 + optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - csrc/moe/ diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 0887a6fb27a4..e1fddcd7cbb5 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -451,3 +451,56 @@ steps: commands: - pytest -v -s kernels/moe/test_moe_layer.py - pytest -v -s kernels/moe/test_deepep_v2_moe.py + +- label: Kernels FusedMoE Layer Test (2xMI355) + key: kernels-fusedmoe-layer-test-2-mi355 + depends_on: + - image-build-amd + timeout_in_minutes: 180 + dind: false + device: mi355_2 + soft_fail: true + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - csrc/moe/ + - csrc/rocm/ + - tests/kernels/moe + - vllm/model_executor/layers/fused_moe/ + - vllm/model_executor/layers/quantization/ + - vllm/distributed/ + - vllm/config/ + - vllm/forward_context.py + - vllm/v1/worker/workspace.py + - vllm/utils/import_utils.py + - vllm/utils/math_utils.py + - vllm/utils/torch_utils.py + - vllm/platforms/ + - vllm/_aiter_ops.py + commands: + - pytest -v -s kernels/moe/test_moe_layer.py + +- label: Kernels FP8 MoE Test (MI355) + key: kernels-fp8-moe-test-mi355 + depends_on: + - image-build-amd + timeout_in_minutes: 180 + dind: false + device: mi355_1 + soft_fail: true + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - csrc/moe/ + - vllm/model_executor/layers/fused_moe/ + - tests/kernels/moe/test_deepep_moe.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + - vllm/envs.py + commands: + - pytest -v -s kernels/moe/test_gpt_oss_triton_kernels.py + - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py + - pytest -v -s kernels/moe/test_moe.py + - pytest -v -s kernels/moe/test_block_int8.py + - pytest -v -s kernels/moe/test_triton_moe_no_act_mul.py + - pytest -v -s kernels/moe/test_triton_moe_ptpc_fp8.py diff --git a/tests/kernels/moe/test_modular_oai_triton_moe.py b/tests/kernels/moe/test_modular_oai_triton_moe.py index b3a42912da7f..0fdce2b82b7b 100644 --- a/tests/kernels/moe/test_modular_oai_triton_moe.py +++ b/tests/kernels/moe/test_modular_oai_triton_moe.py @@ -8,6 +8,7 @@ import pytest import torch +import torch.nn.functional as F from tests.utils import wait_for_gpu_memory_to_clear from vllm.model_executor.layers.fused_moe.activation import MoEActivation @@ -37,6 +38,7 @@ ) from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel from vllm.platforms import current_platform +from vllm.utils.math_utils import round_up from vllm.utils.torch_utils import set_random_seed from .utils import make_dummy_moe_config, shuffle_weight @@ -104,13 +106,30 @@ def make_weights(dtype, k, n, e): w1_tri = shuffle_weight(w1_tri) w1_bias_tri = shuffle_weight(w1_bias_tri) + if current_platform.is_rocm(): + k_align, n2_align = 256, 512 + else: + k_align, n2_align = 64, 128 + + w1_bottom_pad = round_up(w1_tri.shape[1], k_align) - w1_tri.shape[1] + w1_right_pad = round_up(w1_tri.shape[2], n2_align) - w1_tri.shape[2] + w2_bottom_pad = w1_right_pad // 2 + w2_right_pad = w1_bottom_pad + + w1_tri = F.pad(w1_tri, (0, w1_right_pad, 0, w1_bottom_pad, 0, 0)) + w2_tri = F.pad(w2_tri, (0, w2_right_pad, 0, w2_bottom_pad, 0, 0)) + w1_bias_tri = F.pad(w1_bias_tri, (0, w1_right_pad, 0, 0)) + w2_bias_tri = F.pad(w2_bias_tri, (0, w2_right_pad, 0, 0)) + # quant triton_weights w1_tri, w1_scale_tri = downcast_to_mxfp(w1_tri, torch.uint8, axis=1) w1 = upcast_from_mxfp(w1_tri, w1_scale_tri, dtype, axis=1) + w1 = w1[..., :k, : 2 * n] w1 = unshuffle_weight(w1) w2_tri, w2_scale_tri = downcast_to_mxfp(w2_tri, torch.uint8, axis=1) w2 = upcast_from_mxfp(w2_tri, w2_scale_tri, dtype, axis=1) + w2 = w2[..., :n, :k] num_warps = 8 w_layout, w_layout_opts = layout.make_default_matmul_mxfp4_w_layout(mx_axis=1) @@ -150,6 +169,7 @@ def make_weights(dtype, k, n, e): w2_bias_tri, w1_precision_config, w2_precision_config, + w1_bottom_pad, ) @@ -237,7 +257,7 @@ def oai_triton_moe_impl( @pytest.mark.skipif( - not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform." + not current_platform.is_cuda_alike(), reason="Requires CUDA-alike platform." ) @pytest.mark.parametrize("dtype", [torch.bfloat16]) @pytest.mark.parametrize("m,n,k", MNK) @@ -256,6 +276,7 @@ def test_oai_triton_moe( ): wait_for_gpu_memory_to_clear(devices=[0], threshold_ratio=0.1) set_random_seed(0) + ( w1, w2, @@ -267,9 +288,11 @@ def test_oai_triton_moe( w2_bias_tri, w1_precision_config, w2_precision_config, + x_pad, ) = make_weights(dtype, k, n, num_experts) x = torch.randn((m, k), dtype=dtype, device="cuda") + x_tri = F.pad(x, (0, x_pad, 0, 0)) router_logits = torch.randn(m, num_experts, device="cuda", dtype=dtype) topk_weights, topk_ids = torch.topk(router_logits, k=topk, dim=-1, sorted=True) topk_weights = torch.nn.functional.softmax(topk_weights, dim=-1) @@ -278,7 +301,7 @@ def test_oai_triton_moe( out_ref = torch_moe_impl(x, w1, w2, w1_bias, w2_bias, topk_weights, topk_ids) out = oai_triton_moe_impl( - x, + x_tri, w1_tri, w2_tri, w1_precision_config, @@ -290,6 +313,7 @@ def test_oai_triton_moe( topk_ids, unfused, ) + out = out[..., :k] assert_close(ref=out_ref, tri=out, maxtol=0.025, rmstol=0.005) @@ -321,6 +345,7 @@ def test_unfused_oai_triton_experts_apply_direct_deepseek_v4_topology(workspace_ w2_bias_tri, w1_precision_config, w2_precision_config, + _x_pad, ) = make_weights(dtype, k, n, num_experts) x = torch.randn((m, k), dtype=dtype, device="cuda") diff --git a/tests/kernels/moe/test_moe.py b/tests/kernels/moe/test_moe.py index 5db21f671c2d..17b8cae318bd 100644 --- a/tests/kernels/moe/test_moe.py +++ b/tests/kernels/moe/test_moe.py @@ -16,6 +16,7 @@ from torch.nn import functional as F import vllm.model_executor.layers.fused_moe # noqa +import vllm.model_executor.layers.fused_moe.fused_moe as fused_moe_module from tests.kernels.moe.utils import ( fused_moe, make_dummy_moe_config, @@ -68,6 +69,46 @@ DEVICE_TYPE = current_platform.device_type +def test_triton_moe_launcher_passes_scalar_scale_as_pointer(monkeypatch) -> None: + captured: dict[str, torch.Tensor] = {} + + class FakeKernel: + def __getitem__(self, grid): + def launch(*args, **kwargs) -> None: + captured["a_scale"] = args[4] + + return launch + + monkeypatch.setattr(fused_moe_module, "fused_moe_kernel", FakeKernel()) + + a_scale = torch.tensor(0.5) + fused_moe_module.invoke_fused_moe_triton_kernel( + A=torch.ones((1, 1)), + B=torch.ones((1, 1, 1)), + C=torch.empty((1, 1, 1)), + A_scale=a_scale, + B_scale=torch.ones(1), + topk_weights=torch.ones((1, 1)), + sorted_token_ids=None, + expert_ids=torch.zeros(1, dtype=torch.int32), + num_tokens_post_padded=torch.ones(1, dtype=torch.int32), + mul_routed_weight=True, + top_k=1, + config={"BLOCK_SIZE_M": 1, "BLOCK_SIZE_N": 1, "BLOCK_SIZE_K": 1}, + compute_type=tl.float32, + use_fp8_w8a8=True, + use_int8_w8a8=False, + use_int8_w8a16=False, + use_int4_w4a16=False, + per_channel_quant=False, + ) + + captured_scale = captured["a_scale"] + assert a_scale.ndim == 0 + assert captured_scale.shape == (1,) + assert captured_scale.data_ptr() == a_scale.data_ptr() + + def iterative_moe( hidden_states: torch.Tensor, w1: torch.Tensor, diff --git a/tests/kernels/moe/test_moe_layer.py b/tests/kernels/moe/test_moe_layer.py index 74554944a314..8a081784fb1d 100644 --- a/tests/kernels/moe/test_moe_layer.py +++ b/tests/kernels/moe/test_moe_layer.py @@ -7,9 +7,11 @@ import functools import os +import tempfile import traceback import types from collections.abc import Callable +from contextlib import suppress from dataclasses import astuple, dataclass, fields from itertools import product from typing import get_args @@ -71,6 +73,15 @@ is_workspace_manager_initialized, ) + +def on_gfx950() -> bool: + if current_platform.is_rocm(): + from vllm.platforms.rocm import on_gfx950 as rocm_on_gfx950 + + return rocm_on_gfx950() + return False + + fp8_dtype = torch.float8_e4m3fn # current_platform.fp8_dtype SHAPE_COMBOS = [ @@ -114,6 +125,7 @@ BACKENDS += ["nixl_ep"] DEEPEP_BACKENDS = {"deepep_high_throughput", "deepep_low_latency"} +MORI_BACKENDS = {"mori_high_throughput", "mori_low_latency"} QUANT_METHODS = [ None, @@ -484,12 +496,12 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]: "leads to large differences.", ) - # Skip modelopt_fp4 if not on B100+ (compute capability 10.0+) - if ( - config.quantization == "modelopt_fp4" - and not current_platform.has_device_capability(100) + # Skip modelopt_fp4 if not on B100+ (compute capability 10.0+) or gfx950. + if config.quantization == "modelopt_fp4" and not ( + (current_platform.is_rocm() and on_gfx950()) + or current_platform.has_device_capability(100) ): - return False, "modelopt_fp4 not supported on H100+ GPUs" + return False, "modelopt_fp4 requires native NVFP4 or emulation" # Skip flashinfer_nvlink if not on H100+ (compute capability 10.0+) if ( @@ -508,6 +520,17 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]: f"{config.backend} does not support quantization={config.quantization}", ) + if config.backend in MORI_BACKENDS: + if os.environ.get("VLLM_TEST_ENABLE_MORI_MOE_LAYER") != "1": + return False, "mori MoE layer matrix is opt-in" + + from vllm._aiter_ops import rocm_aiter_ops + + if not rocm_aiter_ops.is_fused_moe_enabled(): + return False, "mori requires AITER fused MoE" + if rocm_aiter_ops.is_fusion_moe_shared_experts_enabled(): + return False, "mori does not support AITER shared expert fusion" + if config.backend == "deepep_low_latency": from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll import ( # noqa: E501 DeepEPLLPrepareAndFinalize, @@ -1695,14 +1718,16 @@ def _parallel_worker( cpu_group, test_configs: list[MoETestConfig], verbosity: int, + failure_report_path: str | None = None, **kwargs, ) -> None: set_random_seed(7) + is_logging_rank = pgi.rank == 0 total = 0 passed = 0 failed = 0 - fail_ids = [] + failure_details = [] dp_rank = vllm_config.parallel_config.data_parallel_rank @@ -1717,9 +1742,11 @@ def _parallel_worker( tp_rank = pgi.rank % test_config.tp_size - if verbosity > 0: + if verbosity > 0 and is_logging_rank: print(f"subtest: {test_config.id()}", end="") + local_failed = False + local_error: str | None = None try: _run_one_config( vllm_config, @@ -1743,19 +1770,9 @@ def _parallel_worker( use_gate=test_config.use_gate, use_routed_input_transform=test_config.use_routed_input_transform, ) - if verbosity > 0: - print(" PASSED") - else: - print(".", end="") - passed = passed + 1 - except Exception as ex: - fail_ids.append(test_config.id()) - failed = failed + 1 - if verbosity > 0: - traceback.print_exc() - print(f"\n{str(ex)}\nFAILED") - else: - print("F", end="") + except Exception: + local_failed = True + local_error = traceback.format_exc() finally: # DeepEP managers are not reliably reusable across many subtests in # a single worker process. Tear them down after each DeepEP case so @@ -1771,6 +1788,40 @@ def _parallel_worker( total = total + 1 torch.distributed.barrier() + any_failed_tensor = torch.tensor( + [int(local_failed)], device=pgi.device, dtype=torch.int32 + ) + torch.distributed.all_reduce( + any_failed_tensor, op=torch.distributed.ReduceOp.MAX + ) + any_failed = bool(any_failed_tensor.item()) + + if any_failed: + failed = failed + 1 + + gathered_errors = [None] * pgi.world_size + torch.distributed.all_gather_object( + gathered_errors, local_error, group=cpu_group + ) + first_error = next( + (error for error in gathered_errors if error is not None), + "unknown distributed failure", + ) + assert first_error is not None + failure_details.append(f"{test_config.id()}\n{first_error.rstrip()}") + + if verbosity > 0 and is_logging_rank: + print(" FAILED") + print(first_error.rstrip()) + elif is_logging_rank: + print("F", end="") + else: + passed = passed + 1 + if verbosity > 0 and is_logging_rank: + print(" PASSED") + elif is_logging_rank: + print(".", end="") + skipped = total - (passed + failed) fails = f"{failed} failed" if failed > 0 else "" @@ -1780,17 +1831,24 @@ def _parallel_worker( passes = f"{sep}{passed} passed" if passed > 0 else "" report = ( - f"============= {fails}{skips}{passes} of {total} total tests =============" + f"============= {fails}{skips}{passes} of {total} total subcases =============" ) - sep = "\n" if verbosity == 0 else "" - print(f"{sep}{report}") + if is_logging_rank: + sep = "\n" if verbosity == 0 else "" + print(f"{sep}{report}") if failed > 0: - fail_ids_str = "\n".join(fail_ids) - raise RuntimeError( - f"\n============= Failed subtests =============\n{fail_ids_str}\n{report}" + failure_details_str = "\n\n".join(failure_details) + failure_report = ( + f"\n============= Failed subcases =============\n" + f"{failure_details_str}\n{report}" ) + if is_logging_rank and failure_report_path is not None: + with open(failure_report_path, "w", encoding="utf-8") as report_file: + report_file.write(failure_report) + if is_logging_rank: + raise RuntimeError(failure_report) # TODO: add cudagraphs/torch.compile tests @@ -1836,6 +1894,17 @@ def test_moe_layer( if os.environ.get("VLLM_LOGGING_LEVEL") is None: monkeypatch.setenv("VLLM_LOGGING_LEVEL", "ERROR") + if ( + backend in MORI_BACKENDS + and os.environ.get("VLLM_TEST_ENABLE_MORI_MOE_LAYER") == "1" + ): + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") + monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "1") + monkeypatch.setenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", "0") + from vllm._aiter_ops import rocm_aiter_ops + + rocm_aiter_ops.refresh_env_variables() + # TODO: cover FlashInfer MoE backends via moe_backend, e.g. # moe_backend=flashinfer_trtllm / flashinfer_cutlass / flashinfer_cutedsl # (BF16, FP8 and NVFP4 paths), and VLLM_USE_FLASHINFER_MOE_INT4=1. @@ -1881,6 +1950,11 @@ def test_moe_layer( if len(test_configs) == 0: pytest.skip("No supported configs found for this testpoint.") + with tempfile.NamedTemporaryFile( + prefix="moe-layer-failures-", delete=False + ) as failure_report_file: + failure_report_path = failure_report_file.name + try: parallel_launch_with_config( world_size, @@ -1889,7 +1963,13 @@ def test_moe_layer( None, test_configs, verbosity, + failure_report_path=failure_report_path, ) + if os.path.getsize(failure_report_path) > 0: + with open(failure_report_path, encoding="utf-8") as report_file: + pytest.fail(report_file.read()) finally: + with suppress(FileNotFoundError): + os.remove(failure_report_path) torch.accelerator.synchronize() # TODO: Is this needed? torch.accelerator.empty_cache() diff --git a/tests/kernels/moe/utils.py b/tests/kernels/moe/utils.py index 5fdcb8682f57..ce6e07f122ff 100644 --- a/tests/kernels/moe/utils.py +++ b/tests/kernels/moe/utils.py @@ -33,6 +33,10 @@ ) from vllm.model_executor.layers.fused_moe.router.fused_topk_router import fused_topk from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input +from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( + ref_nvfp4_quant, +) +from vllm.platforms import current_platform from vllm.utils.deep_gemm import per_block_cast_to_fp8 from vllm.utils.math_utils import round_up @@ -294,13 +298,34 @@ def moe_quantize_weights_2d( assert not per_token_quant w_amax = torch.abs(w).max().to(torch.float32) w_gs = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w_amax - w, w_s = ops.scaled_fp4_quant(w, w_gs) + if current_platform.is_rocm(): + w, w_s = _scaled_fp4_quant_emulated(w, w_gs) + else: + w, w_s = ops.scaled_fp4_quant(w, w_gs) else: raise RuntimeError(f"Unsupported quant type {quant_dtype}") return w, w_s, w_gs +def _pack_e2m1_fp4(fp4_values: torch.Tensor) -> torch.Tensor: + assert fp4_values.shape[-1] % 2 == 0 + + abs_values = fp4_values.abs() + codes = torch.empty_like(abs_values, dtype=torch.uint8) + for code, value in enumerate((0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)): + codes[abs_values == value] = code + codes = codes | ((fp4_values < 0).to(torch.uint8) << 3) + return codes[..., 0::2] | (codes[..., 1::2] << 4) + + +def _scaled_fp4_quant_emulated( + w: torch.Tensor, w_gs: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + fp4_values, w_s = ref_nvfp4_quant(w, w_gs, block_size=16) + return _pack_e2m1_fp4(fp4_values), w_s.to(torch.float8_e4m3fn) + + def moe_quantize_weights( w: torch.Tensor, w_s: torch.Tensor | None, diff --git a/tests/quantization/test_fp8.py b/tests/quantization/test_fp8.py index 5ad3b666c2ea..1e96b630faab 100644 --- a/tests/quantization/test_fp8.py +++ b/tests/quantization/test_fp8.py @@ -31,6 +31,9 @@ from vllm.model_executor.layers.quantization.online.fp8 import ( Fp8PerTensorOnlineLinearMethod, ) +from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + process_fp8_input_tensor_strategy_moe, +) from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.platforms import current_platform @@ -47,6 +50,16 @@ ] +def test_static_fp8_moe_input_scales_remain_scalar() -> None: + a1_scale, a2_scale = process_fp8_input_tensor_strategy_moe( + torch.tensor([0.25, 0.5]), + torch.tensor([0.75, 0.6]), + enable_eplb=False, + ) + + assert a1_scale.ndim == a2_scale.ndim == 0 + + @pytest.mark.skipif( not is_quant_method_supported("fp8"), reason="FP8 is not supported on this GPU type.", diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py index be4930052a9a..dc619b2b12d2 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe.py @@ -857,6 +857,12 @@ def invoke_fused_moe_triton_kernel( BLOCK_SIZE_K, ) use_td = False + + # Triton treats 0-D tensor arguments as scalar values, but the kernel + # loads tensor-wise activation scales through a pointer. + if A_scale is not None and A_scale.ndim == 0: + A_scale = A_scale.reshape(1) + fused_moe_kernel[grid]( A, B, diff --git a/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py b/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py index abb2043d7c45..cb76bbcd5c51 100644 --- a/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py +++ b/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py @@ -313,7 +313,7 @@ def break_fp4_bytes(a, dtype): signs = (combined & 0x08).to(torch.bool) # Sign bits abs_vals = (combined & 0x07).to(torch.long) - kE2M1 = kE2M1ToFloat_handle.val + kE2M1 = kE2M1ToFloat_handle.val.to(device=a.device) # Device-aware lookup and sign application values = kE2M1[abs_vals] * torch.where(signs, -1.0, 1.0) # Reshape to final form From 17dbd429306d2fe9cc6be04e57aab314bd697f5e Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Wed, 19 Aug 2026 10:20:05 -0500 Subject: [PATCH 153/839] [ROCm] Add UE8M0 scale packing for Triton silu_mul_quant (#37835) Signed-off-by: Andreas Karatzas --- .../moe/test_silu_mul_fp8_quant_deep_gemm.py | 6 +-- .../experts/batched_deep_gemm_moe.py | 51 ++++++++++++++++--- 2 files changed, 44 insertions(+), 13 deletions(-) diff --git a/tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py b/tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py index e9a246f51b25..85b992ad6544 100644 --- a/tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py +++ b/tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py @@ -239,12 +239,8 @@ def test_silu_mul_fp8_quant_deep_gemm(E: int, T: int, H: int, fp8_type: torch.dt scale_fmts = [ DeepGemmQuantScaleFMT.FLOAT32, DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0, + DeepGemmQuantScaleFMT.UE8M0, ] - # UE8M0 (int32 packed) scales require the C++ kernel which is - # not available on ROCm (#ifndef USE_ROCM). - # https://github.com/ROCm/aiter/issues/2420 - if current_platform.is_cuda(): - scale_fmts.append(DeepGemmQuantScaleFMT.UE8M0) # Run the SiLU V2 kernel for scale_fmt in scale_fmts: diff --git a/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py b/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py index af9d170fcbab..c578fb641b10 100644 --- a/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py @@ -220,6 +220,22 @@ def persistent_masked_m_silu_mul_quant( # Triton fallback for ROCm and XPU -- the C++ kernel is guarded by # #ifndef USE_ROCM in activation_kernels.cu. # https://github.com/ROCm/aiter/issues/2420 + # For UE8M0 (int32 packed scales), compute with float32 scales + # first, then pack afterwards. + is_packed_ue8m0 = quant_scale_fmt == DeepGemmQuantScaleFMT.UE8M0 + if is_packed_ue8m0: + f32_shape, f32_strides, _ = scales_shape_stride_dtype( + E, T, G, DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0 + ) + y_s_f32 = torch.empty_strided( + f32_shape, + f32_strides, + dtype=torch.float32, + device=y.device, + ) + else: + y_s_f32 = y_s + stride_cnt_e = tokens_per_expert.stride()[0] # Static grid over experts and H-groups. @@ -231,14 +247,16 @@ def persistent_masked_m_silu_mul_quant( fp8_min, fp8_max = get_fp8_min_max() eps: float = 1e-10 - assert y_s.dtype == torch.float32, ( - "_silu_mul_fp8_quant_deep_gemm Triton fallback does not " - f"support {y_s.dtype} scales. Only torch.float32 supported." - ) + if not is_packed_ue8m0: + assert y_s.dtype == torch.float32, ( + "_silu_mul_fp8_quant_deep_gemm Triton fallback does not " + f"support {y_s.dtype} scales. Only torch.float32 supported." + ) + f32_strides = y_s_f32.stride() _silu_mul_fp8_quant_deep_gemm[grid]( y, y_q, - y_s, + y_s_f32, tokens_per_expert, H, group_size, @@ -248,9 +266,9 @@ def persistent_masked_m_silu_mul_quant( stride_yq_e, stride_yq_t, stride_yq_h, - ys_strides[0], - ys_strides[1], - ys_strides[2], + f32_strides[0], + f32_strides[1], + f32_strides[2], stride_cnt_e, eps, fp8_min, @@ -261,6 +279,23 @@ def persistent_masked_m_silu_mul_quant( num_warps=1, ) + if is_packed_ue8m0: + # Pack float32 scales into int32 UE8M0 format: + # extract exponent bits (bits 30:23) from float32. + E_dim, T_dim, G_dim = y_s_f32.shape + y_s_cont = y_s_f32.contiguous() + i32_pad = round_up(G_dim, 4) - G_dim + y_s_u8 = (y_s_cont.view(torch.int32) >> 23).to(torch.uint8) + if i32_pad > 0: + y_s_u8 = torch.nn.functional.pad(y_s_u8, (0, i32_pad)) + # y_s has shape (E, T, G//4) with stride (T*G//4, 1, T) + packed = y_s_u8.view(torch.int32) + # Copy with matching strides + for e_idx in range(E_dim): + nt = tokens_per_expert[e_idx].item() + if nt > 0: + y_s[e_idx, :nt].copy_(packed[e_idx, :nt]) + return y_q, y_s From 525b7bbb3a74d811cedcee108425d9bb5c215fdd Mon Sep 17 00:00:00 2001 From: Dinesh Chitlangia Date: Wed, 19 Aug 2026 11:50:59 -0400 Subject: [PATCH 154/839] [Bugfix][CPU] Enable C++ causal_conv1d GDN path and float32 SSM cache on non-AMX AVX-512BF16 CPUs (#49688) Signed-off-by: Dinesh Chitlangia --- tests/kernels/mamba/cpu/test_cpu_gdn_ops.py | 274 +++++++++++++++++- .../layers/mamba/ops/cpu/gdn_attention.py | 14 +- vllm/model_executor/layers/utils.py | 7 +- vllm/platforms/cpu.py | 12 +- 4 files changed, 290 insertions(+), 17 deletions(-) diff --git a/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py b/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py index f43a9b4a58c4..385870e49c20 100644 --- a/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py +++ b/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py @@ -878,17 +878,18 @@ def test_causal_conv1d_update_cpu_rejects_invalid_accepted_count( @pytest.mark.skipif( - not torch.cpu._is_amx_tile_supported(), - reason="requires AMX support", + not torch.cpu._is_avx512_bf16_supported(), + reason="causal_conv1d_fwd_cpu requires AVX-512BF16 (Intel Xeon or AMD EPYC)", ) @pytest.mark.parametrize("total_tokens, split", TWO_CALL_SPLITS) @torch.inference_mode() def test_causal_conv1d_fwd_cpu_two_call_split(total_tokens: int, split: int) -> None: - """AMX prefill conv op must honor ``has_initial_state`` so a two-call split + """C++ prefill conv op must honor ``has_initial_state`` so a two-call split matches the single-call result. Regression test for ``causal_conv1d_fwd_varlen_kernel_impl`` (``conv.cpp``) - ignoring the carried conv state on continued chunks. + ignoring the carried conv state on continued chunks. Runs on any + AVX-512BF16 CPU since conv.cpp uses VDPBF16PS, not AMX tiles. """ state_len = CONV_KERNEL - 1 x, weight, bias = _conv_inputs(total_tokens) @@ -996,3 +997,268 @@ def test_batch_memcpy_cpu_fallback() -> None: for src, dst in zip(srcs, dsts): torch.testing.assert_close(dst, src) + + +# --------------------------------------------------------------------------- +# C++ conv (conv.cpp) uses VDPBF16PS, not AMX tiles, so it runs on any +# AVX-512BF16 CPU; weight is VNNI-packed on this same predicate at load time. +# --------------------------------------------------------------------------- + +_HAS_AVX512_BF16 = torch.cpu._is_avx512_bf16_supported() + +_STATE_LEN = CONV_KERNEL - 1 + +CONV_EQUIV_SEQ_LENS = [[1], [7], [64], [65], [1, 2, 3], [63, 64, 65], [128, 129]] + + +def _conv_fp32_oracle(x, weight, bias, seq_lens, activation="silu"): + """High-precision conv reference: everything in float32, no bf16 rounding. + x: [total_tokens, dim] (no initial state). Returns [total_tokens, dim].""" + xf = x.float() + wf = weight.float().unsqueeze(1) + bf = bias.float() + out = torch.empty_like(xf) + start = 0 + for n in seq_lens: + seg = xf[start : start + n].transpose(0, 1).unsqueeze(0) # [1, dim, n] + conv_in = F.pad(seg, (_STATE_LEN, 0)) + seg_out = F.conv1d(conv_in, wf, bf, padding=0, groups=CONV_DIM)[..., -n:] + if activation in ("silu", "swish"): + seg_out = F.silu(seg_out) + out[start : start + n] = seg_out.squeeze(0).transpose(0, 1) + start += n + return out + + +def _run_prefill_torch(x, weight, bias, seq_lens): + from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( + causal_conv1d_fn_cpu as causal_conv1d_torch, + ) + + num_seqs = len(seq_lens) + conv_states = torch.zeros(num_seqs, CONV_DIM, _STATE_LEN, dtype=x.dtype) + qsl = torch.tensor( + [0, *torch.tensor(seq_lens).cumsum(0).tolist()], dtype=torch.int32 + ) + out = causal_conv1d_torch( + x=x.transpose(0, 1).contiguous(), + weight=weight, + bias=bias, + conv_states=conv_states, + query_start_loc=qsl, + cache_indices=torch.arange(num_seqs, dtype=torch.int32), + has_initial_state=torch.zeros(num_seqs, dtype=torch.bool), + activation="silu", + ) + return out.transpose(0, 1).contiguous(), conv_states + + +def _run_prefill_cpp(x, weight, bias, seq_lens, is_vnni=False): + num_seqs = len(seq_lens) + packed_w = ops.causal_conv1d_weight_pack(weight) if is_vnni else weight + if is_vnni: + # C++-branch layout: kv-cache "SD" [slots, state_len, dim] transposed to + # [slots, dim, state_len] (a non-contiguous view). + conv_state = torch.zeros( + num_seqs, _STATE_LEN, CONV_DIM, dtype=x.dtype + ).transpose(1, 2) + else: + conv_state = torch.zeros(num_seqs, CONV_DIM, _STATE_LEN, dtype=x.dtype) + qsl = torch.tensor( + [0, *torch.tensor(seq_lens).cumsum(0).tolist()], dtype=torch.int32 + ) + out = ops.causal_conv1d_fwd_cpu( + x=x.transpose(0, 1), + weight=packed_w, + bias=bias, + conv_states=conv_state, + query_start_loc=qsl, + cache_indices=torch.arange(num_seqs, dtype=torch.int32), + has_initial_state=torch.zeros(num_seqs, dtype=torch.bool), + silu_activation=True, + is_vnni=is_vnni, + ) + return out.transpose(0, 1).contiguous(), conv_state + + +@pytest.mark.skipif( + not _HAS_AVX512_BF16, reason="C++ causal_conv1d requires AVX-512BF16" +) +@pytest.mark.parametrize("seq_lens", CONV_EQUIV_SEQ_LENS) +@torch.inference_mode() +def test_conv_cpp_matches_torch(seq_lens): + """C++ causal_conv1d_fwd_cpu matches the torch fallback within bf16 tol.""" + x, weight, bias = _conv_inputs(sum(seq_lens)) + out_torch, state_torch = _run_prefill_torch(x, weight, bias, seq_lens) + out_cpp, state_cpp = _run_prefill_cpp(x, weight, bias, seq_lens) + torch.testing.assert_close(out_cpp, out_torch, atol=1e-2, rtol=1e-2) + torch.testing.assert_close(state_cpp, state_torch, atol=1e-2, rtol=1e-2) + + +@pytest.mark.skipif( + not _HAS_AVX512_BF16, reason="C++ causal_conv1d requires AVX-512BF16" +) +@pytest.mark.parametrize("seq_lens", CONV_EQUIV_SEQ_LENS) +@torch.inference_mode() +def test_conv_cpp_no_worse_than_torch_vs_fp32(seq_lens): + """Swapping torch -> C++ conv must not increase error vs an fp32 oracle.""" + x, weight, bias = _conv_inputs(sum(seq_lens)) + oracle = _conv_fp32_oracle(x, weight, bias, seq_lens) + out_torch, _ = _run_prefill_torch(x, weight, bias, seq_lens) + out_cpp, _ = _run_prefill_cpp(x, weight, bias, seq_lens) + err_torch = (out_torch.float() - oracle).abs().mean().item() + err_cpp = (out_cpp.float() - oracle).abs().mean().item() + assert err_cpp <= err_torch + 1e-3, ( + f"C++ conv less accurate than torch: " + f"err_cpp={err_cpp:.2e} err_torch={err_torch:.2e}" + ) + + +@pytest.mark.skipif( + not _HAS_AVX512_BF16, reason="C++ causal_conv1d requires AVX-512BF16" +) +@pytest.mark.parametrize("seq_lens", CONV_EQUIV_SEQ_LENS) +@torch.inference_mode() +def test_conv_cpp_vnni_packed_matches_torch(seq_lens): + """The exact runtime prefill sequence (VNNI-packed weight + SD-layout + conv_state view + is_vnni=True) must match the torch fallback. Validates + the packing + layout handoff on any AVX-512BF16 CPU.""" + x, weight, bias = _conv_inputs(sum(seq_lens)) + out_torch, _ = _run_prefill_torch(x, weight, bias, seq_lens) + out_vnni, _ = _run_prefill_cpp(x, weight, bias, seq_lens, is_vnni=True) + torch.testing.assert_close(out_vnni, out_torch, atol=1e-2, rtol=1e-2) + + +@pytest.mark.skipif( + not _HAS_AVX512_BF16, reason="C++ causal_conv1d requires AVX-512BF16" +) +@pytest.mark.parametrize("batch", DECODE_BATCH_SIZES) +@torch.inference_mode() +def test_conv_update_cpp_matches_torch(batch): + """Decode conv: causal_conv1d_update_cpu matches causal_conv1d_update_torch, + including the in-place conv_state update (the next-step handoff).""" + from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( + causal_conv1d_update_torch, + ) + + x = tensor_cache(batch * CONV_DIM, torch.bfloat16).view(batch, CONV_DIM) + weight = tensor_cache(CONV_DIM * CONV_KERNEL, torch.bfloat16).view( + CONV_DIM, CONV_KERNEL + ) + bias = tensor_cache(CONV_DIM, torch.bfloat16) + conv_state = tensor_cache(batch * CONV_DIM * _STATE_LEN, torch.bfloat16).view( + batch, CONV_DIM, _STATE_LEN + ) + + cs_torch = conv_state.clone() + out_torch = causal_conv1d_update_torch( + x=x.unsqueeze(-1), + conv_state=cs_torch, + weight=weight, + bias=bias, + activation="silu", + ).squeeze(-1) + + cs_cpp = conv_state.clone() + out_cpp = ops.causal_conv1d_update_cpu( + x=x.contiguous(), + conv_states=cs_cpp, + weight=weight, + bias=bias, + silu_activation=True, + conv_state_indices=torch.arange(batch, dtype=torch.int32), + is_vnni=False, + ) + torch.testing.assert_close(out_cpp, out_torch, atol=1e-2, rtol=1e-2) + torch.testing.assert_close(cs_cpp, cs_torch, atol=1e-2, rtol=1e-2) + + +@pytest.mark.skipif( + not _HAS_AVX512_BF16, reason="C++ causal_conv1d requires AVX-512BF16" +) +@torch.inference_mode() +def test_conv_weight_pack_roundtrip_unpacked_matches(): + """`causal_conv1d_weight_pack` repacks the (dim, width) weight, so the C++ conv op + with packed weight + is_vnni=True must equal unpacked weight + is_vnni=False. + Guards the spec-decode contract: the runtime VNNI-packs `layer.conv1d.weight` + in place and stashes the original as `_cpu_unpacked_conv_weight` for the torch + spec-decode path, so both must produce identical math. + """ + seq_lens = [7, 64, 65] + x, weight, bias = _conv_inputs(sum(seq_lens)) + out_unpacked, _ = _run_prefill_cpp(x, weight, bias, seq_lens, is_vnni=False) + out_packed, _ = _run_prefill_cpp(x, weight, bias, seq_lens, is_vnni=True) + torch.testing.assert_close(out_packed, out_unpacked, atol=1e-2, rtol=1e-2) + + +@pytest.mark.skipif( + not _HAS_AVX512_BF16, reason="C++ causal_conv1d requires AVX-512BF16" +) +@torch.inference_mode() +def test_spec_decode_unpacked_conv_weight_stash(): + """Spec-decode correctness: AVX-512BF16 CPUs VNNI-pack ``conv1d.weight`` in place + at load time and stash the original ``(dim, width)`` tensor as + ``_cpu_unpacked_conv_weight``; the spec-decode path must use that stash since + reading the packed weight directly produces garbage. Verifies the recovered + weight equals the original and that torch F.conv1d agrees with the C++ conv + fed the packed weight. + """ + from vllm.model_executor.layers.utils import dispatch_cpu_unquantized_gemm + + torch.manual_seed(0) + # conv1d weight is stored [dim, 1, width]; bias [dim]. + orig_2d = torch.rand(CONV_DIM, CONV_KERNEL, dtype=torch.bfloat16) + bias = torch.rand(CONV_DIM, dtype=torch.bfloat16) + + conv = torch.nn.Module() + conv.weight = torch.nn.Parameter( + orig_2d.view(CONV_DIM, 1, CONV_KERNEL).clone(), requires_grad=False + ) + + # Load-time dispatch: packs weight in place + stashes the unpacked copy. + dispatch_cpu_unquantized_gemm(conv, remove_weight=False) + + # 1. The stash must exist and equal the original (dim, width) weight. + assert hasattr(conv, "_cpu_unpacked_conv_weight"), ( + "dispatch_cpu_unquantized_gemm did not stash _cpu_unpacked_conv_weight " + "on an AVX-512BF16 CPU" + ) + stash = conv._cpu_unpacked_conv_weight + torch.testing.assert_close(stash, orig_2d, atol=0, rtol=0) + + # 2. Mirror _unpacked_conv_weight()'s lookup: stash wins over conv.weight. + recovered = getattr(conv, "_cpu_unpacked_conv_weight", None) + assert recovered is not None + # conv.weight is now packed; using it directly (the bug) would differ. + packed_weight = conv.weight # [dim, 1, width], VNNI-packed contents + + # 3. Spec-path torch conv with the recovered unpacked weight must match the + # nonspec C++ conv with the packed weight (is_vnni=True). + seq_lens = [7, 65] + total = sum(seq_lens) + x = tensor_cache(total * CONV_DIM, torch.bfloat16).view(total, CONV_DIM) + + out_torch, _ = _run_prefill_torch(x, recovered, bias, seq_lens) + + num_seqs = len(seq_lens) + cs = torch.zeros(num_seqs, _STATE_LEN, CONV_DIM, dtype=x.dtype).transpose(1, 2) + qsl = torch.tensor( + [0, *torch.tensor(seq_lens).cumsum(0).tolist()], dtype=torch.int32 + ) + out_cpp = ( + ops.causal_conv1d_fwd_cpu( + x=x.transpose(0, 1), + weight=packed_weight.view(CONV_DIM, CONV_KERNEL), + bias=bias, + conv_states=cs, + query_start_loc=qsl, + cache_indices=torch.arange(num_seqs, dtype=torch.int32), + has_initial_state=torch.zeros(num_seqs, dtype=torch.bool), + silu_activation=True, + is_vnni=True, + ) + .transpose(0, 1) + .contiguous() + ) + + torch.testing.assert_close(out_cpp, out_torch, atol=1e-2, rtol=1e-2) diff --git a/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py b/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py index e84a0c6b723e..3f28babb4427 100644 --- a/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py +++ b/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py @@ -103,13 +103,15 @@ def _cpu_gdn_attention_nonspec( assert state_indices_tensor is not None assert query_start_loc is not None - is_amx = torch.cpu._is_amx_tile_supported() + # C++ conv (conv.cpp) uses VDPBF16PS, not AMX tiles, so it runs on any + # AVX-512BF16 CPU; weight is VNNI-packed on this same predicate at load time. + use_cpp_conv = torch.cpu._is_avx512_bf16_supported() conv_state = layer.kv_cache[0] - if is_amx: - # AMX causal conv requires [num_allocated_slots, kernel - 1, conv_dim]. + if use_cpp_conv: + # C++ conv requires [num_allocated_slots, kernel - 1, conv_dim] (SD). if is_conv_state_dim_first(): - raise RuntimeError("AMX GDN attention requires `SD` conv_state layout.") + raise RuntimeError("C++ CPU GDN attention requires `SD` conv_state layout.") conv_state = conv_state.transpose(1, 2) else: if not is_conv_state_dim_first(): @@ -138,7 +140,7 @@ def _cpu_gdn_attention_nonspec( decode_b = b[:num_decode_tokens] decode_a = a[:num_decode_tokens] decode_state_indices = state_indices_tensor[:num_decodes] - if is_amx: + if use_cpp_conv: decode_mixed_qkv = ops.causal_conv1d_update_cpu( x=decode_mixed_qkv, conv_states=conv_state, @@ -207,7 +209,7 @@ def _cpu_gdn_attention_nonspec( num_decodes : num_decodes + num_prefills ] - if is_amx: + if use_cpp_conv: prefill_mixed_qkv = ops.causal_conv1d_fwd_cpu( x=prefill_mixed_qkv.transpose(0, 1), weight=layer.conv1d.weight, diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py index 16539bdf6559..5cb23fb14b58 100644 --- a/vllm/model_executor/layers/utils.py +++ b/vllm/model_executor/layers/utils.py @@ -280,7 +280,10 @@ def dispatch_cpu_unquantized_gemm( if layer.weight.ndim != 2: # this is not a linear layer # For now it should be a causal_conv1d op or MoE 3D expert weights - if torch.cpu._is_amx_tile_supported() and hasattr( + # The C++ causal_conv1d kernels use VDPBF16PS (no AMX tiles), so the + # VNNI weight prepack applies to any AVX-512BF16 CPU, not just AMX + # (e.g. AMD Zen5/Turin). + if torch.cpu._is_avx512_bf16_supported() and hasattr( ops, "causal_conv1d_weight_pack" ): # prepack conv weight @@ -293,7 +296,7 @@ def dispatch_cpu_unquantized_gemm( .clone() ) # Stash the un-packed (dim, width) weight so the speculative-decode - # GDN path (which uses torch conv, not the AMX kernel) can use it. + # GDN path (which uses torch conv, not the C++ kernel) can use it. layer._cpu_unpacked_conv_weight = unpacked layer.weight.data = ops.causal_conv1d_weight_pack(unpacked) return diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index 02ca99ff33a9..e002c80ad899 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -181,13 +181,13 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None: "otherwise the performance is not optimized." ) - # AMX GDN requires float32 state + # Accelerated GDN (AMX tiles or AVX-512BF16 VDPBF16PS) requires float32 SSM state. if ( - torch.cpu._is_amx_tile_supported() + torch.cpu._is_avx512_bf16_supported() and cache_config.mamba_ssm_cache_dtype != "float32" ): cache_config.mamba_ssm_cache_dtype = "float32" - logger.warning("Reset SSM cache type to float32 for AMX mamba attention.") + logger.warning("Reset SSM cache type to float32 for accelerated GDN mamba attention.") # Lagecy setting env_key = "VLLM_CPU_KVCACHE_SPACE" @@ -290,8 +290,10 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None: # Avoid inductor generates num_thread() and breaks the thread binding os.environ["TORCHINDUCTOR_CPP_DYNAMIC_THREADS"] = "1" - # For efficient conv state memory access - if torch.cpu._is_amx_tile_supported(): + # For efficient conv state memory access. The C++ causal_conv1d + # kernels (VDPBF16PS, no AMX tiles) consume the SD layout on any + # AVX-512BF16 CPU, so apply it beyond AMX (e.g. AMD Zen5/Turin). + if torch.cpu._is_avx512_bf16_supported(): os.environ["VLLM_SSM_CONV_STATE_LAYOUT"] = "SD" ld_preload_str = os.getenv("LD_PRELOAD", "") From 583a00257d4c5d1a54063d956057df1df6822b06 Mon Sep 17 00:00:00 2001 From: akii96 Date: Wed, 19 Aug 2026 19:14:52 +0300 Subject: [PATCH 155/839] [ROCm] [Bugfix] Fix Triton fused shared expert alignment (#51632) Signed-off-by: Aakif Nawaz --- tests/kernels/moe/test_moe.py | 44 +++++++++++++++++++ .../layers/fused_moe/experts/triton_moe.py | 8 +++- 2 files changed, 50 insertions(+), 2 deletions(-) diff --git a/tests/kernels/moe/test_moe.py b/tests/kernels/moe/test_moe.py index 17b8cae318bd..ac73ca941621 100644 --- a/tests/kernels/moe/test_moe.py +++ b/tests/kernels/moe/test_moe.py @@ -459,6 +459,50 @@ def m_fused_moe( ) +def test_fused_shared_expert_alignment(workspace_init): + set_random_seed(7) + m, n, k = 4, 64, 128 + routed_experts = 8 + physical_experts = routed_experts + 1 + dtype = torch.bfloat16 + + a = torch.randn((m, k), device=DEVICE_TYPE, dtype=dtype) / 10 + w1 = torch.randn((physical_experts, 2 * n, k), device=DEVICE_TYPE, dtype=dtype) / 10 + w2 = torch.randn((physical_experts, k, n), device=DEVICE_TYPE, dtype=dtype) / 10 + topk_ids = torch.tensor( + [[0, 8], [1, 8], [2, 8], [3, 8]], device=DEVICE_TYPE, dtype=torch.int32 + ) + topk_weights = torch.tensor( + [[0.5, 1.0]] * m, device=DEVICE_TYPE, dtype=torch.float32 + ) + + moe_config = make_dummy_moe_config( + num_experts=physical_experts, + experts_per_token=2, + hidden_dim=k, + intermediate_size=n, + in_dtype=dtype, + max_num_tokens=m, + ) + modular_moe = modular_triton_fused_moe(moe_config, FUSED_MOE_UNQUANTIZED_CONFIG) + + with set_current_vllm_config(vllm_config): + expected = torch_experts(a, w1, w2, topk_weights, topk_ids) + actual = modular_moe.apply( + hidden_states=a, + w1=w1, + w2=w2, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=MoEActivation.SILU, + global_num_experts=routed_experts, + expert_map=None, + apply_router_weight_on_input=False, + ) + + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=0) + + def test_fused_moe_int64_overflow(workspace_init): """Regression test for int32 overflow in stride*offset products. diff --git a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py index d0ece26cfa01..9887423e7974 100644 --- a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py @@ -315,13 +315,15 @@ def apply( ) intermediate_cache3 = _resize_cache(workspace2, (num_tokens, top_k_num, K)) + # Include fused shared-expert rows while preserving EP remapping. + num_align_experts = w1.shape[0] if expert_map is None else global_num_experts sorted_token_ids, expert_ids, num_tokens_post_padded = ( _prepare_expert_assignment( topk_ids, config, num_tokens, top_k_num, - global_num_experts, + num_align_experts, expert_map, use_int8_w8a16=self.quant_config.use_int8_w8a16, use_int4_w4a16=self.quant_config.use_int4_w4a16, @@ -688,8 +690,10 @@ def apply( ) intermediate_cache3 = _resize_cache(workspace2, (num_tokens, top_k_num, K)) + # Include fused shared-expert rows while preserving EP remapping. + num_align_experts = w1.shape[0] if expert_map is None else global_num_experts sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( - topk_ids, config["BLOCK_SIZE_M"], global_num_experts, expert_map + topk_ids, config["BLOCK_SIZE_M"], num_align_experts, expert_map ) invoke_fused_moe_wna16_triton_kernel( From c676232313fce3d607da7c66e41fe8d4739c194d Mon Sep 17 00:00:00 2001 From: xuebwang-amd Date: Thu, 20 Aug 2026 00:29:29 +0800 Subject: [PATCH 156/839] [Bugfix][Quantization] Fix OCP MX MoE emulation silently skipping mxfp6 activation QDQ (#52704) Signed-off-by: xuebwang-amd --- tests/kernels/moe/test_ocp_mx_moe.py | 148 ++++++++++++++++++ .../fused_moe/experts/ocp_mx_emulation_moe.py | 70 ++++++--- 2 files changed, 195 insertions(+), 23 deletions(-) diff --git a/tests/kernels/moe/test_ocp_mx_moe.py b/tests/kernels/moe/test_ocp_mx_moe.py index 2d3e48d435f3..00d3cf29c8e9 100644 --- a/tests/kernels/moe/test_ocp_mx_moe.py +++ b/tests/kernels/moe/test_ocp_mx_moe.py @@ -13,6 +13,11 @@ is_aiter_found_and_supported, rocm_aiter_ops, ) +from vllm.model_executor.layers.fused_moe.experts.ocp_mx_emulation_moe import ( + activation_quant_dtype, +) +from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input +from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import OCP_MX_Scheme from vllm.platforms import current_platform from vllm.utils.flashinfer import has_flashinfer @@ -1859,3 +1864,146 @@ def is_supported_config( with pytest.raises(NotImplementedError, match="Unsupported reasons"): mxfp4_oracle.select_mxfp4_moe_backend(moe_config) + + +# Every activation-quantizing OCP MX scheme must map to a `quant_dtype` that +# `moe_kernel_quantize_input` actually dispatches on. Its final `else` returns +# the activation untouched, so a name it does not know (e.g. "mxfp6" instead of +# "mxfp6_e3m2") silently skips the fake-quantization the emulation exists for. +@pytest.mark.skipif(not ROCM_AVAILABLE, reason="emulation backend targets ROCm") +@pytest.mark.parametrize("ocp_mx_scheme", list(OCP_MX_Scheme)) +def test_emulation_activation_quant_dtype_is_dispatchable(ocp_mx_scheme): + quant_dtype = activation_quant_dtype(ocp_mx_scheme) + + if "_a_" not in ocp_mx_scheme.value: + assert quant_dtype is None, "weight-only schemes must not quantize activations" + return + + a = torch.randn(64, 128, dtype=torch.bfloat16, device="cuda") + a_scale = torch.ones(1, dtype=torch.float32, device="cuda") + out, _ = moe_kernel_quantize_input( + a, a_scale, quant_dtype, False, None, quantization_emulation=True + ) + assert not torch.equal(out, a), ( + f"{ocp_mx_scheme.value} -> quant_dtype={quant_dtype!r} left the activation" + " unquantized; moe_kernel_quantize_input does not dispatch on it" + ) + + +@pytest.mark.skipif(not ROCM_AVAILABLE, reason="emulation backend targets ROCm") +@torch.inference_mode() +def test_emulation_a_mxfp6_moe_forward_quantizes_activations(): + """The same property observed through a full MoE forward. + + `w_mxfp4_a_mxfp6_e3m2` and the weight-only `w_mxfp4` differ only in whether + activations are fake-quantized, so with identical weights, activations and + routing their layer outputs must differ. When the emulation selects a + `quant_dtype` `moe_kernel_quantize_input` does not dispatch on, the QDQ is + skipped and the two outputs come out bit-identical. + """ + from vllm.config import VllmConfig, set_current_vllm_config + from vllm.model_executor.layers.fused_moe import FusedMoEConfig + from vllm.model_executor.layers.fused_moe.activation import MoEActivation + from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEParallelConfig, + RoutingMethodType, + mxfp4_w4a16_moe_quant_config, + ocp_mx_moe_quant_config, + ) + from vllm.model_executor.layers.fused_moe.experts.ocp_mx_emulation_moe import ( + OCP_MXQuantizationEmulationTritonExperts, + ) + from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( + Mxfp4MoeBackend, + make_mxfp4_moe_kernel, + ) + from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( + mxfp4_quantize, + ) + from vllm.v1.worker.workspace import init_workspace_manager + + init_workspace_manager(torch.accelerator.current_device_index()) + + num_experts, topk = 8, 2 + hidden_size, intermediate_size, num_tokens = 256, 256, 64 + dtype, device = torch.bfloat16, "cuda:0" + + torch.manual_seed(0) + w13, w13_scale = mxfp4_quantize( + torch.randn( + num_experts, 2 * intermediate_size, hidden_size, dtype=dtype, device=device + ) + / 8 + ) + w2, w2_scale = mxfp4_quantize( + torch.randn( + num_experts, hidden_size, intermediate_size, dtype=dtype, device=device + ) + / 8 + ) + w13, w2 = w13.contiguous(), w2.contiguous() + w13_scale, w2_scale = w13_scale.contiguous(), w2_scale.contiguous() + + torch.manual_seed(1) + hidden_states = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device) + topk_weights, topk_ids = torch.topk( + torch.randn(num_tokens, num_experts, dtype=torch.float32, device=device), + k=topk, + dim=-1, + ) + topk_weights = torch.softmax(topk_weights, dim=-1).to(dtype) + topk_ids = topk_ids.to(torch.int32) + + def run(quant_config): + moe_config = FusedMoEConfig( + num_experts=num_experts, + experts_per_token=topk, + hidden_dim=hidden_size, + intermediate_size=intermediate_size, + num_local_experts=num_experts, + num_logical_experts=num_experts, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation.SILU, + in_dtype=dtype, + device=device, + routing_method=RoutingMethodType.Renormalize, + ) + with set_current_vllm_config(VllmConfig()): + kernel = make_mxfp4_moe_kernel( + moe_quant_config=quant_config, + moe_config=moe_config, + mxfp4_backend=Mxfp4MoeBackend.EMULATION, + experts_cls=OCP_MXQuantizationEmulationTritonExperts, + routing_tables=None, + ) + return kernel.apply( + hidden_states=hidden_states, + w1=w13, + w2=w2, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=MoEActivation.SILU, + global_num_experts=num_experts, + expert_map=None, + apply_router_weight_on_input=False, + ) + + # A fresh quant config per run: the experts nulls the weight scales on it. + a_mxfp6 = ocp_mx_moe_quant_config( + quant_dtype="mxfp6_e3m2", + weight_dtype="mxfp4", + w1_scale=w13_scale, + w2_scale=w2_scale, + ) + assert a_mxfp6.ocp_mx_scheme == OCP_MX_Scheme.w_mxfp4_a_mxfp6_e3m2 + out_a_mxfp6 = run(a_mxfp6) + + weight_only = mxfp4_w4a16_moe_quant_config(w1_scale=w13_scale, w2_scale=w2_scale) + assert weight_only.ocp_mx_scheme == OCP_MX_Scheme.w_mxfp4 + out_weight_only = run(weight_only) + + max_diff = (out_a_mxfp6.float() - out_weight_only.float()).abs().max().item() + assert max_diff > 0.0, ( + "w_mxfp4_a_mxfp6_e3m2 output is bit-identical to weight-only w_mxfp4:" + " the emulation never fake-quantized the activations" + ) diff --git a/vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py b/vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py index 9107a52ed812..2c91e22d9108 100644 --- a/vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py @@ -35,6 +35,52 @@ logger = init_logger(__name__) +def activation_quant_dtype( + ocp_mx_scheme: OCP_MX_Scheme | str, +) -> torch.dtype | str | None: + """Activation dtype `moe_kernel_quantize_input` should fake-quantize to. + + Args: + ocp_mx_scheme: The OCP MX scheme the emulated experts run. Accepts the + enum member or its string value. + + Returns: + A `quant_dtype` `moe_kernel_quantize_input` dispatches on, or None for + weight-only schemes, which leave activations untouched. + + Raises: + NotImplementedError: If the scheme has no emulated activation dtype. + """ + if ocp_mx_scheme in { + OCP_MX_Scheme.w_mxfp4, + OCP_MX_Scheme.w_mxfp6_e3m2, + OCP_MX_Scheme.w_mxfp6_e2m3, + }: + return None + elif ocp_mx_scheme == OCP_MX_Scheme.w_mxfp4_a_mxfp4: + return "mxfp4" + elif ocp_mx_scheme in { + OCP_MX_Scheme.w_mxfp4_a_mxfp6_e3m2, + OCP_MX_Scheme.w_mxfp6_e3m2_a_mxfp6_e3m2, + }: + return "mxfp6_e3m2" + elif ocp_mx_scheme in { + OCP_MX_Scheme.w_mxfp4_a_mxfp6_e2m3, + OCP_MX_Scheme.w_mxfp6_e2m3_a_mxfp6_e2m3, + }: + return "mxfp6_e2m3" + elif ocp_mx_scheme in { + OCP_MX_Scheme.w_mxfp4_a_fp8, + OCP_MX_Scheme.w_mxfp6_e3m2_a_fp8, + OCP_MX_Scheme.w_mxfp6_e2m3_a_fp8, + }: + return current_platform.fp8_dtype() + raise NotImplementedError( + f"No emulated activation dtype for OCP MX scheme {ocp_mx_scheme}." + " Please open an issue." + ) + + class OCP_MXQuantizationEmulationTritonExperts(TritonExperts): """ Extension of TritonExperts to support emulated OCP MX MoE experts. @@ -72,29 +118,7 @@ def __init__( self.quantization_emulation = True - if self.ocp_mx_scheme in { - OCP_MX_Scheme.w_mxfp4, - OCP_MX_Scheme.w_mxfp6_e3m2, - OCP_MX_Scheme.w_mxfp6_e2m3, - }: - # Weight-only schemes leave activations unquantized. - self._quant_dtype = None - elif self.ocp_mx_scheme in { - OCP_MX_Scheme.w_mxfp4_a_mxfp4, - }: - self._quant_dtype = "mxfp4" - elif self.ocp_mx_scheme in [ - OCP_MX_Scheme.w_mxfp4_a_mxfp6_e3m2, - OCP_MX_Scheme.w_mxfp4_a_mxfp6_e2m3, - OCP_MX_Scheme.w_mxfp6_e3m2_a_mxfp6_e3m2, - OCP_MX_Scheme.w_mxfp6_e2m3_a_mxfp6_e2m3, - ]: - self._quant_dtype = "mxfp6" - elif self.ocp_mx_scheme in [ - OCP_MX_Scheme.w_mxfp4_a_fp8, - OCP_MX_Scheme.w_mxfp6_e3m2_a_fp8, - ]: - self._quant_dtype = current_platform.fp8_dtype() + self._quant_dtype = activation_quant_dtype(self.ocp_mx_scheme) @staticmethod def is_supported_config( From 54dd98be2830bbd7b70fd60a0e917b7d0df5adab Mon Sep 17 00:00:00 2001 From: Yi Liu Date: Thu, 20 Aug 2026 00:31:32 +0800 Subject: [PATCH 157/839] [CT] Support Humming for WNA16 MoE (#48918) Signed-off-by: yiliu30 Co-authored-by: OpenAI Codex Co-authored-by: OpenAI Codex Co-authored-by: Misha Goin Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/quantization/test_compressed_tensors.py | 44 +++++++ tests/quantization/test_moe_wna16.py | 110 ++++++++++++++++++ .../fused_moe/experts/fused_humming_moe.py | 32 ++++- .../layers/fused_moe/oracle/int_wna16.py | 16 ++- .../compressed_tensors_moe_wna16.py | 56 ++++++--- .../quantization/utils/humming_utils.py | 3 + .../layers/quantization/utils/quant_utils.py | 16 +-- 7 files changed, 253 insertions(+), 24 deletions(-) diff --git a/tests/quantization/test_compressed_tensors.py b/tests/quantization/test_compressed_tensors.py index 3f11cec1c143..a301e8d9cfcc 100644 --- a/tests/quantization/test_compressed_tensors.py +++ b/tests/quantization/test_compressed_tensors.py @@ -43,8 +43,14 @@ find_matched_target, ) from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8 +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + GroupShape, + QuantKey, + ScaleDesc, +) from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead from vllm.platforms import current_platform +from vllm.scalar_type import scalar_types from vllm.v1.attention.backends.fa_utils import get_flash_attn_version # AITER only supports per-channel-per-channel INT8 gemm @@ -971,6 +977,44 @@ def test_wna16_moe_w2_scale_sharding(actorder, group_size, part, full, expected) assert result == expected +@pytest.mark.parametrize("num_bits", range(2, 9)) +def test_humming_supports_compressed_tensors_wna16_quant_key(num_bits): + from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( + HummingExpertsBase, + ) + from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import ( # noqa: E501 + WNA16_SUPPORTED_TYPES_MAP, + ) + + weight_key = QuantKey( + dtype=WNA16_SUPPORTED_TYPES_MAP[num_bits], + scale=ScaleDesc( + dtype=torch.float16, + static=True, + group_shape=GroupShape(row=1, col=128), + ), + symmetric=True, + ) + + assert HummingExpertsBase._supports_quant_scheme(weight_key, None) + + +def test_quant_key_str_supports_scalar_type_dtypes(): + quant_key = QuantKey( + dtype=scalar_types.uint2b2, + scale=ScaleDesc( + dtype=torch.float16, + static=True, + group_shape=GroupShape(row=1, col=128), + ), + symmetric=True, + ) + + assert str(quant_key) == ( + "QuantKey(uint2b2,scale(f16,static,GroupShape(row=1, col=128)),symmetric)" + ) + + @pytest.mark.skipif( not current_platform.is_cuda() or not current_platform.has_device_capability(80), reason="MXFP4 requires ampere or newer", diff --git a/tests/quantization/test_moe_wna16.py b/tests/quantization/test_moe_wna16.py index 8624832e7bcb..49e224fef4d3 100644 --- a/tests/quantization/test_moe_wna16.py +++ b/tests/quantization/test_moe_wna16.py @@ -26,6 +26,7 @@ MoeWNA16Config, MoeWNA16Method, ) +from vllm.platforms import current_platform def test_map_wna16_backend_supports_triton(): @@ -203,3 +204,112 @@ def test_moe_wna16_uses_humming_quant_config(monkeypatch): ) assert method.get_fused_moe_quant_config(layer) is quant_config + + +@pytest.mark.skipif( + not current_platform.is_cuda(), + reason="Compressed-tensors Humming WNA16 MoE requires CUDA", +) +@pytest.mark.parametrize("num_bits", [3, 5, 6, 7]) +def test_compressed_tensors_wna16_moe_create_weights_uses_ceil_packed_shapes( + num_bits, +): + pytest.importorskip("humming") + + from tests.kernels.moe.utils import make_dummy_moe_config + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_wna16 import ( # noqa: E501 + CompressedTensorsWNA16MoEMethod, + ) + + quant_args = QuantizationArgs( + num_bits=num_bits, + type=QuantizationType.INT, + strategy=QuantizationStrategy.GROUP, + symmetric=True, + dynamic=False, + group_size=128, + ) + moe_config = make_dummy_moe_config( + num_experts=2, + hidden_dim=256, + intermediate_size=512, + ) + moe_config.moe_backend = "humming" + method = CompressedTensorsWNA16MoEMethod(quant_args, None, moe_config) + layer = torch.nn.Module() + + method.create_weights( + layer, + num_experts=2, + hidden_size=256, + intermediate_size_per_partition=512, + params_dtype=torch.float16, + intermediate_size_full=512, + ) + + packed_hidden = (256 * num_bits + 31) // 32 + packed_intermediate = (512 * num_bits + 31) // 32 + assert method.wna16_backend == WNA16MoEBackend.HUMMING + assert layer.w13_weight_packed.shape == (2, 1024, packed_hidden) + assert layer.w2_weight_packed.shape == (2, 256, packed_intermediate) + assert layer.w13_weight_scale.shape == (2, 1024, 2) + assert layer.w2_weight_scale.shape == (2, 256, 4) + assert layer.w13_weight_packed.dtype is torch.int32 + assert layer.w2_weight_scale.dtype is torch.float16 + + +@pytest.mark.skipif( + not current_platform.is_cuda(), + reason="Compressed-tensors Humming WNA16 MoE requires CUDA", +) +def test_compressed_tensors_wna16_moe_converts_and_sets_up_humming_kernel(): + pytest.importorskip("humming") + + from tests.kernels.moe.utils import make_dummy_moe_config + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_wna16 import ( # noqa: E501 + CompressedTensorsWNA16MoEMethod, + ) + + quant_args = QuantizationArgs( + num_bits=3, + type=QuantizationType.INT, + strategy=QuantizationStrategy.GROUP, + symmetric=True, + dynamic=False, + group_size=128, + ) + moe_config = make_dummy_moe_config( + num_experts=2, + hidden_dim=256, + intermediate_size=512, + ) + moe_config.moe_backend = "humming" + method = CompressedTensorsWNA16MoEMethod(quant_args, None, moe_config) + layer = torch.nn.Module() + layer.moe_config = moe_config + layer.params_dtype = torch.bfloat16 + layer.layer_name = "test.humming_moe" + layer._expert_routing_tables = lambda: (None, None, None) + + method.create_weights( + layer, + num_experts=2, + hidden_size=256, + intermediate_size_per_partition=512, + intermediate_size_full=512, + params_dtype=torch.bfloat16, + ) + layer.cuda() + for parameter in layer.parameters(): + parameter.data.zero_() + + method.process_weights_after_loading(layer) + + assert method.wna16_backend == WNA16MoEBackend.HUMMING + assert method.moe_kernel is not None + assert set(layer.weight_schemas) == {"w13", "w2"} + assert set(layer.humming_configs) == {"w13", "w2"} + assert not hasattr(layer, "w13_weight_packed") + assert not hasattr(layer, "w2_weight_packed") + assert layer.w13_weight.dtype is torch.int32 + assert layer.w2_weight.dtype is torch.int32 diff --git a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py index 96c42967e9e6..69843c228a07 100644 --- a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py @@ -40,6 +40,9 @@ swiglu_limit_func, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( + INT4_DTYPE, + INT8_DTYPE, + GroupShape, QuantKey, kFp8Dynamic128Sym, kFp8DynamicTokenSym, @@ -58,6 +61,7 @@ kNvfp4Static, ) from vllm.platforms import current_platform +from vllm.scalar_type import ScalarType from vllm.utils.import_utils import has_humming from vllm.v1.worker.workspace import current_workspace_manager @@ -76,6 +80,29 @@ logger = init_logger(__name__) +def _is_supported_wna16_weight_key(weight_key: QuantKey | None) -> bool: + if weight_key is None or weight_key.scale2 is not None: + return False + + group_shape = weight_key.scale.group_shape + if not ( + group_shape == GroupShape.PER_CHANNEL + or (group_shape.row == 1 and group_shape.col > 0) + ): + return False + + dtype = weight_key.dtype + if dtype in (INT4_DTYPE, INT8_DTYPE, torch.uint8): + return True + + return ( + isinstance(dtype, ScalarType) + and dtype.is_integer() + and not dtype.is_signed() + and 2 <= dtype.size_bits <= 8 + ) + + def get_humming_moe_gemm_type() -> str: env_gemm_type: str | None = envs.VLLM_HUMMING_MOE_GEMM_TYPE gemm_type = "indexed" @@ -267,7 +294,10 @@ def _supports_quant_scheme( # mxfp8 (compressed-tensors / modelopt / online) (kMxfp8Static, kMxfp8Dynamic), ] - return (weight_key, activation_key) in SUPPORTED_W_A + return (weight_key, activation_key) in SUPPORTED_W_A or ( + activation_key in (None, kFp8DynamicTokenSym) + and _is_supported_wna16_weight_key(weight_key) + ) @property def expects_unquantized_inputs(self) -> bool: diff --git a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py index 5e4f30b9d197..c56712ac58bc 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -1104,9 +1104,21 @@ def _humming_wna16_weight_schema( "desc_act": quant_config.desc_act, "sym": quant_config.is_sym, } + if isinstance(quant_config, QuantizationArgs): + quant_type = getattr(quant_config.type, "value", quant_config.type) + quant_strategy = getattr(quant_config.strategy, "value", quant_config.strategy) + return { + "quant_method": "compressed-tensors", + "format": "pack-quantized", + "type": str(quant_type), + "num_bits": quant_config.num_bits, + "strategy": str(quant_strategy), + "group_size": quant_config.group_size, + "symmetric": quant_config.symmetric, + } raise TypeError( - "Humming WNA16 checkpoint schema requires AutoAWQConfig or " - "AutoGPTQConfig, " + "Humming WNA16 checkpoint schema requires AutoAWQConfig, " + "AutoGPTQConfig or QuantizationArgs, " f"got {type(quant_config).__name__}." ) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py index 35894b322482..1ab50851453d 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import math +from fractions import Fraction from typing import Any import torch @@ -40,7 +42,9 @@ marlin_make_workspace_new, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( + GroupShape, QuantKey, + ScaleDesc, kInt4Static32GroupScale, kInt4StaticGroupScale, kInt8StaticGroupScale, @@ -64,7 +68,7 @@ def __init__( # Extract properties from weight_quant self.symmetric = weight_quant.symmetric self.num_bits = weight_quant.num_bits - self.packed_factor = 32 // weight_quant.num_bits + self.packed_factor = Fraction(32, weight_quant.num_bits) self.strategy = weight_quant.strategy self.group_size = weight_quant.group_size self.actorder = weight_quant.actorder @@ -84,8 +88,14 @@ def __init__( elif self.num_bits == 8: scale = kInt8StaticGroupScale else: - raise ValueError( - "CompressedTensorsWNA16MoEMethod only supports int4 and int8 now." + scale = ScaleDesc( + dtype=torch.float16, + static=True, + group_shape=( + GroupShape.PER_CHANNEL + if self.group_size == -1 + else GroupShape(row=1, col=self.group_size) + ), ) weight_key = QuantKey(self.quant_type, scale, symmetric=self.symmetric) @@ -109,7 +119,10 @@ def __init__( WNA16MoEBackend.MARLIN, WNA16MoEBackend.BATCHED_MARLIN, ] - self.is_transposed = self.wna16_backend != WNA16MoEBackend.FLASHINFER_TRTLLM + self.is_transposed = self.wna16_backend not in ( + WNA16MoEBackend.FLASHINFER_TRTLLM, + WNA16MoEBackend.HUMMING, + ) if self.is_marlin: assert check_moe_marlin_supports_config( @@ -125,6 +138,9 @@ def __init__( # Non-Marlin WNA16 always uses bf16/fp16 inputs self.input_dtype = torch.bfloat16 + def _packed_dim(self, dim: int) -> int: + return math.ceil(dim * self.num_bits / 32) + def get_weight_shape( self, weight_name: str, @@ -149,16 +165,16 @@ def get_weight_shape( "num_groups_w2 must be provided for weight scales/zero_points" ) w13_num_shards = 2 if self.moe.is_act_and_mul else 1 - shape_map = { + shape_map: dict[str, dict[str, tuple[int, int | None, int | None]]] = { "w13_weight": { "Flashinfer": ( num_experts, w13_num_shards * intermediate_size_per_partition, - hidden_size // self.packed_factor, + self._packed_dim(hidden_size), ), "Marlin": ( num_experts, - hidden_size // self.packed_factor, + self._packed_dim(hidden_size), w13_num_shards * intermediate_size_per_partition, ), }, @@ -178,20 +194,18 @@ def get_weight_shape( "Marlin": ( num_experts, num_groups_w13, - w13_num_shards - * intermediate_size_per_partition - // self.packed_factor, + self._packed_dim(w13_num_shards * intermediate_size_per_partition), ), }, "w2_weight": { "Flashinfer": ( num_experts, hidden_size, - intermediate_size_per_partition // self.packed_factor, + self._packed_dim(intermediate_size_per_partition), ), "Marlin": ( num_experts, - intermediate_size_per_partition // self.packed_factor, + self._packed_dim(intermediate_size_per_partition), hidden_size, ), }, @@ -203,12 +217,14 @@ def get_weight_shape( "Marlin": ( num_experts, num_groups_w2, - hidden_size // self.packed_factor, + self._packed_dim(hidden_size), ), }, } backend_key = "Marlin" if self.is_transposed else "Flashinfer" - return shape_map[weight_name][backend_key] + shape = shape_map[weight_name][backend_key] + assert shape[1] is not None and shape[2] is not None + return shape[0], shape[1], shape[2] @staticmethod def _w2_scale_sharding( @@ -483,6 +499,7 @@ def _setup_kernel(self, layer: RoutedExperts): moe_quant_config=self.moe_quant_config, moe_config=self.moe, experts_cls=self.experts_cls, + backend=self.wna16_backend, routing_tables=layer._expert_routing_tables(), **marlin_args, ) @@ -586,6 +603,17 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: def get_fused_moe_quant_config( self, layer: torch.nn.Module ) -> FusedMoEQuantConfig | None: + if self.wna16_backend == WNA16MoEBackend.HUMMING: + from vllm.model_executor.layers.quantization.utils.humming_utils import ( + get_humming_moe_quant_config, + ) + + return get_humming_moe_quant_config( + layer, + gemm1_alpha=getattr(layer, "swiglu_alpha", None), + gemm1_beta=getattr(layer, "swiglu_beta", None), + gemm1_clamp_limit=getattr(layer, "swiglu_limit", None), + ) return make_wna16_moe_quant_config( w1_scale=layer.w13_weight_scale, w2_scale=layer.w2_weight_scale, diff --git a/vllm/model_executor/layers/quantization/utils/humming_utils.py b/vllm/model_executor/layers/quantization/utils/humming_utils.py index 7447cbef12c9..1303a92d91a9 100644 --- a/vllm/model_executor/layers/quantization/utils/humming_utils.py +++ b/vllm/model_executor/layers/quantization/utils/humming_utils.py @@ -73,6 +73,9 @@ class HummingMoEQuantConfig(FusedMoEQuantConfig): humming_dtypes.uint8: INT8_DTYPE, humming_dtypes.uint2: torch.uint8, humming_dtypes.uint3: torch.uint8, + humming_dtypes.uint5: torch.uint8, + humming_dtypes.uint6: torch.uint8, + humming_dtypes.uint7: torch.uint8, } _HUMMING_TO_SCALE_DTYPE: dict[humming_dtypes.DataType, torch.dtype] = { diff --git a/vllm/model_executor/layers/quantization/utils/quant_utils.py b/vllm/model_executor/layers/quantization/utils/quant_utils.py index 7ec658a375e4..20f33be240f3 100644 --- a/vllm/model_executor/layers/quantization/utils/quant_utils.py +++ b/vllm/model_executor/layers/quantization/utils/quant_utils.py @@ -25,6 +25,13 @@ INT8_DTYPE = scalar_types.uint8b128 +def _dtype_abbr(dtype: torch.dtype | ScalarType) -> str: + """Return a stable short name for torch and ScalarType dtypes.""" + if isinstance(dtype, ScalarType): + return str(dtype) + return fx.graph.dtype_abbrs[dtype] + + def weight_amax( weight: torch.Tensor, *, dim: int | None = None, keepdim: bool = False ) -> torch.Tensor: @@ -152,7 +159,7 @@ def __str__(self): group_shape = d.get(self.group_shape, str(self.group_shape)) return ( - f"{fx.graph.dtype_abbrs[self.dtype]}," + f"{_dtype_abbr(self.dtype)}," f"{'static' if self.static else 'dynamic'},{group_shape}" ) @@ -178,13 +185,8 @@ class QuantKey: def __str__(self): scale2_str = f"scale2({self.scale2})," if self.scale2 else "" - dtype_description = ( - fx.graph.dtype_abbrs[self.dtype] - if isinstance(self.dtype, torch.dtype) - else self.dtype - ) return ( - f"QuantKey({dtype_description}," + f"QuantKey({_dtype_abbr(self.dtype)}," f"scale({self.scale}),{scale2_str}" f"{'a' if not self.symmetric else ''}symmetric)" ) From e9e1630e93b13b241c4cdb52c47e35fd38c6eb46 Mon Sep 17 00:00:00 2001 From: Egor <47443236+Lossfull@users.noreply.github.com> Date: Wed, 19 Aug 2026 19:36:27 +0300 Subject: [PATCH 158/839] =?UTF-8?q?[Model]=20Support=20bidirectional=20(en?= =?UTF-8?q?coder-only)=20attention=20for=20DeepSeek=20e=E2=80=A6=20(#52948?= =?UTF-8?q?)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Egor Signed-off-by: Egor <47443236+Lossfull@users.noreply.github.com> Co-authored-by: Claude --- tests/models/registry.py | 5 +++++ vllm/config/model.py | 6 ++++++ vllm/model_executor/models/deepseek_v2.py | 26 ++++++++++++++++++++--- vllm/model_executor/models/registry.py | 1 + 4 files changed, 35 insertions(+), 3 deletions(-) diff --git a/tests/models/registry.py b/tests/models/registry.py index c862e767cd46..bc3ea0733dd7 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -593,6 +593,11 @@ def check_available_online( hf_overrides={"architectures": ["BertSpladeSparseEmbeddingModel"]}, ), "BgeM3EmbeddingModel": _HfExamplesInfo("BAAI/bge-m3"), + "DeepseekV3BidirectionalModel": _HfExamplesInfo( + "ai-sage/Giga-Embeddings-instruct-10B-A1.8B-0826", + trust_remote_code=True, + hf_overrides={"model_type": "deepseek_v3", "auto_map": None}, + ), "Gemma2Model": _HfExamplesInfo("BAAI/bge-multilingual-gemma2"), "Gemma3TextModel": _HfExamplesInfo("google/embeddinggemma-300m"), "GritLM": _HfExamplesInfo("parasail-ai/GritLM-7B-vllm"), diff --git a/vllm/config/model.py b/vllm/config/model.py index a5f98d8c5cca..16ff966c2d51 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -1877,6 +1877,12 @@ def use_mla(self) -> bool: # kv_lora_rank indicates that a Transformers model implementation uses MLA return getattr(self.hf_text_config, "kv_lora_rank", None) is not None # Manually maintained list of model types for vLLM model implementations + + # Bidirectional DeepSeek variants (is_causal=False, used by some + # embedding models) must use the non-MLA attention path, since the + # MLA kernels only support causal attention. + if not getattr(self.hf_text_config, "is_causal", True): + return False return self.is_deepseek_mla @property diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index d5521a1960b8..d8bc3fb64b6e 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -47,7 +47,11 @@ ) from vllm.logger import init_logger from vllm.model_executor.layers.activation import SiluAndMul -from vllm.model_executor.layers.attention import Attention, RSWAAttention +from vllm.model_executor.layers.attention import ( + Attention, + EncoderOnlyAttention, + RSWAAttention, +) from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.fused_moe import ( FusedMoEFactory, @@ -98,7 +102,7 @@ from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm.utils.torch_utils import direct_register_custom_op -from vllm.v1.attention.backend import AttentionBackend +from vllm.v1.attention.backend import AttentionBackend, AttentionType from vllm.v1.attention.backends.mla.indexer import ( DeepseekV32IndexerBackend, ) @@ -553,7 +557,22 @@ def __init__( mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim)) self.scaling = self.scaling * mscale * mscale - self.attn = Attention( + # DeepSeek is causal by default (decoder-only). Bidirectional + # (encoder-only) attention is used by embedding models derived from + # this architecture, which set `is_causal=False` on the HF config. + # This only applies on the non-MLA path, since the MLA kernels are + # causal-only; `ModelConfig.use_mla` already disables MLA for these + # models. + if getattr(config, "is_causal", True): + attn_type = AttentionType.DECODER + else: + attn_type = AttentionType.ENCODER_ONLY + attn_cls = ( + EncoderOnlyAttention + if attn_type == AttentionType.ENCODER_ONLY + else Attention + ) + self.attn = attn_cls( self.num_local_heads, self.qk_head_dim, self.scaling, @@ -561,6 +580,7 @@ def __init__( cache_config=cache_config, quant_config=quant_config, prefix=f"{prefix}.attn", + attn_type=attn_type, ) def forward( diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 75ef6756ddc9..8912827d0e8d 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -230,6 +230,7 @@ "BertSpladeSparseEmbeddingModel": ("bert", "BertSpladeSparseEmbeddingModel"), "BgeM3EmbeddingModel": ("roberta", "BgeM3EmbeddingModel"), "DeciLMForCausalLM": ("nemotron_nas", "DeciLMForCausalLM"), + "DeepseekV3BidirectionalModel": ("deepseek_v2", "DeepseekV3ForCausalLM"), "Gemma2Model": ("gemma2", "Gemma2ForCausalLM"), "Gemma3TextModel": ("gemma3", "Gemma3Model"), "GlmForCausalLM": ("glm", "GlmForCausalLM"), From 3a386cfaf51f872334c64edd08fbf71dd4bb0120 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Wed, 19 Aug 2026 11:39:59 -0500 Subject: [PATCH 159/839] [ROCm] Give EngineCore cleanup grace after request abort (#52281) Signed-off-by: Andreas Karatzas Signed-off-by: Andreas Karatzas Co-authored-by: OpenAI Codex Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../entrypoints/openai/test_dp_supervisor.py | 34 +++++ .../v1/engine/test_startup_watch_processes.py | 129 ++++++++++++++++++ vllm/entrypoints/openai/dp_supervisor.py | 7 +- vllm/v1/engine/core.py | 18 ++- vllm/v1/engine/utils.py | 36 ++++- 5 files changed, 221 insertions(+), 3 deletions(-) diff --git a/tests/entrypoints/openai/test_dp_supervisor.py b/tests/entrypoints/openai/test_dp_supervisor.py index 980f98712fa4..4bd0b4f7c3fe 100644 --- a/tests/entrypoints/openai/test_dp_supervisor.py +++ b/tests/entrypoints/openai/test_dp_supervisor.py @@ -242,6 +242,40 @@ def test_handles_shutdown_event(): assert supervisor.is_ready is False +@pytest.mark.asyncio +async def test_shutdown_children_uses_engine_process_timeout( + monkeypatch: pytest.MonkeyPatch, +): + supervisor = DPSupervisor(_make_unit_args(shutdown_timeout=0.0)) + supervisor._processes = [ + SimpleNamespace(name="APIServer_DPRank_4", pid=None, is_alive=lambda: False) + ] + calls = [] + timeout_calls = [] + + def get_process_timeout(request_timeout, manager_timeout): + timeout_calls.append((request_timeout, manager_timeout)) + return 15.0 + + monkeypatch.setattr( + dp_sup, + "get_engine_process_shutdown_timeout", + get_process_timeout, + ) + monkeypatch.setattr( + dp_sup, + "_join_processes_with_timeout", + lambda processes, timeout: calls.append((processes, timeout)), + ) + + await supervisor._shutdown_children() + + assert timeout_calls == [(0.0, 0.0)] + assert calls == [ + (supervisor._processes, 15.0 + CHILD_EXIT_GRACE_S), + ] + + @pytest.mark.asyncio async def test_handles_child_exit( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/v1/engine/test_startup_watch_processes.py b/tests/v1/engine/test_startup_watch_processes.py index b7b4b8a5a6e9..0b346efb65d7 100644 --- a/tests/v1/engine/test_startup_watch_processes.py +++ b/tests/v1/engine/test_startup_watch_processes.py @@ -2,14 +2,20 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from multiprocessing import connection +from threading import Event from types import SimpleNamespace import pytest import zmq +import vllm.platforms as platforms +from vllm.v1.engine import core as core_module +from vllm.v1.engine import utils as engine_utils +from vllm.v1.engine.core import EngineCoreProc, EngineShutdownState from vllm.v1.engine.utils import ( CoreEngine, CoreEngineLaunch, + CoreEngineProcManager, EngineZmqAddresses, wait_for_engine_startup, ) @@ -17,6 +23,129 @@ pytestmark = pytest.mark.skip_global_cleanup +@pytest.mark.parametrize( + ("is_rocm", "request_timeout", "manager_timeout", "process_timeout"), + [ + (True, 0, 0, 15.0), + (True, 0, 7, 7), + (True, 0, None, None), + (False, 0, 0, 0), + (True, 7, 0, 0), + ], +) +def test_engine_core_process_shutdown_timeout( + monkeypatch: pytest.MonkeyPatch, + is_rocm: bool, + request_timeout: float | None, + manager_timeout: float | None, + process_timeout: float | None, +): + manager = object.__new__(CoreEngineProcManager) + manager._request_shutdown_timeout = request_timeout + manager.manager_stopped = Event() + manager.processes = [object()] + detach_results = iter((object(), None)) + manager._finalizer = SimpleNamespace(detach=lambda: next(detach_results)) + + shutdown_calls = [] + monkeypatch.setattr( + engine_utils, + "current_platform", + SimpleNamespace(is_rocm=lambda: is_rocm), + ) + monkeypatch.setattr( + engine_utils, + "shutdown", + lambda processes, timeout: shutdown_calls.append((processes, timeout)), + ) + + manager.shutdown(timeout=manager_timeout) + manager.shutdown(timeout=manager_timeout) + + assert manager.manager_stopped.is_set() + assert shutdown_calls == [(manager.processes, process_timeout)] + + +@pytest.mark.parametrize( + ( + "is_rocm", + "shutdown_state", + "has_work", + "shutdown_timeout", + "exit_code", + "expected_calls", + ), + [ + ( + True, + EngineShutdownState.SHUTTING_DOWN, + False, + 0, + None, + ["shutdown", "freeze"], + ), + (False, EngineShutdownState.SHUTTING_DOWN, False, 0, None, ["shutdown"]), + (True, EngineShutdownState.RUNNING, False, 0, None, ["shutdown"]), + (True, EngineShutdownState.SHUTTING_DOWN, True, 0, None, ["shutdown"]), + (True, EngineShutdownState.SHUTTING_DOWN, False, 7, None, ["shutdown"]), + (True, EngineShutdownState.SHUTTING_DOWN, False, 0, 1, ["shutdown"]), + ], +) +def test_freeze_gc_after_clean_rocm_engine_core_shutdown( + monkeypatch: pytest.MonkeyPatch, + is_rocm: bool, + shutdown_state: EngineShutdownState, + has_work: bool, + shutdown_timeout: int, + exit_code: int | None, + expected_calls: list[str], +): + calls: list[str] = [] + vllm_config = SimpleNamespace(shutdown_timeout=shutdown_timeout) + proc = SimpleNamespace( + shutdown_state=EngineShutdownState.RUNNING, + has_work=lambda: has_work, + vllm_config=vllm_config, + ) + + def run_busy_loop(): + proc.shutdown_state = shutdown_state + raise SystemExit(exit_code) + + proc.run_busy_loop = run_busy_loop + proc.shutdown = lambda: calls.append("shutdown") + parallel_config = SimpleNamespace( + data_parallel_size=1, + numa_bind=False, + reconfigure_for_independent_dp_rank=lambda: None, + ) + vllm_config.parallel_config = parallel_config + + for name in ( + "maybe_register_config_serialize_by_value", + "set_process_title", + "maybe_init_worker_tracer", + "decorate_logs", + ): + monkeypatch.setattr(core_module, name, lambda *args, **kwargs: None) + monkeypatch.setattr(core_module, "EngineCoreProc", lambda *args, **kwargs: proc) + monkeypatch.setattr( + core_module, + "SignalCallback", + lambda callback: SimpleNamespace(trigger=lambda: None, stop=lambda: None), + ) + monkeypatch.setattr(core_module.signal, "signal", lambda *args: None) + monkeypatch.setattr( + platforms, "current_platform", SimpleNamespace(is_rocm=lambda: is_rocm) + ) + monkeypatch.setattr(core_module.gc, "freeze", lambda: calls.append("freeze")) + + with pytest.raises(SystemExit): + EngineCoreProc.run_engine_core(vllm_config=vllm_config) + + assert calls == expected_calls + + class _FinishedProcess: name = "RustFrontend" diff --git a/vllm/entrypoints/openai/dp_supervisor.py b/vllm/entrypoints/openai/dp_supervisor.py index c5efe916c025..da757264f324 100644 --- a/vllm/entrypoints/openai/dp_supervisor.py +++ b/vllm/entrypoints/openai/dp_supervisor.py @@ -30,6 +30,7 @@ kill_process_tree, set_process_title, ) +from vllm.v1.engine.utils import get_engine_process_shutdown_timeout logger = init_logger(__name__) @@ -513,7 +514,11 @@ async def _monitor_children(self) -> None: async def _shutdown_children(self) -> None: """Terminate the vLLM DP servers.""" - timeout = self.args.shutdown_timeout + CHILD_EXIT_GRACE_S + process_timeout = get_engine_process_shutdown_timeout( + self.args.shutdown_timeout, self.args.shutdown_timeout + ) + assert process_timeout is not None + timeout = process_timeout + CHILD_EXIT_GRACE_S try: logger.info( diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 71a791d72185..8946271aafaf 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -1276,6 +1276,7 @@ def run_engine_core(*args, dp_rank: int = 0, local_dp_rank: int = 0, **kwargs): engine_core: EngineCoreProc | None = None signal_callback: SignalCallback | None = None + clean_shutdown = False try: vllm_config: VllmConfig = kwargs["vllm_config"] parallel_config: ParallelConfig = vllm_config.parallel_config @@ -1338,8 +1339,15 @@ def signal_handler(signum, frame): engine_core.run_busy_loop() - except SystemExit: + except SystemExit as e: logger.info_once("[shutdown] EngineCore: exiting busy loop") + clean_shutdown = ( + e.code in (None, 0) + and engine_core is not None + and engine_core.shutdown_state == EngineShutdownState.SHUTTING_DOWN + and not engine_core.has_work() + and engine_core.vllm_config.shutdown_timeout == 0 + ) raise except Exception as e: if engine_core is None: @@ -1355,6 +1363,14 @@ def signal_handler(signum, frame): signal_callback.stop() if engine_core is not None: engine_core.shutdown() + if clean_shutdown: + from vllm.platforms import current_platform + + if current_platform.is_rocm(): + # Cleanup above already unfreezes and collects the heap. + # Freeze the surviving graph to skip another slow cyclic-GC + # scan during finalization; process exit reclaims it. + gc.freeze() def _init_data_parallel(self, vllm_config: VllmConfig): pass diff --git a/vllm/v1/engine/utils.py b/vllm/v1/engine/utils.py index 7736062d99cb..9bce2b7687ef 100644 --- a/vllm/v1/engine/utils.py +++ b/vllm/v1/engine/utils.py @@ -40,6 +40,30 @@ logger = init_logger(__name__) STARTUP_POLL_PERIOD_MS = 10000 +ROCM_ENGINE_PROCESS_SHUTDOWN_TIMEOUT_S = 15.0 + + +def get_engine_process_shutdown_timeout( + request_timeout: float | None, + process_timeout: float | None, +) -> float | None: + """Return the EngineCore process-manager shutdown timeout. + + ``VllmConfig.shutdown_timeout`` controls how long in-flight requests may + drain. A value of zero therefore tells EngineCore to abort requests as soon + as it receives SIGTERM. The parent process manager still needs a separate + window in which the EngineCore can release device resources before it is + force-killed. ROCm teardown can take longer than the generic best-effort + window, and force-killing during teardown can leave VRAM resident. + + ``process_timeout`` may be a remaining budget computed by an outer process + manager. Keep it unchanged unless both values are zero: a zero remaining + budget for a positive request timeout must not receive a fresh grace period + because EngineCore relies on that deadline to enforce request draining. + """ + if request_timeout == 0 and process_timeout == 0 and current_platform.is_rocm(): + return ROCM_ENGINE_PROCESS_SHUTDOWN_TIMEOUT_S + return process_timeout class CoreEngineState(Enum): @@ -136,6 +160,7 @@ def __init__( client_handshake_address: str | None = None, tensor_queue: Queue | None = None, ): + self._request_shutdown_timeout = vllm_config.shutdown_timeout context = get_mp_context() common_kwargs = { "vllm_config": vllm_config, @@ -217,7 +242,16 @@ def shutdown(self, timeout: float | None = None) -> None: """Shutdown engine core processes with configurable timeout.""" self.manager_stopped.set() if self._finalizer.detach() is not None: - shutdown(self.processes, timeout=timeout) + process_timeout = get_engine_process_shutdown_timeout( + self._request_shutdown_timeout, timeout + ) + if process_timeout != timeout: + logger.info( + "[shutdown] EngineCore process manager: using %ss ROCm " + "cleanup grace after immediate request abort", + process_timeout, + ) + shutdown(self.processes, timeout=process_timeout) def monitor_engine_liveness(self) -> None: """Monitor engine core process liveness.""" From 0c8c3f41cff4a4d361d4a80804eddd8dfc2a56ce Mon Sep 17 00:00:00 2001 From: lxy <145345338+lxyxinyi@users.noreply.github.com> Date: Thu, 20 Aug 2026 00:44:36 +0800 Subject: [PATCH 160/839] [Bugfix][V1] Sync mamba_block_size via EngineCoreReadyResponse (#50809) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: “新颐” Co-authored-by: “新颐” --- tests/v1/engine/test_engine_core_client.py | 36 ++++++++++++++++++++++ vllm/v1/engine/__init__.py | 1 + vllm/v1/engine/core.py | 1 + vllm/v1/engine/core_client.py | 1 + 4 files changed, 39 insertions(+) diff --git a/tests/v1/engine/test_engine_core_client.py b/tests/v1/engine/test_engine_core_client.py index 1727c51b19e0..7424b7622384 100644 --- a/tests/v1/engine/test_engine_core_client.py +++ b/tests/v1/engine/test_engine_core_client.py @@ -343,6 +343,42 @@ def test_apply_ready_response_syncs_block_size(): assert client.vllm_config.cache_config.block_size == 1056 +def test_apply_ready_response_syncs_mamba_block_size(): + import msgspec + + client = object.__new__(MPClient) + client.vllm_config = SimpleNamespace( + cache_config=SimpleNamespace(block_size=16, num_gpu_blocks=0), + model_config=SimpleNamespace(max_model_len=8192), + ) + client.stats_update_address = None + + payload = msgspec.msgpack.encode( + EngineCoreReadyResponse( + max_model_len=8192, + num_gpu_blocks=100, + block_size=1056, + dp_stats_address=None, + dtype="bfloat16", + vllm_version="test", + world_size=1, + data_parallel_size=1, + tensor_parallel_size=1, + pipeline_parallel_size=1, + decode_context_parallel_size=1, + data_parallel_rank=0, + max_num_seqs=256, + max_num_batched_tokens=8192, + instance_id="test-instance", + supports_lora=False, + max_loras=0, + mamba_block_size=1056, + ) + ) + client._apply_ready_response(payload) + assert client.vllm_config.cache_config.mamba_block_size == 1056 + + def loop_until_done(client: EngineCoreClient, outputs: dict): while True: engine_core_outputs = client.get_output().outputs diff --git a/vllm/v1/engine/__init__.py b/vllm/v1/engine/__init__.py index e07df08f3d9a..0a7f440ef799 100644 --- a/vllm/v1/engine/__init__.py +++ b/vllm/v1/engine/__init__.py @@ -90,6 +90,7 @@ class EngineCoreReadyResponse: instance_id: str supports_lora: bool max_loras: int + mamba_block_size: int | None = None # KV cache capacity (None for encoder-only/attention-free models). kv_cache_size_tokens: int | None = None kv_cache_max_concurrency: float | None = None diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 8946271aafaf..96010465f3d5 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -1636,6 +1636,7 @@ def _make_ready_response(self) -> EngineCoreReadyResponse: max_model_len=self.vllm_config.model_config.max_model_len, num_gpu_blocks=self.vllm_config.cache_config.num_gpu_blocks or 0, block_size=self.vllm_config.cache_config.block_size, + mamba_block_size=self.vllm_config.cache_config.mamba_block_size, dp_stats_address=self.frontend_stats_publish_address, dtype=str(self.vllm_config.model_config.dtype).removeprefix("torch."), vllm_version=VLLM_VERSION, diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index a626af89e1f3..2bf7cf03c188 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -758,6 +758,7 @@ def _apply_ready_response(self, payload: bytes) -> None: # worker for hybrid Mamba models. cache_config = vllm_config.cache_config cache_config.block_size = response.block_size + cache_config.mamba_block_size = response.mamba_block_size # Keep these as per-engine cache_config_info values; do not sum across DP. cache_config.kv_cache_size_tokens = ( getattr(cache_config, "kv_cache_size_tokens", None) From 9b5f3454f2be4c2c0f17c84a78af25a1d550d446 Mon Sep 17 00:00:00 2001 From: linitra24 Date: Thu, 20 Aug 2026 00:48:11 +0800 Subject: [PATCH 161/839] [LoRA] Avoid false target matches for unsupported module types (#52313) Signed-off-by: linitra24 Co-authored-by: Jee Jee Li --- tests/lora/test_lora_utils.py | 32 +------------------------------ vllm/lora/model_manager.py | 22 +++++++++++++++++---- vllm/lora/utils.py | 36 +++++++++++++++++++---------------- 3 files changed, 39 insertions(+), 51 deletions(-) diff --git a/tests/lora/test_lora_utils.py b/tests/lora/test_lora_utils.py index 603ec9297491..771e5b39b309 100644 --- a/tests/lora/test_lora_utils.py +++ b/tests/lora/test_lora_utils.py @@ -1,37 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from vllm.lora.utils import is_in_target_modules, is_supported_lora_module - - -class TestIsSupportedLoraModule: - """Tests for is_supported_lora_module (model-definition check).""" - - def test_suffix_match(self): - assert is_supported_lora_module( - "model.layers.0.self_attn.o_proj", ["o_proj", "q_proj"] - ) - - def test_no_match(self): - assert not is_supported_lora_module( - "model.layers.0.self_attn.o_proj", ["q_proj", "k_proj"] - ) - - def test_exact_match(self): - assert is_supported_lora_module("o_proj", ["o_proj"]) - - def test_regex_suffix_matching(self): - """Regex anchors to end — partial suffix should not match.""" - assert not is_supported_lora_module("model.layers.0.self_attn.o_proj", ["proj"]) - - def test_empty_supported_modules(self): - assert not is_supported_lora_module("model.layers.0.self_attn.o_proj", []) - - def test_multiple_supported_modules(self): - supported = ["q_proj", "k_proj", "v_proj", "o_proj"] - assert is_supported_lora_module("model.layers.0.self_attn.v_proj", supported) - assert not is_supported_lora_module("model.layers.0.mlp.gate_proj", supported) +from vllm.lora.utils import is_in_target_modules class TestIsInTargetModules: diff --git a/vllm/lora/model_manager.py b/vllm/lora/model_manager.py index 5db3cd4faab9..48e31a5c03dc 100644 --- a/vllm/lora/model_manager.py +++ b/vllm/lora/model_manager.py @@ -418,7 +418,16 @@ def _parent_module(module_name: str) -> str: if isinstance(module, PPMissingLayer): continue - if not self._match_target_modules(module_name): + target_modules = self.lora_config.target_modules + is_configured_target = target_modules is not None and is_in_target_modules( + module_name, + target_modules, + self.packed_modules_mapping, + ) + if ( + not self._match_target_modules(module_name, module) + and not is_configured_target + ): continue punica_wrapper = self._get_punica_wrapper(module_name) @@ -565,7 +574,7 @@ def create_dummy_lora( model = LoRAModel(lora_id, rank, {}) for module_name, module in self.model.named_modules(): if ( - not self._match_target_modules(module_name) + not self._match_target_modules(module_name, module) or not isinstance(module, BaseLayerWithLoRA) or self._get_punica_wrapper(module_name) is None ): @@ -701,7 +710,7 @@ def get_dummy_lora_warmup_rank(self, default_rank: int) -> int: ) return adjusted_rank - def _match_target_modules(self, module_name: str) -> bool: + def _match_target_modules(self, module_name: str, module: nn.Module) -> bool: """Check if a module should have LoRA applied. This method first checks if the module is in vLLM's supported LoRA @@ -711,11 +720,16 @@ def _match_target_modules(self, module_name: str) -> bool: Args: module_name: Full dot-separated module name (e.g., "model.layers.0.self_attn.o_proj") + module: Runtime module associated with ``module_name``. Returns: True if LoRA should be applied to this module, False otherwise. """ - if not is_supported_lora_module(module_name, self.supported_lora_modules): + if not is_supported_lora_module( + module_name, + module, + self.supported_lora_modules, + ): return False return is_in_target_modules( module_name, diff --git a/vllm/lora/utils.py b/vllm/lora/utils.py index c5d5765bd078..f2a9dc902fce 100644 --- a/vllm/lora/utils.py +++ b/vllm/lora/utils.py @@ -4,7 +4,6 @@ import os from typing import TYPE_CHECKING -import regex as re from huggingface_hub.utils import HfHubHTTPError, HFValidationError from torch import nn from transformers import PretrainedConfig @@ -33,8 +32,10 @@ RowParallelLinearWithShardedLoRA, VocabParallelEmbeddingWithLoRA, ) +from vllm.model_executor.custom_op import maybe_get_oot_by_class from vllm.model_executor.layers.fused_moe import MoERunner from vllm.model_executor.layers.linear import LinearBase +from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding from vllm.model_executor.utils import get_moe_expert_mapping, get_packed_modules_mapping from vllm.transformers_utils.repo_utils import hf_api @@ -230,11 +231,7 @@ def get_supported_lora_modules(model: nn.Module) -> list[str]: for name in embedding_modules: supported_lora_modules.add(name) - # get all the linear subfixes. - if isinstance(module, (LinearBase,)): - supported_lora_modules.add(name.split(".")[-1]) - - if isinstance(module, (MoERunner,)): + if isinstance(module, (LinearBase, MoERunner)): supported_lora_modules.add(name.split(".")[-1]) return list(supported_lora_modules) @@ -242,29 +239,36 @@ def get_supported_lora_modules(model: nn.Module) -> list[str]: def is_supported_lora_module( module_name: str, + module: nn.Module, supported_lora_modules: list[str], ) -> bool: """Check if a module is in the model's supported LoRA modules. - Uses regex suffix matching against the model-defined supported modules - list (e.g., matching "model.layers.0.self_attn.o_proj" against - "o_proj"). + The module name must match a model-supported suffix, and the runtime + module must belong to a module family handled by LoRA. Args: module_name: Full dot-separated module name. + module: Runtime module associated with ``module_name``. supported_lora_modules: List of module suffixes supported by the model. Returns: True if the module is supported, False otherwise. """ - return any( - re.match( - r".*\.{target_module}$".format(target_module=target_module), - module_name, - ) - or target_module == module_name - for target_module in supported_lora_modules + module_suffix = module_name.rsplit(".", 1)[-1] + if module_suffix not in supported_lora_modules: + return False + + return isinstance( + module, + ( + LinearBase, + MoERunner, + VocabParallelEmbedding, + maybe_get_oot_by_class(VocabParallelEmbedding), + BaseLayerWithLoRA, + ), ) From f76d71d7ceac040d555323f36b2f6afeab0e8df2 Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Wed, 19 Aug 2026 13:32:37 -0400 Subject: [PATCH 162/839] [CI/Build] Fix CPU platform pre-commit formatting (#52981) Signed-off-by: mgoin Co-authored-by: OpenAI Codex --- vllm/platforms/cpu.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index e002c80ad899..724f84012919 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -181,13 +181,16 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None: "otherwise the performance is not optimized." ) - # Accelerated GDN (AMX tiles or AVX-512BF16 VDPBF16PS) requires float32 SSM state. + # Accelerated GDN (AMX tiles or AVX-512BF16 VDPBF16PS) requires + # float32 SSM state. if ( torch.cpu._is_avx512_bf16_supported() and cache_config.mamba_ssm_cache_dtype != "float32" ): cache_config.mamba_ssm_cache_dtype = "float32" - logger.warning("Reset SSM cache type to float32 for accelerated GDN mamba attention.") + logger.warning( + "Reset SSM cache type to float32 for accelerated GDN mamba attention." + ) # Lagecy setting env_key = "VLLM_CPU_KVCACHE_SPACE" From 755492e37d7d7201d93f0effa200acac56601f38 Mon Sep 17 00:00:00 2001 From: Roger Wang Date: Wed, 19 Aug 2026 10:48:46 -0700 Subject: [PATCH 163/839] Revert "[Kernel] Gemma-4 FA4 FP8 Kernel" (#52987) --- cmake/external_projects/vllm_flash_attn.cmake | 2 +- .../layers/attention/attention.py | 14 +---- vllm/platforms/interface.py | 8 ++- vllm/v1/attention/backend.py | 14 ----- vllm/v1/attention/backends/fa_utils.py | 6 +- vllm/v1/attention/backends/flash_attn.py | 59 +------------------ .../gpu/spec_decode/gemma4/speculator.py | 55 ++++------------- vllm/vllm_flash_attn/flash_attn_interface.py | 23 +------- 8 files changed, 29 insertions(+), 152 deletions(-) diff --git a/cmake/external_projects/vllm_flash_attn.cmake b/cmake/external_projects/vllm_flash_attn.cmake index b70db3adeb90..2a8e9cc781bb 100644 --- a/cmake/external_projects/vllm_flash_attn.cmake +++ b/cmake/external_projects/vllm_flash_attn.cmake @@ -39,7 +39,7 @@ else() FetchContent_Declare( vllm-flash-attn GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git - GIT_TAG 617264c1c7955c9e84817654ebeedff069f3c5f1 + GIT_TAG f3e1a4f74c99145c0717709860bf765de1703779 GIT_PROGRESS TRUE # Don't share the vllm-flash-attn build between build types BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn diff --git a/vllm/model_executor/layers/attention/attention.py b/vllm/model_executor/layers/attention/attention.py index b4831e2a0b41..ff0c7a33af59 100644 --- a/vllm/model_executor/layers/attention/attention.py +++ b/vllm/model_executor/layers/attention/attention.py @@ -8,10 +8,7 @@ import vllm.envs as envs from vllm.compilation.breakable_cudagraph import eager_break_during_capture -from vllm.config import ( - CacheConfig, - get_current_vllm_config, -) +from vllm.config import CacheConfig, get_current_vllm_config from vllm.config.vllm import VllmConfig from vllm.forward_context import ForwardContext, get_forward_context from vllm.logger import init_logger @@ -97,7 +94,6 @@ def should_load_quant_weights(quant_method: QuantizeMethodBase | None) -> bool: def _largest_kernel_block_within( attn_backend: "type[AttentionBackend]", - vllm_config: VllmConfig, per_token_bytes: int, page_budget: int | None, fallback: int, @@ -112,7 +108,7 @@ def _largest_kernel_block_within( """ from vllm.v1.attention.backend import MultipleOf - sizes = attn_backend.get_supported_kernel_block_sizes_for_config(vllm_config) + sizes = attn_backend.get_supported_kernel_block_sizes() candidates = [s for s in sizes if isinstance(s, int)] if not candidates: candidates = [s.base for s in sizes if isinstance(s, MultipleOf)] @@ -636,11 +632,7 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: ) ).real_page_size_bytes sw_block_size = _largest_kernel_block_within( - self.attn_backend, - vllm_config, - sw_per_token, - shared_page, - block_size, + self.attn_backend, sw_per_token, shared_page, block_size ) return SlidingWindowSpec( block_size=sw_block_size, diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index c3ade53ca139..5c041f1632b7 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -612,6 +612,7 @@ def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None: For hybrid models, also aligns block_size with mamba page sizes. """ from vllm.config.cache import CacheConfig + from vllm.config.vllm import set_current_vllm_config cache_config = vllm_config.cache_config model_config = vllm_config.model_config @@ -626,9 +627,10 @@ def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None: # Phase 1: Pick block size from backend (skip if user set --block-size) if not cache_config.user_specified_block_size: - preferred = backend_cls.get_preferred_block_size_for_config( - CacheConfig.DEFAULT_BLOCK_SIZE, vllm_config - ) + with set_current_vllm_config(vllm_config): + preferred = backend_cls.get_preferred_block_size( + CacheConfig.DEFAULT_BLOCK_SIZE + ) if preferred != CacheConfig.DEFAULT_BLOCK_SIZE: logger.info( "Setting kv cache block size to %d for %s backend.", diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index daa217f8553d..d0da103fb399 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -69,13 +69,6 @@ class AttentionBackend(ABC): def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: return [MultipleOf(1)] - @classmethod - def get_supported_kernel_block_sizes_for_config( - cls, vllm_config: "VllmConfig" - ) -> list[int | MultipleOf]: - """Return kernel block sizes for a concrete engine configuration.""" - return cls.get_supported_kernel_block_sizes() - @staticmethod @abstractmethod def get_name() -> str: @@ -222,13 +215,6 @@ def get_preferred_block_size(cls, default_block_size: int) -> int: return min(s.base if isinstance(s, MultipleOf) else s for s in supported_sizes) - @classmethod - def get_preferred_block_size_for_config( - cls, default_block_size: int, vllm_config: "VllmConfig" - ) -> int: - """Return the preferred block size for a concrete engine config.""" - return cls.get_preferred_block_size(default_block_size) - @classmethod def indexes_kv_by_block_stride(cls) -> bool: """Whether the backend reads KV pages by the runtime block stride. diff --git a/vllm/v1/attention/backends/fa_utils.py b/vllm/v1/attention/backends/fa_utils.py index 907881e18773..4ebb1d05b904 100644 --- a/vllm/v1/attention/backends/fa_utils.py +++ b/vllm/v1/attention/backends/fa_utils.py @@ -238,9 +238,9 @@ def flash_attn_supports_kv_cache_dtype( head_size_v=head_size_v, has_sinks=has_sinks, ) - return ( - fa_version in (3, 4) and current_platform.is_device_capability_family(90) - ) or (fa_version == 4 and current_platform.is_device_capability_family(100)) + return (fa_version == 3 and current_platform.is_device_capability_family(90)) or ( + fa_version == 4 and current_platform.is_device_capability_family(100) + ) def flash_attn_supports_quant_query_input() -> bool: diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index d176bd5e928a..aa89e3010d68 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -85,58 +85,13 @@ class FlashAttentionBackend(AttentionBackend): ] @staticmethod - def _get_sm90_fa4_fp8_kv_block_size( - vllm_config: VllmConfig | None = None, - ) -> int | None: - if vllm_config is None: - vllm_config = get_current_vllm_config_or_none() - if vllm_config is None or vllm_config.model_config is None: - return None - - head_size = vllm_config.model_config.get_head_size() - if ( - current_platform.is_device_capability_family(90) - and vllm_config.cache_config.cache_dtype in ("fp8", "fp8_e4m3") - and head_size == 512 - and get_flash_attn_version(head_size=head_size) == 4 - ): - # The SM90 FP8-KV-dequant kernel uses a 64-token TMA tile/page. - return 64 - return None - - @classmethod - def get_supported_kernel_block_sizes(cls) -> list[int | MultipleOf]: - if block_size := cls._get_sm90_fa4_fp8_kv_block_size(): - # Sliding-window cache specs select the smallest advertised size. - # Report the kernel's exact page-size contract instead of the - # generic FlashAttention multiple-of-16 capability. - return [block_size] - return [MultipleOf(16)] - - @classmethod - def get_supported_kernel_block_sizes_for_config( - cls, vllm_config: VllmConfig - ) -> list[int | MultipleOf]: - if block_size := cls._get_sm90_fa4_fp8_kv_block_size(vllm_config): - return [block_size] + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: return [MultipleOf(16)] forward_includes_kv_cache_update: bool = False @classmethod def get_preferred_block_size(cls, default_block_size: int) -> int: - if block_size := cls._get_sm90_fa4_fp8_kv_block_size(): - return max(default_block_size, block_size) - if current_platform.is_xpu(): - return max(default_block_size, 64) - return super().get_preferred_block_size(default_block_size) - - @classmethod - def get_preferred_block_size_for_config( - cls, default_block_size: int, vllm_config: VllmConfig - ) -> int: - if block_size := cls._get_sm90_fa4_fp8_kv_block_size(vllm_config): - return max(default_block_size, block_size) if current_platform.is_xpu(): return max(default_block_size, 64) return super().get_preferred_block_size(default_block_size) @@ -939,17 +894,7 @@ def __init__( "heads in the layer" ) - # FA4's SM90 FP8-KV path consumes native FP16/BF16 Q and dequantizes - # FP8 K/V in-kernel. Other FA4 paths (notably SM100) still require Q, - # K, and V to have the same FP8 dtype. - uses_sm90_fa4_fp8_kv_dequant = ( - self.vllm_flash_attn_version == 4 - and current_platform.is_device_capability_family(90) - and self.kv_cache_dtype in ("fp8", "fp8_e4m3") - ) - self.supports_quant_query_input = flash_attn_supports_quant_query_input() and ( - not uses_sm90_fa4_fp8_kv_dequant - ) + self.supports_quant_query_input = flash_attn_supports_quant_query_input() vllm_config = get_current_vllm_config_or_none() dcp_a2a = ( diff --git a/vllm/v1/worker/gpu/spec_decode/gemma4/speculator.py b/vllm/v1/worker/gpu/spec_decode/gemma4/speculator.py index 34fc5f4bef39..dfa2c680109d 100644 --- a/vllm/v1/worker/gpu/spec_decode/gemma4/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/gemma4/speculator.py @@ -23,23 +23,6 @@ logger = init_logger(__name__) -def _copy_target_kv_scales(attn: nn.Module, target_attn: nn.Module) -> None: - """Copy target KV scales while preserving their tensor representation. - - Default attention scales are scalar buffers, while some quantization - methods replace them with length-one or per-head parameters. Re-register - cloned buffers on the draft layer so the shared KV cache is interpreted - with the target's values and shapes without aliasing target parameters. - """ - for scale_name in ("_k_scale", "_v_scale"): - target_scale = getattr(target_attn, scale_name) - attn.register_buffer(scale_name, target_scale.detach().clone()) - for scale_name in ("_k_scale_float", "_v_scale_float"): - setattr(attn, scale_name, getattr(target_attn, scale_name)) - for scale_name in ("_k_scale_cpu", "_v_scale_cpu"): - getattr(attn, scale_name).copy_(getattr(target_attn, scale_name)) - - class Gemma4Speculator(AutoRegressiveSpeculator): @property def advance_draft_positions(self) -> bool: @@ -93,7 +76,7 @@ def _setup_gemma4_kv_sharing( model: nn.Module, target_attn_layer_names: set[str], ) -> None: - """Wire draft layers to share KV and KV scales with the target model. + """Wire draft layers to share KV with the target model. Each draft decoder layer is mapped to the last non-KV-shared target layer of the same attention type (sliding or full). @@ -109,19 +92,15 @@ def _setup_gemma4_kv_sharing( target_num_kv_shared = getattr(target_text_config, "num_kv_shared_layers", 0) num_non_shared = len(target_layer_types) - target_num_kv_shared - target_names_by_index: dict[int, str] = {} - for name in target_attn_layer_names: - _, separator, layer_suffix = name.partition(".layers.") - if not separator: - continue - layer_index, _, _ = layer_suffix.partition(".") - if layer_index.isdigit(): - target_names_by_index[int(layer_index)] = name - - type_to_target_names: dict[str, list[str]] = defaultdict(list) + type_to_target_indices: dict[str, list[int]] = defaultdict(list) for idx, lt in enumerate(target_layer_types[:num_non_shared]): - if target_name := target_names_by_index.get(idx): - type_to_target_names[lt].append(target_name) + type_to_target_indices[lt].append(idx) + + target_prefix = "model.layers" + for name in target_attn_layer_names: + if ".layers." in name: + target_prefix = name.split(".layers.")[0] + ".layers" + break draft_layer_types = getattr(draft_text_config, "layer_types", []) for draft_idx, layer in enumerate(model.model.layers): @@ -136,7 +115,7 @@ def _setup_gemma4_kv_sharing( if draft_idx < len(draft_layer_types) else "full_attention" ) - candidates = type_to_target_names.get(draft_layer_type, []) + candidates = type_to_target_indices.get(draft_layer_type, []) if not candidates: logger.warning( "No target layer of type '%s' for draft layer %d", @@ -145,19 +124,9 @@ def _setup_gemma4_kv_sharing( ) continue - target_layer_name = candidates[-1] + target_idx = candidates[-1] + target_layer_name = f"{target_prefix}.{target_idx}.self_attn.attn" attn.kv_sharing_target_layer_name = target_layer_name - - # KV-cache sharing aliases the cache tensor during allocation, but - # the quantization scales live on the Attention modules themselves. - # The BF16 draft model has no quantization config, so its K/V scales - # otherwise remain at the default 1.0 while it reads the target's - # calibrated FP8 cache. - target_attn = self.vllm_config.compilation_config.static_forward_context[ - target_layer_name - ] - _copy_target_kv_scales(attn, target_attn) - logger.info( "Gemma4 MTP: draft layer %d (%s) -> %s", draft_idx, diff --git a/vllm/vllm_flash_attn/flash_attn_interface.py b/vllm/vllm_flash_attn/flash_attn_interface.py index aba0d297272c..81768e12092e 100644 --- a/vllm/vllm_flash_attn/flash_attn_interface.py +++ b/vllm/vllm_flash_attn/flash_attn_interface.py @@ -404,22 +404,6 @@ def flash_attn_varlen_func( from vllm.vllm_flash_attn.cute.interface import _flash_attn_fwd - # SM90 FA4 fp8-KV path: fp8 e4m3 paged K/V dequantized and the K/V descale folded - # in-kernel; accepts bf16/fp16 Q and writes O in its native dtype (no Q cast, no - # output copy). Only the (batch, num_kv_heads) f32 K/V descales are forwarded. - fa4_fp8_kv_dequant = ( - k.dtype == torch.float8_e4m3fn - and torch.cuda.get_device_capability()[0] == 9 - ) - if fa4_fp8_kv_dequant: - fa4_q_descale = None - fa4_k_descale = k_descale - fa4_v_descale = v_descale - else: - fa4_q_descale = None - fa4_k_descale = None - fa4_v_descale = None - out, softmax_lse, _, _ = _flash_attn_fwd( q, k, @@ -444,11 +428,10 @@ def flash_attn_varlen_func( block_sparse_tensors=block_sparse_tensors, aux_tensors=aux_tensors, aux_tensor_leading_dims=aux_tensor_leading_dims, - q_descale=fa4_q_descale, - k_descale=fa4_k_descale, - v_descale=fa4_v_descale, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, output_scale=output_scale, - fp8_kv_dequant=fa4_fp8_kv_dequant, ) else: raise ValueError(f"Unsupported FA version: {fa_version}") From 480d4f0d154a00b719abf876dd9c1a98bc5acbf7 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Wed, 19 Aug 2026 13:46:12 -0500 Subject: [PATCH 164/839] [ROCm][CI] Enable modular OAI Triton MoE tests (#46434) Signed-off-by: Andreas Karatzas Signed-off-by: Andreas Karatzas --- .../kernels/moe/test_modular_oai_triton_moe.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/tests/kernels/moe/test_modular_oai_triton_moe.py b/tests/kernels/moe/test_modular_oai_triton_moe.py index 0fdce2b82b7b..d0fd4df23ccd 100644 --- a/tests/kernels/moe/test_modular_oai_triton_moe.py +++ b/tests/kernels/moe/test_modular_oai_triton_moe.py @@ -257,7 +257,8 @@ def oai_triton_moe_impl( @pytest.mark.skipif( - not current_platform.is_cuda_alike(), reason="Requires CUDA-alike platform." + not OAITritonExperts._supports_current_device(), + reason="OAI Triton MoE is not supported on this device.", ) @pytest.mark.parametrize("dtype", [torch.bfloat16]) @pytest.mark.parametrize("m,n,k", MNK) @@ -319,7 +320,8 @@ def test_oai_triton_moe( @pytest.mark.skipif( - not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform." + not UnfusedOAITritonExperts._supports_current_device(), + reason="Unfused OAI Triton MoE is not supported on this device.", ) def test_unfused_oai_triton_experts_apply_direct_deepseek_v4_topology(workspace_init): """Exercise ``UnfusedOAITritonExperts.apply`` with explicit workspaces. @@ -345,10 +347,11 @@ def test_unfused_oai_triton_experts_apply_direct_deepseek_v4_topology(workspace_ w2_bias_tri, w1_precision_config, w2_precision_config, - _x_pad, + x_pad, ) = make_weights(dtype, k, n, num_experts) x = torch.randn((m, k), dtype=dtype, device="cuda") + x_tri = F.pad(x, (0, x_pad, 0, 0)) router_logits = torch.randn(m, num_experts, device="cuda", dtype=dtype) topk_weights, topk_ids = torch.topk(router_logits, k=topk, dim=-1, sorted=True) topk_weights = torch.nn.functional.softmax(topk_weights, dim=-1) @@ -367,10 +370,7 @@ def test_unfused_oai_triton_experts_apply_direct_deepseek_v4_topology(workspace_ ) experts = UnfusedOAITritonExperts(moe_config, quant_config) - if not UnfusedOAITritonExperts._supports_current_device(): - pytest.skip("UnfusedOAITritonExperts does not support this device") - - _, _, N, K, top_k = experts.moe_problem_size(x, w1_tri, w2_tri, topk_ids) + _, _, N, K, top_k = experts.moe_problem_size(x_tri, w1_tri, w2_tri, topk_ids) assert top_k == topk ws13_shape, ws2_shape, out_shape = experts.workspace_shapes( m, @@ -390,7 +390,7 @@ def test_unfused_oai_triton_experts_apply_direct_deepseek_v4_topology(workspace_ out_ref = torch_moe_impl(x, w1, w2, w1_bias, w2_bias, topk_weights, topk_ids) experts.apply( output=output, - hidden_states=x, + hidden_states=x_tri, w1=w1_tri, w2=w2_tri, topk_weights=topk_weights, @@ -405,5 +405,6 @@ def test_unfused_oai_triton_experts_apply_direct_deepseek_v4_topology(workspace_ expert_tokens_meta=None, apply_router_weight_on_input=False, ) + output = output[..., :k] assert_close(ref=out_ref, tri=output, maxtol=0.025, rmstol=0.005) From cb58bb9c1e38cd366910858873946c95be6328a9 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Wed, 19 Aug 2026 14:17:17 -0500 Subject: [PATCH 165/839] [CI] Harden RemoteVLLMServer GPU cleanup checks (#52282) Signed-off-by: Andreas Karatzas Signed-off-by: Andreas Karatzas Co-authored-by: OpenAI Codex --- tests/entrypoints/launchers/test_shutdown.py | 18 ++++++---- .../unit_tests/test_remote_vllm_server.py | 33 +++++++++++++++++++ tests/utils.py | 17 ++++++++-- 3 files changed, 59 insertions(+), 9 deletions(-) create mode 100644 tests/entrypoints/unit_tests/test_remote_vllm_server.py diff --git a/tests/entrypoints/launchers/test_shutdown.py b/tests/entrypoints/launchers/test_shutdown.py index e3be2919f5ef..643eff7594de 100644 --- a/tests/entrypoints/launchers/test_shutdown.py +++ b/tests/entrypoints/launchers/test_shutdown.py @@ -272,7 +272,9 @@ async def test_wait_timeout_completes_requests(): @pytest.mark.asyncio @pytest.mark.parametrize("wait_for_engine_idle", [0.0, 2.0]) -async def test_abort_timeout_exits_quickly(wait_for_engine_idle: float): +async def test_abort_timeout_exits_within_cleanup_grace( + wait_for_engine_idle: float, +): server_args = [ "--dtype", "bfloat16", @@ -303,19 +305,21 @@ async def test_abort_timeout_exits_quickly(wait_for_engine_idle: float): # Wait for engine to become idle await asyncio.sleep(wait_for_engine_idle) - start_time = time.time() proc.send_signal(signal.SIGTERM) - # abort timeout (0) should stop the server promptly. + # A zero request timeout aborts requests immediately, but process + # teardown may still use the platform-specific resource cleanup grace. + termination_timeout = remote_server._get_process_termination_timeout() try: - proc.wait(timeout=4.0) + proc.wait(timeout=termination_timeout) except subprocess.TimeoutExpired: proc.kill() proc.wait(timeout=5) - pytest.fail("Process did not exit after SIGTERM with abort timeout") + pytest.fail( + "Process did not exit within the resource cleanup grace " + f"({termination_timeout}s)" + ) - exit_time = time.time() - start_time - assert exit_time < 4.1, f"Default shutdown took too long: {exit_time:.1f}s" assert proc.returncode in (0, -15, None), f"Unexpected: {proc.returncode}" await _assert_children_cleaned_up(child_pids) diff --git a/tests/entrypoints/unit_tests/test_remote_vllm_server.py b/tests/entrypoints/unit_tests/test_remote_vllm_server.py new file mode 100644 index 000000000000..7f22746c88ca --- /dev/null +++ b/tests/entrypoints/unit_tests/test_remote_vllm_server.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from unittest.mock import Mock + +import pytest + +import tests.utils as test_utils +from tests.utils import RemoteOpenAIServer + + +def test_openai_server_shutdown_wait_covers_engine_cleanup( + monkeypatch: pytest.MonkeyPatch, +): + server = object.__new__(RemoteOpenAIServer) + server._request_shutdown_timeout = 0.0 + server.proc = Mock(pid=1234) + engine_timeout_args = [] + + def get_engine_timeout(request_timeout, process_timeout): + engine_timeout_args.append((request_timeout, process_timeout)) + return 60.0 + + monkeypatch.setattr( + test_utils, "get_engine_process_shutdown_timeout", get_engine_timeout + ) + monkeypatch.setattr(test_utils.os, "getpgid", lambda pid: pid) + monkeypatch.setattr(server, "_kill_process_group_survivors", Mock()) + + server._terminate_process_tree() + + assert engine_timeout_args == [(0.0, 0.0)] + server.proc.wait.assert_called_once_with(timeout=75.0) diff --git a/tests/utils.py b/tests/utils.py index 07601b74e486..71f9d4b669c8 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -63,6 +63,7 @@ from vllm.utils.torch_utils import ( set_random_seed, # noqa: F401 - re-exported for use in test files ) +from vllm.v1.engine.utils import get_engine_process_shutdown_timeout logger = init_logger(__name__) @@ -237,6 +238,9 @@ def _pre_download_model(self, model: str, args) -> None: model_loader = get_model_loader(load_config) model_loader.download_model(model_config) + def _get_process_termination_timeout(self) -> float: + return 15.0 + def __init__( self, model: str, @@ -287,6 +291,7 @@ def __init__( self.show_hidden_metrics = ( getattr(args, "show_hidden_metrics_for_version", None) is not None ) + self._request_shutdown_timeout = float(args.shutdown_timeout) with _temporarily_sanitized_pythonpath_env(): self._pre_download_model(model, args) @@ -416,7 +421,7 @@ def _terminate_process_tree(self) -> None: print(f"[RemoteOpenAIServer] Sent SIGTERM to process {pid}") try: - self.proc.wait(timeout=15) + self.proc.wait(timeout=self._get_process_termination_timeout()) print(f"[RemoteOpenAIServer] Server {pid} terminated gracefully") except subprocess.TimeoutExpired: # Phase 2: SIGKILL the entire process group @@ -758,6 +763,14 @@ def get_async_client_anthropic(self, **kwargs): class RemoteOpenAIServer(RemoteVLLMServer): """Launches ``vllm serve`` for testing OpenAI-compatible endpoints.""" + def _get_process_termination_timeout(self) -> float: + engine_timeout = get_engine_process_shutdown_timeout( + self._request_shutdown_timeout, + self._request_shutdown_timeout, + ) + assert engine_timeout is not None + return engine_timeout + super()._get_process_termination_timeout() + def _create_cli_subcommand(self): return ServeSubcommand() @@ -889,7 +902,7 @@ def _terminate_process_tree(self) -> None: self.proc.terminate() print(f"[RemoteOpenAIServerCustom] Sent SIGTERM to process {pid}") - self.proc.join(15) + self.proc.join(self._get_process_termination_timeout()) if self.proc.is_alive(): print( f"[RemoteOpenAIServerCustom] Server {pid} did not respond " From 541c6d64c19b64b970edc9122a56027d41b7a0be Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Wed, 19 Aug 2026 15:41:26 -0400 Subject: [PATCH 166/839] [Bugfix][Quantization] Support CT block FP8 with Marlin (#52966) Signed-off-by: mgoin Co-authored-by: OpenAI Codex --- tests/quantization/test_compressed_tensors.py | 43 +++++++++++++------ .../kernels/linear/scaled_mm/marlin.py | 15 +++++-- 2 files changed, 41 insertions(+), 17 deletions(-) diff --git a/tests/quantization/test_compressed_tensors.py b/tests/quantization/test_compressed_tensors.py index a301e8d9cfcc..304ef1efb0c6 100644 --- a/tests/quantization/test_compressed_tensors.py +++ b/tests/quantization/test_compressed_tensors.py @@ -21,6 +21,9 @@ from vllm.model_executor.kernels.linear import ( Fp8BlockScaledMMLinearKernel, ) +from vllm.model_executor.kernels.linear.scaled_mm import ( + MarlinFP8ScaledMMLinearKernel, +) from vllm.model_executor.layers.fused_moe import UnquantizedFusedMoEMethod from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501 @@ -527,9 +530,14 @@ def test_compressed_tensors_transforms_perplexity( assert perplexity <= exp_perplexity -def test_compressed_tensors_fp8_block_enabled(vllm_runner): +@pytest.mark.parametrize( + "linear_backend", ["auto", "marlin"] if current_platform.is_cuda() else ["auto"] +) +def test_compressed_tensors_fp8_block_enabled(vllm_runner, linear_backend): model_path = "RedHatAI/Qwen3-0.6B-FP8-BLOCK" - with vllm_runner(model_path, enforce_eager=True) as llm: + with vllm_runner( + model_path, enforce_eager=True, linear_backend=linear_backend + ) as llm: fp8_dtype = current_platform.fp8_dtype() def check_model(model): @@ -538,20 +546,29 @@ def check_model(model): qkv_proj = layer.self_attn.qkv_proj assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod) assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8Fp8) - assert isinstance(qkv_proj.scheme.fp8_linear, Fp8BlockScaledMMLinearKernel) - - assert qkv_proj.weight.dtype is fp8_dtype - assert qkv_proj.weight_scale.dtype is torch.float32 + if linear_backend == "marlin": + assert isinstance( + qkv_proj.scheme.fp8_linear, MarlinFP8ScaledMMLinearKernel + ) + assert qkv_proj.weight.dtype is torch.int32 + assert qkv_proj.weight_scale.dtype is qkv_proj.orig_dtype + else: + assert isinstance( + qkv_proj.scheme.fp8_linear, Fp8BlockScaledMMLinearKernel + ) + assert qkv_proj.weight.dtype is fp8_dtype + assert qkv_proj.weight_scale.dtype is torch.float32 assert len(qkv_proj.weight.shape) == 2 assert len(qkv_proj.weight_scale.shape) == 2 - input_quant_op = qkv_proj.scheme.fp8_linear.quant_fp8 - assert isinstance(input_quant_op, QuantFP8) - assert input_quant_op._forward_method in ( - input_quant_op.forward_cuda, - input_quant_op.forward_hip, - input_quant_op.forward_xpu, - ) + if linear_backend == "auto": + input_quant_op = qkv_proj.scheme.fp8_linear.quant_fp8 + assert isinstance(input_quant_op, QuantFP8) + assert input_quant_op._forward_method in ( + input_quant_op.forward_cuda, + input_quant_op.forward_hip, + input_quant_op.forward_xpu, + ) llm.apply_model(check_model) diff --git a/vllm/model_executor/kernels/linear/scaled_mm/marlin.py b/vllm/model_executor/kernels/linear/scaled_mm/marlin.py index 53bdc724bb90..9e148ca308cf 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/marlin.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/marlin.py @@ -57,14 +57,21 @@ def __init__( self.block_quant = self.config.weight_quant_key in {kFp8Static128BlockSym} self.size_k_first = not self.block_quant + @staticmethod + def _block_scale_name(layer: torch.nn.Module) -> str: + if getattr(layer, "weight_scale_inv", None) is not None: + return "weight_scale_inv" + return "weight_scale" + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: if self.block_quant: - weight, weight_scale_inv = process_fp8_weight_block_strategy( - layer.weight, layer.weight_scale_inv + scale_name = self._block_scale_name(layer) + weight, weight_scale = process_fp8_weight_block_strategy( + layer.weight, getattr(layer, scale_name) ) # Update layer with new values replace_parameter(layer, "weight", weight.data) - replace_parameter(layer, "weight_scale_inv", weight_scale_inv.data) + replace_parameter(layer, scale_name, weight_scale.data) # Non-block: callers must pass weight in (K, N) layout. layer.input_scale = None @@ -80,7 +87,7 @@ def apply_weights( bias: torch.Tensor | None = None, ) -> torch.Tensor: if self.block_quant: - weight_scale = layer.weight_scale_inv + weight_scale = getattr(layer, self._block_scale_name(layer)) else: weight_scale = layer.weight_scale return apply_fp8_marlin_linear( From d591d1d511b5f2a70ae34adfd5adc5c1956ecae2 Mon Sep 17 00:00:00 2001 From: vanshbhatia-amd Date: Wed, 19 Aug 2026 15:42:30 -0400 Subject: [PATCH 167/839] [Bugfix] Add Kimi K3 MoE support to benchmark_moe.py (#50082) Signed-off-by: Vansh Bhatia <210711135+vanshbhatia-amd@users.noreply.github.com> Co-authored-by: Vansh Bhatia <210711135+vanshbhatia-amd@users.noreply.github.com> --- benchmarks/kernels/benchmark_moe.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/benchmarks/kernels/benchmark_moe.py b/benchmarks/kernels/benchmark_moe.py index 5d86a7599fca..1fcc7ddde7f8 100644 --- a/benchmarks/kernels/benchmark_moe.py +++ b/benchmarks/kernels/benchmark_moe.py @@ -813,6 +813,19 @@ def get_model_params(config): topk = config.thinker_config.text_config.num_experts_per_tok intermediate_size = config.thinker_config.text_config.moe_intermediate_size hidden_size = config.thinker_config.text_config.hidden_size + elif architecture in ( + "KimiK3ForConditionalGeneration", + "KimiLinearForCausalLM", + ): + # Kimi K3 (multimodal) nests its MoE params in a KimiLinearConfig + # text_config and uses ``num_experts_per_token`` rather than the more + # common ``num_experts_per_tok``. get_text_config() returns the config + # itself for the text-only KimiLinearForCausalLM. + text_config = config.get_text_config() + E = text_config.num_experts + topk = text_config.num_experts_per_token + intermediate_size = text_config.moe_intermediate_size + hidden_size = text_config.hidden_size elif architecture == "PixtralForConditionalGeneration": # Pixtral can contain different LLM architectures, # recurse to get their parameters From c205726108df54bb6fbf15b19e725a4a3add2b18 Mon Sep 17 00:00:00 2001 From: jcotant-inferact Date: Wed, 19 Aug 2026 14:45:44 -0700 Subject: [PATCH 168/839] [CI] Fix and extend PR/issue auto-labeling (#51459) Signed-off-by: Joe Cotant Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .github/ISSUE_TEMPLATE/600-new-model.yml | 40 ----- .github/mergify.yml | 198 ++++++++++++++++++++--- .github/workflows/issue_autolabel.yml | 153 ++++++++++++++++++ .pre-commit-config.yaml | 7 + docs/contributing/labels.md | 95 +++++++++++ tools/pre_commit/check_label_rules.py | 97 +++++++++++ 6 files changed, 528 insertions(+), 62 deletions(-) delete mode 100644 .github/ISSUE_TEMPLATE/600-new-model.yml create mode 100644 docs/contributing/labels.md create mode 100644 tools/pre_commit/check_label_rules.py diff --git a/.github/ISSUE_TEMPLATE/600-new-model.yml b/.github/ISSUE_TEMPLATE/600-new-model.yml deleted file mode 100644 index 5f0125ef9809..000000000000 --- a/.github/ISSUE_TEMPLATE/600-new-model.yml +++ /dev/null @@ -1,40 +0,0 @@ -name: 🤗 Support request for a new model from huggingface -description: Submit a proposal/request for a new model from huggingface -title: "[New Model]: " -labels: ["new model"] - -body: -- type: markdown - attributes: - value: > - #### Before submitting an issue, please make sure the issue hasn't been already addressed by searching through [the existing and past issues](https://github.com/vllm-project/vllm/issues?q=is%3Aissue+sort%3Acreated-desc+). - - #### We also highly recommend you read https://docs.vllm.ai/en/latest/contributing/model/index.html first to understand how to add a new model. -- type: textarea - attributes: - label: The model to consider. - description: > - A huggingface url, pointing to the model, e.g. https://huggingface.co/openai-community/gpt2 . - validations: - required: true -- type: textarea - attributes: - label: The closest model vllm already supports. - description: > - Here is the list of models already supported by vllm: https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/models . Which model is the most similar to the model you want to add support for? -- type: textarea - attributes: - label: What's your difficulty of supporting the model you want? - description: > - For example, any new operators or new architecture? -- type: markdown - attributes: - value: > - Thanks for contributing 🎉! -- type: checkboxes - id: askllm - attributes: - label: Before submitting a new issue... - options: - - label: Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the [documentation page](https://docs.vllm.ai/en/latest/), which can answer lots of frequently asked questions. - required: true diff --git a/.github/mergify.yml b/.github/mergify.yml index 4e6de638ab59..04a623a499d1 100644 --- a/.github/mergify.yml +++ b/.github/mergify.yml @@ -93,6 +93,7 @@ pull_request_rules: - files~=^examples/.*deepseek.*\.py - files~=^tests/.*deepseek.*\.py - files~=^vllm/model_executor/models/.*deepseek.*\.py + - files~=^vllm/models/deepseek.*/ - files~=^vllm/tool_parsers/.*deepseek.*\.py - files~=^vllm/reasoning/.*deepseek.*\.py - files~=^vllm/transformers_utils/.*deepseek.*\.py @@ -133,9 +134,7 @@ pull_request_rules: - files~=^examples/.*llama.*\.py - files~=^tests/.*llama.*\.py - files~=^vllm/model_executor/models/.*llama.*\.py - - files~=^vllm/reasoning/.*llama.*\.py - files~=^vllm/tool_parsers/.*llama.*\.py - - files~=^vllm/transformers_utils/.*llama.*\.py - title~=(?i)llama actions: label: @@ -177,9 +176,17 @@ pull_request_rules: description: Automatically apply new-model label conditions: - label != stale - - and: - - files~=^vllm/model_executor/models/ - - files=vllm/model_executor/models/registry.py + - or: + # A model file is added alongside a registry entry. added-files rather + # than files: touching an existing model file is a fix or a refactor, + # and a PR that removes a model is not a new one. + - and: + - added-files~=^vllm/model_executor/models/.*\.py$ + - files=vllm/model_executor/models/registry.py + # A new top-level package under vllm/models/. common/ is shared kernel + # code rather than a model, and nested packages belong to a model that + # already exists. + - added-files~=^vllm/models/(?!common/)[^/]+/__init__\.py$ actions: label: add: @@ -233,16 +240,12 @@ pull_request_rules: conditions: - label != stale - or: - - files~=^examples/.*gpt[-_]?oss.*\.py - files~=^tests/.*gpt[-_]?oss.*\.py - - files~=^tests/entrypoints/openai/test_response_api_with_harmony.py - - files~=^tests/entrypoints/test_context.py - files~=^vllm/model_executor/models/.*gpt[-_]?oss.*\.py - files~=^vllm/model_executor/layers/.*gpt[-_]?oss.*\.py - files~=^vllm/entrypoints/openai/parser/harmony_utils.py - - files~=^vllm/entrypoints/tool_server.py - - files~=^vllm/entrypoints/tool.py - - files~=^vllm/entrypoints/context.py + # Harmony rendering now lives in the Rust frontend. + - files~=^rust/src/chat/src/.*/harmony/ - title~=(?i)gpt[-_]?oss - title~=(?i)harmony actions: @@ -256,7 +259,6 @@ pull_request_rules: - label != stale - or: - files~=(?i)kimi - - files~=(?i)moonshot - title~=(?i)(?:kimi|moonshot) actions: label: @@ -275,12 +277,73 @@ pull_request_rules: add: - k3 +- name: label-glm + description: Automatically apply glm label + conditions: + - label != stale + - or: + # Anchored to the start of the filename so midashenglm.py (an unrelated + # model) does not match. + - files~=^vllm/model_executor/models/(?:chat)?glm[^/]*\.py + - files~=^tests/.*[/_](?:chat)?glm[^/]*\.py + - files~=^vllm/tool_parsers/.*glm.*\.py + - files~=^vllm/transformers_utils/.*/(?:chat)?glm[^/]*\.py + - title~=(?i)\b(?:chat)?glm + actions: + label: + add: + - glm + +- name: label-minimax + description: Automatically apply minimax label + conditions: + - label != stale + - or: + - files~=^vllm/models/minimax.*/ + - files~=^vllm/model_executor/models/.*minimax.*\.py + - files~=^vllm/tool_parsers/.*minimax.*\.py + - files~=^tests/.*minimax.*\.py + - title~=(?i)minimax + actions: + label: + add: + - minimax + +- name: label-inkling + description: Automatically apply inkling label + conditions: + - label != stale + - or: + - files~=^vllm/models/inkling/ + - files~=^vllm/tool_parsers/inkling.*\.py + - files~=^vllm/renderers/inkling.*\.py + - files~=^tests/.*inkling.*\.py + - title~=(?i)inkling + actions: + label: + add: + - inkling + +- name: label-dsv4 + description: Automatically apply DSv4 label + conditions: + - label != stale + - or: + - files~=^vllm/models/deepseek_v4/ + - title~=(?i)(?:\bDSv4\b|deepseek[-\s]?v4) + actions: + label: + add: + - DSv4 + - name: label-nvidia description: Automatically apply nvidia label conditions: - label != stale - or: - - files~=cuda + # Exclude requirements/: a dependency bump touching requirements/cuda.txt + # is not NVIDIA backend work. + - files~=^(?!requirements/).*cuda - files~=cutlass - files~=flashinfer - files~=trtllm @@ -300,7 +363,7 @@ pull_request_rules: - files~=^csrc/rocm/ - files~=^docker/Dockerfile.rocm - files~=^requirements/rocm.*\.txt - - files~=^vllm/model_executor/layers/fused_moe/rocm.*\.py + - files~=^vllm/model_executor/layers/fused_moe/.*rocm.*\.py - files~=^vllm/v1/attention/backends/rocm.*\.py - files~=^vllm/v1/attention/backends/mla/rocm.*\.py - files~=^vllm/v1/attention/ops/rocm.*\.py @@ -319,8 +382,8 @@ pull_request_rules: - label != stale - or: - files~=^docker/Dockerfile.xpu - - files~=^\\.buildkite/intel_jobs/ - - files=\.buildkite/ci_config_intel.yaml + - files~=^\.buildkite/intel_jobs/ + - files=.buildkite/ci_config_intel.yaml - files=vllm/model_executor/layers/fused_moe/experts/xpu_moe.py - files=vllm/model_executor/kernels/linear/mixed_precision/xpu.py - files=vllm/model_executor/kernels/linear/mxfp8/xpu.py @@ -331,7 +394,6 @@ pull_request_rules: - files=vllm/v1/worker/xpu_worker.py - files=vllm/v1/worker/xpu_model_runner.py - files=vllm/_xpu_ops.py - - files=vllm/kernels/xpu_ops.py - files~=^vllm/lora/ops/xpu_ops - files=vllm/lora/punica_wrapper/punica_xpu.py - files=vllm/platforms/xpu.py @@ -367,7 +429,7 @@ pull_request_rules: - files=benchmarks/benchmark_serving_structured_output.py - files=benchmarks/run_structured_output_benchmark.sh - files=docs/features/structured_outputs.md - - files=^examples/features/structured_outputs/ + - files~=^examples/features/structured_outputs/ - files~=^tests/v1/structured_output/ - files=tests/entrypoints/llm/test_struct_output_generate.py - files~=^vllm/v1/structured_output/ @@ -382,8 +444,11 @@ pull_request_rules: - label != stale - or: - files~=^vllm/v1/spec_decode/ + # Model Runner V2 keeps its speculators in a separate subtree. + - files~=^vllm/v1/worker/gpu/spec_decode/ - files~=^tests/v1/spec_decode/ - - files=^examples/features/speculative_decoding/ + - files~=^tests/v1/e2e/spec_decode/ + - files~=^examples/features/speculative_decoding/ - files~=^vllm/model_executor/models/.*eagle.*\.py - files=vllm/model_executor/models/mlp_speculator.py - files~=^vllm/transformers_utils/configs/(eagle|medusa|mlp_speculator)\.py @@ -392,6 +457,33 @@ pull_request_rules: add: - speculative-decoding +# dflash covers the DFlash and DSpark speculative decoding techniques as one +# workstream, matching how they are tracked in the sprint board. It is applied +# in addition to speculative-decoding, never instead of it. +- name: label-dflash + description: Automatically apply dflash label + conditions: + - label != stale + - or: + - files~=^vllm/v1/spec_decode/dflash\.py + - files~=^vllm/v1/worker/gpu/spec_decode/dflash/ + - files~=^vllm/model_executor/models/.*dflash.*\.py + - files~=^tests/.*dflash.*\.py + - title~=(?i)dflash + - files~=^vllm/v1/worker/gpu/spec_decode/dspark/ + - files~=^vllm/model_executor/models/.*dspark.*\.py + - files~=^vllm/models/deepseek_v4/nvidia/dspark\.py + - files~=^vllm/models/deepseek_v4/amd/dspark\.py + - files~=^vllm/models/deepseek_v4/xpu/dspark\.py + - files~=^vllm/models/kimi_k3/nvidia/dspark_mla\.py + - files~=^vllm/transformers_utils/configs/k3_dspark\.py + - files~=^tests/.*dspark.*\.py + - title~=(?i)dspark + actions: + label: + add: + - dflash + - name: label-mrv2 description: Automatically apply mrv2 label conditions: @@ -416,7 +508,6 @@ pull_request_rules: - files~=_tpu - files~=tpu_ - files~=/tpu/ - - files~=pallas actions: label: add: @@ -432,7 +523,6 @@ pull_request_rules: - -files~=_tpu - -files~=tpu_ - -files~=/tpu/ - - -files~=pallas actions: label: remove: @@ -448,7 +538,6 @@ pull_request_rules: - files~=^tests/parser/ - files~=^tests/reasoning/ - files~=^tests/entrypoints/openai/.*tool.* - - files~=^tests/entrypoints/anthropic/.*tool.* - files~=^vllm/tool_parsers/ - files~=^vllm/parser/ - files~=^vllm/reasoning/ @@ -552,3 +641,68 @@ pull_request_rules: label: add: - kv-connector + +- name: label-scheduler + description: Automatically apply scheduler label + conditions: + - label != stale + - or: + - files~=^vllm/v1/core/sched/ + - files~=^tests/v1/core/test_scheduler.*\.py + actions: + label: + add: + - scheduler + +- name: label-kv-cache-manager + description: Automatically apply kv-cache-manager label + conditions: + - label != stale + - or: + # Scoped to the block/cache allocator in v1/core. KV transfer between + # instances is kv-connector, not this. + - files~=^vllm/v1/core/(?:kv_cache|block_pool|single_type_kv_cache) + - files~=^tests/v1/core/test_(?:kv_cache|single_type_kv_cache|prefix_caching).*\.py + actions: + label: + add: + - kv-cache-manager + +- name: label-torch-compile + description: Automatically apply torch.compile label + conditions: + - label != stale + - or: + - files~=^vllm/compilation/ + - files~=^tests/compile/ + - title~=(?i)torch\.compile + actions: + label: + add: + - torch.compile + +- name: label-ray + description: Automatically apply ray label + conditions: + - label != stale + - or: + - files~=^vllm/v1/executor/ray_ + - files~=^tests/.*ray_.*\.py + - title~=(?i)\bray\b + actions: + label: + add: + - ray + +- name: label-vllm-ir + description: Automatically apply vllm-ir label + conditions: + - label != stale + - or: + - files~=^vllm/ir/ + - files~=^tests/ir/ + - title~=(?i)vllm[-\s]?ir + actions: + label: + add: + - vllm-ir diff --git a/.github/workflows/issue_autolabel.yml b/.github/workflows/issue_autolabel.yml index 4a38aa2fbd2e..fd328e3ba895 100644 --- a/.github/workflows/issue_autolabel.yml +++ b/.github/workflows/issue_autolabel.yml @@ -190,6 +190,159 @@ jobs: }, ], }, + // Model labels. Title-only by design: issue bodies carry pasted + // collect_env output, tracebacks and configs that name unrelated + // models, hardware and libraries. + deepseek: { + keywords: [{ term: "DeepSeek", searchIn: "title" }], + }, + DSv4: { + keywords: [{ term: "DSv4", searchIn: "title" }], + substrings: [{ term: "deepseek-v4", searchIn: "title" }], + }, + llama: { + keywords: [{ term: "Llama", searchIn: "title" }], + }, + mistral: { + keywords: [ + { term: "Mistral", searchIn: "title" }, + { term: "Ministral", searchIn: "title" }, + { term: "Voxtral", searchIn: "title" }, + { term: "Mixtral", searchIn: "title" }, + { term: "Pixtral", searchIn: "title" }, + ], + }, + qwen: { + keywords: [{ term: "Qwen", searchIn: "title" }], + }, + "gpt-oss": { + keywords: [{ term: "harmony", searchIn: "title" }], + substrings: [{ term: "gpt-oss", searchIn: "title" }], + }, + glm: { + keywords: [ + { term: "GLM", searchIn: "title" }, + { term: "ChatGLM", searchIn: "title" }, + ], + }, + minimax: { + substrings: [{ term: "minimax", searchIn: "title" }], + }, + inkling: { + keywords: [{ term: "Inkling", searchIn: "title" }], + }, + // Hardware backends, mirroring the cpu/rocm/intel-gpu rules above. + nvidia: { + keywords: [ + { term: "NVIDIA", searchIn: "title" }, + { term: "CUTLASS", searchIn: "title" }, + ], + }, + tpu: { + keywords: [ + { term: "TPU", searchIn: "title" }, + { term: "XLA", searchIn: "title" }, + ], + }, + // Subsystems. + // + // Single distinctive words go in `keywords` so the word-boundary + // match keeps them from firing inside longer words. Multi-word + // phrases go in `substrings`: they are specific enough on their + // own, and a boundary match would miss ordinary variations -- + // "structured outputs", "tool parsers", "prefix caching". + "speculative-decoding": { + keywords: [ + { term: "eagle", searchIn: "title" }, + { term: "medusa", searchIn: "title" }, + ], + substrings: [ + { term: "speculative decoding", searchIn: "title" }, + { term: "speculative decode", searchIn: "title" }, + { term: "draft model", searchIn: "title" }, + ], + }, + "structured-output": { + substrings: [ + { term: "structured output", searchIn: "title" }, + { term: "guided decoding", searchIn: "title" }, + { term: "JSON schema", searchIn: "title" }, + ], + }, + "tool-calling": { + substrings: [ + { term: "tool calling", searchIn: "title" }, + { term: "tool parser", searchIn: "title" }, + { term: "function calling", searchIn: "title" }, + ], + }, + "kv-connector": { + keywords: [ + { term: "NIXL", searchIn: "title" }, + { term: "LMCache", searchIn: "title" }, + ], + substrings: [ + { term: "KV transfer", searchIn: "title" }, + { term: "disaggregat", searchIn: "title" }, + ], + }, + scheduler: { + keywords: [{ term: "scheduler", searchIn: "title" }], + }, + "kv-cache-manager": { + substrings: [ + { term: "kv cache manager", searchIn: "title" }, + { term: "block manager", searchIn: "title" }, + { term: "prefix cache", searchIn: "title" }, + { term: "prefix caching", searchIn: "title" }, + ], + }, + "torch.compile": { + substrings: [{ term: "torch.compile", searchIn: "title" }], + }, + "multi-modality": { + keywords: [ + { term: "multimodal", searchIn: "title" }, + { term: "multi-modal", searchIn: "title" }, + ], + substrings: [ + { term: "image input", searchIn: "title" }, + { term: "audio input", searchIn: "title" }, + ], + }, + "ci-failure": { + keywords: [{ term: "flaky", searchIn: "title" }], + substrings: [ + { term: "CI failure", searchIn: "title" }, + { term: "test failure", searchIn: "title" }, + ], + }, + rl: { + keywords: [{ term: "RLHF", searchIn: "title" }], + substrings: [ + { term: "reinforcement learning", searchIn: "title" }, + { term: "weight update", searchIn: "title" }, + ], + }, + ray: { + substrings: [ + { term: "Ray cluster", searchIn: "title" }, + { term: "Ray executor", searchIn: "title" }, + { term: "placement group", searchIn: "title" }, + ], + }, + "vllm-ir": { + substrings: [ + { term: "vllm-ir", searchIn: "title" }, + { term: "vllm ir", searchIn: "title" }, + ], + }, + rust: { + keywords: [ + { term: "Rust", searchIn: "title" }, + { term: "PyO3", searchIn: "title" }, + ], + }, // Add more label configurations here as needed // example: { // keywords: [...], diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index c1c5f5f90b05..b9344b6ffc8f 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -165,6 +165,13 @@ repos: language: python entry: python tools/pre_commit/generate_nightly_torch_test.py files: ^requirements/test/cuda\.(in|txt)$ + - id: check-label-rules + name: Check auto-label rules still match real files + language: python + entry: python tools/pre_commit/check_label_rules.py + files: ^\.github/mergify\.yml$ + pass_filenames: false + additional_dependencies: [PyYAML, regex] - id: mypy-3.10 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward name: Run mypy for Python 3.10 entry: python tools/pre_commit/mypy.py "3.10" diff --git a/docs/contributing/labels.md b/docs/contributing/labels.md new file mode 100644 index 000000000000..891f7e4ed5cd --- /dev/null +++ b/docs/contributing/labels.md @@ -0,0 +1,95 @@ +# Labels + +vLLM applies most labels automatically. Three systems do the work: + +| | Config | Applies to | +| --- | --- | --- | +| Mergify | `.github/mergify.yml` | pull requests, by changed file path or title | +| GitHub Action | `.github/workflows/issue_autolabel.yml` | issues, by title keyword | +| Issue templates | `.github/ISSUE_TEMPLATE/*.yml` | issues, by which template was used | + +A label that none of them applies has to be applied by hand, and in practice +that means it mostly doesn't get applied at all. + +A template's `labels:` field must name a label that already exists — GitHub +silently skips one that doesn't, with no error anywhere. + +`mergify.yml` also drives reviewer auto-assignment; see +[Collaboration](../governance/collaboration.md) for adding yourself as a +maintainer of an area. + +## Adding a label + +A new label needs three things. If it can't have all three, don't create it. + +1. **An audience.** Someone has to want to filter on it. "It would be nice to + know" isn't enough. +2. **An owner.** A person or SIG who triages that queue. Without one the label + accumulates issues nobody reads. +3. **A rule, in the same PR.** If you can't write a precise rule, the label + depends on humans remembering it exists. + +## Measure before you argue + +Volume claims should come from the repo, not from intuition. vLLM squash-merges, +so one commit on `main` is one PR: + +```bash +git log --since="6 months ago" --oneline -- | wc -l +``` + +For reference, at the time of writing: `speculative-decoding` 81, `llama` 85, +`mistral` 72, `tpu` 14. A subsystem in that range clears the bar comfortably. +Something in the low teens needs a stronger argument than volume. + +The same applies to claiming a rule is too broad. Check what it actually labels +before rewriting it: + +```bash +gh pr list --repo vllm-project/vllm --state merged --limit 200 \ + --json number,title,labels +``` + +Several rules in this repo have been called over-broad on inspection and turned +out to be accurate when measured. + +## Writing conditions + +**`files=` is exact equality. `files~=` is a regex.** Writing a pattern after +`=` compares the pattern text to each filename, so it can never match: + +```yaml +- files=^examples/features/speculative_decoding/ # never fires +- files~=^examples/features/speculative_decoding/ # correct +``` + +**Anchor file patterns.** An unanchored fragment matches far more than intended +— `files~=cuda` also matches `requirements/cuda.txt`. + +**Watch for names that contain other names.** `midashenglm.py` contains "glm" +but is unrelated to GLM, so the `glm` rule anchors to the start of the filename. + +**Keep an eye on where model code lives.** Newer models live in +`vllm/models//` rather than `vllm/model_executor/models/`. A rule that +only knows the old location silently stops matching. + +**Prefer titles for issue keywords.** Issue bodies contain pasted `collect_env` +output, tracebacks and configs that mention unrelated hardware and libraries. +Body matching tags the reporter's environment rather than the topic. + +**Single words go in `keywords`, phrases in `substrings`.** `keywords` matches on +word boundaries, which stops a term firing inside a longer word but also misses +ordinary variations — `structured output` will not match "structured outputs". + +## Keeping rules honest + +`tools/pre_commit/check_label_rules.py` fails if any file condition in +`mergify.yml` matches nothing in the tree. It runs via pre-commit when that file +changes. This catches conditions that were correct when written and rotted when +the code moved. + +## Not everything needs a label + +Triage and workflow labels — `bug`, `stale`, `ready`, `RFC`, `needs-rebase` — +and the bot provenance tags are applied by other mechanisms or by hand, and are +intentionally outside this system. diff --git a/tools/pre_commit/check_label_rules.py b/tools/pre_commit/check_label_rules.py new file mode 100644 index 000000000000..c5074d21d0e2 --- /dev/null +++ b/tools/pre_commit/check_label_rules.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Check that auto-labeling rules still match real files. + +Mergify silently ignores a condition that can never match, so a file-path +condition keeps passing review long after the file it names has moved or been +deleted. This walks every `files=` / `files~=` condition in `.github/mergify.yml` +and fails if one matches nothing in the tree. + +Usage: + python tools/pre_commit/check_label_rules.py +""" + +import subprocess +import sys + +import regex as re +import yaml + +MERGIFY = ".github/mergify.yml" + + +def tracked_files() -> list[str]: + out = subprocess.run( + ["git", "ls-files"], capture_output=True, text=True, check=True + ).stdout + return [line for line in out.split("\n") if line] + + +def file_conditions(node, found: list[str]) -> list[str]: + """Collect every string condition nested under a rule's `conditions`.""" + if isinstance(node, dict): + for value in node.values(): + file_conditions(value, found) + elif isinstance(node, list): + for value in node: + file_conditions(value, found) + elif isinstance(node, str): + found.append(node) + return found + + +def check(files: list[str]) -> list[tuple[str, str, str]]: + with open(MERGIFY) as f: + rules = yaml.safe_load(f)["pull_request_rules"] + + # removed-files is excluded: it names paths the tree no longer has. + attrs = ("files", "added-files", "modified-files") + + dead = [] + for rule in rules: + for cond in file_conditions(rule.get("conditions", []), []): + cond = cond.strip() + # Negated conditions (label-tpu-remove) are expected to match nothing. + if cond.startswith("-"): + continue + attr = next( + ( + a + for a in attrs + if cond.startswith(f"{a}~=") or cond.startswith(f"{a}=") + ), + None, + ) + if attr is None: + continue + if cond.startswith(f"{attr}~="): + pattern = cond[len(attr) + 2 :] + try: + regex = re.compile(pattern) + except re.error as exc: + dead.append((rule["name"], cond, f"invalid regex: {exc}")) + continue + matched = any(regex.search(f) for f in files) + else: + matched = cond[len(attr) + 1 :] in files + if not matched: + dead.append((rule["name"], cond, "matches no tracked file")) + return dead + + +def main() -> int: + dead = check(tracked_files()) + if not dead: + return 0 + print(f"{len(dead)} label condition(s) match nothing:\n", file=sys.stderr) + for name, cond, why in dead: + print(f" [{name}] {cond}\n {why}", file=sys.stderr) + print( + "\nRepoint the condition at the file's new location, or remove it.", + file=sys.stderr, + ) + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) From 823ec22b78b639f6794df3e5ee16ce4429f2c269 Mon Sep 17 00:00:00 2001 From: Rohan Potdar Date: Wed, 19 Aug 2026 17:39:08 -0500 Subject: [PATCH 169/839] [ROCm]: Bump triton 3.7 commit (#52819) Signed-off-by: Rohan138 --- docker/Dockerfile.rocm_base | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index f72573e21b2c..639dadbb74bf 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -1,5 +1,5 @@ ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2.3-complete -ARG TRITON_BRANCH="532137f" # release/internal/3.7.x as of 08/07 +ARG TRITON_BRANCH="f0b55c0" # release/internal/3.7.x as of 08/18 ARG TRITON_REPO="https://github.com/ROCm/triton.git" ARG PYTORCH_BRANCH="6bbd260" # release/2.12 as of 08/02 ARG PYTORCH_REPO="https://github.com/ROCm/pytorch.git" @@ -58,7 +58,7 @@ RUN apt-get update -y \ && curl -sS https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION} \ && python3 --version && python3 -m pip --version -RUN pip install -U packaging 'cmake<4' ninja wheel 'setuptools<80' pybind11 Cython +RUN pip install -U packaging 'cmake<4' ninja wheel 'setuptools<80' 'pybind11<3.1.0' Cython RUN apt-get update && apt-get install -y libjpeg-dev libsox-dev libsox-fmt-all sox && rm -rf /var/lib/apt/lists/* # Install sccache if USE_SCCACHE is enabled (for release builds) From 5bf0dbd6fd7b0e9d0aa744e483037563111ce887 Mon Sep 17 00:00:00 2001 From: Sage Moore Date: Wed, 19 Aug 2026 17:18:47 -0700 Subject: [PATCH 170/839] [Bugfix] vLLM crashes at startup when DeepEP v2 is used with `--enforce-eager` wiht TRTLLM Bf16 (#51824) --- tests/kernels/moe/parallel_utils.py | 1 + tests/kernels/moe/test_deepep_v2_moe.py | 278 +++++++++++------- .../fused_moe/experts/trtllm_bf16_moe.py | 2 +- 3 files changed, 170 insertions(+), 111 deletions(-) diff --git a/tests/kernels/moe/parallel_utils.py b/tests/kernels/moe/parallel_utils.py index b11d0c0ea503..439e624252e1 100644 --- a/tests/kernels/moe/parallel_utils.py +++ b/tests/kernels/moe/parallel_utils.py @@ -267,6 +267,7 @@ def make_deepep_v2_a2a( hidden=v2_args.hidden_size, num_topk=v2_args.num_topk, use_fp8_dispatch=v2_args.use_fp8_dispatch, + allow_hybrid_mode=False, explicitly_destroy=True, ) return DeepEPV2PrepareAndFinalize( diff --git a/tests/kernels/moe/test_deepep_v2_moe.py b/tests/kernels/moe/test_deepep_v2_moe.py index ab609d7eabbb..f1e34d5bd780 100644 --- a/tests/kernels/moe/test_deepep_v2_moe.py +++ b/tests/kernels/moe/test_deepep_v2_moe.py @@ -11,6 +11,7 @@ import torch.distributed from torch.distributed import ProcessGroup +import vllm.model_executor.layers.fused_moe.modular_kernel as mk from tests.kernels.moe.utils import make_dummy_moe_config, make_test_weights from tests.kernels.utils import torch_experts from vllm.config import VllmConfig, set_current_vllm_config @@ -21,6 +22,8 @@ FusedMoEQuantConfig, ) from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel +from vllm.platforms import current_platform +from vllm.utils.flashinfer import has_flashinfer from vllm.utils.import_utils import has_deep_ep_v2 from vllm.utils.torch_utils import set_random_seed from vllm.v1.worker.workspace import init_workspace_manager @@ -356,18 +359,136 @@ def test_deep_ep_v2_moe( ) -def _deep_ep_v2_moe_cudagraph( +EXPERTS_BACKENDS = [ + "flashinfer_trtllm", + "flashinfer_cutlass", + "trtllm_fp8", +] + + +def _make_experts( + experts_backend: str, + config: TestConfig, + moe_config, + num_local_experts: int, + rank: int, + w1_bf16: torch.Tensor, + w2_bf16: torch.Tensor, + test_tensors: TestTensors, +): + e_start = num_local_experts * rank + e_end = e_start + num_local_experts + + if experts_backend != "trtllm_fp8": + from vllm.model_executor.layers.fused_moe.config import ( + FUSED_MOE_UNQUANTIZED_CONFIG, + ) + from vllm.model_executor.layers.fused_moe.oracle.unquantized import ( + backend_to_kernel_cls, + convert_to_unquantized_kernel_format, + map_unquantized_backend, + ) + + torch_combined = torch_experts( + test_tensors.rank_tokens, + w1_bf16, + w2_bf16, + test_tensors.topk_weights, + test_tensors.topk, + ) + backend = map_unquantized_backend(experts_backend) + w1_ep, w2_ep = convert_to_unquantized_kernel_format( + backend, + moe_config, + w1_bf16[e_start:e_end], + w2_bf16[e_start:e_end], + ) + experts_cls = next( + cls + for cls in backend_to_kernel_cls(backend) + if issubclass(cls, mk.FusedMoEExpertsModular) + ) + fused_experts = experts_cls( + moe_config=moe_config, + quant_config=FUSED_MOE_UNQUANTIZED_CONFIG, + ) + return fused_experts, w1_ep, w2_ep, torch_combined, 1e-1, 2e-1 + + from tests.kernels.moe.test_moe_layer import _quantize_fp8_halves + from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import ( + TrtLlmFp8ExpertsModular, + ) + from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( + Fp8MoeBackend, + convert_to_fp8_moe_kernel_format, + ) + + w1_ref = w1_bf16.to(torch.float8_e4m3fn).to(torch.bfloat16) + w2_ref = w2_bf16.to(torch.float8_e4m3fn).to(torch.bfloat16) + + block_shape = [128, 128] + qw = _quantize_fp8_halves(w1_ref, w2_ref, block_shape) + assert qw.w13_weight_scale is not None + assert qw.w2_weight_scale is not None + + reference_topk_weights = test_tensors.topk_weights.to(torch.bfloat16).to( + torch.float32 + ) + torch_combined = torch_experts( + test_tensors.rank_tokens, + qw.w13_weight, + qw.w2_weight, + reference_topk_weights, + test_tensors.topk, + w1_scale=qw.w13_weight_scale, + w2_scale=qw.w2_weight_scale, + quant_dtype=torch.float8_e4m3fn, + block_shape=block_shape, + ) + + class _MockLayer: + weight_block_size = block_shape + + class moe_config: + is_act_and_mul = True + intermediate_size_per_partition = config.n + + class activation: + is_gated = True + + w1_ep, w2_ep, w1_scale_ep, w2_scale_ep = convert_to_fp8_moe_kernel_format( + fp8_backend=Fp8MoeBackend.FLASHINFER_TRTLLM, + layer=_MockLayer(), + w13=qw.w13_weight[e_start:e_end], + w2=qw.w2_weight[e_start:e_end], + w13_scale=qw.w13_weight_scale[e_start:e_end], + w2_scale=qw.w2_weight_scale[e_start:e_end], + w13_input_scale=None, + w2_input_scale=None, + ) + + fused_experts = TrtLlmFp8ExpertsModular( + moe_config=moe_config, + quant_config=FusedMoEQuantConfig.make( + torch.float8_e4m3fn, + block_shape=block_shape, + w1_scale=w1_scale_ep, + w2_scale=w2_scale_ep, + ), + ) + return fused_experts, w1_ep, w2_ep, torch_combined, 6e-2, 6e-2 + + +def _deep_ep_v2_moe_backends( pgi: ProcessGroupInfo, dp_size: int, config: TestConfig, - w1: torch.Tensor, - w2: torch.Tensor, - w1_scale: torch.Tensor | None, - w2_scale: torch.Tensor | None, + use_cudagraph: bool, + experts_backend: str, ): - """Worker function: verify DeepEP v2 + TrtLLM FP8 with do_expand=False.""" import tempfile + from vllm.config import KernelConfig from vllm.distributed import ( init_distributed_environment, initialize_model_parallel, @@ -402,19 +523,10 @@ def _deep_ep_v2_moe_cudagraph( torch.distributed.broadcast(w1_bf16, src=0, group=pg) torch.distributed.broadcast(w2_bf16, src=0, group=pg) - # Round-trip through FP8 before constructing the reference and kernel weights. - w1_fp8 = w1_bf16.to(torch.float8_e4m3fn) - w2_fp8 = w2_bf16.to(torch.float8_e4m3fn) - w1_ref = w1_fp8.to(torch.bfloat16) - w2_ref = w2_fp8.to(torch.bfloat16) - - from vllm.config import KernelConfig - vllm_cfg = VllmConfig() vllm_cfg.kernel_config = KernelConfig(moe_backend="flashinfer_trtllm") with set_current_vllm_config(vllm_cfg): - # Initialize vLLM parallel state (needed by MoERunner layer) temp_file = tempfile.mktemp() init_distributed_environment( world_size=pgi.world_size, @@ -424,77 +536,7 @@ def _deep_ep_v2_moe_cudagraph( backend="nccl", ) initialize_model_parallel(tensor_model_parallel_size=1) - # Mirror production weight processing: quantize, EP-slice, then - # convert to the TrtLLM BlockMajorK format. - from tests.kernels.moe.test_moe_layer import _quantize_fp8_halves - from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import ( - TrtLlmFp8ExpertsModular, - ) - from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( - Fp8MoeBackend, - convert_to_fp8_moe_kernel_format, - ) - block_shape = [128, 128] - qw = _quantize_fp8_halves(w1_ref, w2_ref, block_shape) - assert qw.w13_weight_scale is not None - assert qw.w2_weight_scale is not None - - # Reference MoE using the same blockwise FP8 quantization scheme as - # the production kernel. torch_experts quantizes activations before - # both GEMMs and dequantizes the operands for the reference matmuls. - reference_topk_weights = test_tensors.topk_weights.to(torch.bfloat16).to( - torch.float32 - ) - torch_combined = torch_experts( - test_tensors.rank_tokens, - qw.w13_weight, - qw.w2_weight, - reference_topk_weights, - test_tensors.topk, - w1_scale=qw.w13_weight_scale, - w2_scale=qw.w2_weight_scale, - quant_dtype=torch.float8_e4m3fn, - block_shape=block_shape, - ) - - # EP-slice before format conversion - e_start = num_local_experts * pgi.rank - e_end = e_start + num_local_experts - w1_ep = qw.w13_weight[e_start:e_end] - w2_ep = qw.w2_weight[e_start:e_end] - w1_scale_ep = qw.w13_weight_scale[e_start:e_end] - w2_scale_ep = qw.w2_weight_scale[e_start:e_end] - - # Convert to TrtLLM format (W31 swap + BlockMajorK shuffle) - class _MockLayer: - weight_block_size = block_shape - - class moe_config: - is_act_and_mul = True - intermediate_size_per_partition = config.n - - class activation: - is_gated = True - - w1_ep, w2_ep, w1_scale_ep, w2_scale_ep = convert_to_fp8_moe_kernel_format( - fp8_backend=Fp8MoeBackend.FLASHINFER_TRTLLM, - layer=_MockLayer(), - w13=w1_ep, - w2=w2_ep, - w13_scale=w1_scale_ep, - w2_scale=w2_scale_ep, - w13_input_scale=None, - w2_input_scale=None, - ) - - # Build TrtLLM expert with correct EP params - quant_config = FusedMoEQuantConfig.make( - torch.float8_e4m3fn, - block_shape=block_shape, - w1_scale=w1_scale_ep, - w2_scale=w2_scale_ep, - ) moe_config = make_dummy_moe_config( num_experts=config.num_experts, num_local_experts=num_local_experts, @@ -502,20 +544,33 @@ class activation: hidden_dim=hidden_size, intermediate_size=config.n, ) - moe_parallel_config = dataclasses.replace( - moe_config.moe_parallel_config, - ep_size=pgi.world_size, - ep_rank=pgi.rank, - use_ep=True, - all2all_backend="deepep_v2", - ) moe_config = dataclasses.replace( moe_config, - moe_parallel_config=moe_parallel_config, + moe_parallel_config=dataclasses.replace( + moe_config.moe_parallel_config, + ep_size=pgi.world_size, + ep_rank=pgi.rank, + use_ep=True, + all2all_backend="deepep_v2", + ), ) - fused_experts = TrtLlmFp8ExpertsModular( - moe_config=moe_config, - quant_config=quant_config, + + ( + fused_experts, + w1_ep, + w2_ep, + torch_combined, + atol, + rtol, + ) = _make_experts( + experts_backend, + config, + moe_config, + num_local_experts, + pgi.rank, + w1_bf16, + w2_bf16, + test_tensors, ) v2_args = DeepEPV2Args( @@ -531,7 +586,7 @@ class activation: pgi=pgi, dp_size=dp_size, v2_args=v2_args, - use_cudagraph=True, + use_cudagraph=use_cudagraph, ) mk_kernel = FusedMoEKernel( prepare_finalize=a2a, @@ -552,33 +607,38 @@ class activation: apply_router_weight_on_input=False, ) - torch.testing.assert_close( - torch_combined, - out, - atol=6e-2, - rtol=6e-2, - ) + torch.testing.assert_close(torch_combined, out, atol=atol, rtol=rtol) @pytest.mark.parametrize("m,n,k", [(32, 256, 1024)]) @pytest.mark.parametrize("num_experts", [32]) @pytest.mark.parametrize("topk", [6]) @pytest.mark.parametrize("world_dp_size", [(2, 1)]) +@pytest.mark.parametrize("experts_backend", EXPERTS_BACKENDS) +@pytest.mark.parametrize("use_cudagraph", [True, False]) @multi_gpu_test(num_gpus=2) @requires_deep_ep_v2 -def test_deep_ep_v2_moe_cudagraph( +@pytest.mark.skipif( + not has_flashinfer() or not current_platform.has_device_capability(100), + reason="Requires FlashInfer TRTLLM fused MoE (SM100)", +) +def test_deep_ep_v2_moe_backends( m: int, n: int, k: int, num_experts: int, topk: int, world_dp_size: tuple[int, int], + experts_backend: str, + use_cudagraph: bool, workspace_init, ): set_random_seed(7) world_size, dp_size = world_dp_size config = TestConfig( - dtype=torch.float8_e4m3fn, + dtype=torch.float8_e4m3fn + if experts_backend == "trtllm_fp8" + else torch.bfloat16, topk=topk, m=m, k=k, @@ -588,11 +648,9 @@ def test_deep_ep_v2_moe_cudagraph( parallel_launch( world_size, - _deep_ep_v2_moe_cudagraph, + _deep_ep_v2_moe_backends, dp_size, config, - None, # weights created inside worker - None, - None, - None, + use_cudagraph, + experts_backend, ) diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py index ae6dc11bebd7..4372634be8d1 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py @@ -207,7 +207,7 @@ def apply( gemm1_weights=w1, gemm2_weights=w2, num_experts=global_num_experts, - top_k=self.topk, + top_k=topk_ids.size(1), n_group=None, topk_group=None, intermediate_size=self.intermediate_size_per_partition, From 0a111cca2cddbfccf5b99f61af9c9bbde53f6d9c Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Thu, 20 Aug 2026 03:24:07 +0300 Subject: [PATCH 171/839] =?UTF-8?q?[kv=5Foffload]=20fix(metrics):=20rename?= =?UTF-8?q?=20kv=5Foffload=5Ftiering=5Fblock=5F{queries,hits}=20=E2=86=92?= =?UTF-8?q?=20chunk=20(#52812)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Ronen Schaffer --- vllm/v1/kv_offload/tiering/base.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index 768d7e328ad7..c1e39811e341 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -45,8 +45,8 @@ class TieringOffloadingMetrics: WRITE_TIME = "vllm:kv_offload_tiering_write_time" PROMOTION_JOB_FAILURES = "vllm:kv_offload_tiering_promotion_job_failures" CASCADE_JOB_FAILURES = "vllm:kv_offload_tiering_cascade_job_failures" - BLOCK_QUERIES = "vllm:kv_offload_tiering_block_queries" - BLOCK_HITS = "vllm:kv_offload_tiering_block_hits" + BLOCK_QUERIES = "vllm:kv_offload_tiering_chunk_queries" + BLOCK_HITS = "vllm:kv_offload_tiering_chunk_hits" PRIMARY_WRITE_USAGE_PERC = "vllm:kv_offload_tiering_primary_write_usage_perc" PRIMARY_READ_USAGE_PERC = "vllm:kv_offload_tiering_primary_read_usage_perc" PROMOTION_ALLOCATION_FAILURES = ( From d4f4d3f40fc5350a71777fcb0e5eb8a57bda631f Mon Sep 17 00:00:00 2001 From: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Date: Wed, 19 Aug 2026 17:39:30 -0700 Subject: [PATCH 172/839] [Model Runner V2][Spec Decode] Fix draft logits cache column stride in gumbel_sample (#53017) --- tests/v1/worker/test_gpu_gumbel_sample.py | 69 +++++++++++++++++++ vllm/v1/worker/gpu/sample/gumbel.py | 22 ++++-- .../spec_decode/rejection_sampler_utils.py | 3 +- 3 files changed, 88 insertions(+), 6 deletions(-) diff --git a/tests/v1/worker/test_gpu_gumbel_sample.py b/tests/v1/worker/test_gpu_gumbel_sample.py index f64a3f1f2481..a87c12416eb5 100644 --- a/tests/v1/worker/test_gpu_gumbel_sample.py +++ b/tests/v1/worker/test_gpu_gumbel_sample.py @@ -291,3 +291,72 @@ def test_logits_cache_stores_input_logits_bitwise( assert torch.equal(_float_bits(stored), _float_bits(logits)), ( "cached logits differ from the input logits" ) + + +@pytest.mark.parametrize("extra_cache_cols", [0, 1]) +def test_logits_cache_columns_stay_separate_across_steps(extra_cache_cols: int): + """Each drafting step must land in its own cache column. + + `extra_cache_cols=1` is the shape a draft produces when it adds an + input-only mask/noise row to its embedding table but keeps a full-width + output head: N-wide logits cached into N+1-wide rows. Striding cache + columns by the logits width instead of the cache's own row width leaves + step 0 correct and silently misaligns every step after it. + """ + torch.manual_seed(0) + num_reqs, vocab_size, num_steps = 4, 1031, 3 + idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=DEVICE) + temp = torch.ones(num_reqs, dtype=torch.float32, device=DEVICE) + seed = torch.arange(num_reqs, dtype=torch.int64, device=DEVICE) + pos = torch.arange(num_reqs, dtype=torch.int64, device=DEVICE) + + cache = torch.zeros( + num_reqs, num_steps, vocab_size + extra_cache_cols, device=DEVICE + ) + cols = torch.arange(num_steps, dtype=torch.int32, device=DEVICE) + per_step = [ + torch.randn(num_reqs, vocab_size, device=DEVICE) for _ in range(num_steps) + ] + + for step, logits in enumerate(per_step): + gumbel_sample( + logits, + idx_mapping, + temp, + seed, + pos, + apply_temperature=True, + logits_cache=cache, + logits_cache_col=cols[step], + ) + + for step, logits in enumerate(per_step): + stored = cache[:, step, :vocab_size] + assert torch.equal(_float_bits(stored), _float_bits(logits)), ( + f"step {step} was overwritten by a later step" + ) + # Columns past the sampled width belong to no step and stay untouched. + assert not cache[:, :, vocab_size:].any() + + +def test_logits_cache_narrower_than_logits_is_rejected(): + """A cache too narrow to hold a step would silently drop its tail.""" + num_reqs, vocab_size, num_steps = 2, 64, 3 + logits = torch.randn(num_reqs, vocab_size, device=DEVICE) + idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=DEVICE) + temp = torch.ones(num_reqs, dtype=torch.float32, device=DEVICE) + seed = torch.zeros(num_reqs, dtype=torch.int64, device=DEVICE) + pos = torch.arange(num_reqs, dtype=torch.int64, device=DEVICE) + cache = torch.zeros(num_reqs, num_steps, vocab_size - 1, device=DEVICE) + + with pytest.raises(AssertionError, match="narrower"): + gumbel_sample( + logits, + idx_mapping, + temp, + seed, + pos, + apply_temperature=True, + logits_cache=cache, + logits_cache_col=torch.tensor(0, dtype=torch.int32, device=DEVICE), + ) diff --git a/vllm/v1/worker/gpu/sample/gumbel.py b/vllm/v1/worker/gpu/sample/gumbel.py index 0a930b8183c0..195b2aadf1bf 100644 --- a/vllm/v1/worker/gpu/sample/gumbel.py +++ b/vllm/v1/worker/gpu/sample/gumbel.py @@ -91,8 +91,10 @@ def gumbel_block_argmax( temp_ptr, seeds_ptr, pos_ptr, + # [max_num_reqs, num_cols, vocab_size] logits_cache_ptr, - logits_cache_stride, + logits_cache_stride_0, + logits_cache_stride_1, logits_cache_col_ptr, vocab_size, APPLY_TEMPERATURE: tl.constexpr, @@ -115,8 +117,8 @@ def gumbel_block_argmax( col = tl.load(logits_cache_col_ptr) tl.store( logits_cache_ptr - + req_state_idx * logits_cache_stride - + col * vocab_size + + req_state_idx * logits_cache_stride_0 + + col * logits_cache_stride_1 + block, logits, mask=mask & is_valid_req, @@ -164,8 +166,10 @@ def _gumbel_sample_kernel( local_argmax_stride, local_max_ptr, local_max_stride, + # [max_num_reqs, num_cols, vocab_size] logits_cache_ptr, - logits_cache_stride, + logits_cache_stride_0, + logits_cache_stride_1, logits_cache_col_ptr, logits_ptr, logits_stride, @@ -200,7 +204,8 @@ def _gumbel_sample_kernel( seeds_ptr, pos_ptr, logits_cache_ptr, - logits_cache_stride, + logits_cache_stride_0, + logits_cache_stride_1, logits_cache_col_ptr, vocab_size, APPLY_TEMPERATURE=APPLY_TEMPERATURE, @@ -229,6 +234,12 @@ def gumbel_sample( if logits_cache_col is not None: logits_cache_col = logits_cache_col.contiguous() num_tokens, vocab_size = logits.shape + if logits_cache is not None: + assert logits_cache.size(-1) >= vocab_size, ( + f"draft logits cache vocab dim ({logits_cache.size(-1)}) is narrower " + f"than the sampled logits ({vocab_size}). Cached logits would be " + "truncated." + ) BLOCK_SIZE = 1024 num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE) local_argmax = logits.new_empty(num_tokens, num_blocks, dtype=torch.int64) @@ -242,6 +253,7 @@ def gumbel_sample( local_max.stride(0), logits_cache, logits_cache.stride(0) if logits_cache is not None else 0, + logits_cache.stride(1) if logits_cache is not None else 0, logits_cache_col, logits, logits.stride(0), diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py index 3ea9f35ee8bc..abbf0ee96ef4 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py +++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py @@ -833,7 +833,8 @@ def _resample_kernel( seed_ptr, pos_ptr, None, # logits_cache_ptr - 0, # logits_cache_stride + 0, # logits_cache_stride_0 + 0, # logits_cache_stride_1 None, # logits_cache_col_ptr vocab_size, APPLY_TEMPERATURE=False, From bf2866f8bf5bb20628e2b93835be3c281a9b4ca4 Mon Sep 17 00:00:00 2001 From: Guanyi Chen <939416532@qq.com> Date: Thu, 20 Aug 2026 08:40:35 +0800 Subject: [PATCH 173/839] [KV Connector] Add decode offloading to Mooncake Store consumers (#52466) Signed-off-by: z-zanez Signed-off-by: Guanyi Chen <939416532@qq.com> Co-authored-by: z-zanez --- .../mooncake_store_connector_usage.md | 23 +++ .../unit/test_mooncake_store_scheduler.py | 189 +++++++++++++++++- .../unit/test_mooncake_store_worker.py | 160 ++++++++++++++- .../v1/mooncake/store/connector.py | 10 +- .../kv_connector/v1/mooncake/store/data.py | 15 +- .../v1/mooncake/store/scheduler.py | 75 ++++--- .../kv_connector/v1/mooncake/store/worker.py | 118 ++++++++--- 7 files changed, 518 insertions(+), 72 deletions(-) diff --git a/docs/features/mooncake_store_connector_usage.md b/docs/features/mooncake_store_connector_usage.md index a07716dea627..800b814250eb 100644 --- a/docs/features/mooncake_store_connector_usage.md +++ b/docs/features/mooncake_store_connector_usage.md @@ -136,6 +136,25 @@ vllm serve meta-llama/Llama-3.1-8B-Instruct \ }' ``` +To also offload newly completed decode KV blocks, add the following extra +configuration to the decoder's `MooncakeStoreConnector` entry. This currently +requires the prefill and decode instances to use the same tensor-parallel size +and compatible KV-cache topology. + +When decode processing starts, the consumer checks the block-aligned prompt +prefix and fills any blocks missing from the Store. Subsequent saves append +newly completed decode blocks. This keeps a complete, reusable prefix in the +Store and also covers deployments where prompt KV is delivered directly by +`MooncakeConnector` instead of through the Store. + +```json +{ + "kv_connector_extra_config": { + "save_decode_cache": true + } +} +``` + **Proxy:** A disaggregation proxy is required to route requests between prefiller and decoder nodes. The proxy assigns `do_remote_prefill=True` / `do_remote_decode=True` to coordinate P2P transfer via `MooncakeConnector`. Refer to the [MooncakeConnector usage guide](mooncake_connector_usage.md) for proxy setup details. @@ -230,6 +249,10 @@ Strict isolation requires a Mooncake master started with `--enable_multi_tenants - `enable_cross_layers_blocks` (bool): Enable cross-layer block packing for reduced store operations. Default: `false`. - `lookup_rpc_port` (int): Custom port for the ZMQ lookup RPC socket. Default: `0`. - `cache_prefix` (str): Namespace prepended to every store key. Lets separate deployments share one Mooncake master without polluting each other — instances configured with different prefixes never see each other's cached blocks, even for identical prompts. All instances that should share a prefix cache must use the same value. Default: `""` (no prefix; keys are byte-identical to the unprefixed format). +- `save_decode_cache` (bool): Enable offloading decode tokens' KV cache. A `kv_consumer` does not save during prefill; when decode starts, it fills any missing block-aligned prompt prefix before appending completed decode blocks. Default: `false`. + +Decode offloading uses the existing TP-rank key namespace. Cross-TP sharing is +not yet supported. ## Notes diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py b/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py index 998a8408e651..1f02d189008c 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py @@ -3,6 +3,8 @@ from types import SimpleNamespace +import pytest + from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( LoadSpec, MooncakeStoreWorkerMetadata, @@ -16,10 +18,16 @@ def _make_bare_scheduler( - *, hash_block_size: int = 16, enable_partial_hash_hits: bool = False + *, + hash_block_size: int = 16, + enable_partial_hash_hits: bool = False, + kv_role: str = "kv_both", + save_decode_cache: bool = False, ) -> MooncakeStoreScheduler: scheduler = object.__new__(MooncakeStoreScheduler) - scheduler.kv_role = "kv_both" + scheduler.kv_role = kv_role + scheduler.save_decode_cache = save_decode_cache + scheduler.enable_kv_events = False scheduler.lookup_async = False scheduler.enable_lookup = True scheduler._block_size = 16 @@ -56,6 +64,74 @@ def _make_scheduler_output(*, scheduled_spec_tokens: list[int] | None): ) +def _make_decode_scheduler_output( + *, num_computed_tokens: int, num_scheduled_tokens: int = 1 +) -> SimpleNamespace: + return SimpleNamespace( + finished_req_ids=set(), + preempted_req_ids=set(), + scheduled_new_reqs=[], + scheduled_cached_reqs=SimpleNamespace( + req_ids=["req-0"], + # The block that becomes full was allocated on an earlier step. + new_block_ids=[()], + num_computed_tokens=[num_computed_tokens], + resumed_req_ids=set(), + ), + num_scheduled_tokens={"req-0": num_scheduled_tokens}, + scheduled_spec_decode_tokens={}, + ) + + +def _make_new_scheduler_output() -> SimpleNamespace: + request = SimpleNamespace( + req_id="req-0", + num_computed_tokens=0, + prompt_token_ids=list(range(32)), + prefill_token_ids=None, + block_ids=([0, 1],), + block_hashes=[b"h0", b"h1"], + ) + return SimpleNamespace( + request=request, + finished_req_ids=set(), + preempted_req_ids=set(), + scheduled_new_reqs=[request], + scheduled_cached_reqs=SimpleNamespace( + req_ids=[], + new_block_ids=[], + num_computed_tokens=[], + resumed_req_ids=set(), + ), + num_scheduled_tokens={"req-0": 32}, + scheduled_spec_decode_tokens={}, + ) + + +def test_scheduler_only_tracks_token_ids_for_kv_events(): + for enable_kv_events in (False, True): + scheduler = _make_bare_scheduler() + scheduler.enable_kv_events = enable_kv_events + scheduler._unfinished_requests["req-0"] = ( + _make_new_scheduler_output().request, + ([0, 1],), + ) + + meta = scheduler.build_connector_meta(_make_new_scheduler_output()) + + tracker = scheduler._request_trackers["req-0"] + req_meta = meta.requests[0] + if enable_kv_events: + assert tracker.token_ids == list(range(32)) + assert req_meta.token_ids == list(range(32)) + assert req_meta.token_ids_start == 0 + tracker.token_ids.append(99) + assert req_meta.token_ids == list(range(32)) + else: + assert tracker.token_ids is None + assert req_meta.token_ids is None + + def _make_preemption_scheduler_output(): return SimpleNamespace( finished_req_ids=set(), @@ -103,6 +179,29 @@ def _add_unfinished_request( ) +def _setup_decode_request( + *, + kv_role: str = "kv_consumer", + save_decode_cache: bool = False, + token_len: int = 47, +) -> tuple[MooncakeStoreScheduler, RequestTracker]: + scheduler = _make_bare_scheduler( + kv_role=kv_role, save_decode_cache=save_decode_cache + ) + token_ids = list(range(token_len + 1)) + _add_unfinished_request( + scheduler, + token_ids=token_ids, + block_hashes=[b"h0", b"h1", b"h2"], + prefill_end_tokens=32, + ) + tracker = scheduler._request_trackers["req-0"] + tracker.token_len = token_len + tracker.allocated_block_ids = ([0, 1, 2],) + tracker.token_ids = token_ids[:token_len] + return scheduler, tracker + + def test_cached_request_with_spec_decode_does_not_save_scheduled_drafts(): # Drafts in scheduled_spec_decode_tokens are not appended to all_token_ids # yet, so the tracker's token_len does not advance and num_tokens_to_save @@ -149,6 +248,92 @@ def test_cached_request_without_spec_decode_keeps_current_step_save_overlap(): assert tracker.num_saved_tokens == 48 +@pytest.mark.parametrize("kv_role", ["kv_consumer", "kv_both"]) +def test_decode_tracking_is_skipped_by_default(kv_role): + scheduler, tracker = _setup_decode_request(kv_role=kv_role) + + meta = scheduler.build_connector_meta( + _make_decode_scheduler_output(num_computed_tokens=47) + ) + + assert meta.requests == [] + assert tracker.token_len == 47 + assert tracker.num_saved_tokens == 32 + assert tracker.allocated_block_ids == ([0, 1, 2],) + assert tracker.token_ids == list(range(47)) + + +def test_fresh_consumer_first_decode_save_can_backfill_missing_prompt(): + scheduler = _make_bare_scheduler( + kv_role="kv_consumer", + save_decode_cache=True, + ) + scheduler.enable_kv_events = True + new_output = _make_new_scheduler_output() + request = new_output.request + request.block_ids = ([0, 1, 2],) + request.block_hashes = [b"h0", b"h1", b"h2"] + request.all_token_ids = list(range(48)) + request.num_output_placeholders = 0 + scheduler._unfinished_requests["req-0"] = (request, request.block_ids) + + # The consumer does not save during prefill, so its first decode save must + # still cover the prompt. The worker deduplicates prompt blocks already in + # the Store and fills any that are missing, preserving a complete prefix. + assert scheduler.build_connector_meta(new_output).requests == [] + tracker = scheduler._request_trackers["req-0"] + assert tracker.num_saved_tokens == 0 + + # The first decode step checks/saves the block-aligned prompt. The worker's + # Store dedup turns this into a no-op when the producer already saved it. + [prompt_meta] = scheduler.build_connector_meta( + _make_decode_scheduler_output(num_computed_tokens=32) + ).requests + assert prompt_meta.token_len_chunk == 32 + assert prompt_meta.token_ids_start == 0 + assert prompt_meta.token_ids == list(range(32)) + + for num_computed_tokens in range(33, 48): + meta = scheduler.build_connector_meta( + _make_decode_scheduler_output( + num_computed_tokens=num_computed_tokens, + ) + ) + if num_computed_tokens < 47: + assert meta.requests == [] + + [req_meta] = meta.requests + assert req_meta.token_len_chunk == 48 + assert req_meta.token_ids_start == 32 + assert req_meta.token_ids == list(range(32, 48)) + assert tracker.num_saved_tokens == 48 + + +@pytest.mark.parametrize("token_len, saved_tokens", [(46, 32), (47, 48)]) +def test_consumer_saves_only_full_decode_blocks(token_len, saved_tokens): + scheduler, tracker = _setup_decode_request( + save_decode_cache=True, token_len=token_len + ) + + meta = scheduler.build_connector_meta( + _make_decode_scheduler_output(num_computed_tokens=token_len) + ) + + assert tracker.token_len == token_len + 1 + assert tracker.num_saved_tokens == saved_tokens + if saved_tokens == 32: + assert meta.requests == [] + else: + [req_meta] = meta.requests + assert req_meta.can_save is True + assert req_meta.token_len_chunk == 48 + assert req_meta.block_ids == ([0, 1, 2],) + assert req_meta.token_ids == list(range(32, 48)) + assert req_meta.token_ids_start == 32 + tracker.token_ids.append(999) + assert req_meta.token_ids == list(range(32, 48)) + + def test_preemption_resets_tracker(): scheduler = _make_bare_scheduler() _add_unfinished_request( diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index 50d2ee16149e..39c0f8eb4ade 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -36,7 +36,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.metrics import ( MooncakeStoreConnectorStats, ) -from vllm.v1.core.kv_cache_utils import BlockHash +from vllm.v1.core.kv_cache_utils import BlockHash, maybe_convert_block_hash class _RecordingBlockHashes: @@ -230,6 +230,7 @@ def _make_vllm_config( extra_config: dict[str, object] | None = None, rank: int = 0, decode_context_parallel_size: int = 1, + kv_role: str = "kv_both", ) -> SimpleNamespace: return SimpleNamespace( model_config=_FakeModelConfig(), @@ -240,7 +241,9 @@ def _make_vllm_config( decode_context_parallel_size=decode_context_parallel_size, prefill_context_parallel_size=1, ), - kv_transfer_config=_FakeKVTransferConfig(extra_config=extra_config), + kv_transfer_config=_FakeKVTransferConfig( + kv_role=kv_role, extra_config=extra_config + ), cache_config=SimpleNamespace(block_size=16, num_gpu_blocks=10), kv_events_config=SimpleNamespace(enable_kv_cache_events=False), speculative_config=None, @@ -943,15 +946,20 @@ def test_stale_store_job_cannot_touch_a_reused_request_id(): live = _make_store_req("req-a", [b"a0", b"a1"]) live.store_job_id = 2 thread.add_request(live) + thread._retry_token_ids["req-a"] = (32, list(range(32, 64))) thread.finish_store_job(stale) thread._record_saved(stale, 64) thread._mark_request_skipped_for_pressure(stale) + thread._update_retry_token_ids(stale, False, 0, list(range(32))) + thread._update_retry_token_ids(stale, True, 0, None) assert thread.is_live_store_job(live) assert not thread.is_live_store_job(stale) assert thread._saved_offset.get("req-a") is None assert "req-a" not in thread._skip_store_requests + assert thread._get_retry_token_ids(stale) is None + assert thread._get_retry_token_ids(live) == (32, list(range(32, 64))) def test_store_recving_thread_reports_failed_block_ids(): @@ -1655,6 +1663,39 @@ def test_requester_worker_init_skips_disk_budget_when_offload_disabled( assert w.disk_offload_buffer_budget_bytes is None +def test_save_decode_cache_keeps_transfer_path_enabled(tmp_path, monkeypatch): + store = MagicMock() + store.setup.return_value = 0 + _install_fake_mooncake(monkeypatch, store) + _patch_worker_runtime(monkeypatch) + monkeypatch.setenv( + "MOONCAKE_CONFIG_PATH", + _write_mooncake_config( + tmp_path, + { + "metadata_server": "http://metadata/endpoint", + "protocol": "tcp", + "device_name": "", + "master_server_address": "10.0.0.7:50051", + }, + ), + ) + + w = worker.MooncakeStoreWorker( + _make_vllm_config( + kv_role="kv_consumer", + extra_config={ + "enable_lookup": False, + "save_decode_cache": True, + }, + ), + _make_kv_cache_config(), + ) + + assert w.can_put is True + assert w._capacity_only is False + + def test_requester_worker_init_builds_replicate_config_for_preferred_segment( tmp_path, monkeypatch, @@ -1962,7 +2003,7 @@ def test_store_sending_thread_only_stores_swa_blocks_in_window(): store = MagicMock() store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) store.batch_put_from_multi_buffers.side_effect = ( - lambda keys, addrs, sizes, *_args: [256] * len(keys) + lambda keys, addrs, sizes, *_args: ([256] * len(keys)) ) full_spec = FullAttentionSpec( @@ -2176,6 +2217,85 @@ def test_store_sending_thread_kv_events_use_group_chunk_metadata(): assert swa_event.block_hashes == [maybe_convert_block_hash(BlockHash(hs[3]))] +def _make_event_store_req(token_len: int, token_ids_start: int = 0) -> ReqMeta: + num_blocks = token_len // 16 + return ReqMeta( + req_id="r0", + token_len_chunk=token_len, + block_ids=(list(range(num_blocks)),), + block_hashes=[f"a{i}".encode() for i in range(num_blocks)], + can_save=True, + token_ids=list(range(token_ids_start, token_len)), + token_ids_start=token_ids_start, + ) + + +@pytest.mark.parametrize( + ("saved_offset", "put_step", "exists", "stored_indices", "parent_indices"), + [ + pytest.param(16, 1, [0, 0, 0], [1, 2, 3], [0, 1, 2], id="suffix"), + pytest.param(0, 1, [1, 0, 1, 0], [1, 3], [0, 2], id="dedup-holes"), + pytest.param(0, 2, [0, 0], [0, 2], [None, 1], id="tp-stride"), + ], +) +def test_store_sending_thread_kv_events_use_request_chain_parents( + saved_offset, put_step, exists, stored_indices, parent_indices +): + store = MagicMock() + store.batch_is_exist.return_value = exists + store.batch_put_from_multi_buffers.return_value = [256] * len(stored_indices) + thread = _make_store_sending_thread(store, put_step=put_step) + thread.enable_kv_event = True + + thread._saved_offset["r0"] = saved_offset + _run_store_req(thread, _make_event_store_req(64, saved_offset)) + + events = thread.get_kv_events() + assert [event.block_hashes for event in events] == [ + [maybe_convert_block_hash(BlockHash(f"a{i}".encode()))] for i in stored_indices + ] + assert [event.parent_block_hash for event in events] == [ + (None if i is None else maybe_convert_block_hash(BlockHash(f"a{i}".encode()))) + for i in parent_indices + ] + assert thread._retry_token_ids == {} + + +def test_store_sending_thread_kv_events_retry_without_covered_tokens(): + store = MagicMock() + store.batch_is_exist.return_value = [0, 0] + store.batch_put_from_multi_buffers.return_value = [256, 256] + thread = _make_store_sending_thread(store) + thread.enable_kv_event = True + + _run_store_req(thread, _make_event_store_req(32, 16)) + + retry_event, suffix_event = thread.get_kv_events() + assert retry_event.token_ids == [] + assert suffix_event.token_ids == list(range(16, 32)) + + +def test_store_sending_thread_kv_events_recover_suffix_after_put_failure(): + store = MagicMock() + store.batch_is_exist.side_effect = ([0], [0, 0]) + store.batch_put_from_multi_buffers.side_effect = ([-1], [256, 256]) + thread = _make_store_sending_thread(store) + thread.enable_kv_event = True + + _run_store_req(thread, _make_event_store_req(16)) + + assert thread.get_kv_events() == [] + assert thread._retry_token_ids["r0"] == (0, list(range(16))) + + _run_store_req(thread, _make_event_store_req(32, 16)) + + retry_event, suffix_event = thread.get_kv_events() + assert retry_event.token_ids == list(range(16)) + assert suffix_event.token_ids == list(range(16, 32)) + assert thread._retry_token_ids == {} + assert thread._saved_offset["r0"] == 32 + + def _auto_set_ready_event(*args, **kwargs): """Side effect for mocked thread constructors that auto-sets ready_event.""" for arg in args: @@ -2214,6 +2334,7 @@ def _make_bare_worker( num_gpu_blocks: int = 10, block_size: int = 16, kv_role: str = "kv_both", + save_decode_cache: bool = False, ) -> mooncake_store_worker.MooncakeStoreWorker: """Construct a MooncakeStoreWorker via __new__, bypassing __init__. @@ -2227,14 +2348,17 @@ def _make_bare_worker( worker.store = MagicMock() worker.store.register_buffer.return_value = 0 worker.kv_role = kv_role + worker.can_put = kv_role in ("kv_producer", "kv_both") or save_decode_cache worker._capacity_only = False worker.block_size = block_size worker.tp_rank = 0 worker.enable_kv_events = False + worker.load_async = True worker.kv_send_thread = None worker.kv_recv_threads = [] worker.num_recv_threads = 1 worker.recv_request_queue = queue.Queue() + worker.finished_store_req = set() worker.tp_size = 1 worker.num_kv_head = 1 worker.pp_size = 1 @@ -2743,6 +2867,36 @@ def exists(keys): # --------------------------------------------------------------------------- +@pytest.mark.parametrize("save_decode_cache", [False, True]) +def test_consumer_starts_send_thread_only_when_put_is_enabled(save_decode_cache): + num_blocks = 10 + tensor = torch.zeros(num_blocks, 64, dtype=torch.float16) + + worker = _make_bare_worker( + kv_role="kv_consumer", save_decode_cache=save_decode_cache + ) + _register_with_mocked_threads(worker, {"layer0": tensor}) + assert (worker.kv_send_thread is not None) == save_decode_cache + + +def test_putting_consumer_queues_decode_save(): + w = _make_bare_worker(kv_role="kv_consumer", save_decode_cache=True) + send_thread = MagicMock() + w.kv_send_thread = send_thread + req = _make_store_req("decode-req", [b"h0", b"h1"]) + req.store_job_id = 1 + meta = mooncake_store_worker.MooncakeStoreConnectorMetadata(set(), set()) + meta.add_request(req) + + with patch.object(torch.cuda, "Event") as event_cls: + event = event_cls.return_value + w.get_finished(set(), meta) + + event.record.assert_called_once_with() + send_thread.add_request.assert_called_once_with(req) + assert req.current_event is event + + def test_register_kv_caches_blocks_first_single_segment(): """Blocks-first layout (FlashInfer/MLA): one segment per layer.""" num_blocks = 10 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py index 297fc22ad8a0..e35eb6adb382 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py @@ -141,12 +141,14 @@ def __init__( assert vllm_config.kv_transfer_config is not None assert kv_cache_config is not None, "kv_cache_config is required" self.kv_role = vllm_config.kv_transfer_config.kv_role + extra_config = vllm_config.kv_transfer_config.kv_connector_extra_config + save_decode_cache = extra_config.get("save_decode_cache", False) # Capacity-only: contributes its segment to the store pool but transfers # no KV, so the KV-cache-shape invariants below cannot be reached. - self._capacity_only = self.kv_role == "kv_consumer" and not ( - vllm_config.kv_transfer_config.kv_connector_extra_config.get( - "enable_lookup", True - ) + self._capacity_only = ( + self.kv_role == "kv_consumer" + and not extra_config.get("enable_lookup", True) + and not save_decode_cache ) if not self._capacity_only: self._validate_kv_cache_config(vllm_config, kv_cache_config) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py index 55ffc040cfd1..725da61e797d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py @@ -362,10 +362,11 @@ class ReqMeta: can_save: bool | None = None load_spec: LoadSpec | None = None - is_last_chunk: bool | None = None current_event: torch.cuda.Event | None = None token_ids: list[int] | None = None + # Absolute request offset represented by token_ids[0]. + token_ids_start: int = 0 num_prompt_tokens: int | None = None # Identifies this store job for the engine's lifetime. A request id cannot # serve that purpose: it is reused once a preempted request resumes, so it @@ -384,14 +385,14 @@ def from_request_tracker( load_spec: LoadSpec | None = None, skip_save: bool | None = False, block_hashes: list[BlockHash] | None = None, - is_last_chunk: bool | None = None, ) -> "ReqMeta | None": """Create ReqMeta from a RequestTracker.""" if block_hashes is None: block_hashes = [] input_token_len = tracker.token_len - chunk_boundary = cdiv(tracker.num_saved_tokens + 1, block_size) * block_size + token_ids_start = tracker.num_saved_tokens + chunk_boundary = cdiv(token_ids_start + 1, block_size) * block_size num_tokens_to_save = input_token_len // block_size * block_size skip_save = skip_save or num_tokens_to_save < chunk_boundary @@ -408,8 +409,10 @@ def from_request_tracker( tracker.num_saved_tokens = num_tokens_to_save token_ids = None - if tracker.token_ids: - token_ids = tracker.token_ids + if tracker.token_ids and not skip_save: + # Scheduler tracking continues while this job is handled by an + # asynchronous worker, so metadata must own a stable snapshot. + token_ids = tracker.token_ids[token_ids_start:num_tokens_to_save] if load_spec is not None and load_spec.can_load: logger.debug( @@ -433,8 +436,8 @@ def from_request_tracker( can_save=not skip_save, load_spec=load_spec, block_hashes=block_hashes, - is_last_chunk=is_last_chunk, token_ids=token_ids, + token_ids_start=token_ids_start, num_prompt_tokens=tracker.prefill_end_tokens, ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py index f583b3fcfdf1..c26e25fc5a7f 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py @@ -63,6 +63,11 @@ def __init__( self.lookup_async = kvc_extra_config.get("lookup_async", False) # Skips lookup CPU cost on instances that never load KV from the store. self.enable_lookup = kvc_extra_config.get("enable_lookup", True) + self.save_decode_cache = kvc_extra_config.get("save_decode_cache", False) + kv_event_config = vllm_config.kv_events_config + self.enable_kv_events = bool( + kv_event_config and kv_event_config.enable_kv_cache_events + ) self.client = LookupKeyClient(vllm_config) # Align with the engine's own scheduler_block_size and hash_block_size. @@ -182,7 +187,8 @@ def build_connector_meta( self, scheduler_output: SchedulerOutput ) -> KVConnectorMetadata: """Build connector metadata for this scheduler step.""" - force_skip_save = self.kv_role == "kv_consumer" + is_consumer = self.kv_role == "kv_consumer" + can_process_cached = not is_consumer or self.save_decode_cache for finished_req_id in scheduler_output.finished_req_ids: self.client.discard(finished_req_id) @@ -227,37 +233,38 @@ def build_connector_meta( token_len=num_tokens_to_compute, allocated_block_ids=unfolded_block_ids, num_saved_tokens=0, - token_ids=prefill_tokens[:num_tokens_to_compute], + token_ids=( + prefill_tokens[:num_tokens_to_compute] + if self.enable_kv_events + else None + ), prefill_end_tokens=len(prefill_tokens), ) self._request_trackers[request.req_id] = request_tracker - last_chunk_tokens_num = ( - len(prefill_tokens) // self._block_size * self._block_size - ) - req_meta = ReqMeta.from_request_tracker( request_tracker, self._block_size, load_spec=load_spec, - skip_save=force_skip_save, + # A consumer may write decode KV without becoming a prefill + # producer. Loads are still carried by the same metadata. + skip_save=is_consumer, block_hashes=request_real.block_hashes, - is_last_chunk=(request_tracker.token_len >= last_chunk_tokens_num), ) if req_meta is not None: meta.add_request(req_meta) # Handle cached (running, or MRV1 resumed-from-preemption) requests cached_reqs = scheduler_output.scheduled_cached_reqs - if not force_skip_save: + if can_process_cached: for i, req_id in enumerate(cached_reqs.req_ids): new_block_ids = cached_reqs.new_block_ids[i] - if not new_block_ids: - continue req_meta = None if req_id in cached_reqs.resumed_req_ids: # Resumed after preemption + if not new_block_ids: + continue if isinstance(new_block_ids, tuple): new_block_ids = tuple(b.copy() for b in new_block_ids) else: @@ -277,27 +284,33 @@ def build_connector_meta( token_len=num_tokens_to_compute, allocated_block_ids=new_block_ids, num_saved_tokens=0, - token_ids=prefill_tokens[:num_tokens_to_compute].copy(), + token_ids=( + prefill_tokens[:num_tokens_to_compute] + if self.enable_kv_events + else None + ), prefill_end_tokens=len(prefill_tokens), ) self._request_trackers[req_id] = request_tracker - last_chunk_tokens_num = ( - len(prefill_tokens) // self._block_size * self._block_size - ) req_meta = ReqMeta.from_request_tracker( request_tracker, self._block_size, load_spec=load_spec, - skip_save=force_skip_save, + skip_save=is_consumer, block_hashes=request_real.block_hashes, - is_last_chunk=( - request_tracker.token_len >= last_chunk_tokens_num - ), ) else: # Decode/chunked request request_tracker = self._request_trackers[req_id] + num_computed_token = cached_reqs.num_computed_tokens[i] + # Use the tracker's snapshot of the prefill range so resumed + # requests keep saving past the original prompt boundary. + prefill_end = request_tracker.prefill_end_tokens + is_decode = num_computed_token >= prefill_end + if is_decode and not self.save_decode_cache: + continue + num_new_tokens = scheduler_output.num_scheduled_tokens[req_id] req_tuple = self._unfinished_requests.get(req_id) if req_tuple: @@ -307,30 +320,26 @@ def build_connector_meta( num_current_tokens : num_current_tokens + num_new_tokens ] request_tracker.token_len += len(new_token_ids) + if request_tracker.token_ids is not None: + request_tracker.token_ids.extend(new_token_ids) else: raise ValueError( f"Request {req_id} is not in _unfinished_requests" ) - num_computed_token = cached_reqs.num_computed_tokens[i] - # Use the tracker's snapshot of the prefill range so resumed - # requests keep saving past the original prompt boundary. - prefill_end = request_tracker.prefill_end_tokens - if num_computed_token >= prefill_end: + # A block is usually allocated before the step that fills + # it, so reaching a save boundary does not imply that this + # step has new block ids. + if new_block_ids: + request_tracker.update(new_block_ids) + if is_consumer and not is_decode: continue - request_tracker.update(new_block_ids) - last_chunk_tokens_num = ( - prefill_end // self._block_size * self._block_size - ) req_meta = ReqMeta.from_request_tracker( request_tracker, self._block_size, load_spec=None, - skip_save=force_skip_save, + skip_save=False, block_hashes=unfinished_req.block_hashes, - is_last_chunk=( - request_tracker.token_len >= last_chunk_tokens_num - ), ) if req_meta is not None: @@ -370,7 +379,7 @@ def build_connector_meta( # emit an offload-only ReqMeta (token_len_chunk=0 skips the normal # save; can_save=True takes the normal enqueue path). step_partial_tails = getattr(scheduler_output, "partial_tail_offloads", None) - if step_partial_tails and not force_skip_save: + if step_partial_tails and not is_consumer: pending = dict(step_partial_tails) for req_meta in meta.requests: if req_meta.can_save: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index 83c6d16e1b9e..efc93f7d9423 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -530,6 +530,9 @@ def __init__( # Per-request high-water mark of tokens actually persisted; the next # batch resumes here, so pressure-skipped or failed ranges are retried. self._saved_offset: dict[str, int] = {} + # Retained only after a failed store so retry events can recover the + # token suffix without full snapshots on the normal path. + self._retry_token_ids: dict[str, tuple[int, list[int]]] = {} def add_request(self, request: ReqMeta) -> None: # Register before enqueueing so a job is never picked up unledgered. @@ -552,6 +555,7 @@ def delete_finished_stored_request(self, req_id: str): del self.stored_requests[req_id] self._skip_store_requests.discard(req_id) self._saved_offset.pop(req_id, None) + self._retry_token_ids.pop(req_id, None) def finish_store_job(self, req_meta: ReqMeta) -> None: """Retire a job from the ledger and report its blocks as no longer read. @@ -585,6 +589,36 @@ def _record_saved(self, req_meta: ReqMeta, token_len: int) -> None: if req_meta.store_job_id in self.stored_requests.get(req_meta.req_id, ()): self._saved_offset[req_meta.req_id] = token_len + def _get_retry_token_ids(self, req_meta: ReqMeta) -> tuple[int, list[int]] | None: + """Return retry state only if this store job is still live.""" + with self.done_task_lock: + if req_meta.store_job_id not in self.stored_requests.get( + req_meta.req_id, () + ): + return None + return self._retry_token_ids.get(req_meta.req_id) + + def _update_retry_token_ids( + self, + req_meta: ReqMeta, + save_completed: bool, + token_ids_start: int, + event_token_ids: list[int] | None, + ) -> None: + """Update retry state without letting a stale job touch a reused ID.""" + with self.done_task_lock: + if req_meta.store_job_id not in self.stored_requests.get( + req_meta.req_id, () + ): + return + if save_completed: + self._retry_token_ids.pop(req_meta.req_id, None) + elif event_token_ids is not None: + self._retry_token_ids[req_meta.req_id] = ( + token_ids_start, + event_token_ids, + ) + def _should_skip_request(self, req_id: str) -> bool: with self.done_task_lock: return self._store_pressure_active and req_id in self._skip_store_requests @@ -785,18 +819,30 @@ def _maybe_offload_partial_tail(self, req_meta: ReqMeta) -> bool: def _handle_request(self, req_meta: ReqMeta): # The single `finally` is the only way out, so the scheduler releases # this job's GPU block references however the job ends. + save_completed = False + token_len = 0 + req_id = req_meta.req_id + event_token_ids = req_meta.token_ids + token_ids_start = req_meta.token_ids_start try: # Cache hits are always a multiple of ``lcm_block_size`` tokens, # which is also ``store_mask``'s precondition. lcm_block_size = self.coord.lcm_block_size token_len = req_meta.token_len_chunk // lcm_block_size * lcm_block_size block_ids_per_group = req_meta.block_ids - req_id = req_meta.req_id current_event = req_meta.current_event if not self.is_live_store_job(req_meta): return + if self.enable_kv_event: + retry_token_ids = self._get_retry_token_ids(req_meta) + if retry_token_ids is not None and event_token_ids is not None: + retry_start, retry_ids = retry_token_ids + if retry_start + len(retry_ids) == token_ids_start: + event_token_ids = retry_ids + event_token_ids + token_ids_start = retry_start + if self._should_skip_request(req_id): logger.debug( "Skipping Mooncake store for request %s while CPU/disk " @@ -852,6 +898,7 @@ def _handle_request(self, req_meta: ReqMeta): if not keys: self._record_saved(req_meta, token_len) + save_completed = True return # Check which blocks already exist (dedup) @@ -878,6 +925,7 @@ def _handle_request(self, req_meta: ReqMeta): if not missing_indices: self._record_saved(req_meta, token_len) + save_completed = True return if len(missing_indices) != len(keys): @@ -927,12 +975,11 @@ def _handle_request(self, req_meta: ReqMeta): addrs.extend(group_addrs) sizes.extend(group_sizes) - # parent_block_hash chains live within a group, not across. if self.enable_kv_event: - prev_key_per_group: dict[int, Any] = {} new_block_hashes = [ maybe_convert_block_hash(bh) for bh in kv_event_block_hashes ] + token_ids_end = token_ids_start + len(event_token_ids or ()) for idx, (s, e, g_idx) in enumerate( zip(starts, ends, group_indices, strict=True) @@ -940,13 +987,24 @@ def _handle_request(self, req_meta: ReqMeta): db = self.token_databases[g_idx] if self.enable_kv_event: token_ids = ( - req_meta.token_ids[s:e] - if req_meta.token_ids is not None - else None + event_token_ids[s - token_ids_start : e - token_ids_start] + if event_token_ids is not None + and token_ids_start <= s + and e <= token_ids_end + else [] ) stored_event = BlockStored( block_hashes=[new_block_hashes[idx]], - parent_block_hash=prev_key_per_group.get(g_idx), + # Derive the direct predecessor from the unfiltered + # request chain. Adjacent PUTs need not be adjacent in + # that chain after Store dedup, masks, or TP striding. + parent_block_hash=( + maybe_convert_block_hash( + req_meta.block_hashes[s // db.hash_block_size - 1] + ) + if s > 0 + else None + ), token_ids=token_ids, block_size=db.block_size, lora_id=None, @@ -955,7 +1013,6 @@ def _handle_request(self, req_meta: ReqMeta): group_idx=g_idx, ) stored_events.append(stored_event) - prev_key_per_group[g_idx] = new_block_hashes[idx] if current_event is not None: current_event.synchronize() @@ -984,6 +1041,13 @@ def _handle_request(self, req_meta: ReqMeta): ) if failed: failed_codes = set(res[i] for i in failed) + if self.enable_kv_event: + failed_indices = set(failed) + stored_events = [ + event + for i, event in enumerate(stored_events) + if i not in failed_indices + ] logger.warning( "batch_put failed: %d/%d keys failed " "(codes=%s, batch_bytes=%d, num_keys=%d), " @@ -1008,6 +1072,7 @@ def _handle_request(self, req_meta: ReqMeta): ) else: self._record_saved(req_meta, token_len) + save_completed = True if self._clear_store_pressure(): logger.info( "Mooncake CPU/disk offloading pressure cleared " @@ -1023,10 +1088,18 @@ def _handle_request(self, req_meta: ReqMeta): num_failed_keys=len(keys), ) logger.error("Failed to put key %s, error: %s", keys, e) + stored_events.clear() if self.enable_kv_event and stored_events: self.update_kv_event(stored_events) finally: + if self.enable_kv_event and token_len: + self._update_retry_token_ids( + req_meta, + save_completed, + token_ids_start, + event_token_ids, + ) self.finish_store_job(req_meta) self.request_queue.task_done() @@ -1274,15 +1347,19 @@ def __init__( self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_size > 1 else 0 assert vllm_config.kv_transfer_config is not None - self.kv_role = vllm_config.kv_transfer_config.kv_role - self.load_async = vllm_config.kv_transfer_config.kv_connector_extra_config.get( - "load_async", True + kv_role = vllm_config.kv_transfer_config.kv_role + assert kv_role is not None + self.kv_role = kv_role + extra_config = vllm_config.kv_transfer_config.kv_connector_extra_config + self.can_put = self.kv_role in ("kv_producer", "kv_both") or ( + extra_config.get("save_decode_cache", False) ) + self.load_async = extra_config.get("load_async", True) # Mirrors MooncakeStoreConnector._capacity_only. - self._capacity_only = self.kv_role == "kv_consumer" and not ( - vllm_config.kv_transfer_config.kv_connector_extra_config.get( - "enable_lookup", True - ) + self._capacity_only = ( + self.kv_role == "kv_consumer" + and not extra_config.get("enable_lookup", True) + and not self.can_put ) self.cache_config = vllm_config.cache_config self.block_size, self.hash_block_size = resolve_kv_cache_block_sizes( @@ -1294,11 +1371,6 @@ def __init__( # Initialize MooncakeDistributedStore with its own TransferEngine store_config = MooncakeStoreConfig.load_from_config() - extra_config = ( - vllm_config.kv_transfer_config.kv_connector_extra_config - if vllm_config.kv_transfer_config - else {} - ) self.store = MooncakeDistributedStore() local_ip = get_ip() local_hostname = rdma_utils.get_requester_local_hostname(local_ip) @@ -1617,7 +1689,7 @@ def _repr_tensor(v: torch.Tensor | list[torch.Tensor]) -> torch.Tensor: db.set_block_len(block_lens) # Start transfer threads - if self.kv_role in ["kv_producer", "kv_both"]: + if self.can_put: ready_event_sending = threading.Event() self.kv_send_thread = KVCacheStoreSendingThread( self.store, @@ -1700,7 +1772,7 @@ def get_finished( assert self.load_async, "load_async must be True for better performance." # Issue stores with CUDA event synchronization. - if self.kv_role in ["kv_producer", "kv_both"]: + if self.can_put: current_event = None for request in meta.requests: if request.can_save: @@ -1714,8 +1786,6 @@ def get_finished( request.current_event = current_event assert self.kv_send_thread is not None self.kv_send_thread.add_request(request) - - if self.kv_role in ["kv_producer", "kv_both"]: self._close_ended_store_requests(finished_req_ids, meta) # Blocks read by a store job are released by the scheduler when the job From 58e5ee0158b6a264c3506f00480e108a34b33ee3 Mon Sep 17 00:00:00 2001 From: Summer Yang Date: Wed, 19 Aug 2026 18:04:11 -0700 Subject: [PATCH 174/839] [refactor] consolidate cp attn ops (#52839) Signed-off-by: Summer Yang Co-authored-by: Woosuk Kwon --- tests/distributed/test_dcp_a2a.py | 34 +- .../test_dcp_direct_a2a_lse_reduce.py | 81 +- tests/v1/attention/test_flashinfer_mla_dcp.py | 2 +- .../v1/attention/test_indexer_dcp_localize.py | 2 +- .../layers/attention/mla_attention.py | 10 +- .../layers/attention/sparse_mla_attention.py | 2 +- .../layers/sparse_attn_indexer.py | 2 +- vllm/models/deepseek_v32/attention.py | 8 +- vllm/models/kimi_k3/nvidia/mla.py | 2 +- vllm/v1/attention/backend.py | 2 +- vllm/v1/attention/backends/flash_attn.py | 6 +- vllm/v1/attention/backends/flashinfer.py | 6 +- vllm/v1/attention/ops/common.py | 299 ---- vllm/v1/attention/ops/cp_common.py | 154 ++ vllm/v1/attention/ops/dcp.py | 1368 +++++++++++++++++ vllm/v1/attention/ops/dcp_alltoall.py | 470 ------ vllm/v1/attention/ops/dcp_utils.py | 740 --------- .../attention => v1/attention/ops}/pcp.py | 0 18 files changed, 1603 insertions(+), 1585 deletions(-) create mode 100644 vllm/v1/attention/ops/cp_common.py create mode 100644 vllm/v1/attention/ops/dcp.py delete mode 100644 vllm/v1/attention/ops/dcp_alltoall.py delete mode 100644 vllm/v1/attention/ops/dcp_utils.py rename vllm/{model_executor/layers/attention => v1/attention/ops}/pcp.py (100%) diff --git a/tests/distributed/test_dcp_a2a.py b/tests/distributed/test_dcp_a2a.py index dfd291742c07..73b6e4a18b11 100644 --- a/tests/distributed/test_dcp_a2a.py +++ b/tests/distributed/test_dcp_a2a.py @@ -44,7 +44,7 @@ def _packed_a2a_reference( h_per_rank: int, is_lse_base_on_e: bool, ) -> tuple[torch.Tensor, torch.Tensor]: - from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine + from vllm.v1.attention.ops.dcp import _lse_weighted_combine B, _H, D = cp_attn_out.shape outputs = ( @@ -157,13 +157,13 @@ class TestLSEWeightedCombine: def test_importable(self): """Verify _lse_weighted_combine is importable.""" - from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine + from vllm.v1.attention.ops.dcp import _lse_weighted_combine assert callable(_lse_weighted_combine) def test_single_rank(self): """Single rank: output unchanged.""" - from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine + from vllm.v1.attention.ops.dcp import _lse_weighted_combine # N=1, B=2, H=4, D=8 outputs = torch.randn(1, 2, 4, 8) @@ -176,7 +176,7 @@ def test_single_rank(self): def test_equal_lse(self): """Equal LSE values: outputs averaged equally.""" - from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine + from vllm.v1.attention.ops.dcp import _lse_weighted_combine _N, B, H, D = 2, 1, 1, 4 outputs = torch.tensor( @@ -200,7 +200,7 @@ def test_equal_lse(self): def test_dominant_rank(self): """Different LSE values: larger LSE gets more weight.""" - from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine + from vllm.v1.attention.ops.dcp import _lse_weighted_combine B, H, D = 1, 1, 2 outputs = torch.tensor( @@ -222,7 +222,7 @@ def test_dominant_rank(self): torch.testing.assert_close(result, outputs[1], atol=1e-5, rtol=1e-5) def test_empty_shard_ignores_undefined_output(self): - from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine + from vllm.v1.attention.ops.dcp import _lse_weighted_combine outputs = torch.tensor([[[[float("nan")]]], [[[3.0]]]]) lses = torch.tensor([[[-float("inf")]], [[0.0]]]) @@ -232,7 +232,7 @@ def test_empty_shard_ignores_undefined_output(self): torch.testing.assert_close(result, outputs[1]) def test_ag_rs_masks_empty_shard_and_padded_lse(self, monkeypatch): - import vllm.v1.attention.ops.common as common + import vllm.v1.attention.ops.dcp as dcp class FakeGroup: world_size = 2 @@ -243,7 +243,7 @@ def all_gather(self, tensor, dim): return torch.cat((tensor, tensor), dim=dim) monkeypatch.setattr( - common, + dcp, "correct_attn_out", lambda output, lses, *args, **kwargs: (output, lses[0]), ) @@ -252,7 +252,7 @@ def all_gather(self, tensor, dim): seq_lens = torch.tensor([0, 2], dtype=torch.int32) query_start_loc = torch.tensor([0, 1, 5], dtype=torch.int32) - _, masked_lse = common._cp_lse_common( + _, masked_lse = dcp._cp_lse_common( output, lse, FakeGroup(), @@ -266,7 +266,7 @@ def all_gather(self, tensor, dim): def test_mathematically_correct(self): """Verify mathematical correctness of LSE combination.""" - from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine + from vllm.v1.attention.ops.dcp import _lse_weighted_combine outputs = torch.tensor( [ @@ -291,7 +291,7 @@ def test_mathematically_correct(self): def test_return_lse(self): """return_lse=True returns global LSE (logsumexp of inputs).""" - from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine + from vllm.v1.attention.ops.dcp import _lse_weighted_combine B, H, D = 1, 1, 2 outputs = torch.tensor( @@ -317,7 +317,7 @@ def test_return_lse(self): def test_base2_return_lse(self): """Base-2 LSE mode returns log2-sum-exp2 global LSE.""" - from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine + from vllm.v1.attention.ops.dcp import _lse_weighted_combine outputs = torch.tensor( [ @@ -354,7 +354,7 @@ def test_base2_return_lse(self): def test_lse_pack_dim(self): """Packed A2A stores one fp32 LSE in output-dtype lanes.""" - from vllm.v1.attention.ops.dcp_alltoall import _dcp_a2a_lse_pack_dim + from vllm.v1.attention.ops.dcp import _dcp_a2a_lse_pack_dim assert _dcp_a2a_lse_pack_dim(torch.bfloat16) == 2 assert _dcp_a2a_lse_pack_dim(torch.float16) == 2 @@ -374,7 +374,7 @@ def test_pack_unpack_combine_matches_reference( return_lse: bool, is_lse_base_on_e: bool, ): - from vllm.v1.attention.ops.dcp_alltoall import ( + from vllm.v1.attention.ops.dcp import ( _dcp_a2a_lse_pack_dim, _dcp_a2a_pack_send, _dcp_a2a_unpack_combine, @@ -421,7 +421,7 @@ def test_pack_unpack_combine_matches_reference( torch.accelerator.device_count() < 1, reason="CUDA is required." ) def test_empty_seq_lens_ignore_undefined_output(self): - from vllm.v1.attention.ops.dcp_alltoall import ( + from vllm.v1.attention.ops.dcp import ( _dcp_a2a_lse_pack_dim, _dcp_a2a_pack_send, _dcp_a2a_unpack_combine, @@ -505,7 +505,7 @@ def _distributed_packed_a2a_worker(env: dict[str, str]) -> None: init_workspace_manager(torch.device(f"cuda:{local_rank}")) try: - from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce + from vllm.v1.attention.ops.dcp import dcp_a2a_lse_reduce dtype = _dtype_from_name(env["TEST_DTYPE"]) return_lse = env["RETURN_LSE"] == "1" @@ -555,7 +555,7 @@ def _distributed_packed_a2a_worker(env: dict[str, str]) -> None: [t[:, rank * h_per_rank : (rank + 1) * h_per_rank] for t in gathered_lse], dim=0, ) - from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine + from vllm.v1.attention.ops.dcp import _lse_weighted_combine expected_out, expected_lse = _lse_weighted_combine( outputs, diff --git a/tests/distributed/test_dcp_direct_a2a_lse_reduce.py b/tests/distributed/test_dcp_direct_a2a_lse_reduce.py index 13561bac7848..e5764954bd53 100644 --- a/tests/distributed/test_dcp_direct_a2a_lse_reduce.py +++ b/tests/distributed/test_dcp_direct_a2a_lse_reduce.py @@ -10,7 +10,8 @@ import torch import torch.distributed as dist -import vllm.v1.attention.ops.dcp_utils as dcp_utils +import vllm.v1.attention.ops.cp_common as cp_common +import vllm.v1.attention.ops.dcp as dcp from vllm.utils.network_utils import get_open_port from vllm.utils.system_utils import update_environment_variables @@ -142,17 +143,17 @@ def rank(self) -> int: class TestDirectDCPGating: def test_env_disabled_returns_none(self, monkeypatch): monkeypatch.setenv("VLLM_USE_DIRECT_DCP_A2A", "0") - dcp_utils.get_direct_dcp_a2a_workspace.cache_clear() - workspace = dcp_utils.get_direct_dcp_a2a_workspace( + dcp.get_direct_dcp_a2a_workspace.cache_clear() + workspace = dcp.get_direct_dcp_a2a_workspace( _FakeGroupCoordinator(), torch.device("cpu"), 16, 2, 32, torch.bfloat16, 1 ) assert workspace is None def test_forced_with_unsupported_dtype_raises(self, monkeypatch): monkeypatch.setenv("VLLM_USE_DIRECT_DCP_A2A", "1") - dcp_utils.get_direct_dcp_a2a_workspace.cache_clear() + dcp.get_direct_dcp_a2a_workspace.cache_clear() with pytest.raises(ValueError, match="does not support"): - dcp_utils.get_direct_dcp_a2a_workspace( + dcp.get_direct_dcp_a2a_workspace( _FakeGroupCoordinator(), torch.device("cpu"), 16, @@ -164,14 +165,14 @@ def test_forced_with_unsupported_dtype_raises(self, monkeypatch): def test_zero_ubatches_raises(self): with pytest.raises(ValueError, match="ubatch"): - dcp_utils.DirectDCPA2AWorkspace( + dcp.DirectDCPA2AWorkspace( None, torch.device("cpu"), 16, 2, 32, torch.bfloat16, num_ubatches=0 ) def test_auto_with_unsupported_dtype_returns_none(self, monkeypatch): monkeypatch.delenv("VLLM_USE_DIRECT_DCP_A2A", raising=False) - dcp_utils.get_direct_dcp_a2a_workspace.cache_clear() - workspace = dcp_utils.get_direct_dcp_a2a_workspace( + dcp.get_direct_dcp_a2a_workspace.cache_clear() + workspace = dcp.get_direct_dcp_a2a_workspace( _FakeGroupCoordinator(), torch.device("cpu"), 16, 2, 32, torch.float32, 1 ) assert workspace is None @@ -179,8 +180,8 @@ def test_auto_with_unsupported_dtype_returns_none(self, monkeypatch): def test_q_gather_env_disabled_returns_none(self, monkeypatch): monkeypatch.setenv("VLLM_USE_DIRECT_DCP_Q_GATHER", "0") monkeypatch.setenv("VLLM_USE_DIRECT_DCP_A2A", "1") - dcp_utils.get_direct_dcp_q_gather_workspace.cache_clear() - workspace = dcp_utils.get_direct_dcp_q_gather_workspace( + dcp.get_direct_dcp_q_gather_workspace.cache_clear() + workspace = dcp.get_direct_dcp_q_gather_workspace( _FakeGroupCoordinator(), torch.device("cpu"), 16, @@ -194,17 +195,17 @@ def test_q_gather_env_disabled_returns_none(self, monkeypatch): def test_q_gather_flag_is_independent(self, monkeypatch): monkeypatch.setenv("VLLM_USE_DIRECT_DCP_Q_GATHER", "1") monkeypatch.setenv("VLLM_USE_DIRECT_DCP_A2A", "0") - monkeypatch.setattr(dcp_utils, "_symm_mem_spans_group", lambda group: True) - dcp_utils.get_direct_dcp_q_gather_workspace.cache_clear() + monkeypatch.setattr(cp_common, "_symm_mem_spans_group", lambda group: True) + dcp.get_direct_dcp_q_gather_workspace.cache_clear() workspace = object() init_workspace = MagicMock(return_value=workspace) monkeypatch.setattr( - dcp_utils, + dcp, "DirectDCPQGatherWorkspace", init_workspace, ) - result = dcp_utils.get_direct_dcp_q_gather_workspace( + result = dcp.get_direct_dcp_q_gather_workspace( _FakeGroupCoordinator(), torch.device("cpu"), 16, @@ -220,8 +221,8 @@ def test_q_gather_flag_is_independent(self, monkeypatch): def test_kv_gather_env_disabled_returns_none(self, monkeypatch): monkeypatch.setenv("VLLM_USE_DIRECT_DCP_KV_GATHER", "0") monkeypatch.setenv("VLLM_USE_DIRECT_DCP_A2A", "1") - dcp_utils.get_direct_dcp_kv_gather_workspace.cache_clear() - workspace = dcp_utils.get_direct_dcp_kv_gather_workspace( + dcp.get_direct_dcp_kv_gather_workspace.cache_clear() + workspace = dcp.get_direct_dcp_kv_gather_workspace( _FakeGroupCoordinator(), torch.device("cpu"), 64, 576, torch.bfloat16, 1 ) assert workspace is None @@ -229,17 +230,17 @@ def test_kv_gather_env_disabled_returns_none(self, monkeypatch): def test_kv_gather_flag_is_independent(self, monkeypatch): monkeypatch.setenv("VLLM_USE_DIRECT_DCP_KV_GATHER", "1") monkeypatch.setenv("VLLM_USE_DIRECT_DCP_A2A", "0") - monkeypatch.setattr(dcp_utils, "_symm_mem_spans_group", lambda group: True) - dcp_utils.get_direct_dcp_kv_gather_workspace.cache_clear() + monkeypatch.setattr(cp_common, "_symm_mem_spans_group", lambda group: True) + dcp.get_direct_dcp_kv_gather_workspace.cache_clear() workspace = object() init_workspace = MagicMock(return_value=workspace) monkeypatch.setattr( - dcp_utils, + dcp, "DirectDCPKVGatherWorkspace", init_workspace, ) - result = dcp_utils.get_direct_dcp_kv_gather_workspace( + result = dcp.get_direct_dcp_kv_gather_workspace( _FakeGroupCoordinator(), torch.device("cpu"), 64, 576, torch.bfloat16, 1 ) @@ -267,9 +268,9 @@ def test_gather_requires_multicast( factory_name, factory_args, ): - factory = getattr(dcp_utils, factory_name) + factory = getattr(dcp, factory_name) monkeypatch.setenv(flag_name, "1") - monkeypatch.setattr(dcp_utils, "_symm_mem_spans_group", lambda group: False) + monkeypatch.setattr(cp_common, "_symm_mem_spans_group", lambda group: False) factory.cache_clear() assert ( @@ -283,25 +284,25 @@ def test_gather_requires_multicast( def test_kv_gather_rejects_invalid_workspace_geometry(self): with pytest.raises(ValueError, match="ubatch"): - dcp_utils.DirectDCPKVGatherWorkspace( + dcp.DirectDCPKVGatherWorkspace( None, torch.device("cpu"), 64, 576, num_ubatches=0 ) with pytest.raises(ValueError, match="divide evenly"): - dcp_utils.DirectDCPKVGatherWorkspace( + dcp.DirectDCPKVGatherWorkspace( _FakeProcessGroup(), torch.device("cpu"), 63, 576 ) with pytest.raises(ValueError, match="16-byte"): - dcp_utils.DirectDCPKVGatherWorkspace( + dcp.DirectDCPKVGatherWorkspace( _FakeProcessGroup(), torch.device("cpu"), 64, 3 ) def test_q_gather_rejects_invalid_workspace_geometry(self): with pytest.raises(ValueError, match="ubatch"): - dcp_utils.DirectDCPQGatherWorkspace( + dcp.DirectDCPQGatherWorkspace( None, torch.device("cpu"), 16, 2, 32, num_ubatches=0 ) with pytest.raises(ValueError, match="padded heads"): - dcp_utils.DirectDCPQGatherWorkspace( + dcp.DirectDCPQGatherWorkspace( _FakeProcessGroup(), torch.device("cpu"), 16, @@ -310,7 +311,7 @@ def test_q_gather_rejects_invalid_workspace_geometry(self): padded_num_heads=7, ) with pytest.raises(ValueError, match="16-byte"): - dcp_utils.DirectDCPQGatherWorkspace( + dcp.DirectDCPQGatherWorkspace( _FakeProcessGroup(), torch.device("cpu"), 16, @@ -332,7 +333,7 @@ def _manager_config(dcp_comm_backend: str = "a2a"): def test_mla_dcp_manager_selects_direct_backends(monkeypatch): - import vllm.v1.attention.ops.dcp_utils as dcp_manager + import vllm.v1.attention.ops.dcp as dcp_manager group = MagicMock(world_size=2) monkeypatch.setattr(dcp_manager, "get_dcp_group", lambda: group) @@ -391,7 +392,7 @@ def test_mla_dcp_manager_selects_direct_backends(monkeypatch): def test_mla_dcp_manager_selects_fallback_backends(monkeypatch): - import vllm.v1.attention.ops.dcp_utils as dcp_manager + import vllm.v1.attention.ops.dcp as dcp_manager group = MagicMock(world_size=2) gathered_query = torch.empty(1, 4, 8) @@ -460,11 +461,11 @@ def test_dcp_workspace_covers_parallel_drafting(): config.num_speculative_tokens = 3 config.speculative_config = MagicMock(parallel_drafting=True) - assert dcp_utils.get_dcp_workspace_max_num_tokens(config) == 28 + assert dcp.get_dcp_workspace_max_num_tokens(config) == 28 def test_mla_dcp_manager_selects_pcp_combine(monkeypatch): - import vllm.v1.attention.ops.dcp_utils as dcp_manager + import vllm.v1.attention.ops.dcp as dcp_manager monkeypatch.setattr(dcp_manager, "get_dcp_group", lambda: MagicMock(world_size=2)) manager = dcp_manager.MLADCPManager( @@ -521,7 +522,7 @@ def test_sparse_mla_builder_initializes_dcp_manager(monkeypatch): lambda _: MagicMock(kv_lora_rank=8, qk_rope_head_dim=4), ) - manager = object.__new__(dcp_utils.MLADCPManager) + manager = object.__new__(dcp.MLADCPManager) manager.init_kv_gather = MagicMock() layer = MagicMock(dcp_manager=manager) config = MagicMock() @@ -580,12 +581,12 @@ def _distributed_direct_q_gather_worker(env: dict[str, str]) -> None: heads_per_rank, head_dim, max_num_tokens = 6, 576, 128 padded_num_heads = 128 if world_size == 4 else None active_ubatch = [0] - dcp_utils.dbo_current_ubatch_id = lambda: active_ubatch[0] + dcp.dbo_current_ubatch_id = lambda: active_ubatch[0] for dtype_idx, dtype_name in enumerate( ("bfloat16", "float8_e4m3fn", "float32") ): dtype = _dtype_from_name(dtype_name) - workspace = dcp_utils.DirectDCPQGatherWorkspace( + workspace = dcp.DirectDCPQGatherWorkspace( dist.group.WORLD, device, max_num_tokens, @@ -747,11 +748,11 @@ def _distributed_direct_kv_gather_worker(env: dict[str, str]) -> None: token_dim = 576 max_gathered_tokens = 128 * world_size active_ubatch = [0] - dcp_utils.dbo_current_ubatch_id = lambda: active_ubatch[0] + dcp.dbo_current_ubatch_id = lambda: active_ubatch[0] for dtype_idx, dtype_name in enumerate(("bfloat16", "float16")): dtype = _dtype_from_name(dtype_name) - workspace = dcp_utils.DirectDCPKVGatherWorkspace( + workspace = dcp.DirectDCPKVGatherWorkspace( dist.group.WORLD, device, max_gathered_tokens, @@ -818,7 +819,7 @@ def _distributed_direct_a2a_worker(env: dict[str, str]) -> None: torch.accelerator.set_device_index(local_rank) dist.init_process_group(backend="nccl") try: - from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine + from vllm.v1.attention.ops.dcp import _lse_weighted_combine rank = dist.get_rank() world_size = dist.get_world_size() @@ -829,8 +830,8 @@ def _distributed_direct_a2a_worker(env: dict[str, str]) -> None: heads_per_rank, head_dim, max_num_tokens = 6, 512, 128 total_heads = world_size * heads_per_rank active_ubatch = [0] - dcp_utils.dbo_current_ubatch_id = lambda: active_ubatch[0] - workspace = dcp_utils.DirectDCPA2AWorkspace( + dcp.dbo_current_ubatch_id = lambda: active_ubatch[0] + workspace = dcp.DirectDCPA2AWorkspace( dist.group.WORLD, device, max_num_tokens, diff --git a/tests/v1/attention/test_flashinfer_mla_dcp.py b/tests/v1/attention/test_flashinfer_mla_dcp.py index f3a46bbb1900..1841b25d490d 100644 --- a/tests/v1/attention/test_flashinfer_mla_dcp.py +++ b/tests/v1/attention/test_flashinfer_mla_dcp.py @@ -15,7 +15,7 @@ def test_mla_dcp_gathered_query_reserves_backend_head_storage(): - from vllm.v1.attention.ops.dcp_utils import reserve_query_head_storage + from vllm.v1.attention.ops.dcp import reserve_query_head_storage query = torch.randn(3, 24, 576, dtype=torch.bfloat16) diff --git a/tests/v1/attention/test_indexer_dcp_localize.py b/tests/v1/attention/test_indexer_dcp_localize.py index 60468cd33771..48f3ffd5ac16 100644 --- a/tests/v1/attention/test_indexer_dcp_localize.py +++ b/tests/v1/attention/test_indexer_dcp_localize.py @@ -12,7 +12,7 @@ triton_filter_and_convert_dcp_index, ) from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens -from vllm.v1.attention.ops.common import CPTritonContext, correct_attn_out +from vllm.v1.attention.ops.dcp import CPTritonContext, correct_attn_out def _local_count(length: int, rank: int, world: int, interleave: int) -> int: diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 37dc2c95ff49..b41c810d7646 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -242,10 +242,6 @@ from vllm.model_executor.layers.attention.kv_transfer_utils import ( maybe_transfer_kv_layer, ) -from vllm.model_executor.layers.attention.pcp import ( - finalize_mla_pcp_decode, - maybe_gather_mla_latent_cache_inputs, -) from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.linear import ( ColumnParallelLinear, @@ -293,8 +289,12 @@ get_num_attention_heads_from_layers, split_decodes_and_prefills, ) -from vllm.v1.attention.ops.dcp_utils import MLADCPManager +from vllm.v1.attention.ops.dcp import MLADCPManager from vllm.v1.attention.ops.merge_attn_states import merge_attn_states +from vllm.v1.attention.ops.pcp import ( + finalize_mla_pcp_decode, + maybe_gather_mla_latent_cache_inputs, +) from vllm.v1.attention.selector import get_attn_backend from vllm.v1.kv_cache_interface import ( AttentionSpec, diff --git a/vllm/model_executor/layers/attention/sparse_mla_attention.py b/vllm/model_executor/layers/attention/sparse_mla_attention.py index 85e09cd5d49b..1fe267f3c847 100644 --- a/vllm/model_executor/layers/attention/sparse_mla_attention.py +++ b/vllm/model_executor/layers/attention/sparse_mla_attention.py @@ -32,7 +32,7 @@ from vllm.v1.attention.backend import AttentionMetadata, AttentionMetadataBuilder from vllm.v1.attention.backends.fa_utils import get_flash_attn_version from vllm.v1.attention.backends.utils import split_decodes_and_prefills -from vllm.v1.attention.ops.dcp_utils import MLADCPManager +from vllm.v1.attention.ops.dcp import MLADCPManager from vllm.v1.attention.ops.merge_attn_states import merge_attn_states if TYPE_CHECKING: diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 4f0e00cfabff..8c7332aac78b 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -13,7 +13,6 @@ from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.custom_op import CustomOp -from vllm.model_executor.layers.attention.pcp import maybe_gather_indexer_k from vllm.model_executor.layers.quantization.utils.quant_utils import ( get_fp8_min_max, ) @@ -35,6 +34,7 @@ DeepseekV32IndexerMetadata, ) from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton +from vllm.v1.attention.ops.pcp import maybe_gather_indexer_k from vllm.v1.worker.workspace import current_workspace_manager logger = init_logger(__name__) diff --git a/vllm/models/deepseek_v32/attention.py b/vllm/models/deepseek_v32/attention.py index 02582810ca02..4edcae4640af 100644 --- a/vllm/models/deepseek_v32/attention.py +++ b/vllm/models/deepseek_v32/attention.py @@ -13,10 +13,6 @@ from vllm.forward_context import get_forward_context from vllm.model_executor.layers.attention import MLAAttention from vllm.model_executor.layers.attention.attention import get_attention_context -from vllm.model_executor.layers.attention.pcp import ( - finalize_mla_pcp_decode, - maybe_gather_mla_latent_cache_inputs, -) from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm from vllm.model_executor.layers.linear import ( ColumnParallelLinear, @@ -42,6 +38,10 @@ from vllm.models.deepseek_v32.common.kernels import fused_norm_rope, fused_q from vllm.platforms import current_platform from vllm.utils.torch_utils import is_quantized_kv_cache +from vllm.v1.attention.ops.pcp import ( + finalize_mla_pcp_decode, + maybe_gather_mla_latent_cache_inputs, +) if TYPE_CHECKING: from vllm.model_executor.layers.attention.mla_attention import MLACommonMetadata diff --git a/vllm/models/kimi_k3/nvidia/mla.py b/vllm/models/kimi_k3/nvidia/mla.py index 2c67a8a4c4a6..07f51a92508c 100644 --- a/vllm/models/kimi_k3/nvidia/mla.py +++ b/vllm/models/kimi_k3/nvidia/mla.py @@ -92,7 +92,7 @@ MLAAttentionImpl, ) from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend -from vllm.v1.attention.ops.dcp_utils import MLADCPManager +from vllm.v1.attention.ops.dcp import MLADCPManager from vllm.v1.attention.ops.merge_attn_states import merge_attn_states from vllm.v1.attention.selector import get_attn_backend from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec, get_kv_quant_mode diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index d0da103fb399..04e91abb03d0 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -886,7 +886,7 @@ class AttentionImplBase(ABC, Generic[T]): # False => base 2 (lse = log2(sum(exp(qk)))) # -- e.g. FlashInfer trtllm-gen MLA # The DCP combine kernel (cp_lse_ag_out_rs / dcp_a2a_lse_reduce in - # vllm/v1/attention/ops/common.py) branches on this via its IS_BASE_E + # vllm/v1/attention/ops/dcp.py) branches on this via its IS_BASE_E # constexpr; getting it wrong silently corrupts the cross-shard # softmax denominator. lse_base_on_e: bool = True diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index aa89e3010d68..7ccbddf1246b 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -34,8 +34,10 @@ fill_mm_prefix_query_ranges, get_dcp_local_seq_lens, ) -from vllm.v1.attention.ops.common import cp_lse_ag_out_rs -from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce +from vllm.v1.attention.ops.dcp import ( + cp_lse_ag_out_rs, + dcp_a2a_lse_reduce, +) from vllm.v1.attention.ops.merge_attn_states import merge_attn_states from vllm.v1.worker.workspace import current_workspace_manager diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index de51e2198ba2..fa8faf2a58f9 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -75,8 +75,10 @@ infer_global_hyperparameters, split_decodes_and_prefills, ) -from vllm.v1.attention.ops.common import cp_lse_ag_out_rs -from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce +from vllm.v1.attention.ops.dcp import ( + cp_lse_ag_out_rs, + dcp_a2a_lse_reduce, +) from vllm.v1.attention.ops.merge_attn_states import merge_attn_states from vllm.v1.kv_cache_interface import ( AttentionSpec, diff --git a/vllm/v1/attention/ops/common.py b/vllm/v1/attention/ops/common.py index cbe24cf87ce4..f3316ac557eb 100644 --- a/vllm/v1/attention/ops/common.py +++ b/vllm/v1/attention/ops/common.py @@ -2,308 +2,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch -from vllm.distributed.parallel_state import GroupCoordinator from vllm.triton_utils import tl, triton -def mask_dcp_empty_shards_( - lse: torch.Tensor, - seq_lens: torch.Tensor | None, - query_start_loc: torch.Tensor | None, -) -> None: - if seq_lens is None and query_start_loc is None: - return - if seq_lens is None or query_start_loc is None: - raise ValueError("seq_lens and query_start_loc must be provided together") - if ( - seq_lens.ndim != 1 - or query_start_loc.ndim != 1 - or query_start_loc.shape[0] != seq_lens.shape[0] + 1 - ): - raise ValueError("query_start_loc must contain one boundary per sequence") - - row_indices = torch.arange( - lse.shape[0], device=lse.device, dtype=query_start_loc.dtype - ) - sequence_indices = torch.searchsorted( - query_start_loc[1:], row_indices, right=True - ).clamp_max(seq_lens.shape[0] - 1) - empty_rows = (row_indices >= query_start_loc[-1]) | ( - seq_lens[sequence_indices] == 0 - ) - lse.masked_fill_(empty_rows[:, None], float("-inf")) - - -@triton.jit -def _correct_attn_cp_out_kernel( - outputs_ptr, - new_output_ptr, - lses_ptr, - vlse_ptr, - outputs_stride_B, - outputs_stride_H, - outputs_stride_D, - lses_stride_N, - lses_stride_B, - lses_stride_H, - lse_idx, - HEAD_DIM: tl.constexpr, - N_ROUNDED: tl.constexpr, - IS_BASE_E: tl.constexpr, -): - """ - Apply the all-gathered lses to correct each local rank's attention - output. we still need perform a cross-rank reduction to obtain the - final attention output. - - Args: - outputs_ptr (triton.PointerType): - Pointer to input tensor of shape [ B, H, D ] - lses_ptr (triton.PointerType): - Pointer to input tensor of shape [ N, B, H ] - new_output_ptr (triton.PointerType): - Pointer to output tensor of shape [ B, H, D ] - vlse_ptr (triton.PointerType): - Pointer to output tensor of shape [ B, H ] - """ - batch_idx = tl.program_id(axis=0).to(tl.int64) - head_idx = tl.program_id(axis=1).to(tl.int64) - d_offsets = tl.arange(0, HEAD_DIM) - num_n_offsets = tl.arange(0, N_ROUNDED) - - # shape = [N] - lse_offsets = ( - num_n_offsets * lses_stride_N - + batch_idx * lses_stride_B - + head_idx * lses_stride_H - ) - - # calc final lse - lse = tl.load(lses_ptr + lse_offsets).to(tl.float32) - lse = tl.where((lse != lse) | (lse == float("inf")), -float("inf"), lse) - lse_max = tl.max(lse, axis=0) - lse_max = tl.where(lse_max == -float("inf"), 0, lse_max) - lse -= lse_max - if IS_BASE_E: - lse_exp = tl.exp(lse) - lse_acc = tl.sum(lse_exp, axis=0) - lse = tl.log(lse_acc) - else: - lse_exp = tl.exp2(lse) - lse_acc = tl.sum(lse_exp, axis=0) - lse = tl.log2(lse_acc) - lse += lse_max - - lse_offsets = batch_idx * lses_stride_B + head_idx * lses_stride_H - tl.store(vlse_ptr + lse_offsets, lse) - - # shape = [D] - output_offsets = ( - batch_idx * outputs_stride_B - + head_idx * outputs_stride_H - + d_offsets * outputs_stride_D - ) - - # correct output - lse_offset = ( - lse_idx * lses_stride_N + batch_idx * lses_stride_B + head_idx * lses_stride_H - ) - lse_tmp = tl.load(lses_ptr + lse_offset).to(tl.float32) - lse_finally = lse_tmp - lse - lse_finally = tl.where( - (lse_finally != lse_finally) | (lse_finally == float("inf")), - -float("inf"), - lse_finally, - ) - factor = tl.exp(lse_finally) if IS_BASE_E else tl.exp2(lse_finally) - output = tl.load(outputs_ptr + output_offsets) - output = output * factor - output = tl.where(factor == 0.0, 0.0, output) - - tl.store(new_output_ptr + output_offsets, output) - - -class CPTritonContext: - """The CPTritonContext is used to avoid recompilation of the Triton JIT.""" - - def __init__(self): - self.inner_kernel = None - - def call_kernel(self, kernel, grid, *regular_args, **const_args): - if self.inner_kernel is None: - self.inner_kernel = kernel[grid](*regular_args, **const_args) - else: - self.inner_kernel[grid](*regular_args) - - -def correct_attn_out( - out: torch.Tensor, - lses: torch.Tensor, - cp_rank: int, - ctx: CPTritonContext, - is_lse_base_on_e: bool = True, -) -> tuple[torch.Tensor, torch.Tensor]: - """Correct the attention output using the all-gathered lses. - - Args: - out: Tensor of shape [ B, H, D ] - lses: Tensor of shape [ N, B, H ] - cp_rank: Current rank in the context-parallel group - ctx: Triton context to avoid recompilation - - Returns: - Tuple of (out, lse) with corrected attention and final log-sum-exp. - """ - if ctx is None: - ctx = CPTritonContext() - - # --- Normalize to 3D views --- - if out.ndim == 4 and out.shape[1] == 1: - out = out.squeeze(1) - assert out.ndim == 3, f"expected out [B,H,D] or [B,1,H,D], got {tuple(out.shape)}" - - if lses.ndim == 4 and lses.shape[-1] == 1: - lses = lses.squeeze(-1) - if lses.ndim == 4 and lses.shape[1] == 1: - lses = lses.squeeze(1) - assert lses.ndim == 3, ( - f"expected lses [N,B,H] (optionally with a 1-sized extra dim), " - f"got {tuple(lses.shape)}" - ) - - B, H, D = out.shape - N = lses.shape[0] - - # Strides after we normalized shapes to 3-D views. The kernel computes - # offsets for `vlse_ptr` using lses_stride_B/H, so the output buffer must - # have the same B/H stride layout as a slice of `lses`. - o_sB, o_sH, o_sD = out.stride() - l_sN, l_sB, l_sH = lses.stride() - - # Allocate LSE with the same B/H strides as `lses` so writes land correctly - # even when `lses` is a non-contiguous view (e.g., 4-D to 3-D squeeze). - lse = torch.empty_strided( - (B, H), (l_sB, l_sH), device=lses.device, dtype=lses.dtype - ) - - # Kernel launch config - grid = (B, H, 1) - - regular_args = ( - out, - out, - lses, - lse, - o_sB, - o_sH, - o_sD, - l_sN, - l_sB, - l_sH, - cp_rank, - ) - const_args = {"HEAD_DIM": D, "N_ROUNDED": N, "IS_BASE_E": is_lse_base_on_e} - ctx.call_kernel(_correct_attn_cp_out_kernel, grid, *regular_args, **const_args) - return out, lse - - -def _cp_lse_common( - cp_attn_out: torch.Tensor, - cp_attn_lse: torch.Tensor, - cp_group: GroupCoordinator, - ctx: CPTritonContext | None = None, - is_lse_base_on_e=True, - seq_lens: torch.Tensor | None = None, - query_start_loc: torch.Tensor | None = None, -): - """ - cp_attn_out: [ B, H, D ] - cp_attn_lse: [ B, H ] - """ - if cp_group.world_size == 1: - return cp_attn_out - - if ctx is None: - ctx = CPTritonContext() - - cp_attn_lse = cp_attn_lse.contiguous() - mask_dcp_empty_shards_(cp_attn_lse, seq_lens, query_start_loc) - lses = cp_group.all_gather(cp_attn_lse, dim=0).reshape( - (cp_group.world_size,) + cp_attn_lse.shape - ) - out, lse = correct_attn_out( - cp_attn_out, - lses, - cp_group.rank_in_group, - ctx, - is_lse_base_on_e=is_lse_base_on_e, - ) - return out, lse - - -def cp_lse_ag_out_rs( - cp_attn_out: torch.Tensor, - cp_attn_lse: torch.Tensor, - cp_group: GroupCoordinator, - ctx: CPTritonContext | None = None, - return_lse: bool = False, - is_lse_base_on_e=True, - seq_lens: torch.Tensor | None = None, - query_start_loc: torch.Tensor | None = None, -): - """ - cp_attn_out: [ B, H, D ] - cp_attn_lse: [ B, H ] - """ - out, lse = _cp_lse_common( - cp_attn_out, - cp_attn_lse, - cp_group, - ctx=ctx, - is_lse_base_on_e=is_lse_base_on_e, - seq_lens=seq_lens, - query_start_loc=query_start_loc, - ) - out = cp_group.reduce_scatter(out, dim=1) - - if return_lse: - cp_num_heads = lse.shape[1] // cp_group.world_size - cp_rank = cp_group.rank_in_group - lse = lse[:, cp_num_heads * cp_rank : cp_num_heads * (cp_rank + 1)] - return out, lse - return out - - -def cp_lse_ag_out_ar( - cp_attn_out: torch.Tensor, - cp_attn_lse: torch.Tensor, - cp_group: GroupCoordinator, - ctx: CPTritonContext | None = None, - return_lse: bool = False, - is_lse_base_on_e=True, - seq_lens: torch.Tensor | None = None, - query_start_loc: torch.Tensor | None = None, -): - """ - cp_attn_out: [ B, H, D ] - cp_attn_lse: [ B, H ] - """ - out, lse = _cp_lse_common( - cp_attn_out, - cp_attn_lse, - cp_group, - ctx=ctx, - is_lse_base_on_e=is_lse_base_on_e, - seq_lens=seq_lens, - query_start_loc=query_start_loc, - ) - out = cp_group.all_reduce(out) - - if return_lse: - return out, lse - return out - - @triton.jit def _pack_seq_kernel( x_ptr, # [N, D] diff --git a/vllm/v1/attention/ops/cp_common.py b/vllm/v1/attention/ops/cp_common.py new file mode 100644 index 000000000000..3ee29b93e62e --- /dev/null +++ b/vllm/v1/attention/ops/cp_common.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared symmetric-memory infrastructure for context-parallel attention.""" + +from __future__ import annotations + +import functools +from typing import TYPE_CHECKING, Any + +import torch + +from vllm.distributed.parallel_state import in_the_same_node_as +from vllm.logger import init_logger +from vllm.platforms import current_platform + +if TYPE_CHECKING: + from torch.distributed import ProcessGroup + + from vllm.distributed.parallel_state import GroupCoordinator + +logger = init_logger(__name__) + +try: + import torch.distributed._symmetric_memory as symm_mem + + symm_mem_available = True +except ImportError: + symm_mem = None # type: ignore[assignment] + symm_mem_available = False + + +@functools.cache +def _symm_mem_spans_group(group: GroupCoordinator) -> bool: + """Probe whether the group has NVLS symmetric memory.""" + if not symm_mem_available: + return False + try: + from torch._C._autograd import DeviceType + from torch._C._distributed_c10d import _SymmetricMemory + + device = torch.device("cuda", torch.accelerator.current_device_index()) + if not _SymmetricMemory.has_multicast_support(DeviceType.CUDA, device.index): + return False + probe = symm_mem.empty(8, dtype=torch.uint8, device=device) + probe.zero_() + torch.accelerator.synchronize() + handle = symm_mem.rendezvous(probe, group.device_group.group_name) + spans = handle is not None and handle.multicast_ptr != 0 + except Exception as error: + logger.debug("Direct CP symmetric-memory probe failed: %s", error) + return False + logger.debug_once( + "Direct CP symmetric memory across %d ranks: %s", + group.world_size, + "available" if spans else "unavailable", + ) + return spans + + +def direct_cp_enabled( + group: GroupCoordinator, + dtype: torch.dtype, + use_direct: bool | None, + supported_dtypes: tuple[torch.dtype, ...] | None = None, +) -> bool: + if use_direct is not None: + return use_direct + return ( + symm_mem_available + and current_platform.is_cuda() + and (supported_dtypes is None or dtype in supported_dtypes) + and ( + all(in_the_same_node_as(group.cpu_group, source_rank=0)) + or _symm_mem_spans_group(group) + ) + ) + + +def direct_cp_multicast_enabled( + group: GroupCoordinator, + dtype: torch.dtype, + use_direct: bool | None, + supported_dtypes: tuple[torch.dtype, ...] | None = None, +) -> bool: + return direct_cp_enabled( + group, dtype, use_direct, supported_dtypes + ) and _symm_mem_spans_group(group) + + +class DirectCPWorkspace: + def __init__( + self, + group: ProcessGroup, + device: torch.device, + num_ubatches: int, + ) -> None: + self.group = group + self.world_size = group.size() + self.rank = group.rank() + self.device = torch.device(device) + self.num_ubatches = num_ubatches + self.epoch = torch.zeros(num_ubatches, dtype=torch.int64, device=self.device) + self._allocations: list[tuple[torch.Tensor, Any, list[torch.Tensor]]] = [] + + def _allocate( + self, shape: tuple[int, ...], dtype: torch.dtype + ) -> tuple[torch.Tensor, torch.Tensor]: + storage = symm_mem.empty(shape, device=self.device, dtype=dtype) + storage.zero_() + torch.accelerator.synchronize() + handle = symm_mem.rendezvous(storage, self.group.group_name) + assert handle is not None, "CP symmetric memory rendezvous returned None" + handle.barrier() + views = [ + handle.get_buffer(peer, list(shape), dtype, 0) + for peer in range(self.world_size) + ] + self.device = storage.device + peer_ptrs = torch.tensor( + [ + [view[ubatch].data_ptr() for view in views] + for ubatch in range(self.num_ubatches) + ], + dtype=torch.int64, + device=self.device, + ) + self._allocations.append((storage, handle, views)) + return storage, peer_ptrs + + def _multicast_ptrs(self, storage: torch.Tensor) -> list[int]: + disabled = [0] * self.num_ubatches + for allocated, handle, _ in self._allocations: + if allocated is storage: + break + else: + return disabled + try: + from torch._C._autograd import DeviceType + from torch._C._distributed_c10d import _SymmetricMemory + + if not _SymmetricMemory.has_multicast_support( + DeviceType.CUDA, storage.device.index + ): + return disabled + multicast_base = handle.multicast_ptr + except Exception: + return disabled + if not multicast_base: + return disabled + storage_base = storage.data_ptr() + return [ + multicast_base + (storage[ubatch].data_ptr() - storage_base) + for ubatch in range(self.num_ubatches) + ] diff --git a/vllm/v1/attention/ops/dcp.py b/vllm/v1/attention/ops/dcp.py new file mode 100644 index 000000000000..6381bfa487e6 --- /dev/null +++ b/vllm/v1/attention/ops/dcp.py @@ -0,0 +1,1368 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MLA DCP collective selection and direct symmetric-memory implementations.""" + +from __future__ import annotations + +import functools +from collections.abc import Callable +from typing import TYPE_CHECKING, Protocol + +import torch +import torch.distributed as dist + +import vllm.envs as envs +from vllm.config import VllmConfig +from vllm.distributed import get_dcp_group +from vllm.logger import init_logger +from vllm.triton_utils import tl, triton +from vllm.v1.attention.ops.cp_common import ( + DirectCPWorkspace, + direct_cp_enabled, + direct_cp_multicast_enabled, +) +from vllm.v1.worker.ubatching import dbo_current_ubatch_id + +logger = init_logger(__name__) + +if TYPE_CHECKING: + from torch.distributed import ProcessGroup + + from vllm.distributed.parallel_state import GroupCoordinator + + +# LSE/output combine + + +def mask_dcp_empty_shards_( + lse: torch.Tensor, + seq_lens: torch.Tensor | None, + query_start_loc: torch.Tensor | None, +) -> None: + if seq_lens is None and query_start_loc is None: + return + if seq_lens is None or query_start_loc is None: + raise ValueError("seq_lens and query_start_loc must be provided together") + if ( + seq_lens.ndim != 1 + or query_start_loc.ndim != 1 + or query_start_loc.shape[0] != seq_lens.shape[0] + 1 + ): + raise ValueError("query_start_loc must contain one boundary per sequence") + + row_indices = torch.arange( + lse.shape[0], device=lse.device, dtype=query_start_loc.dtype + ) + sequence_indices = torch.searchsorted( + query_start_loc[1:], row_indices, right=True + ).clamp_max(seq_lens.shape[0] - 1) + empty_rows = (row_indices >= query_start_loc[-1]) | ( + seq_lens[sequence_indices] == 0 + ) + lse.masked_fill_(empty_rows[:, None], float("-inf")) + + +# AG + RS/AR implementation + + +@triton.jit +def _correct_attn_cp_out_kernel( + outputs_ptr, + new_output_ptr, + lses_ptr, + vlse_ptr, + outputs_stride_B, + outputs_stride_H, + outputs_stride_D, + lses_stride_N, + lses_stride_B, + lses_stride_H, + lse_idx, + HEAD_DIM: tl.constexpr, + N_ROUNDED: tl.constexpr, + IS_BASE_E: tl.constexpr, +): + """ + Apply the all-gathered lses to correct each local rank's attention + output. we still need perform a cross-rank reduction to obtain the + final attention output. + + Args: + outputs_ptr (triton.PointerType): + Pointer to input tensor of shape [ B, H, D ] + lses_ptr (triton.PointerType): + Pointer to input tensor of shape [ N, B, H ] + new_output_ptr (triton.PointerType): + Pointer to output tensor of shape [ B, H, D ] + vlse_ptr (triton.PointerType): + Pointer to output tensor of shape [ B, H ] + """ + batch_idx = tl.program_id(axis=0).to(tl.int64) + head_idx = tl.program_id(axis=1).to(tl.int64) + d_offsets = tl.arange(0, HEAD_DIM) + num_n_offsets = tl.arange(0, N_ROUNDED) + + # shape = [N] + lse_offsets = ( + num_n_offsets * lses_stride_N + + batch_idx * lses_stride_B + + head_idx * lses_stride_H + ) + + # calc final lse + lse = tl.load(lses_ptr + lse_offsets).to(tl.float32) + lse = tl.where((lse != lse) | (lse == float("inf")), -float("inf"), lse) + lse_max = tl.max(lse, axis=0) + lse_max = tl.where(lse_max == -float("inf"), 0, lse_max) + lse -= lse_max + if IS_BASE_E: + lse_exp = tl.exp(lse) + lse_acc = tl.sum(lse_exp, axis=0) + lse = tl.log(lse_acc) + else: + lse_exp = tl.exp2(lse) + lse_acc = tl.sum(lse_exp, axis=0) + lse = tl.log2(lse_acc) + lse += lse_max + + lse_offsets = batch_idx * lses_stride_B + head_idx * lses_stride_H + tl.store(vlse_ptr + lse_offsets, lse) + + # shape = [D] + output_offsets = ( + batch_idx * outputs_stride_B + + head_idx * outputs_stride_H + + d_offsets * outputs_stride_D + ) + + # correct output + lse_offset = ( + lse_idx * lses_stride_N + batch_idx * lses_stride_B + head_idx * lses_stride_H + ) + lse_tmp = tl.load(lses_ptr + lse_offset).to(tl.float32) + lse_finally = lse_tmp - lse + lse_finally = tl.where( + (lse_finally != lse_finally) | (lse_finally == float("inf")), + -float("inf"), + lse_finally, + ) + factor = tl.exp(lse_finally) if IS_BASE_E else tl.exp2(lse_finally) + output = tl.load(outputs_ptr + output_offsets) + output = output * factor + output = tl.where(factor == 0.0, 0.0, output) + + tl.store(new_output_ptr + output_offsets, output) + + +class CPTritonContext: + """The CPTritonContext is used to avoid recompilation of the Triton JIT.""" + + def __init__(self): + self.inner_kernel = None + + def call_kernel(self, kernel, grid, *regular_args, **const_args): + if self.inner_kernel is None: + self.inner_kernel = kernel[grid](*regular_args, **const_args) + else: + self.inner_kernel[grid](*regular_args) + + +def correct_attn_out( + out: torch.Tensor, + lses: torch.Tensor, + cp_rank: int, + ctx: CPTritonContext, + is_lse_base_on_e: bool = True, +) -> tuple[torch.Tensor, torch.Tensor]: + """Correct the attention output using the all-gathered lses. + + Args: + out: Tensor of shape [ B, H, D ] + lses: Tensor of shape [ N, B, H ] + cp_rank: Current rank in the context-parallel group + ctx: Triton context to avoid recompilation + + Returns: + Tuple of (out, lse) with corrected attention and final log-sum-exp. + """ + if ctx is None: + ctx = CPTritonContext() + + # --- Normalize to 3D views --- + if out.ndim == 4 and out.shape[1] == 1: + out = out.squeeze(1) + assert out.ndim == 3, f"expected out [B,H,D] or [B,1,H,D], got {tuple(out.shape)}" + + if lses.ndim == 4 and lses.shape[-1] == 1: + lses = lses.squeeze(-1) + if lses.ndim == 4 and lses.shape[1] == 1: + lses = lses.squeeze(1) + assert lses.ndim == 3, ( + f"expected lses [N,B,H] (optionally with a 1-sized extra dim), " + f"got {tuple(lses.shape)}" + ) + + B, H, D = out.shape + N = lses.shape[0] + + # Strides after we normalized shapes to 3-D views. The kernel computes + # offsets for `vlse_ptr` using lses_stride_B/H, so the output buffer must + # have the same B/H stride layout as a slice of `lses`. + o_sB, o_sH, o_sD = out.stride() + l_sN, l_sB, l_sH = lses.stride() + + # Allocate LSE with the same B/H strides as `lses` so writes land correctly + # even when `lses` is a non-contiguous view (e.g., 4-D to 3-D squeeze). + lse = torch.empty_strided( + (B, H), (l_sB, l_sH), device=lses.device, dtype=lses.dtype + ) + + # Kernel launch config + grid = (B, H, 1) + + regular_args = ( + out, + out, + lses, + lse, + o_sB, + o_sH, + o_sD, + l_sN, + l_sB, + l_sH, + cp_rank, + ) + const_args = {"HEAD_DIM": D, "N_ROUNDED": N, "IS_BASE_E": is_lse_base_on_e} + ctx.call_kernel(_correct_attn_cp_out_kernel, grid, *regular_args, **const_args) + return out, lse + + +def _cp_lse_common( + cp_attn_out: torch.Tensor, + cp_attn_lse: torch.Tensor, + cp_group: GroupCoordinator, + ctx: CPTritonContext | None = None, + is_lse_base_on_e=True, + seq_lens: torch.Tensor | None = None, + query_start_loc: torch.Tensor | None = None, +): + """ + cp_attn_out: [ B, H, D ] + cp_attn_lse: [ B, H ] + """ + if cp_group.world_size == 1: + return cp_attn_out + + if ctx is None: + ctx = CPTritonContext() + + cp_attn_lse = cp_attn_lse.contiguous() + mask_dcp_empty_shards_(cp_attn_lse, seq_lens, query_start_loc) + lses = cp_group.all_gather(cp_attn_lse, dim=0).reshape( + (cp_group.world_size,) + cp_attn_lse.shape + ) + out, lse = correct_attn_out( + cp_attn_out, + lses, + cp_group.rank_in_group, + ctx, + is_lse_base_on_e=is_lse_base_on_e, + ) + return out, lse + + +def cp_lse_ag_out_rs( + cp_attn_out: torch.Tensor, + cp_attn_lse: torch.Tensor, + cp_group: GroupCoordinator, + ctx: CPTritonContext | None = None, + return_lse: bool = False, + is_lse_base_on_e=True, + seq_lens: torch.Tensor | None = None, + query_start_loc: torch.Tensor | None = None, +): + """ + cp_attn_out: [ B, H, D ] + cp_attn_lse: [ B, H ] + """ + out, lse = _cp_lse_common( + cp_attn_out, + cp_attn_lse, + cp_group, + ctx=ctx, + is_lse_base_on_e=is_lse_base_on_e, + seq_lens=seq_lens, + query_start_loc=query_start_loc, + ) + out = cp_group.reduce_scatter(out, dim=1) + + if return_lse: + cp_num_heads = lse.shape[1] // cp_group.world_size + cp_rank = cp_group.rank_in_group + lse = lse[:, cp_num_heads * cp_rank : cp_num_heads * (cp_rank + 1)] + return out, lse + return out + + +def cp_lse_ag_out_ar( + cp_attn_out: torch.Tensor, + cp_attn_lse: torch.Tensor, + cp_group: GroupCoordinator, + ctx: CPTritonContext | None = None, + return_lse: bool = False, + is_lse_base_on_e=True, + seq_lens: torch.Tensor | None = None, + query_start_loc: torch.Tensor | None = None, +): + """ + cp_attn_out: [ B, H, D ] + cp_attn_lse: [ B, H ] + """ + out, lse = _cp_lse_common( + cp_attn_out, + cp_attn_lse, + cp_group, + ctx=ctx, + is_lse_base_on_e=is_lse_base_on_e, + seq_lens=seq_lens, + query_start_loc=query_start_loc, + ) + out = cp_group.all_reduce(out) + + if return_lse: + return out, lse + return out + + +# Standard A2A implementation + + +def _lse_weighted_combine( + outputs: torch.Tensor, + lses: torch.Tensor, + return_lse: bool = False, + is_lse_base_on_e: bool = True, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """ + CPU reference implementation for LSE-weighted combination. + + This is a pure PyTorch implementation used for testing and validation. + + Args: + outputs: Partial attention outputs [N, B, H, D] + N = number of KV shards (ranks) + B = batch size (num_tokens) + H = number of heads per rank + D = head dimension + lses: Log-sum-exp values [N, B, H] + return_lse: If True, also return the global LSE + is_lse_base_on_e: If True, LSE is base e; if False, base 2 + + Returns: + Combined output [B, H, D], and optionally global LSE [B, H] + """ + N, B, H, D = outputs.shape + + # Handle NaN and inf in LSEs + lses = torch.where( + torch.isnan(lses) | torch.isinf(lses), + torch.tensor(float("-inf"), device=lses.device, dtype=lses.dtype), + lses, + ) + + # Compute max LSE for numerical stability + lse_max, _ = lses.max(dim=0) # [B, H] + lse_max = torch.where( + lse_max == float("-inf"), + torch.zeros_like(lse_max), + lse_max, + ) + + # Compute weights: softmax over the N dimension + if is_lse_base_on_e: + weights = torch.exp(lses - lse_max.unsqueeze(0)) # [N, B, H] + else: + weights = torch.pow(2.0, lses - lse_max.unsqueeze(0)) # [N, B, H] + + # Handle NaN weights + weights = torch.where(torch.isnan(weights), torch.zeros_like(weights), weights) + + # Normalize weights + weight_sum = weights.sum(dim=0, keepdim=True) # [1, B, H] + weights = weights / weight_sum.clamp(min=1e-10) # [N, B, H] + + # Weighted combination: sum over N dimension + weights = weights.unsqueeze(-1) + outputs = torch.where(weights == 0, torch.zeros_like(outputs), outputs) + result = (outputs * weights).sum(dim=0) # [B, H, D] + + if return_lse: + if is_lse_base_on_e: + global_lse = torch.log(weight_sum.squeeze(0)) + lse_max # [B, H] + else: + global_lse = torch.log2(weight_sum.squeeze(0)) + lse_max # [B, H] + return result, global_lse + + return result + + +def _dcp_a2a_lse_pack_dim(output_dtype: torch.dtype) -> int: + bits = torch.finfo(output_dtype).bits + if bits == 16: + return 2 + if bits == 32: + return 1 + raise ValueError(f"Cannot pack fp32 LSE into output dtype {output_dtype}.") + + +def _dcp_a2a_send_recv_buffers( + shape: tuple[int, ...], + device: torch.device, + dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor]: + # Don't use the shared WorkspaceManager here. A FULL cudagraph bakes in the + # buffer address at capture, but the workspace is growable and sized only to + # the largest *captured* batch (the cudagraph capture cap). Any eager a2a + # with a bigger batch regrows it, freeing that address and poisoning every + # captured graph -> illegal memory access on replay. This bites the very + # first request: the post-capture warmup runs an eager decode at + # max_num_seqs (> the cap), so the graphs are already dangling before the + # server is ready. torch.empty buffers instead live in the graph's private + # pool and stay valid for its lifetime (as _dcp_a2a_unpack_combine and the + # AG+RS combine path already rely on). + return ( + torch.empty(shape, device=device, dtype=dtype), + torch.empty(shape, device=device, dtype=dtype), + ) + + +@triton.jit +def _dcp_a2a_pack_send_kernel( + out_ptr, + lse_ptr, + send_ptr, + out_stride_B, + out_stride_H, + out_stride_D, + lse_stride_B, + lse_stride_H, + send_stride_N, + send_stride_B, + send_stride_H, + send_stride_D, + N: tl.constexpr, + HEAD_DIM: tl.constexpr, + H_PER_RANK: tl.constexpr, + LSE_PACK_DIM: tl.constexpr, +): + batch_idx = tl.program_id(0).to(tl.int64) + local_head_idx = tl.program_id(1).to(tl.int64) + d_offsets = tl.arange(0, HEAD_DIM) + + for rank_idx in tl.static_range(N): + src_head_idx = rank_idx * H_PER_RANK + local_head_idx + send_base = ( + rank_idx * send_stride_N + + batch_idx * send_stride_B + + local_head_idx * send_stride_H + ) + + out_offsets = ( + batch_idx * out_stride_B + + src_head_idx * out_stride_H + + d_offsets * out_stride_D + ) + tl.store( + send_ptr + send_base + d_offsets * send_stride_D, + tl.load(out_ptr + out_offsets), + ) + + lse_val = tl.load( + lse_ptr + batch_idx * lse_stride_B + src_head_idx * lse_stride_H + ).to(tl.float32) + if LSE_PACK_DIM == 1: + tl.store( + send_ptr + send_base + HEAD_DIM * send_stride_D, + lse_val.to(send_ptr.dtype.element_ty), + ) + else: + lse_bits = lse_val.to(tl.uint32, bitcast=True) + lo = (lse_bits & 0xFFFF).to(tl.uint16) + hi = ((lse_bits >> 16) & 0xFFFF).to(tl.uint16) + tl.store( + send_ptr + send_base + HEAD_DIM * send_stride_D, + lo.to(send_ptr.dtype.element_ty, bitcast=True), + ) + tl.store( + send_ptr + send_base + (HEAD_DIM + 1) * send_stride_D, + hi.to(send_ptr.dtype.element_ty, bitcast=True), + ) + + +@triton.jit +def _dcp_a2a_unpack_combine_kernel( + recv_ptr, + out_ptr, + out_lse_ptr, + recv_stride_N, + recv_stride_B, + recv_stride_H, + recv_stride_D, + out_stride_B, + out_stride_H, + out_stride_D, + out_lse_stride_B, + out_lse_stride_H, + N: tl.constexpr, + HEAD_DIM: tl.constexpr, + IS_BASE_E: tl.constexpr, + RETURN_LSE: tl.constexpr, + LSE_PACK_DIM: tl.constexpr, +): + batch_idx = tl.program_id(0).to(tl.int64) + head_idx = tl.program_id(1).to(tl.int64) + d_offsets = tl.arange(0, HEAD_DIM) + + lse_max = -float("inf") + for rank_idx in tl.static_range(N): + recv_base = ( + rank_idx * recv_stride_N + + batch_idx * recv_stride_B + + head_idx * recv_stride_H + ) + if LSE_PACK_DIM == 1: + lse_val = tl.load(recv_ptr + recv_base + HEAD_DIM * recv_stride_D).to( + tl.float32 + ) + else: + lo_raw = tl.load(recv_ptr + recv_base + HEAD_DIM * recv_stride_D) + hi_raw = tl.load(recv_ptr + recv_base + (HEAD_DIM + 1) * recv_stride_D) + lo = lo_raw.to(tl.uint16, bitcast=True).to(tl.uint32) + hi = hi_raw.to(tl.uint16, bitcast=True).to(tl.uint32) + lse_val = (lo | (hi << 16)).to(tl.float32, bitcast=True) + lse_val = tl.where( + (lse_val != lse_val) | (lse_val == float("inf")), + -float("inf"), + lse_val, + ) + lse_max = tl.maximum(lse_max, lse_val) + + lse_max = tl.where(lse_max == -float("inf"), 0.0, lse_max) + + lse_sum = 0.0 + for rank_idx in tl.static_range(N): + recv_base = ( + rank_idx * recv_stride_N + + batch_idx * recv_stride_B + + head_idx * recv_stride_H + ) + if LSE_PACK_DIM == 1: + lse_val = tl.load(recv_ptr + recv_base + HEAD_DIM * recv_stride_D).to( + tl.float32 + ) + else: + lo_raw = tl.load(recv_ptr + recv_base + HEAD_DIM * recv_stride_D) + hi_raw = tl.load(recv_ptr + recv_base + (HEAD_DIM + 1) * recv_stride_D) + lo = lo_raw.to(tl.uint16, bitcast=True).to(tl.uint32) + hi = hi_raw.to(tl.uint16, bitcast=True).to(tl.uint32) + lse_val = (lo | (hi << 16)).to(tl.float32, bitcast=True) + lse_val = tl.where( + (lse_val != lse_val) | (lse_val == float("inf")), + -float("inf"), + lse_val, + ) + if IS_BASE_E: + lse_sum += tl.exp(lse_val - lse_max) + else: + lse_sum += tl.exp2(lse_val - lse_max) + + if IS_BASE_E: # noqa: SIM108 + global_lse = tl.log(lse_sum) + lse_max + else: + global_lse = tl.log2(lse_sum) + lse_max + + acc = tl.zeros([HEAD_DIM], dtype=tl.float32) + for rank_idx in tl.static_range(N): + recv_base = ( + rank_idx * recv_stride_N + + batch_idx * recv_stride_B + + head_idx * recv_stride_H + ) + if LSE_PACK_DIM == 1: + lse_val = tl.load(recv_ptr + recv_base + HEAD_DIM * recv_stride_D).to( + tl.float32 + ) + else: + lo_raw = tl.load(recv_ptr + recv_base + HEAD_DIM * recv_stride_D) + hi_raw = tl.load(recv_ptr + recv_base + (HEAD_DIM + 1) * recv_stride_D) + lo = lo_raw.to(tl.uint16, bitcast=True).to(tl.uint32) + hi = hi_raw.to(tl.uint16, bitcast=True).to(tl.uint32) + lse_val = (lo | (hi << 16)).to(tl.float32, bitcast=True) + lse_val = tl.where( + (lse_val != lse_val) | (lse_val == float("inf")), + -float("inf"), + lse_val, + ) + if IS_BASE_E: + weight = tl.exp(lse_val - global_lse) + else: + weight = tl.exp2(lse_val - global_lse) + weight = tl.where(weight != weight, 0.0, weight) + partial = tl.load(recv_ptr + recv_base + d_offsets * recv_stride_D).to( + tl.float32 + ) + partial = tl.where(weight == 0.0, 0.0, partial) + acc += partial * weight + + final_offsets = ( + batch_idx * out_stride_B + head_idx * out_stride_H + d_offsets * out_stride_D + ) + tl.store(out_ptr + final_offsets, acc) + + if RETURN_LSE: + out_lse_offset = batch_idx * out_lse_stride_B + head_idx * out_lse_stride_H + tl.store(out_lse_ptr + out_lse_offset, global_lse) + + +def _dcp_a2a_pack_send( + cp_attn_out: torch.Tensor, + cp_attn_lse: torch.Tensor, + send_buffer: torch.Tensor, + world_size: int, + h_per_rank: int, + head_dim: int, + lse_pack_dim: int, + seq_lens: torch.Tensor | None = None, + query_start_loc: torch.Tensor | None = None, +) -> None: + mask_dcp_empty_shards_(cp_attn_lse, seq_lens, query_start_loc) + grid = (cp_attn_out.shape[0], h_per_rank, 1) + _dcp_a2a_pack_send_kernel[grid]( + cp_attn_out, + cp_attn_lse, + send_buffer, + cp_attn_out.stride(0), + cp_attn_out.stride(1), + cp_attn_out.stride(2), + cp_attn_lse.stride(0), + cp_attn_lse.stride(1), + send_buffer.stride(0), + send_buffer.stride(1), + send_buffer.stride(2), + send_buffer.stride(3), + N=world_size, + HEAD_DIM=head_dim, + H_PER_RANK=h_per_rank, + LSE_PACK_DIM=lse_pack_dim, + ) + + +def _dcp_a2a_unpack_combine( + recv_buffer: torch.Tensor, + head_dim: int, + lse_pack_dim: int, + return_lse: bool, + is_lse_base_on_e: bool, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + world_size, num_tokens, h_per_rank, _ = recv_buffer.shape + out = torch.empty( + (num_tokens, h_per_rank, head_dim), + device=recv_buffer.device, + dtype=recv_buffer.dtype, + ) + out_lse = torch.empty( + (num_tokens, h_per_rank) if return_lse else (1, 1), + device=recv_buffer.device, + dtype=torch.float32 if return_lse else recv_buffer.dtype, + ) + grid = (num_tokens, h_per_rank, 1) + _dcp_a2a_unpack_combine_kernel[grid]( + recv_buffer, + out, + out_lse, + recv_buffer.stride(0), + recv_buffer.stride(1), + recv_buffer.stride(2), + recv_buffer.stride(3), + out.stride(0), + out.stride(1), + out.stride(2), + out_lse.stride(0), + out_lse.stride(1), + N=world_size, + HEAD_DIM=head_dim, + IS_BASE_E=is_lse_base_on_e, + RETURN_LSE=return_lse, + LSE_PACK_DIM=lse_pack_dim, + ) + if return_lse: + return out, out_lse + return out + + +def dcp_a2a_lse_reduce( + cp_attn_out: torch.Tensor, + cp_attn_lse: torch.Tensor, + cp_group: GroupCoordinator, + ctx: CPTritonContext | None = None, + return_lse: bool = False, + is_lse_base_on_e: bool = True, + seq_lens: torch.Tensor | None = None, + query_start_loc: torch.Tensor | None = None, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """ + Combine partial attention outputs across DCP ranks using All-to-All. + + The output and LSE are packed into a single output-dtype buffer, sent + with one All-to-All, then unpacked and combined with exact LSE weighting. + + Args: + cp_attn_out: [B, H, D] where B=num_tokens, H=total_heads, D=head_dim + cp_attn_lse: [B, H] floating-point log-sum-exp values + cp_group: GroupCoordinator for DCP communication + ctx: CPTritonContext (unused, for signature compatibility) + return_lse: If True, also return the combined global LSE + is_lse_base_on_e: If True, LSE is base e; if False, base 2 + seq_lens: Local KV lengths. Empty shards contribute zero weight. + query_start_loc: Cumulative query-token offsets for each request. + + Returns: + Combined output [B, H/N, D] (head-scattered) + If return_lse=True, also returns global_lse [B, H/N] + """ + world_size = cp_group.world_size + + if world_size == 1: + if return_lse: + return cp_attn_out, cp_attn_lse + return cp_attn_out + + B, H, D = cp_attn_out.shape + if H % world_size != 0: + raise ValueError(f"H={H} must be divisible by DCP world size {world_size}.") + H_per_rank = H // world_size + lse_pack_dim = _dcp_a2a_lse_pack_dim(cp_attn_out.dtype) + + send_buffer, recv_buffer = _dcp_a2a_send_recv_buffers( + (world_size, B, H_per_rank, D + lse_pack_dim), + device=cp_attn_out.device, + dtype=cp_attn_out.dtype, + ) + + _dcp_a2a_pack_send( + cp_attn_out, + cp_attn_lse, + send_buffer, + world_size, + H_per_rank, + D, + lse_pack_dim, + seq_lens=seq_lens, + query_start_loc=query_start_loc, + ) + + work = dist.all_to_all_single( + recv_buffer.view(-1), + send_buffer.view(-1), + group=cp_group.device_group, + async_op=True, + ) + work.wait() + + return _dcp_a2a_unpack_combine( + recv_buffer, D, lse_pack_dim, return_lse, is_lse_base_on_e + ) + + +def get_dcp_workspace_max_num_tokens(vllm_config: VllmConfig) -> int: + scheduler_config = vllm_config.scheduler_config + speculative_config = vllm_config.speculative_config + speculative_tokens = vllm_config.num_speculative_tokens + tokens_per_seq = ( + 1 + + ( + 2 + if speculative_config is not None and speculative_config.parallel_drafting + else 1 + ) + * speculative_tokens + ) + return min( + scheduler_config.max_num_batched_tokens, + max( + scheduler_config.max_num_seqs * tokens_per_seq, + vllm_config.compilation_config.max_cudagraph_capture_size or 0, + ), + ) + + +def reserve_query_head_storage( + query: torch.Tensor, padded_num_heads: int +) -> torch.Tensor: + """Reserve backing storage for fixed-head decode kernels.""" + assert query.ndim == 3 + assert query.shape[1] <= padded_num_heads + padded = query.new_empty((query.shape[0], padded_num_heads, query.shape[2])) + padded.resize_(query.shape) + padded.copy_(query) + return padded + + +# Symmetric-memory A2A implementation + + +_A2A_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16) + + +class DirectDCPA2AWorkspace(DirectCPWorkspace): + """Persistent symmetric buffers for direct DCP output exchange.""" + + def __init__( + self, + group: ProcessGroup, + device: torch.device, + max_num_tokens: int, + heads_per_rank: int, + head_dim: int, + dtype: torch.dtype = torch.bfloat16, + num_ubatches: int = 1, + ) -> None: + if dtype not in _A2A_SUPPORTED_DTYPES: + raise ValueError(f"Direct DCP A2A does not support {dtype}") + if num_ubatches < 1: + raise ValueError( + f"Direct DCP A2A requires at least one ubatch slot, got {num_ubatches}" + ) + super().__init__(group, device, num_ubatches) + self.max_num_tokens = max_num_tokens + self.heads_per_rank = heads_per_rank + self.head_dim = head_dim + + output_shape = ( + num_ubatches, + 2, + self.world_size, + max_num_tokens, + heads_per_rank, + head_dim, + ) + lse_shape = ( + num_ubatches, + 2, + self.world_size, + max_num_tokens, + heads_per_rank, + ) + signal_shape = (num_ubatches, 2, self.world_size) + self.received_output, self.peer_output_ptrs = self._allocate( + output_shape, dtype + ) + self.received_lse, self.peer_lse_ptrs = self._allocate(lse_shape, torch.float32) + self.received_signal, self.peer_signal_ptrs = self._allocate( + signal_shape, torch.int32 + ) + + def lse_reduce( + self, + partial_output: torch.Tensor, + partial_lse: torch.Tensor, + is_lse_base_on_e: bool, + seq_lens: torch.Tensor | None = None, + query_start_loc: torch.Tensor | None = None, + ) -> torch.Tensor: + ubatch = dbo_current_ubatch_id() + num_tokens = partial_output.shape[0] + output = partial_output.new_empty( + (num_tokens, self.heads_per_rank, self.head_dim) + ) + torch.ops._C.direct_dcp_a2a_lse_reduce( + partial_output, + partial_lse, + seq_lens, + query_start_loc, + self.peer_output_ptrs[ubatch], + self.peer_lse_ptrs[ubatch], + self.peer_signal_ptrs[ubatch], + self.received_output[ubatch], + self.received_lse[ubatch], + self.received_signal[ubatch], + self.epoch[ubatch : ubatch + 1], + output, + self.world_size, + self.rank, + self.max_num_tokens, + is_lse_base_on_e, + ) + return output + + +@functools.cache +def get_direct_dcp_a2a_workspace( + group: GroupCoordinator, + device: torch.device, + max_num_tokens: int, + heads_per_rank: int, + head_dim: int, + dtype: torch.dtype, + num_ubatches: int, +) -> DirectDCPA2AWorkspace | None: + if not direct_cp_enabled( + group, dtype, envs.VLLM_USE_DIRECT_DCP_A2A, _A2A_SUPPORTED_DTYPES + ): + return None + return DirectDCPA2AWorkspace( + group.device_group, + device, + max_num_tokens, + heads_per_rank, + head_dim, + dtype, + num_ubatches, + ) + + +# Q gather + +# Symmetric-memory implementation + + +def _q_gather_layout_supported( + world_size: int, + heads_per_rank: int, + head_dim: int, + dtype: torch.dtype, + padded_num_heads: int | None, +) -> bool: + element_size = torch.empty((), dtype=dtype).element_size() + gathered_num_heads = world_size * heads_per_rank + storage_num_heads = ( + gathered_num_heads if padded_num_heads is None else padded_num_heads + ) + return ( + heads_per_rank * head_dim * element_size % 16 == 0 + and storage_num_heads * head_dim * element_size % 16 == 0 + ) + + +class DirectDCPQGatherWorkspace(DirectCPWorkspace): + """Publish query shards directly into the consumer-final symmetric buffer. + + The final buffer is reusable after the downstream DCP output combine. That + combine orders all ranks after attention has consumed the gathered query. + """ + + def __init__( + self, + group: ProcessGroup, + device: torch.device, + max_num_tokens: int, + heads_per_rank: int, + head_dim: int, + dtype: torch.dtype = torch.bfloat16, + num_ubatches: int = 1, + padded_num_heads: int | None = None, + ) -> None: + if num_ubatches < 1: + raise ValueError( + "Direct DCP q-gather requires at least one ubatch slot, " + f"got {num_ubatches}" + ) + if max_num_tokens < 1 or heads_per_rank < 1 or head_dim < 1: + raise ValueError( + "Direct DCP q-gather dimensions must be positive, got " + f"T={max_num_tokens}, H={heads_per_rank}, D={head_dim}" + ) + gathered_num_heads = group.size() * heads_per_rank + if not _q_gather_layout_supported( + group.size(), heads_per_rank, head_dim, dtype, padded_num_heads + ): + raise ValueError("Direct DCP q-gather requires 16-byte-aligned query rows.") + super().__init__(group, device, num_ubatches) + if self.world_size <= 1: + raise ValueError("Direct DCP q-gather requires at least two ranks") + self.max_num_tokens = max_num_tokens + self.heads_per_rank = heads_per_rank + self.gathered_num_heads = gathered_num_heads + self.padded_num_heads = ( + self.gathered_num_heads if padded_num_heads is None else padded_num_heads + ) + if self.padded_num_heads < self.gathered_num_heads: + raise ValueError( + "Direct DCP q-gather padded heads must cover gathered heads: " + f"{self.padded_num_heads} < {self.gathered_num_heads}" + ) + self.head_dim = head_dim + + query_shape = ( + num_ubatches, + max_num_tokens, + self.padded_num_heads, + head_dim, + ) + signal_shape = (num_ubatches, 2, self.world_size) + self.final_query, _ = self._allocate(query_shape, dtype) + self.received_signal, _ = self._allocate(signal_shape, torch.int32) + query_multicast_ptrs = self._multicast_ptrs(self.final_query) + signal_multicast_ptrs = self._multicast_ptrs(self.received_signal) + self.multicast_ptrs = list( + zip(query_multicast_ptrs, signal_multicast_ptrs, strict=True) + ) + if not all( + query_ptr and signal_ptr for query_ptr, signal_ptr in self.multicast_ptrs + ): + raise RuntimeError( + "Direct DCP q-gather requires NVLS symmetric-memory multicast." + ) + self.completion = self.received_signal.new_zeros((num_ubatches, 1)) + torch.accelerator.synchronize() + + def gather(self, local_query: torch.Tensor) -> torch.Tensor: + ubatch = dbo_current_ubatch_id() + if not 0 <= ubatch < self.num_ubatches: + raise ValueError( + f"DCP q-gather ubatch {ubatch} exceeds {self.num_ubatches} slots" + ) + if local_query.ndim == 3 and local_query.shape[1] != self.heads_per_rank: + raise ValueError( + f"DCP q-gather expected {self.heads_per_rank} local query heads, " + f"got {local_query.shape[1]}" + ) + + num_tokens = local_query.shape[0] + output = torch.as_strided( + self.final_query[ubatch], + size=(num_tokens, self.gathered_num_heads, self.head_dim), + stride=( + self.gathered_num_heads * self.head_dim, + self.head_dim, + 1, + ), + ) + query_multicast_ptr, signal_multicast_ptr = self.multicast_ptrs[ubatch] + torch.ops._C.direct_dcp_q_gather( + local_query, + output, + self.received_signal[ubatch], + self.completion[ubatch], + self.epoch[ubatch : ubatch + 1], + self.world_size, + self.rank, + self.max_num_tokens, + self.padded_num_heads, + query_multicast_ptr, + signal_multicast_ptr, + ) + return output + + +@functools.cache +def get_direct_dcp_q_gather_workspace( + group: GroupCoordinator, + device: torch.device, + max_num_tokens: int, + heads_per_rank: int, + head_dim: int, + dtype: torch.dtype, + num_ubatches: int, + padded_num_heads: int | None = None, +) -> DirectDCPQGatherWorkspace | None: + if not direct_cp_multicast_enabled(group, dtype, envs.VLLM_USE_DIRECT_DCP_Q_GATHER): + return None + if not _q_gather_layout_supported( + group.world_size, heads_per_rank, head_dim, dtype, padded_num_heads + ): + return None + return DirectDCPQGatherWorkspace( + group.device_group, + device, + max_num_tokens, + heads_per_rank, + head_dim, + dtype, + num_ubatches, + padded_num_heads, + ) + + +# KV gather + +# Symmetric-memory implementation + + +_KV_GATHER_SUPPORTED_DTYPES = ( + torch.float16, + torch.bfloat16, + torch.float8_e4m3fn, +) + + +def _kv_gather_layout_supported(token_dim: int, dtype: torch.dtype) -> bool: + return token_dim * torch.empty((), dtype=dtype).element_size() % 16 == 0 + + +class DirectDCPKVGatherWorkspace(DirectCPWorkspace): + """Persistent symmetric buffers for direct DCP KV gather.""" + + def __init__( + self, + group: ProcessGroup, + device: torch.device, + max_gathered_tokens: int, + token_dim: int, + dtype: torch.dtype = torch.bfloat16, + num_ubatches: int = 1, + ) -> None: + if dtype not in _KV_GATHER_SUPPORTED_DTYPES: + raise ValueError(f"Direct DCP kv-gather does not support {dtype}") + if num_ubatches < 1: + raise ValueError( + "Direct DCP kv-gather requires at least one ubatch slot, " + f"got {num_ubatches}" + ) + if max_gathered_tokens < 1 or token_dim < 1: + raise ValueError( + "Direct DCP kv-gather dimensions must be positive, got " + f"T={max_gathered_tokens}, D={token_dim}" + ) + if not _kv_gather_layout_supported(token_dim, dtype): + raise ValueError("Direct DCP kv-gather requires 16-byte-aligned KV rows.") + super().__init__(group, device, num_ubatches) + if self.world_size <= 1: + raise ValueError("Direct DCP kv-gather requires at least two ranks") + if max_gathered_tokens % self.world_size != 0: + raise ValueError( + "Direct DCP kv-gather capacity must divide evenly across " + f"ranks: {max_gathered_tokens} % {self.world_size} != 0" + ) + self.max_gathered_tokens = max_gathered_tokens + + kv_shape = (num_ubatches, 2, max_gathered_tokens, token_dim) + signal_shape = (num_ubatches, 2, self.world_size) + self.received_kv, _ = self._allocate(kv_shape, dtype) + self.received_signal, _ = self._allocate(signal_shape, torch.int32) + kv_multicast_ptrs = self._multicast_ptrs(self.received_kv) + signal_multicast_ptrs = self._multicast_ptrs(self.received_signal) + self.multicast_ptrs = list( + zip(kv_multicast_ptrs, signal_multicast_ptrs, strict=True) + ) + if not all(kv_ptr and signal_ptr for kv_ptr, signal_ptr in self.multicast_ptrs): + raise RuntimeError( + "Direct DCP kv-gather requires NVLS symmetric-memory multicast." + ) + self.completion = self.received_signal.new_zeros((num_ubatches, 2)) + torch.accelerator.synchronize() + + def gather(self, gathered_kv: torch.Tensor, local_kv: torch.Tensor) -> None: + ubatch = dbo_current_ubatch_id() + if not 0 <= ubatch < self.num_ubatches: + raise ValueError( + f"DCP kv-gather ubatch {ubatch} exceeds {self.num_ubatches} slots" + ) + kv_multicast_ptr, signal_multicast_ptr = self.multicast_ptrs[ubatch] + torch.ops._C.direct_dcp_kv_gather( + local_kv, + self.received_kv[ubatch], + self.received_signal[ubatch], + self.completion[ubatch], + self.epoch[ubatch : ubatch + 1], + gathered_kv, + self.world_size, + self.rank, + self.max_gathered_tokens, + kv_multicast_ptr, + signal_multicast_ptr, + ) + + +@functools.cache +def get_direct_dcp_kv_gather_workspace( + group: GroupCoordinator, + device: torch.device, + max_gathered_tokens: int, + token_dim: int, + dtype: torch.dtype, + num_ubatches: int, +) -> DirectDCPKVGatherWorkspace | None: + if not direct_cp_multicast_enabled( + group, + dtype, + envs.VLLM_USE_DIRECT_DCP_KV_GATHER, + _KV_GATHER_SUPPORTED_DTYPES, + ): + return None + if not _kv_gather_layout_supported(token_dim, dtype): + return None + return DirectDCPKVGatherWorkspace( + group.device_group, + device, + max_gathered_tokens, + token_dim, + dtype, + num_ubatches, + ) + + +# MLA DCP backend selection + + +class DCPCombine(Protocol): + def __call__( + self, + partial_output: torch.Tensor, + partial_lse: torch.Tensor, + *, + seq_lens: torch.Tensor, + query_start_loc: torch.Tensor, + ) -> torch.Tensor: ... + + +class MLADCPManager: + """Select and own layer-level collective implementations for MLA DCP.""" + + _kv_gather: Callable[[torch.Tensor, torch.Tensor], object] + + def __init__( + self, + vllm_config: VllmConfig, + device: torch.device, + num_heads: int, + query_head_dim: int, + output_head_dim: int, + query_dtype: torch.dtype, + output_dtype: torch.dtype, + padded_num_heads: int | None, + is_lse_base_on_e: bool, + use_pcp: bool, + ) -> None: + parallel_config = vllm_config.parallel_config + self.group = get_dcp_group() + self.device = torch.device(device) + self.num_ubatches = max(parallel_config.num_ubatches, 1) + self.max_num_tokens = get_dcp_workspace_max_num_tokens(vllm_config) + self.use_a2a = parallel_config.dcp_comm_backend == "a2a" + self.padded_num_heads = padded_num_heads + + self.combine = self._init_combine( + num_heads, + output_head_dim, + output_dtype, + is_lse_base_on_e, + use_pcp, + ) + self.query_gather = ( + None + if use_pcp + else self._init_query_gather( + num_heads, + query_head_dim, + query_dtype, + ) + ) + + def _init_combine( + self, + num_heads: int, + head_dim: int, + dtype: torch.dtype, + is_lse_base_on_e: bool, + use_pcp: bool, + ) -> DCPCombine: + direct_workspace = None + if self.use_a2a: + direct_workspace = get_direct_dcp_a2a_workspace( + self.group, + self.device, + self.max_num_tokens, + num_heads, + head_dim, + dtype, + self.num_ubatches, + ) + if direct_workspace is not None: + logger.info_once("Using direct symmetric-memory DCP A2A for MLA.") + return functools.partial( + direct_workspace.lse_reduce, + is_lse_base_on_e=is_lse_base_on_e, + ) + + combine_fn = ( + dcp_a2a_lse_reduce + if self.use_a2a + else cp_lse_ag_out_ar + if use_pcp + else cp_lse_ag_out_rs + ) + return functools.partial( + combine_fn, + cp_group=self.group, + is_lse_base_on_e=is_lse_base_on_e, + ) + + def _init_query_gather( + self, + num_heads: int, + head_dim: int, + dtype: torch.dtype, + ) -> Callable[[torch.Tensor], torch.Tensor]: + direct_workspace = get_direct_dcp_q_gather_workspace( + self.group, + self.device, + self.max_num_tokens, + num_heads, + head_dim, + dtype, + self.num_ubatches, + self.padded_num_heads, + ) + if direct_workspace is not None: + logger.info_once("Using direct symmetric-memory DCP query gather for MLA.") + return direct_workspace.gather + return self._gather_query + + def _gather_query(self, query: torch.Tensor) -> torch.Tensor: + query = self.group.all_gather(query, dim=1) + if self.padded_num_heads is not None: + query = reserve_query_head_storage(query, self.padded_num_heads) + return query + + def init_kv_gather( + self, + workspace: torch.Tensor, + max_gathered_tokens: int, + ) -> None: + world_size = self.group.world_size + assert max_gathered_tokens > 0 + assert max_gathered_tokens % world_size == 0 + assert workspace.ndim == 2 + assert workspace.is_contiguous() + assert workspace.shape[0] == ( + max_gathered_tokens + max_gathered_tokens // world_size + ) + assert workspace.shape[1] > 0 + + direct_workspace = get_direct_dcp_kv_gather_workspace( + self.group, + workspace.device, + max_gathered_tokens, + workspace.shape[1], + workspace.dtype, + self.num_ubatches, + ) + if direct_workspace is not None: + logger.info_once( + "Using direct symmetric-memory DCP chunked-context KV gather for MLA." + ) + self._kv_gather = direct_workspace.gather + else: + self._kv_gather = functools.partial( + torch.distributed.all_gather_into_tensor, + group=self.group.device_group, + ) + + def kv_gather( + self, + gathered_kv: torch.Tensor, + local_kv: torch.Tensor, + ) -> object: + return self._kv_gather(gathered_kv, local_kv) diff --git a/vllm/v1/attention/ops/dcp_alltoall.py b/vllm/v1/attention/ops/dcp_alltoall.py deleted file mode 100644 index 5939b61a8d35..000000000000 --- a/vllm/v1/attention/ops/dcp_alltoall.py +++ /dev/null @@ -1,470 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -DCP All-to-All communication backend for attention. - -Provides All-to-All (A2A) communication as an alternative to -AllGather + ReduceScatter (AG+RS) for Decode Context Parallel (DCP). -Instead of gathering the full Q tensor and scattering partial outputs, -A2A exchanges partial attention outputs and their LSE values across -ranks, then combines them with exact LSE-weighted reduction. - -This reduces the number of NCCL calls per attention layer by exchanging -the partial output and LSE in a single packed All-to-All payload. - -Usage: - vllm serve model --tp 16 --dcp 16 --dcp-comm-backend a2a - -Reference: https://arxiv.org/abs/2507.07120 -""" - -from __future__ import annotations - -from typing import TYPE_CHECKING - -import torch -import torch.distributed as dist - -from vllm.triton_utils import tl, triton -from vllm.v1.attention.ops.common import mask_dcp_empty_shards_ - -if TYPE_CHECKING: - from vllm.distributed.parallel_state import GroupCoordinator - from vllm.v1.attention.ops.common import CPTritonContext - - -def _lse_weighted_combine( - outputs: torch.Tensor, - lses: torch.Tensor, - return_lse: bool = False, - is_lse_base_on_e: bool = True, -) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - """ - CPU reference implementation for LSE-weighted combination. - - This is a pure PyTorch implementation used for testing and validation. - - Args: - outputs: Partial attention outputs [N, B, H, D] - N = number of KV shards (ranks) - B = batch size (num_tokens) - H = number of heads per rank - D = head dimension - lses: Log-sum-exp values [N, B, H] - return_lse: If True, also return the global LSE - is_lse_base_on_e: If True, LSE is base e; if False, base 2 - - Returns: - Combined output [B, H, D], and optionally global LSE [B, H] - """ - N, B, H, D = outputs.shape - - # Handle NaN and inf in LSEs - lses = torch.where( - torch.isnan(lses) | torch.isinf(lses), - torch.tensor(float("-inf"), device=lses.device, dtype=lses.dtype), - lses, - ) - - # Compute max LSE for numerical stability - lse_max, _ = lses.max(dim=0) # [B, H] - lse_max = torch.where( - lse_max == float("-inf"), - torch.zeros_like(lse_max), - lse_max, - ) - - # Compute weights: softmax over the N dimension - if is_lse_base_on_e: - weights = torch.exp(lses - lse_max.unsqueeze(0)) # [N, B, H] - else: - weights = torch.pow(2.0, lses - lse_max.unsqueeze(0)) # [N, B, H] - - # Handle NaN weights - weights = torch.where(torch.isnan(weights), torch.zeros_like(weights), weights) - - # Normalize weights - weight_sum = weights.sum(dim=0, keepdim=True) # [1, B, H] - weights = weights / weight_sum.clamp(min=1e-10) # [N, B, H] - - # Weighted combination: sum over N dimension - weights = weights.unsqueeze(-1) - outputs = torch.where(weights == 0, torch.zeros_like(outputs), outputs) - result = (outputs * weights).sum(dim=0) # [B, H, D] - - if return_lse: - if is_lse_base_on_e: - global_lse = torch.log(weight_sum.squeeze(0)) + lse_max # [B, H] - else: - global_lse = torch.log2(weight_sum.squeeze(0)) + lse_max # [B, H] - return result, global_lse - - return result - - -def _dcp_a2a_lse_pack_dim(output_dtype: torch.dtype) -> int: - bits = torch.finfo(output_dtype).bits - if bits == 16: - return 2 - if bits == 32: - return 1 - raise ValueError(f"Cannot pack fp32 LSE into output dtype {output_dtype}.") - - -def _dcp_a2a_send_recv_buffers( - shape: tuple[int, ...], - device: torch.device, - dtype: torch.dtype, -) -> tuple[torch.Tensor, torch.Tensor]: - # Don't use the shared WorkspaceManager here. A FULL cudagraph bakes in the - # buffer address at capture, but the workspace is growable and sized only to - # the largest *captured* batch (the cudagraph capture cap). Any eager a2a - # with a bigger batch regrows it, freeing that address and poisoning every - # captured graph -> illegal memory access on replay. This bites the very - # first request: the post-capture warmup runs an eager decode at - # max_num_seqs (> the cap), so the graphs are already dangling before the - # server is ready. torch.empty buffers instead live in the graph's private - # pool and stay valid for its lifetime (as _dcp_a2a_unpack_combine and the - # AG+RS combine path already rely on). - return ( - torch.empty(shape, device=device, dtype=dtype), - torch.empty(shape, device=device, dtype=dtype), - ) - - -@triton.jit -def _dcp_a2a_pack_send_kernel( - out_ptr, - lse_ptr, - send_ptr, - out_stride_B, - out_stride_H, - out_stride_D, - lse_stride_B, - lse_stride_H, - send_stride_N, - send_stride_B, - send_stride_H, - send_stride_D, - N: tl.constexpr, - HEAD_DIM: tl.constexpr, - H_PER_RANK: tl.constexpr, - LSE_PACK_DIM: tl.constexpr, -): - batch_idx = tl.program_id(0).to(tl.int64) - local_head_idx = tl.program_id(1).to(tl.int64) - d_offsets = tl.arange(0, HEAD_DIM) - - for rank_idx in tl.static_range(N): - src_head_idx = rank_idx * H_PER_RANK + local_head_idx - send_base = ( - rank_idx * send_stride_N - + batch_idx * send_stride_B - + local_head_idx * send_stride_H - ) - - out_offsets = ( - batch_idx * out_stride_B - + src_head_idx * out_stride_H - + d_offsets * out_stride_D - ) - tl.store( - send_ptr + send_base + d_offsets * send_stride_D, - tl.load(out_ptr + out_offsets), - ) - - lse_val = tl.load( - lse_ptr + batch_idx * lse_stride_B + src_head_idx * lse_stride_H - ).to(tl.float32) - if LSE_PACK_DIM == 1: - tl.store( - send_ptr + send_base + HEAD_DIM * send_stride_D, - lse_val.to(send_ptr.dtype.element_ty), - ) - else: - lse_bits = lse_val.to(tl.uint32, bitcast=True) - lo = (lse_bits & 0xFFFF).to(tl.uint16) - hi = ((lse_bits >> 16) & 0xFFFF).to(tl.uint16) - tl.store( - send_ptr + send_base + HEAD_DIM * send_stride_D, - lo.to(send_ptr.dtype.element_ty, bitcast=True), - ) - tl.store( - send_ptr + send_base + (HEAD_DIM + 1) * send_stride_D, - hi.to(send_ptr.dtype.element_ty, bitcast=True), - ) - - -@triton.jit -def _dcp_a2a_unpack_combine_kernel( - recv_ptr, - out_ptr, - out_lse_ptr, - recv_stride_N, - recv_stride_B, - recv_stride_H, - recv_stride_D, - out_stride_B, - out_stride_H, - out_stride_D, - out_lse_stride_B, - out_lse_stride_H, - N: tl.constexpr, - HEAD_DIM: tl.constexpr, - IS_BASE_E: tl.constexpr, - RETURN_LSE: tl.constexpr, - LSE_PACK_DIM: tl.constexpr, -): - batch_idx = tl.program_id(0).to(tl.int64) - head_idx = tl.program_id(1).to(tl.int64) - d_offsets = tl.arange(0, HEAD_DIM) - - lse_max = -float("inf") - for rank_idx in tl.static_range(N): - recv_base = ( - rank_idx * recv_stride_N - + batch_idx * recv_stride_B - + head_idx * recv_stride_H - ) - if LSE_PACK_DIM == 1: - lse_val = tl.load(recv_ptr + recv_base + HEAD_DIM * recv_stride_D).to( - tl.float32 - ) - else: - lo_raw = tl.load(recv_ptr + recv_base + HEAD_DIM * recv_stride_D) - hi_raw = tl.load(recv_ptr + recv_base + (HEAD_DIM + 1) * recv_stride_D) - lo = lo_raw.to(tl.uint16, bitcast=True).to(tl.uint32) - hi = hi_raw.to(tl.uint16, bitcast=True).to(tl.uint32) - lse_val = (lo | (hi << 16)).to(tl.float32, bitcast=True) - lse_val = tl.where( - (lse_val != lse_val) | (lse_val == float("inf")), - -float("inf"), - lse_val, - ) - lse_max = tl.maximum(lse_max, lse_val) - - lse_max = tl.where(lse_max == -float("inf"), 0.0, lse_max) - - lse_sum = 0.0 - for rank_idx in tl.static_range(N): - recv_base = ( - rank_idx * recv_stride_N - + batch_idx * recv_stride_B - + head_idx * recv_stride_H - ) - if LSE_PACK_DIM == 1: - lse_val = tl.load(recv_ptr + recv_base + HEAD_DIM * recv_stride_D).to( - tl.float32 - ) - else: - lo_raw = tl.load(recv_ptr + recv_base + HEAD_DIM * recv_stride_D) - hi_raw = tl.load(recv_ptr + recv_base + (HEAD_DIM + 1) * recv_stride_D) - lo = lo_raw.to(tl.uint16, bitcast=True).to(tl.uint32) - hi = hi_raw.to(tl.uint16, bitcast=True).to(tl.uint32) - lse_val = (lo | (hi << 16)).to(tl.float32, bitcast=True) - lse_val = tl.where( - (lse_val != lse_val) | (lse_val == float("inf")), - -float("inf"), - lse_val, - ) - if IS_BASE_E: - lse_sum += tl.exp(lse_val - lse_max) - else: - lse_sum += tl.exp2(lse_val - lse_max) - - if IS_BASE_E: # noqa: SIM108 - global_lse = tl.log(lse_sum) + lse_max - else: - global_lse = tl.log2(lse_sum) + lse_max - - acc = tl.zeros([HEAD_DIM], dtype=tl.float32) - for rank_idx in tl.static_range(N): - recv_base = ( - rank_idx * recv_stride_N - + batch_idx * recv_stride_B - + head_idx * recv_stride_H - ) - if LSE_PACK_DIM == 1: - lse_val = tl.load(recv_ptr + recv_base + HEAD_DIM * recv_stride_D).to( - tl.float32 - ) - else: - lo_raw = tl.load(recv_ptr + recv_base + HEAD_DIM * recv_stride_D) - hi_raw = tl.load(recv_ptr + recv_base + (HEAD_DIM + 1) * recv_stride_D) - lo = lo_raw.to(tl.uint16, bitcast=True).to(tl.uint32) - hi = hi_raw.to(tl.uint16, bitcast=True).to(tl.uint32) - lse_val = (lo | (hi << 16)).to(tl.float32, bitcast=True) - lse_val = tl.where( - (lse_val != lse_val) | (lse_val == float("inf")), - -float("inf"), - lse_val, - ) - if IS_BASE_E: - weight = tl.exp(lse_val - global_lse) - else: - weight = tl.exp2(lse_val - global_lse) - weight = tl.where(weight != weight, 0.0, weight) - partial = tl.load(recv_ptr + recv_base + d_offsets * recv_stride_D).to( - tl.float32 - ) - partial = tl.where(weight == 0.0, 0.0, partial) - acc += partial * weight - - final_offsets = ( - batch_idx * out_stride_B + head_idx * out_stride_H + d_offsets * out_stride_D - ) - tl.store(out_ptr + final_offsets, acc) - - if RETURN_LSE: - out_lse_offset = batch_idx * out_lse_stride_B + head_idx * out_lse_stride_H - tl.store(out_lse_ptr + out_lse_offset, global_lse) - - -def _dcp_a2a_pack_send( - cp_attn_out: torch.Tensor, - cp_attn_lse: torch.Tensor, - send_buffer: torch.Tensor, - world_size: int, - h_per_rank: int, - head_dim: int, - lse_pack_dim: int, - seq_lens: torch.Tensor | None = None, - query_start_loc: torch.Tensor | None = None, -) -> None: - mask_dcp_empty_shards_(cp_attn_lse, seq_lens, query_start_loc) - grid = (cp_attn_out.shape[0], h_per_rank, 1) - _dcp_a2a_pack_send_kernel[grid]( - cp_attn_out, - cp_attn_lse, - send_buffer, - cp_attn_out.stride(0), - cp_attn_out.stride(1), - cp_attn_out.stride(2), - cp_attn_lse.stride(0), - cp_attn_lse.stride(1), - send_buffer.stride(0), - send_buffer.stride(1), - send_buffer.stride(2), - send_buffer.stride(3), - N=world_size, - HEAD_DIM=head_dim, - H_PER_RANK=h_per_rank, - LSE_PACK_DIM=lse_pack_dim, - ) - - -def _dcp_a2a_unpack_combine( - recv_buffer: torch.Tensor, - head_dim: int, - lse_pack_dim: int, - return_lse: bool, - is_lse_base_on_e: bool, -) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - world_size, num_tokens, h_per_rank, _ = recv_buffer.shape - out = torch.empty( - (num_tokens, h_per_rank, head_dim), - device=recv_buffer.device, - dtype=recv_buffer.dtype, - ) - out_lse = torch.empty( - (num_tokens, h_per_rank) if return_lse else (1, 1), - device=recv_buffer.device, - dtype=torch.float32 if return_lse else recv_buffer.dtype, - ) - grid = (num_tokens, h_per_rank, 1) - _dcp_a2a_unpack_combine_kernel[grid]( - recv_buffer, - out, - out_lse, - recv_buffer.stride(0), - recv_buffer.stride(1), - recv_buffer.stride(2), - recv_buffer.stride(3), - out.stride(0), - out.stride(1), - out.stride(2), - out_lse.stride(0), - out_lse.stride(1), - N=world_size, - HEAD_DIM=head_dim, - IS_BASE_E=is_lse_base_on_e, - RETURN_LSE=return_lse, - LSE_PACK_DIM=lse_pack_dim, - ) - if return_lse: - return out, out_lse - return out - - -def dcp_a2a_lse_reduce( - cp_attn_out: torch.Tensor, - cp_attn_lse: torch.Tensor, - cp_group: GroupCoordinator, - ctx: CPTritonContext | None = None, - return_lse: bool = False, - is_lse_base_on_e: bool = True, - seq_lens: torch.Tensor | None = None, - query_start_loc: torch.Tensor | None = None, -) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - """ - Combine partial attention outputs across DCP ranks using All-to-All. - - The output and LSE are packed into a single output-dtype buffer, sent - with one All-to-All, then unpacked and combined with exact LSE weighting. - - Args: - cp_attn_out: [B, H, D] where B=num_tokens, H=total_heads, D=head_dim - cp_attn_lse: [B, H] floating-point log-sum-exp values - cp_group: GroupCoordinator for DCP communication - ctx: CPTritonContext (unused, for signature compatibility) - return_lse: If True, also return the combined global LSE - is_lse_base_on_e: If True, LSE is base e; if False, base 2 - seq_lens: Local KV lengths. Empty shards contribute zero weight. - query_start_loc: Cumulative query-token offsets for each request. - - Returns: - Combined output [B, H/N, D] (head-scattered) - If return_lse=True, also returns global_lse [B, H/N] - """ - world_size = cp_group.world_size - - if world_size == 1: - if return_lse: - return cp_attn_out, cp_attn_lse - return cp_attn_out - - B, H, D = cp_attn_out.shape - if H % world_size != 0: - raise ValueError(f"H={H} must be divisible by DCP world size {world_size}.") - H_per_rank = H // world_size - lse_pack_dim = _dcp_a2a_lse_pack_dim(cp_attn_out.dtype) - - send_buffer, recv_buffer = _dcp_a2a_send_recv_buffers( - (world_size, B, H_per_rank, D + lse_pack_dim), - device=cp_attn_out.device, - dtype=cp_attn_out.dtype, - ) - - _dcp_a2a_pack_send( - cp_attn_out, - cp_attn_lse, - send_buffer, - world_size, - H_per_rank, - D, - lse_pack_dim, - seq_lens=seq_lens, - query_start_loc=query_start_loc, - ) - - work = dist.all_to_all_single( - recv_buffer.view(-1), - send_buffer.view(-1), - group=cp_group.device_group, - async_op=True, - ) - work.wait() - - return _dcp_a2a_unpack_combine( - recv_buffer, D, lse_pack_dim, return_lse, is_lse_base_on_e - ) diff --git a/vllm/v1/attention/ops/dcp_utils.py b/vllm/v1/attention/ops/dcp_utils.py deleted file mode 100644 index b28d337a1046..000000000000 --- a/vllm/v1/attention/ops/dcp_utils.py +++ /dev/null @@ -1,740 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""MLA DCP collective selection and direct symmetric-memory implementations.""" - -from __future__ import annotations - -import functools -from collections.abc import Callable -from typing import TYPE_CHECKING, Any, Protocol - -import torch - -import vllm.envs as envs -from vllm.config import VllmConfig -from vllm.distributed import get_dcp_group -from vllm.distributed.parallel_state import in_the_same_node_as -from vllm.logger import init_logger -from vllm.platforms import current_platform -from vllm.v1.attention.ops.common import cp_lse_ag_out_ar, cp_lse_ag_out_rs -from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce -from vllm.v1.worker.ubatching import dbo_current_ubatch_id - -logger = init_logger(__name__) - -if TYPE_CHECKING: - from torch.distributed import ProcessGroup - - from vllm.distributed.parallel_state import GroupCoordinator - -try: - import torch.distributed._symmetric_memory as symm_mem - - symm_mem_available = True -except ImportError: - symm_mem = None # type: ignore[assignment] - symm_mem_available = False - - -@functools.cache -def _symm_mem_spans_group(group: GroupCoordinator) -> bool: - """Probe whether the group has NVLS symmetric memory.""" - if not symm_mem_available: - return False - try: - from torch._C._autograd import DeviceType - from torch._C._distributed_c10d import _SymmetricMemory - - device = torch.device("cuda", torch.accelerator.current_device_index()) - if not _SymmetricMemory.has_multicast_support(DeviceType.CUDA, device.index): - return False - probe = symm_mem.empty(8, dtype=torch.uint8, device=device) - probe.zero_() - torch.accelerator.synchronize() - handle = symm_mem.rendezvous(probe, group.device_group.group_name) - spans = handle is not None and handle.multicast_ptr != 0 - except Exception as error: - logger.debug("Direct DCP symmetric-memory probe failed: %s", error) - return False - logger.debug_once( - "Direct DCP symmetric memory across %d ranks: %s", - group.world_size, - "available" if spans else "unavailable", - ) - return spans - - -def _direct_dcp_enabled( - group: GroupCoordinator, - dtype: torch.dtype, - use_direct: bool | None, - supported_dtypes: tuple[torch.dtype, ...] | None = None, -) -> bool: - if use_direct is not None: - return use_direct - return ( - symm_mem_available - and current_platform.is_cuda() - and (supported_dtypes is None or dtype in supported_dtypes) - and ( - all(in_the_same_node_as(group.cpu_group, source_rank=0)) - or _symm_mem_spans_group(group) - ) - ) - - -def _direct_dcp_multicast_enabled( - group: GroupCoordinator, - dtype: torch.dtype, - use_direct: bool | None, - supported_dtypes: tuple[torch.dtype, ...] | None = None, -) -> bool: - return _direct_dcp_enabled( - group, dtype, use_direct, supported_dtypes - ) and _symm_mem_spans_group(group) - - -def get_dcp_workspace_max_num_tokens(vllm_config: VllmConfig) -> int: - scheduler_config = vllm_config.scheduler_config - speculative_config = vllm_config.speculative_config - speculative_tokens = vllm_config.num_speculative_tokens - tokens_per_seq = ( - 1 - + ( - 2 - if speculative_config is not None and speculative_config.parallel_drafting - else 1 - ) - * speculative_tokens - ) - return min( - scheduler_config.max_num_batched_tokens, - max( - scheduler_config.max_num_seqs * tokens_per_seq, - vllm_config.compilation_config.max_cudagraph_capture_size or 0, - ), - ) - - -class _DirectDCPWorkspace: - def __init__( - self, - group: ProcessGroup, - device: torch.device, - num_ubatches: int, - ) -> None: - self.group = group - self.world_size = group.size() - self.rank = group.rank() - self.device = torch.device(device) - self.num_ubatches = num_ubatches - self.epoch = torch.zeros(num_ubatches, dtype=torch.int64, device=self.device) - self._allocations: list[tuple[torch.Tensor, Any, list[torch.Tensor]]] = [] - - def _allocate( - self, shape: tuple[int, ...], dtype: torch.dtype - ) -> tuple[torch.Tensor, torch.Tensor]: - storage = symm_mem.empty(shape, device=self.device, dtype=dtype) - storage.zero_() - torch.accelerator.synchronize() - handle = symm_mem.rendezvous(storage, self.group.group_name) - assert handle is not None, "DCP symmetric memory rendezvous returned None" - handle.barrier() - views = [ - handle.get_buffer(peer, list(shape), dtype, 0) - for peer in range(self.world_size) - ] - self.device = storage.device - peer_ptrs = torch.tensor( - [ - [view[ubatch].data_ptr() for view in views] - for ubatch in range(self.num_ubatches) - ], - dtype=torch.int64, - device=self.device, - ) - self._allocations.append((storage, handle, views)) - return storage, peer_ptrs - - def _multicast_ptrs(self, storage: torch.Tensor) -> list[int]: - disabled = [0] * self.num_ubatches - for allocated, handle, _ in self._allocations: - if allocated is storage: - break - else: - return disabled - try: - from torch._C._autograd import DeviceType - from torch._C._distributed_c10d import _SymmetricMemory - - if not _SymmetricMemory.has_multicast_support( - DeviceType.CUDA, storage.device.index - ): - return disabled - multicast_base = handle.multicast_ptr - except Exception: - return disabled - if not multicast_base: - return disabled - storage_base = storage.data_ptr() - return [ - multicast_base + (storage[ubatch].data_ptr() - storage_base) - for ubatch in range(self.num_ubatches) - ] - - -def reserve_query_head_storage( - query: torch.Tensor, padded_num_heads: int -) -> torch.Tensor: - """Reserve backing storage for fixed-head decode kernels.""" - assert query.ndim == 3 - assert query.shape[1] <= padded_num_heads - padded = query.new_empty((query.shape[0], padded_num_heads, query.shape[2])) - padded.resize_(query.shape) - padded.copy_(query) - return padded - - -_A2A_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16) - - -class DirectDCPA2AWorkspace(_DirectDCPWorkspace): - """Persistent symmetric buffers for direct DCP output exchange.""" - - def __init__( - self, - group: ProcessGroup, - device: torch.device, - max_num_tokens: int, - heads_per_rank: int, - head_dim: int, - dtype: torch.dtype = torch.bfloat16, - num_ubatches: int = 1, - ) -> None: - if dtype not in _A2A_SUPPORTED_DTYPES: - raise ValueError(f"Direct DCP A2A does not support {dtype}") - if num_ubatches < 1: - raise ValueError( - f"Direct DCP A2A requires at least one ubatch slot, got {num_ubatches}" - ) - super().__init__(group, device, num_ubatches) - self.max_num_tokens = max_num_tokens - self.heads_per_rank = heads_per_rank - self.head_dim = head_dim - - output_shape = ( - num_ubatches, - 2, - self.world_size, - max_num_tokens, - heads_per_rank, - head_dim, - ) - lse_shape = ( - num_ubatches, - 2, - self.world_size, - max_num_tokens, - heads_per_rank, - ) - signal_shape = (num_ubatches, 2, self.world_size) - self.received_output, self.peer_output_ptrs = self._allocate( - output_shape, dtype - ) - self.received_lse, self.peer_lse_ptrs = self._allocate(lse_shape, torch.float32) - self.received_signal, self.peer_signal_ptrs = self._allocate( - signal_shape, torch.int32 - ) - - def lse_reduce( - self, - partial_output: torch.Tensor, - partial_lse: torch.Tensor, - is_lse_base_on_e: bool, - seq_lens: torch.Tensor | None = None, - query_start_loc: torch.Tensor | None = None, - ) -> torch.Tensor: - ubatch = dbo_current_ubatch_id() - num_tokens = partial_output.shape[0] - output = partial_output.new_empty( - (num_tokens, self.heads_per_rank, self.head_dim) - ) - torch.ops._C.direct_dcp_a2a_lse_reduce( - partial_output, - partial_lse, - seq_lens, - query_start_loc, - self.peer_output_ptrs[ubatch], - self.peer_lse_ptrs[ubatch], - self.peer_signal_ptrs[ubatch], - self.received_output[ubatch], - self.received_lse[ubatch], - self.received_signal[ubatch], - self.epoch[ubatch : ubatch + 1], - output, - self.world_size, - self.rank, - self.max_num_tokens, - is_lse_base_on_e, - ) - return output - - -@functools.cache -def get_direct_dcp_a2a_workspace( - group: GroupCoordinator, - device: torch.device, - max_num_tokens: int, - heads_per_rank: int, - head_dim: int, - dtype: torch.dtype, - num_ubatches: int, -) -> DirectDCPA2AWorkspace | None: - if not _direct_dcp_enabled( - group, dtype, envs.VLLM_USE_DIRECT_DCP_A2A, _A2A_SUPPORTED_DTYPES - ): - return None - return DirectDCPA2AWorkspace( - group.device_group, - device, - max_num_tokens, - heads_per_rank, - head_dim, - dtype, - num_ubatches, - ) - - -def _q_gather_layout_supported( - world_size: int, - heads_per_rank: int, - head_dim: int, - dtype: torch.dtype, - padded_num_heads: int | None, -) -> bool: - element_size = torch.empty((), dtype=dtype).element_size() - gathered_num_heads = world_size * heads_per_rank - storage_num_heads = ( - gathered_num_heads if padded_num_heads is None else padded_num_heads - ) - return ( - heads_per_rank * head_dim * element_size % 16 == 0 - and storage_num_heads * head_dim * element_size % 16 == 0 - ) - - -class DirectDCPQGatherWorkspace(_DirectDCPWorkspace): - """Publish query shards directly into the consumer-final symmetric buffer. - - The final buffer is reusable after the downstream DCP output combine. That - combine orders all ranks after attention has consumed the gathered query. - """ - - def __init__( - self, - group: ProcessGroup, - device: torch.device, - max_num_tokens: int, - heads_per_rank: int, - head_dim: int, - dtype: torch.dtype = torch.bfloat16, - num_ubatches: int = 1, - padded_num_heads: int | None = None, - ) -> None: - if num_ubatches < 1: - raise ValueError( - "Direct DCP q-gather requires at least one ubatch slot, " - f"got {num_ubatches}" - ) - if max_num_tokens < 1 or heads_per_rank < 1 or head_dim < 1: - raise ValueError( - "Direct DCP q-gather dimensions must be positive, got " - f"T={max_num_tokens}, H={heads_per_rank}, D={head_dim}" - ) - gathered_num_heads = group.size() * heads_per_rank - if not _q_gather_layout_supported( - group.size(), heads_per_rank, head_dim, dtype, padded_num_heads - ): - raise ValueError("Direct DCP q-gather requires 16-byte-aligned query rows.") - super().__init__(group, device, num_ubatches) - if self.world_size <= 1: - raise ValueError("Direct DCP q-gather requires at least two ranks") - self.max_num_tokens = max_num_tokens - self.heads_per_rank = heads_per_rank - self.gathered_num_heads = gathered_num_heads - self.padded_num_heads = ( - self.gathered_num_heads if padded_num_heads is None else padded_num_heads - ) - if self.padded_num_heads < self.gathered_num_heads: - raise ValueError( - "Direct DCP q-gather padded heads must cover gathered heads: " - f"{self.padded_num_heads} < {self.gathered_num_heads}" - ) - self.head_dim = head_dim - - query_shape = ( - num_ubatches, - max_num_tokens, - self.padded_num_heads, - head_dim, - ) - signal_shape = (num_ubatches, 2, self.world_size) - self.final_query, _ = self._allocate(query_shape, dtype) - self.received_signal, _ = self._allocate(signal_shape, torch.int32) - query_multicast_ptrs = self._multicast_ptrs(self.final_query) - signal_multicast_ptrs = self._multicast_ptrs(self.received_signal) - self.multicast_ptrs = list( - zip(query_multicast_ptrs, signal_multicast_ptrs, strict=True) - ) - if not all( - query_ptr and signal_ptr for query_ptr, signal_ptr in self.multicast_ptrs - ): - raise RuntimeError( - "Direct DCP q-gather requires NVLS symmetric-memory multicast." - ) - self.completion = self.received_signal.new_zeros((num_ubatches, 1)) - torch.accelerator.synchronize() - - def gather(self, local_query: torch.Tensor) -> torch.Tensor: - ubatch = dbo_current_ubatch_id() - if not 0 <= ubatch < self.num_ubatches: - raise ValueError( - f"DCP q-gather ubatch {ubatch} exceeds {self.num_ubatches} slots" - ) - if local_query.ndim == 3 and local_query.shape[1] != self.heads_per_rank: - raise ValueError( - f"DCP q-gather expected {self.heads_per_rank} local query heads, " - f"got {local_query.shape[1]}" - ) - - num_tokens = local_query.shape[0] - output = torch.as_strided( - self.final_query[ubatch], - size=(num_tokens, self.gathered_num_heads, self.head_dim), - stride=( - self.gathered_num_heads * self.head_dim, - self.head_dim, - 1, - ), - ) - query_multicast_ptr, signal_multicast_ptr = self.multicast_ptrs[ubatch] - torch.ops._C.direct_dcp_q_gather( - local_query, - output, - self.received_signal[ubatch], - self.completion[ubatch], - self.epoch[ubatch : ubatch + 1], - self.world_size, - self.rank, - self.max_num_tokens, - self.padded_num_heads, - query_multicast_ptr, - signal_multicast_ptr, - ) - return output - - -@functools.cache -def get_direct_dcp_q_gather_workspace( - group: GroupCoordinator, - device: torch.device, - max_num_tokens: int, - heads_per_rank: int, - head_dim: int, - dtype: torch.dtype, - num_ubatches: int, - padded_num_heads: int | None = None, -) -> DirectDCPQGatherWorkspace | None: - if not _direct_dcp_multicast_enabled( - group, dtype, envs.VLLM_USE_DIRECT_DCP_Q_GATHER - ): - return None - if not _q_gather_layout_supported( - group.world_size, heads_per_rank, head_dim, dtype, padded_num_heads - ): - return None - return DirectDCPQGatherWorkspace( - group.device_group, - device, - max_num_tokens, - heads_per_rank, - head_dim, - dtype, - num_ubatches, - padded_num_heads, - ) - - -_KV_GATHER_SUPPORTED_DTYPES = ( - torch.float16, - torch.bfloat16, - torch.float8_e4m3fn, -) - - -def _kv_gather_layout_supported(token_dim: int, dtype: torch.dtype) -> bool: - return token_dim * torch.empty((), dtype=dtype).element_size() % 16 == 0 - - -class DirectDCPKVGatherWorkspace(_DirectDCPWorkspace): - """Persistent symmetric buffers for direct DCP KV gather.""" - - def __init__( - self, - group: ProcessGroup, - device: torch.device, - max_gathered_tokens: int, - token_dim: int, - dtype: torch.dtype = torch.bfloat16, - num_ubatches: int = 1, - ) -> None: - if dtype not in _KV_GATHER_SUPPORTED_DTYPES: - raise ValueError(f"Direct DCP kv-gather does not support {dtype}") - if num_ubatches < 1: - raise ValueError( - "Direct DCP kv-gather requires at least one ubatch slot, " - f"got {num_ubatches}" - ) - if max_gathered_tokens < 1 or token_dim < 1: - raise ValueError( - "Direct DCP kv-gather dimensions must be positive, got " - f"T={max_gathered_tokens}, D={token_dim}" - ) - if not _kv_gather_layout_supported(token_dim, dtype): - raise ValueError("Direct DCP kv-gather requires 16-byte-aligned KV rows.") - super().__init__(group, device, num_ubatches) - if self.world_size <= 1: - raise ValueError("Direct DCP kv-gather requires at least two ranks") - if max_gathered_tokens % self.world_size != 0: - raise ValueError( - "Direct DCP kv-gather capacity must divide evenly across " - f"ranks: {max_gathered_tokens} % {self.world_size} != 0" - ) - self.max_gathered_tokens = max_gathered_tokens - - kv_shape = (num_ubatches, 2, max_gathered_tokens, token_dim) - signal_shape = (num_ubatches, 2, self.world_size) - self.received_kv, _ = self._allocate(kv_shape, dtype) - self.received_signal, _ = self._allocate(signal_shape, torch.int32) - kv_multicast_ptrs = self._multicast_ptrs(self.received_kv) - signal_multicast_ptrs = self._multicast_ptrs(self.received_signal) - self.multicast_ptrs = list( - zip(kv_multicast_ptrs, signal_multicast_ptrs, strict=True) - ) - if not all(kv_ptr and signal_ptr for kv_ptr, signal_ptr in self.multicast_ptrs): - raise RuntimeError( - "Direct DCP kv-gather requires NVLS symmetric-memory multicast." - ) - self.completion = self.received_signal.new_zeros((num_ubatches, 2)) - torch.accelerator.synchronize() - - def gather(self, gathered_kv: torch.Tensor, local_kv: torch.Tensor) -> None: - ubatch = dbo_current_ubatch_id() - if not 0 <= ubatch < self.num_ubatches: - raise ValueError( - f"DCP kv-gather ubatch {ubatch} exceeds {self.num_ubatches} slots" - ) - kv_multicast_ptr, signal_multicast_ptr = self.multicast_ptrs[ubatch] - torch.ops._C.direct_dcp_kv_gather( - local_kv, - self.received_kv[ubatch], - self.received_signal[ubatch], - self.completion[ubatch], - self.epoch[ubatch : ubatch + 1], - gathered_kv, - self.world_size, - self.rank, - self.max_gathered_tokens, - kv_multicast_ptr, - signal_multicast_ptr, - ) - - -@functools.cache -def get_direct_dcp_kv_gather_workspace( - group: GroupCoordinator, - device: torch.device, - max_gathered_tokens: int, - token_dim: int, - dtype: torch.dtype, - num_ubatches: int, -) -> DirectDCPKVGatherWorkspace | None: - if not _direct_dcp_multicast_enabled( - group, - dtype, - envs.VLLM_USE_DIRECT_DCP_KV_GATHER, - _KV_GATHER_SUPPORTED_DTYPES, - ): - return None - if not _kv_gather_layout_supported(token_dim, dtype): - return None - return DirectDCPKVGatherWorkspace( - group.device_group, - device, - max_gathered_tokens, - token_dim, - dtype, - num_ubatches, - ) - - -class DCPCombine(Protocol): - def __call__( - self, - partial_output: torch.Tensor, - partial_lse: torch.Tensor, - *, - seq_lens: torch.Tensor, - query_start_loc: torch.Tensor, - ) -> torch.Tensor: ... - - -class MLADCPManager: - """Select and own layer-level collective implementations for MLA DCP.""" - - _kv_gather: Callable[[torch.Tensor, torch.Tensor], object] - - def __init__( - self, - vllm_config: VllmConfig, - device: torch.device, - num_heads: int, - query_head_dim: int, - output_head_dim: int, - query_dtype: torch.dtype, - output_dtype: torch.dtype, - padded_num_heads: int | None, - is_lse_base_on_e: bool, - use_pcp: bool, - ) -> None: - parallel_config = vllm_config.parallel_config - self.group = get_dcp_group() - self.device = torch.device(device) - self.num_ubatches = max(parallel_config.num_ubatches, 1) - self.max_num_tokens = get_dcp_workspace_max_num_tokens(vllm_config) - self.use_a2a = parallel_config.dcp_comm_backend == "a2a" - self.padded_num_heads = padded_num_heads - - self.combine = self._init_combine( - num_heads, - output_head_dim, - output_dtype, - is_lse_base_on_e, - use_pcp, - ) - self.query_gather = ( - None - if use_pcp - else self._init_query_gather( - num_heads, - query_head_dim, - query_dtype, - ) - ) - - def _init_combine( - self, - num_heads: int, - head_dim: int, - dtype: torch.dtype, - is_lse_base_on_e: bool, - use_pcp: bool, - ) -> DCPCombine: - direct_workspace = None - if self.use_a2a: - direct_workspace = get_direct_dcp_a2a_workspace( - self.group, - self.device, - self.max_num_tokens, - num_heads, - head_dim, - dtype, - self.num_ubatches, - ) - if direct_workspace is not None: - logger.info_once("Using direct symmetric-memory DCP A2A for MLA.") - return functools.partial( - direct_workspace.lse_reduce, - is_lse_base_on_e=is_lse_base_on_e, - ) - - combine_fn = ( - dcp_a2a_lse_reduce - if self.use_a2a - else cp_lse_ag_out_ar - if use_pcp - else cp_lse_ag_out_rs - ) - return functools.partial( - combine_fn, - cp_group=self.group, - is_lse_base_on_e=is_lse_base_on_e, - ) - - def _init_query_gather( - self, - num_heads: int, - head_dim: int, - dtype: torch.dtype, - ) -> Callable[[torch.Tensor], torch.Tensor]: - direct_workspace = get_direct_dcp_q_gather_workspace( - self.group, - self.device, - self.max_num_tokens, - num_heads, - head_dim, - dtype, - self.num_ubatches, - self.padded_num_heads, - ) - if direct_workspace is not None: - logger.info_once("Using direct symmetric-memory DCP query gather for MLA.") - return direct_workspace.gather - return self._gather_query - - def _gather_query(self, query: torch.Tensor) -> torch.Tensor: - query = self.group.all_gather(query, dim=1) - if self.padded_num_heads is not None: - query = reserve_query_head_storage(query, self.padded_num_heads) - return query - - def init_kv_gather( - self, - workspace: torch.Tensor, - max_gathered_tokens: int, - ) -> None: - world_size = self.group.world_size - assert max_gathered_tokens > 0 - assert max_gathered_tokens % world_size == 0 - assert workspace.ndim == 2 - assert workspace.is_contiguous() - assert workspace.shape[0] == ( - max_gathered_tokens + max_gathered_tokens // world_size - ) - assert workspace.shape[1] > 0 - - direct_workspace = get_direct_dcp_kv_gather_workspace( - self.group, - workspace.device, - max_gathered_tokens, - workspace.shape[1], - workspace.dtype, - self.num_ubatches, - ) - if direct_workspace is not None: - logger.info_once( - "Using direct symmetric-memory DCP chunked-context KV gather for MLA." - ) - self._kv_gather = direct_workspace.gather - else: - self._kv_gather = functools.partial( - torch.distributed.all_gather_into_tensor, - group=self.group.device_group, - ) - - def kv_gather( - self, - gathered_kv: torch.Tensor, - local_kv: torch.Tensor, - ) -> object: - return self._kv_gather(gathered_kv, local_kv) diff --git a/vllm/model_executor/layers/attention/pcp.py b/vllm/v1/attention/ops/pcp.py similarity index 100% rename from vllm/model_executor/layers/attention/pcp.py rename to vllm/v1/attention/ops/pcp.py From 6a962071bdad41c2d739bb2e879e72bbe9df9e63 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Wed, 19 Aug 2026 18:17:43 -0700 Subject: [PATCH 175/839] [Distributed] Enable FlashInfer all-reduce by default (#52998) Signed-off-by: Woosuk Kwon Co-authored-by: OpenAI Codex --- vllm/distributed/device_communicators/cuda_communicator.py | 5 ++++- vllm/envs.py | 4 ++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index 06b441c5a416..f8d4e9fb28cc 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -58,7 +58,10 @@ def __init__( use_custom_allreduce = _ENABLE_CUSTOM_ALL_REDUCE use_torch_symm_mem = envs.VLLM_ALLREDUCE_USE_SYMM_MEM - use_flashinfer_allreduce = envs.VLLM_ALLREDUCE_USE_FLASHINFER + # FlashInfer all-reduce does not provide a fixed reduction order. + use_flashinfer_allreduce = ( + envs.VLLM_ALLREDUCE_USE_FLASHINFER and not envs.VLLM_BATCH_INVARIANT + ) use_aiter_allreduce = use_custom_allreduce and bool( rocm_aiter_ops.is_custom_all_reduce_enabled() ) diff --git a/vllm/envs.py b/vllm/envs.py index b49841d9fff0..a928461da210 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -260,7 +260,7 @@ VLLM_HAS_FLASHINFER_CUBIN: bool = False VLLM_ROCM_FP8_MFMA_PAGE_ATTN: bool = False VLLM_ALLREDUCE_USE_SYMM_MEM: bool = True - VLLM_ALLREDUCE_USE_FLASHINFER: bool = False + VLLM_ALLREDUCE_USE_FLASHINFER: bool = True VLLM_TUNED_CONFIG_FOLDER: str | None = None VLLM_ENABLE_STARTUP_PLAN: bool = False VLLM_GPT_OSS_SYSTEM_TOOL_MCP_LABELS: set[str] = set() @@ -1857,7 +1857,7 @@ def _resolve_rust_cli_path() -> str | None: ), # Whether to use FlashInfer allreduce "VLLM_ALLREDUCE_USE_FLASHINFER": lambda: bool( - int(os.getenv("VLLM_ALLREDUCE_USE_FLASHINFER", "0")) + int(os.getenv("VLLM_ALLREDUCE_USE_FLASHINFER", "1")) ), # Experimental: use this to enable MCP tool calling for non harmony models "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT": lambda: bool( From fbb4c04db42ff0c6a98af909321bf841798a20c2 Mon Sep 17 00:00:00 2001 From: Micah Williamson Date: Wed, 19 Aug 2026 21:20:21 -0500 Subject: [PATCH 176/839] [ROCm][CI] Speed up `test_rocm_aiter_qk_norm_rope_kvcache_fusion` (#53004) Signed-off-by: Micah Williamson Signed-off-by: Andreas Karatzas Co-authored-by: Andreas Karatzas Co-authored-by: Andreas Karatzas --- ..._rocm_aiter_qk_norm_rope_kvcache_fusion.py | 50 +++++++++++-------- 1 file changed, 30 insertions(+), 20 deletions(-) diff --git a/tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py b/tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py index 900f90eb426f..e385999db33a 100644 --- a/tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py +++ b/tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py @@ -2,9 +2,11 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import os +from importlib.metadata import version import pytest import torch +from packaging.version import Version import vllm.config from tests.compile.backend import TestBackend @@ -40,10 +42,27 @@ from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.kv_cache_interface import AttentionSpec +pytestmark = pytest.mark.skip_global_cleanup + +if not is_aiter_found_and_supported(): + pytest.skip( + "ROCm with supported AITER is required", + allow_module_level=True, + ) + INDEX_SELECT_OP = torch.ops.aten.index.Tensor FP8_DTYPE = current_platform.fp8_dtype() +@pytest.fixture(scope="module", autouse=True) +def module_global_cleanup(): + # Cleanup once at the end of the module + from vllm.distributed import cleanup_dist_env_and_memory + + yield + cleanup_dist_env_and_memory() + + class QKNormRoPEKVCacheTestModel(torch.nn.Module): """Minimal model that reproduces the QK-norm + RoPE + KV cache update pattern matched by QkNormRopeKvCacheFusionPass: @@ -416,8 +435,7 @@ def _run_qk_norm_rope_kvcache_fusion_test( AttentionBackendEnum.ROCM_AITER_FA, ], ) -@pytest.mark.parametrize("num_tokens", [5, 16, 2048]) -@pytest.mark.parametrize("use_shuffle_kv_layout", ["1", "0"]) +@pytest.mark.parametrize("num_tokens", [5, 2048]) @pytest.mark.parametrize( "kv_stride_order", # FA/unified use the 4D packed cache (num_blocks, num_kv_heads, block_size, @@ -428,11 +446,17 @@ def _run_qk_norm_rope_kvcache_fusion_test( @pytest.mark.parametrize("block_size", [16, 32, 64]) @pytest.mark.parametrize("dtype", [torch.bfloat16]) @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) -@pytest.mark.parametrize("rms_norm_eps", [1e-5, 1e-6]) +@pytest.mark.parametrize("rms_norm_eps", [1e-6]) @pytest.mark.parametrize("custom_op", ["+rotary_embedding", "+rms_norm"]) +# The fused_qk_norm_rope_cache kernel used by this test aborts on AITER < 0.1.20 +# (fused_qk_norm_rope_cache_quant.cu: "k_cache/v_cache must be contiguous within +# a block") @pytest.mark.skipif( - not is_aiter_found_and_supported(), - reason="Only test on ROCm with AITER installed and supported", + not Version(version("amd_aiter")) >= Version("0.1.20"), + reason=( + "Requires AITER >= 0.1.20; older kernels abort on " + "fused_qk_norm_rope_cache (k_cache/v_cache contiguity)" + ), ) def test_qk_norm_rope_kvcache_fusion( num_tokens: int, @@ -443,7 +467,6 @@ def test_qk_norm_rope_kvcache_fusion( is_neox: bool, attn_backend: AttentionBackendEnum, enable_aiter_triton_rope: bool, - use_shuffle_kv_layout: str, kv_stride_order: tuple[int, ...], block_size: int, dtype: torch.dtype, @@ -452,19 +475,6 @@ def test_qk_norm_rope_kvcache_fusion( custom_op: str, monkeypatch: pytest.MonkeyPatch, ): - if ( - attn_backend == AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN - and use_shuffle_kv_layout == "1" - ): - pytest.skip("ROCM_AITER_UNIFIED_ATTN is NHD-only; shuffle env is ignored") - if ( - attn_backend == AttentionBackendEnum.ROCM_AITER_FA - and use_shuffle_kv_layout == "1" - ): - pytest.skip( - "ROCM_AITER_FA gates the qk_norm+rope+kvcache fusion off under shuffle " - "layout (defers to reshape_and_cache_shuffle_triton), so nothing fuses" - ) _run_qk_norm_rope_kvcache_fusion_test( attn_backend=attn_backend, enable_aiter_triton_rope=enable_aiter_triton_rope, @@ -475,7 +485,7 @@ def test_qk_norm_rope_kvcache_fusion( rotary_dim=rotary_dim, block_size=block_size, is_neox=is_neox, - use_shuffle_kv_layout=use_shuffle_kv_layout, + use_shuffle_kv_layout="0", kv_stride_order=kv_stride_order, dtype=dtype, kv_cache_dtype=kv_cache_dtype, From c0233fcf0100c07820d7f7a725de458a0e14df64 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Wed, 19 Aug 2026 19:39:52 -0700 Subject: [PATCH 177/839] [Model] Remove unused DeepseekV32Indexer forward (#53021) Signed-off-by: Woosuk Kwon Co-authored-by: OpenAI Codex --- vllm/models/deepseek_v32/attention.py | 51 --------------------------- 1 file changed, 51 deletions(-) diff --git a/vllm/models/deepseek_v32/attention.py b/vllm/models/deepseek_v32/attention.py index 4edcae4640af..87e7ec24e329 100644 --- a/vllm/models/deepseek_v32/attention.py +++ b/vllm/models/deepseek_v32/attention.py @@ -21,9 +21,6 @@ RowParallelLinear, ) from vllm.model_executor.layers.quantization import QuantizationConfig -from vllm.model_executor.layers.quantization.utils.fp8_utils import ( - per_token_group_quant_fp8, -) from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.sparse_attn_indexer import ( SparseAttnIndexer, @@ -117,54 +114,6 @@ def __init__( self.topk_indices_buffer, ) - def forward( - self, - hidden_states: torch.Tensor, - qr: torch.Tensor, - positions: torch.Tensor, - rotary_emb: nn.Module, - ) -> torch.Tensor: - q, _ = self.wq_b(qr) - q = q.view(-1, self.n_head, self.head_dim) - - q_pe, q_nope = torch.split( - q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 - ) - - kw, _ = self.wk_weights_proj(hidden_states) - k = kw[:, : self.head_dim] - weights = kw[:, self.head_dim :] - - k = self.k_norm(k) - k_pe, k_nope = torch.split( - k, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 - ) - - q_pe, k_pe = rotary_emb(positions, q_pe, k_pe.unsqueeze(1)) - - q_pe = q_pe.reshape(-1, self.n_head, self.rope_dim) - k_pe = k_pe.reshape(-1, 1, self.rope_dim) - - q = torch.cat([q_pe, q_nope], dim=-1) - k = torch.cat([k_pe.squeeze(-2), k_nope], dim=-1) - - q = q.view(-1, self.head_dim) - q_fp8, q_scale = per_token_group_quant_fp8( - q, - self.quant_block_size, - column_major_scales=False, - use_ue8m0=self.scale_fmt is not None, - ) - q_fp8 = q_fp8.view(-1, self.n_head, self.head_dim) - q_scale = q_scale.view(-1, self.n_head, 1) - - weights = ( - weights.unsqueeze(-1) * q_scale * self.softmax_scale * self.n_head**-0.5 - ) - weights = weights.squeeze(-1) - - return self.indexer_op(hidden_states, q_fp8, k, weights) - class DeepseekV32Attention(MLAAttention): indexer: "DeepseekV32Indexer | None" From 76fb6d210a57c7efbb61fa6cb029d1fa1911bfb5 Mon Sep 17 00:00:00 2001 From: 92hyungjun <78071764+92hyungjun@users.noreply.github.com> Date: Thu, 20 Aug 2026 11:49:23 +0900 Subject: [PATCH 178/839] [Bugfix][Core] Reserve the KV null block when validating max_model_len (#47272) Signed-off-by: hyungjun oh <92hyungjun@gmail.com> Signed-off-by: Nick Hill Co-authored-by: Claude Co-authored-by: Nick Hill --- tests/v1/core/test_kv_cache_utils.py | 74 ++++++++++++++++++- tests/v1/e2e/general/test_async_scheduling.py | 11 +-- tests/v1/e2e/general/test_context_length.py | 8 +- tests/v1/engine/test_init_error_messaging.py | 5 +- tests/v1/sample/test_logprobs.py | 2 +- vllm/v1/core/kv_cache_utils.py | 25 +++++-- 6 files changed, 109 insertions(+), 16 deletions(-) diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 829f45608a48..48943aeb6016 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -27,6 +27,7 @@ BlockHash, FreeKVCacheBlockQueue, KVCacheBlock, + check_enough_kv_cache_memory, estimate_max_model_len, generate_block_hash_extra_keys, generate_scheduler_kv_cache_config, @@ -2685,7 +2686,9 @@ def test_auto_fit_max_model_len_with_hybrid(): "layer_2": new_kv_cache_spec(), } - available_memory = mem_per_block_per_layer * (1024 // 16 + 1 + gamma) + # One extra block on top of what a 1024-token request needs: the pool + # reserves one block as the null block. + available_memory = mem_per_block_per_layer * (1024 // 16 + 1 + gamma + 1) _kv_cache_configs = get_kv_cache_configs( vllm_config, [kv_cache_specs], [available_memory] ) @@ -3095,3 +3098,72 @@ def test_resolve_block_hashes_rejects_mismatched_view(): mismatched = BlockHashListWithBlockSize(raw, 2, 8) with pytest.raises(AssertionError): resolve_block_hashes(mismatched, 2, 4) + + +@pytest.mark.parametrize("use_override", [True, False]) +def test_kv_cache_reserves_null_block_for_max_model_len(use_override): + """A KV cache sized to exactly the blocks a max_model_len request needs must + be rejected. BlockPool keeps one block as the null block, so only + num_blocks - 1 are usable; accepting num_blocks == needed would leave the + request unschedulable and hang the engine. Covers both the + num_gpu_blocks_override path and the memory-derived block count. + """ + block_size = 16 + max_model_len = 512 # needs 512 / 16 = 32 blocks + vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len)) + spec = new_kv_cache_spec(block_size=block_size) + + # 32 blocks -> only 31 usable after the null block: one short -> reject. + if use_override: + # Ample raw memory; the override alone constrains the block count. + vllm_config.cache_config.num_gpu_blocks_override = 32 + available_memory = [spec.page_size_bytes * 1024] + else: + available_memory = [spec.page_size_bytes * 32] + + with pytest.raises(ValueError, match="max seq len"): + get_kv_cache_configs(vllm_config, [{"layer1": spec}], available_memory) + + +def test_auto_fit_max_model_len_reserves_null_block(): + """Auto-fit (max_model_len=-1) must size max_model_len against usable + blocks, not the total pool. With memory for exactly 64 blocks, one is the + null block, so auto-fit must settle on 63 * block_size; picking the full + pool would leave a max-length request unschedulable and hang the engine. + """ + block_size = 16 + model_config = ModelConfig(max_model_len=1024) + model_config.original_max_model_len = -1 + vllm_config = VllmConfig(model_config=model_config) + spec = new_kv_cache_spec(block_size=block_size) + + # Exactly the 1024 / 16 = 64 blocks a full-length request would need. + available_memory = [spec.page_size_bytes * 64] + + get_kv_cache_configs(vllm_config, [{"layer1": spec}], available_memory) + + assert vllm_config.model_config.max_model_len == 63 * block_size + + +def test_check_enough_kv_cache_memory_reserves_null_block(): + """The public admission check must reject a KV cache sized to exactly the + blocks a max_model_len request needs. BlockPool keeps one block as the + null block, so only num_blocks - 1 are usable; accepting + num_blocks == needed would leave the request unschedulable and hang the + engine. + """ + block_size = 16 + max_model_len = 512 # needs 512 / 16 = 32 blocks + vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len)) + spec = new_kv_cache_spec(block_size=block_size) + + # 32 blocks -> only 31 usable after the null block: one short -> reject. + with pytest.raises(ValueError, match="max seq len"): + check_enough_kv_cache_memory( + vllm_config, {"layer1": spec}, spec.page_size_bytes * 32 + ) + + # 33 blocks -> 32 usable after the null block -> accept. + check_enough_kv_cache_memory( + vllm_config, {"layer1": spec}, spec.page_size_bytes * 33 + ) diff --git a/tests/v1/e2e/general/test_async_scheduling.py b/tests/v1/e2e/general/test_async_scheduling.py index 7f5a11514563..1024829307c7 100644 --- a/tests/v1/e2e/general/test_async_scheduling.py +++ b/tests/v1/e2e/general/test_async_scheduling.py @@ -332,11 +332,12 @@ def run_test( ): spec_decoding = spec_config is not None cache_arg: dict[str, Any] = ( - # Force preemptions: with 32 blocks the cache holds at most a single - # max-length request, so the ~34 concurrent prompts contend and trigger - # preemption. (Prompts here are << max_model_len, so dropping - # max_model_len from 4096 to 512 doesn't change generation behavior.) - dict(num_gpu_blocks_override=32, max_model_len=512) + # Force preemptions: with 33 blocks (one is the reserved null block) + # the cache holds at most a single max-length request, so the ~34 + # concurrent prompts contend and trigger preemption. (Prompts here are + # << max_model_len, so dropping max_model_len from 4096 to 512 doesn't + # change generation behavior.) + dict(num_gpu_blocks_override=33, max_model_len=512) if test_preemption else dict(gpu_memory_utilization=0.9, max_model_len=4096) ) diff --git a/tests/v1/e2e/general/test_context_length.py b/tests/v1/e2e/general/test_context_length.py index 955e5c5bf025..7e77b31fb1b7 100644 --- a/tests/v1/e2e/general/test_context_length.py +++ b/tests/v1/e2e/general/test_context_length.py @@ -77,8 +77,12 @@ def test_auto_fit_max_model_len_rejects_oversized_input( # Use a small KV cache budget to force auto-fit to a small # max_model_len. Pin block_size=16 so the budget is independent - # of the platform's default block size. - kv_cache_bytes = 1_000_000 # 1 MB + # of the platform's default block size. One block for this model is + # 2 (K/V) * 16 (block) * 12 (heads) * 64 (head_dim) * 2 (fp16) + # * 12 (layers) = 589,824 bytes, and the pool must cover at least + # two: one is reserved as the null block, so a 1 MB single-block + # pool cannot serve any tokens at all. + kv_cache_bytes = 2_000_000 # 2 MB = 3 blocks -> 2 usable with vllm_runner( model_name=model, diff --git a/tests/v1/engine/test_init_error_messaging.py b/tests/v1/engine/test_init_error_messaging.py index bc23a68f9deb..5626dec17043 100644 --- a/tests/v1/engine/test_init_error_messaging.py +++ b/tests/v1/engine/test_init_error_messaging.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest +import torch from vllm.v1.core.kv_cache_utils import check_enough_kv_cache_memory from vllm.v1.kv_cache_interface import FullAttentionSpec @@ -18,7 +19,7 @@ def test_kv_cache_oom_no_memory(): block_size=16, num_kv_heads=8, head_size=128, - dtype="float16", + dtype=torch.float16, ) } @@ -46,7 +47,7 @@ def test_kv_cache_oom_insufficient_memory(monkeypatch): block_size=16, num_kv_heads=8, head_size=128, - dtype="float16", + dtype=torch.float16, ) } diff --git a/tests/v1/sample/test_logprobs.py b/tests/v1/sample/test_logprobs.py index 05a1e9361824..f1e88bec8aed 100644 --- a/tests/v1/sample/test_logprobs.py +++ b/tests/v1/sample/test_logprobs.py @@ -1244,7 +1244,7 @@ def test_prompt_logprobs_with_chunking_and_preemption(): max_model_len=512, enable_chunked_prefill=True, max_num_batched_tokens=48, # Force prefill chunking - num_gpu_blocks_override=32, # Force preemptions + num_gpu_blocks_override=33, # Force preemptions (32 usable + null block) disable_log_stats=False, gpu_memory_utilization=0.25, ) as vllm_model: diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 4ac55dbbf70d..46e9b8e96b2f 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -892,8 +892,17 @@ def check_enough_kv_cache_memory( # No need to check for available memory if the kv_cache_spec is empty if kv_cache_spec: + # Reserve the null block BlockPool permanently holds back, so the check + # plans against usable blocks, as in get_kv_cache_configs. Group a copy + # of the specs since grouping may unify them in-place. + groups = get_kv_cache_groups(vllm_config, dict(kv_cache_spec)) + check_memory = ( + available_memory - _pool_bytes_per_block(vllm_config, groups) + if groups + else available_memory + ) _check_enough_kv_cache_memory( - available_memory, + check_memory, lambda: max_memory_usage_bytes(vllm_config, kv_cache_spec.values()), vllm_config.model_config.max_model_len, lambda am: estimate_max_model_len(vllm_config, kv_cache_spec, am), @@ -2228,13 +2237,19 @@ def get_kv_cache_configs( adjusted_memory.append(override * bytes_per_block) available_memory = adjusted_memory + # Reserve the null block BlockPool permanently holds back, so auto-fit and + # the capacity check both plan against usable blocks. Allocation below + # still uses the full memory. + check_memory = [ + avail_mem - _pool_bytes_per_block(vllm_config, groups) if groups else avail_mem + for groups, avail_mem in zip(projected_groups_per_worker, available_memory) + ] + if vllm_config.model_config.original_max_model_len == -1: - _auto_fit_max_model_len( - vllm_config, projected_groups_per_worker, available_memory - ) + _auto_fit_max_model_len(vllm_config, projected_groups_per_worker, check_memory) # Check if the available memory is enough per worker. - for groups, avail_mem in zip(projected_groups_per_worker, available_memory): + for groups, avail_mem in zip(projected_groups_per_worker, check_memory): if not groups: continue _check_enough_kv_cache_memory( From f85e060ec24ad5b67b47690b89bd0a5f1def05f4 Mon Sep 17 00:00:00 2001 From: pmanczak Date: Thu, 20 Aug 2026 04:59:22 +0200 Subject: [PATCH 179/839] [Doc] Update Gaudi HPU committers (#52726) Signed-off-by: pmanczak Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- docs/governance/committers.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/governance/committers.md b/docs/governance/committers.md index 1737d0c0e185..0ca6e24ded0a 100644 --- a/docs/governance/committers.md +++ b/docs/governance/committers.md @@ -193,5 +193,5 @@ If you have PRs touching the area, please feel free to ping the area owner for r ### Ecosystem Projects - Ascend NPU: [@wangxiyuan](https://github.com/wangxiyuan) and [see more details](https://vllm-ascend.readthedocs.io/en/latest/community/contributors.html#maintainers) -- Intel Gaudi HPU [@xuechendi](https://github.com/xuechendi) and [@kzawora-intel](https://github.com/kzawora-intel) +- Intel Gaudi HPU [@xuechendi](https://github.com/xuechendi) and [@iboiko-habana](https://github.com/iboiko-habana) - Semantic Router: [@xunzhuo](https://github.com/xunzhuo), [@rootfs](https://github.com/rootfs) and [see more details](https://vllm-semantic-router.com/community/team) From e85dd21497e46f739911ee1cef4f6d6c68da1e62 Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Thu, 20 Aug 2026 11:23:43 +0800 Subject: [PATCH 180/839] fix: report stop_sequence stop_reason in Anthropic Messages API (#45807) Signed-off-by: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Co-authored-by: Chauncey --- .../test_anthropic_messages_conversion.py | 99 +++++++++++++++++++ vllm/entrypoints/anthropic/serving.py | 28 +++++- 2 files changed, 122 insertions(+), 5 deletions(-) diff --git a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py index 1cd7227bdc99..1491b986a541 100644 --- a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py +++ b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py @@ -975,6 +975,7 @@ def _make_stream_chunk( *, delta: DeltaMessage | None = None, finish_reason: str | None = None, + stop_reason: int | str | None = None, choices: list[ChatCompletionResponseStreamChoice] | None = None, usage: UsageInfo | None = None, ) -> str: @@ -984,6 +985,7 @@ def _make_stream_chunk( index=0, delta=delta or DeltaMessage(), finish_reason=finish_reason, + stop_reason=stop_reason, ) ] chunk = ChatCompletionStreamResponse( @@ -1463,6 +1465,103 @@ def test_empty_cache_salt_returns_bad_request(self): handler.create_messages.assert_not_awaited() +class TestStopSequenceReason: + """When generation stops because a configured stop string matched, the + Anthropic Messages API must report ``stop_reason="stop_sequence"`` and echo + the matched string in ``stop_sequence``. vLLM surfaces the matched string in + the OpenAI choice's ``stop_reason`` field (a str) while ``finish_reason`` + stays ``"stop"``. A natural EOS (stop_reason None) or a stop token id (int) + must still map to ``end_turn``. + """ + + def test_non_streaming_stop_string_maps_to_stop_sequence(self): + converter = _make_full_converter() + response = ChatCompletionResponse( + id="chatcmpl-test", + model="test-model", + choices=[ + ChatCompletionResponseChoice( + index=0, + message=ChatMessage(role="assistant", content="hello"), + finish_reason="stop", + stop_reason="", + ) + ], + usage=UsageInfo(prompt_tokens=5, total_tokens=8, completion_tokens=3), + ) + + result = converter.messages_full_converter(response) + + assert result.stop_reason == "stop_sequence" + assert result.stop_sequence == "" + + def test_non_streaming_natural_eos_maps_to_end_turn(self): + converter = _make_full_converter() + response = ChatCompletionResponse( + id="chatcmpl-test", + model="test-model", + choices=[ + ChatCompletionResponseChoice( + index=0, + message=ChatMessage(role="assistant", content="hello"), + finish_reason="stop", + stop_reason=None, + ) + ], + usage=UsageInfo(prompt_tokens=5, total_tokens=8, completion_tokens=3), + ) + + result = converter.messages_full_converter(response) + + assert result.stop_reason == "end_turn" + assert result.stop_sequence is None + + def test_non_streaming_stop_token_id_maps_to_end_turn(self): + converter = _make_full_converter() + response = ChatCompletionResponse( + id="chatcmpl-test", + model="test-model", + choices=[ + ChatCompletionResponseChoice( + index=0, + message=ChatMessage(role="assistant", content="hello"), + finish_reason="stop", + stop_reason=128009, + ) + ], + usage=UsageInfo(prompt_tokens=5, total_tokens=8, completion_tokens=3), + ) + + result = converter.messages_full_converter(response) + + assert result.stop_reason == "end_turn" + assert result.stop_sequence is None + + @pytest.mark.asyncio + async def test_streaming_stop_string_maps_to_stop_sequence(self): + async def sse_input(): + yield _make_stream_chunk( + delta=DeltaMessage(content="hi"), + usage=UsageInfo(prompt_tokens=5, total_tokens=5, completion_tokens=0), + ) + yield _make_stream_chunk(finish_reason="stop", stop_reason="") + yield _make_stream_chunk( + choices=[], + usage=UsageInfo(prompt_tokens=5, total_tokens=8, completion_tokens=3), + ) + yield "data: [DONE]" + + converter = _make_stream_converter() + output = [] + async for event in converter.message_stream_converter(sse_input()): + output.append(event) + + events = _parse_sse_events(output) + msg_deltas = [data for ev_type, data in events if ev_type == "message_delta"] + assert msg_deltas[0]["delta"]["stop_reason"] == "stop_sequence" + assert msg_deltas[0]["delta"]["stop_sequence"] == "" + + # ====================================================================== # Client-caused errors are 4xx, not 500 (Issue #52088) # ====================================================================== diff --git a/vllm/entrypoints/anthropic/serving.py b/vllm/entrypoints/anthropic/serving.py index 4ec5c14e019b..7e18fe1ef350 100644 --- a/vllm/entrypoints/anthropic/serving.py +++ b/vllm/entrypoints/anthropic/serving.py @@ -631,7 +631,13 @@ def messages_full_converter( ) choice = generator.choices[0] if choice.finish_reason == "stop": - result.stop_reason = "end_turn" + # vLLM reports the matched stop string in stop_reason (a str); + # an int stop-token-id or None (natural EOS) maps to end_turn. + if isinstance(choice.stop_reason, str): + result.stop_reason = "stop_sequence" + result.stop_sequence = choice.stop_reason + else: + result.stop_reason = "end_turn" elif choice.finish_reason == "length": result.stop_reason = "max_tokens" elif choice.finish_reason == "tool_calls": @@ -715,6 +721,9 @@ def start(self, block: AnthropicContentBlock) -> None: first_item = True finish_reason = None + # Matched stop string, when generation stopped on one (a str); + # int stop-token-id / None are not stop sequences. + stop_sequence: int | str | None = None state = _ActiveBlockState() # Map from tool call index to tool_use_id tool_index_to_id: dict[int, str] = {} @@ -824,12 +833,20 @@ def stop_and_flush() -> list[str]: if len(origin_chunk.choices) == 0: for event in stop_and_flush(): yield event - stop_reason = self.stop_reason_map.get( - finish_reason or "stop" - ) + if isinstance(stop_sequence, str): + stop_delta = AnthropicDelta( + stop_reason="stop_sequence", + stop_sequence=stop_sequence, + ) + else: + stop_delta = AnthropicDelta( + stop_reason=self.stop_reason_map.get( + finish_reason or "stop" + ) + ) chunk = AnthropicStreamEvent( type="message_delta", - delta=AnthropicDelta(stop_reason=stop_reason), + delta=stop_delta, usage=_build_anthropic_usage(origin_chunk.usage), ) data = chunk.model_dump_json(exclude_unset=True) @@ -838,6 +855,7 @@ def stop_and_flush() -> list[str]: if origin_chunk.choices[0].finish_reason is not None: finish_reason = origin_chunk.choices[0].finish_reason + stop_sequence = origin_chunk.choices[0].stop_reason # continue # thinking / text content From 4f66bc3e0d92aca0e8d1a29e22425327a07c916e Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Wed, 19 Aug 2026 22:33:25 -0500 Subject: [PATCH 181/839] [CI][ROCm] Standardize AMD test job labels by device (#52976) Signed-off-by: Andreas Karatzas Signed-off-by: Andreas Karatzas --- .buildkite/test-amd.yaml | 708 ++++++++++--------- .buildkite/test_areas/distributed.yaml | 6 +- .buildkite/test_areas/models_multimodal.yaml | 2 +- tests/models/test_vision.py | 8 +- 4 files changed, 384 insertions(+), 340 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index fdeaf788e0b8..e21c8f59e627 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -13,7 +13,7 @@ # soft_fail(bool): allow this step to fail without failing the entire pipeline (useful for flaky or experimental tests). # command(str): the single command to run for tests. incompatible with commands. # commands(list): the list of commands to run for the test. incompatible with command. -# mirror_hardwares(list): the list of hardware to run the test on as well. currently only supports [amdexperimental] +# mirror_hardwares(list): selector tags for AMD pipelines that should include the test. # dind(bool): when false, run the job directly in the AMD Kubernetes pod. # When true or omitted, use the legacy Docker-in-Docker path. # num_gpus(int): override the number of GPUs for the test. defaults to 1 GPU. currently supports 2,4. @@ -105,6 +105,156 @@ steps: +######################################################################################################################################### +# # +# AMD CPU tests # +# # +######################################################################################################################################### + +- label: ":computer: (CPU) Basic Models Other" # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + agent_pool: mi250_1 + no_gpu: true + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/models/test_utils.py + - tests/models/test_vision.py + - tests/models/test_adapters.py + - tests/models/transformers/fusers/ + commands: + - pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/transformers/fusers/ + +- label: ":computer: (CPU) Multimodal Processor Shard %N" # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + agent_pool: mi250_1 + no_gpu: true + parallelism: 6 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/models/multimodal + - tests/models/registry.py + commands: + - pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB + +- label: ":computer: (CPU) V1 Others" # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + agent_pool: mi250_1 + no_gpu: true + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/v1 + commands: + - pytest -v -s -m 'cpu_test' v1/core + - pytest -v -s v1/structured_output + - pytest -v -s v1/test_serial_utils.py + - pytest -v -s v1/test_kv_cache_spec_registry.py + - pytest -v -s v1/cudagraph/test_cudagraph_manager.py + - pytest -v -s -m 'cpu_test' v1/kv_connector/unit + - pytest -v -s -m 'cpu_test' v1/ec_connector/unit + - pytest -v -s -m 'cpu_test' v1/metrics + +- label: ":computer: (CPU) Async Engine, Inputs, Utils, Worker, Config" # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + agent_pool: mi250_1 + no_gpu: true + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/test_envs.py + - tests/test_inputs.py + - tests/test_outputs.py + - tests/test_pooling_params.py + - tests/test_ray_env.py + - tests/test_sampling_params.py + - tests/multimodal + - tests/renderers + - tests/standalone_tests/lazy_imports.py + - tests/tokenizers_ + - tests/reasoning + - tests/tool_parsers + - tests/parser + - tests/transformers_utils + - tests/config + commands: + - python3 standalone_tests/lazy_imports.py + - pytest -v -s test_envs.py + - pytest -v -s test_inputs.py + - pytest -v -s test_outputs.py + - pytest -v -s test_pooling_params.py + - pytest -v -s test_ray_env.py + - pytest -v -s test_sampling_params.py + - pytest -v -s -m 'cpu_test' multimodal + - pytest -v -s renderers + - pytest -v -s tokenizers_ + - pytest -v -s reasoning + - pytest -v -s tool_parsers + - pytest -v -s parser + - pytest -v -s transformers_utils + - pytest -v -s config --ignore=config/test_mp_reducer.py + +- label: ":computer: (CPU) Rust Frontend Cargo Style + Clippy" # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + agent_pool: mi250_1 + no_gpu: true + optional: true + working_dir: "/vllm-workspace" + source_file_dependencies: + - rust/ + - rust-toolchain.toml + - .buildkite/test_areas/rust_frontend_cargo.yaml + - .buildkite/scripts/run-rust-frontend-cargo-ci.sh + commands: + - bash .buildkite/scripts/run-rust-frontend-cargo-ci.sh style-clippy + +- label: ":computer: (CPU) Rust Frontend Cargo" # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + agent_pool: mi250_1 + no_gpu: true + optional: true + working_dir: "/vllm-workspace" + source_file_dependencies: + - rust/ + - rust-toolchain.toml + - .buildkite/test_areas/rust_frontend_cargo.yaml + - .buildkite/scripts/run-rust-frontend-cargo-ci.sh + commands: + - bash .buildkite/scripts/run-rust-frontend-cargo-ci.sh test + +- label: ":computer: (CPU) Docker Build Metadata" # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + agent_pool: mi250_1 + no_gpu: true + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - .buildkite/scripts/docker-build-metadata-args.sh + - .buildkite/scripts/ci-bake-rocm.sh + - .buildkite/release-pipeline.yaml + - docker/Dockerfile + - docker/Dockerfile.cpu + - docker/Dockerfile.rocm + - docker/Dockerfile.rocm_base + - docker/ci-rocm.hcl + - docker/docker-bake.hcl + - docker/docker-bake-rocm.hcl + - tests/tools/test_docker_build_metadata_args.py + commands: + - pytest -v -s tools/test_docker_build_metadata_args.py + ######################################################################################################################################### # # # MI250 (gfx90a) tests # @@ -113,7 +263,7 @@ steps: #---------------------------------------------------------- mi250 · compile ----------------------------------------------------------# -- label: PyTorch Fullgraph Smoke Test # TBD +- label: ":amd: (MI250) PyTorch Fullgraph Smoke" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 @@ -132,7 +282,7 @@ steps: #-------------------------------------------------------- mi250 · distributed --------------------------------------------------------# -- label: Pipeline + Context Parallelism (4 GPUs) # TBD +- label: ":amd: (MI250) Pipeline + Context Parallelism" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_4 @@ -155,7 +305,7 @@ steps: #---------------------------------------------------------- mi250 · kernels ----------------------------------------------------------# -- label: Kernels Helion Test # TBD +- label: ":amd: (MI250) Helion Kernels" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 @@ -170,25 +320,9 @@ steps: - pip install helion==1.4.0 - pytest -v -s kernels/helion/ -#------------------------------------------------------ mi250 · models / basic -------------------------------------------------------# - -- label: Basic Models Test (Other CPU) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - no_gpu: true - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/test_utils.py - - tests/models/test_vision.py - commands: - - pytest -v -s models/test_utils.py models/test_vision.py - #----------------------------------------------------- mi250 · models / language -----------------------------------------------------# -- label: Language Models Test (MTEB) # TBD +- label: ":amd: (MI250) Language Models (MTEB)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 @@ -200,7 +334,7 @@ steps: commands: - pytest -v -s models/language/pooling_mteb_test -- label: Language Models Test (PPL) # TBD +- label: ":amd: (MI250) Language Models (PPL)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 @@ -213,7 +347,7 @@ steps: #---------------------------------------------------- mi250 · models / multimodal ----------------------------------------------------# -- label: "Multi-Modal Models (Standard) 3: llava + qwen2_vl" # TBD +- label: ":amd: (MI250) Multimodal Models (Standard) 3: llava + qwen2_vl" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 @@ -226,23 +360,9 @@ steps: - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma" - pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model -- label: Multi-Modal Processor (CPU) %N # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - parallelism: 6 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/multimodal - - tests/models/registry.py - commands: - - pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB - #------------------------------------------------------------ mi250 · v1 -------------------------------------------------------------# -- label: Batch Invariance (H100-MI250) # TBD +- label: ":amd: (MI250) Batch Invariance" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 @@ -260,7 +380,7 @@ steps: - pytest -v -s v1/determinism/test_batch_invariance.py - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py -- label: Cudagraph # TBD +- label: ":amd: (MI250) CUDAGraph" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 @@ -276,7 +396,7 @@ steps: - pytest -v -s v1/cudagraph/test_cudagraph_dispatch.py - pytest -v -s v1/cudagraph/test_cudagraph_mode.py -- label: e2e Core (1 GPU) # TBD +- label: ":amd: (MI250) E2E Core" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 @@ -289,7 +409,7 @@ steps: commands: - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py -- label: e2e Scheduling (1 GPU) # TBD +- label: ":amd: (MI250) E2E Scheduling" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 @@ -302,7 +422,7 @@ steps: commands: - pytest -v -s v1/e2e/general/test_async_scheduling.py -- label: Engine (1 GPU) # TBD +- label: ":amd: (MI250) V1 Engine" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 @@ -311,12 +431,19 @@ steps: source_file_dependencies: - vllm/v1/ - tests/v1/engine/ + - tests/v1/test_tensor_ipc_queue.py + - tests/config/test_mp_reducer.py + - vllm/config/ + - vllm/engine/arg_utils.py + - vllm/transformers_utils/config.py - vllm/platforms/rocm.py commands: - pytest -v -s v1/engine/test_preprocess_error_handling.py - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py + - pytest -v -s v1/test_tensor_ipc_queue.py + - pytest -v -s config/test_mp_reducer.py -- label: Spec Decode Draft Model # TBD +- label: ":amd: (MI250) Spec Decode Draft Model" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 @@ -333,7 +460,7 @@ steps: commands: - pytest -v -s v1/e2e/spec_decode/draft_model/ -- label: Spec Decode Speculators + MTP # TBD +- label: ":amd: (MI250) Spec Decode Speculators + MTP" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 @@ -352,64 +479,9 @@ steps: - pytest -v -s v1/e2e/spec_decode/speculators/ - pytest -v -s v1/e2e/spec_decode/mtp/ -- label: V1 others (CPU) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/v1 - commands: - - pytest -v -s -m 'cpu_test' v1/core - - pytest -v -s v1/structured_output - - pytest -v -s v1/test_serial_utils.py - - pytest -v -s -m 'cpu_test' v1/kv_connector/unit - - pytest -v -s -m 'cpu_test' v1/metrics - #------------------------------------------------------------- mi250 · misc ------------------------------------------------------------# -- label: Async Engine, Inputs, Utils, Worker, Config (CPU) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - optional: true - no_gpu: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/test_inputs.py - - tests/test_outputs.py - - tests/test_pooling_params.py - - tests/test_ray_env.py - - tests/test_sampling_params.py - - tests/multimodal - - tests/renderers - - tests/standalone_tests/lazy_imports.py - - tests/tokenizers_ - - tests/reasoning - - tests/tool_parsers - - tests/parser - - tests/transformers_utils - - tests/config - commands: - - python3 standalone_tests/lazy_imports.py - - pytest -v -s test_inputs.py - - pytest -v -s test_outputs.py - - pytest -v -s test_pooling_params.py - - pytest -v -s test_ray_env.py - - pytest -v -s test_sampling_params.py - - pytest -v -s -m 'cpu_test' multimodal - - pytest -v -s renderers - - pytest -v -s tokenizers_ - - pytest -v -s reasoning - - pytest -v -s tool_parsers - - pytest -v -s parser - - pytest -v -s transformers_utils - - pytest -v -s config - -- label: Python-only Installation # TBD +- label: ":amd: (MI300) Python-only Installation" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -423,59 +495,6 @@ steps: commands: - bash standalone_tests/python_only_compile.sh -#------------------------------------------------------------ mi250 · rust -----------------------------------------------------------# - -- label: Rust Frontend Cargo Style + Clippy # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - no_gpu: true - working_dir: "/vllm-workspace" - source_file_dependencies: - - rust/ - - rust-toolchain.toml - - .buildkite/test_areas/rust_frontend_cargo.yaml - - .buildkite/scripts/run-rust-frontend-cargo-ci.sh - commands: - - bash .buildkite/scripts/run-rust-frontend-cargo-ci.sh style-clippy - -- label: Rust Frontend Cargo Tests # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - no_gpu: true - working_dir: "/vllm-workspace" - source_file_dependencies: - - rust/ - - rust-toolchain.toml - - .buildkite/test_areas/rust_frontend_cargo.yaml - - .buildkite/scripts/run-rust-frontend-cargo-ci.sh - commands: - - bash .buildkite/scripts/run-rust-frontend-cargo-ci.sh test - -#----------------------------------------------------------- mi250 · docker ----------------------------------------------------------# - -- label: Docker Build Metadata (ROCm) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - no_gpu: true - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - .buildkite/scripts/docker-build-metadata-args.sh - - .buildkite/scripts/ci-bake-rocm.sh - - docker/Dockerfile - - docker/Dockerfile.cpu - - docker/Dockerfile.rocm - - docker/Dockerfile.rocm_base - - docker/ci-rocm.hcl - - docker/docker-bake.hcl - - docker/docker-bake-rocm.hcl - - tests/tools/test_docker_build_metadata_args.py - commands: - - pytest -v -s tools/test_docker_build_metadata_args.py - ######################################################################################################################################### # # # MI300 (gfx942) tests # @@ -484,7 +503,7 @@ steps: #----------------------------------------------------- mi300 · basic_correctness -----------------------------------------------------# -- label: Basic Correctness # TBD +- label: ":amd: (MI300) Basic Correctness" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -502,7 +521,7 @@ steps: - VLLM_TARGET_TEST_SUITE=MI300 pytest -v -s basic_correctness/test_basic_correctness.py - pytest -v -s basic_correctness/test_cpu_offload.py -- label: Distributed Model Tests (2 GPUs) # TBD +- label: ":amd: (MI300) Distributed Models" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -529,10 +548,11 @@ steps: - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_phi4siglip.py - pytest models/multimodal/generation/test_phi4siglip.py -v -s -m 'distributed(num_gpus=2)' - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest models/multimodal/generation/test_whisper.py -v -s -m 'distributed(num_gpus=2)' + - pytest -v -s models/test_vision.py -m 'distributed(num_gpus=2)' #-------------------------------------------------------- mi300 · benchmarks ---------------------------------------------------------# -- label: Benchmarks CLI Test # TBD +- label: ":amd: (MI300) Benchmarks CLI" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -547,7 +567,7 @@ steps: #---------------------------------------------------------- mi300 · compile ----------------------------------------------------------# -- label: PyTorch Compilation Unit Tests # TBD +- label: ":amd: (MI300) PyTorch Compilation" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -567,7 +587,7 @@ steps: commands: - "find compile/ -maxdepth 1 -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'" -- label: Fusion E2E Config Sweep (H100-MI300) # TBD +- label: ":amd: (MI300) Fusion E2E Config Sweep" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -588,7 +608,7 @@ steps: - rocm-smi - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "llama-3" -- label: Fusion E2E Quick (H100-MI300) # TBD +- label: ":amd: (MI300) Fusion E2E Quick" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -610,7 +630,7 @@ steps: # Different from CUDA, Qwen requires +rms_norm and +quant_fp8 as rms+quant fusion is only supported on AITER - "pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k 'inductor_partition and +rms_norm and +quant_fp8 and qwen3'" -- label: PyTorch Compilation Passes Unit Tests # TBD +- label: ":amd: (MI300) PyTorch Compilation Passes" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -622,7 +642,7 @@ steps: commands: - pytest -s -v compile/passes --ignore compile/passes/distributed -- label: PyTorch Fullgraph # TBD +- label: ":amd: (MI300) PyTorch Fullgraph" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -640,7 +660,7 @@ steps: commands: - pytest -v -s compile/fullgraph/test_full_graph.py -- label: Pytorch Nightly Dependency Override Check # TBD +- label: ":amd: (MI300) PyTorch Nightly Dependency Override Check" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -654,7 +674,7 @@ steps: commands: - bash standalone_tests/pytorch_nightly_dependency.sh -- label: Distributed Compile Unit Tests (2xH100-2xMI300) # TBD +- label: ":amd: (MI300) Distributed Compile" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -673,7 +693,7 @@ steps: - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/passes/distributed/test_async_tp.py - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py::test_tp2_ar_rms_fusions -- label: Distributed Compile + RPC Tests (2 GPUs) # TBD +- label: ":amd: (MI300) Distributed Compile + RPC" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -699,7 +719,7 @@ steps: #----------------------------------------------------------- mi300 · cuda ------------------------------------------------------------# -- label: Platform Tests # TBD +- label: ":amd: (MI300) CUDA Platform" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -714,7 +734,7 @@ steps: #-------------------------------------------------------- mi300 · detokenizer --------------------------------------------------------# -- label: Async Engine, Inputs, Utils, Worker # TBD +- label: ":amd: (MI300) Async Engine, Inputs, Utils, Worker" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -729,11 +749,12 @@ steps: commands: - pytest -v -s detokenizer - pytest -v -s -m 'not cpu_test' multimodal + - pytest -v -s multimodal/test_cache.py::test_sleep_wake_preserves_mm_cache_consistency - pytest -v -s utils_ #-------------------------------------------------------- mi300 · distributed --------------------------------------------------------# -- label: Distributed Comm Ops # TBD +- label: ":amd: (MI300) Distributed Comm Ops" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -750,7 +771,7 @@ steps: - pytest -v -s distributed/test_shm_buffer.py - pytest -v -s distributed/test_shm_storage.py -- label: EPLB Algorithm # TBD +- label: ":amd: (MI300) EPLB Algorithm" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -766,7 +787,7 @@ steps: - pytest -v -s distributed/test_eplb_algo.py - pytest -v -s distributed/test_eplb_utils.py -- label: EPLB Execution # TBD +- label: ":amd: (MI300) EPLB Execution" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -782,7 +803,7 @@ steps: - pytest -v -s distributed/test_eplb_execute.py - pytest -v -s distributed/test_eplb_spec_decode.py -- label: Distributed Tests (2xH100-2xMI300) # TBD +- label: ":amd: (MI300) Distributed Features" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -809,7 +830,7 @@ steps: - VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py - pytest -v -s tests/distributed/test_packed_tensor.py -- label: Distributed Tests (4xA100-4xMI300) # TBD +- label: ":amd: (MI300) Distributed" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -819,13 +840,15 @@ steps: working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ + - tests/models/test_vision.py commands: - pytest -v -s distributed/test_custom_all_reduce.py - torchrun --nproc_per_node=2 distributed/test_ca_buffer_sharing.py - TARGET_TEST_SUITE=MI300 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' - pytest -v -s -x lora/test_mixtral.py + - pytest -v -s models/test_vision.py -m 'distributed(num_gpus=4)' -- label: Distributed Torchrun + Examples (4 GPUs) # TBD +- label: ":amd: (MI300) Distributed Torchrun + Examples" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -851,7 +874,7 @@ steps: - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 ../examples/rl/rlhf_http_nccl.py - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 ../examples/rl/rlhf_http_ipc.py -- label: Elastic EP Scaling Test # TBD +- label: ":amd: (MI300) Elastic EP Scaling" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -869,7 +892,7 @@ steps: commands: - pytest -v -s distributed/test_elastic_ep.py -- label: RayExecutorV2 (4 GPUs) # TBD +- label: ":amd: (MI300) RayExecutorV2" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -892,7 +915,7 @@ steps: - pytest -v -s distributed/test_pipeline_parallel.py -k "ray" - TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/test_basic_correctness.py -k "ray" -- label: Distributed Tests (8xH100-8xMI300) # TBD +- label: ":amd: (MI300) Distributed DP + EP" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -911,7 +934,7 @@ steps: commands: - torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep -- label: Distributed Torchrun + Shutdown Tests (2 GPUs) # TBD +- label: ":amd: (MI300) Distributed Torchrun + Shutdown" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -935,7 +958,7 @@ steps: - HIP_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown - pytest -v -s v1/worker/test_worker_memory_snapshot.py -- label: Distributed Compile + Comm (4 GPUs) # TBD +- label: ":amd: (MI300) Distributed Compile + Comm" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -959,7 +982,7 @@ steps: #---------------------------------------------------------- mi300 · engine -----------------------------------------------------------# -- label: Engine # TBD +- label: ":amd: (MI300) Engine Core" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -973,12 +996,14 @@ steps: - tests/test_config - tests/test_logger - tests/test_vllm_port + - tests/jit_monitor/test_hooks.py + - tests/jit_monitor/test_hooks_gpu.py commands: - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py jit_monitor/test_hooks.py jit_monitor/test_hooks_gpu.py #-------------------------------------------------------- mi300 · entrypoints --------------------------------------------------------# -- label: Entrypoints Unit Tests # TBD +- label: ":amd: (MI300) Entrypoints Unit" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -997,7 +1022,7 @@ steps: - pytest -v -s entrypoints/weight_transfer - pytest -v -s entrypoints/launchers -- label: Entrypoints Integration (LLM) # TBD +- label: ":amd: (MI300) Entrypoints Integration (LLM)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1014,7 +1039,7 @@ steps: - pytest -v -s entrypoints/llm/test_generate.py # it needs a clean process - pytest -v -s entrypoints/llm/offline_mode # Needs to avoid interference with other tests -- label: Entrypoints Integration (API Server) # TBD +- label: ":amd: (MI300) Entrypoints Integration (API Server)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1032,7 +1057,7 @@ steps: - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc - pytest -v -s entrypoints/scale_out -- label: Entrypoints Integration (API Server OpenAI - Part 1) # TBD +- label: ":amd: (MI300) Entrypoints Integration (API Server OpenAI - Part 1)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1048,7 +1073,7 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/openai/ --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness -- label: Entrypoints Integration (API Server OpenAI - Part 2) # TBD +- label: ":amd: (MI300) Entrypoints Integration (API Server OpenAI - Part 2)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1065,7 +1090,7 @@ steps: - pytest -v -s entrypoints/openai/chat_completion - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py -- label: Entrypoints Integration (API Server Generate) # TBD +- label: ":amd: (MI300) Entrypoints Integration (API Server Generate)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1086,7 +1111,7 @@ steps: - pytest -v -s entrypoints/generate - pytest -v -s entrypoints/anthropic -- label: Entrypoints Integration (Responses API) # TBD +- label: ":amd: (MI300) Entrypoints Integration (Responses API)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1101,7 +1126,7 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/openai/responses -- label: Entrypoints Integration (Speech to Text) # TBD +- label: ":amd: (MI300) Entrypoints Integration (Speech to Text)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1115,7 +1140,7 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/speech_to_text -- label: Entrypoints Integration (Multimodal) +- label: ":amd: (MI300) Entrypoints Integration (Multimodal)" timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1130,7 +1155,7 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/multimodal -- label: Entrypoints Integration (Pooling) # TBD +- label: ":amd: (MI300) Entrypoints Integration (Pooling)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1145,7 +1170,7 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/pooling -- label: OpenAI API correctness # TBD +- label: ":amd: (MI300) OpenAI API Correctness" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1167,7 +1192,7 @@ steps: #----------------------------------------------------------- mi300 · evals -----------------------------------------------------------# -- label: DeepSeek V2-Lite Prefetch Offload Accuracy (H100-MI300) # TBD +- label: ":amd: (MI300) DeepSeek V2-Lite Prefetch Offload Accuracy" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1188,7 +1213,7 @@ steps: commands: - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_prefetch_offload.sh 0.25 200 8030 -- label: LM Eval Small Models # TBD +- label: ":amd: (MI300) LM Eval Small Models" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1206,7 +1231,7 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt -- label: MRCR Eval Small Models # TBD +- label: ":amd: (MI300) MRCR Eval Small Models" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1224,7 +1249,7 @@ steps: commands: - pytest -s -v evals/mrcr/test_mrcr_correctness.py --config-list-file=evals/mrcr/configs/models-small.txt -- label: LM Eval Small Models (MI300) # TBD +- label: ":amd: (MI300) LM Eval Small Models Harness" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1243,7 +1268,7 @@ steps: commands: - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-small-rocm.txt -- label: Multi-Modal Accuracy Eval (Small Models) # TBD +- label: ":amd: (MI300) Multimodal Accuracy Eval (Small Models)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1259,7 +1284,7 @@ steps: commands: - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-mm-small.txt --tp-size=1 -- label: GPQA Eval (GPT-OSS) (2xH100-2xMI300) # TBD +- label: ":amd: (MI300) GPQA Eval (GPT-OSS)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1282,7 +1307,7 @@ steps: - uv pip install --system 'gpt-oss[eval]==0.0.5' - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-gfx942.txt -- label: LM Eval Small Models (2xB200-2xMI300) # TBD +- label: ":amd: (MI300) LM Eval Small Models FP8 + Mixed" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1301,7 +1326,7 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt -- label: DeepSeek V2-Lite Sync EPLB Accuracy (4xH100-4xMI300) # TBD +- label: ":amd: (MI300) DeepSeek V2-Lite Sync EPLB Accuracy" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1324,7 +1349,7 @@ steps: commands: - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh 0.25 200 8010 -- label: LM Eval Large Models (4xA100-4xMI300) # TBD +- label: ":amd: (MI300) LM Eval Large Models Harness" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1345,7 +1370,7 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4 -- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (4xH100-4xMI300) # TBD +- label: ":amd: (MI300) Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1367,7 +1392,7 @@ steps: commands: - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 -- label: Qwen3-30B-A3B-FP8 DP4 Async EPLB Accuracy (4xH100-4xMI300) # TBD +- label: ":amd: (MI300) Qwen3-30B-A3B-FP8 DP4 Async EPLB Accuracy" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1389,7 +1414,7 @@ steps: commands: - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_dp4_async_eplb.sh 0.8 200 8050 -- label: Qwen3-Next-80B-A3B-Instruct MTP Async EPLB Accuracy # TBD +- label: ":amd: (MI300) Qwen3-Next-80B-A3B-Instruct MTP Async EPLB Accuracy" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1412,7 +1437,7 @@ steps: commands: - bash .buildkite/scripts/scheduled_integration_test/qwen3_next_mtp_async_eplb.sh 0.8 1319 8040 -- label: LM Eval Large Models (8xH200-8xMI300) # TBD +- label: ":amd: (MI300) LM Eval Large Models" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1433,9 +1458,10 @@ steps: - tests/evals/ commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - export PYTORCH_ROCM_ARCH=gfx942 - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt -- label: ROCm LM Eval Large Models (8 GPUs) # TBD +- label: ":amd: (MI300) LM Eval Large Models ROCm Harness" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1457,7 +1483,7 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm.txt --tp-size=8 -- label: LM Eval Large Models (4xH100-4xMI300) # TBD +- label: ":amd: (MI300) LM Eval Large Models FP8" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1480,7 +1506,7 @@ steps: #--------------------------------------------------------- mi300 · examples ----------------------------------------------------------# -- label: Examples # TBD +- label: ":amd: (MI300) Examples" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1517,7 +1543,7 @@ steps: #---------------------------------------------------------- mi300 · kernels ----------------------------------------------------------# -- label: vLLM IR Tests # TBD +- label: ":amd: (MI300) vLLM IR" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1533,7 +1559,7 @@ steps: - pytest -v -s tests/ir - pytest -v -s tests/kernels/ir -- label: Kernels MLA (MI300) # TBD +- label: ":amd: (MI300) MLA Kernels" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1555,7 +1581,7 @@ steps: - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_op_registration.py -- label: Kernels Attention Test %N # TBD +- label: ":amd: (MI300) Attention Kernels Shard %N" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1574,7 +1600,7 @@ steps: commands: - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT -- label: Kernels Core Operation Test %N # TBD +- label: ":amd: (MI300) Core Operation Kernels Shard %N" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1592,7 +1618,7 @@ steps: commands: - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT -- label: Kernels KDA Test # TBD +- label: ":amd: (MI300) KDA Kernels" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1612,7 +1638,7 @@ steps: commands: - pytest -v -s models/kimi_k3/test_kda.py models/kimi_k3/test_kda_metadata.py -- label: Kernels Mamba Test # TBD +- label: ":amd: (MI300) Mamba Kernels" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1627,7 +1653,7 @@ steps: commands: - pytest -v -s kernels/mamba -- label: Kernels MoE Test %N # TBD +- label: ":amd: (MI300) MoE Kernels Shard %N" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1653,7 +1679,7 @@ steps: - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT -- label: Kernels Quantization Test %N # TBD +- label: ":amd: (MI300) Quantization Kernels" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1677,7 +1703,7 @@ steps: commands: - pytest -v -s kernels/quantization -- label: Kernels FP8 MoE Test (2xH100-2xMI300) # TBD +- label: ":amd: (MI300) DeepEP FP8 MoE Kernels" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1698,7 +1724,7 @@ steps: #----------------------------------------------------------- mi300 · lora ------------------------------------------------------------# -- label: LoRA %N # TBD +- label: ":amd: (MI300) LoRA Shard %N" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1713,7 +1739,7 @@ steps: commands: - pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_qwen3_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py -- label: LoRA TP (Distributed) # TBD +- label: ":amd: (MI300) LoRA TP (Distributed)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1722,6 +1748,7 @@ steps: working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/lora + - vllm/model_executor/layers/fused_moe/ - tests/lora - vllm/platforms/rocm.py commands: @@ -1731,10 +1758,11 @@ steps: - pytest -v -s -x lora/test_olmoe_tp.py - pytest -v -s -x lora/test_gptoss_tp.py - pytest -v -s -x lora/test_qwen35_densemodel_lora.py + - pytest -v -s -x lora/test_gemma4_tp.py #------------------------------------------------------ mi300 · model_executor -------------------------------------------------------# -- label: Model Executor # TBD +- label: ":amd: (MI300) Model Executor" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1757,7 +1785,7 @@ steps: #---------------------------------------------------- mi300 · model_runner_v2 -------------------------------------------------------# -- label: Model Runner V2 Core Tests # TBD +- label: ":amd: (MI300) Model Runner V2 Core" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1782,7 +1810,7 @@ steps: - pytest -v -s v1/e2e/general/test_min_tokens.py - pytest -v -s entrypoints/llm/test_struct_output_generate.py -k "xgrammar and not speculative_config6 and not speculative_config7 and not speculative_config8 and not speculative_config0" -- label: Model Runner V2 Examples # TBD +- label: ":amd: (MI300) Model Runner V2 Examples" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1817,7 +1845,7 @@ steps: - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048 - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536 -- label: Model Runner V2 Distributed (2 GPUs) # TBD +- label: ":amd: (MI300) Model Runner V2 Distributed" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1839,7 +1867,7 @@ steps: - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py -k "not ray" - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py -- label: Model Runner V2 Pipeline Parallelism (4 GPUs) # TBD +- label: ":amd: (MI300) Model Runner V2 Pipeline Parallelism" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1859,7 +1887,7 @@ steps: - pytest -v -s distributed/test_pipeline_parallel.py -k "not ray" - pytest -v -s distributed/test_pp_cudagraph.py -k "not ray" -- label: Model Runner V2 Spec Decode # TBD +- label: ":amd: (MI300) Model Runner V2 Spec Decode" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1886,7 +1914,7 @@ steps: #------------------------------------------------------ mi300 · models / basic -------------------------------------------------------# -- label: Basic Models Tests (Extra Initialization) %N # TBD +- label: ":amd: (MI300) Basic Models (Extra Initialization) Shard %N" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1903,7 +1931,7 @@ steps: commands: - pytest -v -s models/test_initialization.py -k 'not test_can_initialize_small_subset' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB -- label: Basic Models Tests (Initialization) # TBD +- label: ":amd: (MI300) Basic Models (Initialization)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1916,7 +1944,7 @@ steps: commands: - pytest -v -s models/test_initialization.py::test_can_initialize_small_subset -- label: Basic Models Tests (Other) # TBD +- label: ":amd: (MI300) Basic Models (Other)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1928,12 +1956,14 @@ steps: - tests/models/test_terratorch.py - tests/models/transformers/test_backend.py - tests/models/test_registry.py + - tests/models/test_vision.py commands: - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py + - pytest -v -s models/test_vision.py::test_simple_mrope_vision_model_spatial_merge #----------------------------------------------------- mi300 · models / language -----------------------------------------------------# -- label: Language Models Test (Extended Pooling) # TBD +- label: ":amd: (MI300) Language Models (Extended Pooling)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1946,7 +1976,7 @@ steps: commands: - pytest -v -s models/language/pooling -m 'not core_model' -- label: Language Models Tests (Standard) # TBD +- label: ":amd: (MI300) Language Models (Standard)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1960,7 +1990,7 @@ steps: - pip freeze | grep -E 'torch' - pytest -v -s models/language -m 'core_model and (not slow_test)' -- label: Language Models Tests (Extra Standard) %N # TBD +- label: ":amd: (MI300) Language Models (Extra Standard) Shard %N" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1983,7 +2013,7 @@ steps: - pip freeze | grep -E 'torch' - pytest -v -s models/language -m 'core_model and slow_test' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB -- label: Language Models Test (Extended Generation) # TBD +- label: ":amd: (MI300) Language Models (Extended Generation)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1998,7 +2028,7 @@ steps: - CAUSAL_CONV1D_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)' -- label: Language Models Tests (Hybrid) %N # TBD +- label: ":amd: (MI300) Language Models (Hybrid) Shard %N" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2016,7 +2046,7 @@ steps: #---------------------------------------------------- mi300 · models / multimodal ----------------------------------------------------# -- label: Multi-Modal Models (Extended Generation 1) # TBD +- label: ":amd: (MI300) Multimodal Models (Extended Generation 1)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2031,7 +2061,7 @@ steps: - pytest -v -s models/multimodal/generation -m 'not core_model' --ignore models/multimodal/generation/test_common.py - pytest -v -s models/multimodal/test_mapping.py -- label: Multi-Modal Models (Extended Generation 2) # TBD +- label: ":amd: (MI300) Multimodal Models (Extended Generation 2)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2045,7 +2075,7 @@ steps: - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model' -- label: Multi-Modal Models (Extended Generation 3) # TBD +- label: ":amd: (MI300) Multimodal Models (Extended Generation 3)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2058,7 +2088,7 @@ steps: commands: - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model' -- label: "Multi-Modal Models (Standard) 1: qwen2" # TBD +- label: ":amd: (MI300) Multimodal Models (Standard) 1: qwen2" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2072,7 +2102,7 @@ steps: - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen2" - pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model -- label: "Multi-Modal Models (Standard) 3: llava + qwen2_vl" # TBD +- label: ":amd: (MI300) Multimodal Models (Standard) 3: llava + qwen2_vl" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2087,7 +2117,7 @@ steps: - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma" - pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model -- label: "Multi-Modal Models (Standard) 4: other + whisper" # TBD +- label: ":amd: (MI300) Multimodal Models (Standard) 4: other + whisper" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2095,14 +2125,14 @@ steps: working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/models/multimodal/generation - - tests/models/multimodal/test_mapping.py + - tests/models/multimodal commands: - - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/generation/test_vit_cudagraph.py --ignore models/multimodal/processing + - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_mm_prefix_lm.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/generation/test_vit_cudagraph.py --ignore models/multimodal/processing + - pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model - pytest -v -s models/multimodal/generation/test_memory_leak.py -m core_model - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model -- label: Multi-Modal Processor # 1h 42m +- label: ":amd: (MI300) Multimodal Processor" # 1h 42m timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2116,7 +2146,7 @@ steps: commands: - pytest -v -s models/multimodal/processing/test_tensor_schema.py -- label: Multi-Modal Models (Extended Pooling) # TBD +- label: ":amd: (MI300) Multimodal Models (Extended Pooling)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2129,7 +2159,7 @@ steps: commands: - pytest -v -s models/multimodal/pooling -m 'not core_model' -- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma" # TBD +- label: ":amd: (MI300) Multimodal Models (Standard) 2: qwen3 + gemma" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2141,11 +2171,12 @@ steps: - tests/models/multimodal commands: - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma" + - pytest -v -s models/multimodal/generation/test_mm_prefix_lm.py -m core_model - pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model #----------------------------------------------------- mi300 · models / quantized -----------------------------------------------------# -- label: Quantized Models Test # TBD +- label: ":amd: (MI300) Quantized Models" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2162,7 +2193,7 @@ steps: #-------------------------------------------------- mi300 · models / transformers ---------------------------------------------------# -- label: Transformers Nightly Models (Initialization) %N # TBD +- label: ":amd: (MI300) Transformers Nightly Models (Initialization) Shard %N" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2184,7 +2215,7 @@ steps: - pip install --upgrade git+https://github.com/huggingface/transformers - pytest -v -s tests/models/test_initialization.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB -- label: Transformers Nightly Models (Processing) %N # TBD +- label: ":amd: (MI300) Transformers Nightly Models (Processing) Shard %N" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2206,7 +2237,7 @@ steps: - pip install --upgrade git+https://github.com/huggingface/transformers - pytest -v -s tests/models/multimodal/processing/ --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB -- label: Transformers Nightly Models (Single) # TBD +- label: ":amd: (MI300) Transformers Nightly Models (Single)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2234,7 +2265,7 @@ steps: #---------------------------------------------------------- mi300 · plugins ----------------------------------------------------------# -- label: Plugin Tests (2 GPUs) # TBD +- label: ":amd: (MI300) Plugin Integration" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2286,7 +2317,7 @@ steps: - pytest -v -s plugins_tests/test_oot_registration_offline.py # it needs a clean process - pytest -v -s plugins_tests/lora_resolvers # unit tests for in-tree lora resolver plugins -- label: GGUF Plugin # TBD +- label: ":amd: (MI300) GGUF Plugin" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2305,7 +2336,7 @@ steps: #------------------------------------------------------- mi300 · rust_frontend -------------------------------------------------------# -- label: Rust Frontend OpenAI Coverage # TBD +- label: ":amd: (MI300) Rust Frontend OpenAI Coverage" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2338,7 +2369,7 @@ steps: - pytest -v -s entrypoints/serve/instrumentator/test_uds.py - pytest -v -s v1/sample/test_logprobs_e2e.py -k "test_prompt_logprobs_e2e_server" -- label: Rust Frontend Serve Admin Coverage # TBD +- label: ":amd: (MI300) Rust Frontend Serve + Admin Coverage" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2366,7 +2397,7 @@ steps: - pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist" - pytest -v -s entrypoints/serve/tokenize/test_tokenization.py -k "not tokenizer_info" -- label: Rust Frontend Core Correctness # TBD +- label: ":amd: (MI300) Rust Frontend Core Correctness" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2384,7 +2415,7 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine -- label: Rust Frontend Tool Use # TBD +- label: ":amd: (MI300) Rust Frontend Tool Use" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2403,7 +2434,7 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s tool_use --ignore=tool_use/mistral --models llama3.2 -k "not test_response_format_with_tool_choice_required and not test_parallel_tool_calls_false and not test_tool_call_and_choice" -- label: Rust Frontend Distributed # TBD +- label: ":amd: (MI300) Rust Frontend Distributed" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2434,7 +2465,7 @@ steps: #------------------------------------------------------- mi300 · quantization --------------------------------------------------------# -- label: Quantization # TBD +- label: ":amd: (MI300) Quantization" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2461,7 +2492,7 @@ steps: #--------------------------------------------------------- mi300 · samplers ----------------------------------------------------------# -- label: Samplers Test # TBD +- label: ":amd: (MI300) Samplers" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2482,7 +2513,7 @@ steps: #------------------------------------------------------------ mi300 · misc ------------------------------------------------------------# -- label: Regression # TBD +- label: ":amd: (MI300) Regression" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2497,7 +2528,7 @@ steps: #--------------------------------------------------------- mi300 · ray_compat ---------------------------------------------------------# -- label: Ray Dependency Compatibility Check # TBD +- label: ":amd: (MI300) Ray Dependency Compatibility Check" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2512,7 +2543,7 @@ steps: #------------------------------------------------------------ mi300 · v1 -------------------------------------------------------------# -- label: Acceptance Length Test (Large Models) # TBD +- label: ":amd: (MI300) Acceptance Length (Large Models)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2528,7 +2559,7 @@ steps: - export VLLM_ALLOW_INSECURE_SERIALIZATION=1 - pytest -v -s v1/spec_decode/test_acceptance_length.py -m slow_test -- label: e2e Scheduling (1 GPU) # TBD +- label: ":amd: (MI300) E2E Scheduling" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2542,7 +2573,7 @@ steps: commands: - pytest -v -s v1/e2e/general/test_async_scheduling.py -- label: Engine (1 GPU) # TBD +- label: ":amd: (MI300) V1 Engine" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2557,7 +2588,7 @@ steps: - pytest -v -s v1/engine/test_preprocess_error_handling.py - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py -- label: Spec Decode Draft Model # TBD +- label: ":amd: (MI300) Spec Decode Draft Model" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2575,7 +2606,7 @@ steps: commands: - pytest -v -s v1/e2e/spec_decode/draft_model/ -- label: Spec Decode Eagle # TBD +- label: ":amd: (MI300) Spec Decode Eagle" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2593,7 +2624,7 @@ steps: commands: - pytest -v -s v1/e2e/spec_decode/eagle/ -- label: Spec Decode Ngram + Suffix # TBD +- label: ":amd: (MI300) Spec Decode N-Gram + Suffix" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2607,11 +2638,13 @@ steps: - vllm/v1/sample/ - vllm/model_executor/layers/ - tests/v1/e2e/spec_decode/ + - tests/spec_decode/ - vllm/platforms/rocm.py commands: - pytest -v -s v1/e2e/spec_decode/ngram_suffix/ + - python3 spec_decode/test_custom_proposer.py -- label: Spec Decode Speculators + MTP # TBD +- label: ":amd: (MI300) Spec Decode Speculators + MTP" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2631,7 +2664,7 @@ steps: - pytest -v -s v1/e2e/spec_decode/speculators/ - pytest -v -s v1/e2e/spec_decode/mtp/ -- label: Speculators Correctness # TBD +- label: ":amd: (MI300) Speculators Correctness" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2663,7 +2696,7 @@ steps: - export VLLM_ALLOW_INSECURE_SERIALIZATION=1 - pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test -- label: V1 Spec Decode # TBD +- label: ":amd: (MI300) V1 Spec Decode" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2673,9 +2706,10 @@ steps: - vllm/ - tests/v1/spec_decode commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s -m 'not slow_test' v1/spec_decode -- label: Extract Hidden States Integration # TBD +- label: ":amd: (MI300) Extract Hidden States Integration" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2707,7 +2741,7 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s v1/kv_connector/extract_hidden_states_integration -- label: V1 attention (H100-MI300) %N # TBD +- label: ":amd: (MI300) V1 Attention Shard %N" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2726,7 +2760,7 @@ steps: commands: - pytest -v -s v1/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT -- label: V1 Core + KV + Metrics # TBD +- label: ":amd: (MI300) V1 Core + KV + Metrics" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2737,23 +2771,30 @@ steps: - tests/v1/core - tests/v1/executor - tests/v1/kv_offload + - tests/v1/simple_kv_offload - tests/v1/worker + - tests/v1/streaming_input - tests/v1/kv_connector/unit + - tests/v1/ec_connector/unit - tests/v1/metrics - tests/entrypoints/openai/correctness/test_lmeval.py commands: - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s -m 'not cpu_test' v1/core - pytest -v -s v1/executor - pytest -v -s v1/kv_offload + - pytest -v -s v1/simple_kv_offload - pytest -v -s v1/worker + - pytest -v -s v1/streaming_input - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit + - pytest -v -s -m 'not cpu_test' v1/ec_connector/unit - pytest -v -s -m 'not cpu_test' v1/metrics - pip install -U git+https://github.com/vllm-project/lm-evaluation-harness.git@streaming-api # - export HSA_NO_SCRATCH_RECLAIM=1 - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine -- label: V1 Sample + Logits # TBD +- label: ":amd: (MI300) V1 Sample + Logits" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2768,13 +2809,14 @@ steps: - tests/v1/test_request.py - tests/v1/test_outputs.py commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s v1/sample - pytest -v -s v1/logits_processors - pytest -v -s v1/test_oracle.py - pytest -v -s v1/test_request.py - pytest -v -s v1/test_outputs.py -- label: Distributed DP Tests (2 GPUs) # TBD +- label: ":amd: (MI300) Distributed DP Basic" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2790,16 +2832,16 @@ steps: - vllm/v1/engine/ - vllm/v1/worker/ - tests/v1/distributed - - tests/entrypoints/openai/test_multi_api_servers.py + - tests/entrypoints/launchers/api_server/test_multi_api_servers.py - vllm/platforms/rocm.py commands: - export NCCL_CUMEM_HOST_ENABLE=0 - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py - - DP_SIZE=2 pytest -v -s entrypoints/openai/test_multi_api_servers.py + - DP_SIZE=2 pytest -v -s entrypoints/launchers/api_server/test_multi_api_servers.py -- label: Metrics, Tracing (2 GPUs) # TBD +- label: ":amd: (MI300) Metrics, Tracing" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2810,6 +2852,7 @@ steps: source_file_dependencies: - vllm/ - tests/v1/tracing + - tests/tracing/ commands: - "pip install \ 'opentelemetry-sdk>=1.26.0' \ @@ -2817,8 +2860,9 @@ steps: 'opentelemetry-exporter-otlp>=1.26.0' \ 'opentelemetry-semantic-conventions-ai>=0.4.1'" - pytest -v -s v1/tracing + - pytest -v -s tracing -- label: V1 e2e (2 GPUs) # TBD +- label: ":amd: (MI300) V1 E2E" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2835,7 +2879,7 @@ steps: v1/e2e/spec_decode/draft_model/test_draft_model.py::test_draft_model_tensor_parallelism v1/e2e/spec_decode/draft_model/test_draft_model.py::test_draft_model_engine_args_tensor_parallelism -- label: NixlConnector PD + Spec Decode acceptance (2 GPUs) # TBD +- label: ":amd: (MI300) NixlConnector PD + Spec Decode acceptance" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2852,7 +2896,7 @@ steps: - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - KV_CACHE_MEMORY_BYTES=8G ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh -- label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs) # TBD +- label: ":amd: (MI300) CrossLayer KV layout Distributed NixlConnector PD accuracy" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2868,7 +2912,7 @@ steps: - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - CROSS_LAYERS_BLOCKS=True ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh -- label: Distributed DP Tests (4 GPUs) # TBD +- label: ":amd: (MI300) Distributed DP Extended" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2892,7 +2936,7 @@ steps: - pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp - pytest -v -s distributed/test_utils.py -- label: Distributed NixlConnector PD accuracy (4 GPUs) # TBD +- label: ":amd: (MI300) Distributed NixlConnector PD accuracy" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2908,7 +2952,7 @@ steps: - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh -- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs) # TBD +- label: ":amd: (MI300) DP EP Distributed NixlConnector PD accuracy" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2924,7 +2968,7 @@ steps: - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - DP_EP=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh -- label: Hybrid SSM NixlConnector PD accuracy tests (4 GPUs) # TBD +- label: ":amd: (MI300) Hybrid SSM NixlConnector PD accuracy" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2940,7 +2984,7 @@ steps: - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - HYBRID_SSM=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh -- label: Hybrid SSM NixlConnector PD prefix cache test (2 GPUs) # TBD +- label: ":amd: (MI300) Hybrid SSM NixlConnector PD prefix cache" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2958,7 +3002,7 @@ steps: - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh -- label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs) # TBD +- label: ":amd: (MI300) MultiConnector (Nixl+Offloading) PD accuracy" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2977,7 +3021,7 @@ steps: - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh -- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs) # TBD +- label: ":amd: (MI300) MultiConnector (Nixl+Offloading) PD edge cases" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2996,7 +3040,7 @@ steps: - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh -- label: V1 e2e (4xH100-4xMI300) # TBD +- label: ":amd: (MI300) V1 E2E Hybrid Chunked Prefill" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -3012,7 +3056,7 @@ steps: #------------------------------------------------------ mi300 · weight_loading -------------------------------------------------------# -- label: Weight Loading Multiple GPU # TBD +- label: ":amd: (MI300) Weight Loading Multi-GPU" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -3025,7 +3069,7 @@ steps: commands: - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-amd.txt -- label: Weight Loading Multiple GPU - Large Models # TBD +- label: ":amd: (MI300) Weight Loading Multi-GPU (Large Models)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -3047,7 +3091,7 @@ steps: #-------------------------------------------------------- mi355 · benchmarks ---------------------------------------------------------# -- label: Attention Benchmarks Smoke Test (B200-MI355) # TBD +- label: ":amd: (MI355) Attention Benchmark Smoke" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3064,7 +3108,7 @@ steps: #-------------------------------------------------------- mi355 · distributed --------------------------------------------------------# -- label: Distributed Tests (2xH100-2xMI355) # TBD +- label: ":amd: (MI355) Distributed Features" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3109,7 +3153,7 @@ steps: #-------------------------------------------------------- mi355 · entrypoints --------------------------------------------------------# -- label: Entrypoints Integration (API Server) # TBD +- label: ":amd: (MI355) Entrypoints Integration (API Server)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3127,7 +3171,7 @@ steps: - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc - pytest -v -s entrypoints/scale_out -- label: Entrypoints Integration (API Server OpenAI - Part 1) # TBD +- label: ":amd: (MI355) Entrypoints Integration (API Server OpenAI - Part 1)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3143,7 +3187,7 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness -- label: Entrypoints Integration (API Server OpenAI - Part 2) # TBD +- label: ":amd: (MI355) Entrypoints Integration (API Server OpenAI - Part 2)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3160,7 +3204,7 @@ steps: - pytest -v -s entrypoints/openai/chat_completion - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py -- label: Entrypoints Integration (API Server Generate) # TBD +- label: ":amd: (MI355) Entrypoints Integration (API Server Generate)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3181,7 +3225,7 @@ steps: - pytest -v -s entrypoints/generate - pytest -v -s entrypoints/anthropic -- label: Entrypoints Integration (Speech to Text) # TBD +- label: ":amd: (MI355) Entrypoints Integration (Speech to Text)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi355] dind: false @@ -3195,7 +3239,7 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/speech_to_text -- label: Entrypoints Integration (Multimodal) +- label: ":amd: (MI355) Entrypoints Integration (Multimodal)" timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi355] dind: false @@ -3209,7 +3253,7 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/multimodal -- label: Entrypoints Integration (Pooling) # TBD +- label: ":amd: (MI355) Entrypoints Integration (Pooling)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3225,7 +3269,7 @@ steps: #----------------------------------------------------------- mi355 · evals -----------------------------------------------------------# -- label: GPQA Eval (GPT-OSS) (2xB200-2xMI355) # TBD +- label: ":amd: (MI355) GPQA Eval (GPT-OSS)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3248,7 +3292,7 @@ steps: - uv pip install --system 'gpt-oss[eval]==0.0.5' - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-gfx950.txt -- label: LM Eval Qwen3-5 Models (B200-MI355) # TBD +- label: ":amd: (MI355) LM Eval Qwen3-5 Models" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3271,7 +3315,7 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-qwen35-mi355.txt -- label: LM Eval Small Models (2xB200-2xMI355) # TBD +- label: ":amd: (MI355) LM Eval Small Models" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3291,7 +3335,7 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt -- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (B200-MI355) # TBD +- label: ":amd: (MI355) Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3312,7 +3356,7 @@ steps: commands: - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 2 1 -- label: LM Eval Large Models (4xH100-4xMI355) # TBD +- label: ":amd: (MI355) LM Eval Large Models FP8" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3333,7 +3377,7 @@ steps: - export VLLM_USE_DEEP_GEMM=0 - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm-fp8.txt --tp-size=4 -- label: LM Eval Large Models (8xB200-8xMI355) # TBD +- label: ":amd: (MI355) LM Eval Large Models" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3358,7 +3402,7 @@ steps: #--------------------------------------------------------- mi355 · examples ----------------------------------------------------------# -- label: Examples # TBD +- label: ":amd: (MI355) Examples" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3394,7 +3438,7 @@ steps: #---------------------------------------------------------- mi355 · kernels ----------------------------------------------------------# -- label: Kernels Core Operation Test %N # TBD +- label: ":amd: (MI355) Core Operation Kernels Shard %N" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3412,7 +3456,7 @@ steps: commands: - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT -- label: Kernels (B200-MI355) # TBD +- label: ":amd: (MI355) Kernels" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3438,7 +3482,7 @@ steps: - pytest -v -s tests/kernels/attention/test_attention_selector.py - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py -- label: Kernels MLA (MI355) # TBD +- label: ":amd: (MI355) MLA Kernels" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3460,7 +3504,7 @@ steps: - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_op_registration.py -- label: Kernels Attention Test %N # TBD +- label: ":amd: (MI355) Attention Kernels Shard %N" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3478,7 +3522,7 @@ steps: commands: - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT -- label: Kernels MoE Test %N # TBD +- label: ":amd: (MI355) MoE Kernels Shard %N" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3512,7 +3556,7 @@ steps: --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT -- label: Kernels FusedMoE Layer Test (2xB200-2xMI355) # TBD +- label: ":amd: (MI355) FusedMoE Layer Kernels" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3538,7 +3582,7 @@ steps: commands: - pytest -v -s kernels/moe/test_moe_layer.py -- label: Kernels Quantization Test %N # TBD +- label: ":amd: (MI355) Quantization Kernels" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3561,7 +3605,7 @@ steps: commands: - pytest -v -s kernels/quantization -- label: Kernels FP8 MoE Test (2xH100-1xMI355) # TBD +- label: ":amd: (MI355) FP8 MoE Kernels" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3582,7 +3626,7 @@ steps: - pytest -v -s kernels/moe/test_triton_moe_no_act_mul.py - pytest -v -s kernels/moe/test_triton_moe_ptpc_fp8.py -- label: Kernels FP8 MoE Test (2xH100-2xMI355) # TBD +- label: ":amd: (MI355) DeepEP FP8 MoE Kernels" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3603,7 +3647,7 @@ steps: #----------------------------------------------------- mi355 · models / language -----------------------------------------------------# -- label: Language Models Test (Extended Generation) # TBD +- label: ":amd: (MI355) Language Models (Extended Generation)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3617,7 +3661,7 @@ steps: - CAUSAL_CONV1D_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)' -- label: Language Models Test (Extended Pooling) # TBD +- label: ":amd: (MI355) Language Models (Extended Pooling)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3630,7 +3674,7 @@ steps: commands: - pytest -v -s models/language/pooling -m 'not core_model' -- label: Language Models Test (PPL) # TBD +- label: ":amd: (MI355) Language Models (PPL)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3659,7 +3703,7 @@ steps: commands: - pytest -v -s models/language/generation_ppl_test -- label: Language Models Tests (Standard) # TBD +- label: ":amd: (MI355) Language Models (Standard)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3674,7 +3718,7 @@ steps: #---------------------------------------------------- mi355 · models / multimodal ----------------------------------------------------# -- label: Multi-Modal Models (Extended Generation 1) # TBD +- label: ":amd: (MI355) Multimodal Models (Extended Generation 1)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3689,7 +3733,7 @@ steps: - pytest -v -s models/multimodal/generation -m 'not core_model' --ignore models/multimodal/generation/test_common.py - pytest -v -s models/multimodal/test_mapping.py -- label: Multi-Modal Models (Extended Generation 3) # TBD +- label: ":amd: (MI355) Multimodal Models (Extended Generation 3)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3702,7 +3746,7 @@ steps: commands: - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model' -- label: Multi-Modal Models (Extended Pooling) # TBD +- label: ":amd: (MI355) Multimodal Models (Extended Pooling)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3715,7 +3759,7 @@ steps: commands: - pytest -v -s models/multimodal/pooling -m 'not core_model' -- label: "Multi-Modal Models (Standard) 1: qwen2" # TBD +- label: ":amd: (MI355) Multimodal Models (Standard) 1: qwen2" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3729,7 +3773,7 @@ steps: - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen2" - pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model -- label: "Multi-Modal Models (Standard) 4: other + whisper" # TBD +- label: ":amd: (MI355) Multimodal Models (Standard) 4: other + whisper" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3738,16 +3782,16 @@ steps: working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/models/multimodal/generation + - tests/models/multimodal commands: - - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/generation/test_vit_cudagraph.py --ignore models/multimodal/processing + - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_mm_prefix_lm.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/generation/test_vit_cudagraph.py --ignore models/multimodal/processing - pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model - pytest -v -s models/multimodal/generation/test_memory_leak.py -m core_model - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model #----------------------------------------------------- mi355 · models / quantized -----------------------------------------------------# -- label: Quantized Models Test # TBD +- label: ":amd: (MI355) Quantized Models" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3764,7 +3808,7 @@ steps: #------------------------------------------------------- mi355 · quantization --------------------------------------------------------# -- label: Quantization # TBD +- label: ":amd: (MI355) Quantization" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3810,7 +3854,7 @@ steps: #------------------------------------------------------------ mi355 · v1 -------------------------------------------------------------# -- label: V1 attention (B200-MI355) %N # TBD +- label: ":amd: (MI355) V1 Attention Shard %N" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3828,7 +3872,7 @@ steps: commands: - pytest -v -s v1/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT -- label: V1 Core + KV + Metrics # TBD +- label: ":amd: (MI355) V1 Core + KV + Metrics" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3857,7 +3901,7 @@ steps: - pip install -U git+https://github.com/vllm-project/lm-evaluation-harness.git@streaming-api - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine -- label: V1 Sample + Logits # TBD +- label: ":amd: (MI355) V1 Sample + Logits" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3878,7 +3922,7 @@ steps: - pytest -v -s v1/test_request.py - pytest -v -s v1/test_outputs.py -- label: V1 Spec Decode # TBD +- label: ":amd: (MI355) V1 Spec Decode" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3892,7 +3936,7 @@ steps: #------------------------------------------------------ mi355 · weight_loading -------------------------------------------------------# -- label: Weight Loading Multiple GPU # TBD +- label: ":amd: (MI355) Weight Loading Multi-GPU" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3905,7 +3949,7 @@ steps: commands: - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-amd.txt -- label: Weight Loading Multiple GPU - Large Models # TBD +- label: ":amd: (MI355) Weight Loading Multi-GPU (Large Models)" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3921,7 +3965,7 @@ steps: #----------------------------------------------------------- mi355 · misc ------------------------------------------------------------# -- label: Regression # TBD +- label: ":amd: (MI355) Regression" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index b4b2b5ddc605..771e43bcee66 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -31,14 +31,14 @@ steps: - vllm/v1/engine/ - vllm/v1/worker/ - tests/v1/distributed - - tests/entrypoints/openai/test_multi_api_servers.py + - tests/entrypoints/launchers/api_server/test_multi_api_servers.py commands: # https://github.com/NVIDIA/nccl/issues/1838 - export NCCL_CUMEM_HOST_ENABLE=0 - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py - - DP_SIZE=2 pytest -v -s entrypoints/openai/test_multi_api_servers.py + - DP_SIZE=2 pytest -v -s entrypoints/launchers/api_server/test_multi_api_servers.py mirror: amd: label: ":amd: (MI300) Distributed DP Basic" @@ -56,7 +56,7 @@ steps: - vllm/v1/engine/ - vllm/v1/worker/ - tests/v1/distributed - - tests/entrypoints/openai/test_multi_api_servers.py + - tests/entrypoints/launchers/api_server/test_multi_api_servers.py - vllm/platforms/rocm.py - label: ":nvidia: (L4) Distributed Compile + RPC" diff --git a/.buildkite/test_areas/models_multimodal.yaml b/.buildkite/test_areas/models_multimodal.yaml index 62883c341e37..340c97527cbd 100644 --- a/.buildkite/test_areas/models_multimodal.yaml +++ b/.buildkite/test_areas/models_multimodal.yaml @@ -80,7 +80,7 @@ steps: label: ":amd: (MI300) Multimodal Models (Standard) 4: other + whisper" dind: false device: mi300_1 - timeout_in_minutes: 50 + timeout_in_minutes: 70 depends_on: - image-build-amd diff --git a/tests/models/test_vision.py b/tests/models/test_vision.py index 987324eb8677..24260e2ef804 100644 --- a/tests/models/test_vision.py +++ b/tests/models/test_vision.py @@ -102,7 +102,7 @@ def run_dp_sharded_vision_model_vs_direct( # Set random seed for reproducibility set_random_seed(0) - device = f"{current_platform.device_name}:{local_rank}" + device = f"{current_platform.device_type}:{local_rank}" torch.accelerator.set_device_index(device) torch.set_default_device(device) @@ -288,7 +288,7 @@ def run_dp_sharded_mrope_vision_model_vs_direct( """ # Set random seed for reproducibility set_random_seed(0) - device = f"{current_platform.device_name}:{local_rank}" + device = f"{current_platform.device_type}:{local_rank}" torch.accelerator.set_device_index(device) torch.set_default_device(device) @@ -365,7 +365,7 @@ def run_dp_sharded_mrope_vision_model_empty_input_worker( ): """Test run_dp_sharded_mrope_vision_model with empty input.""" # Set up distributed environment - device = f"{current_platform.device_name}:{local_rank}" + device = f"{current_platform.device_type}:{local_rank}" torch.accelerator.set_device_index(device) torch.set_default_device(device) @@ -414,7 +414,7 @@ def run_dp_sharded_mrope_vision_model_uneven_load_worker( """Test run_dp_sharded_mrope_vision_model with uneven load distribution.""" # Set up distributed environment set_random_seed(123) - device = f"{current_platform.device_name}:{local_rank}" + device = f"{current_platform.device_type}:{local_rank}" torch.accelerator.set_device_index(device) torch.set_default_device(device) From 11baa0ebd88480b0a3e30138e4ee189420cd33be Mon Sep 17 00:00:00 2001 From: Xin Yang <105740670+xyang16@users.noreply.github.com> Date: Wed, 19 Aug 2026 20:52:42 -0700 Subject: [PATCH 182/839] [Attention] Avoid redundant mask compute in GDN metadata build (#52078) Signed-off-by: Xin Yang --- vllm/v1/attention/backends/gdn_attn.py | 24 +++++++++++------------- 1 file changed, 11 insertions(+), 13 deletions(-) diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 1a07a6957e3d..0e843e31baa7 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -228,20 +228,17 @@ def build( # type: ignore[override] ) spec_sequence_masks_cpu: torch.Tensor | None = None - if ( - not self.use_spec_decode - or num_decode_draft_tokens_cpu is None - or num_decode_draft_tokens_cpu[num_decode_draft_tokens_cpu >= 0] - .sum() - .item() - == 0 - ): + if not self.use_spec_decode or num_decode_draft_tokens_cpu is None: spec_sequence_masks = None num_spec_decodes = 0 else: spec_sequence_masks_cpu = num_decode_draft_tokens_cpu >= 0 num_spec_decodes = spec_sequence_masks_cpu.sum().item() - if num_spec_decodes == 0: + if ( + num_spec_decodes == 0 + or num_decode_draft_tokens_cpu[spec_sequence_masks_cpu].sum().item() + == 0 + ): spec_sequence_masks = None spec_sequence_masks_cpu = None else: @@ -265,10 +262,11 @@ def build( # type: ignore[override] else: query_lens = query_start_loc[1:] - query_start_loc[:-1] assert spec_sequence_masks_cpu is not None + non_spec_sequence_masks_cpu = ~spec_sequence_masks_cpu query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] # Use CPU tensors to avoid CPU-GPU sync - non_spec_query_lens_cpu = query_lens_cpu[~spec_sequence_masks_cpu] + non_spec_query_lens_cpu = query_lens_cpu[non_spec_sequence_masks_cpu] num_decodes = (non_spec_query_lens_cpu == 1).sum().item() # Exclude zero-length padded sequences from prefill count. num_zero_len = (non_spec_query_lens_cpu == 0).sum().item() @@ -330,7 +328,7 @@ def build( # type: ignore[override] spec_sequence_masks_cpu, : self.num_spec + 1 ] non_spec_state_indices_tensor = block_table_tensor[ - ~spec_sequence_masks_cpu, 0 + non_spec_sequence_masks_cpu, 0 ] spec_query_start_loc = torch.zeros( @@ -349,7 +347,7 @@ def build( # type: ignore[override] device=query_start_loc.device, ) torch.cumsum( - query_lens[~spec_sequence_masks_cpu], + query_lens[non_spec_sequence_masks_cpu], dim=0, out=non_spec_query_start_loc[1:], ) @@ -358,7 +356,7 @@ def build( # type: ignore[override] dtype=torch.int32, ) torch.cumsum( - query_lens_cpu[~spec_sequence_masks_cpu], + query_lens_cpu[non_spec_sequence_masks_cpu], dim=0, out=non_spec_query_start_loc_cpu[1:], ) From 9b2aef512440f024fc98cd0ac954cd452613ad5b Mon Sep 17 00:00:00 2001 From: Amit Moryossef Date: Thu, 20 Aug 2026 06:23:09 +0200 Subject: [PATCH 183/839] [Bugfix] Video loading: sample over presentable frames, not header sample count (MP4 edit-list trims) (#48608) Signed-off-by: AmitMY Signed-off-by: Isotr0py Co-authored-by: Claude Fable 5 Co-authored-by: Isotr0py --- tests/multimodal/test_video.py | 52 +++++++++++++++++++++++- tests/multimodal/utils.py | 37 +++++++++++++++++ vllm/multimodal/video_decoders/opencv.py | 21 ++++++++++ vllm/multimodal/video_decoders/pyav.py | 8 ++++ 4 files changed, 117 insertions(+), 1 deletion(-) diff --git a/tests/multimodal/test_video.py b/tests/multimodal/test_video.py index 9ab718d9457d..4020900696f3 100644 --- a/tests/multimodal/test_video.py +++ b/tests/multimodal/test_video.py @@ -40,7 +40,11 @@ from vllm.platforms import current_platform from vllm.transformers_utils.processor import get_video_processor_cls_name_from_config -from .utils import create_long_gop_video, create_video_from_image +from .utils import ( + create_edit_list_trimmed_video, + create_long_gop_video, + create_video_from_image, +) pytestmark = pytest.mark.cpu_test @@ -862,6 +866,44 @@ def test_video_backend_handles_broken_frames(monkeypatch: pytest.MonkeyPatch): ) +def test_video_backend_handles_edit_list_trimmed_video( + monkeypatch: pytest.MonkeyPatch, +): + """ + An mp4 edit list (e.g. from a lossless ``ffmpeg -ss ... -c copy`` cut) + hides the decode lead-in: the header still counts every physical sample + while sequential decode only yields the visible frames. Sampling over the + header count used to collapse such videos to the few indices below the + visible count (a single frame for the Qwen loaders) — the loader must + resample over the true stream length instead. + """ + with monkeypatch.context() as m: + m.setenv("VLLM_VIDEO_LOADER_BACKEND", "opencv") + + video_data, num_visible = create_edit_list_trimmed_video( + num_frames=90, trim_start_frame=60 + ) + + loader = VIDEO_LOADER_REGISTRY.load("opencv") + for backend in ["opencv", "pyav"]: + frames, metadata = loader.load_bytes( + video_data, num_frames=-1, backend=backend + ) + assert metadata["total_num_frames"] == num_visible, backend + assert frames.shape[0] == num_visible, backend + assert len(metadata["frames_indices"]) == num_visible, backend + # The green channel encodes the source frame index: the visible + # frames must be the trailing ones (60..89), not the lead-in. + mean_green = frames[..., 1].reshape(frames.shape[0], -1).mean(axis=1) + assert abs(mean_green[0] - 60) <= 5, backend + assert abs(mean_green[-1] - 89) <= 5, backend + + # The Qwen samplers used to degenerate to a single frame here. + qwen_frames, qwen_metadata = Qwen2VLVideoBackend.load_bytes(video_data) + assert qwen_metadata["total_num_frames"] == num_visible + assert qwen_frames.shape[0] >= 4 + + # ============================================================================ # Frame Recovery Tests # ============================================================================ @@ -912,6 +954,14 @@ def retrieve(self): def get(self, prop): return self._cap.get(prop) + def set(self, prop, value): + # get_video_metadata probes the stream end with a seek; + # keep the simulated frame counter aligned with rewinds. + result = self._cap.set(prop, value) + if prop == cv2.CAP_PROP_POS_FRAMES: + self._current_frame = int(value) - 1 + return result + def isOpened(self): return self._cap.isOpened() diff --git a/tests/multimodal/utils.py b/tests/multimodal/utils.py index bae0a9d2942c..23a05a094422 100644 --- a/tests/multimodal/utils.py +++ b/tests/multimodal/utils.py @@ -103,6 +103,43 @@ def create_long_gop_video( return buf.getvalue() +def create_edit_list_trimmed_video( + num_frames: int = 90, + trim_start_frame: int = 60, + fps: int = 30, +) -> tuple[bytes, int]: + """Stream-copy-cut a long-GOP clip so an mp4 edit list hides the lead-in. + + Remuxes the tail of a single-keyframe clip without re-encoding, rebasing + timestamps at ``trim_start_frame``: the lead-in packets are still needed + to decode, so they stay in the file at negative pts and the muxer records + an edit list. The header sample count stays ``num_frames`` while only the + trailing frames are presentable — the shape produced by lossless trims + (e.g. ``ffmpeg -ss ... -c copy``). Returns (video_bytes, visible_frames). + """ + import io + + import av + + src = create_long_gop_video(num_frames=num_frames, fps=fps) + buf = io.BytesIO() + with ( + av.open(io.BytesIO(src)) as source, + av.open(buf, mode="w", format="mp4") as out, + ): + in_stream = source.streams.video[0] + out_stream = out.add_stream_from_template(in_stream) + start_pts = round(trim_start_frame / fps / in_stream.time_base) + for packet in source.demux(in_stream): + if packet.pts is None or packet.dts is None: + continue + packet.pts -= start_pts + packet.dts -= start_pts + packet.stream = out_stream + out.mux(packet) + return buf.getvalue(), num_frames - trim_start_frame + + def cosine_similarity(A: npt.NDArray, B: npt.NDArray, axis: int = -1) -> npt.NDArray: """Compute cosine similarity between two vectors.""" return np.sum(A * B, axis=axis) / ( diff --git a/vllm/multimodal/video_decoders/opencv.py b/vllm/multimodal/video_decoders/opencv.py index e5e474eab606..e46979dd1d4c 100644 --- a/vllm/multimodal/video_decoders/opencv.py +++ b/vllm/multimodal/video_decoders/opencv.py @@ -78,6 +78,27 @@ def open_video_capture(cls, data: bytes) -> "cv2.VideoCapture": def get_video_metadata(cap: "cv2.VideoCapture") -> VideoSourceMetadata: total_frames_num = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) original_fps = cap.get(cv2.CAP_PROP_FPS) + # CAP_PROP_FRAME_COUNT counts every physical sample in the container, + # overstating the presentable frames when an mp4 edit list hides the + # decode lead-in (e.g. a lossless `ffmpeg -ss ... -c copy` trim). + # Cross-check with an end-of-stream seek: the FFMPEG backend reports + # the resulting position against the edit-list-aware timeline, which + # exposes the true frame count without decoding anything. The seek + # may land mid-stream on corrupt files, so only trust it when the + # stream really ends there. A one-frame tolerance keeps mere header + # miscounts from being "corrected". + if total_frames_num > 0 and cap.set(cv2.CAP_PROP_POS_AVI_RATIO, 1.0): + visible_frames = int(cap.get(cv2.CAP_PROP_POS_FRAMES)) + at_stream_end = not cap.grab() + cap.set(cv2.CAP_PROP_POS_FRAMES, 0) + if at_stream_end and 0 < visible_frames < total_frames_num - 1: + logger.warning( + "Video header claims %d frames but only %d are " + "presentable; sampling over the true frame count.", + total_frames_num, + visible_frames, + ) + total_frames_num = visible_frames duration = total_frames_num / original_fps if original_fps > 0 else 0 return VideoSourceMetadata( total_frames_num=total_frames_num, diff --git a/vllm/multimodal/video_decoders/pyav.py b/vllm/multimodal/video_decoders/pyav.py index dc82b01230d1..80c282966668 100644 --- a/vllm/multimodal/video_decoders/pyav.py +++ b/vllm/multimodal/video_decoders/pyav.py @@ -62,6 +62,14 @@ def get_metadata( duration = float(stream.duration * stream.time_base) if stream.duration else 0.0 if total_frames == 0 and duration > 0 and fps > 0: total_frames = int(duration * fps) + elif duration > 0 and fps > 0 and total_frames > round(duration * fps) + 1: + # The header sample count can exceed the frames the presentation + # timeline actually holds — e.g. an mp4 edit list hides the decode + # lead-in of a stream-copied cut, but ``stream.frames`` still + # counts the hidden samples. Sampling indices past the visible + # range would silently collapse onto the last visible frame, so + # trust the (edit-list-aware) duration instead. + total_frames = round(duration * fps) return VideoSourceMetadata(total_frames, fps, duration) @staticmethod From cd5035379c4c565fc9265cf121bc1c8364bc5eca Mon Sep 17 00:00:00 2001 From: akii96 Date: Thu, 20 Aug 2026 07:33:37 +0300 Subject: [PATCH 184/839] [ROCm] [Bugfix] Preserve CPU query offsets during capture (#51585) Signed-off-by: Aakif Nawaz --- vllm/v1/attention/backends/rocm_attn.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/vllm/v1/attention/backends/rocm_attn.py b/vllm/v1/attention/backends/rocm_attn.py index 331b5c3fb352..f73e558ec263 100644 --- a/vllm/v1/attention/backends/rocm_attn.py +++ b/vllm/v1/attention/backends/rocm_attn.py @@ -103,11 +103,9 @@ def build_for_cudagraph_capture( # slow, so here we set it to 1. attn_metadata.seq_lens.fill_(1) - # Here we set the query start locs to 0. This is to - # cover up an invalid memory access in the prefix_prefil kernel - # that we run into during graph capture (#25985) + # Zero device query start locations to avoid invalid memory access in + # the prefix prefill kernel during graph capture (#25985). common_attn_metadata.query_start_loc.zero_() - common_attn_metadata.query_start_loc_cpu.zero_() return attn_metadata From d626108b1841888ec90aced33367149a6bbc7e4b Mon Sep 17 00:00:00 2001 From: Shanshan Shen <467638484@qq.com> Date: Thu, 20 Aug 2026 12:51:04 +0800 Subject: [PATCH 185/839] [ROCm][Perf] Fuse DeepSeek-V4 mHC post/pre and RMSNorm with AITER (#52737) Signed-off-by: shen-shanshan <467638484@qq.com> Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com> --- vllm/_aiter_ops.py | 96 +++++++++++++++++++++ vllm/model_executor/kernels/mhc/aiter.py | 102 +++++++++++++++++++++++ vllm/model_executor/layers/mhc.py | 30 ++++++- vllm/models/deepseek_v4/amd/dspark.py | 4 +- vllm/models/deepseek_v4/amd/model.py | 50 +++++++++-- 5 files changed, 272 insertions(+), 10 deletions(-) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index b6a504d0cfbe..94ca4f7fe86e 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -3299,6 +3299,8 @@ def mhc_pre( hc_sinkhorn_eps: float, hc_post_mult_value: float, sinkhorn_repeat: int, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 0.0, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ Forward pass for mHC pre block. @@ -3313,6 +3315,8 @@ def mhc_pre( hc_sinkhorn_eps: sinkhorn epsilon hc_post_mult_value: post-mix multiplier value sinkhorn_repeat: number of sinkhorn iterations + norm_weight: optional RMSNorm weight fused into the pre kernel + norm_eps: epsilon for the fused RMSNorm when norm_weight is set Returns: post_mix: shape (..., hc_mult), dtype torch.float32 @@ -3380,6 +3384,8 @@ def mhc_pre( hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat, + norm_weight, + norm_eps, ) return ( post_mix.view(*outer_shape, hc_mult, 1), @@ -3465,5 +3471,95 @@ def mhc_post( ) return out.view_as(residual) + @staticmethod + def mhc_fused_post_pre( + x: torch.Tensor, + residual: torch.Tensor, + post_layer_mix: torch.Tensor, + comb_res_mix: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 0.0, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Fused mHC post + next mHC pre via AITER. + + Returns residual_cur, post_mix, comb_mix, layer_input in vLLM order + (AITER returns post_mix, comb_mix, layer_input, next_residual). + """ + from aiter.ops.mhc import mhc_fused_post_pre + + assert x.dtype == torch.bfloat16 + assert residual.dtype == torch.bfloat16 + assert fn.dtype == torch.float32 + assert hc_scale.dtype == torch.float32 + assert hc_base.dtype == torch.float32 + + hc_mult = residual.shape[-2] + hidden_size = residual.shape[-1] + outer_shape = residual.shape[:-2] + + residual_flat = residual.view(-1, hc_mult, hidden_size) + num_tokens = residual_flat.shape[0] + x_flat = x.view(num_tokens, hidden_size) + post_flat = post_layer_mix.view(num_tokens, hc_mult, 1) + comb_flat = comb_res_mix.view(num_tokens, hc_mult, hc_mult) + + if num_tokens == 0: + return ( + torch.empty_like(residual_flat).view_as(residual), + torch.empty( + *outer_shape, + hc_mult, + 1, + dtype=torch.float32, + device=residual.device, + ), + torch.empty( + *outer_shape, + hc_mult, + hc_mult, + dtype=torch.float32, + device=residual.device, + ), + torch.empty( + *outer_shape, + hidden_size, + dtype=torch.bfloat16, + device=residual.device, + ), + ) + + with torch.device(residual_flat.device): + post_mix, comb_mix, layer_input, next_residual = mhc_fused_post_pre( + x_flat, + residual_flat, + post_flat, + comb_flat, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + norm_weight, + norm_eps, + ) + + return ( + next_residual.view_as(residual), + post_mix.view(*outer_shape, hc_mult, 1), + comb_mix.view(*outer_shape, hc_mult, hc_mult), + layer_input.view(*outer_shape, hidden_size), + ) + rocm_aiter_ops.register_ops_once() diff --git a/vllm/model_executor/kernels/mhc/aiter.py b/vllm/model_executor/kernels/mhc/aiter.py index e844b4191e30..ce7690ec1d21 100644 --- a/vllm/model_executor/kernels/mhc/aiter.py +++ b/vllm/model_executor/kernels/mhc/aiter.py @@ -16,6 +16,8 @@ def mhc_pre_aiter( hc_post_mult_value: float, sinkhorn_repeat: int, n_splits: int = 1, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 0.0, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ Forward pass for mHC pre block. @@ -31,6 +33,8 @@ def mhc_pre_aiter( hc_post_mult_value: post-mix multiplier value sinkhorn_repeat: number of sinkhorn iterations n_splits: split-k factor; + norm_weight: optional RMSNorm weight fused into the pre kernel + norm_eps: epsilon for the fused RMSNorm when norm_weight is set Returns: post_mix: shape (..., hc_mult), dtype torch.float32 @@ -52,6 +56,8 @@ def mhc_pre_aiter( hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat, + norm_weight, + norm_eps, ) @@ -66,6 +72,8 @@ def _mhc_pre_aiter_fake( hc_post_mult_value: float, sinkhorn_repeat: int, n_splits: int = 1, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 0.0, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: hc_mult = residual.shape[-2] hidden_size = residual.shape[-1] @@ -124,6 +132,94 @@ def _mhc_post_aiter_fake( return torch.empty_like(residual) +def mhc_fused_post_pre_aiter( + x: torch.Tensor, + residual: torch.Tensor, + post_layer_mix: torch.Tensor, + comb_res_mix: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + n_splits: int = 1, + tile_n: int = 1, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 0.0, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Fused mHC post + next mHC pre on ROCm via AITER. + + Returns residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur. + """ + hidden_size = residual.shape[-1] + assert hidden_size % 256 == 0 + from vllm._aiter_ops import rocm_aiter_ops + + return rocm_aiter_ops.mhc_fused_post_pre( + x, + residual, + post_layer_mix, + comb_res_mix, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + norm_weight, + norm_eps, + ) + + +def _mhc_fused_post_pre_aiter_fake( + x: torch.Tensor, + residual: torch.Tensor, + post_layer_mix: torch.Tensor, + comb_res_mix: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + n_splits: int = 1, + tile_n: int = 1, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 0.0, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + hc_mult = residual.shape[-2] + hidden_size = residual.shape[-1] + outer_shape = residual.shape[:-2] + post_mix = torch.empty( + *outer_shape, + hc_mult, + 1, + dtype=torch.float32, + device=residual.device, + ) + comb_mix = torch.empty( + *outer_shape, + hc_mult, + hc_mult, + dtype=torch.float32, + device=residual.device, + ) + layer_input = torch.empty( + *outer_shape, + hidden_size, + dtype=torch.bfloat16, + device=residual.device, + ) + return torch.empty_like(residual), post_mix, comb_mix, layer_input + + direct_register_custom_op( op_name="mhc_pre_aiter", op_func=mhc_pre_aiter, @@ -136,3 +232,9 @@ def _mhc_post_aiter_fake( mutates_args=[], fake_impl=_mhc_post_aiter_fake, ) +direct_register_custom_op( + op_name="mhc_fused_post_pre_aiter", + op_func=mhc_fused_post_pre_aiter, + mutates_args=[], + fake_impl=_mhc_fused_post_pre_aiter_fake, +) diff --git a/vllm/model_executor/layers/mhc.py b/vllm/model_executor/layers/mhc.py index da6bbf2e0084..8f1c22f16ac3 100644 --- a/vllm/model_executor/layers/mhc.py +++ b/vllm/model_executor/layers/mhc.py @@ -90,7 +90,8 @@ def forward_hip( norm_eps: float = 0.0, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: hidden_size = residual.shape[-1] - if HAS_AITER_MHC and hidden_size % 256 == 0: + hc_mult = residual.shape[-2] + if HAS_AITER_MHC and hidden_size % 256 == 0 and hc_mult == 4: return torch.ops.vllm.mhc_pre_aiter( residual, fn, @@ -101,6 +102,9 @@ def forward_hip( hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat, + n_splits, + norm_weight, + norm_eps, ) elif HAS_TILELANG_MHC: return torch.ops.vllm.mhc_pre_tilelang( @@ -222,7 +226,8 @@ def forward_hip( comb_res_mix: torch.Tensor, ) -> torch.Tensor: hidden_size = residual.shape[-1] - if HAS_AITER_MHC and hidden_size % 256 == 0: + hc_mult = residual.shape[-2] + if HAS_AITER_MHC and hidden_size % 256 == 0 and hc_mult == 4: return torch.ops.vllm.mhc_post_aiter( x, residual, @@ -444,6 +449,27 @@ def forward_hip( norm_weight: torch.Tensor | None = None, norm_eps: float = 0.0, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + hidden_size = residual.shape[-1] + hc_mult = residual.shape[-2] + if HAS_AITER_MHC and hidden_size % 256 == 0 and hc_mult == 4: + return torch.ops.vllm.mhc_fused_post_pre_aiter( + x, + residual, + post_layer_mix, + comb_res_mix, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + n_splits, + tile_n, + norm_weight, + norm_eps, + ) if HAS_TILELANG_MHC: return torch.ops.vllm.mhc_fused_post_pre_tilelang( x, diff --git a/vllm/models/deepseek_v4/amd/dspark.py b/vllm/models/deepseek_v4/amd/dspark.py index 28056cb550db..edca59032aa3 100644 --- a/vllm/models/deepseek_v4/amd/dspark.py +++ b/vllm/models/deepseek_v4/amd/dspark.py @@ -9,8 +9,8 @@ attention + MHC CustomOp path) instead of the nvidia one; * route the MHC head through the ``HCHeadOp`` CustomOp dispatcher (aiter / tilelang / triton / torch) instead of calling the tilelang kernels directly, - and gate the trailing ``mhc_post`` on ``use_fused_mhc`` (False on the aiter - path, where the decoder layer already applies hc_post in-layer); + and gate the trailing ``mhc_post`` on ``use_fused_mhc`` (True when AITER + or TileLang fused MHC is available; False only on the torch fallback); * drop the mega-MoE weight path (``make_deepseek_v4_expert_params_mapping`` / ``use_mega_moe`` / ``finalize_mega_moe_weights`` do not exist in amd/model.py). diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index 145640860572..f28df04d4b8f 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -305,6 +305,10 @@ def forward( return final_hidden_states.view(org_shape) +# Hidden sizes supported by AITER mhc_pre_big_fuse_rmsnorm. +_AITER_MHC_FUSED_RMSNORM_SIZES = frozenset({1280, 2560, 4096, 7168}) + + class DeepseekV4DecoderLayer(nn.Module): def __init__( self, @@ -384,9 +388,18 @@ def __init__( self.mhc_pre = MHCPreOp() self.mhc_post = MHCPostOp() self.mhc_fused_post_pre = MHCFusedPostPreOp() - self.use_fused_mhc = HAS_TILELANG_MHC and not ( - HAS_AITER_MHC and self.hidden_size % 256 == 0 + # AITER mhc kernels (pre/post/fused) require hc_mult == 4. + use_aiter_mhc = ( + HAS_AITER_MHC and self.hidden_size % 256 == 0 and self.hc_mult == 4 ) + # Prefer AITER fused post+pre when eligible; otherwise TileLang. + self.use_fused_mhc = use_aiter_mhc or HAS_TILELANG_MHC + # Fold attn/ffn RMSNorm into MHC only when the active backend's + # fused-rmsnorm path supports this hidden size. + if use_aiter_mhc: + self.fuse_mhc_rmsnorm = self.hidden_size in _AITER_MHC_FUSED_RMSNORM_SIZES + else: + self.fuse_mhc_rmsnorm = HAS_TILELANG_MHC and self.use_fused_mhc def hc_pre( self, @@ -394,7 +407,13 @@ def hc_pre( hc_fn: torch.Tensor, hc_scale: torch.Tensor, hc_base: torch.Tensor, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 0.0, ): + """Reduce HC residual streams to the next sub-layer input. + + When ``norm_weight`` is set, RMSNorm is fused into the pre kernel. + """ post_mix, res_mix, layer_input = self.mhc_pre( residual=x, fn=hc_fn, @@ -405,6 +424,8 @@ def hc_pre( hc_sinkhorn_eps=self.hc_eps, hc_post_mult_value=self.hc_post_alpha, sinkhorn_repeat=self.hc_sinkhorn_iters, + norm_weight=norm_weight, + norm_eps=norm_eps, ) return layer_input, post_mix, res_mix @@ -426,11 +447,20 @@ def _forward_fused_post_pre( res_mix: torch.Tensor | None = None, residual: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + attn_norm_weight = self.attn_norm.weight if self.fuse_mhc_rmsnorm else None + attn_norm_eps = ( + self.attn_norm.variance_epsilon if self.fuse_mhc_rmsnorm else 0.0 + ) if residual is None: # Run standalone hc_pre on first layer residual = x x, post_mix, res_mix = self.hc_pre( - x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base + x, + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, + norm_weight=attn_norm_weight, + norm_eps=attn_norm_eps, ) else: residual, post_mix, res_mix, x = self.mhc_fused_post_pre( @@ -446,11 +476,16 @@ def _forward_fused_post_pre( self.hc_eps, self.hc_post_alpha, self.hc_sinkhorn_iters, + norm_weight=attn_norm_weight, + norm_eps=attn_norm_eps, ) - x = self.attn_norm(x) + if not self.fuse_mhc_rmsnorm: + x = self.attn_norm(x) x = self.attn(positions, x, None) + ffn_norm_weight = self.ffn_norm.weight if self.fuse_mhc_rmsnorm else None + ffn_norm_eps = self.ffn_norm.variance_epsilon if self.fuse_mhc_rmsnorm else 0.0 residual, post_mix, res_mix, x = self.mhc_fused_post_pre( x, residual, @@ -464,8 +499,11 @@ def _forward_fused_post_pre( self.hc_eps, self.hc_post_alpha, self.hc_sinkhorn_iters, + norm_weight=ffn_norm_weight, + norm_eps=ffn_norm_eps, ) - x = self.ffn_norm(x) + if not self.fuse_mhc_rmsnorm: + x = self.ffn_norm(x) x = self.ffn(x, input_ids) return x, residual, post_mix, res_mix @@ -678,7 +716,7 @@ def forward( residual, ) if (idx + 1) in self.aux_hidden_state_layers: - # On the unfused (aiter) path the layer already applied hc_post, + # On the unfused path the layer already applied hc_post, # so hidden_states is the reconstructed stream; on the fused # path reconstruct it via hc_post before averaging. if layer.use_fused_mhc: From 754e1c3de4e35c44da76beddf33ca71081d15d3c Mon Sep 17 00:00:00 2001 From: Qiming Zhang Date: Wed, 19 Aug 2026 23:56:09 -0700 Subject: [PATCH 186/839] [CI][XPU] Skip test_fused_shared_expert.py on XPU (#53035) Signed-off-by: mayuyuace --- tests/model_executor/layers/test_fused_shared_expert.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/model_executor/layers/test_fused_shared_expert.py b/tests/model_executor/layers/test_fused_shared_expert.py index c667f326f6da..3ec77d36ffa7 100644 --- a/tests/model_executor/layers/test_fused_shared_expert.py +++ b/tests/model_executor/layers/test_fused_shared_expert.py @@ -22,9 +22,15 @@ from vllm.model_executor.models.utils import PPMissingLayer from vllm.models.deepseek_v4 import quant_config as deepseek_v4_quant_config from vllm.models.minimax_m3.amd import model as minimax_m3_model +from vllm.platforms import current_platform from vllm.transformers_utils.configs.minimax_m3 import MiniMaxM3TextConfig from vllm.transformers_utils.configs.qwen3_5_moe import Qwen3_5MoeTextConfig +pytestmark = pytest.mark.skipif( + current_platform.is_xpu(), + reason="ROCm-specific aiter ops are not supported on XPU", +) + _QUARK_FSE_CONFIG: dict[str, Any] = { "global_quant_config": { "input_tensors": { From a1c5b1fd9f6ef06a4fa236b7d48350115e5688b9 Mon Sep 17 00:00:00 2001 From: Zhixia Liu Date: Thu, 20 Aug 2026 15:12:43 +0800 Subject: [PATCH 187/839] [Core][V1] Support trace_decode_token_ids for deterministic decode replay (#46701) Signed-off-by: ianzxliu Signed-off-by: aoshen02 Signed-off-by: Zhixia Liu Co-authored-by: ianzxliu Co-authored-by: Claude Co-authored-by: aoshen02 Co-authored-by: Codex Co-authored-by: Nick Hill --- docs/serving/online_serving/trace_replay.md | 60 ++++++ examples/generate/trace_replay_offline.py | 195 ++++++++++++++++++ .../test_input_processor_trace_replay.py | 110 ++++++++++ tests/v1/sample/test_trace_replay_params.py | 111 ++++++++++ tests/v1/worker/test_gpu_trace_replay.py | 195 ++++++++++++++++++ vllm/config/model.py | 7 + vllm/config/vllm.py | 8 + vllm/engine/arg_utils.py | 5 + vllm/sampling_params.py | 80 ++++++- vllm/v1/engine/input_processor.py | 41 +++- vllm/v1/worker/gpu/model_runner.py | 1 + vllm/v1/worker/gpu/sample/sampler.py | 14 ++ vllm/v1/worker/gpu/sample/trace_replay.py | 113 ++++++++++ 13 files changed, 931 insertions(+), 9 deletions(-) create mode 100644 docs/serving/online_serving/trace_replay.md create mode 100644 examples/generate/trace_replay_offline.py create mode 100644 tests/v1/engine/test_input_processor_trace_replay.py create mode 100644 tests/v1/sample/test_trace_replay_params.py create mode 100644 tests/v1/worker/test_gpu_trace_replay.py create mode 100644 vllm/v1/worker/gpu/sample/trace_replay.py diff --git a/docs/serving/online_serving/trace_replay.md b/docs/serving/online_serving/trace_replay.md new file mode 100644 index 000000000000..86ed13644c7a --- /dev/null +++ b/docs/serving/online_serving/trace_replay.md @@ -0,0 +1,60 @@ +# Trace Replay + +Trace replay forces the engine to emit a predetermined token sequence during decoding while computing real logprobs from the model's unmodified logit distribution. The primary use case is comparing logprob distributions across different configurations: + +- **Inference config diff**: compare logprobs between different quantization schemes, tensor parallelism layouts, or attention backends for the same model. +- **Train vs inference diff**: replay a training-time token sequence through the inference engine to detect logprob divergence caused by numerical differences between training and serving frameworks. + +## Requirements + +Trace replay reserves a per-request trace buffer, so it is off by default. Enable it with `--enable-trace-replay` (or `LLM(..., enable_trace_replay=True)`). It is only supported by model runner V2; requests are rejected when either condition is not met. + +## Usage + +```python +from vllm import LLM, SamplingParams + +llm = LLM(model="Qwen/Qwen3-0.6B", enable_trace_replay=True) + +# Token sequence captured from a previous run or training log +trace_tokens = [15, 284, 1026, 374] + +params = SamplingParams( + trace_decode_token_ids=trace_tokens, + logprobs=5, +) +outputs = llm.generate(["Once upon a time"], sampling_params=params) + +for token, logprob in zip( + outputs[0].outputs[0].token_ids, + outputs[0].outputs[0].logprobs, +): + print(f"token={token} logprob={logprob[token].logprob:.4f}") +``` + +The output tokens will always be `[15, 284, 1026, 374]`. The logprobs reflect the model's true probability for each forced token under the current inference configuration. + +To compare configurations, run the same trace against two engine setups and diff the per-token logprobs. + +## Behavior + +When `trace_decode_token_ids` is set: + +- `max_tokens` is automatically set to the trace length. +- The trace is truncated if it does not fit within `max_model_len` given the prompt length. +- All stop conditions (EOS, stop strings, stop token IDs) are disabled. +- Generation produces exactly the trace tokens, then stops. + +## Limitations + +`trace_decode_token_ids` is incompatible with the following features (raises `ValueError`): + +| Feature | Reason | +| --- | --- | +| `n > 1` | Trace replay produces a single deterministic sequence | +| `prompt_logprobs` | Expanded logit layout conflicts with trace kernel indexing | +| Speculative decoding | Multi-token speculation conflicts with single-token trace stepping | +| Structured outputs | Grammar constraints conflict with forced token injection | +| `repetition_detection` | Would terminate the request on repeated patterns in the trace | +| `thinking_token_budget` | Logit masking corrupts logprobs for trace tokens | +| `bad_words` | Logit masking corrupts logprobs for trace tokens | diff --git a/examples/generate/trace_replay_offline.py b/examples/generate/trace_replay_offline.py new file mode 100644 index 000000000000..2baa03307d19 --- /dev/null +++ b/examples/generate/trace_replay_offline.py @@ -0,0 +1,195 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Trace-replay with vLLM offline inference. + +Trace-replay lets you supply a known sequence of decode token IDs alongside +the prompt. Instead of sampling from the model distribution, the engine +injects each decode token deterministically, step by step. All other +outputs — logprobs, token ranks, text decoding — are computed faithfully +from the real logit distribution for that token. + +How it works: + 1. Start the engine with ``--enable-trace-replay`` (or + ``LLM(..., enable_trace_replay=True)``). + 2. Set ``SamplingParams.trace_decode_token_ids`` to the list of decode + token IDs you want to force. + 3. The engine will output exactly those tokens and stop. ``max_tokens`` + is overwritten with the trace length, and the trace is truncated if it + does not fit within ``max_model_len``. EOS tokens inside the trace + sequence do **not** halt generation early. + +Requires model runner V2. Requests are rejected with ``ValueError`` when +combined with any of: + * n > 1 * Speculative decoding + * prompt_logprobs * Structured outputs + * repetition_detection * thinking_token_budget + * bad_words + +Typical use-cases: + * Reproduce exact outputs from a previous run for benchmarking. + * Compute logprobs for an already-known output (e.g. reference answers). + * Dataset annotation: given (prompt, response) pairs, obtain per-token + logprob scores without altering the response. + +Usage: + python examples/generate/trace_replay_offline.py + python examples/generate/trace_replay_offline.py --model facebook/opt-125m +""" + +import argparse + +from vllm import LLM, SamplingParams + +DEFAULT_PROMPT = "Hello, my name is" + + +def build_llm(args: argparse.Namespace) -> LLM: + """Construct an LLM from common CLI args.""" + llm_kwargs: dict = { + "model": args.model, + "trust_remote_code": args.trust_remote_code, + "tensor_parallel_size": args.tensor_parallel_size, + "enforce_eager": args.enforce_eager, + "gpu_memory_utilization": args.gpu_memory_utilization, + "max_num_seqs": args.max_num_seqs, + "enable_trace_replay": True, + } + if args.max_model_len is not None: + llm_kwargs["max_model_len"] = args.max_model_len + return LLM(**llm_kwargs) + + +def run_normal_generation(llm: LLM, prompt: str, max_tokens: int) -> list[int]: + """Run a standard greedy generation and return the output token IDs.""" + sampling_params = SamplingParams( + temperature=0.0, + max_tokens=max_tokens, + # Request logprobs for the greedy token at each step so we can + # compare them against the trace-replay logprobs below. + logprobs=1, + ) + outputs = llm.generate([prompt], sampling_params=sampling_params) + result = outputs[0].outputs[0] + output_token_ids = list(result.token_ids) + print("[Normal generation]") + print(f" Prompt : {prompt!r}") + print(f" Output token IDs : {output_token_ids}") + print(f" Output text : {result.text!r}") + return output_token_ids + + +def run_trace_replay(llm: LLM, prompt: str, decode_token_ids: list[int]) -> None: + """Replay a known decode sequence and print per-token logprobs.""" + sampling_params = SamplingParams( + # Provide the decode tokens to replay. + trace_decode_token_ids=decode_token_ids, + # Request top-5 logprobs so we can inspect the distribution. + logprobs=5, + ) + outputs = llm.generate([prompt], sampling_params=sampling_params) + result = outputs[0].outputs[0] + replayed_ids = list(result.token_ids) + + print("\n[Trace-replay]") + print(f" Requested decode token IDs : {decode_token_ids}") + print(f" Replayed output token IDs : {replayed_ids}") + print(f" Replayed output text : {result.text!r}") + + # Verify the replayed tokens match exactly. + assert replayed_ids == decode_token_ids, ( + f"Mismatch!\n expected: {decode_token_ids}\n got: {replayed_ids}" + ) + print(" Replayed tokens match the requested trace exactly.") + + # Show per-token logprobs (computed from the real distribution). + if result.logprobs: + print("\n Per-token logprobs (trace token):") + for step, (token_id, logprob_dict) in enumerate( + zip(replayed_ids, result.logprobs) + ): + sampled_lp = logprob_dict.get(token_id) + lp_value = f"{sampled_lp.logprob:.4f}" if sampled_lp is not None else "n/a" + rank = sampled_lp.rank if sampled_lp is not None else "n/a" + print( + f" step {step:2d}: token_id={token_id:6d} " + f"logprob={lp_value} rank={rank}" + ) + + +def run_demo(args: argparse.Namespace) -> None: + print(f"Loading model: {args.model}\n") + llm = build_llm(args) + prompt = args.prompt + + print("=" * 60) + print("Step 1 — Normal greedy generation (captures decode tokens)") + print("=" * 60) + decode_token_ids = run_normal_generation(llm, prompt, max_tokens=8) + + print("\n" + "=" * 60) + print("Step 2 — Trace-replay with the captured tokens") + print("=" * 60) + run_trace_replay(llm, prompt, decode_token_ids) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Trace-replay with vLLM offline inference" + ) + parser.add_argument( + "--prompt", + type=str, + default=DEFAULT_PROMPT, + help="Text prompt to use for generation (default: %(default)r)", + ) + parser.add_argument( + "--model", + type=str, + default="facebook/opt-125m", + help="Name or path of the HuggingFace model to use", + ) + parser.add_argument( + "--trust-remote-code", + action="store_true", + help="Trust remote code from HuggingFace", + ) + parser.add_argument( + "--tensor-parallel-size", + type=int, + default=1, + help="Number of tensor parallel replicas", + ) + parser.add_argument( + "--enforce-eager", + action="store_true", + help="Always use eager-mode PyTorch (disable CUDA graph)", + ) + parser.add_argument( + "--gpu-memory-utilization", + type=float, + default=0.9, + help="Fraction of GPU memory to use", + ) + parser.add_argument( + "--max-model-len", + type=int, + default=None, + help="Model context length", + ) + parser.add_argument( + "--max-num-seqs", + type=int, + default=256, + help="Maximum number of sequences per iteration", + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + run_demo(args) + + +if __name__ == "__main__": + main() diff --git a/tests/v1/engine/test_input_processor_trace_replay.py b/tests/v1/engine/test_input_processor_trace_replay.py new file mode 100644 index 000000000000..894293bc3db1 --- /dev/null +++ b/tests/v1/engine/test_input_processor_trace_replay.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import time +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +from vllm import SamplingParams +from vllm.config import VllmConfig +from vllm.exceptions import VLLMValidationError +from vllm.v1.core.sched.utils import check_stop +from vllm.v1.engine.input_processor import InputProcessor +from vllm.v1.request import Request, RequestStatus + +pytestmark = pytest.mark.skip_global_cleanup + + +def _make_request(params: SamplingParams) -> Request: + return Request( + request_id="req-0", + client_index=0, + prompt_token_ids=[1, 2, 3], + sampling_params=params, + pooling_params=None, + arrival_time=time.time(), + ) + + +def _normalize(params: SamplingParams, prompt_len: int, max_model_len: int) -> None: + processor = SimpleNamespace( + model_config=SimpleNamespace(max_model_len=max_model_len) + ) + InputProcessor._normalize_trace_replay_params(processor, params, prompt_len) + + +def test_normalize_trace_replay_params(): + params = SamplingParams( + max_tokens=16, + min_tokens=10, + stop_token_ids=[20], + trace_decode_token_ids=[10, 20, 30], + ) + params.update_from_generation_config({}, eos_token_id=20) + + _normalize(params, prompt_len=3, max_model_len=128) + + assert params.max_tokens == 3 + assert params.min_tokens == 0 + assert params.ignore_eos is True + assert params.eos_token_id is None + assert params.stop_token_ids == [] + assert params.all_stop_token_ids == set() + + request = _make_request(params) + assert request.max_tokens == 3 + request.append_output_token_ids([10, 20, 30]) + assert check_stop(request, max_model_len=128) + assert request.status == RequestStatus.FINISHED_LENGTH_CAPPED + + +def test_trace_longer_than_remaining_context_is_truncated(): + """The trace is staged into a max_model_len-wide row, so it must be cut. + + An explicitly set max_tokens is never clamped to max_model_len, so it cannot + bound the staged write. + """ + params = SamplingParams(max_tokens=100, trace_decode_token_ids=list(range(20))) + + _normalize(params, prompt_len=6, max_model_len=10) + + assert params.trace_decode_token_ids == [0, 1, 2, 3] + assert params.max_tokens == 4 + + +def _validate(enable_trace_replay: bool) -> None: + """Run _validate_params' trace gating with the rest of verify() stubbed.""" + processor = SimpleNamespace( + model_config=SimpleNamespace( + return_sampling_mask=False, + enable_trace_replay=enable_trace_replay, + ), + vllm_config=SimpleNamespace(reasoning_config=None), + speculative_config=None, + structured_outputs_config=None, + tokenizer=None, + ) + params = SamplingParams(trace_decode_token_ids=[1, 2, 3]) + with patch.object(SamplingParams, "verify"): + InputProcessor._validate_params(processor, params, ("generate",)) + + +def test_trace_request_rejected_when_feature_disabled(): + with pytest.raises(VLLMValidationError, match="--enable-trace-replay"): + _validate(enable_trace_replay=False) + + +def test_trace_request_accepted_when_feature_enabled(): + _validate(enable_trace_replay=True) + + +def test_trace_replay_requires_v2_model_runner(): + config = SimpleNamespace( + model_config=SimpleNamespace(enable_trace_replay=True), + use_v2_model_runner=False, + ) + + with pytest.raises(ValueError, match="trace replay requires Model Runner V2"): + VllmConfig._verify_trace_replay_config(config) diff --git a/tests/v1/sample/test_trace_replay_params.py b/tests/v1/sample/test_trace_replay_params.py new file mode 100644 index 000000000000..6e0d72647017 --- /dev/null +++ b/tests/v1/sample/test_trace_replay_params.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from vllm import SamplingParams +from vllm.exceptions import VLLMValidationError +from vllm.sampling_params import StructuredOutputsParams + +pytestmark = pytest.mark.skip_global_cleanup + +# --------------------------------------------------------------------------- +# SamplingParams field +# --------------------------------------------------------------------------- + + +def test_sampling_params_trace_field_defaults_to_none(): + params = SamplingParams(max_tokens=10) + assert params.trace_decode_token_ids is None + + +def test_sampling_params_trace_field_accepts_list(): + ids = [100, 200, 300] + params = SamplingParams(trace_decode_token_ids=ids) + assert params.trace_decode_token_ids == ids + + +def test_sampling_params_trace_field_preserved_by_clone(): + ids = [1, 2, 3] + params = SamplingParams(trace_decode_token_ids=ids) + cloned = params.clone() + assert cloned.trace_decode_token_ids == ids + assert cloned.trace_decode_token_ids is not params.trace_decode_token_ids + + +def test_sampling_params_trace_field_rejects_empty_list(): + params = SamplingParams(trace_decode_token_ids=[]) + with pytest.raises(ValueError, match="non-empty"): + params._validate_trace_replay( + _make_model_config(vocab_size=100), speculative_config=None + ) + + +def test_sampling_params_trace_field_requires_single_output(): + params = SamplingParams(n=2, trace_decode_token_ids=[1]) + with pytest.raises(ValueError, match="requires n=1"): + params._validate_trace_replay( + _make_model_config(vocab_size=100), speculative_config=None + ) + + +@pytest.mark.parametrize("invalid_ids", [[-1, 5], [1, "2"]]) +def test_sampling_params_trace_field_rejects_invalid_token_ids(invalid_ids): + params = SamplingParams(trace_decode_token_ids=invalid_ids) + with pytest.raises(ValueError, match="non-negative integers"): + params._validate_trace_replay( + _make_model_config(vocab_size=100), speculative_config=None + ) + + +def _make_model_config(vocab_size: int): + from unittest.mock import Mock + + model_config = Mock() + model_config.get_vocab_size = lambda: vocab_size + return model_config + + +def test_validate_trace_replay_accepts_in_vocab(): + params = SamplingParams(trace_decode_token_ids=[0, 50, 99]) + # Should not raise. + params._validate_trace_replay( + _make_model_config(vocab_size=100), speculative_config=None + ) + + +def test_validate_trace_replay_rejects_out_of_vocab(): + # The non-negative check passes at construction, but the token id exceeds + # the vocabulary; verify() must reject it before it reaches the sampler. + params = SamplingParams(trace_decode_token_ids=[0, 100]) + with pytest.raises(VLLMValidationError, match="out-of-vocab"): + params._validate_trace_replay( + _make_model_config(vocab_size=100), speculative_config=None + ) + + +def test_validate_trace_replay_noop_when_unset(): + params = SamplingParams(max_tokens=4) + # Should not raise when the field is unset. + params._validate_trace_replay( + _make_model_config(vocab_size=100), speculative_config=None + ) + + +def test_trace_decode_token_ids_rejects_speculative_decoding(): + params = SamplingParams(trace_decode_token_ids=[1]) + with pytest.raises(ValueError, match="not supported with speculative decoding"): + params._validate_trace_replay( + _make_model_config(vocab_size=100), speculative_config=object() + ) + + +def test_trace_decode_token_ids_rejects_structured_outputs(): + params = SamplingParams( + trace_decode_token_ids=[1], + structured_outputs=StructuredOutputsParams(json_object=True), + ) + with pytest.raises(ValueError, match="not supported with structured outputs"): + params._validate_trace_replay( + _make_model_config(vocab_size=100), speculative_config=None + ) diff --git a/tests/v1/worker/test_gpu_trace_replay.py b/tests/v1/worker/test_gpu_trace_replay.py new file mode 100644 index 000000000000..9d418101fce4 --- /dev/null +++ b/tests/v1/worker/test_gpu_trace_replay.py @@ -0,0 +1,195 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for Model Runner V2 inference trace-replay. + +Trace-replay forces the sampler to emit a predetermined sequence of decode +token IDs. The replay step for a request is derived from GPU state as +``total_len - prompt_len`` (output tokens produced so far), and the trace token +overwrites the sampled token in place before logprobs are computed. +""" + +from types import SimpleNamespace +from typing import cast + +import pytest +import torch + +pytest.importorskip("triton") +if not torch.cuda.is_available(): + pytest.skip("CUDA required for trace-replay tests", allow_module_level=True) + +from vllm.sampling_params import SamplingParams +from vllm.v1.worker.gpu.sample.trace_replay import ( + TraceReplayState, + apply_trace_tokens, +) +from vllm.v1.worker.gpu.states import RequestState + +DEVICE = "cuda" +TEST_MAX_MODEL_LEN = 2048 + + +def _i32(x) -> torch.Tensor: + return torch.tensor(x, dtype=torch.int32, device=DEVICE) + + +def _i64(x) -> torch.Tensor: + return torch.tensor(x, dtype=torch.int64, device=DEVICE) + + +def _trace_state(max_num_reqs: int) -> TraceReplayState: + """Build a state whose req_states exposes the buffers apply_trace reads. + + total_len/prompt_len are written by the caller to position each request at + the desired replay step. + """ + req_states = cast( + RequestState, + SimpleNamespace( + max_num_reqs=max_num_reqs, + max_model_len=TEST_MAX_MODEL_LEN, + device=torch.device(DEVICE), + total_len=SimpleNamespace( + gpu=torch.zeros(max_num_reqs, dtype=torch.int32, device=DEVICE) + ), + prompt_len=SimpleNamespace( + gpu=torch.zeros(max_num_reqs, dtype=torch.int32, device=DEVICE) + ), + ), + ) + return TraceReplayState(req_states) + + +def _set_lens(state: TraceReplayState, total_len, prompt_len) -> None: + """Position each request at a replay step via total_len - prompt_len.""" + state.req_states.total_len.gpu[: len(total_len)] = _i32(total_len) + state.req_states.prompt_len.gpu[: len(prompt_len)] = _i32(prompt_len) + + +# ------------------------------ Kernel ------------------------------------ + + +def test_replay_overwrites_sampled_at_each_step(): + """The trace token for the current step replaces the sampled token.""" + trace = [[100, 101, 102], [200, 201, 202]] + trace_len = _i32([3, 3]) + trace_token_ids = torch.zeros( + 2, TEST_MAX_MODEL_LEN, dtype=torch.int32, device=DEVICE + ) + for i, t in enumerate(trace): + trace_token_ids[i, : len(t)] = _i32(t) + prompt_len = _i32([5, 8]) + idx_mapping = _i32([0, 1]) + + for step in range(3): + sampled = _i64([-7, -7]) # sentinel that must be overwritten + total_len = _i32([5 + step, 8 + step]) + apply_trace_tokens( + sampled, idx_mapping, trace_token_ids, trace_len, total_len, prompt_len + ) + assert sampled.tolist() == [trace[0][step], trace[1][step]] + + +def test_past_end_of_trace_leaves_sampled_untouched(): + """Once step >= trace_len, the sampler's own token is kept.""" + trace_token_ids = torch.zeros( + 1, TEST_MAX_MODEL_LEN, dtype=torch.int32, device=DEVICE + ) + trace_token_ids[0, :2] = _i32([100, 101]) + trace_len = _i32([2]) + prompt_len = _i32([4]) + idx_mapping = _i32([0]) + + sampled = _i64([999]) + total_len = _i32([4 + 2]) # step == 2 == trace_len -> out of range + apply_trace_tokens( + sampled, idx_mapping, trace_token_ids, trace_len, total_len, prompt_len + ) + assert sampled.tolist() == [999] + + +def test_non_trace_request_untouched(): + """trace_len == 0 means the request never uses replay.""" + trace_token_ids = torch.zeros( + 1, TEST_MAX_MODEL_LEN, dtype=torch.int32, device=DEVICE + ) + trace_len = _i32([0]) + prompt_len = _i32([3]) + idx_mapping = _i32([0]) + + sampled = _i64([42]) + total_len = _i32([3]) + apply_trace_tokens( + sampled, idx_mapping, trace_token_ids, trace_len, total_len, prompt_len + ) + assert sampled.tolist() == [42] + + +def test_idx_mapping_indirection_and_negative_skip(): + """batch_idx -> req_state_idx indirection, and negative entries are skipped.""" + # req_state 0: no trace; req_state 1: trace [500, 501]. + trace_token_ids = torch.zeros( + 2, TEST_MAX_MODEL_LEN, dtype=torch.int32, device=DEVICE + ) + trace_token_ids[1, :2] = _i32([500, 501]) + trace_len = _i32([0, 2]) + prompt_len = _i32([10, 6]) + + # batch: [maps to state 1, masked (-1), maps to state 0]. + idx_mapping = _i32([1, -1, 0]) + sampled = _i64([1, 2, 3]) + # Indexed by req_state_idx: state 0 at total_len=10, state 1 at step 1. + total_len = _i32([10, 6 + 1]) + apply_trace_tokens( + sampled, idx_mapping, trace_token_ids, trace_len, total_len, prompt_len + ) + # batch 0 -> state 1 step 1 -> 501; batch 1 masked; batch 2 -> state 0 no trace. + assert sampled.tolist() == [501, 2, 3] + + +# --------------------------- TraceReplayState ----------------------------- + + +def test_state_end_to_end(): + """add_request -> apply_staged_writes -> apply_trace overwrites correctly.""" + state = _trace_state(4) + state.add_request(0, SamplingParams(trace_decode_token_ids=[11, 22, 33])) + state.add_request(1, SamplingParams()) # no trace + state.apply_staged_writes() + + idx_mapping = _i32([0, 1]) + sampled = _i64([-1, -1]) + _set_lens(state, [7 + 1, 7, 0, 0], [7, 7, 0, 0]) # req 0 at step 1 + state.apply_trace(sampled, idx_mapping) + assert sampled.tolist() == [22, -1] + + +def test_state_leaves_non_trace_batch_unchanged(): + """Requests with trace_len == 0 remain unchanged after a trace was seen.""" + state = _trace_state(4) + state.add_request(0, SamplingParams(trace_decode_token_ids=[11, 22])) + state.add_request(1, SamplingParams()) + state.apply_staged_writes() + + # Only the non-trace request (state 1) is in this batch. + idx_mapping = _i32([1]) + sampled = _i64([555]) + _set_lens(state, [0, 7, 0, 0], [7, 7, 0, 0]) + state.apply_trace(sampled, idx_mapping) + assert sampled.tolist() == [555] + + +def test_slot_reuse_clears_trace(): + """Reusing a slot for a non-trace request must not replay stale tokens.""" + state = _trace_state(2) + state.add_request(0, SamplingParams(trace_decode_token_ids=[11, 22])) + state.apply_staged_writes() + # Slot 0 reused by a request without a trace. + state.add_request(0, SamplingParams()) + state.apply_staged_writes() + + idx_mapping = _i32([0]) + sampled = _i64([888]) + _set_lens(state, [3, 0], [3, 0]) + state.apply_trace(sampled, idx_mapping) + assert sampled.tolist() == [888] diff --git a/vllm/config/model.py b/vllm/config/model.py index 16ff966c2d51..56e533ea5305 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -262,6 +262,12 @@ class ModelConfig: equivalent exponential-race sampling. FP64 preserves lower-tail sampling events that fp32 uniform/exponential draws can truncate, at the cost of significantly lower throughput on most GPUs.""" + enable_trace_replay: bool = False + """Whether to allow requests to set + `SamplingParams.trace_decode_token_ids`, which forces decoding to follow a + predetermined token sequence while still computing real logprobs. Reserved + for debugging and RL workflows: enabling it reserves a per-request trace + buffer, so it is off by default.""" disable_sliding_window: bool = False """Whether to disable sliding window. If True, we will disable the sliding window functionality of the model, capping to sliding window size. If the @@ -424,6 +430,7 @@ def compute_hash(self) -> str: "return_sampling_mask", "logprobs_mode", "use_fp64_gumbel", + "enable_trace_replay", "disable_cascade_attn", "skip_tokenizer_init", "served_model_name", diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index ade8a666cf27..a271a654134a 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1091,6 +1091,13 @@ def _verify_sampling_replay_config(self) -> None: "are normalized over the same nucleus as the sampling mask" ) + def _verify_trace_replay_config(self) -> None: + model_config = self.model_config + if model_config is None or not model_config.enable_trace_replay: + return + if not self.use_v2_model_runner: + raise ValueError("trace replay requires Model Runner V2") + def __post_init__(self): """Verify configs are valid & consistent with each other.""" @@ -1141,6 +1148,7 @@ def __post_init__(self): ) self._verify_sampling_replay_config() + self._verify_trace_replay_config() if self.lora_config is not None: self.lora_config.verify_with_model_config(self.model_config) diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 1fb1388f100c..035595d5a8df 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -543,6 +543,7 @@ class EngineArgs: max_logprobs: int = ModelConfig.max_logprobs logprobs_mode: LogprobsMode = ModelConfig.logprobs_mode use_fp64_gumbel: bool = ModelConfig.use_fp64_gumbel + enable_trace_replay: bool = ModelConfig.enable_trace_replay disable_log_stats: bool = False aggregate_engine_logging: bool = False revision: str | None = ModelConfig.revision @@ -884,6 +885,9 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: model_group.add_argument("--max-logprobs", **model_kwargs["max_logprobs"]) model_group.add_argument("--logprobs-mode", **model_kwargs["logprobs_mode"]) model_group.add_argument("--use-fp64-gumbel", **model_kwargs["use_fp64_gumbel"]) + model_group.add_argument( + "--enable-trace-replay", **model_kwargs["enable_trace_replay"] + ) model_group.add_argument( "--disable-sliding-window", **model_kwargs["disable_sliding_window"] ) @@ -1760,6 +1764,7 @@ def create_model_config(self) -> ModelConfig: max_logprobs=self.max_logprobs, logprobs_mode=self.logprobs_mode, use_fp64_gumbel=self.use_fp64_gumbel, + enable_trace_replay=self.enable_trace_replay, disable_sliding_window=self.disable_sliding_window, disable_cascade_attn=self.disable_cascade_attn, skip_tokenizer_init=self.skip_tokenizer_init, diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index 2f0af9e1b43b..f0e325c255a6 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -346,14 +346,6 @@ class SamplingParams( """Arbitrary additional args, that can be used by custom sampling implementations, plugins, etc. Not used by any in-tree sampling implementations.""" - routed_experts_prompt_start: int = 0 - """When enable_return_routed_experts is active, skip the first - routed_experts_prompt_start prompt tokens from the returned routing - data. In multi-turn agent scenarios, set this to the length of the - already-returned prefix to avoid duplicating routing for prompt tokens - covered by earlier turns. Default 0 returns routing for all prompt - tokens.""" - # Fields used for bad words bad_words: list[str] | None = None """Words that are not allowed to be generated. More precisely, only the @@ -373,6 +365,19 @@ class SamplingParams( '\\emoji \\emoji \\emoji ...'). This feature can detect such behavior and terminate early, saving time and tokens.""" + # Debugging / RL-specific parameters. Not intended for production serving. + routed_experts_prompt_start: int = 0 + """When enable_return_routed_experts is active, skip the first + routed_experts_prompt_start prompt tokens from the returned routing + data. In multi-turn agent scenarios, set this to the length of the + already-returned prefix to avoid duplicating routing for prompt tokens + covered by earlier turns. Default 0 returns routing for all prompt + tokens.""" + trace_decode_token_ids: list[int] | None = None + """If provided, forces the engine to emit this predetermined sequence of + token IDs during decoding instead of sampling randomly. Real logprobs are + still computed. Conflict checking is performed at the engine level.""" + @staticmethod def from_optional( n: int | None = 1, @@ -407,6 +412,8 @@ def from_optional( repetition_detection: RepetitionDetectionParams | None = None, logprob_token_ids: list[int] | None = None, routed_experts_prompt_start: int = 0, + # Debugging / RL-specific parameters. + trace_decode_token_ids: list[int] | None = None, ) -> "SamplingParams": if logit_bias is not None: # Fast path uses a dict comprehension; on failure we iterate once @@ -470,6 +477,7 @@ def from_optional( skip_clone=skip_clone, repetition_detection=repetition_detection, routed_experts_prompt_start=routed_experts_prompt_start, + trace_decode_token_ids=trace_decode_token_ids, ) def __post_init__(self) -> None: @@ -780,6 +788,7 @@ def verify( ) -> None: self._validate_logprobs(model_config) self._validate_logit_bias(model_config) + self._validate_trace_replay(model_config, speculative_config) self._validate_logits_processors(model_config) self._validate_allowed_token_ids(tokenizer) self._validate_spec_decode(speculative_config) @@ -870,6 +879,61 @@ def _validate_logit_bias(self, model_config: ModelConfig) -> None: value=invalid_token_ids, ) + def _validate_trace_replay( + self, + model_config: ModelConfig, + speculative_config: SpeculativeConfig | None, + ) -> None: + """Validate trace replay request compatibility.""" + if self.trace_decode_token_ids is None: + return + + if len(self.trace_decode_token_ids) == 0: + raise ValueError("trace_decode_token_ids must be a non-empty list.") + if self.n != 1: + raise ValueError("trace_decode_token_ids requires n=1.") + if not all(isinstance(t, int) and t >= 0 for t in self.trace_decode_token_ids): + raise ValueError( + "trace_decode_token_ids must contain non-negative integers." + ) + + if self.prompt_logprobs is not None: + raise ValueError( + "trace_decode_token_ids is not supported with prompt_logprobs." + ) + if speculative_config is not None: + raise ValueError( + "trace_decode_token_ids is not supported with speculative decoding." + ) + if self.structured_outputs is not None: + raise ValueError( + "trace_decode_token_ids is not supported with structured outputs." + ) + if self.repetition_detection is not None: + raise ValueError( + "trace_decode_token_ids is not supported with repetition_detection." + ) + if self.thinking_token_budget is not None: + raise ValueError( + "trace_decode_token_ids is not supported with thinking_token_budget." + ) + if self.bad_words: + raise ValueError("trace_decode_token_ids is not supported with bad_words.") + + vocab_size = model_config.get_vocab_size() + invalid_token_ids = [ + token_id + for token_id in self.trace_decode_token_ids + if token_id < 0 or token_id >= vocab_size + ] + if invalid_token_ids: + raise VLLMValidationError( + f"token_id(s) {invalid_token_ids} in trace_decode_token_ids " + f"contain out-of-vocab token ids. Vocabulary size: {vocab_size}", + parameter="trace_decode_token_ids", + value=invalid_token_ids, + ) + def _validate_logits_processors(self, model_config: ModelConfig) -> None: from vllm.v1.sample.logits_processor import ( validate_logits_processors_parameters, diff --git a/vllm/v1/engine/input_processor.py b/vllm/v1/engine/input_processor.py index ac870c931a32..744da670a04f 100644 --- a/vllm/v1/engine/input_processor.py +++ b/vllm/v1/engine/input_processor.py @@ -51,7 +51,6 @@ def __init__( self.speculative_config = vllm_config.speculative_config self.structured_outputs_config = vllm_config.structured_outputs_config self.observability_config = vllm_config.observability_config - self.use_v2_model_runner = vllm_config.use_v2_model_runner self.generation_config_fields = model_config.try_get_generation_config() @@ -128,6 +127,15 @@ def _validate_params( "not configured. Please set --reasoning-parser " "and/or --reasoning-config to use thinking_token_budget." ) + if ( + params.trace_decode_token_ids + and not self.model_config.enable_trace_replay + ): + raise VLLMValidationError( + "trace_decode_token_ids is set but trace replay is not " + "enabled. Start the engine with --enable-trace-replay " + "to use it." + ) elif isinstance(params, PoolingParams): supported_pooling_tasks = [ task for task in supported_tasks if task in POOLING_TASKS @@ -156,6 +164,30 @@ def _validate_params( f"but got {type(params).__name__}" ) + def _normalize_trace_replay_params( + self, sampling_params: SamplingParams, prompt_len: int + ) -> None: + """Apply trace replay's generation semantics to request-local params.""" + trace_token_ids = sampling_params.trace_decode_token_ids + assert trace_token_ids + assert sampling_params.max_tokens is not None + + max_trace_len = max(self.model_config.max_model_len - prompt_len, 1) + trace_token_ids = trace_token_ids[:max_trace_len] + sampling_params.trace_decode_token_ids = trace_token_ids + + # Apply this after the generation config so its EOS token cannot stop + # replay before the trace is exhausted. + sampling_params.max_tokens = min( + len(trace_token_ids), sampling_params.max_tokens + ) + sampling_params.min_tokens = 0 + sampling_params.ignore_eos = True + sampling_params._eos_token_id = None + sampling_params.stop = [] + sampling_params.stop_token_ids = [] + sampling_params._all_stop_token_ids = set() + def _validate_lora(self, lora_request: LoRARequest | None) -> None: if lora_request is None: return @@ -340,6 +372,13 @@ def process_inputs( ) if self.tokenizer is not None: sampling_params.update_from_tokenizer(self.tokenizer) + if sampling_params.trace_decode_token_ids: + self._normalize_trace_replay_params( + sampling_params, + length_from_prompt_token_ids_or_embeds( + prompt_token_ids, prompt_embeds + ), + ) else: pooling_params = params.clone() diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 3861bb52649d..66b3c8db4a1e 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -408,6 +408,7 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None: logprobs_mode=self.model_config.logprobs_mode, num_speculative_tokens=self.decode_query_len, use_fp64_gumbel=self.model_config.use_fp64_gumbel, + enable_trace_replay=self.model_config.enable_trace_replay, reasoning_config=self.vllm_config.reasoning_config, return_sampling_mask=self.model_config.return_sampling_mask, ) diff --git a/vllm/v1/worker/gpu/sample/sampler.py b/vllm/v1/worker/gpu/sample/sampler.py index 26371f014e33..9e65d9a3041a 100644 --- a/vllm/v1/worker/gpu/sample/sampler.py +++ b/vllm/v1/worker/gpu/sample/sampler.py @@ -26,6 +26,7 @@ from vllm.v1.worker.gpu.sample.penalties import PenaltiesState from vllm.v1.worker.gpu.sample.states import NO_LOGPROBS, SamplingStates from vllm.v1.worker.gpu.sample.thinking_budget import ThinkingBudgetState +from vllm.v1.worker.gpu.sample.trace_replay import TraceReplayState from vllm.v1.worker.gpu.states import RequestState @@ -39,6 +40,7 @@ def __init__( logprobs_mode: LogprobsMode = "raw_logprobs", num_speculative_tokens: int = 1, use_fp64_gumbel: bool = False, + enable_trace_replay: bool = False, reasoning_config: ReasoningConfig | None = None, return_sampling_mask: bool = False, ): @@ -53,6 +55,9 @@ def __init__( self.bad_words_state = BadWordsState(req_states) self.logprob_token_ids_state = LogprobTokenIdsState(max_num_reqs, device) self.thinking_budget_state = ThinkingBudgetState(req_states, reasoning_config) + self.trace_replay_state = ( + TraceReplayState(req_states) if enable_trace_replay else None + ) self.needs_logits_processing = np.zeros(max_num_reqs, dtype=bool) self.num_speculative_tokens = num_speculative_tokens self.return_sampling_mask = return_sampling_mask @@ -69,6 +74,8 @@ def add_request( self.bad_words_state.add_request(req_idx, sampling_params) self.logprob_token_ids_state.add_request(req_idx, sampling_params) self.thinking_budget_state.add_request(req_idx, sampling_params) + if self.trace_replay_state is not None: + self.trace_replay_state.add_request(req_idx, sampling_params) states = self.sampling_states temperature = states.temperature.np[req_idx] @@ -93,6 +100,8 @@ def apply_staged_writes(self) -> None: self.bad_words_state.apply_staged_writes() self.logprob_token_ids_state.apply_staged_writes() self.thinking_budget_state.apply_staged_writes() + if self.trace_replay_state is not None: + self.trace_replay_state.apply_staged_writes() def __call__( self, @@ -128,6 +137,11 @@ def __call__( return_logprobs=return_logprobs, ) + if self.trace_replay_state is not None: + # Overwrite sampled tokens with the replay trace up-front so that + # computed logprobs reflect the real distribution of the forced token. + self.trace_replay_state.apply_trace(sampled, idx_mapping) + if return_logprobs: if self.logprobs_mode in PROCESSED_LOGPROBS_MODES: logits = processed_logits diff --git a/vllm/v1/worker/gpu/sample/trace_replay.py b/vllm/v1/worker/gpu/sample/trace_replay.py new file mode 100644 index 000000000000..9b58d4cf2923 --- /dev/null +++ b/vllm/v1/worker/gpu/sample/trace_replay.py @@ -0,0 +1,113 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch + +from vllm.sampling_params import SamplingParams +from vllm.triton_utils import tl, triton +from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor +from vllm.v1.worker.gpu.states import RequestState + + +class TraceReplayState: + """Per-request state for inference trace-replay. + + When a request carries ``SamplingParams.trace_decode_token_ids``, the + sampler overwrites the sampled token at each decode step with the + predetermined trace token, while real logprobs and ranks are still computed + from the unmodified logit distribution. The replay step for a request is + derived entirely from GPU state (``total_len - prompt_len``), so no CPU + synchronization or async placeholder handling is needed. + """ + + def __init__(self, req_states: RequestState): + self.max_num_reqs = req_states.max_num_reqs + self.device = req_states.device + self.req_states = req_states + self.trace_token_ids = StagedWriteTensor( + (self.max_num_reqs, req_states.max_model_len), + dtype=torch.int32, + device=self.device, + uva_instead_of_gpu=True, + ) + self.trace_len = UvaBackedTensor(self.max_num_reqs, dtype=torch.int32) + + def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None: + trace = sampling_params.trace_decode_token_ids + if trace is not None: + self.trace_len.np[req_idx] = len(trace) + self.trace_token_ids.stage_write(req_idx, 0, trace) + else: + self.trace_len.np[req_idx] = 0 + + def apply_staged_writes(self) -> None: + self.trace_len.copy_to_uva() + self.trace_token_ids.apply_write() + + def apply_trace( + self, + sampled: torch.Tensor, + idx_mapping: torch.Tensor, + ) -> None: + apply_trace_tokens( + sampled, + idx_mapping, + self.trace_token_ids.gpu, + self.trace_len.gpu, + self.req_states.total_len.gpu, + self.req_states.prompt_len.gpu, + ) + + +@triton.jit +def _trace_replay_kernel( + sampled_ptr, # [num_reqs], int64, mutated in place + idx_mapping_ptr, # [num_reqs] batch_idx -> req_state_idx + trace_token_ids_ptr, # [max_num_reqs, max_model_len], int32 + trace_token_ids_stride, + trace_len_ptr, # [max_num_reqs], int32 + total_len_ptr, # [max_num_reqs], int32 + prompt_len_ptr, # [max_num_reqs], int32 +): + batch_idx = tl.program_id(0) + req_state_idx = tl.load(idx_mapping_ptr + batch_idx) + if req_state_idx < 0: + return + + trace_len = tl.load(trace_len_ptr + req_state_idx) + if trace_len <= 0: + return + + # The token being sampled now is output token number + # (total_len - prompt_len): total_len reflects tokens committed through the + # previous step (post_update runs after sampling). + step = tl.load(total_len_ptr + req_state_idx) - tl.load( + prompt_len_ptr + req_state_idx + ) + if step < 0 or step >= trace_len: + return + + token_id = tl.load( + trace_token_ids_ptr + req_state_idx * trace_token_ids_stride + step + ) + tl.store(sampled_ptr + batch_idx, token_id.to(tl.int64)) + + +def apply_trace_tokens( + sampled: torch.Tensor, + idx_mapping: torch.Tensor, + trace_token_ids: torch.Tensor, + trace_len: torch.Tensor, + total_len: torch.Tensor, + prompt_len: torch.Tensor, +) -> None: + """Overwrite ``sampled`` in place with trace tokens for the current step.""" + num_reqs = idx_mapping.shape[0] + _trace_replay_kernel[(num_reqs,)]( + sampled, + idx_mapping, + trace_token_ids, + trace_token_ids.stride(0), + trace_len, + total_len, + prompt_len, + ) From d66300a1baa7779c68c7dfa4e51eee2502b48017 Mon Sep 17 00:00:00 2001 From: Ankit Nakhawa <40123208+AnkitNakhawa@users.noreply.github.com> Date: Thu, 20 Aug 2026 03:30:40 -0400 Subject: [PATCH 188/839] [Bugfix][EPD] Fix encoder round-robin fan-out (#52491) Signed-off-by: AnkitNakhawa --- .../disaggregated_encoder/disagg_epd_proxy.py | 30 ++++++- .../unit/test_epd_proxy_round_robin.py | 82 +++++++++++++++++++ 2 files changed, 110 insertions(+), 2 deletions(-) create mode 100644 tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py diff --git a/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py b/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py index 9a1bd517215b..267176fd3e83 100644 --- a/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py +++ b/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py @@ -51,6 +51,11 @@ prefill_session: aiohttp.ClientSession | None = None decode_session: aiohttp.ClientSession | None = None +# Cursor for round-robin encoder assignment, shared across requests so the +# fan-out doesn't restart from e_urls[0] every time. +encoder_rr_idx = 0 +encoder_rr_lock = asyncio.Lock() + ############################################################################### # Utils ############################################################################### @@ -58,6 +63,21 @@ MM_TYPES = {"image_url", "audio_url", "input_audio"} + +def encoder_rr_assignment( + e_urls: list[str], start: int, count: int +) -> tuple[list[str], int]: + """Assign `count` items to encoder URLs starting from cursor `start`. + + Returns the per-item URL list and the cursor value the next call should + start from, so the assignment is contiguous across calls instead of + restarting at e_urls[0] every time. + """ + urls = [e_urls[(start + i) % len(e_urls)] for i in range(count)] + next_start = (start + count) % len(e_urls) + return urls, next_start + + # Diagnostic switch: forward the original request to the decoder so the # only difference from the rewrite path is the rewrite itself. NO_REWRITE = False @@ -189,8 +209,14 @@ async def fanout_encoder_primer( item_uuids: dict[int, str] = {} item_meta: dict[int, dict] = {} - # Round-robin over encode servers to distribute load a bit - url_cycle = (e_urls[i % len(e_urls)] for i in range(len(mm_items))) + # Round-robin over encode servers to distribute load a bit. The cursor + # persists across requests so fan-out doesn't restart at e_urls[0] every + # time (which would hot-spot the first encoder for single-item requests). + global encoder_rr_idx + async with encoder_rr_lock: + url_cycle, encoder_rr_idx = encoder_rr_assignment( + e_urls, encoder_rr_idx, len(mm_items) + ) for idx, (item, target_url) in enumerate(zip(mm_items, url_cycle)): # Derive a *child* request id: :: diff --git a/tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py b/tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py new file mode 100644 index 000000000000..48123235a72a --- /dev/null +++ b/tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Encoder fan-out fairness in the disaggregated EPD proxy. + +Exercises the REAL ``encoder_rr_assignment`` helper loaded from its +``examples/`` path -- no routing logic is re-implemented here, so a future +change to that function is what these tests exercise. + +Regression target: the fan-out cursor used to restart at e_urls[0] on every +incoming request (``e_urls[i % len(e_urls)] for i in range(len(mm_items))``), +so single-item requests -- the common case -- always hit the first encoder +instance and left the rest idle. ``encoder_rr_assignment`` threads a cursor +across calls so consecutive requests keep advancing through every URL. +""" + +import importlib.util +from collections import Counter +from pathlib import Path + +import pytest + +PROXY_REL = "examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py" + + +def _load_proxy_module(): + path = Path(__file__).parents[4] / PROXY_REL + spec = importlib.util.spec_from_file_location("disagg_epd_proxy_under_test", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.fixture(scope="module") +def assign(): + return _load_proxy_module().encoder_rr_assignment + + +def _drive(assign, e_urls, counts): + """Feed a sequence of request item-counts through the cursor, as + fanout_encoder_primer does one request at a time.""" + cursor = 0 + all_urls = [] + for count in counts: + urls, cursor = assign(e_urls, cursor, count) + assert len(urls) == count + all_urls.append(urls) + return all_urls + + +@pytest.mark.parametrize("n_urls", [1, 2, 3, 5]) +def test_full_url_space_is_covered_uniformly(assign, n_urls): + e_urls = [f"E{i}" for i in range(n_urls)] + # Three full cycles of single-item requests. + urls = _drive(assign, e_urls, counts=[1] * (n_urls * 3)) + hits = Counter(u for req in urls for u in req) + assert set(hits) == set(e_urls) + assert max(hits.values()) == min(hits.values()) + + +def test_single_item_requests_rotate_through_all_encoders(assign): + # Exact scenario from the bug report: 3 encoders, one image/request. + e_urls = ["E0", "E1", "E2"] + routed = _drive(assign, e_urls, counts=[1] * 6) + assert [urls[0] for urls in routed] == ["E0", "E1", "E2", "E0", "E1", "E2"] + + +def test_cursor_stays_contiguous_across_varying_item_counts(assign): + e_urls = ["E0", "E1", "E2"] + routed = _drive(assign, e_urls, counts=[2, 1, 3, 1]) + assert routed == [ + ["E0", "E1"], + ["E2"], + ["E0", "E1", "E2"], + ["E0"], + ] + + +def test_single_encoder_always_resolves_to_it(assign): + urls, next_cursor = assign(["E0"], 0, 4) + assert urls == ["E0"] * 4 + assert next_cursor == 0 From 16cfe728d8d0bc3cd4a8397db0f392dd52a2c109 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <45557362+qgallouedec@users.noreply.github.com> Date: Thu, 20 Aug 2026 03:41:06 -0400 Subject: [PATCH 189/839] [Bugfix][Rust Frontend] Reject n > 1 in the `/inference/v1/generate` route (#52844) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Bugen Zhao Signed-off-by: Quentin Gallouédec Signed-off-by: Bugen Zhao --- .../server/src/routes/inference/generate.rs | 2 +- .../src/routes/inference/generate/convert.rs | 6 ++-- .../src/routes/inference/generate/types.rs | 19 +++++++++++- .../src/routes/inference/generate/validate.rs | 30 +++++++++++++++++-- rust/src/server/src/routes/render.rs | 7 +++-- rust/src/server/src/routes/tests.rs | 30 +++++++++++++++++++ 6 files changed, 85 insertions(+), 9 deletions(-) diff --git a/rust/src/server/src/routes/inference/generate.rs b/rust/src/server/src/routes/inference/generate.rs index 61fed421d38c..aa84913f23d5 100644 --- a/rust/src/server/src/routes/inference/generate.rs +++ b/rust/src/server/src/routes/inference/generate.rs @@ -26,11 +26,11 @@ use vllm_llm::{ }; use self::convert::{ResponseOptions, prepare_generate_request}; -pub(crate) use self::types::GenerateRequest; use self::types::{ GenerateLogprob, GenerateResponse, GenerateResponseChoice, GenerateResponseStreamChoice, GenerateStreamResponse, }; +pub(crate) use self::types::{GenerateRequest, GenerateSamplingParams}; pub(crate) use self::validate::validate_request_compat; use crate::config::ApiServerOptions; use crate::error::{ApiError, bail_server_error, server_error, text_submit_error}; diff --git a/rust/src/server/src/routes/inference/generate/convert.rs b/rust/src/server/src/routes/inference/generate/convert.rs index 5c141079890d..09c735238b2e 100644 --- a/rust/src/server/src/routes/inference/generate/convert.rs +++ b/rust/src/server/src/routes/inference/generate/convert.rs @@ -54,9 +54,9 @@ pub(super) fn prepare_generate_request( .as_ref() .and_then(|options| options.continuous_usage_stats) .unwrap_or(false); - let include_logprobs = request.sampling_params.logprobs.is_some(); - let include_prompt_logprobs = request.sampling_params.prompt_logprobs.is_some(); - let mut sampling_params = request.sampling_params; + let include_logprobs = request.sampling_params.inner.logprobs.is_some(); + let include_prompt_logprobs = request.sampling_params.inner.prompt_logprobs.is_some(); + let mut sampling_params = request.sampling_params.inner; sampling_params.vllm_xargs = merge_kv_transfer_params( sampling_params.vllm_xargs, request.kv_transfer_params.as_ref(), diff --git a/rust/src/server/src/routes/inference/generate/types.rs b/rust/src/server/src/routes/inference/generate/types.rs index 200209e095e1..b9220aa7b767 100644 --- a/rust/src/server/src/routes/inference/generate/types.rs +++ b/rust/src/server/src/routes/inference/generate/types.rs @@ -11,6 +11,23 @@ use vllm_text::SamplingParams; use crate::routes::openai::utils::types::{ChatLogProbs, Normalizable, StreamOptions, Usage}; +/// Sampling parameters for the token-in/token-out generate API. +/// +/// Wraps [`SamplingParams`] to additionally capture `n`, which the shared +/// northbound type intentionally omits (parallel sampling is handled by +/// higher layers, and the Rust frontend does not implement it). Capturing it +/// here lets validation reject `n > 1` explicitly instead of silently +/// dropping the key and returning a single choice. +#[serde_with::skip_serializing_none] +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct GenerateSamplingParams { + /// Number of output sequences to generate. Only `1` is supported. + pub n: Option, + /// The supported sampling parameters, lowered to the engine. + #[serde(flatten)] + pub inner: SamplingParams, +} + /// vLLM-compatible request type for the token-in/token-out generate API. #[serde_with::skip_serializing_none] #[derive(Debug, Clone, Deserialize, Serialize, Validate)] @@ -18,7 +35,7 @@ pub struct GenerateRequest { pub request_id: Option, pub model: Option, pub token_ids: Vec, - pub sampling_params: SamplingParams, + pub sampling_params: GenerateSamplingParams, #[serde(default)] pub stream: bool, pub stream_options: Option, diff --git a/rust/src/server/src/routes/inference/generate/validate.rs b/rust/src/server/src/routes/inference/generate/validate.rs index 99e2b23158cc..9ff8d0460ef6 100644 --- a/rust/src/server/src/routes/inference/generate/validate.rs +++ b/rust/src/server/src/routes/inference/generate/validate.rs @@ -23,6 +23,10 @@ pub(crate) fn validate_request_compat( ); } + if request.sampling_params.n.unwrap_or(1) != 1 { + bail_invalid_request!(param = "n", "Only n=1 is supported."); + } + if request.token_ids.is_empty() { bail_invalid_request!( param = "token_ids", @@ -30,14 +34,14 @@ pub(crate) fn validate_request_compat( ); } - if request.sampling_params.max_tokens == Some(0) { + if request.sampling_params.inner.max_tokens == Some(0) { bail_invalid_request!( param = "sampling_params", "max_tokens must be greater than 0." ); } - if let Some(prompt_logprobs) = request.sampling_params.prompt_logprobs { + if let Some(prompt_logprobs) = request.sampling_params.inner.prompt_logprobs { if prompt_logprobs < 0 && prompt_logprobs != -1 { bail_invalid_request!( param = "sampling_params", @@ -98,6 +102,28 @@ mod tests { assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err()); } + #[test] + fn validate_request_compat_rejects_parallel_sampling() { + let request: GenerateRequest = serde_json::from_value(json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "token_ids": [11, 22], + "sampling_params": {"n": 4} + })) + .expect("parse request"); + assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err()); + } + + #[test] + fn validate_request_compat_accepts_explicit_n_one() { + let request: GenerateRequest = serde_json::from_value(json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "token_ids": [11, 22], + "sampling_params": {"n": 1} + })) + .expect("parse request"); + assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_ok()); + } + #[test] fn validate_request_compat_rejects_empty_token_ids() { let request = GenerateRequest { diff --git a/rust/src/server/src/routes/render.rs b/rust/src/server/src/routes/render.rs index 1972c02d0976..b2d7601ddca9 100644 --- a/rust/src/server/src/routes/render.rs +++ b/rust/src/server/src/routes/render.rs @@ -17,7 +17,7 @@ use crate::error::{ApiError, text_submit_error}; use crate::lora::LoraModelResolution; use crate::render::RenderState; use crate::routes::inference::generate::{ - GenerateRequest, validate_request_compat as validate_generate_request, + GenerateRequest, GenerateSamplingParams, validate_request_compat as validate_generate_request, }; use crate::routes::openai::utils::types::{ListModelsResponse, ModelObject, StreamOptions}; use crate::routes::openai::utils::validated_json::ValidatedJson; @@ -86,7 +86,10 @@ fn lower_render_request( request_id: Some(text_request.request_id), model: Some(model), token_ids, - sampling_params: text_request.sampling_params, + sampling_params: GenerateSamplingParams { + n: None, + inner: text_request.sampling_params, + }, stream, stream_options, cache_salt: text_request.cache_salt, diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 9ac21ae07f9f..d57792b0bc28 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -4316,6 +4316,36 @@ async fn raw_generate_rejects_empty_token_ids() { assert_eq!(json["error"]["param"], "token_ids"); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn raw_generate_rejects_parallel_sampling() { + let mut app = test_app().await; + + let response = app + .call( + Request::builder() + .method("POST") + .uri("/inference/v1/generate") + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "token_ids": [11, 22], + "sampling_params": {"n": 4} + }) + .to_string(), + )) + .expect("build request"), + ) + .await + .expect("call app"); + + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body"); + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json"); + assert_eq!(json["error"]["param"], "n"); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn raw_generate_rejects_streaming_prompt_logprobs() { From a34fd69106080aed8eab1490883942f2fe8264b7 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Thu, 20 Aug 2026 01:10:06 -0700 Subject: [PATCH 190/839] [CI][Bugfix] Update distributed DP API server test path (#52939) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> From 14617c2b6c1257ac0d6c7b5e05b195ca30013827 Mon Sep 17 00:00:00 2001 From: Anton A Date: Thu, 20 Aug 2026 04:11:44 -0400 Subject: [PATCH 191/839] [Bugfix] DeepEP-V2: expert_tokens_meta must be None on the decode/cudagraph path (empty recv_expert_num_tokens) (#52632) Signed-off-by: Anton Alexander Signed-off-by: Roger Wang Co-authored-by: Roger Wang --- .../fused_moe/prepare_finalize/deepep_v2.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py index 5e950fdf606e..916cebd74c0b 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py @@ -210,14 +210,23 @@ def _receiver( else: expert_x, expert_x_scale = recv_x, None - expert_tokens_meta = mk.ExpertTokensMetadata.make_from_list( - recv_expert_num_tokens, - device=expert_x.device, - ) + if recv_expert_num_tokens: + expert_tokens_meta = mk.ExpertTokensMetadata.make_from_list( + recv_expert_num_tokens, + device=expert_x.device, + ) + else: + # Decode/cudagraph path (do_cpu_sync=False) skips the CPU sync and + # leaves recv_expert_num_tokens empty. A present-but-empty + # ExpertTokensMetadata violates the decode-mode contract above + # (expert_tokens_meta must be None) and crashes DeepEP combine + # during profile_run when CUDA graphs are enabled. + expert_tokens_meta = None if recv_topk_idx is None: # do_expand=True (prefill mode): build topk_ids from # per-expert token counts. + assert expert_tokens_meta is not None total_tokens = sum(recv_expert_num_tokens) if total_tokens > 0: recv_topk_idx = torch.repeat_interleave( From 4666a8ba9ed5d271bf751cf1384cab2d2dcc0cfa Mon Sep 17 00:00:00 2001 From: Cyrus Leung Date: Thu, 20 Aug 2026 16:39:48 +0800 Subject: [PATCH 192/839] [Refactor] Remove InputPreprocessor (#53064) Signed-off-by: DarkLight1337 --- .buildkite/test-amd.yaml | 2 - .buildkite/test_areas/misc.yaml | 2 - .github/CODEOWNERS | 1 - .../test_process_multi_modal_uuids.py | 3 +- tests/test_inputs.py | 31 -- vllm/inputs/preprocess.py | 291 ------------------ vllm/renderers/base.py | 7 +- vllm/v1/engine/input_processor.py | 59 ++-- 8 files changed, 29 insertions(+), 367 deletions(-) delete mode 100644 tests/test_inputs.py delete mode 100644 vllm/inputs/preprocess.py diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index e21c8f59e627..e789829fe941 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -172,7 +172,6 @@ steps: source_file_dependencies: - vllm/ - tests/test_envs.py - - tests/test_inputs.py - tests/test_outputs.py - tests/test_pooling_params.py - tests/test_ray_env.py @@ -189,7 +188,6 @@ steps: commands: - python3 standalone_tests/lazy_imports.py - pytest -v -s test_envs.py - - pytest -v -s test_inputs.py - pytest -v -s test_outputs.py - pytest -v -s test_pooling_params.py - pytest -v -s test_ray_env.py diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index a73ca9bb8192..b564896fbdf0 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -358,7 +358,6 @@ steps: - vllm/utils/ - vllm/v1/ - tests/test_envs.py - - tests/test_inputs.py - tests/test_outputs.py - tests/test_pooling_params.py - tests/test_ray_env.py @@ -376,7 +375,6 @@ steps: commands: - python3 standalone_tests/lazy_imports.py - pytest -v -s test_envs.py - - pytest -v -s test_inputs.py - pytest -v -s test_outputs.py - pytest -v -s test_pooling_params.py - pytest -v -s test_ray_env.py diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index a9a21b5ef17d..95ea9046d5a9 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -107,7 +107,6 @@ /tests/multimodal @DarkLight1337 @ywang96 @NickLucche @shen-shanshan /tests/multimodal/media @Isotr0py /tests/quantization @mgoin @robertgshaw2-redhat @yewentao256 @pavanimajety @zyongye @AndreasKaratzas -/tests/test_inputs.py @DarkLight1337 @ywang96 /tests/entrypoints/llm/test_struct_output_generate.py @mgoin @russellb @aarnphm /tests/v1/structured_output @mgoin @russellb @aarnphm /tests/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery @ivanium diff --git a/tests/renderers/test_process_multi_modal_uuids.py b/tests/renderers/test_process_multi_modal_uuids.py index c811630c859f..76b394897ee1 100644 --- a/tests/renderers/test_process_multi_modal_uuids.py +++ b/tests/renderers/test_process_multi_modal_uuids.py @@ -54,11 +54,10 @@ def test_text_only_model_mm_data_maps_to_bad_request(): with pytest.raises(ValueError, match="text-only") as exc_info: renderer._process_multimodal( - prompt="What is in this image?", + prompt=[1], mm_data={"image": [cherry_pil_image]}, mm_uuids=None, mm_processor_kwargs=None, - tokenization_kwargs=None, ) error_response = create_error_response(exc_info.value) diff --git a/tests/test_inputs.py b/tests/test_inputs.py deleted file mode 100644 index fb1bbd21eacd..000000000000 --- a/tests/test_inputs.py +++ /dev/null @@ -1,31 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import pytest - -from vllm.config import ModelConfig, VllmConfig -from vllm.inputs.preprocess import InputPreprocessor - -pytestmark = pytest.mark.cpu_test - - -@pytest.mark.parametrize("model_id", ["facebook/chameleon-7b"]) -@pytest.mark.parametrize("prompt", ["", {"prompt_token_ids": []}]) -@pytest.mark.skip( - reason=( - "Applying huggingface processor on text inputs results in " - "significant performance regression for multimodal models. " - "See https://github.com/vllm-project/vllm/issues/26320" - ) -) -def test_preprocessor_always_mm_code_path(model_id, prompt): - model_config = ModelConfig(model=model_id) - vllm_config = VllmConfig(model_config=model_config) - input_preprocessor = InputPreprocessor(vllm_config) - - # HF processor adds sep token - tokenizer = input_preprocessor.get_tokenizer() - sep_token_id = tokenizer.vocab[tokenizer.sep_token] - - processed_inputs = input_preprocessor.preprocess(prompt) - assert sep_token_id in processed_inputs["prompt_token_ids"] diff --git a/vllm/inputs/preprocess.py b/vllm/inputs/preprocess.py deleted file mode 100644 index 7722014f9c0f..000000000000 --- a/vllm/inputs/preprocess.py +++ /dev/null @@ -1,291 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from collections.abc import Mapping -from typing import Any, overload - -from typing_extensions import assert_never - -from vllm.config import VllmConfig -from vllm.inputs import build_enc_dec_input -from vllm.logger import init_logger -from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry -from vllm.renderers import BaseRenderer, renderer_from_config -from vllm.renderers.inputs import ( - DecoderDictPrompt, - DecoderOnlyDictPrompt, - EncoderDecoderDictPrompt, - EncoderDictPrompt, - SingletonDictPrompt, -) -from vllm.renderers.inputs.preprocess import parse_dec_only_prompt, parse_enc_dec_prompt -from vllm.tokenizers import TokenizerLike - -from .engine import ( - DecoderEngineInput, - DecoderOnlyEngineInput, - EmbedsInput, - EncoderDecoderInput, - EncoderInput, - EngineInput, - MultiModalInput, - SingletonInput, - TokensInput, - tokens_input, -) -from .llm import ( - EmbedsPrompt, - MultiModalDataDict, - MultiModalUUIDDict, - PromptType, - TextPrompt, - TokensPrompt, -) - -logger = init_logger(__name__) - - -class InputPreprocessor: - def __init__( - self, - vllm_config: VllmConfig, - renderer: BaseRenderer | None = None, - mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY, - ) -> None: - super().__init__() - - self.model_config = vllm_config.model_config - self.renderer = renderer or renderer_from_config(vllm_config) - self.mm_registry = mm_registry - - @property - def tokenizer(self) -> TokenizerLike | None: - return self.renderer.tokenizer - - def get_tokenizer(self) -> TokenizerLike: - return self.renderer.get_tokenizer() - - def _tokenize_prompt( - self, - prompt: str, - tokenization_kwargs: dict[str, Any] | None = None, - ) -> list[int]: - """ - Apply the model's tokenizer to a text prompt, returning the - corresponding token IDs. - """ - renderer = self.renderer - - tok_params = renderer.default_cmpl_tok_params.with_kwargs( - **(tokenization_kwargs or {}) - ) - - tok_prompt = renderer._tokenize_singleton_prompt( - TextPrompt(prompt=prompt), - tok_params, - ) - - return tok_prompt["prompt_token_ids"] - - def _process_multimodal( - self, - prompt: str | list[int], - mm_data: MultiModalDataDict, - mm_processor_kwargs: Mapping[str, object] | None = None, - tokenization_kwargs: dict[str, Any] | None = None, - *, - mm_uuids: MultiModalUUIDDict | None = None, - ) -> MultiModalInput: - """ - Apply the model's multi-modal processor to a multi-modal prompt, - returning the corresponding token IDs and metadata. - """ - return self.renderer._process_multimodal( - prompt, - mm_data, - mm_uuids=mm_uuids, - mm_processor_kwargs=mm_processor_kwargs, - tokenization_kwargs=tokenization_kwargs, - ) - - def _process_embeds( - self, - parsed_content: EmbedsPrompt, - ) -> EmbedsInput: - return self.renderer._process_embeds(parsed_content) - - def _truncate_inputs( - self, inputs: list[int], tokenization_kwargs: dict[str, Any] | None = None - ) -> list[int]: - renderer = self.renderer - - tok_params = renderer.default_cmpl_tok_params.with_kwargs( - **(tokenization_kwargs or {}) - ) - - tok_prompt = renderer._tokenize_singleton_prompt( - TokensPrompt(prompt_token_ids=inputs), - tok_params, - ) - - return tok_prompt["prompt_token_ids"] - - def _process_tokens( - self, - parsed_content: TokensPrompt, - tokenization_kwargs: dict[str, Any] | None = None, - ) -> TokensInput | MultiModalInput: - prompt_token_ids = self._truncate_inputs( - parsed_content["prompt_token_ids"], tokenization_kwargs - ) - - inputs: TokensInput | MultiModalInput - if multi_modal_data := parsed_content.get("multi_modal_data"): - inputs = self._process_multimodal( - prompt_token_ids, - multi_modal_data, - parsed_content.get("mm_processor_kwargs"), - tokenization_kwargs=tokenization_kwargs, - mm_uuids=parsed_content.get("multi_modal_uuids"), - ) - else: - inputs = tokens_input(prompt_token_ids) - - if prompt_text := parsed_content.get("prompt"): - inputs["prompt"] = prompt_text - if cache_salt := parsed_content.get("cache_salt"): - inputs["cache_salt"] = cache_salt - - return inputs - - def _process_text( - self, - parsed_content: TextPrompt, - tokenization_kwargs: dict[str, Any] | None = None, - ) -> TokensInput | MultiModalInput: - prompt_text = parsed_content["prompt"] - - inputs: TokensInput | MultiModalInput - if multi_modal_data := parsed_content.get("multi_modal_data"): - inputs = self._process_multimodal( - prompt_text, - multi_modal_data, - parsed_content.get("mm_processor_kwargs") or {}, - tokenization_kwargs=tokenization_kwargs, - ) - else: - prompt_token_ids = self._tokenize_prompt( - prompt_text, - tokenization_kwargs=tokenization_kwargs, - ) - inputs = tokens_input(prompt_token_ids) - - inputs["prompt"] = prompt_text - - if cache_salt := parsed_content.get("cache_salt"): - inputs["cache_salt"] = cache_salt - - return inputs - - @overload - def _prompt_to_llm_inputs( - self, - prompt: EncoderDictPrompt, - tokenization_kwargs: dict[str, Any] | None = None, - ) -> EncoderInput: ... - - @overload - def _prompt_to_llm_inputs( # type: ignore[misc] - self, - prompt: DecoderDictPrompt, - tokenization_kwargs: dict[str, Any] | None = None, - ) -> DecoderEngineInput: ... - - @overload - def _prompt_to_llm_inputs( # type: ignore[misc] - self, - prompt: DecoderOnlyDictPrompt, - tokenization_kwargs: dict[str, Any] | None = None, - ) -> DecoderOnlyEngineInput: ... - - def _prompt_to_llm_inputs( - self, - prompt: SingletonDictPrompt, - tokenization_kwargs: dict[str, Any] | None = None, - ) -> SingletonInput: - if "prompt_embeds" in prompt: - return self._process_embeds(prompt) # type: ignore[arg-type] - - if "prompt_token_ids" in prompt: - return self._process_tokens(prompt) # type: ignore[arg-type] - - if "prompt" in prompt: - return self._process_text( - prompt, # type: ignore[arg-type] - tokenization_kwargs=tokenization_kwargs, - ) - - assert_never(prompt) # type: ignore[arg-type] - - def _process_encoder_decoder_prompt( - self, - prompt: EncoderDecoderDictPrompt, - tokenization_kwargs: dict[str, Any] | None = None, - ) -> EncoderDecoderInput: - encoder_prompt = prompt["encoder_prompt"] - decoder_prompt = prompt["decoder_prompt"] - - skip_decoder_start_token = False - if self.renderer.mm_processor is not None: - from vllm.multimodal.processing import EncDecMultiModalProcessor - - if isinstance(self.renderer.mm_processor, EncDecMultiModalProcessor): - skip_decoder_start_token = ( - self.renderer.mm_processor.skip_decoder_start_token - ) - - return build_enc_dec_input( - encoder_input=self._prompt_to_llm_inputs( - encoder_prompt, - tokenization_kwargs=tokenization_kwargs, - ), - decoder_input=( - None - if decoder_prompt is None - else self._prompt_to_llm_inputs( - decoder_prompt, - tokenization_kwargs=tokenization_kwargs, - ) - ), - decoder_start_token_id=self.renderer.get_dec_start_token_id(), - skip_decoder_start_token=skip_decoder_start_token, - ) - - def _process_decoder_only_prompt( - self, - prompt: DecoderOnlyDictPrompt, - tokenization_kwargs: dict[str, Any] | None = None, - ) -> DecoderOnlyEngineInput: - return self._prompt_to_llm_inputs( - prompt, - tokenization_kwargs=tokenization_kwargs, - ) - - def preprocess( - self, - prompt: PromptType, - tokenization_kwargs: dict[str, Any] | None = None, - ) -> EngineInput: - """Preprocess the input prompt.""" - if self.model_config.is_encoder_decoder: - # Encoder-decoder model requires special mapping of - # input prompts to encoder & decoder. - return self._process_encoder_decoder_prompt( - parse_enc_dec_prompt(prompt), - tokenization_kwargs, - ) - - return self._process_decoder_only_prompt( - parse_dec_only_prompt(prompt), - tokenization_kwargs=tokenization_kwargs, - ) diff --git a/vllm/renderers/base.py b/vllm/renderers/base.py index 832f954eaf9c..aaf89e6edd4b 100644 --- a/vllm/renderers/base.py +++ b/vllm/renderers/base.py @@ -725,14 +725,12 @@ def _process_mm_uuids( return mm_uuid_items - # TODO: Remove str and tokenization_kwargs after deprecating InputPreprocessor def _process_multimodal( self, - prompt: list[int] | str, + prompt: list[int], mm_data: MultiModalDataDict, mm_uuids: MultiModalUUIDDict | None, mm_processor_kwargs: Mapping[str, object] | None, - tokenization_kwargs: dict[str, Any] | None, *, skip_mm_cache: bool = False, ) -> "MultiModalInput": @@ -755,7 +753,6 @@ def _process_multimodal( mm_data_items, mm_uuid_items, hf_processor_mm_kwargs=mm_processor_kwargs or {}, - tokenization_kwargs=tokenization_kwargs or {}, ) mm_timing_ctx = self._mm_timing_registry.get(mm_req_id) @@ -783,7 +780,6 @@ def _process_tokens( prompt_token_ids, multi_modal_data, mm_processor_kwargs=prompt.get("mm_processor_kwargs"), - tokenization_kwargs=None, # Tokenization already done in Step 2 mm_uuids=prompt.get("multi_modal_uuids"), skip_mm_cache=skip_mm_cache, ) @@ -846,7 +842,6 @@ async def _process_tokens_async( prompt_token_ids, multi_modal_data, mm_processor_kwargs=prompt.get("mm_processor_kwargs"), - tokenization_kwargs=None, mm_uuids=prompt.get("multi_modal_uuids"), skip_mm_cache=skip_mm_cache, ) diff --git a/vllm/v1/engine/input_processor.py b/vllm/v1/engine/input_processor.py index 744da670a04f..2b2ff4bab89c 100644 --- a/vllm/v1/engine/input_processor.py +++ b/vllm/v1/engine/input_processor.py @@ -14,7 +14,6 @@ SingletonInput, split_enc_dec_input, ) -from vllm.inputs.preprocess import InputPreprocessor from vllm.logger import init_logger from vllm.lora.request import LoRARequest from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry @@ -24,6 +23,7 @@ from vllm.platforms import current_platform from vllm.pooling_params import PoolingParams from vllm.renderers import BaseRenderer, renderer_from_config +from vllm.renderers.inputs.preprocess import parse_model_prompt from vllm.sampling_params import SamplingParams from vllm.tasks import GENERATION_TASKS, POOLING_TASKS, SupportedTask from vllm.tokenizers import TokenizerLike @@ -67,12 +67,6 @@ def __init__( ) mm_budget.reset_cache() # Not used anymore - self.input_preprocessor = InputPreprocessor( - vllm_config, - renderer=renderer, - mm_registry=mm_registry, - ) - # Raw-prompt preprocessing (tokenization and multimodal processing) # is blocking, so async callers should run it on the renderer's # thread pool to keep their event loop responsive. @@ -313,44 +307,45 @@ def process_inputs( ) if isinstance(prompt, dict) and "type" in prompt: - if tokenization_kwargs: - logger.warning_once( - "Passing tokenization_kwargs to InputProcessor is deprecated " - "and will be removed in v0.18. You should instead pass " - "them to Renderer.render_cmpl() or Renderer.render_chat()." - ) - if arrival_time is None: arrival_time = prompt.get("arrival_time", time.time()) # type: ignore[assignment] - processed_inputs: EngineInput = prompt # type: ignore[assignment] + engine_input: EngineInput = prompt # type: ignore[assignment] else: logger.warning_once( "Passing raw prompts to InputProcessor is deprecated " - "and will be removed in v0.18. You should instead pass " + "and will be removed in the future. You should instead pass " "the outputs of Renderer.render_cmpl() or Renderer.render_chat()." ) if arrival_time is None: arrival_time = time.time() - processed_inputs = self.input_preprocessor.preprocess( - prompt, - tokenization_kwargs=tokenization_kwargs, + renderer = self.renderer + model_config = self.model_config + + parsed_prompt = parse_model_prompt(model_config, prompt) + tok_params = renderer.default_cmpl_tok_params.with_kwargs( + **(tokenization_kwargs or {}) + ) + + (engine_input,) = renderer.render_cmpl( + [parsed_prompt], + tok_params, ) - current_platform.validate_request(processed_inputs, params) + current_platform.validate_request(engine_input, params) - encoder_inputs, decoder_inputs = split_enc_dec_input(processed_inputs) - self._validate_model_inputs(encoder_inputs, decoder_inputs) + encoder_input, decoder_input = split_enc_dec_input(engine_input) + self._validate_model_inputs(encoder_input, decoder_input) # Mypy can be conservative for TypedDict unions; normalize access. - if decoder_inputs["type"] == "embeds": - prompt_embeds = decoder_inputs["prompt_embeds"] - prompt_token_ids = decoder_inputs.get("prompt_token_ids") - prompt_is_token_ids = decoder_inputs.get("is_token_ids") + if decoder_input["type"] == "embeds": + prompt_embeds = decoder_input["prompt_embeds"] + prompt_token_ids = decoder_input.get("prompt_token_ids") + prompt_is_token_ids = decoder_input.get("is_token_ids") else: - prompt_token_ids = decoder_inputs["prompt_token_ids"] + prompt_token_ids = decoder_input["prompt_token_ids"] prompt_embeds = None prompt_is_token_ids = None @@ -385,10 +380,10 @@ def process_inputs( # Multimodal related. mm_features: list[MultiModalFeatureSpec] | None = None - if decoder_inputs["type"] == "multimodal": - decoder_mm_inputs = decoder_inputs["mm_kwargs"] - decoder_mm_positions = decoder_inputs["mm_placeholders"] - decoder_mm_hashes = decoder_inputs["mm_hashes"] + if decoder_input["type"] == "multimodal": + decoder_mm_inputs = decoder_input["mm_kwargs"] + decoder_mm_positions = decoder_input["mm_placeholders"] + decoder_mm_hashes = decoder_input["mm_hashes"] if not all( isinstance(leaf, str) for leaf in json_iter_leaves(decoder_mm_hashes) @@ -430,7 +425,7 @@ def process_inputs( pooling_params=pooling_params, arrival_time=arrival_time, lora_request=lora_request, - cache_salt=decoder_inputs.get("cache_salt"), + cache_salt=decoder_input.get("cache_salt"), priority=priority, data_parallel_rank=data_parallel_rank, trace_headers=trace_headers, From 963fcfa48c6306cb459655cae0fbaa3a3d1040e7 Mon Sep 17 00:00:00 2001 From: Roy Wang Date: Thu, 20 Aug 2026 16:50:51 +0800 Subject: [PATCH 193/839] [Rust][Benchmark] Align speed-bench CLI flags with Python and add flag parity test (#51592) Co-authored-by: Claude Fable 5 Co-authored-by: Bugen Zhao Signed-off-by: esmeetu Signed-off-by: Bugen Zhao --- .buildkite/test_areas/benchmarks.yaml | 1 + rust/src/bench/README.md | 10 +- rust/src/bench/src/benchmark.rs | 81 ++++++++++-- rust/src/bench/src/cli.rs | 13 +- rust/src/bench/src/config.rs | 27 ++++ rust/src/bench/tests/cli_parity.rs | 119 ++++++++++++++++++ rust/src/bench/tests/python_serve_flags.txt | 110 ++++++++++++++++ .../benchmarks/test_rust_bench_cli_parity.py | 61 +++++++++ 8 files changed, 406 insertions(+), 16 deletions(-) create mode 100644 rust/src/bench/tests/cli_parity.rs create mode 100644 rust/src/bench/tests/python_serve_flags.txt create mode 100644 tests/benchmarks/test_rust_bench_cli_parity.py diff --git a/.buildkite/test_areas/benchmarks.yaml b/.buildkite/test_areas/benchmarks.yaml index 4010b8d0b8d8..94d9999b81f2 100644 --- a/.buildkite/test_areas/benchmarks.yaml +++ b/.buildkite/test_areas/benchmarks.yaml @@ -9,6 +9,7 @@ steps: source_file_dependencies: - vllm/ - "!vllm/distributed/kv_transfer/" + - rust/src/bench/tests/python_serve_flags.txt - tests/benchmarks/ commands: - pytest -v -s benchmarks/ diff --git a/rust/src/bench/README.md b/rust/src/bench/README.md index 62681254afb0..ec58c40e5a3a 100644 --- a/rust/src/bench/README.md +++ b/rust/src/bench/README.md @@ -483,7 +483,7 @@ Run `vllm-bench --help` for the authoritative list. Grouped reference below. | `--model` | Auto-detect | Model name (fetched from `/v1/models` if omitted) | | `--served-model-name` | — | Model name used in API requests | | `--tokenizer` | Same as model | Tokenizer name or path (supports HF, tiktoken, server fallback) | -| `--tokenizer-mode` | `auto` | Tokenizer mode (`auto`, `hf`, `slow`, `mistral`) | +| `--tokenizer-mode` | `auto` | Accepted for Python CLI compatibility but ignored (warns on non-`auto`): resolution always follows the HF `tokenizer.json` → tiktoken → server-side `/tokenize` chain, so `mistral_common` (tekken) tokenizers fall back to server-side tokenization | | `--trust-remote-code` | `false` | Trust remote code for tokenizer | | `--skip-tokenizer-init` | `false` | Skip tokenizer initialization | @@ -521,7 +521,8 @@ Run `vllm-bench --help` for the authoritative list. Grouped reference below. | `--sonnet-output-len` | `150` | Output tokens per request | | `--sonnet-prefix-len` | `200` | Prefix tokens shared across requests | | **SPEED-Bench** | | | -| `--speed-bench-config` | `qualitative` | Split (`qualitative`, `throughput_1k`/`2k`/`8k`/`16k`/`32k`) | +| `--speed-bench-config` | `qualitative` | Split (`qualitative`, `throughput_1k`/`2k`/`8k`/`16k`/`32k`); alias: `--speed-bench-dataset-subset` (Python name) | +| `--speed-bench-output-len` | `4096` | Output tokens per request (matches Python default) | | `--speed-bench-category` | — | Filter by category (`low_entropy`, `high_entropy`, `mixed_entropy`, `coding`, `math`, …) | | `--speed-bench-max-input-len` | — | Truncate prompts to at most N tokens | | **HuggingFace** | | | @@ -794,6 +795,11 @@ The Rust implementation matches Python `vllm bench serve` in: - Rate control (Gamma distribution, normalization, burstiness, linear/exponential ramp-up) - Metrics (TTFT/TPOT/ITL/E2EL percentiles, peak tokens/sec, peak concurrency, goodput) - Sampling parameters merged into the request body via `extra_body` (same precedence rules) +- CLI flags: a superset of Python's, enforced by `tests/cli_parity.rs` against a + snapshot of the Python parser (`tests/python_serve_flags.txt`, regenerated by + `tests/benchmarks/test_rust_bench_cli_parity.py` in the repo root) with an + explicit allowlist for the Python-only remainder — so `VLLM_USE_RUST_BENCH=1` + delegation cannot silently reject documented flags ## Environment Variables diff --git a/rust/src/bench/src/benchmark.rs b/rust/src/bench/src/benchmark.rs index d2a7bf84802d..feb49cf84d91 100644 --- a/rust/src/bench/src/benchmark.rs +++ b/rust/src/bench/src/benchmark.rs @@ -134,6 +134,17 @@ pub(crate) async fn fetch_spec_decode_metrics( Err(_) => return None, }; + parse_spec_decode_metrics(&text) +} + +/// Parse spec decode counters from Prometheus text exposition format. +/// +/// Matches on the metric name (the part before any labels) rather than the +/// whole line so label values cannot be mistaken for metric names, and only +/// reads `_total` counter samples — skipping e.g. `_created` timestamp series, +/// which would otherwise be summed into the counts (mirrors the Python fix +/// in vllm#41916). +pub(crate) fn parse_spec_decode_metrics(text: &str) -> Option { let mut num_drafts: u64 = 0; let mut num_draft_tokens: u64 = 0; let mut num_accepted_tokens: u64 = 0; @@ -145,24 +156,25 @@ pub(crate) async fn fetch_spec_decode_metrics( if line.is_empty() || line.starts_with('#') { continue; } - if !line.starts_with("vllm:spec_decode") { + let first_token = match line.split_whitespace().next() { + Some(t) => t, + None => continue, + }; + let metric_name = first_token.split('{').next().unwrap_or(first_token); + if !metric_name.starts_with("vllm:spec_decode") || !metric_name.ends_with("_total") { continue; } found_spec_decode = true; - let parts: Vec<&str> = line.split_whitespace().collect(); - if parts.is_empty() { - continue; - } - let val = match parts.last().and_then(|s| s.parse::().ok()) { + let val = match line.split_whitespace().last().and_then(|s| s.parse::().ok()) { Some(v) => v as u64, None => continue, }; - if line.contains("num_drafts") { + if metric_name.contains("num_drafts") { num_drafts += val; - } else if line.contains("num_draft_tokens") { + } else if metric_name.contains("num_draft_tokens") { num_draft_tokens += val; - } else if line.contains("num_accepted_tokens_per_pos") { + } else if metric_name.contains("num_accepted_tokens_per_pos") { // Parse position label: position="N" if let Some(start) = line.find("position=\"") { let start = start + "position=\"".len(); @@ -172,7 +184,7 @@ pub(crate) async fn fetch_spec_decode_metrics( *accepted_per_pos.entry(pos).or_insert(0) += val; } } - } else if line.contains("num_accepted_tokens") { + } else if metric_name.contains("num_accepted_tokens") { num_accepted_tokens += val; } } @@ -502,12 +514,11 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { downloaded.as_str() } }; - let output_len = config.sharegpt_output_len.unwrap_or(config.random_output_len); crate::datasets::speed_bench::load_speed_bench_dataset( tok, path, config.num_prompts, - output_len, + config.speed_bench_output_len, config.seed, &config.request_id_prefix, config.speed_bench_category.as_deref(), @@ -1875,6 +1886,52 @@ mod tests { } } + #[test] + fn test_parse_spec_decode_metrics_sums_counters_across_engines() { + let text = "\ +# TYPE vllm:spec_decode_num_drafts_total counter +vllm:spec_decode_num_drafts_total{model_name=\"m\",engine=\"0\"} 100.0 +vllm:spec_decode_num_drafts_total{model_name=\"m\",engine=\"1\"} 50.0 +vllm:spec_decode_num_draft_tokens_total{model_name=\"m\",engine=\"0\"} 700.0 +vllm:spec_decode_num_accepted_tokens_total{model_name=\"m\",engine=\"0\"} 300.0 +vllm:spec_decode_num_accepted_tokens_per_pos_total{position=\"0\",engine=\"0\"} 80.0 +vllm:spec_decode_num_accepted_tokens_per_pos_total{position=\"1\",engine=\"0\"} 40.0 +vllm:num_requests_running{model_name=\"m\"} 0 +"; + let m = parse_spec_decode_metrics(text).unwrap(); + assert_eq!(m.num_drafts, 150); + assert_eq!(m.num_draft_tokens, 700); + assert_eq!(m.num_accepted_tokens, 300); + assert_eq!(m.accepted_per_pos, [(0, 80), (1, 40)].into_iter().collect()); + } + + #[test] + fn test_parse_spec_decode_metrics_skips_created_series_and_label_values() { + // `_created` timestamps must not be summed into counters, and metric + // matching must key off the metric name, not substrings inside label + // values (vllm#41916). + let text = "\ +vllm:spec_decode_num_drafts_total{model_name=\"m\"} 100.0 +vllm:spec_decode_num_drafts_created{model_name=\"m\"} 1.7863e+09 +vllm:spec_decode_num_draft_tokens_total{model_name=\"m\"} 700.0 +vllm:spec_decode_num_draft_tokens_created{model_name=\"m\"} 1.7863e+09 +vllm:spec_decode_num_accepted_tokens_total{model_name=\"num_drafts\"} 300.0 +"; + let m = parse_spec_decode_metrics(text).unwrap(); + assert_eq!(m.num_drafts, 100); + assert_eq!(m.num_draft_tokens, 700); + assert_eq!(m.num_accepted_tokens, 300); + } + + #[test] + fn test_parse_spec_decode_metrics_none_without_spec_metrics() { + let text = "\ +vllm:num_requests_running{model_name=\"m\"} 0 +vllm:prefix_cache_queries_total{model_name=\"m\"} 0 +"; + assert!(parse_spec_decode_metrics(text).is_none()); + } + #[test] fn test_compute_spec_decode_stats_basic() { let before = make_metrics(100, 300, 200, &[(0, 90), (1, 70), (2, 50)]); diff --git a/rust/src/bench/src/cli.rs b/rust/src/bench/src/cli.rs index 06e32e14ce3b..24a7b3e2512a 100644 --- a/rust/src/bench/src/cli.rs +++ b/rust/src/bench/src/cli.rs @@ -176,7 +176,8 @@ pub struct BenchServeArgs { #[arg(long)] pub tokenizer: Option, - /// Tokenizer mode (auto, hf, slow, mistral). + /// Tokenizer mode (auto, hf, slow, mistral). Accepted for Python CLI + /// compatibility; non-auto values are ignored with a warning. #[arg(long, default_value = "auto")] pub tokenizer_mode: String, @@ -535,9 +536,17 @@ pub struct BenchServeArgs { /// SPEED-Bench config/split (qualitative, throughput_1k, throughput_2k, throughput_8k, /// throughput_16k, throughput_32k). - #[arg(long, default_value = "qualitative")] + #[arg( + long, + visible_alias = "speed-bench-dataset-subset", + default_value = "qualitative" + )] pub speed_bench_config: SpeedBenchConfig, + /// Number of output tokens per request (SPEED-Bench dataset). + #[arg(long, default_value_t = 4096)] + pub speed_bench_output_len: usize, + /// Filter SPEED-Bench by category (e.g. low_entropy, high_entropy, coding, math). #[arg(long)] pub speed_bench_category: Option, diff --git a/rust/src/bench/src/config.rs b/rust/src/bench/src/config.rs index 0700cf1fad3e..3a0936b5f4c1 100644 --- a/rust/src/bench/src/config.rs +++ b/rust/src/bench/src/config.rs @@ -196,6 +196,7 @@ pub struct BenchConfig { pub speed_bench_config: SpeedBenchConfig, pub speed_bench_category: Option, pub speed_bench_max_input_len: Option, + pub speed_bench_output_len: usize, pub hf_split: Option, pub hf_subset: Option, pub hf_output_len: Option, @@ -640,6 +641,15 @@ impl BenchConfig { )); } + if args.tokenizer_mode != "auto" { + tracing::warn!( + mode = %args.tokenizer_mode, + "--tokenizer-mode is ignored by the Rust client; tokenizer resolution always \ + follows the HF tokenizer.json -> tiktoken -> server-side /tokenize fallback \ + chain (mistral_common tokenizers are not supported locally)" + ); + } + Ok(BenchConfig { backend: args.backend, base_url, @@ -727,6 +737,7 @@ impl BenchConfig { speed_bench_config: args.speed_bench_config, speed_bench_category: args.speed_bench_category.clone(), speed_bench_max_input_len: args.speed_bench_max_input_len, + speed_bench_output_len: args.speed_bench_output_len, hf_split: args.hf_split.clone(), hf_subset: args.hf_subset.clone(), hf_output_len: args.hf_output_len, @@ -881,6 +892,22 @@ mod tests { ] } + #[test] + fn test_speed_bench_flags_match_python() { + let args = parse_args(vec![ + "vllm-bench", + "--model", + "test-model", + "--speed-bench-dataset-subset", + "throughput_8k", + ]); + assert!(matches!( + args.speed_bench_config, + crate::cli::SpeedBenchConfig::Throughput8k + )); + assert_eq!(args.speed_bench_output_len, 4096); + } + #[test] fn test_prefix_sharing_defaults_to_zero() { let args = base_multi_turn_args(); diff --git a/rust/src/bench/tests/cli_parity.rs b/rust/src/bench/tests/cli_parity.rs new file mode 100644 index 000000000000..b0fd09ab9e36 --- /dev/null +++ b/rust/src/bench/tests/cli_parity.rs @@ -0,0 +1,119 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Flag parity with the Python `vllm bench serve` parser. +//! +//! `VLLM_USE_RUST_BENCH=1` makes `vllm bench serve` exec this crate's binary +//! with the raw argv, so every documented Python flag must either parse here +//! or be a deliberate, allowlisted gap. The snapshot is kept current by +//! `tests/benchmarks/test_rust_bench_cli_parity.py` in the repo root. + +use std::collections::HashSet; + +use clap::{CommandFactory, Parser}; +use vllm_bench::BenchServeArgs; + +#[derive(Parser)] +struct TestCli { + #[command(flatten)] + args: BenchServeArgs, +} + +const PYTHON_FLAGS: &str = include_str!("python_serve_flags.txt"); + +/// Python-only flags: `vllm bench serve` features vllm-bench deliberately does +/// not implement. Adding an entry is a decision to let the flag fail under +/// Rust delegation; remove the entry once the flag gains Rust support. +const PYTHON_ONLY: &[&str] = &[ + // asr dataset + "--asr-max-audio-len-sec", + "--asr-min-audio-len-sec", + // bfcl dataset + "--bfcl-categories", + // blazedit dataset + "--blazedit-max-distance", + "--blazedit-min-distance", + // spec_bench dataset + "--spec-bench-category", + "--spec-bench-output-len", + // timed-trace dataset + "--timed-trace-chunk-hash-size", + "--timed-trace-label-hash-ids", + "--timed-trace-label-input-length", + "--timed-trace-label-output-length", + "--timed-trace-label-timestamp", + "--timed-trace-sec-multiplier", + // client-side chat templating / request shaping + "--chat-template-kwargs", + "--custom-ensure-client-side-data", + "--use-beam-search", + // result post-processing and plotting + "--plot-dataset-stats", + "--plot-timeline", + "--timeline-itl-thresholds", + // misc python-side controls + "--hf-name", + "--no-self-timed", + "--self-timed", + "--no-stream", + "--probe-request-rate", +]; + +fn rust_flags() -> HashSet { + let cmd = TestCli::command(); + let mut flags = HashSet::new(); + for arg in cmd.get_arguments() { + if let Some(long) = arg.get_long() { + flags.insert(format!("--{long}")); + } + if let Some(aliases) = arg.get_all_aliases() { + for alias in aliases { + flags.insert(format!("--{alias}")); + } + } + } + flags +} + +fn python_flags() -> Vec<&'static str> { + PYTHON_FLAGS + .lines() + .map(str::trim) + .filter(|line| !line.is_empty() && !line.starts_with('#')) + .collect() +} + +#[test] +fn python_serve_flags_parse_or_are_allowlisted() { + let known = rust_flags(); + let missing: Vec<&str> = python_flags() + .into_iter() + .filter(|flag| !known.contains(*flag) && !PYTHON_ONLY.contains(flag)) + .collect(); + assert!( + missing.is_empty(), + "Python `vllm bench serve` flags unknown to vllm-bench: {missing:?}. \ + Support them in src/cli.rs (a clap alias is enough for renames) or \ + add them to PYTHON_ONLY above as a deliberate gap." + ); +} + +#[test] +fn python_only_allowlist_is_not_stale() { + let known = rust_flags(); + let stale: Vec<&&str> = PYTHON_ONLY.iter().filter(|f| known.contains(**f)).collect(); + assert!( + stale.is_empty(), + "PYTHON_ONLY entries now supported by vllm-bench, remove them: {stale:?}" + ); +} + +#[test] +fn python_only_entries_exist_in_snapshot() { + let snapshot: HashSet<&str> = python_flags().into_iter().collect(); + let gone: Vec<&&str> = PYTHON_ONLY.iter().filter(|f| !snapshot.contains(**f)).collect(); + assert!( + gone.is_empty(), + "PYTHON_ONLY entries no longer exist in the Python parser, remove them: {gone:?}" + ); +} diff --git a/rust/src/bench/tests/python_serve_flags.txt b/rust/src/bench/tests/python_serve_flags.txt new file mode 100644 index 000000000000..d200e5bf245a --- /dev/null +++ b/rust/src/bench/tests/python_serve_flags.txt @@ -0,0 +1,110 @@ +# Long CLI flags of the Python `vllm bench serve` parser +# (vllm.benchmarks.serve.add_cli_args), sorted; --help excluded. +# Regenerate: .venv/bin/python tests/benchmarks/test_rust_bench_cli_parity.py +# Consumed by rust/src/bench/tests/cli_parity.rs, which requires every flag +# below to be a known vllm-bench flag/alias or a PYTHON_ONLY allowlist entry. +--append-result +--asr-max-audio-len-sec +--asr-min-audio-len-sec +--backend +--base-url +--bfcl-categories +--blazedit-max-distance +--blazedit-min-distance +--burstiness +--chat-template-kwargs +--custom-ensure-client-side-data +--custom-output-len +--dataset-name +--dataset-path +--disable-shuffle +--disable-tqdm +--enable-multimodal-chat +--endpoint +--extra-body +--frequency-penalty +--goodput +--header +--hf-name +--hf-output-len +--hf-split +--hf-subset +--host +--ignore-eos +--input-len +--insecure +--label +--logprobs +--lora-assignment +--lora-modules +--max-concurrency +--metadata +--metric-percentiles +--min-p +--model +--no-oversample +--no-reranker +--no-self-timed +--no-stream +--num-prompts +--num-warmups +--output-len +--percentile-metrics +--plot-dataset-stats +--plot-timeline +--port +--prefix-repetition-num-prefixes +--prefix-repetition-output-len +--prefix-repetition-prefix-len +--prefix-repetition-suffix-len +--presence-penalty +--probe-request-rate +--profile +--ramp-up-end-rps +--ramp-up-start-rps +--ramp-up-strategy +--random-batch-size +--random-input-len +--random-mm-base-items-per-request +--random-mm-bucket-config +--random-mm-limit-mm-per-prompt +--random-mm-num-mm-items-range-ratio +--random-output-len +--random-prefix-len +--random-range-ratio +--ready-check-timeout-sec +--repetition-penalty +--request-id-prefix +--request-rate +--result-dir +--result-filename +--save-detailed +--save-result +--seed +--self-timed +--served-model-name +--sharegpt-output-len +--skip-chat-template +--skip-tokenizer-init +--sonnet-input-len +--sonnet-output-len +--sonnet-prefix-len +--spec-bench-category +--spec-bench-output-len +--speed-bench-category +--speed-bench-dataset-subset +--speed-bench-output-len +--temperature +--timed-trace-chunk-hash-size +--timed-trace-label-hash-ids +--timed-trace-label-input-length +--timed-trace-label-output-length +--timed-trace-label-timestamp +--timed-trace-sec-multiplier +--timeline-itl-thresholds +--tokenizer +--tokenizer-mode +--top-k +--top-p +--trust-remote-code +--use-beam-search diff --git a/tests/benchmarks/test_rust_bench_cli_parity.py b/tests/benchmarks/test_rust_bench_cli_parity.py new file mode 100644 index 000000000000..fac03f6b4d1f --- /dev/null +++ b/tests/benchmarks/test_rust_bench_cli_parity.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Keep the `vllm bench serve` flag snapshot in sync for the Rust CLI parity test. + +`rust/src/bench/tests/cli_parity.rs` asserts that every flag in the snapshot is +either a known `vllm-bench` clap flag/alias or explicitly allowlisted as +Python-only, so `VLLM_USE_RUST_BENCH=1` delegation cannot silently reject +documented flags. This test keeps the snapshot itself current. + +Run this file directly to regenerate the snapshot: + + .venv/bin/python tests/benchmarks/test_rust_bench_cli_parity.py +""" + +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[2] +SNAPSHOT = REPO_ROOT / "rust" / "src" / "bench" / "tests" / "python_serve_flags.txt" +HEADER = """\ +# Long CLI flags of the Python `vllm bench serve` parser +# (vllm.benchmarks.serve.add_cli_args), sorted; --help excluded. +# Regenerate: .venv/bin/python tests/benchmarks/test_rust_bench_cli_parity.py +# Consumed by rust/src/bench/tests/cli_parity.rs, which requires every flag +# below to be a known vllm-bench flag/alias or a PYTHON_ONLY allowlist entry. +""" + + +def _current_flags() -> list[str]: + from vllm.benchmarks.serve import add_cli_args + from vllm.utils.argparse_utils import FlexibleArgumentParser + + parser = FlexibleArgumentParser() + add_cli_args(parser) + return sorted( + { + opt + for action in parser._actions + for opt in action.option_strings + if opt.startswith("--") and opt != "--help" + } + ) + + +def _snapshot_flags() -> list[str]: + lines = SNAPSHOT.read_text().splitlines() + return [line.strip() for line in lines if line.strip() and not line.startswith("#")] + + +def test_serve_flag_snapshot_is_current(): + assert SNAPSHOT.is_file(), f"missing snapshot {SNAPSHOT}" + assert _current_flags() == _snapshot_flags(), ( + "`vllm bench serve` flags changed. Regenerate the snapshot with " + "`.venv/bin/python tests/benchmarks/test_rust_bench_cli_parity.py`, " + "then for any added flag either support it in rust/src/bench/src/cli.rs " + "or add it to PYTHON_ONLY in rust/src/bench/tests/cli_parity.rs" + ) + + +if __name__ == "__main__": + SNAPSHOT.write_text(HEADER + "\n".join(_current_flags()) + "\n") + print(f"wrote {SNAPSHOT}") From c0a25c089a1827a4a6f304fd10896b1b79585ef1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Juan=20P=C3=A9rez=20de=20Algaba?= <124347725+jperezdealgaba@users.noreply.github.com> Date: Thu, 20 Aug 2026 11:09:49 +0200 Subject: [PATCH 194/839] [Bugfix][Security] Guard _load_ov2_processor with resolve_trust_remote_code (#52952) Signed-off-by: jperezde --- vllm/model_executor/models/llava_onevision2.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/models/llava_onevision2.py b/vllm/model_executor/models/llava_onevision2.py index c09b477f5333..f094a168fb21 100644 --- a/vllm/model_executor/models/llava_onevision2.py +++ b/vllm/model_executor/models/llava_onevision2.py @@ -39,7 +39,10 @@ import torch.nn.functional as F from PIL import Image from transformers import AutoProcessor, AutoTokenizer, BatchFeature -from transformers.dynamic_module_utils import get_class_from_dynamic_module +from transformers.dynamic_module_utils import ( + get_class_from_dynamic_module, + resolve_trust_remote_code, +) from transformers.models.qwen2_vl import Qwen2VLImageProcessor from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize @@ -129,17 +132,22 @@ def _load_ov2_processor( path = convert_model_repo_to_path(model) revision = revision or "main" + resolve_trust_remote_code( + trust_remote_code, + model, + has_local_code=False, + has_remote_code=True, + ) + processor_cls = get_class_from_dynamic_module( "processing_llava_onevision2.LlavaOnevision2Processor", path, revision=revision, - trust_remote_code=trust_remote_code, ) video_processor_cls = get_class_from_dynamic_module( "video_processing_llava_onevision2.LlavaOnevision2VideoProcessor", path, revision=revision, - trust_remote_code=trust_remote_code, ) # Slow Qwen2VLImageProcessor mirrors the remote processor (the Fast variant From 5d4d470470aed5a4a2f2be1268176c3199cc4f8e Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Thu, 20 Aug 2026 04:11:00 -0500 Subject: [PATCH 195/839] [CI] Fix nonexistent dependency for data-parallel example test selection (#53026) Signed-off-by: Taneem Ibrahim --- .buildkite/test-amd.yaml | 2 +- .buildkite/test_areas/distributed.yaml | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index e789829fe941..97281445af17 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -858,7 +858,7 @@ steps: - tests/distributed/test_torchrun_example.py - tests/distributed/test_torchrun_example_moe.py - examples/rl/ - - tests/examples/features/data_parallel/data_parallel_offline.py + - examples/features/data_parallel/data_parallel_offline.py - vllm/platforms/rocm.py commands: - torchrun --nproc-per-node=4 distributed/test_torchrun_example.py diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index 771e43bcee66..ad8e05343b5e 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -116,7 +116,7 @@ steps: - tests/distributed/test_torchrun_example.py - tests/distributed/test_torchrun_example_moe.py - examples/rl/ - - tests/examples/features/data_parallel/data_parallel_offline.py + - examples/features/data_parallel/data_parallel_offline.py commands: # https://github.com/NVIDIA/nccl/issues/1838 - export NCCL_CUMEM_HOST_ENABLE=0 @@ -265,7 +265,7 @@ steps: - vllm/executor/ - vllm/model_executor/models/ - tests/distributed/ - - tests/examples/features/data_parallel/data_parallel_offline.py + - examples/features/data_parallel/data_parallel_offline.py commands: - ./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 2 $IMAGE_TAG "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/features/data_parallel/data_parallel_offline.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=0 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_multi_node_assignment.py && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_pipeline_parallel.py" "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/features/data_parallel/data_parallel_offline.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=1 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code" From c8de519917ce549f72132952116185e38b37c95d Mon Sep 17 00:00:00 2001 From: "rongfu.leng" Date: Thu, 20 Aug 2026 17:15:47 +0800 Subject: [PATCH 196/839] [Kernel][Kimi] fused vision q/k roper kernel (#50400) Signed-off-by: rongfu.leng Co-authored-by: Isotr0py --- vllm/model_executor/models/kimi_k25_vit.py | 40 ++++++++-------------- 1 file changed, 14 insertions(+), 26 deletions(-) diff --git a/vllm/model_executor/models/kimi_k25_vit.py b/vllm/model_executor/models/kimi_k25_vit.py index 20f5b105b966..5e858a129aab 100644 --- a/vllm/model_executor/models/kimi_k25_vit.py +++ b/vllm/model_executor/models/kimi_k25_vit.py @@ -29,6 +29,7 @@ RowParallelLinear, ) from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.rotary_embedding.common import ApplyRotaryEmb from vllm.model_executor.models.utils import maybe_prefix from vllm.model_executor.models.vision import ( is_vit_use_data_parallel, @@ -79,29 +80,6 @@ def get_rope_shape(org, interpolation_mode, shape): ) -def apply_rope( - xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor -) -> tuple[torch.Tensor, torch.Tensor]: - """ - Args: (The leading dimensions of all inputs should be the same) - xq: query, tensor of shape (..., num_heads, head_dim) - xk: key, tensor of shape (..., num_heads, head_dim) - freqs_cis: tensor of shape (..., head_dim/2), dtype=torch.complex64. - Returns: - xq_out, xk_out: tensors of shape (..., num_heads, head_dim) - """ - _apply_rope_input_validation(xq, freqs_cis) - _apply_rope_input_validation(xk, freqs_cis) - - freqs_cis = freqs_cis.unsqueeze(-2) # ..., 1, head_dim/2 - # ..., num_heads, head_dim/2 - xq_ = torch.view_as_complex(xq.float().view(*xq.shape[:-1], -1, 2)) - xk_ = torch.view_as_complex(xk.float().view(*xq.shape[:-1], -1, 2)) - xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(-2) # ..., num_heads, head_dim - xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(-2) # ..., num_heads, head_dim - return xq_out.type_as(xq), xk_out.type_as(xk) - - def get_1d_sincos_pos_embed_from_grid(embed_dim, pos): """Generate 1D sincos positional embedding from grid positions.""" assert embed_dim % 2 == 0 @@ -451,12 +429,17 @@ def __init__( scale=self.hidden_size_per_attention_head**-0.5, prefix=f"{prefix}.attn", ) + self.apply_rotary_emb = ApplyRotaryEmb( + enforce_enable=True, + is_neox_style=False, + enable_fp32_compute=True, + ) def attention_qkvpacked( self, x: torch.Tensor, cu_seqlens: torch.Tensor, - rope_freqs_cis: torch.Tensor | None = None, + rope_freqs_cis: torch.Tensor, max_seqlen: torch.Tensor | None = None, sequence_lengths: torch.Tensor | None = None, ): @@ -478,7 +461,12 @@ def attention_qkvpacked( xqkv = xqkv.view(*qkv_shape) xq, xk, xv = torch.unbind(xqkv, dim=-3) - xq, xk = apply_rope(xq, xk, rope_freqs_cis) + _apply_rope_input_validation(xq, rope_freqs_cis) + _apply_rope_input_validation(xk, rope_freqs_cis) + rope_cos = rope_freqs_cis.real.contiguous() + rope_sin = rope_freqs_cis.imag.contiguous() + xq = self.apply_rotary_emb(xq, rope_cos, rope_sin) + xk = self.apply_rotary_emb(xk, rope_cos, rope_sin) if max_seqlen is None: max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max() @@ -502,7 +490,7 @@ def forward( self, hidden_states: torch.Tensor, cu_seqlens: torch.Tensor, - rope_freqs_cis: torch.Tensor | None = None, + rope_freqs_cis: torch.Tensor, max_seqlen: torch.Tensor | None = None, sequence_lengths: torch.Tensor | None = None, ): From 38e9cefdefe2e3562ceb6e8decf5101fd6d5c15c Mon Sep 17 00:00:00 2001 From: JC-ut0 <809602657@qq.com> Date: Thu, 20 Aug 2026 09:23:41 +0000 Subject: [PATCH 197/839] [Bugfix] Return HTTP 400 instead of 501 for unknown chat roles in DeepSeek encoders (#53071) Signed-off-by: JC-ut0 <809602657@qq.com> --- tests/test_request_input_bounds.py | 16 ++++++++++++++++ tests/tokenizers_/test_deepseek_v4.py | 11 +++++++++++ vllm/tokenizers/deepseek_v32_encoding.py | 2 +- vllm/tokenizers/deepseek_v4_encoding.py | 2 +- 4 files changed, 29 insertions(+), 2 deletions(-) diff --git a/tests/test_request_input_bounds.py b/tests/test_request_input_bounds.py index ecab38c922b1..395abe50c6f2 100644 --- a/tests/test_request_input_bounds.py +++ b/tests/test_request_input_bounds.py @@ -271,3 +271,19 @@ def test_encode_messages_preserves_small_chat_prompt(encoding_module): "<｜begin▁of▁sentence｜><｜User｜>Hello<｜Assistant｜>" "Hi<｜end▁of▁sentence｜>Again<｜end▁of▁sentence｜>" ) + + +@pytest.mark.parametrize( + "encoding_module", + ENCODING_MODULES, + ids=["deepseek_v32", "deepseek_v4"], +) +def test_encode_messages_unknown_role_raises_value_error(encoding_module): + # An invalid role (e.g. uppercase "SYSTEM") is a client error and must be + # raised as ValueError so the OpenAI serving layer maps it to HTTP 400 + # instead of NotImplementedError, which would map to HTTP 501. + with pytest.raises(ValueError, match="Invalid role: SYSTEM"): + encoding_module.encode_messages( + [{"role": "SYSTEM", "content": "Hello"}], + thinking_mode="chat", + ) diff --git a/tests/tokenizers_/test_deepseek_v4.py b/tests/tokenizers_/test_deepseek_v4.py index ea6b30fb70ba..fd1ba9405463 100644 --- a/tests/tokenizers_/test_deepseek_v4.py +++ b/tests/tokenizers_/test_deepseek_v4.py @@ -113,6 +113,17 @@ def test_deepseek_v4_explicitly_disables_thinking(kwargs): assert prompt == ("<｜begin▁of▁sentence｜><｜User｜>Hello<｜Assistant｜>") +def test_deepseek_v4_unknown_role_raises_value_error(): + # Invalid roles are client errors: they must surface as ValueError + # (mapped to HTTP 400 by the OpenAI serving layer), not + # NotImplementedError (mapped to HTTP 501). + with pytest.raises(ValueError, match="Invalid role: SYSTEM"): + _tokenizer().apply_chat_template( + [{"role": "SYSTEM", "content": "Hello"}], + tokenize=False, + ) + + def test_deepseek_v4_uses_v4_tool_prompt_from_request_tools(): tools = [ { diff --git a/vllm/tokenizers/deepseek_v32_encoding.py b/vllm/tokenizers/deepseek_v32_encoding.py index dfe4e478f0cb..75e84888061c 100644 --- a/vllm/tokenizers/deepseek_v32_encoding.py +++ b/vllm/tokenizers/deepseek_v32_encoding.py @@ -258,7 +258,7 @@ def render_message( tool_calls=tool_calls_content, ) else: - raise NotImplementedError(f"Unknown role: {role}") + raise ValueError(f"Invalid role: {role}") return prompt diff --git a/vllm/tokenizers/deepseek_v4_encoding.py b/vllm/tokenizers/deepseek_v4_encoding.py index 461f7f391df0..ceb702d3a078 100644 --- a/vllm/tokenizers/deepseek_v4_encoding.py +++ b/vllm/tokenizers/deepseek_v4_encoding.py @@ -353,7 +353,7 @@ def render_message( tool_calls=tc_content, ) else: - raise NotImplementedError(f"Unknown role: {role}") + raise ValueError(f"Invalid role: {role}") # Append transition tokens based on what follows if index + 1 < len(messages) and messages[index + 1].get("role") not in ["assistant", "latest_reminder"]: From 44cf3f046604188729572e22cabd747357098c2a Mon Sep 17 00:00:00 2001 From: pmanczak Date: Thu, 20 Aug 2026 11:25:41 +0200 Subject: [PATCH 198/839] [XPU][Tests] Make tests device-agnostic (#51968) Signed-off-by: pmanczak --- tests/conftest.py | 13 ++- .../attention/test_merge_attn_states.py | 88 ++++++++++--------- ...ayssm_prefill_decode_equivalence_mamba2.py | 8 +- tests/kernels/moe/test_batched_moe.py | 51 ++++++----- .../moe/test_count_expert_num_tokens.py | 12 ++- ..._mul_per_token_group_quant_fp8_colmajor.py | 7 +- .../kernels/moe/test_triton_moe_no_act_mul.py | 28 +++--- tests/kernels/moe/utils.py | 16 ++-- tests/kernels/quant_utils.py | 3 +- .../test_per_token_group_quant.py | 12 ++- .../test_fused_recurrent_packed_decode.py | 14 ++- .../fused_moe/experts/fused_batched_moe.py | 7 +- 12 files changed, 148 insertions(+), 111 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 53d448b9b530..a2d910c694e3 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -284,19 +284,18 @@ def cleanup_fixture(should_do_global_cleanup_after_test: bool): def workspace_init(): """Initialize the workspace manager for tests that need it. - This fixture initializes the workspace manager with a CUDA device - if available, and resets it after the test completes. Tests that - create a full vLLM engine should NOT use this fixture as the engine - will initialize the workspace manager itself. + This fixture initializes the workspace manager with the current + platform's accelerator device if available, and resets it after the test + completes. Tests that create a full vLLM engine should NOT use this + fixture as the engine will initialize the workspace manager itself. """ from vllm.v1.worker.workspace import ( init_workspace_manager, reset_workspace_manager, ) - if torch.cuda.is_available(): - device = torch.device("cuda:0") - init_workspace_manager(device) + if torch.accelerator.is_available(): + init_workspace_manager(torch.device(0)) yield reset_workspace_manager() diff --git a/tests/kernels/attention/test_merge_attn_states.py b/tests/kernels/attention/test_merge_attn_states.py index eeadf0da0bfa..1b89e4db66b3 100644 --- a/tests/kernels/attention/test_merge_attn_states.py +++ b/tests/kernels/attention/test_merge_attn_states.py @@ -5,7 +5,7 @@ import torch from vllm._custom_ops import ( - merge_attn_states as merge_attn_states_cuda, + merge_attn_states as merge_attn_states_native, ) from vllm._custom_ops import ( scaled_fp8_quant, @@ -15,11 +15,13 @@ merge_attn_states as merge_attn_states_triton, ) +DEVICE = current_platform.device_type + pytestmark = [ pytest.mark.skip_global_cleanup, pytest.mark.skipif( - not current_platform.is_cuda_alike(), - reason="merge_attn_states kernels require CUDA or ROCm.", + not (current_platform.is_cuda_alike() or current_platform.is_xpu()), + reason="merge_attn_states kernels require CUDA, ROCm or XPU.", ), ] @@ -81,20 +83,22 @@ def merge_attn_states_torch( all_case_info: list[tuple] = [] -@pytest.mark.parametrize("merge_fn", [merge_attn_states_cuda, merge_attn_states_triton]) +@pytest.mark.parametrize( + "merge_fn", [merge_attn_states_native, merge_attn_states_triton] +) @pytest.mark.parametrize("output_dtype", [torch.float32, torch.half, torch.bfloat16]) def test_merge_attn_states_both_empty(merge_fn, output_dtype) -> None: """When a token is empty on both sides (both LSE -inf), the 0/0 softmax scales must not surface as NaN in the merged output.""" num_tokens, num_heads, head_size = 6, 8, 128 prefix_output = torch.zeros( - num_tokens, num_heads, head_size, device="cuda", dtype=output_dtype + num_tokens, num_heads, head_size, device=DEVICE, dtype=output_dtype ) - prefix_lse = torch.randn(num_heads, num_tokens, device="cuda") + prefix_lse = torch.randn(num_heads, num_tokens, device=DEVICE) suffix_output = torch.zeros( - num_tokens, num_heads, head_size, device="cuda", dtype=output_dtype + num_tokens, num_heads, head_size, device=DEVICE, dtype=output_dtype ) - suffix_lse = torch.randn(num_heads, num_tokens, device="cuda") + suffix_lse = torch.randn(num_heads, num_tokens, device=DEVICE) # Tokens 2 and 3 are empty on both sides. empty = slice(2, 4) @@ -111,7 +115,7 @@ def generate_markdown_table(): global all_case_info table_header = ( "| tokens | heads | headsize | dtype " - "| device | torch | triton | cuda | speedup |" + "| device | torch | triton | native | speedup |" ) table_separator = "| --- | --- | --- | --- | --- | --- | --- | --- | --- |" @@ -132,7 +136,7 @@ def shortly_device(device: str) -> str: device, avg_time_torch_kernel, avg_time_triton_kernel, - avg_time_cuda_kernel, + avg_time_native_kernel, performance_improved, ) = info dtype = shortly_dtype(dtype) @@ -141,7 +145,7 @@ def shortly_device(device: str) -> str: f"| {num_tokens} | {num_heads} | {head_size} " f"| {dtype} | {device} | {avg_time_torch_kernel:.5f}ms " f"| {avg_time_triton_kernel:.5f}ms " - f"| {avg_time_cuda_kernel:.5f}ms " + f"| {avg_time_native_kernel:.5f}ms " f"| {performance_improved:.4f}x |" ) @@ -171,7 +175,7 @@ def test_merge_attn_states( output_scale = None if use_fp8: output_dtype = current_platform.fp8_dtype() - output_scale = torch.tensor([0.05], dtype=torch.float32, device="cuda") + output_scale = torch.tensor([0.05], dtype=torch.float32, device=DEVICE) print( f"\nNUM_TOKENS:{NUM_TOKENS}, NUM_HEADS:{NUM_HEADS}, " @@ -182,8 +186,8 @@ def test_merge_attn_states( ) # prefix_lse and suffix_lse contain inf and normal values - prefix_lse = torch.randn(NUM_HEADS, NUM_TOKENS, dtype=torch.float32, device="cuda") - suffix_lse = torch.randn(NUM_HEADS, NUM_TOKENS, dtype=torch.float32, device="cuda") + prefix_lse = torch.randn(NUM_HEADS, NUM_TOKENS, dtype=torch.float32, device=DEVICE) + suffix_lse = torch.randn(NUM_HEADS, NUM_TOKENS, dtype=torch.float32, device=DEVICE) # Generate boolean masks mask_prefix = torch.rand(NUM_HEADS, NUM_TOKENS) < 0.1 @@ -199,16 +203,16 @@ def test_merge_attn_states( # Other input tensors (need to be initialized but # no actual calculation needed) output = torch.zeros( - (NUM_TOKENS, NUM_HEADS, HEAD_SIZE), dtype=output_dtype, device="cuda" + (NUM_TOKENS, NUM_HEADS, HEAD_SIZE), dtype=output_dtype, device=DEVICE ) output_lse = torch.zeros( - (NUM_HEADS, NUM_TOKENS), dtype=torch.float32, device="cuda" + (NUM_HEADS, NUM_TOKENS), dtype=torch.float32, device=DEVICE ) prefix_output = torch.randn( - (NUM_TOKENS, NUM_HEADS, HEAD_SIZE), dtype=input_dtype, device="cuda" + (NUM_TOKENS, NUM_HEADS, HEAD_SIZE), dtype=input_dtype, device=DEVICE ) suffix_output = torch.randn( - (NUM_TOKENS, NUM_HEADS, HEAD_SIZE), dtype=input_dtype, device="cuda" + (NUM_TOKENS, NUM_HEADS, HEAD_SIZE), dtype=input_dtype, device=DEVICE ) warmup_times = 2 @@ -293,19 +297,19 @@ def test_merge_attn_states( avg_time_triton_kernel = total_time_triton_kernel / repeat_times - # 2. Run the CUDA kernel - total_time_cuda_kernel = 0 - output_cuda = output.clone() - output_lse_cuda = output_lse.clone() + # 2. Run the native (CUDA/XPU) kernel + total_time_native_kernel = 0 + output_native = output.clone() + output_lse_native = output_lse.clone() for _ in range(warmup_times): - merge_attn_states_cuda( - output_cuda, + merge_attn_states_native( + output_native, prefix_output, prefix_lse, suffix_output, suffix_lse, - output_lse_cuda, + output_lse_native, prefill_tokens_with_context, output_scale, ) @@ -313,28 +317,28 @@ def test_merge_attn_states( for _ in range(repeat_times): start.record() - merge_attn_states_cuda( - output_cuda, + merge_attn_states_native( + output_native, prefix_output, prefix_lse, suffix_output, suffix_lse, - output_lse_cuda, + output_lse_native, prefill_tokens_with_context, output_scale, ) end.record() torch.accelerator.synchronize() - total_time_cuda_kernel += start.elapsed_time(end) + total_time_native_kernel += start.elapsed_time(end) - avg_time_cuda_kernel = total_time_cuda_kernel / repeat_times + avg_time_native_kernel = total_time_native_kernel / repeat_times # 3. Performance compare - performance_improved = avg_time_triton_kernel / avg_time_cuda_kernel + performance_improved = avg_time_triton_kernel / avg_time_native_kernel print(f" Torch time: {avg_time_torch_kernel:.6f}ms") print(f"Triton time: {avg_time_triton_kernel:.6f}ms") print( - f" CUDA time: {avg_time_cuda_kernel:.6f}ms, " + f"Native time: {avg_time_native_kernel:.6f}ms, " f"Performance: {performance_improved:.5f}x" ) print("-" * 100) @@ -362,12 +366,12 @@ def diff(a: torch.Tensor, b: torch.Tensor): return max_diff # Use Triton output as reference because we want to replace - # the Triton kernel with custom CUDA kernel for merge attn + # the Triton kernel with the custom native kernel for merge attn # states operation. output_ref = output_ref_triton output_lse_ref = output_lse_ref_triton torch.testing.assert_close( - output_cuda.float() * scale, + output_native.float() * scale, output_ref.float() * scale, atol=atol, rtol=rtol, @@ -379,19 +383,19 @@ def diff(a: torch.Tensor, b: torch.Tensor): ) _diff = diff(output_ref.float() * scale, output_torch.float() * scale) print(f"(Triton vs Torch) : {_diff}") - _diff = diff(output_torch.float() * scale, output_cuda.float() * scale) - print(f" (CUDA vs Torch) : {_diff}") - _diff = diff(output_ref.float() * scale, output_cuda.float() * scale) - print(f" (CUDA vs Triton): {_diff}") + _diff = diff(output_torch.float() * scale, output_native.float() * scale) + print(f"(Native vs Torch) : {_diff}") + _diff = diff(output_ref.float() * scale, output_native.float() * scale) + print(f"(Native vs Triton): {_diff}") print("-" * 100) torch.testing.assert_close( - output_lse_cuda.float(), output_lse_ref.float(), atol=atol, rtol=rtol + output_lse_native.float(), output_lse_ref.float(), atol=atol, rtol=rtol ) print("Output LSE all match, max abs diff:") print(f"(Triton vs Torch) : {diff(output_lse_torch, output_lse_ref)}") - print(f" (CUDA vs Torch) : {diff(output_lse_torch, output_lse_cuda)}") - print(f" (CUDA vs Triton): {diff(output_lse_ref, output_lse_cuda)}") + print(f"(Native vs Torch) : {diff(output_lse_torch, output_lse_native)}") + print(f"(Native vs Triton): {diff(output_lse_ref, output_lse_native)}") print("-" * 100) print( @@ -410,7 +414,7 @@ def diff(a: torch.Tensor, b: torch.Tensor): device, avg_time_torch_kernel, avg_time_triton_kernel, - avg_time_cuda_kernel, + avg_time_native_kernel, performance_improved, ) ) diff --git a/tests/kernels/mamba/test_replayssm_prefill_decode_equivalence_mamba2.py b/tests/kernels/mamba/test_replayssm_prefill_decode_equivalence_mamba2.py index 5319d8964d34..063ff17dd337 100644 --- a/tests/kernels/mamba/test_replayssm_prefill_decode_equivalence_mamba2.py +++ b/tests/kernels/mamba/test_replayssm_prefill_decode_equivalence_mamba2.py @@ -36,6 +36,7 @@ from vllm.model_executor.layers.mamba.ops.ssd_combined import ( mamba_chunk_scan_combined_varlen, ) +from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed from vllm.v1.attention.backends.mamba2_attn import compute_varlen_chunk_metadata @@ -68,7 +69,7 @@ def _run_prefill_decode_equivalence( flow (raw dt + softplus + a per-head dt_bias), so this also checks prefill and decode apply the softplus/bias preprocessing consistently. ``state_dtype`` is the recurrent-state precision; ``act_dtype`` the activation/buffer one.""" - device = "cuda" + device = current_platform.device_type rtol, atol = _prefill_tolerances(act_dtype) set_random_seed(seed) @@ -203,7 +204,10 @@ def _run_prefill_decode_equivalence( ) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") +@pytest.mark.skipif( + not (current_platform.is_cuda_alike() or current_platform.is_xpu()), + reason="Mamba2 ReplaySSM kernels require a CUDA-alike or XPU device.", +) @pytest.mark.parametrize( "precision", # fp32 state is the default; bf16/fp16 are reduced-footprint configs. fp16 diff --git a/tests/kernels/moe/test_batched_moe.py b/tests/kernels/moe/test_batched_moe.py index 3c605ced688c..19375e6ce960 100644 --- a/tests/kernels/moe/test_batched_moe.py +++ b/tests/kernels/moe/test_batched_moe.py @@ -23,6 +23,13 @@ from vllm.triton_utils import tl from vllm.utils.torch_utils import set_random_seed +DEVICE = current_platform.device_type + +pytestmark = pytest.mark.skipif( + not (current_platform.is_cuda_alike() or current_platform.is_xpu()), + reason="Triton MoE kernels require CUDA/ROCm/XPU.", +) + MNK_FACTORS = [ (1, 128, 128), (1, 512, 512), @@ -71,19 +78,19 @@ def make_tensors(config: BatchedMMConfig): A = ( torch.randn( (config.num_experts, config.max_tokens_per_expert, config.K), - device="cuda", + device=DEVICE, dtype=config.in_dtype, ) / 10 ) B = torch.randn( (config.num_experts, config.N, config.K), - device="cuda", + device=DEVICE, dtype=config.in_dtype, ) C = torch.zeros( (config.num_experts, config.max_tokens_per_expert, config.N), - device="cuda", + device=DEVICE, dtype=config.out_dtype, ) @@ -91,7 +98,7 @@ def make_tensors(config: BatchedMMConfig): low=0, high=config.max_tokens_per_expert, size=(config.num_experts,), - device="cuda", + device=DEVICE, dtype=torch.int32, ) @@ -120,8 +127,10 @@ def test_batched_mm( use_fp8_w8a8 = dtype == torch.float8_e4m3fn - if (dtype == torch.float8_e4m3fn) and not current_platform.has_device_capability( - 89 + if ( + dtype == torch.float8_e4m3fn + and current_platform.is_cuda_alike() + and not current_platform.has_device_capability(89) ): pytest.skip( "Triton limitation: fp8e4nv data type is not supported on CUDA arch < 89" @@ -144,7 +153,7 @@ def test_batched_mm( low=0, high=max_tokens_per_expert, size=(num_experts,), - device="cuda", + device=DEVICE, dtype=torch.int32, ) @@ -169,9 +178,9 @@ def test_batched_mm( ) out_shape = (num_experts, max_tokens_per_expert, N) - test_output = torch.zeros(out_shape, dtype=act_dtype, device="cuda") - ref_output = torch.zeros(out_shape, dtype=act_dtype, device="cuda") - q_ref_output = torch.zeros(out_shape, dtype=act_dtype, device="cuda") + test_output = torch.zeros(out_shape, dtype=act_dtype, device=DEVICE) + ref_output = torch.zeros(out_shape, dtype=act_dtype, device=DEVICE) + q_ref_output = torch.zeros(out_shape, dtype=act_dtype, device=DEVICE) compute_tl_dtype = { torch.float16: tl.float16, @@ -257,8 +266,10 @@ def test_fused_moe_batched_experts( use_fp8_w8a8 = dtype == torch.float8_e4m3fn - if (dtype == torch.float8_e4m3fn) and not current_platform.has_device_capability( - 89 + if ( + dtype == torch.float8_e4m3fn + and current_platform.is_cuda_alike() + and not current_platform.has_device_capability(89) ): pytest.skip( "Triton limitation: fp8e4nv data type is not supported on CUDA arch < 89" @@ -273,8 +284,8 @@ def test_fused_moe_batched_experts( if per_act_token_quant and block_shape is not None: pytest.skip("Skip illegal quantization test.") - a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10 - score = torch.randn((m, e), device="cuda", dtype=torch.bfloat16) + a = torch.randn((m, k), device=DEVICE, dtype=torch.bfloat16) / 10 + score = torch.randn((m, e), device=DEVICE, dtype=torch.bfloat16) if dtype.itemsize == 1: act_dtype = torch.bfloat16 @@ -294,8 +305,8 @@ def test_fused_moe_batched_experts( ) if input_scales and quant_dtype is not None: - a1_scale = torch.tensor(1, device="cuda", dtype=torch.float32) - a2_scale = torch.tensor(1, device="cuda", dtype=torch.float32) + a1_scale = torch.tensor(1, device=DEVICE, dtype=torch.float32) + a2_scale = torch.tensor(1, device=DEVICE, dtype=torch.float32) else: a1_scale = None a2_scale = None @@ -489,10 +500,7 @@ def test_batched_triton_experts_supports_current_device(): BatchedTritonExperts, ) - if current_platform.is_xpu() or current_platform.is_cuda_alike(): - assert BatchedTritonExperts._supports_current_device() - else: - pytest.skip("No GPU device available") + assert BatchedTritonExperts._supports_current_device() @pytest.mark.parametrize("m,n,k,e,topk", [(32, 512, 512, 8, 2), (45, 1024, 128, 8, 1)]) @@ -500,9 +508,6 @@ def test_batched_experts_end_to_end(m, n, k, e, topk): """End-to-end BatchedTritonExperts via the reference (no-comms) BatchedPrepareAndFinalize, validated against a torch reference. Exercises the device-enablement path.""" - if not (current_platform.is_xpu() or current_platform.is_cuda_alike()): - pytest.skip("No GPU device available") - from vllm.v1.worker.workspace import init_workspace_manager set_random_seed(7) diff --git a/tests/kernels/moe/test_count_expert_num_tokens.py b/tests/kernels/moe/test_count_expert_num_tokens.py index 39138be83bcc..8014076192e2 100644 --- a/tests/kernels/moe/test_count_expert_num_tokens.py +++ b/tests/kernels/moe/test_count_expert_num_tokens.py @@ -10,6 +10,14 @@ import torch from vllm.model_executor.layers.fused_moe.utils import count_expert_num_tokens +from vllm.platforms import current_platform + +DEVICE = current_platform.device_type + +pytestmark = pytest.mark.skipif( + not (current_platform.is_cuda_alike() or current_platform.is_xpu()), + reason="Triton MoE kernels require CUDA/ROCm/XPU.", +) @dataclasses.dataclass @@ -88,9 +96,9 @@ def do_test_compute_expert_num_tokens( (num_local_experts), device="cpu", dtype=torch.int32 ) ref_impl(tt_rank, ref_expert_num_tokens) - ref_expert_num_tokens = ref_expert_num_tokens.to("cuda") + ref_expert_num_tokens = ref_expert_num_tokens.to(DEVICE) - tt_rank.to_device("cuda") + tt_rank.to_device(DEVICE) # Test with expert_map triton_expert_num_tokens_w_emap = count_expert_num_tokens( tt_rank.topk_ids, num_local_experts, tt_rank.expert_map diff --git a/tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py b/tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py index cb01db44f8cb..bfe763b44ab2 100644 --- a/tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py +++ b/tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py @@ -13,6 +13,7 @@ from vllm.utils.deep_gemm import is_deep_gemm_e8m0_used from vllm.utils.torch_utils import set_random_seed +DEVICE = current_platform.device_type FLOAT8_DTYPE = torch.float8_e4m3fn GROUP_SIZE = 128 @@ -61,7 +62,7 @@ def reference_quant(x: torch.Tensor, use_ue8m0: bool): def reference(x: torch.Tensor, use_ue8m0: bool) -> tuple[torch.Tensor, torch.Tensor]: T, N = x.size() - ref_act_out = torch.empty((T, N // 2), dtype=torch.bfloat16, device="cuda") + ref_act_out = torch.empty((T, N // 2), dtype=torch.bfloat16, device=DEVICE) torch.ops._C.silu_and_mul(ref_act_out, x) return reference_quant(ref_act_out, use_ue8m0) @@ -93,7 +94,7 @@ def reference_with_clamp( def test_silu_mul_fp8_quant_deep_gemm(T: int, N: int): set_random_seed(42) - input = torch.rand((T, N), dtype=torch.bfloat16, device="cuda") + input = torch.rand((T, N), dtype=torch.bfloat16, device=DEVICE) use_ue8m0 = is_deep_gemm_e8m0_used() @@ -122,7 +123,7 @@ def test_silu_mul_fp8_quant_deep_gemm_clamp(T: int, N: int, clamp_limit: float): # Use a wide distribution so values routinely exceed both clamp limits and # the clamp branch is actually exercised (uniform [0, 1) inputs would never # trigger it). - input = torch.randn((T, N), dtype=torch.bfloat16, device="cuda") * 8.0 + input = torch.randn((T, N), dtype=torch.bfloat16, device=DEVICE) * 8.0 use_ue8m0 = is_deep_gemm_e8m0_used() diff --git a/tests/kernels/moe/test_triton_moe_no_act_mul.py b/tests/kernels/moe/test_triton_moe_no_act_mul.py index e9f6f51c220c..1171061a6f43 100644 --- a/tests/kernels/moe/test_triton_moe_no_act_mul.py +++ b/tests/kernels/moe/test_triton_moe_no_act_mul.py @@ -23,6 +23,20 @@ from vllm.model_executor.layers.fused_moe.experts.triton_moe import TritonExperts from vllm.platforms import current_platform +DEVICE = current_platform.device_type + +pytestmark = [ + pytest.mark.skipif( + not (current_platform.is_cuda_alike() or current_platform.is_xpu()), + reason="Triton MoE kernels require a CUDA-alike or XPU device.", + ), + pytest.mark.skipif( + current_platform.is_cuda_alike() + and not current_platform.has_device_capability(80), + reason="Triton MoE kernels require compute capability >= 8.0 on CUDA/ROCm.", + ), +] + # Test parameters M_SIZES = [1, 16, 64] N_SIZES = [128, 256] @@ -86,7 +100,7 @@ def make_test_tensors( num_experts: int, topk: int, dtype: torch.dtype = torch.bfloat16, - device: str = "cuda", + device: str = DEVICE, ): """Create test tensors for MoE with non-gated activation. @@ -106,10 +120,6 @@ def make_test_tensors( return hidden_states, w1, w2, topk_weights, topk_ids -@pytest.mark.skipif( - not current_platform.has_device_capability(80), - reason="Requires compute capability >= 8.0", -) @pytest.mark.parametrize("m", M_SIZES) @pytest.mark.parametrize("n", N_SIZES) @pytest.mark.parametrize("k", K_SIZES) @@ -192,10 +202,6 @@ def test_triton_experts_no_mul_activation( assert output.abs().sum() > 0, "Output is all zeros" -@pytest.mark.skipif( - not current_platform.has_device_capability(80), - reason="Requires compute capability >= 8.0", -) @torch.inference_mode() def test_workspace_shapes_no_mul_vs_gated(): """Test that workspace shapes differ correctly between gated and non-gated.""" @@ -233,10 +239,6 @@ def test_workspace_shapes_no_mul_vs_gated(): assert out_no_mul == out_gated == (M, K) -@pytest.mark.skipif( - not current_platform.has_device_capability(80), - reason="Requires compute capability >= 8.0", -) @torch.inference_mode() def test_adjust_n_for_activation(): """Test the adjust_N_for_activation method.""" diff --git a/tests/kernels/moe/utils.py b/tests/kernels/moe/utils.py index ce6e07f122ff..e82d31571542 100644 --- a/tests/kernels/moe/utils.py +++ b/tests/kernels/moe/utils.py @@ -40,6 +40,8 @@ from vllm.utils.deep_gemm import per_block_cast_to_fp8 from vllm.utils.math_utils import round_up +DEVICE = current_platform.device_type + def shuffle_weight(w: torch.Tensor) -> torch.Tensor: """Fold weights to adjacent locations for Triton MoE / SwiGLU kernel layout.""" @@ -80,7 +82,7 @@ def make_dummy_moe_config( moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), activation=activation, in_dtype=in_dtype, - device="cuda", + device=DEVICE, routing_method=RoutingMethodType.TopK, max_num_tokens=max_num_tokens, ) @@ -238,7 +240,7 @@ def make_quantized_test_activations( block_shape: list[int] | None = None, per_act_token_quant: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: - a = torch.randn((E, m, k), device="cuda", dtype=in_dtype) / 10 + a = torch.randn((E, m, k), device=DEVICE, dtype=in_dtype) / 10 a_q = a a_scale = None @@ -369,7 +371,7 @@ def make_test_weight( block_shape: list[int] | None = None, per_out_ch_quant: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]: - w_16 = torch.randn((e, rows, cols), device="cuda", dtype=in_dtype) / 15 + w_16 = torch.randn((e, rows, cols), device=DEVICE, dtype=in_dtype) / 15 if quant_dtype is not None: w, w_s, w_gs = moe_quantize_weights( @@ -449,8 +451,8 @@ def make_test_quant_config( a1_gscale: torch.Tensor | None = None a2_gscale: torch.Tensor | None = None if quant_dtype == "nvfp4": - a1_gscale = torch.ones((e,), device="cuda", dtype=torch.float32) - a2_gscale = torch.ones((e,), device="cuda", dtype=torch.float32) + a1_gscale = torch.ones((e,), device=DEVICE, dtype=torch.float32) + a2_gscale = torch.ones((e,), device=DEVICE, dtype=torch.float32) a1_scale = a1_gscale a2_scale = a2_gscale else: @@ -551,8 +553,8 @@ def make_naive_shared_experts( K: int, in_dtype: torch.dtype = torch.bfloat16, ) -> torch.nn.Module: - w1 = torch.randn((K, N * 2), device="cuda", dtype=in_dtype) / 15 - w2 = torch.randn((N, K), device="cuda", dtype=in_dtype) / 15 + w1 = torch.randn((K, N * 2), device=DEVICE, dtype=in_dtype) / 15 + w2 = torch.randn((N, K), device=DEVICE, dtype=in_dtype) / 15 return TestMLP(w1, w2, out_dtype=in_dtype) diff --git a/tests/kernels/quant_utils.py b/tests/kernels/quant_utils.py index a67cb8fd78d8..7e14ec95d367 100644 --- a/tests/kernels/quant_utils.py +++ b/tests/kernels/quant_utils.py @@ -13,10 +13,11 @@ from vllm.utils.math_utils import round_up FP8_DTYPE = current_platform.fp8_dtype() +DEVICE = current_platform.device_type def as_float32_tensor(x: float | torch.Tensor) -> torch.Tensor: - return torch.as_tensor(x, dtype=torch.float32, device="cuda") + return torch.as_tensor(x, dtype=torch.float32, device=DEVICE) def ref_dynamic_per_token_quant( diff --git a/tests/kernels/quantization/test_per_token_group_quant.py b/tests/kernels/quantization/test_per_token_group_quant.py index 0d9b6c0c3e8f..623c9b6f1364 100644 --- a/tests/kernels/quantization/test_per_token_group_quant.py +++ b/tests/kernels/quantization/test_per_token_group_quant.py @@ -18,19 +18,20 @@ @pytest.mark.parametrize("scale_ue8m0", [False, True]) @pytest.mark.parametrize("group_size", [64, 128]) @pytest.mark.skipif( - not current_platform.is_cuda_alike(), reason="Only test on CUDA/ROCm." + not (current_platform.is_cuda_alike() or current_platform.is_xpu()), + reason="Only test on CUDA/ROCm/XPU.", ) def test_per_token_group_quant_fp8( shape, column_major: bool, tma_aligned: bool, scale_ue8m0: bool, group_size: int ): - device = "cuda" + device = current_platform.device_type torch.manual_seed(42) num_tokens, hidden_dim = shape x = torch.randn((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) * 8 - # cuda path + # native kernel path out_q, scale = fp8_utils.per_token_group_quant_fp8( x, group_size, @@ -40,7 +41,10 @@ def test_per_token_group_quant_fp8( ) # triton ref - with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False): + with ( + patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False), + patch("vllm.platforms.current_platform.is_xpu", return_value=False), + ): ref_q, ref_s = fp8_utils.per_token_group_quant_fp8( x, group_size, diff --git a/tests/kernels/test_fused_recurrent_packed_decode.py b/tests/kernels/test_fused_recurrent_packed_decode.py index 33a66034b422..5e83537d945a 100644 --- a/tests/kernels/test_fused_recurrent_packed_decode.py +++ b/tests/kernels/test_fused_recurrent_packed_decode.py @@ -4,13 +4,20 @@ import pytest import torch +from vllm.platforms import current_platform from vllm.third_party.flash_linear_attention.ops import ( fused_recurrent_gated_delta_rule, fused_recurrent_gated_delta_rule_packed_decode, ) +DEVICE = current_platform.device_type + +pytestmark = pytest.mark.skipif( + not (current_platform.is_cuda_alike() or current_platform.is_xpu()), + reason="Gated delta rule Triton kernels require a CUDA-alike or XPU device.", +) + -@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) @pytest.mark.parametrize("strided_mixed_qkv", [False, True]) def test_fused_recurrent_packed_decode_matches_reference( @@ -26,7 +33,7 @@ def test_fused_recurrent_packed_decode_matches_reference( V = 128 qkv_dim = 2 * (H * K) + (HV * V) - device = torch.device("cuda") + device = torch.device(DEVICE) if strided_mixed_qkv: # Simulate a packed view into a larger projection buffer: @@ -102,10 +109,9 @@ def test_fused_recurrent_packed_decode_matches_reference( torch.testing.assert_close(state_packed, state_ref, rtol=rtol, atol=atol) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") def test_packed_decode_supports_large_batch_head_grid(): B, H, HV, K, V = 1024, 8, 64, 1, 1 - device = torch.device("cuda") + device = torch.device(DEVICE) gates = torch.empty((B, HV), device=device) params = torch.empty((HV,), device=device) out = torch.empty((B, 1, HV, V), device=device) diff --git a/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py b/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py index d35c3d0d71ad..40d17b14dcab 100644 --- a/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py @@ -37,9 +37,10 @@ def _is_capturing_or_compiling() -> bool: - # torch.cuda.is_current_stream_capturing() is unavailable on non-CUDA (XPU) torch. - return torch.compiler.is_compiling() or ( - current_platform.is_cuda_alike() and torch.cuda.is_current_stream_capturing() + # The accelerator API dispatches to the active backend. + return ( + torch.compiler.is_compiling() + or torch.accelerator.current_stream().is_capturing() ) From 6b68db441e7689676d52a64400f367f06276f2bb Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Thu, 20 Aug 2026 04:34:00 -0500 Subject: [PATCH 199/839] [ROCm] Fix DeepSeek V4 indexer numerics and coverage (#50803) Signed-off-by: Andreas Karatzas --- tests/kernels/test_compressor_kv_cache.py | 17 ++++++++++++++++- .../kernels/test_fused_indexer_q_rope_quant.py | 18 +++++++++++++++++- .../deepseek_v4/common/ops/fused_indexer_q.py | 11 +++++++++-- 3 files changed, 42 insertions(+), 4 deletions(-) diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index b300d0ebeccb..c1049eb65811 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -812,7 +812,22 @@ def _reference_kv_compress_norm_rope( @pytest.mark.parametrize("num_tokens", [1, 7, 32]) @pytest.mark.parametrize("kv_block_size", [16, 32]) -@pytest.mark.parametrize("use_fp4", [False, True]) +@pytest.mark.parametrize( + "use_fp4", + [ + False, + pytest.param( + True, + marks=pytest.mark.skipif( + not ( + current_platform.is_cuda() + and current_platform.is_device_capability_family(100) + ), + reason="MXFP4 indexer cache requires an SM100-family GPU", + ), + ), + ], +) def test_fused_kv_insert_indexer(num_tokens: int, kv_block_size: int, use_fp4: bool): """Fused K compress+norm+rope+quant+insert for the indexer KV cache.""" HEAD_DIM = 128 diff --git a/tests/kernels/test_fused_indexer_q_rope_quant.py b/tests/kernels/test_fused_indexer_q_rope_quant.py index 6114b7efd6e7..5b568c938eec 100644 --- a/tests/kernels/test_fused_indexer_q_rope_quant.py +++ b/tests/kernels/test_fused_indexer_q_rope_quant.py @@ -24,6 +24,7 @@ per_token_group_quant_fp8, ) from vllm.models.deepseek_v4.common.ops import fused_indexer_q_rope_quant +from vllm.platforms import current_platform from vllm.utils.import_utils import has_cutedsl HEAD_DIM = 128 @@ -126,7 +127,22 @@ def _reference( @pytest.mark.parametrize("num_tokens", [1, 7, 32, 257, 1023]) @pytest.mark.parametrize("cache_dtype", [torch.float32, torch.bfloat16]) -@pytest.mark.parametrize("use_fp4", [False, True]) +@pytest.mark.parametrize( + "use_fp4", + [ + False, + pytest.param( + True, + marks=pytest.mark.skipif( + not ( + current_platform.is_cuda() + and current_platform.is_device_capability_family(100) + ), + reason="MXFP4 indexer cache requires an SM100-family GPU", + ), + ), + ], +) @pytest.mark.parametrize("use_cutedsl", [False, True]) @torch.inference_mode() def test_fused_indexer_q_rope_quant_matches_unfused( diff --git a/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py b/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py index 3ec0c2b29d5d..9ae34de1a0a6 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py +++ b/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py @@ -91,6 +91,7 @@ def _fused_indexer_q_rope_quant_kernel( index_weights_out_stride, FP8_MAX: tl.constexpr = 448.0, USE_FNUZ: tl.constexpr = False, + USE_EXPLICIT_FMA: tl.constexpr = False, ): # Layout matches the unfused reference (DeepseekV4ScalingRotaryEmbedding # + per_token_group_quant_fp8): GPT-J interleaved RoPE applied to the @@ -118,8 +119,13 @@ def _fused_indexer_q_rope_quant_kernel( rot_base = base_ptr + INDEX_Q_NOPE_DIM x_even = tl.load(rot_base + half_offset * 2).to(tl.float32) x_odd = tl.load(rot_base + half_offset * 2 + 1).to(tl.float32) - r_even = x_even * cos - x_odd * sin - r_odd = x_odd * cos + x_even * sin + if USE_EXPLICIT_FMA: + # Match HIP rotary_embedding contraction before bf16 materialization. + r_even = tl.fma(x_even, cos, -(x_odd * sin)) + r_odd = tl.fma(x_odd, cos, x_even * sin) + else: + r_even = x_even * cos - x_odd * sin + r_odd = x_odd * cos + x_even * sin # Match reference numerics: fp32 → bf16 → fp32 before the ue8m0 absmax. # Same pattern as the K-side compressor kernel (fused_compress_quant_cache.py). @@ -467,6 +473,7 @@ def fused_indexer_q_rope_quant( index_weights_out.stride(0), FP8_MAX=fp8_max, USE_FNUZ=use_fnuz, + USE_EXPLICIT_FMA=current_platform.is_rocm(), num_warps=1, # TODO: Tune this ) return index_q_fp8, index_weights_out From 30e2394c83afe066e91fcc709ac47935a57bbe3c Mon Sep 17 00:00:00 2001 From: Eoin-Houstoun Date: Thu, 20 Aug 2026 10:36:00 +0100 Subject: [PATCH 200/839] [Bugfix] Record non-ImportError attention backend probe failures instead of crashing engine init (#51703) Signed-off-by: Eoin Co-authored-by: Claude Fable 5 --- .../test_cuda_backend_probe_errors.py | 104 ++++++++++++++++++ vllm/platforms/cuda.py | 14 ++- 2 files changed, 114 insertions(+), 4 deletions(-) create mode 100644 tests/v1/attention/test_cuda_backend_probe_errors.py diff --git a/tests/v1/attention/test_cuda_backend_probe_errors.py b/tests/v1/attention/test_cuda_backend_probe_errors.py new file mode 100644 index 000000000000..dd05a0d4ba5f --- /dev/null +++ b/tests/v1/attention/test_cuda_backend_probe_errors.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for error handling in the CUDA attention backend probe. + +Environment-shaped probe failures (missing packages, unreadable caches, +broken driver installs) must mark the backend unavailable rather than +terminating engine init; programming errors must still propagate. +See https://github.com/vllm-project/vllm/issues/51658. +""" + +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from vllm.platforms import current_platform +from vllm.platforms.cuda import CudaPlatform +from vllm.platforms.interface import DeviceCapability +from vllm.v1.attention.backends.registry import AttentionBackendEnum +from vllm.v1.attention.selector import AttentionSelectorConfig + +pytestmark = pytest.mark.skipif( + not current_platform.is_cuda(), reason="CUDA-specific tests" +) + +SELECTOR_CONFIG = AttentionSelectorConfig( + head_size=64, + dtype=torch.float16, + kv_cache_dtype=None, + block_size=16, +) + +SM90 = DeviceCapability(major=9, minor=0) + + +@pytest.mark.parametrize( + "exc", + [ + ImportError("flashinfer is not installed"), + PermissionError(13, "Permission denied", "/root/.cache/flashinfer"), + OSError("libcuda.so.1: cannot open shared object file"), + ], +) +def test_get_valid_backends_records_environment_failure(exc): + with patch("vllm.platforms.cuda._get_attn_backend_class", side_effect=exc): + valid, invalid_reasons = CudaPlatform.get_valid_backends( + device_capability=SM90, + attn_selector_config=SELECTOR_CONFIG, + num_heads=32, + ) + assert valid == [] + assert invalid_reasons + for _priority, reasons in invalid_reasons.values(): + assert reasons == [f"{type(exc).__name__}: {exc}"] + + +@pytest.mark.parametrize( + "exc", + [ + RuntimeError("CUDA error: no kernel image is available"), + AttributeError("module has no attribute 'get_builder_cls'"), + KeyboardInterrupt(), + ], +) +def test_get_valid_backends_propagates_unexpected_errors(exc): + with ( + patch("vllm.platforms.cuda._get_attn_backend_class", side_effect=exc), + pytest.raises(type(exc)), + ): + CudaPlatform.get_valid_backends( + device_capability=SM90, + attn_selector_config=SELECTOR_CONFIG, + num_heads=32, + ) + + +def test_get_valid_backends_keeps_probing_after_failure(): + healthy = MagicMock() + healthy.validate_configuration.return_value = [] + side_effects = [OSError("probe failed")] + [healthy] * 32 + + with patch("vllm.platforms.cuda._get_attn_backend_class", side_effect=side_effects): + valid, invalid_reasons = CudaPlatform.get_valid_backends( + device_capability=SM90, + attn_selector_config=SELECTOR_CONFIG, + num_heads=32, + ) + assert len(invalid_reasons) == 1 + assert valid + + +def test_selected_backend_probe_failure_raises_value_error_with_cause(): + exc = OSError("libcuda.so.1: cannot open shared object file") + with ( + patch("vllm.platforms.cuda._get_attn_backend_class", side_effect=exc), + patch.object(CudaPlatform, "get_device_capability", return_value=SM90), + pytest.raises(ValueError, match="OSError") as excinfo, + ): + CudaPlatform.get_attn_backend_cls( + selected_backend=AttentionBackendEnum.FLASH_ATTN, + attn_selector_config=SELECTOR_CONFIG, + num_heads=32, + ) + assert excinfo.value.__cause__ is exc diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index bfe08a577213..0aff4ff9bef7 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -386,8 +386,11 @@ def get_valid_backends( device_capability=device_capability, **attn_selector_config._asdict(), ) - except ImportError: - invalid_reasons_i = ["ImportError"] + except (ImportError, OSError) as e: + logger.debug( + "Attention backend %s is unavailable", backend.name, exc_info=True + ) + invalid_reasons_i = [f"{type(e).__name__}: {e}"] if invalid_reasons_i: invalid_reasons[backend] = (priority, invalid_reasons_i) else: @@ -415,8 +418,11 @@ def get_attn_backend_cls( device_capability=device_capability, **attn_selector_config._asdict(), ) - except ImportError: - invalid_reasons = ["ImportError"] + except (ImportError, OSError) as e: + raise ValueError( + f"Selected backend {selected_backend} is not valid for " + f"this configuration. Reason: [{type(e).__name__}: {e}]" + ) from e if invalid_reasons: raise ValueError( f"Selected backend {selected_backend} is not valid for " From 727274a75ba7b2668cd2dc04ad61bbb7c5292bf4 Mon Sep 17 00:00:00 2001 From: Sage <80211083+sagearc@users.noreply.github.com> Date: Thu, 20 Aug 2026 12:39:05 +0300 Subject: [PATCH 201/839] [Rust Frontend] Fix Qwen parser auto-detection (#51169) Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Bugen Zhao Signed-off-by: Sage Ahrac --- rust/src/chat/src/parser/reasoning/mod.rs | 3 +- rust/src/chat/src/parser/reasoning/tests.rs | 17 ++++++++ rust/src/chat/src/parser/tool/mod.rs | 5 ++- rust/src/chat/src/parser/tool/tests.rs | 47 +++++++++++++++------ 4 files changed, 57 insertions(+), 15 deletions(-) diff --git a/rust/src/chat/src/parser/reasoning/mod.rs b/rust/src/chat/src/parser/reasoning/mod.rs index 9a1e37533b45..f85c4a930d56 100644 --- a/rust/src/chat/src/parser/reasoning/mod.rs +++ b/rust/src/chat/src/parser/reasoning/mod.rs @@ -85,7 +85,8 @@ impl ReasoningParserFactory { .register_pattern("deepseek-v3", names::DEEPSEEK_V3) .register_pattern("gemma-4", names::GEMMA4) .register_pattern("gemma4", names::GEMMA4) - .register_pattern("qwen", names::QWEN3) + .register_pattern("qwq", names::DEEPSEEK_R1) + .register_pattern("qwen3", names::QWEN3) .register_pattern("glm-5", names::GLM45) .register_pattern("glm-4.7", names::GLM45) .register_pattern("glm-4.6", names::GLM45) diff --git a/rust/src/chat/src/parser/reasoning/tests.rs b/rust/src/chat/src/parser/reasoning/tests.rs index 7b7f1fe94a53..37a20f8e4f26 100644 --- a/rust/src/chat/src/parser/reasoning/tests.rs +++ b/rust/src/chat/src/parser/reasoning/tests.rs @@ -37,6 +37,23 @@ fn factory_resolves_deepseek_v4_to_qwen3_alias() { ); } +#[test] +fn factory_distinguishes_qwen_model_families() { + let factory = ReasoningParserFactory::new(); + assert_eq!( + factory.resolve_name_for_model("Qwen/QwQ-32B"), + Some(names::DEEPSEEK_R1) + ); + assert_eq!( + factory.resolve_name_for_model("Qwen/Qwen3-8B"), + Some(names::QWEN3) + ); + assert_eq!( + factory.resolve_name_for_model("Qwen/Qwen2.5-0.5B-Instruct"), + None + ); +} + #[test] fn factory_routes_step3p5_models_to_dedicated_parser() { let factory = ReasoningParserFactory::new(); diff --git a/rust/src/chat/src/parser/tool/mod.rs b/rust/src/chat/src/parser/tool/mod.rs index e42b45a48628..cc8b160d937e 100644 --- a/rust/src/chat/src/parser/tool/mod.rs +++ b/rust/src/chat/src/parser/tool/mod.rs @@ -94,9 +94,10 @@ impl ToolParserFactory { .register_pattern("mistral-", names::MISTRAL) .register_pattern("mixtral-", names::MISTRAL) .register_pattern("qwen3-coder", names::QWEN3_CODER) - .register_pattern("qwen2.5-coder", names::QWEN3_CODER) .register_pattern("qwen3.5", names::QWEN3_CODER) - .register_pattern("qwen", names::QWEN3_XML) + .register_pattern("qwq", names::HERMES) + .register_pattern("qwen2.5", names::HERMES) + .register_pattern("qwen3", names::QWEN3_XML) .register_pattern("hermes", names::HERMES) .register_pattern("hy3", names::HY_V3) .register_pattern("hy_v3", names::HY_V3) diff --git a/rust/src/chat/src/parser/tool/tests.rs b/rust/src/chat/src/parser/tool/tests.rs index c89a50b22455..d398d84e351a 100644 --- a/rust/src/chat/src/parser/tool/tests.rs +++ b/rust/src/chat/src/parser/tool/tests.rs @@ -80,22 +80,45 @@ fn factory_creates_registered_parser_for_model() { factory.create_for_model("my-fake-model-v1", &[]).unwrap(); } +#[test] +fn factory_distinguishes_qwen_tool_formats() { + let factory = ToolParserFactory::new(); + + for (model, expected) in [ + ("Qwen/Qwen3.5-0.8B", Some(names::QWEN3_CODER)), + ("Qwen/Qwen3-0.6B", Some(names::QWEN3_XML)), + ("Qwen/Qwen3-Coder-30B", Some(names::QWEN3_CODER)), + ("Qwen/QwQ-32B", Some(names::HERMES)), + ("Qwen/Qwen2.5-0.5B-Instruct", Some(names::HERMES)), + ("Qwen/Qwen2-1.5B-Instruct", None), + ] { + assert_eq!(factory.resolve_name_for_model(model), expected, "{model}"); + } +} + +#[test] +fn factory_parses_qwen2_5_coder_template_with_hermes() { + let factory = ToolParserFactory::new(); + let tool_call = r#" +{"name":"get_weather","arguments":{"location":"Tokyo"}} +"#; + + let mut parser = factory.create_for_model("Qwen/Qwen2.5-Coder-7B-Instruct", &[]).unwrap(); + let mut output = ToolParserOutput::default(); + parser.parse_into(tool_call, &mut output).unwrap(); + output.append(parser.finish().unwrap()); + let output = output.coalesce(); + + assert!(output.normal_text().is_empty()); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); + assert_eq!(output.calls()[0].arguments, r#"{"location":"Tokyo"}"#); +} + #[test] fn factory_new_resolves_default_patterns() { let factory = ToolParserFactory::new(); - assert_eq!( - factory.resolve_name_for_model("Qwen/Qwen3.5-0.8B"), - Some(names::QWEN3_CODER) - ); - assert_eq!( - factory.resolve_name_for_model("Qwen/Qwen3-0.6B"), - Some(names::QWEN3_XML) - ); - assert_eq!( - factory.resolve_name_for_model("Qwen/Qwen3-Coder-30B"), - Some(names::QWEN3_CODER) - ); assert_eq!( factory.resolve_name_for_model("meta-llama-4-maverick"), Some(names::LLAMA4_JSON) From 5b1e7a812b1358125b6d5291c51bbc65264f044d Mon Sep 17 00:00:00 2001 From: Ziming Huang Date: Thu, 20 Aug 2026 18:05:43 +0800 Subject: [PATCH 202/839] [BugFix][Mooncake] Fix Mooncake saves from sparse Mamba block tables (#51362) Signed-off-by: ZeldaHuang Signed-off-by: Ziming Huang --- .../unit/test_mooncake_store_hma_e2e.py | 3 +- .../unit/test_mooncake_store_worker.py | 104 ++++++++++++++---- .../kv_connector/v1/mooncake/store/worker.py | 13 +++ 3 files changed, 99 insertions(+), 21 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py index 32457f13d525..cd23186906bd 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py @@ -225,7 +225,8 @@ def _fake_thread_init(*args, **kwargs): save_req = ReqMeta( req_id="r0", token_len_chunk=64, - block_ids=([0, 1, 2, 3], [0, 1, 2, 3]), + # Block 0 is reserved as NULL_BLOCK_ID by the production block pool. + block_ids=([1, 2, 3, 4], [1, 2, 3, 4]), block_hashes=hs, can_save=True, store_job_id=1, diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index 39c0f8eb4ade..55081f90bfc8 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -36,6 +36,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.metrics import ( MooncakeStoreConnectorStats, ) +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID from vllm.v1.core.kv_cache_utils import BlockHash, maybe_convert_block_hash @@ -173,7 +174,7 @@ def _make_store_req(req_id: str, block_hashes: list[bytes]) -> ReqMeta: return ReqMeta( req_id=req_id, token_len_chunk=32, - block_ids=([0, 1],), + block_ids=([1, 2],), block_hashes=block_hashes, can_save=True, ) @@ -183,7 +184,7 @@ def _make_multi_group_store_req(req_id: str, block_hashes: list[bytes]) -> ReqMe return ReqMeta( req_id=req_id, token_len_chunk=32, - block_ids=([0, 1], [2, 3]), + block_ids=([1, 2], [3, 4]), block_hashes=block_hashes, can_save=True, ) @@ -517,7 +518,7 @@ def test_store_sending_thread_delta_saves_only_new_full_attention_chunks(): ReqMeta( req_id="req-a", token_len_chunk=64, - block_ids=([0, 1, 2, 3],), + block_ids=([1, 2, 3, 4],), block_hashes=[b"a0", b"a1", b"a2", b"a3"], can_save=True, ), @@ -543,7 +544,7 @@ def test_store_sending_thread_delta_strides_with_local_phase(): ReqMeta( req_id="req-a", token_len_chunk=64, - block_ids=([0, 1, 2, 3],), + block_ids=([1, 2, 3, 4],), block_hashes=[b"a0", b"a1", b"a2", b"a3"], can_save=True, ), @@ -569,7 +570,7 @@ def test_tp_sharded_group_saves_every_block_on_every_rank(): ReqMeta( req_id="req-a", token_len_chunk=64, - block_ids=([0, 1, 2, 3],), + block_ids=([1, 2, 3, 4],), block_hashes=[b"a0", b"a1", b"a2", b"a3"], can_save=True, ), @@ -595,7 +596,7 @@ def test_store_sending_thread_retries_skipped_range_after_pressure(): ReqMeta( req_id="req-a", token_len_chunk=16, - block_ids=([0],), + block_ids=([1],), block_hashes=[b"a0"], can_save=True, ), @@ -614,7 +615,7 @@ def test_store_sending_thread_retries_skipped_range_after_pressure(): ReqMeta( req_id="req-a", token_len_chunk=64, - block_ids=([0, 1, 2, 3],), + block_ids=([1, 2, 3, 4],), block_hashes=[b"a0", b"a1", b"a2", b"a3"], can_save=True, ), @@ -686,6 +687,69 @@ def test_partial_tail_offload_skips_null_source_blocks(): assert addrs == [[0x1000 + 2 * 256], [0x1000 + 3 * 256]] +def test_store_sending_thread_skips_null_sparse_group_blocks(): + from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheGroupSpec, + MambaSpec, + ) + + store = MagicMock() + store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) + store.batch_put_from_multi_buffers.side_effect = ( + lambda keys, addrs, sizes, replicate_config: [256] * len(keys) + ) + full = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None) + mamba = MambaSpec( + block_size=16, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ) + coord = mooncake_store_worker.MooncakeStoreCoordinator( + [ + KVCacheGroupSpec(["full"], full), + KVCacheGroupSpec(["mamba"], mamba), + ], + scheduler_block_size=16, + hash_block_size=16, + ) + + db_full = ChunkedTokenDatabase( + KeyMetadata("test-model", 0, 0, 0, 0, group_id=0), + block_size=16, + ) + db_full.set_kv_caches_base_addr([0x1000]) + db_full.set_block_len([256]) + db_mamba = ChunkedTokenDatabase( + KeyMetadata("test-model", 0, 0, 0, 0, group_id=1), + block_size=16, + ) + db_mamba.set_kv_caches_base_addr([0x2000]) + db_mamba.set_block_len([512]) + + thread = _make_store_sending_thread( + store, + coord=coord, + token_databases=[db_full, db_mamba], + ) + _run_store_req( + thread, + ReqMeta( + req_id="req-a", + token_len_chunk=48, + block_ids=([1, 2, 3], [NULL_BLOCK_ID, 4, NULL_BLOCK_ID]), + block_hashes=[b"a0", b"a1", b"a2"], + can_save=True, + ), + ) + + keys, addrs, _, _ = store.batch_put_from_multi_buffers.call_args.args + mamba_keys = [key for key in keys if "@group:1" in key] + assert [key.rsplit("@", 1)[-1] for key in mamba_keys] == [b"a1".hex()] + assert all(addr[0] >= 0x2000 for key, addr in zip(keys, addrs) if "@group:1" in key) + + def test_partial_tail_offload_replaces_stale_group_ids_after_filtering(): store = MagicMock() store.batch_is_exist.return_value = [1, 0, 0] @@ -752,7 +816,7 @@ def test_store_sending_thread_delta_start_rank_saves_second_local_chunk(): ReqMeta( req_id="req-a", token_len_chunk=64, - block_ids=([0, 1, 2, 3],), + block_ids=([1, 2, 3, 4],), block_hashes=[b"a0", b"a1", b"a2", b"a3"], can_save=True, ), @@ -859,7 +923,7 @@ def test_store_sending_thread_prepares_missing_chunks_once_per_group(): ReqMeta( req_id="req-a", token_len_chunk=48, - block_ids=([0, 1, 2], [2, 1, 0]), + block_ids=([1, 2, 3], [3, 2, 1]), block_hashes=[b"a0", b"a1", b"a2"], can_save=True, ), @@ -867,8 +931,8 @@ def test_store_sending_thread_prepares_missing_chunks_once_per_group(): db0.prepare_value.assert_not_called() db1.prepare_value.assert_not_called() - db0.prepare_values.assert_called_once_with([(0, 16), (32, 48)], [0, 1, 2]) - db1.prepare_values.assert_called_once_with([(16, 32), (32, 48)], [2, 1, 0]) + db0.prepare_values.assert_called_once_with([(0, 16), (32, 48)], [1, 2, 3]) + db1.prepare_values.assert_called_once_with([(16, 32), (32, 48)], [3, 2, 1]) keys, addrs, sizes, _ = store.batch_put_from_multi_buffers.call_args.args assert [key.rsplit("@", 1)[-1] for key in keys] == [ @@ -877,7 +941,7 @@ def test_store_sending_thread_prepares_missing_chunks_once_per_group(): "6131", "6132", ] - assert addrs == [[0x1000], [0x1200], [0x2200], [0x2000]] + assert addrs == [[0x1100], [0x1300], [0x2400], [0x2200]] assert sizes == [[256], [256], [512], [512]] @@ -1250,7 +1314,7 @@ def test_store_sending_thread_multiple_segments_share_logical_group_id(): "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6130", "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6131", ] - assert addrs == [[0x1000, 0x2000], [0x1100, 0x2100]] + assert addrs == [[0x1100, 0x2100], [0x1200, 0x2200]] assert sizes == [[256, 256], [256, 256]] assert config.group_ids == [ "vllm-mooncake-store:test-model@6130", @@ -1301,7 +1365,7 @@ def test_store_sending_thread_group_ids_share_across_kv_cache_groups(): "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:1@6130", "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:1@6131", ] - assert addrs == [[0x1000], [0x1100], [0x3200], [0x3300]] + assert addrs == [[0x1100], [0x1200], [0x3300], [0x3400]] assert sizes == [[256], [256], [256], [256]] # Different vLLM KV cache groups for the same prefix chunk share the # same Mooncake lifecycle group id. @@ -1935,7 +1999,7 @@ def test_store_sending_thread_clamps_token_len_to_lcm(): ReqMeta( req_id="r0", token_len_chunk=33, - block_ids=([0, 1, 2],), + block_ids=([1, 2, 3],), block_hashes=[b"a0", b"a1", b"a2"], can_save=True, ), @@ -1974,7 +2038,7 @@ def test_store_sending_thread_skips_when_token_len_below_lcm(): ReqMeta( req_id="r0", token_len_chunk=32, - block_ids=([0, 1],), + block_ids=([1, 2],), block_hashes=[b"a0", b"a1"], can_save=True, ), @@ -2050,7 +2114,7 @@ def test_store_sending_thread_only_stores_swa_blocks_in_window(): ReqMeta( req_id="r0", token_len_chunk=64, - block_ids=([0, 1], list(range(8))), + block_ids=([1, 2], list(range(1, 9))), block_hashes=hs, can_save=True, ), @@ -2126,7 +2190,7 @@ def test_store_sending_thread_delta_saves_only_new_swa_boundary_chunks(): ReqMeta( req_id="r0", token_len_chunk=64, - block_ids=([0, 1], list(range(8))), + block_ids=([1, 2], list(range(1, 9))), block_hashes=hs, can_save=True, ), @@ -2196,7 +2260,7 @@ def test_store_sending_thread_kv_events_use_group_chunk_metadata(): ReqMeta( req_id="r0", token_len_chunk=32, - block_ids=([0], list(range(4))), + block_ids=([1], list(range(1, 5))), block_hashes=hs, can_save=True, token_ids=list(range(32)), @@ -2222,7 +2286,7 @@ def _make_event_store_req(token_len: int, token_ids_start: int = 0) -> ReqMeta: return ReqMeta( req_id="r0", token_len_chunk=token_len, - block_ids=(list(range(num_blocks)),), + block_ids=(list(range(1, num_blocks + 1)),), block_hashes=[f"a{i}".encode() for i in range(num_blocks)], can_save=True, token_ids=list(range(token_ids_start, token_len)), diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index efc93f7d9423..c9fba1c69937 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -889,6 +889,19 @@ def _handle_request(self, req_meta: ReqMeta): put_step=put_step, put_step_rank=put_step_rank, ): + block_idx = start // db.block_size + group_blocks = block_ids_per_group[g_idx] + if block_idx >= len(group_blocks) or ( + group_blocks[block_idx] == NULL_BLOCK_ID + ): + logger.debug( + "Skipping unavailable Mooncake store source block " + "(req=%s, group=%d, block=%d)", + req_id, + g_idx, + block_idx, + ) + continue starts.append(start) ends.append(end) keys.append(db.key_for(block_hash)) From 1eab6fef01b78ec4eab6b7156bbf5f120e48d381 Mon Sep 17 00:00:00 2001 From: Turner Jabbour Date: Thu, 20 Aug 2026 04:09:59 -0600 Subject: [PATCH 203/839] [CI] replace shellcheck script with shellcheck-py hook (#52572) Signed-off-by: Turner Signed-off-by: Turner Jabbour Co-authored-by: Claude Sonnet 5 --- .buildkite/scripts/ci-fetch-log.sh | 1 + .../scripts/docker-build-metadata-args.sh | 2 +- .../scripts/hardware_ci/run-amd-test.sh | 19 ++- .buildkite/scripts/publish-release-images.sh | 132 +++++++++--------- .buildkite/scripts/run-multi-node-test.sh | 29 ++-- .buildkite/scripts/tool_call/run-bfcl-eval.sh | 1 + .../scripts/xpu/create-xpu-ecr-manifest.sh | 4 +- .github/workflows/matchers/shellcheck.json | 18 +++ .github/workflows/pre-commit.yml | 5 +- .pre-commit-config.yaml | 10 +- build_vllm_ppc64le.sh | 2 +- .../test_vllm_nonroot_entrypoint.sh | 6 +- docs/pre_run_check.sh | 1 + .../run_mamba_prefix_cache_test.sh | 8 +- .../run_xpu_disagg_accuracy_test.sh | 85 ++++++----- .../spec_decode_acceptance_test.sh | 42 +++--- tools/install_torchcodec_rocm.sh | 1 + tools/pre_commit/shellcheck.sh | 24 ---- tools/setup_deepgemm_pythons.sh | 3 +- vllm/utils/numa_wrapper.sh | 2 + 20 files changed, 201 insertions(+), 194 deletions(-) create mode 100644 .github/workflows/matchers/shellcheck.json delete mode 100755 tools/pre_commit/shellcheck.sh diff --git a/.buildkite/scripts/ci-fetch-log.sh b/.buildkite/scripts/ci-fetch-log.sh index 4830135a1120..eced685f47d3 100755 --- a/.buildkite/scripts/ci-fetch-log.sh +++ b/.buildkite/scripts/ci-fetch-log.sh @@ -143,6 +143,7 @@ fi # Build-wide mode: fetch finished jobs matching $SCOPE. [ -z "$OUT" ] || die "[output_file] is only valid when fetching a single job." +# shellcheck disable=SC2016 # single-quoted: these are awk field refs ($3/$4/$5), not shell vars case "$SCOPE" in failed) FILTER='$3 == "True" && $4 == "False" && $5 == "True"' ;; soft) FILTER='$3 == "True" && $5 == "True"' ;; diff --git a/.buildkite/scripts/docker-build-metadata-args.sh b/.buildkite/scripts/docker-build-metadata-args.sh index 9aa6fa9314f7..6986b5e7ad1c 100644 --- a/.buildkite/scripts/docker-build-metadata-args.sh +++ b/.buildkite/scripts/docker-build-metadata-args.sh @@ -29,7 +29,7 @@ if [[ -n "${BUILDKITE:-}" || -n "${BUILDKITE_COMMIT:-}" ]]; then image_tag="${image_name}:nightly-${tag_commit}" elif [[ "${variant}" == cu* ]]; then cuda_variant="${variant%%-*}" - remaining_variant="${variant#${cuda_variant}}" + remaining_variant="${variant#"${cuda_variant}"}" image_tag="${image_name}:${cuda_variant}-nightly-${tag_commit}${remaining_variant}" else image_tag="${image_name}:nightly-${tag_commit}${variant_suffix}" diff --git a/.buildkite/scripts/hardware_ci/run-amd-test.sh b/.buildkite/scripts/hardware_ci/run-amd-test.sh index ff1fd0dd833c..b9674c108b13 100755 --- a/.buildkite/scripts/hardware_ci/run-amd-test.sh +++ b/.buildkite/scripts/hardware_ci/run-amd-test.sh @@ -1,5 +1,11 @@ #!/bin/bash +# shellcheck disable=SC2329 # Every function in this file is only ever reached +# transitively through handle_amd_runner_exit, the EXIT trap handler - shellcheck's +# unused-function check doesn't do full reachability analysis from an indirectly +# invoked entry point, so it flags the whole diagnostic-collection call chain as +# dead code even though it runs live on every nonzero exit. + # This script runs ROCm tests either directly in a native CI pod or inside the # corresponding Docker container. Multi-node tests continue to use Docker. # @@ -1523,13 +1529,15 @@ fi clear_ci_orchestration_env if is_multi_node "$commands"; then echo "--- Multi-node job detected" - export DCKR_VER=$(docker --version | sed 's/Docker version \(.*\), build .*/\1/') + DCKR_VER=$(docker --version | sed 's/Docker version \(.*\), build .*/\1/') + export DCKR_VER # Parse the bracket syntax: prefix ; [node0_cmds] && [node1_cmds] # BASH_REMATCH[1] = prefix (everything before first bracket) # BASH_REMATCH[2] = comma-separated node0 commands # BASH_REMATCH[3] = comma-separated node1 commands if [[ "$commands" =~ ^(.*)\[(.*)"] && ["(.*)\]$ ]]; then + # shellcheck disable=SC2001 # verified equivalent behavior; TODO: switch to param expansion in a follow-up cleanup PR prefix=$(echo "${BASH_REMATCH[1]}" | sed 's/;//g') echo "PREFIX: ${prefix}" @@ -1545,7 +1553,9 @@ if is_multi_node "$commands"; then fi for i in "${!node0[@]}"; do + # shellcheck disable=SC2001 # verified equivalent behavior; TODO: switch to param expansion in a follow-up cleanup PR command_node_0=$(echo "${node0[i]}" | sed 's/\"//g') + # shellcheck disable=SC2001 # verified equivalent behavior; TODO: switch to param expansion in a follow-up cleanup PR command_node_1=$(echo "${node1[i]}" | sed 's/\"//g') step_cmd="./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 2 ${image_name} '${command_node_0}' '${command_node_1}'" @@ -1581,14 +1591,15 @@ else ulimit_core_hard="-1" fi # Disable core dumps in the ROCm test container unless the ROCm debug agent is enabled - coredump_flags="--ulimit core=0:$ulimit_core_hard" + coredump_flags=(--ulimit "core=0:$ulimit_core_hard") if [[ "$commands" == *"ROCm debug agent enabled"* ]]; then # Works around https://github.com/rocm/rocm-systems/issues/6206 - coredump_flags='-e HSA_COREDUMP_PATTERN="/tmp/gpucore.%p"' + coredump_flags=(-e 'HSA_COREDUMP_PATTERN="/tmp/gpucore.%p"') else echo "ROCm debug agent not enabled, coredumps are disabled in the test container." fi + # shellcheck disable=SC2086 # word splitting is intentional: both hold multiple docker flags docker run \ "${docker_run_terminal_args[@]}" \ --device /dev/kfd $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES \ @@ -1597,7 +1608,7 @@ else --shm-size=16gb \ --group-add "$render_gid" \ --rm \ - $coredump_flags \ + "${coredump_flags[@]}" \ -e HF_TOKEN \ -e "HF_HUB_DOWNLOAD_TIMEOUT=${HF_HUB_DOWNLOAD_TIMEOUT}" \ -e "HF_HUB_ETAG_TIMEOUT=${HF_HUB_ETAG_TIMEOUT}" \ diff --git a/.buildkite/scripts/publish-release-images.sh b/.buildkite/scripts/publish-release-images.sh index dd54b1717e10..514ee99b3892 100755 --- a/.buildkite/scripts/publish-release-images.sh +++ b/.buildkite/scripts/publish-release-images.sh @@ -42,97 +42,97 @@ aws ecr-public get-login-password --region us-east-1 | \ # ---- CUDA (default: 13.0) ---- if target_enabled cuda-13-0; then - docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64 - docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64 + docker pull "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64" + docker pull "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64" - docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64 vllm/vllm-openai:latest-x86_64 - docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 + docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64" vllm/vllm-openai:latest-x86_64 + docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64" "vllm/vllm-openai:v${RELEASE_VERSION}-x86_64" docker push vllm/vllm-openai:latest-x86_64 - docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 + docker push "vllm/vllm-openai:v${RELEASE_VERSION}-x86_64" - docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64 vllm/vllm-openai:latest-aarch64 - docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64 + docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64" vllm/vllm-openai:latest-aarch64 + docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64" "vllm/vllm-openai:v${RELEASE_VERSION}-aarch64" docker push vllm/vllm-openai:latest-aarch64 - docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64 + docker push "vllm/vllm-openai:v${RELEASE_VERSION}-aarch64" docker manifest rm vllm/vllm-openai:latest || true - docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION} || true + docker manifest rm "vllm/vllm-openai:v${RELEASE_VERSION}" || true docker manifest create vllm/vllm-openai:latest vllm/vllm-openai:latest-x86_64 vllm/vllm-openai:latest-aarch64 - docker manifest create vllm/vllm-openai:v${RELEASE_VERSION} vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64 + docker manifest create "vllm/vllm-openai:v${RELEASE_VERSION}" "vllm/vllm-openai:v${RELEASE_VERSION}-x86_64" "vllm/vllm-openai:v${RELEASE_VERSION}-aarch64" docker manifest push vllm/vllm-openai:latest - docker manifest push vllm/vllm-openai:v${RELEASE_VERSION} + docker manifest push "vllm/vllm-openai:v${RELEASE_VERSION}" fi # ---- CUDA 12.9 ---- if target_enabled cuda-12-9; then - docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129 - docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129 + docker pull "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129" + docker pull "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129" - docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129 vllm/vllm-openai:latest-x86_64-cu129 - docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 + docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129" vllm/vllm-openai:latest-x86_64-cu129 + docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129" "vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129" docker push vllm/vllm-openai:latest-x86_64-cu129 - docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 + docker push "vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129" - docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129 vllm/vllm-openai:latest-aarch64-cu129 - docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129 + docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129" vllm/vllm-openai:latest-aarch64-cu129 + docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129" "vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129" docker push vllm/vllm-openai:latest-aarch64-cu129 - docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129 + docker push "vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129" docker manifest rm vllm/vllm-openai:latest-cu129 || true - docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION}-cu129 || true + docker manifest rm "vllm/vllm-openai:v${RELEASE_VERSION}-cu129" || true docker manifest create vllm/vllm-openai:latest-cu129 vllm/vllm-openai:latest-x86_64-cu129 vllm/vllm-openai:latest-aarch64-cu129 - docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129 + docker manifest create "vllm/vllm-openai:v${RELEASE_VERSION}-cu129" "vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129" "vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129" docker manifest push vllm/vllm-openai:latest-cu129 - docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu129 + docker manifest push "vllm/vllm-openai:v${RELEASE_VERSION}-cu129" fi # ---- Ubuntu 24.04 (CUDA 13.0) ---- if target_enabled cuda-13-0-ubuntu-24-04; then - docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404 - docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404 + docker pull "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404" + docker pull "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404" - docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404 vllm/vllm-openai:latest-x86_64-ubuntu2404 - docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404 + docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404" vllm/vllm-openai:latest-x86_64-ubuntu2404 + docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404" "vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404" docker push vllm/vllm-openai:latest-x86_64-ubuntu2404 - docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404 + docker push "vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404" - docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404 vllm/vllm-openai:latest-aarch64-ubuntu2404 - docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404 + docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404" vllm/vllm-openai:latest-aarch64-ubuntu2404 + docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404" "vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404" docker push vllm/vllm-openai:latest-aarch64-ubuntu2404 - docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404 + docker push "vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404" docker manifest rm vllm/vllm-openai:latest-ubuntu2404 || true - docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404 || true + docker manifest rm "vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404" || true docker manifest create vllm/vllm-openai:latest-ubuntu2404 vllm/vllm-openai:latest-x86_64-ubuntu2404 vllm/vllm-openai:latest-aarch64-ubuntu2404 - docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404 + docker manifest create "vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404" "vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404" "vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404" docker manifest push vllm/vllm-openai:latest-ubuntu2404 - docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404 + docker manifest push "vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404" fi # ---- Ubuntu 24.04 (CUDA 12.9) ---- if target_enabled cuda-12-9-ubuntu-24-04; then - docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404 - docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404 + docker pull "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404" + docker pull "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404" - docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404 vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404 - docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404 + docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404" vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404 + docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404" "vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404" docker push vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404 - docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404 + docker push "vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404" - docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404 vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404 - docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404 + docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404" vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404 + docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404" "vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404" docker push vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404 - docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404 + docker push "vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404" docker manifest rm vllm/vllm-openai:latest-cu129-ubuntu2404 || true - docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404 || true + docker manifest rm "vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404" || true docker manifest create vllm/vllm-openai:latest-cu129-ubuntu2404 vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404 vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404 - docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404 + docker manifest create "vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404" "vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404" "vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404" docker manifest push vllm/vllm-openai:latest-cu129-ubuntu2404 - docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404 + docker manifest push "vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404" fi # ---- ROCm ---- @@ -141,36 +141,36 @@ if target_enabled rocm; then ROCM_BASE_CACHE_KEY=$(.buildkite/scripts/cache-rocm-base-wheels.sh key) echo "ROCm base cache key: ${ROCM_BASE_CACHE_KEY}" - docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm - docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base + docker pull "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm" + docker pull "public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base" - docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm vllm/vllm-openai-rocm:latest - docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm vllm/vllm-openai-rocm:v${RELEASE_VERSION} + docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm" vllm/vllm-openai-rocm:latest + docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm" "vllm/vllm-openai-rocm:v${RELEASE_VERSION}" docker push vllm/vllm-openai-rocm:latest - docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION} + docker push "vllm/vllm-openai-rocm:v${RELEASE_VERSION}" - docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base vllm/vllm-openai-rocm:latest-base - docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base + docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base" vllm/vllm-openai-rocm:latest-base + docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base" "vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base" docker push vllm/vllm-openai-rocm:latest-base - docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base + docker push "vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base" fi # ---- XPU ---- if target_enabled xpu; then - docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-xpu + docker pull "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-xpu" - docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-xpu vllm/vllm-openai-xpu:latest-x86_64 - docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-xpu vllm/vllm-openai-xpu:v${RELEASE_VERSION}-x86_64 + docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-xpu" vllm/vllm-openai-xpu:latest-x86_64 + docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-xpu" "vllm/vllm-openai-xpu:v${RELEASE_VERSION}-x86_64" docker push vllm/vllm-openai-xpu:latest-x86_64 - docker push vllm/vllm-openai-xpu:v${RELEASE_VERSION}-x86_64 + docker push "vllm/vllm-openai-xpu:v${RELEASE_VERSION}-x86_64" docker manifest rm vllm/vllm-openai-xpu:latest || true - docker manifest rm vllm/vllm-openai-xpu:v${RELEASE_VERSION} || true + docker manifest rm "vllm/vllm-openai-xpu:v${RELEASE_VERSION}" || true docker manifest create vllm/vllm-openai-xpu:latest vllm/vllm-openai-xpu:latest-x86_64 --amend - docker manifest create vllm/vllm-openai-xpu:v${RELEASE_VERSION} vllm/vllm-openai-xpu:v${RELEASE_VERSION}-x86_64 --amend + docker manifest create "vllm/vllm-openai-xpu:v${RELEASE_VERSION}" "vllm/vllm-openai-xpu:v${RELEASE_VERSION}-x86_64" --amend docker manifest push vllm/vllm-openai-xpu:latest - docker manifest push vllm/vllm-openai-xpu:v${RELEASE_VERSION} + docker manifest push "vllm/vllm-openai-xpu:v${RELEASE_VERSION}" fi # ---- CPU ---- @@ -192,22 +192,22 @@ if target_enabled cpu; then if [ "$CPU_X86_AVAILABLE" = "true" ] && [ "$CPU_ARM_AVAILABLE" = "true" ]; then docker pull "${CPU_X86_TAG}" docker tag "${CPU_X86_TAG}" vllm/vllm-openai-cpu:latest-x86_64 - docker tag "${CPU_X86_TAG}" vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64 + docker tag "${CPU_X86_TAG}" "vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64" docker push vllm/vllm-openai-cpu:latest-x86_64 - docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64 + docker push "vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64" docker pull "${CPU_ARM_TAG}" docker tag "${CPU_ARM_TAG}" vllm/vllm-openai-cpu:latest-arm64 - docker tag "${CPU_ARM_TAG}" vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64 + docker tag "${CPU_ARM_TAG}" "vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64" docker push vllm/vllm-openai-cpu:latest-arm64 - docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64 + docker push "vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64" docker manifest rm vllm/vllm-openai-cpu:latest || true - docker manifest rm vllm/vllm-openai-cpu:v${RELEASE_VERSION} || true + docker manifest rm "vllm/vllm-openai-cpu:v${RELEASE_VERSION}" || true docker manifest create vllm/vllm-openai-cpu:latest vllm/vllm-openai-cpu:latest-x86_64 vllm/vllm-openai-cpu:latest-arm64 - docker manifest create vllm/vllm-openai-cpu:v${RELEASE_VERSION} vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64 + docker manifest create "vllm/vllm-openai-cpu:v${RELEASE_VERSION}" "vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64" "vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64" docker manifest push vllm/vllm-openai-cpu:latest - docker manifest push vllm/vllm-openai-cpu:v${RELEASE_VERSION} + docker manifest push "vllm/vllm-openai-cpu:v${RELEASE_VERSION}" elif [ "$CPU_X86_AVAILABLE" = "false" ] && [ "$CPU_ARM_AVAILABLE" = "false" ]; then echo "WARNING: Neither CPU image found in ECR, skipping CPU publish (ensure block-cpu-release-image-build and block-arm64-cpu-release-image-build were unblocked and the builds finished pushing)" else diff --git a/.buildkite/scripts/run-multi-node-test.sh b/.buildkite/scripts/run-multi-node-test.sh index cf54986e0151..2b8516b1d690 100755 --- a/.buildkite/scripts/run-multi-node-test.sh +++ b/.buildkite/scripts/run-multi-node-test.sh @@ -44,20 +44,20 @@ start_network() { start_nodes() { for node in $(seq 0 $(($NUM_NODES-1))); do - if [ "$IS_ROCM" -eq 1 ]; then - GPU_DEVICES='--device /dev/kfd --device /dev/dri -e HIP_VISIBLE_DEVICES=' - else - GPU_DEVICES='--gpus "device=' - fi + DEVICE_LIST="" for node_gpu in $(seq 0 $(($NUM_GPUS - 1))); do DEVICE_NUM=$(($node * $NUM_GPUS + $node_gpu)) - GPU_DEVICES+=$(($DEVICE_NUM)) + DEVICE_LIST+=$(($DEVICE_NUM)) if [ "$node_gpu" -lt $(($NUM_GPUS - 1)) ]; then - GPU_DEVICES+=',' + DEVICE_LIST+=',' fi done - if [ "$IS_ROCM" -eq 0 ]; then - GPU_DEVICES+='"' + if [ "$IS_ROCM" -eq 1 ]; then + GPU_DEVICES=(--device /dev/kfd --device /dev/dri -e "HIP_VISIBLE_DEVICES=${DEVICE_LIST}") + else + # The literal quotes around device=... are required by docker's + # --gpus value parser when the device list itself contains commas. + GPU_DEVICES=(--gpus "\"device=${DEVICE_LIST}\"") fi # start the container in detached mode @@ -67,7 +67,7 @@ start_nodes() { # 3. map the huggingface cache directory to the container # 3. assign ip addresses to the containers (head node: 192.168.10.10, worker nodes: # starting from 192.168.10.11) - docker run -d $GPU_DEVICES --shm-size=10.24gb -e HF_TOKEN \ + docker run -d "${GPU_DEVICES[@]}" --shm-size=10.24gb -e HF_TOKEN \ -v ~/.cache/huggingface:/root/.cache/huggingface --name "node$node" \ --network docker-net --ip 192.168.10.$((10 + $node)) --rm "$DOCKER_IMAGE" \ /bin/bash -c "tail -f /dev/null" @@ -96,16 +96,15 @@ run_nodes() { # we start the worker nodes first, in detached mode, and then start the head node # in the foreground, so that the output of the head node is visible in the buildkite logs for node in $(seq $(($NUM_NODES - 1)) -1 0); do - GPU_DEVICES='"device=' + DEVICE_LIST="" for node_gpu in $(seq 0 $(($NUM_GPUS - 1))); do DEVICE_NUM=$(($node * $NUM_GPUS + $node_gpu)) - GPU_DEVICES+=$(($DEVICE_NUM)) + DEVICE_LIST+=$(($DEVICE_NUM)) if [ "$node_gpu" -lt $(($NUM_GPUS - 1)) ]; then - GPU_DEVICES+=',' + DEVICE_LIST+=',' fi done - GPU_DEVICES+='"' - echo "Running node$node with GPU devices: $GPU_DEVICES" + echo "Running node$node with GPU devices: $DEVICE_LIST" if [ "$node" -ne 0 ]; then docker exec -d "node$node" /bin/bash -c "cd $WORKING_DIR ; ${COMMANDS[$node]}" else diff --git a/.buildkite/scripts/tool_call/run-bfcl-eval.sh b/.buildkite/scripts/tool_call/run-bfcl-eval.sh index d50767ef0f22..1703a5b4fea4 100755 --- a/.buildkite/scripts/tool_call/run-bfcl-eval.sh +++ b/.buildkite/scripts/tool_call/run-bfcl-eval.sh @@ -75,6 +75,7 @@ fi # ---- Cleanup handler ---- SERVER_PID="" +# shellcheck disable=SC2329 # invoked via `trap cleanup EXIT` below cleanup() { if [ -n "$SERVER_PID" ]; then echo "Stopping vLLM server (pid=$SERVER_PID)..." diff --git a/.buildkite/scripts/xpu/create-xpu-ecr-manifest.sh b/.buildkite/scripts/xpu/create-xpu-ecr-manifest.sh index 9f21118d5019..f2bfc6439d0a 100644 --- a/.buildkite/scripts/xpu/create-xpu-ecr-manifest.sh +++ b/.buildkite/scripts/xpu/create-xpu-ecr-manifest.sh @@ -7,7 +7,7 @@ REPO="vllm-release-repo" ARCH_TAG="${BUILDKITE_COMMIT}-$(uname -m)-xpu" PLATFORM_TAG="${BUILDKITE_COMMIT}-xpu" -aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin ${REGISTRY} +aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "${REGISTRY}" docker manifest rm "${REGISTRY}/${REPO}:${PLATFORM_TAG}" || true docker manifest create "${REGISTRY}/${REPO}:${PLATFORM_TAG}" "${REGISTRY}/${REPO}:${ARCH_TAG}" --amend -docker manifest push ${REGISTRY}/${REPO}:${PLATFORM_TAG} \ No newline at end of file +docker manifest push "${REGISTRY}/${REPO}:${PLATFORM_TAG}" \ No newline at end of file diff --git a/.github/workflows/matchers/shellcheck.json b/.github/workflows/matchers/shellcheck.json new file mode 100644 index 000000000000..b4c95f0d5691 --- /dev/null +++ b/.github/workflows/matchers/shellcheck.json @@ -0,0 +1,18 @@ +{ + "problemMatcher": [ + { + "owner": "shellcheck", + "pattern": [ + { + "regexp": "^(.+):(\\d+):(\\d+): (error|warning|note|info): (.+) \\[(SC\\d+)\\]$", + "file": 1, + "line": 2, + "column": 3, + "severity": 4, + "message": 5, + "code": 6 + } + ] + } + ] +} diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml index 695cb0289440..84848dfb34af 100644 --- a/.github/workflows/pre-commit.yml +++ b/.github/workflows/pre-commit.yml @@ -53,13 +53,10 @@ jobs: - uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # v6.1.0 with: python-version: "3.12" - # Provide shellcheck on PATH so tools/pre_commit/shellcheck.sh skips its - # wget + tar -xJ self-download, which the self-hosted runner image lacks - # (no wget/xz). Pinned to shellcheck 0.10.0 to match the script's "stable". - - run: python -m pip install shellcheck-py==0.10.0.1 - run: echo "::add-matcher::.github/workflows/matchers/actionlint.json" - run: echo "::add-matcher::.github/workflows/matchers/markdownlint.json" - run: echo "::add-matcher::.github/workflows/matchers/mypy.json" + - run: echo "::add-matcher::.github/workflows/matchers/shellcheck.json" - uses: pre-commit/action@2c7b3805fd2a0fd8c1884dcaebf91fc102a13ecd # v3.0.1 with: extra_args: --all-files --hook-stage manual diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index b9344b6ffc8f..0954c53cb290 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -35,6 +35,11 @@ repos: rev: v1.7.7 hooks: - id: actionlint +- repo: https://github.com/shellcheck-py/shellcheck-py + rev: 745eface02aef23e168a8afb6b5737818efbea95 # v0.11.0.1 + hooks: + - id: shellcheck + args: ["-s", "bash", "--format=gcc"] - repo: https://github.com/astral-sh/uv-pre-commit rev: 0.11.1 hooks: @@ -195,11 +200,6 @@ repos: entry: python tools/pre_commit/mypy.py "3.13" <<: *mypy_common stages: [manual] # Only run in CI - - id: shellcheck - name: Lint shell scripts - entry: tools/pre_commit/shellcheck.sh - language: script - types: [shell] - id: png-lint name: Lint PNG exports from excalidraw entry: tools/pre_commit/png-lint.sh diff --git a/build_vllm_ppc64le.sh b/build_vllm_ppc64le.sh index 3c0b74cc74d3..f2fbad8e2199 100644 --- a/build_vllm_ppc64le.sh +++ b/build_vllm_ppc64le.sh @@ -130,7 +130,7 @@ TEMP_BUILD_DIR=$(mktemp -d) cd "${TEMP_BUILD_DIR}" export BUILD_SOX=1 BUILD_KALDI=1 BUILD_RNNT=1 USE_FFMPEG=0 USE_ROCM=0 USE_CUDA=0 export TORCHAUDIO_TEST_ALLOW_SKIP_IF_NO_FFMPEG=1 -git clone --recursive https://github.com/pytorch/audio.git -b v${TORCHAUDIO_VERSION} +git clone --recursive https://github.com/pytorch/audio.git -b "v${TORCHAUDIO_VERSION}" cd audio #patching sed -i ' diff --git a/docker/entrypoints/test_vllm_nonroot_entrypoint.sh b/docker/entrypoints/test_vllm_nonroot_entrypoint.sh index c136f0549199..5f3a5d2e3b60 100755 --- a/docker/entrypoints/test_vllm_nonroot_entrypoint.sh +++ b/docker/entrypoints/test_vllm_nonroot_entrypoint.sh @@ -43,6 +43,7 @@ run_wrapper() { _env="$_env $1"; shift done shift + # shellcheck disable=SC2086 # word splitting is intentional: $_env holds multiple NAME=value assignments env -i PATH="$WORKDIR/bin:/usr/bin:/bin" $_env "$WRAPPER" "$@" > "$_out" } @@ -140,9 +141,8 @@ fake_passwd="$WORKDIR/fake-passwd-prepopulated" printf 'vllm:x:%s:%s:vllm:/home/vllm:/bin/bash\n' "$current_uid" "$current_gid" > "$fake_passwd" out="$WORKDIR/case6.out" run_wrapper "$out" "HOME=$case5_home" "VLLM_PASSWD_FILE=$fake_passwd" -- --model foo -line_count="$(wc -l < "$fake_passwd")" -# NOTE: wc may count 0 or 1 depending on trailing newline; accept 1. -# More robust: count lines matching our UID. +# wc may count 0 or 1 depending on trailing newline, so count lines +# matching our UID instead for a robust check. uid_lines="$(grep -c ":${current_uid}:" "$fake_passwd" || true)" [ "$uid_lines" = "1" ] \ || { echo "FAIL: case6: expected exactly one entry for UID $current_uid, got $uid_lines"; cat "$fake_passwd"; exit 1; } diff --git a/docs/pre_run_check.sh b/docs/pre_run_check.sh index d611e616071e..d7cf214993a9 100644 --- a/docs/pre_run_check.sh +++ b/docs/pre_run_check.sh @@ -1,3 +1,4 @@ +# shellcheck disable=SC2317 # code after the exit/exit-183 gate below is intentionally dead pending re-enablement (see line comment) if [ "$READTHEDOCS_VERSION_TYPE" != "external" ]; then echo "Not a PR build (version type=$READTHEDOCS_VERSION_TYPE); skipping pre-run-check gate." exit 0 diff --git a/tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh b/tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh index d3fee8d6ba58..ddd8d06b0198 100755 --- a/tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh +++ b/tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh @@ -52,10 +52,10 @@ CUDA_VISIBLE_DEVICES=$PREFILL_GPU_ID \ VLLM_SSM_CONV_STATE_LAYOUT=DS \ VLLM_KV_CACHE_LAYOUT=HND \ VLLM_NIXL_SIDE_CHANNEL_PORT=5559 \ -vllm serve $MODEL \ +vllm serve "$MODEL" \ --port $PREFILL_PORT \ --enforce-eager \ - --gpu-memory-utilization $GPU_MEMORY_UTILIZATION \ + --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \ --max-model-len 16384 \ --block-size 128 \ --trust-remote-code \ @@ -70,10 +70,10 @@ CUDA_VISIBLE_DEVICES=$DECODE_GPU_ID \ VLLM_SSM_CONV_STATE_LAYOUT=DS \ VLLM_KV_CACHE_LAYOUT=HND \ VLLM_NIXL_SIDE_CHANNEL_PORT=6000 \ -vllm serve $MODEL \ +vllm serve "$MODEL" \ --port $DECODE_PORT \ --enforce-eager \ - --gpu-memory-utilization $GPU_MEMORY_UTILIZATION \ + --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \ --max-model-len 16384 \ --block-size 128 \ --trust-remote-code \ diff --git a/tests/v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh b/tests/v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh index 4d4512b19c4f..87603d42384b 100644 --- a/tests/v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh +++ b/tests/v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh @@ -86,10 +86,10 @@ launch_baseline() { --gpu-memory-utilization ${GPU_MEMORY_UTILIZATION} \ --dtype float16 \ --enforce-eager" - echo ${BASELINE_BASE_CMD} + echo "${BASELINE_BASE_CMD}" bash -c "${BASELINE_BASE_CMD}" & sleep 10 - wait_for_server ${BASELINE_HOST} ${BASELINE_PORT} + wait_for_server "${BASELINE_HOST}" "${BASELINE_PORT}" } launch_pd() { @@ -97,53 +97,50 @@ launch_pd() { PREFILL_GPU_ID=$(compute_gpu_id 0 "${PREFILLER_TP_SIZE}") local DECODE_GPU_ID DECODE_GPU_ID=$(compute_gpu_id "${PREFILLER_TP_SIZE}" "${DECODER_TP_SIZE}") + local kv_config="{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"${KV_BUFFER_DEVICE}\"}" - PREFILL_BASE_CMD=" - ZE_AFFINITY_MASK=$PREFILL_GPU_ID \ + sleep 2 + + echo "Starting prefill on GPU ${PREFILL_GPU_ID}, port ${PREFILL_PORT}" + ZE_AFFINITY_MASK="${PREFILL_GPU_ID}" \ VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=200 \ - VLLM_NIXL_SIDE_CHANNEL_HOST=${PREFILL_HOST} \ - VLLM_NIXL_SIDE_CHANNEL_PORT=${PREFILL_NIXL_SIDE_PORT} \ + VLLM_NIXL_SIDE_CHANNEL_HOST="${PREFILL_HOST}" \ + VLLM_NIXL_SIDE_CHANNEL_PORT="${PREFILL_NIXL_SIDE_PORT}" \ VLLM_WORKER_MULTIPROC_METHOD=spawn \ - VLLM_ENABLE_V1_MULTIPROCESSING=1 vllm serve $MODEL_NAME \ - --host ${PREFILL_HOST} \ - --port ${PREFILL_PORT} \ - --max-model-len ${MAX_MODEL_LEN}\ - --seed 42 \ - --block-size ${BLOCK_SIZE} \ - --enforce-eager \ - --dtype float16 \ - -tp ${PREFILLER_TP_SIZE} \ - --gpu-memory-utilization ${GPU_MEMORY_UTILIZATION} \ - --kv-transfer-config '{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\"}'" - - - DECODE_BASE_CMD=" - ZE_AFFINITY_MASK=$DECODE_GPU_ID \ + VLLM_ENABLE_V1_MULTIPROCESSING=1 \ + vllm serve "$MODEL_NAME" \ + --host "${PREFILL_HOST}" \ + --port "${PREFILL_PORT}" \ + --max-model-len "${MAX_MODEL_LEN}" \ + --seed 42 \ + --block-size "${BLOCK_SIZE}" \ + --enforce-eager \ + --dtype float16 \ + -tp "${PREFILLER_TP_SIZE}" \ + --gpu-memory-utilization "${GPU_MEMORY_UTILIZATION}" \ + --kv-transfer-config "${kv_config}" & + + echo "Starting decode on GPU ${DECODE_GPU_ID}, port ${DECODE_PORT}" + ZE_AFFINITY_MASK="${DECODE_GPU_ID}" \ VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=200 \ VLLM_WORKER_MULTIPROC_METHOD=spawn \ - VLLM_ENABLE_V1_MULTIPROCESSING=1 vllm serve $MODEL_NAME \ - --host ${DECODE_HOST} \ - --port ${DECODE_PORT} \ - --max-model-len ${MAX_MODEL_LEN}\ - --seed 42 \ - --block-size ${BLOCK_SIZE} \ - --enforce-eager \ - -tp ${DECODER_TP_SIZE} \ - --dtype float16 \ - --gpu-memory-utilization ${GPU_MEMORY_UTILIZATION} \ - --kv-transfer-config '{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\"}'" - - echo ${PREFILL_BASE_CMD} - echo ${DECODE_BASE_CMD} - sleep 2 + VLLM_ENABLE_V1_MULTIPROCESSING=1 \ + vllm serve "$MODEL_NAME" \ + --host "${DECODE_HOST}" \ + --port "${DECODE_PORT}" \ + --max-model-len "${MAX_MODEL_LEN}" \ + --seed 42 \ + --block-size "${BLOCK_SIZE}" \ + --enforce-eager \ + -tp "${DECODER_TP_SIZE}" \ + --dtype float16 \ + --gpu-memory-utilization "${GPU_MEMORY_UTILIZATION}" \ + --kv-transfer-config "${kv_config}" & - # execute on hosts - bash -c "${PREFILL_BASE_CMD}" & - bash -c "${DECODE_BASE_CMD}" & sleep 1 - wait_for_server ${PREFILL_HOST} ${PREFILL_PORT} + wait_for_server "${PREFILL_HOST}" "${PREFILL_PORT}" sleep 1 - wait_for_server ${DECODE_HOST} ${DECODE_PORT} + wait_for_server "${DECODE_HOST}" "${DECODE_PORT}" sleep 1 } @@ -153,7 +150,7 @@ launch_pd_proxy(){ --prefiller-host ${PREFILL_HOST} --prefiller-port ${PREFILL_PORT} \ --decoder-host ${DECODE_HOST} --decoder-port ${DECODE_PORT} \ --host=${PROXY_HOST} --port ${PROXY_PORT}" - echo ${PROXY_BASE_CMD} + echo "${PROXY_BASE_CMD}" bash -c "${PROXY_BASE_CMD}" & sleep 2 } @@ -161,7 +158,7 @@ launch_pd_proxy(){ run_tests(){ local service_url=$1 local mode=$2 - python3 ${EXP_ROOT}/test_disagg_accuracy.py --service_url=${service_url} --model_name=${MODEL_NAME} --mode=${mode} --file_name=${OUTPUT_FILE} + python3 "${EXP_ROOT}/test_disagg_accuracy.py" --service_url="${service_url}" --model_name="${MODEL_NAME}" --mode="${mode}" --file_name="${OUTPUT_FILE}" } @@ -178,7 +175,7 @@ launch_pd_proxy run_tests "http://${PROXY_HOST}:${PROXY_PORT}" "disagg" echo "-----P/D success----" -rm ${OUTPUT_FILE} +rm "${OUTPUT_FILE}" cleanup exit 0 diff --git a/tests/v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh b/tests/v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh index 501ec332d69e..ec98ed3008ce 100755 --- a/tests/v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh +++ b/tests/v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh @@ -103,8 +103,10 @@ fi cleanup_instances() { echo "" echo "Cleaning up..." + # shellcheck disable=SC2046 # word splitting is intentional for multiple PIDs kill $(jobs -pr) 2>/dev/null || true sleep 1 + # shellcheck disable=SC2046 kill -9 $(jobs -pr) 2>/dev/null || true pkill -9 -f "vllm serve.*${MODEL_NAME}" 2>/dev/null || true pkill -9 -f "toy_proxy_server.*8192" 2>/dev/null || true @@ -122,7 +124,7 @@ wait_for_server() { local deadline=${5:-600} local elapsed=0 echo "Waiting for ${server_name} on port ${port}..." - while [ $elapsed -lt $deadline ]; do + while [ "$elapsed" -lt "$deadline" ]; do if ! ps -p "$server_pid" > /dev/null 2>&1; then local status=0 wait "$server_pid" || status=$? @@ -186,7 +188,7 @@ else else num=1 fi - for (( g=0; g/dev/null | head -1) if [ -z "$BUILT_WHEEL" ]; then echo "Error: No wheel produced" diff --git a/tools/pre_commit/shellcheck.sh b/tools/pre_commit/shellcheck.sh deleted file mode 100755 index 557f41f293b7..000000000000 --- a/tools/pre_commit/shellcheck.sh +++ /dev/null @@ -1,24 +0,0 @@ -#!/bin/bash -set -euo pipefail - -scversion="stable" - -if [ -d "shellcheck-${scversion}" ]; then - export PATH="$PATH:$(pwd)/shellcheck-${scversion}" -fi - -if ! [ -x "$(command -v shellcheck)" ]; then - if [ "$(uname -s)" != "Linux" ] || [ "$(uname -m)" != "x86_64" ]; then - echo "Please install shellcheck: https://github.com/koalaman/shellcheck?tab=readme-ov-file#installing" - exit 1 - fi - - # automatic local install if linux x86_64 - wget -qO- "https://github.com/koalaman/shellcheck/releases/download/${scversion?}/shellcheck-${scversion?}.linux.x86_64.tar.xz" | tar -xJv - export PATH="$PATH:$(pwd)/shellcheck-${scversion}" -fi - -# TODO - fix warnings in .buildkite/scripts/hardware_ci/run-amd-test.sh -find . -path ./.git -prune -o -name "*.sh" \ - -not -path "./.buildkite/scripts/hardware_ci/run-amd-test.sh" -print0 | \ - xargs -0 sh -c "for f in \"\$@\"; do git check-ignore -q \"\$f\" || shellcheck -s bash \"\$f\"; done" -- diff --git a/tools/setup_deepgemm_pythons.sh b/tools/setup_deepgemm_pythons.sh index d98ea7e0370d..c6095a36b40c 100755 --- a/tools/setup_deepgemm_pythons.sh +++ b/tools/setup_deepgemm_pythons.sh @@ -17,7 +17,8 @@ if [ "$#" -eq 0 ]; then | grep -oE '>=3\.[0-9]+,<3\.[0-9]+') lo=${spec#>=3.}; lo=${lo%%,*} hi=${spec##*<3.} - set -- $(seq "$lo" $((hi - 1)) | sed 's/^/3./') + readarray -t versions < <(seq "$lo" $((hi - 1)) | sed 's/^/3./') + set -- "${versions[@]}" fi prefix="${DEEPGEMM_VENV_PREFIX:-/tmp/dgenv}" diff --git a/vllm/utils/numa_wrapper.sh b/vllm/utils/numa_wrapper.sh index 541801ed5df5..3c8b38919c4f 100755 --- a/vllm/utils/numa_wrapper.sh +++ b/vllm/utils/numa_wrapper.sh @@ -15,6 +15,7 @@ if ! command -v numactl >/dev/null 2>&1; then exit 1 fi +# shellcheck disable=SC1001 # verified equivalent behavior; TODO: revisit escaping in a follow-up cleanup PR case "${_VLLM_INTERNAL_NUMACTL_ARGS}" in *[![:alnum:]\ \-\_=,./]*) echo "Invalid characters in _VLLM_INTERNAL_NUMACTL_ARGS" >&2 @@ -22,4 +23,5 @@ case "${_VLLM_INTERNAL_NUMACTL_ARGS}" in ;; esac +# shellcheck disable=SC2086 # word splitting is intentional: this expands into multiple numactl flags exec numactl ${_VLLM_INTERNAL_NUMACTL_ARGS} "${_VLLM_INTERNAL_NUMACTL_PYTHON_EXECUTABLE}" "$@" From 6e85feb1b7592e2bb73ff9227b2e529b8a270ec2 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Thu, 20 Aug 2026 11:24:41 +0100 Subject: [PATCH 204/839] [3/N] Harden Transformers modelling backend multi-modal path (#51827) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- .../processing/test_transformers_audio.py | 62 +- .../processing/test_transformers_image.py | 247 ++++++- .../processing/transformers_backend.py | 51 ++ .../models/transformers/multimodal.py | 639 +++++++++++++++--- 4 files changed, 838 insertions(+), 161 deletions(-) create mode 100644 tests/models/multimodal/processing/transformers_backend.py diff --git a/tests/models/multimodal/processing/test_transformers_audio.py b/tests/models/multimodal/processing/test_transformers_audio.py index bc0dedcd1d27..61517ea0af19 100644 --- a/tests/models/multimodal/processing/test_transformers_audio.py +++ b/tests/models/multimodal/processing/test_transformers_audio.py @@ -4,8 +4,11 @@ import pytest from vllm.config import ModelConfig +from vllm.model_executor.models.transformers.multimodal import LegacyMultiModalProcessor from vllm.multimodal import MULTIMODAL_REGISTRY +from .transformers_backend import PROCESSOR_CLASSES, create_processor + AUDIO_MODEL_SETTINGS = { "ibm-granite/granite-speech-3.3-2b": { "prompt": ( @@ -53,6 +56,7 @@ } +@pytest.mark.parametrize("processor_cls", PROCESSOR_CLASSES) @pytest.mark.parametrize( "model_id", [ @@ -61,12 +65,12 @@ pytest.param( "mistralai/Voxtral-Mini-3B-2507", marks=pytest.mark.xfail( - reason="MistralCommonBackend.encode does not produce the audio " - "placeholder token (ID 24) from raw text. apply_chat_template " - "yields token IDs with placeholders, but MultiModalProcessor." - "apply() decodes the prompt back to text and re-tokenizes, at " - "which point the placeholders are lost. Fix belongs in " - "mistral_common or in the Voxtral-specific path.", + reason="Voxtral's mistral_common processor does not compose with " + "the Transformers modelling backend. Loading it currently fails " + "outright, because MistralCommonBackend.from_pretrained rejects " + "the kwargs vLLM passes, and it implements no " + "`replace_audio_token`, so it would report no replacement " + "offsets. Both fixes belong in mistral_common or transformers.", strict=False, ), ), @@ -74,15 +78,10 @@ "zai-org/GLM-ASR-Nano-2512", ], ) -def test_audio_multimodal_processor(model_id): +def test_audio_multimodal_processor(model_id, processor_cls): settings = AUDIO_MODEL_SETTINGS[model_id] - model_config = ModelConfig( - model=model_id, - model_impl="transformers", - ) - - mm_processor = MULTIMODAL_REGISTRY.create_processor(model_config) + mm_processor = create_processor(model_id, processor_cls) audio = np.zeros(16000, dtype=np.float32) mm_data = {"audio": (audio, 16000)} @@ -113,10 +112,13 @@ def test_audio_multimodal_processor(model_id): ) -def _process_granite_speech(separator: str): +@pytest.mark.parametrize("processor_cls", PROCESSOR_CLASSES) +@pytest.mark.parametrize("separator", [" and ", ""]) +def test_audio_multiple_inputs(separator, processor_cls): + """Multiple audios per prompt are each detected as a separate placeholder + and multi-modal item by the Transformers modelling backend.""" model_id = "ibm-granite/granite-speech-3.3-2b" - model_config = ModelConfig(model=model_id, model_impl="transformers") - mm_processor = MULTIMODAL_REGISTRY.create_processor(model_config) + mm_processor = create_processor(model_id, processor_cls) audio_token = mm_processor.info.get_hf_processor().audio_token # One token per audio; the processor expands each to its placeholder run. @@ -126,17 +128,21 @@ def _process_granite_speech(separator: str): ) audios = [np.zeros(16000, dtype=np.float32), np.zeros(24000, dtype=np.float32)] - return mm_processor( - prompt=prompt, - mm_items=mm_processor.info.parse_mm_data({"audio": audios}), - hf_processor_mm_kwargs={}, - ) + def process(): + return mm_processor( + prompt=prompt, + mm_items=mm_processor.info.parse_mm_data({"audio": audios}), + hf_processor_mm_kwargs={}, + ) + # The legacy path reads placeholders off contiguous runs of the audio token, so + # it cannot tell adjacent ones apart and says so instead of merging them + if processor_cls is LegacyMultiModalProcessor and not separator: + with pytest.raises(ValueError, match="Separate them in the prompt"): + process() + return -def test_audio_multiple_inputs(): - """Multiple audios per prompt are each detected as a separate placeholder - and multi-modal item by the Transformers modelling backend.""" - result = _process_granite_speech(separator=" and ") + result = process() assert len(result["mm_placeholders"]["audio"]) == 2 assert len(result["mm_kwargs"]["audio"]) == 2 @@ -167,9 +173,3 @@ def test_unclaimed_fields_warn_rather_than_raise(): assert owned["audio"] == ["input_features"] assert owned["image"] == [] - - -def test_audio_adjacent_inputs(): - """Adjacent audios are rejected rather than silently merged into one placeholder.""" - with pytest.raises(ValueError, match="told apart"): - _process_granite_speech(separator="") diff --git a/tests/models/multimodal/processing/test_transformers_image.py b/tests/models/multimodal/processing/test_transformers_image.py index 6aa42fc71c98..f5de7d66f4ff 100644 --- a/tests/models/multimodal/processing/test_transformers_image.py +++ b/tests/models/multimodal/processing/test_transformers_image.py @@ -1,20 +1,27 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from unittest.mock import patch + import pytest from vllm.assets.image import ImageAsset -from vllm.config import ModelConfig -from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.model_executor.models.transformers.multimodal import ( + LegacyMultiModalProcessor, + OffsetsMultiModalProcessor, +) +from .transformers_backend import ( + PROCESSOR_CLASSES, + create_cached_processor, + create_processor, + offsets_only, +) -@pytest.mark.parametrize("model_id", ["llava-hf/llava-onevision-qwen2-0.5b-ov-hf"]) -def test_multimodal_processor(model_id): - model_config = ModelConfig( - model=model_id, - model_impl="transformers", - ) - mm_processor = MULTIMODAL_REGISTRY.create_processor(model_config) +@pytest.mark.parametrize("processor_cls", PROCESSOR_CLASSES) +@pytest.mark.parametrize("model_id", ["llava-hf/llava-onevision-qwen2-0.5b-ov-hf"]) +def test_multimodal_processor(model_id, processor_cls): + mm_processor = create_processor(model_id, processor_cls) image_pil = ImageAsset("cherry_blossom").pil_image mm_data = {"image": image_pil} @@ -56,10 +63,9 @@ def test_multimodal_processor(model_id): ) -def _process_two_images(separator: str): +def _process_two_images(processor_cls, separator: str): model_id = "llava-hf/llava-onevision-qwen2-0.5b-ov-hf" - model_config = ModelConfig(model=model_id, model_impl="transformers") - mm_processor = MULTIMODAL_REGISTRY.create_processor(model_config) + mm_processor = create_processor(model_id, processor_cls) image = ImageAsset("cherry_blossom").pil_image prompt = ( @@ -74,28 +80,147 @@ def _process_two_images(separator: str): ) -def test_image_multiple_inputs(): +@pytest.mark.parametrize("processor_cls", PROCESSOR_CLASSES) +def test_image_multiple_inputs(processor_cls): """Multiple images per prompt are each detected as a separate placeholder and multi-modal item by the Transformers modelling backend.""" - result = _process_two_images(separator="\n and ") + result = _process_two_images(processor_cls, separator="\n and ") assert len(result["mm_placeholders"]["image"]) == 2 assert len(result["mm_kwargs"]["image"]) == 2 -def test_image_adjacent_inputs(): +@pytest.mark.parametrize("processor_cls", PROCESSOR_CLASSES) +def test_image_adjacent_inputs(processor_cls): """Adjacent images stay separate placeholders rather than merging into one.""" - result = _process_two_images(separator="") + result = _process_two_images(processor_cls, separator="") assert len(result["mm_placeholders"]["image"]) == 2 assert len(result["mm_kwargs"]["image"]) == 2 -def test_text_only_prompt(): +@pytest.mark.parametrize("processor_cls", PROCESSOR_CLASSES) +def test_batch_padding_removed_from_image_items(processor_cls): + """Emu3 pads every image up to the largest in the batch, which would leave an + item's data dependent on what it was processed with and so uncacheable.""" + mm_processor = create_processor("BAAI/Emu3-Chat-hf", processor_cls) + image_token = mm_processor.info.get_hf_processor().image_token + + images = [ + ImageAsset("cherry_blossom").pil_image, + ImageAsset("cherry_blossom").pil_image.resize((256, 1024)), + ] + result = mm_processor( + prompt=f"{image_token} and {image_token}", + mm_items=mm_processor.info.parse_mm_data({"image": images}), + hf_processor_mm_kwargs={}, + ) + + items = result["mm_kwargs"]["image"] + shapes = set() + for item in items: + height, width = item["image_sizes"].data.flatten().tolist() + pixel_values = item["pixel_values"].data + assert tuple(pixel_values.shape[-2:]) == (height, width) + shapes.add(tuple(pixel_values.shape)) + + # Both images would have been padded to a common shape had they been kept + assert len(shapes) == 2 + + +def _process_one_gemma3_image(processor_cls): + mm_processor = create_processor("google/gemma-3-4b-it", processor_cls) + hf_processor = mm_processor.info.get_hf_processor() + result = mm_processor( + prompt=f"{hf_processor.boi_token} What is this?", + mm_items=mm_processor.info.parse_mm_data( + {"image": ImageAsset("cherry_blossom").pil_image} + ), + hf_processor_mm_kwargs={}, + ) + return hf_processor, result + + +@offsets_only +def test_non_embedding_tokens_excluded_from_placeholders(): + """Gemma3 wraps each image in text that carries no embeddings, which must be + inside the placeholder range but masked out of it.""" + _, result = _process_one_gemma3_image(OffsetsMultiModalProcessor) + + (placeholder,) = result["mm_placeholders"]["image"] + assert placeholder.is_embed is not None + assert 0 < int(placeholder.is_embed.sum()) < placeholder.length + + +def test_legacy_placeholders_hold_only_image_tokens(): + """The legacy path spans whatever `mm_token_type_ids` attributes to the image, + which for Gemma3 excludes the text wrapping it, unlike the replacement the offsets + path spans. Gemma3 is also the sharp case for the mask: its `image_token_id` is the + marker in the unexpanded prompt, not the token the expansion repeats.""" + hf_processor, result = _process_one_gemma3_image(LegacyMultiModalProcessor) + + (placeholder,) = result["mm_placeholders"]["image"] + assert placeholder.length == hf_processor.image_seq_length + prompt_ids = result["prompt_token_ids"] + covered = prompt_ids[placeholder.offset : placeholder.offset + placeholder.length] + assert set(covered) == {hf_processor.tokenizer.image_token_id} + + +@pytest.mark.parametrize("processor_cls", PROCESSOR_CLASSES) +def test_tokens_structuring_an_image_are_masked_not_dropped(processor_cls): + """SmolVLM splits each image into tiles introduced by tokens carrying no + embeddings. Those belong inside the placeholder and masked out, because the token + count the processor reports is over the whole span. Idefics3 also refuses a prompt + holding `` when no images are passed, which is how the offsets path has to + tokenize it before splicing in the expansion.""" + mm_processor = create_processor( + "HuggingFaceTB/SmolVLM-256M-Instruct", processor_cls + ) + result = mm_processor( + prompt="What is this?", + mm_items=mm_processor.info.parse_mm_data( + {"image": ImageAsset("cherry_blossom").pil_image} + ), + hf_processor_mm_kwargs={}, + ) + + (placeholder,) = result["mm_placeholders"]["image"] + assert placeholder.is_embed is not None + assert 0 < int(placeholder.is_embed.sum()) < placeholder.length + + +@offsets_only +def test_missing_replacement_offsets_names_the_processor(): + """A processor that reports no replacement offsets cannot be served, which must + be said plainly rather than surfacing later as a field config mismatch.""" + model_id = "llava-hf/llava-onevision-qwen2-0.5b-ov-hf" + mm_processor = create_processor(model_id, OffsetsMultiModalProcessor) + hf_processor_cls = type(mm_processor.info.get_hf_processor()) + hf_call = hf_processor_cls.__call__ + + def without_offsets(self, *args, **kwargs): + hf_inputs = hf_call(self, *args, **kwargs) + hf_inputs.pop("text_replacement_offsets", None) + return hf_inputs + + with ( + patch.object(hf_processor_cls, "__call__", without_offsets), + pytest.raises(ValueError, match="LlavaOnevisionProcessor returned no"), + ): + mm_processor( + prompt="\nWhat is the content of this image?", + mm_items=mm_processor.info.parse_mm_data( + {"image": ImageAsset("cherry_blossom").pil_image} + ), + hf_processor_mm_kwargs={}, + ) + + +@pytest.mark.parametrize("processor_cls", PROCESSOR_CLASSES) +def test_text_only_prompt(processor_cls): """An image model still accepts a prompt with no images.""" model_id = "llava-hf/llava-onevision-qwen2-0.5b-ov-hf" - model_config = ModelConfig(model=model_id, model_impl="transformers") - mm_processor = MULTIMODAL_REGISTRY.create_processor(model_config) + mm_processor = create_processor(model_id, processor_cls) result = mm_processor( prompt="<|im_start|>user Hello!<|im_end|><|im_start|>assistant\n", @@ -105,3 +230,87 @@ def test_text_only_prompt(): assert len(result["prompt_token_ids"]) > 0 assert not result["mm_placeholders"] + + +@offsets_only +def test_repeated_image_hits_the_processor_cache(): + """Check that mm caching is actually working.""" + mm_processor, cache = create_cached_processor( + "llava-hf/llava-onevision-qwen2-0.5b-ov-hf", OffsetsMultiModalProcessor + ) + image = ImageAsset("cherry_blossom").pil_image + + def process(): + return mm_processor( + prompt="\nWhat is this?", + mm_items=mm_processor.info.parse_mm_data({"image": image}), + hf_processor_mm_kwargs={}, + ) + + first, second = process(), process() + + assert cache.make_stats().hits > 0 + assert first["prompt_token_ids"] == second["prompt_token_ids"] + assert first["mm_hashes"] == second["mm_hashes"] + + +@offsets_only +@pytest.mark.parametrize( + ("model_id", "prompt"), + [ + ("llava-hf/llava-onevision-qwen2-0.5b-ov-hf", "\nWhat is this?"), + ("google/gemma-3-4b-it", " What is this?"), + ("HuggingFaceTB/SmolVLM-256M-Instruct", "What is this?"), + pytest.param( + "BAAI/Emu3-Chat-hf", + " and more text", + marks=pytest.mark.xfail( + reason="Emu3Processor prepends its BOS token only when images are " + "passed, so the unexpanded prompt the offsets path tokenizes never " + "gets one. Fixed by huggingface/transformers#47924, unreleased.", + strict=False, + ), + ), + ], +) +def test_spliced_prompt_matches_hf_expansion(model_id, prompt): + """The offsets path splices the expansion into a prompt tokenized without any + multi-modal data, so its token ids have to come out the same as the ones the HF + processor produces itself, which is what the legacy path returns.""" + prompt_ids = [] + for processor_cls in (LegacyMultiModalProcessor, OffsetsMultiModalProcessor): + mm_processor = create_processor(model_id, processor_cls) + prompt_ids.append( + mm_processor( + prompt=prompt, + mm_items=mm_processor.info.parse_mm_data( + {"image": ImageAsset("cherry_blossom").pil_image} + ), + hf_processor_mm_kwargs={}, + )["prompt_token_ids"] + ) + + legacy_ids, offsets_ids = prompt_ids + assert legacy_ids == offsets_ids + + +@pytest.mark.parametrize("processor_cls", PROCESSOR_CLASSES) +def test_nested_image_fields_split_per_image(processor_cls): + """Idefics3 returns image fields with a leading batch dimension, putting the rows + belonging to each image one dimension further in. Slicing the batch dimension + instead handed the first image every row and the second an empty tensor.""" + mm_processor = create_processor( + "HuggingFaceTB/SmolVLM-256M-Instruct", processor_cls + ) + image = ImageAsset("cherry_blossom").pil_image + result = mm_processor( + prompt=" and ", + mm_items=mm_processor.info.parse_mm_data({"image": [image, image]}), + hf_processor_mm_kwargs={}, + ) + + items = result["mm_kwargs"]["image"] + assert len(items) == 2 + for item in items: + pixel_values = item["pixel_values"].data + assert pixel_values.shape[1] == int(item["num_image_patches"].data) diff --git a/tests/models/multimodal/processing/transformers_backend.py b/tests/models/multimodal/processing/transformers_backend.py new file mode 100644 index 000000000000..dc49036e8352 --- /dev/null +++ b/tests/models/multimodal/processing/transformers_backend.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Helpers for testing both processors of the Transformers modelling backend.""" + +import pytest + +from vllm.config import ModelConfig +from vllm.model_executor.models.transformers.multimodal import ( + LegacyMultiModalProcessor, + MultiModalDummyInputsBuilder, + MultiModalProcessingInfo, + MultiModalProcessor, + OffsetsMultiModalProcessor, +) +from vllm.multimodal.cache import MultiModalProcessorOnlyCache +from vllm.multimodal.processing import InputProcessingContext +from vllm.tokenizers.registry import cached_tokenizer_from_config + +offsets_only = pytest.mark.skipif( + MultiModalProcessor is not OffsetsMultiModalProcessor, + reason="Replacement offsets are only used from transformers 5.15.0 onwards", +) + +PROCESSOR_CLASSES = [ + pytest.param(LegacyMultiModalProcessor, id="legacy"), + pytest.param(OffsetsMultiModalProcessor, id="offsets", marks=offsets_only), +] + + +def create_processor(model_id: str, processor_cls): + """Build a processor directly, because the registry only ever builds the one the + installed transformers version selects, leaving the other path untested.""" + model_config = ModelConfig(model=model_id, model_impl="transformers") + ctx = InputProcessingContext( + model_config, cached_tokenizer_from_config(model_config) + ) + info = MultiModalProcessingInfo(ctx) + return processor_cls(info, MultiModalDummyInputsBuilder(info)) + + +def create_cached_processor(model_id: str, processor_cls): + """Build a processor backed by a real multi-modal processor cache, and hand the + cache back so a test can check it was actually used.""" + model_config = ModelConfig(model=model_id, model_impl="transformers") + model_config.multimodal_config.mm_processor_cache_gb = 4 + ctx = InputProcessingContext( + model_config, cached_tokenizer_from_config(model_config) + ) + info = MultiModalProcessingInfo(ctx) + cache = MultiModalProcessorOnlyCache(model_config) + return processor_cls(info, MultiModalDummyInputsBuilder(info), cache=cache), cache diff --git a/vllm/model_executor/models/transformers/multimodal.py b/vllm/model_executor/models/transformers/multimodal.py index e94490232c08..8b6568bf3c60 100644 --- a/vllm/model_executor/models/transformers/multimodal.py +++ b/vllm/model_executor/models/transformers/multimodal.py @@ -22,6 +22,8 @@ from typing import TYPE_CHECKING, Any import torch +import transformers +from packaging.version import Version from vllm.compilation.decorators import should_torch_compile_mm_encoder from vllm.config.utils import getattr_iter @@ -50,7 +52,9 @@ BaseMultiModalProcessor, BaseProcessingInfo, ProcessorInputs, + PromptReplacement, PromptUpdate, + PromptUpdateDetails, TimingContext, ) from vllm.sequence import IntermediateTensors @@ -65,6 +69,12 @@ logger = init_logger(__name__) _MODALITY_TO_TOKEN_TYPE_ID = {"image": 1, "video": 2, "audio": 3} +_MODALITY_SIZE_KEYS = {"audio": "num_audio_tokens", "image": "num_image_patches"} + + +def _get_embed_token_id(replacement_ids: torch.Tensor) -> int: + """The token an expansion repeats is the one holding the embeddings.""" + return int(replacement_ids.mode().values) class MultiModalProcessingInfo(BaseProcessingInfo): @@ -94,21 +104,6 @@ def _get_supported_modalities(self) -> list[str]: ) return modalities - def _get_audio_token_id(self) -> int: - processor = self.get_hf_processor() - if hasattr(processor, "audio_token_id"): - return processor.audio_token_id - config = self.get_hf_config() - val = getattr_iter(config, ("audio_token_id", "audio_token_index")) - if val is not None: - return val - if hasattr(processor, "audio_token"): - tokenizer = self.get_tokenizer() - vocab = tokenizer.get_vocab() - if processor.audio_token in vocab: - return vocab[processor.audio_token] - raise ValueError("Cannot find audio_token_id on processor or model config") - def _get_audio_sampling_rate(self) -> float: sub = self._get_audio_processor() if sub is not None and hasattr(sub, "sampling_rate"): @@ -169,7 +164,7 @@ def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: if self.info._is_audio_model() and (num_audios := mm_counts.get("audio", 0)): processor = self.info.get_hf_processor() audio_token = getattr(processor, "audio_token", "") - # Separated so that `_apply_audio` can tell the placeholders apart + # Separated so that adjacent placeholders stay distinguishable text += " ".join([audio_token] * num_audios) if self.info._is_image_model() and (num_images := mm_counts.get("image", 0)): processor = self.info.get_hf_processor() @@ -210,23 +205,26 @@ def get_dummy_mm_data( return data -class MultiModalProcessor(BaseMultiModalProcessor[MultiModalProcessingInfo]): - def _get_prompt_updates( +class _MultiModalProcessorBase(BaseMultiModalProcessor[MultiModalProcessingInfo]): + """Processing common to both Transformers backend processors: calling the HF + processor, sizing images, and attributing its outputs to a modality. + + Subclasses add the strategy for locating placeholders in the prompt. + """ + + def _get_hf_mm_data( self, mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - out_mm_kwargs: MultiModalKwargsItems, - ) -> Sequence[PromptUpdate]: - """No updates: `apply` locates placeholders via `mm_token_type_ids` instead. - - HF processors have no generic contract for the token sequence they insert, - so it cannot be expressed as a `PromptUpdate`. Returning nothing is only - safe because `apply` is overridden; the base class would reject it. - """ - return [] + ) -> tuple[Mapping[str, object], Mapping[str, object]]: + """Rename the parser's `audios` key to the `audio` argument HF audio + processors take.""" + processor_data, passthrough_data = super()._get_hf_mm_data(mm_items) + if self.info._is_audio_model() and "audios" in processor_data: + processor_data["audio"] = processor_data.pop("audios") + return processor_data, passthrough_data def _get_modality_field_names(self, modality: str) -> set[str]: - """Field names the sub-processor for `modality` produces.""" + """Names of the fields the sub-processor for `modality` produces.""" # TODO: use else branch only once huggingface/transformers#44394 lands. if modality == "audio": sub_processor = self.info._get_audio_processor() @@ -273,6 +271,18 @@ def _partition_keys_by_modality( return owned + def _get_slice_dim(self, data: Any, total_rows: int) -> int: + """Which dimension of a field holds the rows belonging to each item. + + Some processors (e.g., Idefics3) return image fields with a leading batch + dimension, putting the rows one dimension further in. + """ + if not isinstance(data, torch.Tensor) or data.ndim < 2: + return 0 + if data.shape[0] != total_rows and data.shape[1] == total_rows: + return 1 + return 0 + def _get_mm_fields_config( self, hf_inputs: "BatchFeature", @@ -281,29 +291,50 @@ def _get_mm_fields_config( # HF Processors always return a mask but vLLM doesn't need it hf_inputs.pop("attention_mask", None) - # Written by `_apply_audio`/`_apply_vision`; absent if the modality had no items + # Absent if the modality had no items sizes = { - "audio": hf_inputs.get("num_audio_tokens"), - "image": hf_inputs.get("num_image_patches"), + modality: hf_inputs.get(key) + for modality, key in _MODALITY_SIZE_KEYS.items() } modalities = [m for m, size in sizes.items() if size is not None] - size_keys = {"num_audio_tokens", "num_image_patches"} - keys = [key for key in hf_inputs if key not in size_keys] + # Keys we wrote ourselves, rather than ones a sub-processor produced + own_keys = set(_MODALITY_SIZE_KEYS.values()) | { + f"{modality}_replacement_{suffix}" + for modality in modalities + for suffix in ("ids", "sizes") + } + keys = [key for key in hf_inputs if key not in own_keys] owned = self._partition_keys_by_modality(keys, modalities) + # Un-padded fields are already one entry per item, so index rather than slice mm_fields: dict[str, MultiModalFieldConfig] = { - key: MultiModalFieldConfig.flat_from_sizes(modality, sizes[modality]) + key: MultiModalFieldConfig.batched(modality) + if modality == "audio" or isinstance(hf_inputs[key], list) + else MultiModalFieldConfig.flat_from_sizes( + modality, + sizes[modality], + dim=self._get_slice_dim(hf_inputs[key], int(sizes[modality].sum())), + ) for modality in modalities for key in owned[modality] } - # Keep these as batched, as they always have batch size as first dim - if "audio" in modalities: - mm_fields["num_audio_tokens"] = MultiModalFieldConfig.batched( - "audio", keep_on_cpu=True + for modality in modalities: + # One row per item, and only ever read on the CPU + mm_fields[_MODALITY_SIZE_KEYS[modality]] = MultiModalFieldConfig.batched( + modality, keep_on_cpu=True ) + replacement_sizes = hf_inputs.get(f"{modality}_replacement_sizes") + if replacement_sizes is not None: + mm_fields[f"{modality}_replacement_ids"] = ( + MultiModalFieldConfig.flat_from_sizes( + modality, replacement_sizes, keep_on_cpu=True + ) + ) + if "image" in modalities: + # Always one row per item, whatever they describe mm_fields["image_grid_thw"] = MultiModalFieldConfig.batched( "image", keep_on_cpu=True ) @@ -311,25 +342,154 @@ def _get_mm_fields_config( mm_fields["video_grid_thw"] = MultiModalFieldConfig.batched( "image", keep_on_cpu=True ) - mm_fields["num_image_patches"] = MultiModalFieldConfig.batched( - "image", keep_on_cpu=True - ) + return mm_fields - def _get_hf_mm_data( + def _call_hf_processor( self, - mm_items: MultiModalDataItems, - ) -> tuple[Mapping[str, object], Mapping[str, object]]: + prompt: str, + mm_data: Mapping[str, object], + mm_kwargs: Mapping[str, object], + tok_kwargs: Mapping[str, object], + ) -> "BatchFeature": + """Run the HF processor and unpad inputs.""" + try: + hf_inputs = super()._call_hf_processor( + prompt, mm_data, mm_kwargs, tok_kwargs + ) + except ValueError: + if any(mm_data.values()): + raise + # Some processors reject a prompt holding placeholders with + # no data to go with them, so tokenize it without them + prompt_ids = self.info.get_tokenizer().encode(prompt, **tok_kwargs) + hf_inputs = transformers.BatchFeature( + dict(input_ids=[prompt_ids]), tensor_type="pt" + ) + self._unpad_images(hf_inputs) + self._unpad_audios(hf_inputs, mm_data, mm_kwargs, tok_kwargs) + return hf_inputs + + def _unpad_audios( + self, + hf_inputs: "BatchFeature", + mm_data: Mapping[str, object], + mm_kwargs: Mapping[str, object], + tok_kwargs: Mapping[str, object], + ) -> None: + """Replace the audio fields with each audio processed on its own. + + Processors pad every audio up to the longest in the call, which would leave + an item's data dependent on what it was processed with. Unlike images, + nothing in the output states how long each one really is, and processors + pad a lone audio too, so the only way to know what an audio produces by + itself is to process it by itself. + """ + audios = mm_data.get("audio") + if not audios or len({len(audio) for audio in audios}) == 1: + return + + alone = [ + self.info.ctx.call_hf_processor( + self.info.get_hf_processor(**mm_kwargs), + dict( + text=self.dummy_inputs.get_dummy_text({"audio": 1}), audio=[audio] + ), + dict(**mm_kwargs, **tok_kwargs), + ) + for audio in audios + ] + + for key in self._get_modality_field_names("audio"): + if isinstance(hf_inputs.get(key), torch.Tensor): + hf_inputs[key] = [output[key][0] for output in alone] + + def _unpad_images(self, hf_inputs: "BatchFeature") -> None: + """Trim each image back to its own size when the processor padded them all + to the largest in the batch. + + An image's data has to depend on nothing but that image, or the multi-modal + processor cache would store it under that image's hash and later reuse it + beside a different neighbour. Padding is re-applied when the encoder runs. """ - In contrast to the base class, this method requests - `return_mm_token_type_ids` and remaps the `audios` key to `audio` for - audio models. + pixel_values = hf_inputs.get("pixel_values") + image_sizes = hf_inputs.get("image_sizes") + if not isinstance(pixel_values, torch.Tensor): + return + if not isinstance(image_sizes, torch.Tensor): + return + if pixel_values.ndim != 4 or len(pixel_values) != len(image_sizes): + return + + # The sizes describe the trailing dimensions only if the largest of them is + # what the batch was padded up to. Otherwise they mean something else, as + # in llava-onevision, where they are the sizes before any processing. + maxima = image_sizes.max(dim=0).values + if maxima.tolist() != list(pixel_values.shape[-2:]): + return + + hf_inputs["pixel_values"] = [ + image[..., :height, :width] + for image, (height, width) in zip(pixel_values, image_sizes.tolist()) + ] + + +class LegacyMultiModalProcessor(_MultiModalProcessorBase): + """Locates placeholders by searching the prompt the HF processor has already + expanded for the tokens of each modality. + + Serves transformers versions with no `return_text_replacement_offsets`. + Placeholders found this way cannot be rebuilt from an unexpanded prompt, so + this processor overrides `apply` and gets no multi-modal processor cache. + Remove it once `requirements/common.txt` requires `transformers>=5.15.0`. + """ + + def _get_prompt_updates( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + out_mm_kwargs: MultiModalKwargsItems, + ) -> Sequence[PromptUpdate]: + """Empty, because `apply` writes the placeholder ranges itself rather than + deriving them from updates.""" + return [] + + def _call_hf_processor( + self, + prompt: str, + mm_data: Mapping[str, object], + mm_kwargs: Mapping[str, object], + tok_kwargs: Mapping[str, object], + ) -> "BatchFeature": + """Ask which modality owns each token of the expanded prompt, which is the + only marking of the tokens an expansion adds around an item.""" + if any(mm_data.values()): + mm_data = {**mm_data, "return_mm_token_type_ids": True} + return super()._call_hf_processor(prompt, mm_data, mm_kwargs, tok_kwargs) + + def _get_mm_token_ids(self, modality: str) -> list[int]: + """Token ids marking where `modality` sits in the prompt, which for some + processors differ from the placeholder written into it. + + The expanded prompt is all this path has to go on, so it takes the + processor at its word about which tokens belong to the modality. """ - processor_data, passthrough_data = super()._get_hf_mm_data(mm_items) - if self.info._is_audio_model() and "audios" in processor_data: - processor_data["audio"] = processor_data.pop("audios") - processor_data["return_mm_token_type_ids"] = True - return processor_data, passthrough_data + info = self.info + processor = info.get_hf_processor() + declared = getattr(processor, f"{modality}_token_ids", None) or () + if ids := [token_id for token_id in declared if token_id is not None]: + return ids + config = info.get_hf_config() + names = (f"{modality}_token_id", f"{modality}_token_index") + token_id = getattr(processor, names[0], getattr_iter(config, names)) + if token_id is None: + token = getattr(processor, f"{modality}_token", None) + token_id = info.get_tokenizer().get_vocab().get(token) + if token_id is None: + raise ValueError( + f"Cannot find {modality}_token_id on processor or model config" + ) + return [token_id] def _apply_audio( self, @@ -337,12 +497,17 @@ def _apply_audio( processed_data: "BatchFeature", num_audios: int, ) -> dict[str, list[PlaceholderRange]]: - audio_token_id = self.info._get_audio_token_id() + """Take each contiguous run of the audio token as one item's placeholder, + and record how many tokens the run holds.""" + audio_token_ids = self._get_mm_token_ids("audio") prompt_tensor = torch.tensor(prompt_ids) - is_audio = prompt_tensor == audio_token_id + is_audio = torch.isin(prompt_tensor, torch.tensor(audio_token_ids)) if not is_audio.any(): - return {} + raise ValueError( + f"{num_audios} audio item(s) were passed but the prompt " + "contains no audio token. Add one placeholder per audio item." + ) padded = torch.cat([torch.tensor([False]), is_audio, torch.tensor([False])]) transitions = padded.int().diff() @@ -365,6 +530,23 @@ def _apply_audio( processed_data["num_audio_tokens"] = lengths return {"audio": ranges} + def _get_num_multimodal_tokens( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + ) -> Mapping[str, list[int]]: + """Ask the HF processor how many tokens and patches each image expands to.""" + processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) + images = mm_items.get_items("image", ImageProcessorItems) + image_sizes = [ + (size.height, size.width) + for size in map(images.get_image_size, range(len(images))) + ] + return processor._get_num_multimodal_tokens( + image_sizes=image_sizes, + **self.info.ctx.get_merged_mm_kwargs({}), + ) + def _apply_vision( self, prompt_ids: list[int], @@ -373,47 +555,46 @@ def _apply_vision( hf_processor_mm_kwargs: Mapping[str, object], mm_token_type_ids: torch.Tensor | None, ) -> dict[str, list[PlaceholderRange]]: - # Placeholders can't be located without them, so give up rather than guess - if mm_token_type_ids is None: - return {} - + """Split the positions the processor marks as image into one placeholder per + item, sized by the token count it reports for each image.""" hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) + if mm_token_type_ids is None: + raise ValueError( + f"{type(hf_processor).__name__} returned no `mm_token_type_ids`, so " + "the Transformers modeling backend cannot locate the placeholder of " + "each image." + ) - # We can infer vLLM style placeholder from token type ids, if we split - # it for each input `mm_data`. - mm_positions = torch.where(mm_token_type_ids == 1)[1] - images = mm_items.get_items("image", ImageProcessorItems) - image_sizes = [] - for item_idx in range(len(images)): - image_size = images.get_image_size(item_idx) - image_sizes.append((image_size.height, image_size.width)) - - mm_tokens_per_modality = hf_processor._get_num_multimodal_tokens( - image_sizes=image_sizes, - **self.info.ctx.get_merged_mm_kwargs({}), + # These mark the tokens structuring an image as well as the image itself, + # which is what the counts below are over + mm_positions = torch.where(mm_token_type_ids[0] == 1)[0] + mm_tokens_per_modality = self._get_num_multimodal_tokens( + mm_items, hf_processor_mm_kwargs ) mm_placeholders: dict[str, list[PlaceholderRange]] = {} split_sizes = mm_tokens_per_modality["num_image_tokens"] - if split_sizes: - image_token_ids = getattr(hf_processor, "image_token_ids", None) - if image_token_ids is None: - # Transformers <5.10.0 - image_token_ids = [hf_processor.image_token_id] - image_token_ids = torch.tensor( - [i for i in image_token_ids if i is not None] + if sum(split_sizes) != len(mm_positions): + raise ValueError( + f"The expanded prompt holds {len(mm_positions)} image token(s) but " + f"{type(hf_processor).__name__} accounts for {sum(split_sizes)} " + f"across {mm_items.get_count('image')} image item(s)." ) - chunked_mm_positions = torch.split(mm_positions, split_sizes) - mm_tokens = torch.tensor(prompt_ids)[mm_token_type_ids[0].bool()] - chunked_mm_tokens = torch.split(mm_tokens, split_sizes) + if split_sizes: + image_token_ids = torch.tensor(self._get_mm_token_ids("image")) + mm_tokens = torch.tensor(prompt_ids)[mm_positions] ranges = [ PlaceholderRange( offset=positions[0].item(), length=positions.shape[0], - is_embed=torch.isin(mm_tokens, image_token_ids), + # Only some of the span carries embeddings + is_embed=torch.isin(tokens, image_token_ids), + ) + for positions, tokens in zip( + torch.split(mm_positions, split_sizes), + torch.split(mm_tokens, split_sizes), ) - for positions, mm_tokens in zip(chunked_mm_positions, chunked_mm_tokens) ] mm_placeholders = {"image": ranges} @@ -427,12 +608,8 @@ def apply( inputs: ProcessorInputs, timing_ctx: TimingContext, ) -> MultiModalInput: - """ - Process multi-modal inputs to be used in vLLM. - - Apply HF Processor on prompt text and multi-modal data together, - outputting token IDs and processed tensors. - """ + """Process the prompt and every multi-modal item in one HF processor call, + then read the placeholder ranges out of the token ids it returns.""" prompt = inputs.prompt mm_items = inputs.mm_data_items hf_processor_mm_kwargs = inputs.hf_processor_mm_kwargs @@ -468,7 +645,7 @@ def apply( self.info.ctx.get_mm_config().mm_hasher_algorithm, ) - # For gemma3 we check `token_type_ids` as the key + # Gemma3 reports them under the key the model uses for its own token types mm_token_type_ids = processed_data.pop("token_type_ids", None) mm_token_type_ids = processed_data.pop("mm_token_type_ids", mm_token_type_ids) @@ -516,6 +693,201 @@ def apply( ) +class OffsetsMultiModalProcessor(_MultiModalProcessorBase): + """Locates placeholders from the `text_replacement_offsets` the HF processor + reports, expressing each one as a `PromptUpdate`. + + Stating the expansion as an update is what lets it be rebuilt from an + unexpanded prompt, so this processor takes the base class's processing path + and with it the multi-modal processor cache. + """ + + def _get_prompt_updates( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + out_mm_kwargs: MultiModalKwargsItems, + ) -> Sequence[PromptUpdate]: + """Replace each modality's placeholder token with the token ids that item's + replacement text encodes to, marking which of them hold embeddings.""" + hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) + updates = [] + for modality, items in out_mm_kwargs.items(): + # Popped so they are neither cached nor sent to the model; the updates + # they produce are cached alongside the item instead + replacements = [ + PromptUpdateDetails.select_token_id( + (ids := item.pop(f"{modality}_replacement_ids").data).tolist(), + _get_embed_token_id(ids), + ) + for item in items + ] + updates.append( + PromptReplacement( + modality=modality, + target=getattr(hf_processor, f"{modality}_token"), + replacement=replacements.__getitem__, + ) + ) + return updates + + def _apply_hf_processor_main( + self, + prompt: str | list[int], + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + tokenization_kwargs: Mapping[str, object], + *, + enable_hf_prompt_update: bool, + ) -> tuple[list[int], "BatchFeature", bool]: + """Tokenize the prompt unexpanded, leaving the expansion for vLLM to splice + in, whatever `enable_hf_prompt_update` asks for. + + This differs from `super()` only for a text prompt with + `enable_hf_prompt_update`, where `super()` would keep the expanded token + ids the HF processor returns. Both routes give the same ids, so forcing + this one costs an extra processor call and buys the guarantee that a + request bypassing the cache is identical to a cached one. + + A text prompt only arrives with `enable_hf_prompt_update` when there is + no multi-modal processor cache. Requests carrying pre-computed + embeddings bypass the cache too, but `--enable-mm-embeds` is unsupported + here, so they fail validation shortly afterwards. + """ + if isinstance(prompt, str) and enable_hf_prompt_update: + logger.warning_once( + "Disabling the multi-modal processor cache is extra slow with the " + "Transformers modeling backend: the prompt is still tokenized " + "unexpanded and the expansion spliced in, to keep the token ids " + "identical to what the cache produces, which costs an extra HF " + "processor call per request." + ) + return super()._apply_hf_processor_main( + prompt, + mm_items, + hf_processor_mm_kwargs, + tokenization_kwargs, + enable_hf_prompt_update=False, + ) + + def _get_num_image_patches( + self, + hf_inputs: "BatchFeature", + mm_data: Mapping[str, object], + num_images: int, + ) -> torch.Tensor: + """How many rows of the image fields belong to each image. + + Taken from whichever per-image count the processor reports, + and checked against the data it has to slice. + """ + if (grid := hf_inputs.get("image_grid_thw")) is not None: + num_patches = grid.prod(-1) + elif (counts := self._get_num_patches_per_image(mm_data)) is not None: + num_patches = torch.tensor(counts) + else: + num_patches = torch.ones(num_images, dtype=torch.long) + + image_data = hf_inputs.get("pixel_values", hf_inputs.get("image_patches")) + if isinstance(image_data, torch.Tensor): + total = int(num_patches.sum()) + rows = image_data.shape[self._get_slice_dim(image_data, total)] + if rows != total: + raise ValueError( + f"{type(self.info.get_hf_processor()).__name__} returned " + f"{rows} row(s) of image data for {num_images} image(s), which " + f"cannot be split into the {num_patches.tolist()} row(s) per " + "image derived from its outputs, so the rows cannot be " + "attributed to an image. Gemma3 does this when " + "`do_pan_and_scan` crops an image." + ) + return num_patches + + def _get_num_patches_per_image( + self, mm_data: Mapping[str, object] + ) -> list[int] | None: + """Ask the HF processor how many rows of image data each image produces.""" + images = mm_data.get("images") + if not images: + return None + try: + sizes = [(image.height, image.width) for image in images] + mm_tokens = self.info.get_hf_processor()._get_num_multimodal_tokens( + image_sizes=sizes, **self.info.ctx.get_merged_mm_kwargs({}) + ) + return list(mm_tokens["num_image_patches"]) + except (AttributeError, KeyError, TypeError): + return None + + def _call_hf_processor( + self, + prompt: str, + mm_data: Mapping[str, object], + mm_kwargs: Mapping[str, object], + tok_kwargs: Mapping[str, object], + ) -> "BatchFeature": + """Ask for the replacement each placeholder expands to, and record it as + per-item fields: its token ids, and the tokens or patches behind them.""" + if has_mm_data := any(mm_data.values()): + mm_data = {**mm_data, "return_text_replacement_offsets": True} + hf_inputs = super()._call_hf_processor(prompt, mm_data, mm_kwargs, tok_kwargs) + # Drop the inputs the model would reject + hf_inputs.pop("mm_token_type_ids", None) + hf_inputs.pop("token_type_ids", None) + + offsets = hf_inputs.pop("text_replacement_offsets", None) + # Some processors return an empty batch as a tensor rather than a list + if offsets is None or len(offsets) == 0 or len(offsets[0]) == 0: + if has_mm_data: + raise ValueError( + f"{type(self.info.get_hf_processor()).__name__} returned no " + "text replacement offsets, so the Transformers modeling backend " + "cannot locate the placeholder of each item. Its `__call__` has " + "to reach `ProcessorMixin.get_text_with_replacements` with one " + "replacement per item, which usually means implementing " + "`replace__token`. Please report this to transformers " + "so it can be fixed, and install `transformers<5.15.0` in the " + "meantime to locate placeholders in the expanded prompt instead." + ) + return hf_inputs + + tokenizer = self.info.get_tokenizer() + replacements = defaultdict[str, list[list[int]]](list) + for entry in offsets[0]: + replacements[entry["type"]].append( + tokenizer.encode(entry["replacement"], add_special_tokens=False) + ) + + for modality, seqs in replacements.items(): + hf_inputs[f"{modality}_replacement_ids"] = torch.tensor( + [token_id for seq in seqs for token_id in seq] + ) + hf_inputs[f"{modality}_replacement_sizes"] = torch.tensor( + [len(seq) for seq in seqs] + ) + if modality == "image": + hf_inputs["num_image_patches"] = self._get_num_image_patches( + hf_inputs, mm_data, len(seqs) + ) + elif modality == "audio": + counts = [] + for seq in seqs: + ids = torch.tensor(seq) + counts.append(int(ids.eq(_get_embed_token_id(ids)).sum())) + hf_inputs["num_audio_tokens"] = torch.tensor(counts) + + return hf_inputs + + +# From this version on, a processor reporting no offsets is an error rather than a +# fallback to searching the expanded prompt +MultiModalProcessor = ( + LegacyMultiModalProcessor + if Version(transformers.__version__) < Version("5.15.0") + else OffsetsMultiModalProcessor +) + + class MultiModalMixin(SupportsMultiModal, SupportsMRoPE): def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""): # Skip SupportsMRoPE.__init__ and call the next class in MRO @@ -733,18 +1105,32 @@ def _process_audio_input(self, **kwargs) -> list[torch.Tensor] | None: kwargs.pop("token_type_ids", None) kwargs.pop("mm_token_type_ids", None) - # HuggingFace's `get_audio_features` implementations branch on - # per-sample feature lengths internally. - with gpu_sync_allowed(): - audio_output = self.model.get_audio_features( - input_features, return_dict=True, **kwargs - ) - audio_embeddings = audio_output.pooler_output - # Per-audio token counts are needed as Python ints to split. with gpu_sync_allowed(): split_sizes = num_audio_tokens.flatten().tolist() - return self._split_embeddings(audio_embeddings, split_sizes) + if isinstance(input_features, torch.Tensor): + # HuggingFace's `get_audio_features` implementations branch on + # per-sample feature lengths internally. + with gpu_sync_allowed(): + audio_output = self.model.get_audio_features( + input_features, return_dict=True, **kwargs + ) + return self._split_embeddings(audio_output.pooler_output, split_sizes) + + # Audios the processor left un-padded arrive as a list once their + # lengths differ. Encode them one at a time so that none of them is + # padded to match another. + embeddings: list[torch.Tensor] = [] + for index, features in enumerate(input_features): + audio_output = self.model.get_audio_features( + features.unsqueeze(0), + return_dict=True, + **self._select_item_kwargs(kwargs, index, len(input_features)), + ) + embeddings.extend( + self._split_embeddings(audio_output.pooler_output, [split_sizes[index]]) + ) + return embeddings def _process_image_input(self, **kwargs) -> list[torch.Tensor] | None: pixel_values: torch.Tensor | None = kwargs.pop("pixel_values", None) @@ -761,6 +1147,42 @@ def _process_image_input(self, **kwargs) -> list[torch.Tensor] | None: num_image_patches = kwargs.pop("num_image_patches") + split_sizes = num_image_patches.flatten().tolist() + if isinstance(pixel_values, torch.Tensor): + vision_embeddings = self._get_image_features(pixel_values, **kwargs) + if isinstance(vision_embeddings, torch.Tensor): + return self._split_embeddings(vision_embeddings, split_sizes) + return list(vision_embeddings) + + # Images the processor left un-padded arrive as a list once their + # shapes differ. Encode them one at a time so that none of them is + # padded to match another. + embeddings: list[torch.Tensor] = [] + for index, image in enumerate(pixel_values): + features = self._get_image_features( + image.unsqueeze(0), + **self._select_item_kwargs(kwargs, index, len(pixel_values)), + ) + # Encoders which return one entry per image return a single entry + if not isinstance(features, torch.Tensor): + features = torch.cat(list(features)) + embeddings.extend(self._split_embeddings(features, [split_sizes[index]])) + return embeddings + + def _select_item_kwargs( + self, kwargs: dict[str, Any], index: int, num_items: int + ) -> dict[str, Any]: + """Narrow the entries of `kwargs` that hold one row per item down to the item + at `index`. Length is all there is to match on, so an unrelated entry of the + same length is narrowed too.""" + return { + key: value[index : index + 1] + if isinstance(value, (torch.Tensor, list)) and len(value) == num_items + else value + for key, value in kwargs.items() + } + + def _get_image_features(self, pixel_values: torch.Tensor, **kwargs) -> Any: # grid_thw fields are registered keep_on_cpu; restore the on-device # placement that HF get_image_features implementations expect. for key, value in kwargs.items(): @@ -772,22 +1194,17 @@ def _process_image_input(self, **kwargs) -> list[torch.Tensor] | None: # padding images via boolean-mask indexing, LlavaOnevision # branches on per-sample batch counts). with gpu_sync_allowed(): - vision_embeddings = self.model.get_image_features(pixel_values, **kwargs) + features = self.model.get_image_features(pixel_values, **kwargs) # Transformers `v5`, `self.get_image_features` returns a tuple # containing the features and optionally attentions/hidden_states # After v5 is settled, we can enable qwen3-vl with several outputs # from `self.get_image_features` - if isinstance(vision_embeddings, tuple): - vision_embeddings = vision_embeddings[0] - elif isinstance(vision_embeddings, dict): - vision_embeddings = vision_embeddings.pooler_output - - if isinstance(vision_embeddings, torch.Tensor): - split_sizes = num_image_patches.flatten().tolist() - return self._split_embeddings(vision_embeddings, split_sizes) - - return list(vision_embeddings) + if isinstance(features, tuple): + return features[0] + if isinstance(features, dict): + return features.pooler_output + return features def embed_multimodal(self, **kwargs) -> MultiModalEmbeddings: # Each helper detects its own inputs. We are called once per modality, so the From f0c14b4f776bb976e654b430442067539fbdb2ea Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Thu, 20 Aug 2026 11:31:55 +0100 Subject: [PATCH 205/839] Fix weight tying (#51665) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- .../model_loader/test_weight_tying.py | 86 ++++++++++++++++ tests/models/test_utils.py | 78 +++++++++++++++ tests/test_config.py | 39 ++++++++ tests/transformers_utils/test_config.py | 26 ++++- tests/transformers_utils/test_repo_utils.py | 11 +++ vllm/config/model.py | 33 +++++++ vllm/config/vllm.py | 10 ++ vllm/model_executor/model_loader/utils.py | 5 + .../model_loader/weight_tying.py | 99 +++++++++++++++++++ vllm/model_executor/models/apertus.py | 5 +- vllm/model_executor/models/arcee.py | 6 +- vllm/model_executor/models/arctic.py | 5 +- vllm/model_executor/models/bailing_moe.py | 5 +- vllm/model_executor/models/bloom.py | 2 +- vllm/model_executor/models/chameleon.py | 3 +- vllm/model_executor/models/cohere2_moe.py | 1 - vllm/model_executor/models/commandr.py | 3 - vllm/model_executor/models/ernie45_moe.py | 5 +- vllm/model_executor/models/ernie_mtp.py | 3 +- vllm/model_executor/models/exaone.py | 8 +- vllm/model_executor/models/exaone4.py | 8 +- vllm/model_executor/models/exaone_moe.py | 7 +- vllm/model_executor/models/fairseq2_llama.py | 5 +- vllm/model_executor/models/falcon.py | 5 +- vllm/model_executor/models/falcon_h1.py | 5 +- vllm/model_executor/models/gemma.py | 7 +- vllm/model_executor/models/gemma2.py | 7 +- vllm/model_executor/models/gemma3.py | 5 +- vllm/model_executor/models/gemma4.py | 3 - vllm/model_executor/models/glm4.py | 3 +- vllm/model_executor/models/gpt_oss.py | 5 +- vllm/model_executor/models/granite.py | 7 +- vllm/model_executor/models/granitemoe.py | 3 +- .../model_executor/models/granitemoeshared.py | 3 +- vllm/model_executor/models/hrm_text.py | 3 +- vllm/model_executor/models/hunyuan_v1.py | 5 +- vllm/model_executor/models/hunyuan_vision.py | 5 +- vllm/model_executor/models/hy_v3.py | 5 +- vllm/model_executor/models/hyperclovax.py | 5 +- vllm/model_executor/models/internlm2.py | 5 +- .../model_executor/models/iquest_loopcoder.py | 5 +- vllm/model_executor/models/jais2.py | 5 +- vllm/model_executor/models/laguna.py | 5 +- vllm/model_executor/models/lfm2.py | 5 +- vllm/model_executor/models/lfm2_moe.py | 5 +- vllm/model_executor/models/llama.py | 5 +- vllm/model_executor/models/llama4.py | 5 +- vllm/model_executor/models/mimo.py | 4 +- vllm/model_executor/models/minicpm.py | 5 +- vllm/model_executor/models/minicpm_eagle.py | 5 +- vllm/model_executor/models/mistral.py | 5 +- vllm/model_executor/models/mpt.py | 1 - vllm/model_executor/models/nemotron_nas.py | 5 +- vllm/model_executor/models/olmo3.py | 7 +- vllm/model_executor/models/olmo_hybrid.py | 7 +- vllm/model_executor/models/openpangu.py | 5 +- vllm/model_executor/models/opt.py | 7 +- vllm/model_executor/models/plamo3.py | 5 +- vllm/model_executor/models/qwen2.py | 5 +- vllm/model_executor/models/qwen3.py | 5 +- vllm/model_executor/models/rnj1.py | 5 +- vllm/model_executor/models/sarvam.py | 5 +- vllm/model_executor/models/seed_oss.py | 5 +- vllm/model_executor/models/telechat2.py | 5 +- .../models/transformers/causal.py | 4 - vllm/model_executor/models/utils.py | 60 +++++++++-- vllm/models/kimi_k3/amd/linear.py | 5 +- vllm/models/kimi_k3/nvidia/model.py | 5 +- vllm/transformers_utils/config.py | 22 ++++- vllm/transformers_utils/repo_utils.py | 13 +++ 70 files changed, 526 insertions(+), 238 deletions(-) create mode 100644 tests/model_executor/model_loader/test_weight_tying.py create mode 100644 vllm/model_executor/model_loader/weight_tying.py diff --git a/tests/model_executor/model_loader/test_weight_tying.py b/tests/model_executor/model_loader/test_weight_tying.py new file mode 100644 index 000000000000..8d68da5676bf --- /dev/null +++ b/tests/model_executor/model_loader/test_weight_tying.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.weight_tying import maybe_retie_word_embeddings + +VOCAB_SIZE = 16 +HIDDEN_SIZE = 4 + + +class UntiedModel(nn.Module): + """Nests the head like a multimodal model, which is the harder case.""" + + def __init__(self): + super().__init__() + self.language_model = nn.Module() + self.language_model.model = nn.Module() + self.language_model.model.embed_tokens = VocabParallelEmbedding( + VOCAB_SIZE, HIDDEN_SIZE + ) + self.language_model.lm_head = ParallelLMHead(VOCAB_SIZE, HIDDEN_SIZE) + self.embed_tokens.weight.data.fill_(1.0) + self.lm_head.weight.data.fill_(1.0) + + @property + def embed_tokens(self) -> VocabParallelEmbedding: + return self.language_model.model.embed_tokens + + @property + def lm_head(self) -> ParallelLMHead: + return self.language_model.lm_head + + +def make_model_config(untied_by_checkpoint=False): + return SimpleNamespace( + model="dummy-model", + word_embeddings_untied_by_checkpoint=untied_by_checkpoint, + ) + + +@pytest.mark.cpu_test +@pytest.mark.usefixtures("dist_init") +@pytest.mark.parametrize("identical", [True, False]) +def test_retie_only_when_identical(identical: bool): + """A redundant copy of a tied lm_head is shared again to reclaim memory.""" + model = UntiedModel() + if not identical: + model.lm_head.weight.data.fill_(2.0) + + maybe_retie_word_embeddings(model, make_model_config(untied_by_checkpoint=True)) + + assert (model.lm_head.weight is model.embed_tokens.weight) is identical + if not identical: + assert torch.all(model.lm_head.weight == 2.0) + + +@pytest.mark.cpu_test +@pytest.mark.usefixtures("dist_init") +def test_quantized_lm_head_is_left_alone(): + """A quantized head may store its weights packed under another name.""" + model = UntiedModel() + model.lm_head.quant_method = SimpleNamespace() + + maybe_retie_word_embeddings(model, make_model_config(untied_by_checkpoint=True)) + + assert model.lm_head.weight is not model.embed_tokens.weight + + +@pytest.mark.cpu_test +@pytest.mark.usefixtures("dist_init") +def test_no_retie_without_checkpoint_override(): + """Word embeddings the config genuinely unties are left alone.""" + model = UntiedModel() + + maybe_retie_word_embeddings(model, make_model_config()) + + assert model.lm_head.weight is not model.embed_tokens.weight diff --git a/tests/models/test_utils.py b/tests/models/test_utils.py index 8d47b4436575..2b32a63ac6e1 100644 --- a/tests/models/test_utils.py +++ b/tests/models/test_utils.py @@ -4,6 +4,10 @@ import pytest import torch +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) from vllm.model_executor.models.utils import ( AutoWeightsLoader, _merge_multimodal_embeddings, @@ -165,6 +169,80 @@ def weight_generator(): assert new_mod.nested_mod.bn.num_batches_tracked.item() == 1 +VOCAB_SIZE = 16 +HIDDEN_SIZE = 2 + + +class ModuleWithTiedWeights(torch.nn.Module): + """Mimics how models tie `lm_head` to the input embeddings.""" + + def __init__(self, tie: bool): + super().__init__() + self.model = torch.nn.Module() + self.model.embed_tokens = VocabParallelEmbedding(VOCAB_SIZE, HIDDEN_SIZE) + self.lm_head = ParallelLMHead(VOCAB_SIZE, HIDDEN_SIZE) + if tie: + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) + + +def make_embedding_weights(value: float) -> torch.Tensor: + return torch.full((VOCAB_SIZE, HIDDEN_SIZE), value) + + +@pytest.mark.cpu_test +@pytest.mark.usefixtures("dist_init") +@pytest.mark.parametrize("tie", [True, False]) +def test_module_skip_tied_weights(tie: bool): + """Tied weights must be loaded once, under the first of their names.""" + mod = ModuleWithTiedWeights(tie) + + weights = [ + ("model.embed_tokens.weight", make_embedding_weights(1.0)), + ("lm_head.weight", make_embedding_weights(2.0)), + ] + loaded = AutoWeightsLoader(mod).load_weights(iter(weights)) + + if tie: + assert loaded == {"model.embed_tokens.weight"} + assert torch.all(mod.lm_head.weight[:VOCAB_SIZE] == 1.0) + else: + assert loaded == {"model.embed_tokens.weight", "lm_head.weight"} + assert torch.all(mod.lm_head.weight[:VOCAB_SIZE] == 2.0) + + +@pytest.mark.cpu_test +@pytest.mark.usefixtures("dist_init") +def test_module_skip_tied_weights_without_canonical(): + """Skipping a tied weight must not leave the shared weight uninitialized.""" + mod = ModuleWithTiedWeights(tie=True) + + weights = [("lm_head.weight", make_embedding_weights(2.0))] + with pytest.raises(ValueError, match="model.embed_tokens.weight"): + AutoWeightsLoader(mod).load_weights(iter(weights)) + + +class ModuleWithSharedParam(torch.nn.Module): + """Mimics an MoE router shared between the MLP and its fused experts.""" + + def __init__(self): + super().__init__() + self.experts = torch.nn.Module() + self.experts.gate = torch.nn.Linear(2, 2, bias=False) + self.gate = self.experts.gate + + +@pytest.mark.cpu_test +def test_module_load_shared_params_that_are_not_tied_embeddings(): + """Only tied embeddings are skipped; other shared params must still load.""" + mod = ModuleWithSharedParam() + + weights = [("gate.weight", torch.Tensor([[1, 2], [3, 4]]))] + loaded = AutoWeightsLoader(mod).load_weights(iter(weights)) + + assert loaded == {"gate.weight"} + assert torch.all(mod.gate.weight == torch.Tensor([[1, 2], [3, 4]])) + + class raise_if_cuda_sync: def __enter__(self): self.previous_debug_mode = torch.cuda.get_sync_debug_mode() diff --git a/tests/test_config.py b/tests/test_config.py index 2d9c3b83f760..7ec18b3e400e 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -604,6 +604,45 @@ def test_with_hf_config_leaves_unknown_model_type_without_architectures( assert updated.model_config.hf_config.architectures is None +@pytest.mark.parametrize( + "checkpoint_tensors,tied", + [ + # The checkpoint has an lm_head of its own, so it must win over the config + (["model.embed_tokens.weight", "lm_head.weight"], False), + (["model.embed_tokens.weight"], True), + # Contents unknown (not safetensors), so the config must be left alone + ([], True), + ], +) +def test_maybe_untie_word_embeddings(tmp_path, checkpoint_tensors, tied): + import torch + from safetensors.torch import save_file + + if checkpoint_tensors: + save_file( + {name: torch.zeros(2, 2) for name in checkpoint_tensors}, + tmp_path / "model.safetensors", + ) + + text_config = SimpleNamespace(tie_word_embeddings=True) + model_config = SimpleNamespace( + model=str(tmp_path), + revision=None, + hf_config=SimpleNamespace( + tie_word_embeddings=True, + get_text_config=lambda: text_config, + ), + word_embeddings_untied_by_checkpoint=False, + ) + + ModelConfig.maybe_untie_word_embeddings(model_config) + + # Both levels must agree, since different callers read different ones + assert model_config.hf_config.tie_word_embeddings is tied + assert text_config.tie_word_embeddings is tied + assert model_config.word_embeddings_untied_by_checkpoint is not tied + + def test_async_scheduling_with_pipeline_parallelism_is_allowed(): cfg = VllmConfig( scheduler_config=SchedulerConfig( diff --git a/tests/transformers_utils/test_config.py b/tests/transformers_utils/test_config.py index 42261feb536f..35b5a697dbd8 100644 --- a/tests/transformers_utils/test_config.py +++ b/tests/transformers_utils/test_config.py @@ -8,14 +8,17 @@ from types import SimpleNamespace from typing import cast -from unittest.mock import patch +from unittest.mock import MagicMock, patch from transformers import PretrainedConfig from vllm.config.model import ModelConfig from vllm.tokenizers import get_tokenizer from vllm.transformers_utils import config as config_module -from vllm.transformers_utils.config import try_get_generation_config +from vllm.transformers_utils.config import ( + get_safetensors_params_metadata, + try_get_generation_config, +) def test_get_llama3_eos_token(): @@ -77,3 +80,22 @@ def test_model_config_generation_fallback_forwards_code_revision(): config_format="auto", token=None, ) + + +def test_safetensors_metadata_of_repo_without_safetensors(): + """A repo storing its weights in another format is an answer, not a failure, + so it must not be retried.""" + from huggingface_hub.errors import LocalEntryNotFoundError, NotASafetensorsRepoError + + get_safetensors_metadata = MagicMock( + side_effect=NotASafetensorsRepoError("not a safetensors repo") + ) + api = SimpleNamespace( + get_safetensors_metadata=get_safetensors_metadata, + snapshot_download=MagicMock(side_effect=LocalEntryNotFoundError("no cache")), + ) + + with patch.object(config_module, "hf_api", lambda: api): + assert get_safetensors_params_metadata("some/pytorch-only-model") == {} + + get_safetensors_metadata.assert_called_once() diff --git a/tests/transformers_utils/test_repo_utils.py b/tests/transformers_utils/test_repo_utils.py index 36d0acccd6b7..4a32734c534f 100644 --- a/tests/transformers_utils/test_repo_utils.py +++ b/tests/transformers_utils/test_repo_utils.py @@ -14,6 +14,7 @@ get_hf_file_to_dict, is_mistral_model_repo, list_filtered_repo_files, + with_retry, ) @@ -185,3 +186,13 @@ def _glob_path() -> list[str]: repo_type="model", token="token", ) + + +def test_with_retry_does_not_retry_fatal_errors(): + """A definitive answer must not be retried, which costs a delay per attempt.""" + func = MagicMock(side_effect=FileNotFoundError("no such file")) + + with pytest.raises(FileNotFoundError): + with_retry(func, "Error", fatal_errors=(FileNotFoundError,)) + + func.assert_called_once() diff --git a/vllm/config/model.py b/vllm/config/model.py index 56e533ea5305..37d2b7a7541a 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -30,6 +30,7 @@ from vllm.tasks import PoolingTask, ScoreType, SupportedTask from vllm.transformers_utils.config import ( ConfigFormat, + checkpoint_has_lm_head, get_config, get_hf_image_processor_config, get_hf_text_config, @@ -185,6 +186,10 @@ class ModelConfig: """The Hugging Face config of the model.""" hf_text_config: PretrainedConfig = field(init=False) """The Hugging Face config of the text model (same as hf_config for text models).""" + word_embeddings_untied_by_checkpoint: bool = field(default=False, init=False) + """Whether `tie_word_embeddings` was overridden to `False` because the checkpoint + contains an `lm_head` of its own. The two may still turn out to be identical, in + which case they are re-tied once the weights have been loaded.""" hf_config_path: str | None = None """Name or path of the Hugging Face config to use. If unspecified, model name or path will be used.""" @@ -1109,6 +1114,34 @@ def maybe_pull_model_tokenizer_for_runai(self, model: str, tokenizer: str) -> No def _get_encoder_config(self) -> dict[str, Any] | None: return get_sentence_transformer_tokenizer_config(self.model, self.revision) + def maybe_untie_word_embeddings(self) -> None: + """Stop trusting `tie_word_embeddings` when the checkpoint disagrees. + + A config may claim the word embeddings are tied while the checkpoint + ships an `lm_head` of its own. Tying regardless would silently discard + that tensor, so build the `lm_head` as if untied and let it load. The + two are compared once loaded, and re-tied if they turn out to match, by + [maybe_retie_word_embeddings][vllm.model_executor.model_loader.weight_tying.maybe_retie_word_embeddings]. + + Transformers makes the same decision in `PreTrainedModel.tie_weights`, + where it can compare the two tensors directly. + """ + if not getattr(self.hf_config, "tie_word_embeddings", False): + return + if not checkpoint_has_lm_head(self.model, revision=self.revision): + return + + logger.debug( + "The config for %s says the word embeddings are tied, but the checkpoint " + "contains an lm_head. Loading it to find out whether they really are tied.", + self.model, + ) + # Both levels must be updated: `VllmConfig.with_hf_config` reads the top + # level, and the text config is what the language model itself sees. + self.hf_config.tie_word_embeddings = False + self.hf_config.get_text_config().tie_word_embeddings = False + self.word_embeddings_untied_by_checkpoint = True + def _get_default_runner_type( self, architectures: list[str], diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index a271a654134a..ca009b91840a 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1170,6 +1170,16 @@ def __post_init__(self): self.model_config, self.load_config ) + # "dummy" reads no weights at all, and the sharded formats read a vLLM + # state dict, which stores tied word embeddings under the lm_head only. + # Neither can tell us what the original checkpoint contained. + if self.model_config is not None and self.load_config.load_format not in ( + "dummy", + "sharded_state", + "runai_streamer_sharded", + ): + self.model_config.maybe_untie_word_embeddings() + if ( self.quant_config is not None and self.model_config is not None diff --git a/vllm/model_executor/model_loader/utils.py b/vllm/model_executor/model_loader/utils.py index b7ccb78c04ee..d28051a54892 100644 --- a/vllm/model_executor/model_loader/utils.py +++ b/vllm/model_executor/model_loader/utils.py @@ -24,6 +24,7 @@ record_metadata_for_reloading, set_torchao_reload_attrs, ) +from vllm.model_executor.model_loader.weight_tying import maybe_retie_word_embeddings from vllm.model_executor.models.interfaces import SupportsQuant from vllm.tracing import instrument from vllm.utils.mem_utils import release_device_memory_under_pressure @@ -96,6 +97,10 @@ def initialize_model( def process_weights_after_loading( model: nn.Module, model_config: ModelConfig, target_device: torch.device ) -> None: + # Reclaim memory when an explicit lm_head has been + # loaded, but it is identical to the input embeddings. + maybe_retie_word_embeddings(model, model_config) + for _, module in model.named_modules(): quant_method = getattr(module, "quant_method", None) if isinstance(quant_method, QuantizeMethodBase): diff --git a/vllm/model_executor/model_loader/weight_tying.py b/vllm/model_executor/model_loader/weight_tying.py new file mode 100644 index 000000000000..d0c309d2c4ad --- /dev/null +++ b/vllm/model_executor/model_loader/weight_tying.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Reconcile word embedding tying with what the checkpoint actually contains.""" + +from dataclasses import dataclass + +import torch +from torch import nn + +from vllm.config import ModelConfig +from vllm.logger import init_logger +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + UnquantizedEmbeddingMethod, + VocabParallelEmbedding, +) + +logger = init_logger(__name__) + + +@dataclass +class _UntiedLMHead: + name: str + lm_head: ParallelLMHead + embed_tokens: VocabParallelEmbedding + + @property + def weight_name(self) -> str: + return f"{self.name}.weight" + + def tie(self, model: nn.Module) -> None: + parent_name, _, attr = self.name.rpartition(".") + parent = model.get_submodule(parent_name) if parent_name else model + setattr(parent, attr, self.lm_head.tie_weights(self.embed_tokens)) + + +def _get_untied_lm_head(model: nn.Module) -> _UntiedLMHead | None: + """Locate an `lm_head` that could be tied to the input embeddings. + + Returns `None` unless the model has exactly one of each, so that models with + several heads (such as MTP) or with already tied weights are left alone. + Both are found by type rather than by name, because `get_input_embeddings` + takes different arguments on multimodal models and the `lm_head` of a + multimodal model is nested inside its language model. + + A quantized `lm_head` is also left alone. Its weights may be packed under + another name, and online quantization creates them after loading, so + neither their contents nor whether they were loaded can be established here. + Note that this is a property of the layer, not of the model: most quantized + models leave the `lm_head` and the input embeddings unquantized. + """ + heads = list[tuple[str, ParallelLMHead]]() + embeddings = list[VocabParallelEmbedding]() + for name, module in model.named_modules(): + if not isinstance(module, VocabParallelEmbedding): + continue + if not isinstance(module.quant_method, UnquantizedEmbeddingMethod): + continue + if isinstance(module, ParallelLMHead): + heads.append((name, module)) + else: + embeddings.append(module) + + if len(heads) != 1 or len(embeddings) != 1: + return None + + (name, lm_head), embed_tokens = heads[0], embeddings[0] + if lm_head.weight.shape != embed_tokens.weight.shape: + return None + return _UntiedLMHead(name, lm_head, embed_tokens) + + +def maybe_retie_word_embeddings(model: nn.Module, model_config: ModelConfig) -> None: + """Re-tie word embeddings that + [ModelConfig.maybe_untie_word_embeddings][vllm.config.ModelConfig.maybe_untie_word_embeddings] + untied, if the loaded `lm_head` turned out to be identical to the input embeddings + after all. + + Checkpoints produced by quantization or fine-tuning tooling often keep a redundant + copy of the tied `lm_head`. Sharing the storage again reclaims the memory it would + otherwise cost.""" + if not model_config.word_embeddings_untied_by_checkpoint: + return + if (untied := _get_untied_lm_head(model)) is None: + return + + # On device, torch.equal segfaults on ROCm when sleep mode is enabled + if not torch.equal(untied.lm_head.weight.cpu(), untied.embed_tokens.weight.cpu()): + logger.warning( + "The config for %s says the word embeddings are tied, but the checkpoint " + "contains a different %s, which has been used instead of tying. " + "Set `tie_word_embeddings=False` in the config to silence this warning.", + model_config.model, + untied.weight_name, + ) + return + + logger.debug("Re-tying %s, which is identical to the input embeddings", untied.name) + untied.tie(model) diff --git a/vllm/model_executor/models/apertus.py b/vllm/model_executor/models/apertus.py index 74b17129c989..b904802a3a65 100644 --- a/vllm/model_executor/models/apertus.py +++ b/vllm/model_executor/models/apertus.py @@ -491,8 +491,5 @@ def compute_logits( return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/arcee.py b/vllm/model_executor/models/arcee.py index c32a903bba8b..81cf3c5e0f64 100644 --- a/vllm/model_executor/models/arcee.py +++ b/vllm/model_executor/models/arcee.py @@ -356,11 +356,7 @@ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load weights into the model (delegates to inner model and handles tied embeddings).""" - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - skip_substrs=["gate_proj"], - ) + loader = AutoWeightsLoader(self, skip_substrs=["gate_proj"]) # AutoWeightLoader handles weight name remapping, including fusing # separate q_proj, k_proj, v_proj into qkv_proj return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/arctic.py b/vllm/model_executor/models/arctic.py index ccb3f1976cf6..555e58aa13d0 100644 --- a/vllm/model_executor/models/arctic.py +++ b/vllm/model_executor/models/arctic.py @@ -589,8 +589,5 @@ def compute_logits( return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights) diff --git a/vllm/model_executor/models/bailing_moe.py b/vllm/model_executor/models/bailing_moe.py index a3da7e39e78f..33749818b111 100644 --- a/vllm/model_executor/models/bailing_moe.py +++ b/vllm/model_executor/models/bailing_moe.py @@ -557,10 +557,7 @@ def _normalize_lm_head( yield name, loaded_weight def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(self._normalize_lm_head(weights)) diff --git a/vllm/model_executor/models/bloom.py b/vllm/model_executor/models/bloom.py index 1347b0e60538..8fb6bb4aa4eb 100644 --- a/vllm/model_executor/models/bloom.py +++ b/vllm/model_executor/models/bloom.py @@ -360,7 +360,7 @@ def compute_logits( return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self, skip_prefixes=["lm_head.weight"]) + loader = AutoWeightsLoader(self) weights = _add_transformer_prefix(weights) return loader.load_weights(weights) diff --git a/vllm/model_executor/models/chameleon.py b/vllm/model_executor/models/chameleon.py index 03c9fe1179e4..517697669c8e 100644 --- a/vllm/model_executor/models/chameleon.py +++ b/vllm/model_executor/models/chameleon.py @@ -1059,6 +1059,5 @@ def compute_logits( return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - skip_prefixes = ["lm_head."] if self.config.tie_word_embeddings else None - loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/cohere2_moe.py b/vllm/model_executor/models/cohere2_moe.py index 3873b6ee004f..9af1a3439b90 100644 --- a/vllm/model_executor/models/cohere2_moe.py +++ b/vllm/model_executor/models/cohere2_moe.py @@ -493,7 +493,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): config = vllm_config.model_config.hf_config quant_config = vllm_config.quant_config self.config = config - assert getattr(config, "tie_word_embeddings", True) self.unpadded_vocab_size = config.vocab_size self.quant_config = quant_config self.logits_scale = config.logit_scale diff --git a/vllm/model_executor/models/commandr.py b/vllm/model_executor/models/commandr.py index 5481fb0abcee..ca2543055db4 100644 --- a/vllm/model_executor/models/commandr.py +++ b/vllm/model_executor/models/commandr.py @@ -429,9 +429,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): quant_config = vllm_config.quant_config self.config = config - # currently all existing command R models have `tie_word_embeddings` - # enabled - assert config.tie_word_embeddings self.quant_config = quant_config self.logits_processor = LogitsProcessor( diff --git a/vllm/model_executor/models/ernie45_moe.py b/vllm/model_executor/models/ernie45_moe.py index 5cc351a6aa75..c7f26b6dd4a2 100644 --- a/vllm/model_executor/models/ernie45_moe.py +++ b/vllm/model_executor/models/ernie45_moe.py @@ -629,8 +629,5 @@ def compute_logits( return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights) diff --git a/vllm/model_executor/models/ernie_mtp.py b/vllm/model_executor/models/ernie_mtp.py index 88aadbdb73c1..0baf80e8ccde 100644 --- a/vllm/model_executor/models/ernie_mtp.py +++ b/vllm/model_executor/models/ernie_mtp.py @@ -214,6 +214,5 @@ def _filter( if any(k in name for k in ("mtp", "embed_tokens", "lm_head")): yield name, weight - skip_prefixes = ["lm_head"] if self.config.tie_word_embeddings else [] - loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes) + loader = AutoWeightsLoader(self) return loader.load_weights(_filter(weights), mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/exaone.py b/vllm/model_executor/models/exaone.py index 0742ae870b4b..fef05d386a9b 100644 --- a/vllm/model_executor/models/exaone.py +++ b/vllm/model_executor/models/exaone.py @@ -446,11 +446,5 @@ def compute_logits( return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - # With tie_word_embeddings, we can skip lm_head.weight - # The weight might appear unnecessarily in the files if the model is - # processed with quantization, LoRA, fine-tuning, etc. - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/exaone4.py b/vllm/model_executor/models/exaone4.py index 571aa3a42645..31ea665e0aa1 100644 --- a/vllm/model_executor/models/exaone4.py +++ b/vllm/model_executor/models/exaone4.py @@ -446,11 +446,5 @@ def compute_logits( return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - # With tie_word_embeddings, we can skip lm_head.weight - # The weight might appear unnecessarily in the files if the model is - # processed with quantization, LoRA, fine-tuning, etc. - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/exaone_moe.py b/vllm/model_executor/models/exaone_moe.py index ef01330f0c56..229a741dcaaf 100644 --- a/vllm/model_executor/models/exaone_moe.py +++ b/vllm/model_executor/models/exaone_moe.py @@ -560,12 +560,7 @@ def compute_logits( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader( self, - # With tie_word_embeddings, we can skip lm_head.weight - # The weight might appear unnecessarily in the files if the model is - # processed with quantization, LoRA, fine-tuning, etc. - skip_prefixes=( - ["lm_head.", "mtp."] if self.config.tie_word_embeddings else ["mtp."] - ), + skip_prefixes=["mtp."], # Skip loading extra parameters for GPTQ/modelopt models. ignore_unexpected_suffixes=[ ".bias", diff --git a/vllm/model_executor/models/fairseq2_llama.py b/vllm/model_executor/models/fairseq2_llama.py index e898034fbfa5..97afb760cf9e 100644 --- a/vllm/model_executor/models/fairseq2_llama.py +++ b/vllm/model_executor/models/fairseq2_llama.py @@ -74,10 +74,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: params = dict(self.named_parameters()) - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) return loader.load_weights( self.reshape_fairseq2_weights(name, loaded_weight, params) for name, loaded_weight in weights diff --git a/vllm/model_executor/models/falcon.py b/vllm/model_executor/models/falcon.py index cdc388e529b4..97e133a03c08 100644 --- a/vllm/model_executor/models/falcon.py +++ b/vllm/model_executor/models/falcon.py @@ -535,8 +535,5 @@ def compute_logits( return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights) diff --git a/vllm/model_executor/models/falcon_h1.py b/vllm/model_executor/models/falcon_h1.py index 3c96d00c2897..0dca542ff358 100644 --- a/vllm/model_executor/models/falcon_h1.py +++ b/vllm/model_executor/models/falcon_h1.py @@ -643,8 +643,5 @@ def compute_logits( return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/gemma.py b/vllm/model_executor/models/gemma.py index 949799fa654d..0f1282fd969c 100644 --- a/vllm/model_executor/models/gemma.py +++ b/vllm/model_executor/models/gemma.py @@ -346,8 +346,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): quant_config = vllm_config.quant_config self.config = config - # currently all existing Gemma models have `tie_word_embeddings` enabled - assert config.tie_word_embeddings self.quant_config = quant_config self.model = GemmaModel( @@ -381,8 +379,5 @@ def compute_logits( return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/gemma2.py b/vllm/model_executor/models/gemma2.py index da5161ffa019..0bdf9b1392ba 100644 --- a/vllm/model_executor/models/gemma2.py +++ b/vllm/model_executor/models/gemma2.py @@ -335,8 +335,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() self.config = config - # currently all existing Gemma models have `tie_word_embeddings` enabled - assert config.tie_word_embeddings self.quant_config = quant_config self.model = Gemma2Model( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") @@ -371,8 +369,5 @@ def compute_logits( return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/gemma3.py b/vllm/model_executor/models/gemma3.py index 4f370fec1477..e886c2ccdd07 100644 --- a/vllm/model_executor/models/gemma3.py +++ b/vllm/model_executor/models/gemma3.py @@ -442,8 +442,5 @@ def compute_logits( return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/gemma4.py b/vllm/model_executor/models/gemma4.py index f0ac972b27f0..e8bb72ac3703 100644 --- a/vllm/model_executor/models/gemma4.py +++ b/vllm/model_executor/models/gemma4.py @@ -1694,15 +1694,12 @@ def _weight_iterator(): yield name, weight # Skip multimodal weights — handled by the multimodal wrapper. - # Also skip lm_head when weights are tied. skip = [ "audio_tower.", "vision_tower.", "embed_audio.", "embed_vision.", ] - if self.config.tie_word_embeddings: - skip.append("lm_head.") loader = AutoWeightsLoader(self, skip_substrs=skip) return loader.load_weights(_weight_iterator()) diff --git a/vllm/model_executor/models/glm4.py b/vllm/model_executor/models/glm4.py index 4c30e30008fe..5d95175ab03e 100644 --- a/vllm/model_executor/models/glm4.py +++ b/vllm/model_executor/models/glm4.py @@ -292,11 +292,10 @@ def compute_logits( return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - skip_prefixes = ["lm_head."] if self.config.tie_word_embeddings else [] # Skip the speculative (MTP) layers, which are loaded by the # draft model instead. num_nextn_layers = getattr(self.config, "num_nextn_predict_layers", 0) - skip_prefixes += [ + skip_prefixes = [ f"model.layers.{self.config.num_hidden_layers + i}." for i in range(num_nextn_layers) ] diff --git a/vllm/model_executor/models/gpt_oss.py b/vllm/model_executor/models/gpt_oss.py index 27982eb1f3b1..d5fbee384a48 100644 --- a/vllm/model_executor/models/gpt_oss.py +++ b/vllm/model_executor/models/gpt_oss.py @@ -1242,8 +1242,5 @@ def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/granite.py b/vllm/model_executor/models/granite.py index df813fd1e918..7b8f42fdd136 100644 --- a/vllm/model_executor/models/granite.py +++ b/vllm/model_executor/models/granite.py @@ -467,10 +467,5 @@ def make_empty_intermediate_tensors( ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - # With tie_word_embeddings, we can skip lm_head.weight - # The weight might appear unnecessarily in the files if the model is - # processed with quantization, LoRA, fine-tuning, etc. - skip_prefixes = ["lm_head."] if self.config.tie_word_embeddings else None - - loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes) + loader = AutoWeightsLoader(self) return loader.load_weights(weights) diff --git a/vllm/model_executor/models/granitemoe.py b/vllm/model_executor/models/granitemoe.py index d775de50ca4c..3a9872a60573 100644 --- a/vllm/model_executor/models/granitemoe.py +++ b/vllm/model_executor/models/granitemoe.py @@ -452,6 +452,5 @@ def make_empty_intermediate_tensors( ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - skip_prefixes = ["lm_head."] if self.config.tie_word_embeddings else None - loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes) + loader = AutoWeightsLoader(self) return loader.load_weights(weights) diff --git a/vllm/model_executor/models/granitemoeshared.py b/vllm/model_executor/models/granitemoeshared.py index ecd942da4042..d77f58a5b3c9 100644 --- a/vllm/model_executor/models/granitemoeshared.py +++ b/vllm/model_executor/models/granitemoeshared.py @@ -289,6 +289,5 @@ def make_empty_intermediate_tensors( ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - skip_prefixes = ["lm_head."] if self.config.tie_word_embeddings else None - loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes) + loader = AutoWeightsLoader(self) return loader.load_weights(weights) diff --git a/vllm/model_executor/models/hrm_text.py b/vllm/model_executor/models/hrm_text.py index dd7ba93e7b8a..29825ece5a0d 100644 --- a/vllm/model_executor/models/hrm_text.py +++ b/vllm/model_executor/models/hrm_text.py @@ -521,6 +521,5 @@ def compute_logits( return self.logits_processor(self.lm_head, hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - skip_prefixes = ["lm_head."] if self.config.tie_word_embeddings else None - loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/hunyuan_v1.py b/vllm/model_executor/models/hunyuan_v1.py index b90d463627d9..cf1933bd7dd5 100644 --- a/vllm/model_executor/models/hunyuan_v1.py +++ b/vllm/model_executor/models/hunyuan_v1.py @@ -970,10 +970,7 @@ def make_empty_intermediate_tensors( ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: diff --git a/vllm/model_executor/models/hunyuan_vision.py b/vllm/model_executor/models/hunyuan_vision.py index 74e4c1a6cec4..252556e86c43 100644 --- a/vllm/model_executor/models/hunyuan_vision.py +++ b/vllm/model_executor/models/hunyuan_vision.py @@ -1034,10 +1034,7 @@ def compute_logits( return self.language_model.compute_logits(hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) def get_mm_mapping(self) -> MultiModelKeys: diff --git a/vllm/model_executor/models/hy_v3.py b/vllm/model_executor/models/hy_v3.py index 0248ac3d9c02..8b77d4b1e83e 100644 --- a/vllm/model_executor/models/hy_v3.py +++ b/vllm/model_executor/models/hy_v3.py @@ -710,10 +710,7 @@ def _filter_weights(weights): continue yield name, weight - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(_filter_weights(weights)) def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: diff --git a/vllm/model_executor/models/hyperclovax.py b/vllm/model_executor/models/hyperclovax.py index b6d66698c92d..66c3b4a270a3 100644 --- a/vllm/model_executor/models/hyperclovax.py +++ b/vllm/model_executor/models/hyperclovax.py @@ -478,8 +478,5 @@ def load_weights( self, weights: Iterable[tuple[str, torch.Tensor]], ) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=["lm_head."] if self.config.tie_word_embeddings else None, - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/internlm2.py b/vllm/model_executor/models/internlm2.py index ffc750c7c565..4996723b3c6b 100644 --- a/vllm/model_executor/models/internlm2.py +++ b/vllm/model_executor/models/internlm2.py @@ -381,10 +381,7 @@ def compute_logits( return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["output."] if self.config.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights) diff --git a/vllm/model_executor/models/iquest_loopcoder.py b/vllm/model_executor/models/iquest_loopcoder.py index fdb90b962d08..d3b670658a5f 100644 --- a/vllm/model_executor/models/iquest_loopcoder.py +++ b/vllm/model_executor/models/iquest_loopcoder.py @@ -575,8 +575,5 @@ def compute_logits( return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights) diff --git a/vllm/model_executor/models/jais2.py b/vllm/model_executor/models/jais2.py index 95b8c3ee44fe..84254e1ee5f4 100644 --- a/vllm/model_executor/models/jais2.py +++ b/vllm/model_executor/models/jais2.py @@ -437,8 +437,5 @@ def compute_logits( return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/laguna.py b/vllm/model_executor/models/laguna.py index 9431f9a0ec43..115000588db0 100644 --- a/vllm/model_executor/models/laguna.py +++ b/vllm/model_executor/models/laguna.py @@ -763,8 +763,5 @@ def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights) diff --git a/vllm/model_executor/models/lfm2.py b/vllm/model_executor/models/lfm2.py index 69a0e23be497..6f03150b8461 100644 --- a/vllm/model_executor/models/lfm2.py +++ b/vllm/model_executor/models/lfm2.py @@ -508,8 +508,5 @@ def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights) diff --git a/vllm/model_executor/models/lfm2_moe.py b/vllm/model_executor/models/lfm2_moe.py index 966a458dea55..c2a21e52d18c 100644 --- a/vllm/model_executor/models/lfm2_moe.py +++ b/vllm/model_executor/models/lfm2_moe.py @@ -671,8 +671,5 @@ def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights) diff --git a/vllm/model_executor/models/llama.py b/vllm/model_executor/models/llama.py index cbf75593e0b2..e3c5272224d4 100644 --- a/vllm/model_executor/models/llama.py +++ b/vllm/model_executor/models/llama.py @@ -534,10 +534,7 @@ def compute_logits( return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights) diff --git a/vllm/model_executor/models/llama4.py b/vllm/model_executor/models/llama4.py index 9577c27f5494..2a073d2dc48c 100644 --- a/vllm/model_executor/models/llama4.py +++ b/vllm/model_executor/models/llama4.py @@ -792,10 +792,7 @@ def _init_model( ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) # Use a generator (not a list comprehension) so the weights iterator is # consumed lazily by AutoWeightsLoader. Materializing it here would hold # the entire language-model checkpoint in host memory at once, which can diff --git a/vllm/model_executor/models/mimo.py b/vllm/model_executor/models/mimo.py index 56250ba2f4e8..7be514e990a5 100644 --- a/vllm/model_executor/models/mimo.py +++ b/vllm/model_executor/models/mimo.py @@ -119,10 +119,8 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - skip_prefixes = ["lm_head."] if self.config.tie_word_embeddings else [] # MTP layers are loaded by the draft model, not the main model. - skip_prefixes.append("model.mtp_layers.") - loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes) + loader = AutoWeightsLoader(self, skip_prefixes=["model.mtp_layers."]) return loader.load_weights(weights) def compute_logits( diff --git a/vllm/model_executor/models/minicpm.py b/vllm/model_executor/models/minicpm.py index ee9b31f140bd..699e003592b0 100644 --- a/vllm/model_executor/models/minicpm.py +++ b/vllm/model_executor/models/minicpm.py @@ -648,8 +648,5 @@ def compute_logits( return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights) diff --git a/vllm/model_executor/models/minicpm_eagle.py b/vllm/model_executor/models/minicpm_eagle.py index 890d52961c4e..89ce01d6117c 100644 --- a/vllm/model_executor/models/minicpm_eagle.py +++ b/vllm/model_executor/models/minicpm_eagle.py @@ -387,8 +387,5 @@ def transform(inputs): process_eagle_weight(self, name) return name, loaded_weight - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(map(transform, weights)) diff --git a/vllm/model_executor/models/mistral.py b/vllm/model_executor/models/mistral.py index ce1332d0c9d1..86c4b171e5f1 100644 --- a/vllm/model_executor/models/mistral.py +++ b/vllm/model_executor/models/mistral.py @@ -277,10 +277,7 @@ def _init_model( ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) return loader.load_weights( self.maybe_remap_mistral(name, loaded_weight) for name, loaded_weight in weights diff --git a/vllm/model_executor/models/mpt.py b/vllm/model_executor/models/mpt.py index 8e509fbcb4c6..b8b4c5670ed6 100644 --- a/vllm/model_executor/models/mpt.py +++ b/vllm/model_executor/models/mpt.py @@ -279,7 +279,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): config = vllm_config.model_config.hf_config quant_config = vllm_config.quant_config self.config = config - assert config.tie_word_embeddings self.quant_config = quant_config self.transformer = MPTModel( diff --git a/vllm/model_executor/models/nemotron_nas.py b/vllm/model_executor/models/nemotron_nas.py index 5a5f0e77739f..df3eb60dfb00 100644 --- a/vllm/model_executor/models/nemotron_nas.py +++ b/vllm/model_executor/models/nemotron_nas.py @@ -412,8 +412,5 @@ def compute_logits( return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/olmo3.py b/vllm/model_executor/models/olmo3.py index 9ac4e61cf47d..a54dd5749e25 100644 --- a/vllm/model_executor/models/olmo3.py +++ b/vllm/model_executor/models/olmo3.py @@ -408,10 +408,5 @@ def compute_logits( return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): - loader = AutoWeightsLoader( - self, - skip_prefixes=( - ["lm_head.weight"] if self.config.tie_word_embeddings else None - ), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/olmo_hybrid.py b/vllm/model_executor/models/olmo_hybrid.py index 122a1c1c5b25..9a24df5ec9a6 100644 --- a/vllm/model_executor/models/olmo_hybrid.py +++ b/vllm/model_executor/models/olmo_hybrid.py @@ -474,10 +474,5 @@ def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, MambaStateCopyFu return MambaStateCopyFuncCalculator.gated_delta_net_state_copy_func() def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): - loader = AutoWeightsLoader( - self, - skip_prefixes=( - ["lm_head.weight"] if self.config.tie_word_embeddings else None - ), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights) diff --git a/vllm/model_executor/models/openpangu.py b/vllm/model_executor/models/openpangu.py index 4a2030cfddc8..60b5b301d56e 100644 --- a/vllm/model_executor/models/openpangu.py +++ b/vllm/model_executor/models/openpangu.py @@ -1153,10 +1153,7 @@ def compute_logits( return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights) diff --git a/vllm/model_executor/models/opt.py b/vllm/model_executor/models/opt.py index d78e51e3dde5..1efe9eb8a8ca 100644 --- a/vllm/model_executor/models/opt.py +++ b/vllm/model_executor/models/opt.py @@ -384,10 +384,5 @@ def compute_logits( return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=( - ["lm_head.weight"] if self.config.tie_word_embeddings else None - ), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/plamo3.py b/vllm/model_executor/models/plamo3.py index 2ba38a7b1f8f..3f9553f64adf 100644 --- a/vllm/model_executor/models/plamo3.py +++ b/vllm/model_executor/models/plamo3.py @@ -430,8 +430,5 @@ def compute_logits( return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights) diff --git a/vllm/model_executor/models/qwen2.py b/vllm/model_executor/models/qwen2.py index 3ec8ac931913..3820f5e39b42 100644 --- a/vllm/model_executor/models/qwen2.py +++ b/vllm/model_executor/models/qwen2.py @@ -500,8 +500,5 @@ def compute_logits( return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights) diff --git a/vllm/model_executor/models/qwen3.py b/vllm/model_executor/models/qwen3.py index 7db73a749f75..79434bfac964 100644 --- a/vllm/model_executor/models/qwen3.py +++ b/vllm/model_executor/models/qwen3.py @@ -335,8 +335,5 @@ def compute_logits( return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights) diff --git a/vllm/model_executor/models/rnj1.py b/vllm/model_executor/models/rnj1.py index 2bcd27919813..b184e3fcade4 100644 --- a/vllm/model_executor/models/rnj1.py +++ b/vllm/model_executor/models/rnj1.py @@ -395,8 +395,5 @@ def compute_logits( return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/sarvam.py b/vllm/model_executor/models/sarvam.py index 04590a2a913a..9b5a784cefd1 100644 --- a/vllm/model_executor/models/sarvam.py +++ b/vllm/model_executor/models/sarvam.py @@ -679,10 +679,7 @@ def load_weights( self, weights: Iterable[tuple[str, torch.Tensor]], ) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights) diff --git a/vllm/model_executor/models/seed_oss.py b/vllm/model_executor/models/seed_oss.py index f3fa8b2d5de2..df0ed7e72e5a 100644 --- a/vllm/model_executor/models/seed_oss.py +++ b/vllm/model_executor/models/seed_oss.py @@ -428,8 +428,5 @@ def compute_logits( return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/telechat2.py b/vllm/model_executor/models/telechat2.py index 42fa6d6871a6..637c456dbe3b 100644 --- a/vllm/model_executor/models/telechat2.py +++ b/vllm/model_executor/models/telechat2.py @@ -122,8 +122,5 @@ def _init_model( return TeleChat2Model(vllm_config=vllm_config, prefix=prefix) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/transformers/causal.py b/vllm/model_executor/models/transformers/causal.py index a32e2b54ba42..29e780557cb7 100644 --- a/vllm/model_executor/models/transformers/causal.py +++ b/vllm/model_executor/models/transformers/causal.py @@ -40,11 +40,7 @@ def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""): vllm_config=vllm_config, prefix=prefix ) - # Tell `Base.load_weights` to skip - # `lm_head` if the model has tied word embeddings tie_word_embeddings = self._get_tie_word_embeddings() - if tie_word_embeddings: - self.skip_prefixes.append("lm_head.") if self.pp_group.is_last_rank: self.lm_head = ParallelLMHead( diff --git a/vllm/model_executor/models/utils.py b/vllm/model_executor/models/utils.py index 5e1236a7ebdf..2f1d6f3b5c17 100644 --- a/vllm/model_executor/models/utils.py +++ b/vllm/model_executor/models/utils.py @@ -18,6 +18,7 @@ get_tensor_model_parallel_world_size, ) from vllm.logger import init_logger +from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding from vllm.model_executor.model_loader.reload import ( support_quantized_model_reload_from_hp_weights, ) @@ -169,6 +170,20 @@ def get_unstacked_mapper(self) -> "WeightsMapper": return replace(self, orig_to_new_stacked={}) +def _get_tied_embedding_params(module: nn.Module) -> dict[str, str]: + """Map each tied word embedding qualname to the first name it aliases.""" + canonical = dict[int, str]() + aliased = dict[str, str]() + for prefix, submodule in module.named_modules(remove_duplicate=False): + if not isinstance(submodule, VocabParallelEmbedding): + continue + for name, param in submodule.named_parameters(remove_duplicate=False): + qualname = f"{prefix}.{name}" if prefix else name + if (first_name := canonical.setdefault(id(param), qualname)) != qualname: + aliased[qualname] = first_name + return aliased + + class AutoWeightsLoader: """ Helper class to load weights into a [`torch.nn.Module`][]. It is able @@ -213,6 +228,14 @@ def __init__( # update default skip_substrs self.skip_substrs += self.ROTARY_EMBEDS_UNUSED_WEIGHTS + # Weight tying makes two qualnames point at the same `nn.Parameter` + # (e.g. `lm_head.weight` and `model.embed_tokens.weight`). Loading both + # would write the same storage twice, so only the first name reached by + # module traversal is loaded and the rest are skipped. + self.aliased_params = _get_tied_embedding_params(module) + self._skipped_aliases = dict[str, str]() + self._loaded_params_are_complete = True + def _groupby_prefix( self, weights: Iterable[tuple[str, torch.Tensor]], @@ -242,8 +265,10 @@ def _get_qualname(self, prefix: str, rest: str) -> str: return ".".join((prefix, rest)) def _can_skip(self, qualname: str) -> bool: - return any(qualname.startswith(p) for p in self.skip_prefixes) or any( - substr in qualname for substr in self.skip_substrs + return ( + qualname in self.aliased_params + or any(qualname.startswith(p) for p in self.skip_prefixes) + or any(substr in qualname for substr in self.skip_substrs) ) def _can_ignore_unexpected(self, qualname: str) -> bool: @@ -332,6 +357,7 @@ def _load_module( logger.warning( "Unable to collect loaded parameters for module %s", module ) + self._loaded_params_are_complete = False else: yield from map( lambda x: self._get_qualname(base_prefix, x), @@ -415,14 +441,36 @@ def load_weights( self.ignore_unexpected_suffixes.extend(ignore_unexpected_suffixes) if mapper is not None: weights = mapper.apply(weights) - # filter out weights with first-prefix/substr to skip in name - weights = ( - (name, weight) for name, weight in weights if not self._can_skip(name) - ) + weights = self._filter_skipped(weights) autoloaded_weights = set(self._load_module("", self.module, weights)) + self._check_skipped_aliases(autoloaded_weights) return autoloaded_weights + def _filter_skipped( + self, weights: Iterable[tuple[str, torch.Tensor]] + ) -> Iterable[tuple[str, torch.Tensor]]: + for name, weight in weights: + if (canonical := self.aliased_params.get(name)) is not None: + self._skipped_aliases[name] = canonical + if self._can_skip(name): + continue + + yield name, weight + + def _check_skipped_aliases(self, autoloaded_weights: set[str]) -> None: + """Guard against skipping an alias whose canonical name never loads.""" + if not self._loaded_params_are_complete: + return + for alias, canonical in self._skipped_aliases.items(): + if canonical not in autoloaded_weights: + raise ValueError( + f"{alias!r} was skipped because it is tied to {canonical!r} " + f"in {self.module._get_name()}, but {canonical!r} was not " + "found in the checkpoint, so the tied weight is " + "uninitialized." + ) + def maybe_fuse_shared_experts( weights: Iterable[tuple[str, torch.Tensor]], diff --git a/vllm/models/kimi_k3/amd/linear.py b/vllm/models/kimi_k3/amd/linear.py index 5cb916ef54ee..835c3be4dbee 100644 --- a/vllm/models/kimi_k3/amd/linear.py +++ b/vllm/models/kimi_k3/amd/linear.py @@ -1082,8 +1082,5 @@ def compute_logits( return self.logits_processor(self.lm_head, hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights) diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index f1166695fdeb..472272a35741 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -1676,10 +1676,7 @@ def compute_logits( return self.logits_processor(self.lm_head, hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + loader = AutoWeightsLoader(self) loaded = loader.load_weights(weights) self.model.finalize_mega_moe_weights() # The fused MultiHeadLatentAttention's process_weights_after_loading diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index 87271c258825..72fe24302e64 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -1112,7 +1112,9 @@ def try_get_safetensors_metadata( try: return with_retry( - get_safetensors_metadata_partial, "Error retrieving safetensors" + get_safetensors_metadata_partial, + "Error retrieving safetensors", + fatal_errors=(huggingface_hub.errors.NotASafetensorsRepoError,), ) except Exception: return None @@ -1218,6 +1220,24 @@ def get_safetensors_params_metadata( return _read_safetensors_metadata_in_dir(Path(local_dir)) +@cache +def checkpoint_has_lm_head(model: str, *, revision: str | None = None) -> bool | None: + """Whether the checkpoint contains an `lm_head` tensor of its own. + + Args: + model: Name or path of the model repository. + revision: The specific model version to use. + + Returns: + `None` if the checkpoint contents could not be determined, for example + because it is not stored as safetensors. + """ + metadata = get_safetensors_params_metadata(model, revision=revision) + if not metadata: + return None + return any(name.endswith("lm_head.weight") for name in metadata) + + def _download_mistral_config_file(model, revision) -> dict: config_file_name = "params.json" config_dict = get_hf_file_to_dict(config_file_name, model, revision) diff --git a/vllm/transformers_utils/repo_utils.py b/vllm/transformers_utils/repo_utils.py index 4066edbbf094..ca20390aad51 100644 --- a/vllm/transformers_utils/repo_utils.py +++ b/vllm/transformers_utils/repo_utils.py @@ -95,10 +95,23 @@ def with_retry( log_msg: str, max_retries: int = 2, retry_delay: int = 2, + fatal_errors: tuple[type[Exception], ...] = (), ) -> _R: + """Call `func`, retrying transient failures. + + Args: + func: The call to make. + log_msg: Prefix for the message logged when a call fails. + max_retries: How many times to call `func` before giving up. + retry_delay: Seconds to wait after the first failure, doubled each time. + fatal_errors: Exceptions that are a definitive answer rather than a + transient failure. These are raised immediately, without logging. + """ for attempt in range(max_retries): try: return func() + except fatal_errors: + raise except Exception as e: if attempt == max_retries - 1: logger.error("%s: %s", log_msg, e) From 6df7adc17f7af8ca3f5b3e6f5ccd48960e95eacb Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Thu, 20 Aug 2026 03:33:41 -0700 Subject: [PATCH 206/839] [Bugfix][GDN] Reset speculative decode count for an empty draft schedule (#53077) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> --- vllm/v1/attention/backends/gdn_attn.py | 1 + 1 file changed, 1 insertion(+) diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 0e843e31baa7..27df94d7bd65 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -239,6 +239,7 @@ def build( # type: ignore[override] or num_decode_draft_tokens_cpu[spec_sequence_masks_cpu].sum().item() == 0 ): + num_spec_decodes = 0 spec_sequence_masks = None spec_sequence_masks_cpu = None else: From 6259572b283a4df3d0e8690aad5da003b012c103 Mon Sep 17 00:00:00 2001 From: Thien Tran Date: Thu, 20 Aug 2026 19:08:50 +0800 Subject: [PATCH 207/839] [Docs] Use incremental builds for C++ changes in `AGENTS.md` (#53098) Signed-off-by: Thien Tran Co-authored-by: OpenAI Codex --- AGENTS.md | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index a53b81873cf0..a1060975b0b4 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -60,13 +60,14 @@ pre-commit install ### Installing dependencies ```bash -# If you are only making Python changes: +# Start with precompiled artifacts for an editable install: VLLM_USE_PRECOMPILED=1 uv pip install -e . --torch-backend=auto - -# If you are also making C/C++ changes: -uv pip install -e . --torch-backend=auto ``` +For C/C++ or CUDA changes, follow the +[incremental compilation workflow](docs/contributing/incremental_build.md) to +configure and perform incremental builds. + ### Tests > Requires [Environment setup](#environment-setup) and [Installing dependencies](#installing-dependencies). From df1376907b16b8b57a6c08fe074f62015d81cde3 Mon Sep 17 00:00:00 2001 From: Zupeng Wang Date: Thu, 20 Aug 2026 20:33:43 +0800 Subject: [PATCH 208/839] [Model] Add tower and connector LoRA support for LFM2-VL (#51498) Signed-off-by: zupengwang <71580390+zupengwang@users.noreply.github.com> --- vllm/model_executor/models/lfm2_siglip2.py | 8 +++++--- vllm/model_executor/models/lfm2_vl.py | 19 +++++++++++++++++-- 2 files changed, 22 insertions(+), 5 deletions(-) diff --git a/vllm/model_executor/models/lfm2_siglip2.py b/vllm/model_executor/models/lfm2_siglip2.py index f1679af813c4..31e7c0f14438 100644 --- a/vllm/model_executor/models/lfm2_siglip2.py +++ b/vllm/model_executor/models/lfm2_siglip2.py @@ -20,6 +20,7 @@ from vllm.model_executor.layers.linear import ( ColumnParallelLinear, QKVParallelLinear, + ReplicatedLinear, RowParallelLinear, ) from vllm.model_executor.layers.quantization import QuantizationConfig @@ -37,9 +38,10 @@ def __init__(self, config: Siglip2VisionConfig): self.config = config self.embed_dim = config.hidden_size self.patch_size = config.patch_size - self.patch_embedding = nn.Linear( - in_features=config.num_channels * self.patch_size * self.patch_size, - out_features=self.embed_dim, + self.patch_embedding = ReplicatedLinear( + input_size=config.num_channels * self.patch_size * self.patch_size, + output_size=self.embed_dim, + return_bias=False, ) self.num_patches = config.num_patches self.position_embedding_size = int(self.num_patches**0.5) diff --git a/vllm/model_executor/models/lfm2_vl.py b/vllm/model_executor/models/lfm2_vl.py index ce60f2d236d9..d247fef4d7e3 100644 --- a/vllm/model_executor/models/lfm2_vl.py +++ b/vllm/model_executor/models/lfm2_vl.py @@ -22,6 +22,7 @@ from vllm.config.multimodal import BaseDummyOptions from vllm.forward_context import set_forward_context from vllm.inputs import MultiModalDataDict +from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.layers.mamba.mamba_utils import ( MambaStateCopyFunc, MambaStateCopyFuncCalculator, @@ -499,16 +500,20 @@ def __init__( self.projector_use_layernorm = config.projector_use_layernorm if self.projector_use_layernorm: self.layer_norm = nn.LayerNorm(in_channels) - self.linear_1 = nn.Linear( + self.linear_1 = ReplicatedLinear( in_channels, config.projector_hidden_size, bias=config.projector_bias, + prefix=maybe_prefix(prefix, "linear_1"), + return_bias=False, ) self.act = ACT2FN[config.projector_hidden_act] - self.linear_2 = nn.Linear( + self.linear_2 = ReplicatedLinear( config.projector_hidden_size, config.text_config.hidden_size, bias=config.projector_bias, + prefix=maybe_prefix(prefix, "linear_2"), + return_bias=False, ) def forward( @@ -1259,3 +1264,13 @@ def get_mm_mapping(self) -> MultiModelKeys: connector="multi_modal_projector", tower_model="vision_tower", ) + + def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int: + downsample_factor = self.config.downsample_factor + + return num_image_tokens * downsample_factor**2 + + def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int: + downsample_factor = self.config.downsample_factor + + return num_vision_tokens // downsample_factor**2 From 4b7cb949a9067e3906e0b89df115b22100036763 Mon Sep 17 00:00:00 2001 From: Sandeep Maddipatla Date: Thu, 20 Aug 2026 06:03:32 -0700 Subject: [PATCH 209/839] [Docker] Update to nixl-1.3.2 (#51777) Signed-off-by: Sandeep Maddipatla Co-authored-by: Kunshang Ji --- docker/Dockerfile.xpu | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/docker/Dockerfile.xpu b/docker/Dockerfile.xpu index 456be41f6f2a..0811d64dff26 100644 --- a/docker/Dockerfile.xpu +++ b/docker/Dockerfile.xpu @@ -162,8 +162,11 @@ RUN --mount=type=cache,target=/root/.cache/uv \ FROM vllm-base AS ucx-nixl-build -ARG UCX_VERSION=v1.21.0-rc2 -ARG NIXL_VERSION=v1.2.0 +# Keep UCX_VERSION aligned with the UCX that NIXL itself builds against for this +# NIXL_VERSION (nixl contrib/Dockerfile.manylinux: ARG UCX_REF), so XPU runs the +# same UCX commit as the released nixl-cu* wheels rather than an untested pairing. +ARG UCX_VERSION=v1.21.x +ARG NIXL_VERSION=v1.3.2 # Build-time only: compiler, autotools, and verbs dev headers RUN apt-get update -y && apt-get install -y --no-install-recommends \ @@ -198,7 +201,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \ FROM vllm-base AS vllm-openai -ARG NIXL_VERSION=v1.2.0 +# Must match the ucx-nixl-build stage above: ARG does not cross stages, so a stale +# value here installs a meta package that disagrees with the wheel built there. +ARG NIXL_VERSION=v1.3.2 # Copy compiled UCX runtime libraries and the pre-built NIXL wheel. # No compiler or autotools are installed in this stage. From de216b6e6487b8b452805b582b242616ee616682 Mon Sep 17 00:00:00 2001 From: Prudhvi Vuda <53619858+Prudhvivuda@users.noreply.github.com> Date: Thu, 20 Aug 2026 09:42:50 -0400 Subject: [PATCH 210/839] [Bugfix] Skip MM processor cache inserts larger than capacity (#53016) Signed-off-by: Prudhvivuda --- docs/configuration/optimization.md | 3 ++ tests/multimodal/test_cache.py | 29 +++++++++++++++++++ tests/utils_/test_cache.py | 15 ++++++++++ vllm/config/multimodal.py | 3 ++ vllm/multimodal/cache.py | 46 ++++++++++++++++++++++++++---- vllm/utils/cache.py | 17 +++++++++++ 6 files changed, 108 insertions(+), 5 deletions(-) diff --git a/docs/configuration/optimization.md b/docs/configuration/optimization.md index efa0f8b9046d..18122ea2c8bc 100644 --- a/docs/configuration/optimization.md +++ b/docs/configuration/optimization.md @@ -376,6 +376,9 @@ processes. ### Configuration You can adjust the size of the cache by setting the value of `mm_processor_cache_gb` (default 4 GiB). +A processed item larger than this budget is served uncached (with a warning) +instead of failing engine startup; raise `mm_processor_cache_gb` if you want +those items cached. If you do not benefit much from the cache, you can disable both IPC and processor caching completely via `mm_processor_cache_gb=0`. diff --git a/tests/multimodal/test_cache.py b/tests/multimodal/test_cache.py index e5845b12b14d..c20e554e5794 100644 --- a/tests/multimodal/test_cache.py +++ b/tests/multimodal/test_cache.py @@ -17,6 +17,7 @@ MultiModalProcessorCacheInItem, MultiModalProcessorCacheItem, MultiModalProcessorCacheItemMetadata, + MultiModalProcessorOnlyCache, MultiModalProcessorSenderCache, MultiModalReceiverCache, ShmObjectStoreReceiverCache, @@ -247,6 +248,34 @@ def get_multimodal_config(self) -> MultiModalConfig: return self._mm_config +@pytest.mark.skip_global_cleanup +def test_oversized_item_is_served_uncached(): + """Items larger than the processor cache must not crash insert. + + cachetools.LRUCache raises ValueError("value too large") when a single + item exceeds maxsize. Engine profiling uses a max-size dummy item, so this + used to abort EngineCore startup (vllm-project/vllm#52835). Skip the + insert and serve the item uncached instead. + """ + # 1 KiB cache; a 4 KiB item cannot fit even if the cache is empty. + model_config = _StubModelConfig(mm_processor_cache_gb=1024 / GiB_bytes) + item = MultiModalKwargsItem.dummy(nbytes=4096) + small = MultiModalKwargsItem.dummy(nbytes=64) + + p0_only = MultiModalProcessorOnlyCache(model_config) # type: ignore[arg-type] + assert p0_only.get_and_update_item((item, []), "big")[0] is item + assert not p0_only.is_cached_item("big") + assert p0_only.get_and_update_item((small, []), "small")[0] is small + assert p0_only.is_cached_item("small") + + p0 = MultiModalProcessorSenderCache(model_config) # type: ignore[arg-type] + p1 = MultiModalReceiverCache(model_config) # type: ignore[arg-type] + assert p0.get_and_update_item((item, []), "big")[0] is item + assert not p0.is_cached_item("big") + assert p1.get_and_update_item(item, "big") is item + assert "big" not in p1._cache + + def test_mm_cache_miss_raises_and_recovers(): """A P0/P1 multimodal cache drift must be recoverable, not a hard crash. diff --git a/tests/utils_/test_cache.py b/tests/utils_/test_cache.py index e361006fd8e6..8d09b224151b 100644 --- a/tests/utils_/test_cache.py +++ b/tests/utils_/test_cache.py @@ -1,5 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest + from vllm.utils.cache import CacheInfo, LRUCache @@ -123,3 +125,16 @@ def test_lru_cache(): assert 2 in cache assert 4 in cache assert 6 in cache + + +def test_lru_cache_put_if_fits(): + cache = LRUCache(10, getsizeof=lambda x: x) + + assert cache.put_if_fits("ok", 4) is True + assert cache["ok"] == 4 + + assert cache.put_if_fits("big", 11) is False + assert "big" not in cache + + with pytest.raises(ValueError, match="value too large"): + cache.put("big", 11) diff --git a/vllm/config/multimodal.py b/vllm/config/multimodal.py index 04e94735a1cd..ce2217d9fefd 100644 --- a/vllm/config/multimodal.py +++ b/vllm/config/multimodal.py @@ -157,6 +157,9 @@ class MultiModalConfig: resulting in a total memory usage of `mm_processor_cache_gb * (api_server_count + data_parallel_size)`. + A single processed item larger than this budget is served uncached + (with a warning) instead of failing. Raise this value to cache such items. + Set to `0` to disable this cache completely (not recommended).""" mm_processor_cache_type: MMCacheType = "lru" """Type of cache to use for the multi-modal preprocessor/mapper. If `shm`, diff --git a/vllm/multimodal/cache.py b/vllm/multimodal/cache.py index 54d86b57e5a8..5ec160d2f875 100644 --- a/vllm/multimodal/cache.py +++ b/vllm/multimodal/cache.py @@ -266,6 +266,42 @@ def get_and_update( for mm_item, mm_hash in zip(mm_items, mm_hashes) ] + def cache_if_fits( + self, + cache: LRUCache[str, _V], + key: str, + value: _V, + ) -> bool: + """Insert `value` if it fits in `cache`. + + `cachetools.Cache` raises `ValueError("value too large")` when a + single item exceeds `maxsize`. An item bigger than the whole + processor cache can never be a hit, so skip the insert and serve it + uncached instead of aborting engine startup. + + LRU P0/P1 caches call this so they stay mirrored. SHM subclasses do + not use it; they already skip oversize items in `put()`. + + Args: + cache: The LRU cache to update. + key: Cache key (typically the multi-modal item hash). + value: Value to insert. + + Returns: + `True` if the item was cached, otherwise `False`. + """ + if cache.put_if_fits(key, value): + return True + logger.warning_once( + "Skipping multi-modal processor cache insert for an item of " + "%s GiB because it exceeds --mm-processor-cache-gb=%s. " + "The item will be processed uncached; increase " + "--mm-processor-cache-gb to cache items of this size.", + format_gib(int(cache.getsizeof(value))), + format_gib(int(cache.maxsize)), + ) + return False + @abstractmethod def clear_cache(self) -> None: """Clear the underlying cache.""" @@ -389,8 +425,7 @@ def get_and_update_item( assert mm_item is not None, f"Expected a cached item for {mm_hash=}" - self._cache[mm_hash] = MultiModalProcessorCacheItem(*mm_item) - + self.cache_if_fits(self._cache, mm_hash, MultiModalProcessorCacheItem(*mm_item)) return mm_item @override @@ -447,8 +482,9 @@ def get_and_update_item( assert mm_item is not None, f"Expected a cached item for {mm_hash=}" - self._cache[mm_hash] = MultiModalProcessorCacheItemMetadata(*mm_item) - + self.cache_if_fits( + self._cache, mm_hash, MultiModalProcessorCacheItemMetadata(*mm_item) + ) return mm_item @override @@ -709,7 +745,7 @@ def get_and_update_item( if mm_item is None: raise MultiModalCacheMissError([mm_hash]) - self._cache[mm_hash] = mm_item + self.cache_if_fits(self._cache, mm_hash, mm_item) return mm_item @override diff --git a/vllm/utils/cache.py b/vllm/utils/cache.py index e45c2f6a8bab..791b8d848d4f 100644 --- a/vllm/utils/cache.py +++ b/vllm/utils/cache.py @@ -157,6 +157,23 @@ def pop(self, key: _K, default: _V | _T | None = None) -> _V | _T | None: def put(self, key: _K, value: _V) -> None: self.__setitem__(key, value) + def put_if_fits(self, key: _K, value: _V) -> bool: + """Insert `value` if it is not larger than the cache capacity. + + Unlike `put`, this does not raise when a single item exceeds + `maxsize`. Size is computed once inside the insert path. + + Returns: + `True` if the item was cached, otherwise `False`. + """ + try: + self[key] = value + except ValueError as e: + if str(e) != "value too large": + raise + return False + return True + def pin(self, key: _K) -> None: """ Pins a key in the cache preventing it from being From bd8865a299c4a68cff9b6443b9fd795f4c4735f6 Mon Sep 17 00:00:00 2001 From: Seonjin Date: Thu, 20 Aug 2026 07:07:22 -0700 Subject: [PATCH 211/839] [Kernel] Add FlashInfer TRTLLM MXFP8 linear backend (#52204) Signed-off-by: seonjinn Signed-off-by: Misha Goin Co-authored-by: Misha Goin --- docs/features/quantization/modelopt.md | 5 + .../test_flashinfer_mxfp8_trtllm.py | 123 ++++++++++++++++++ .../model_executor/kernels/linear/__init__.py | 4 + .../kernels/linear/mxfp8/flashinfer.py | 100 ++++++++++++++ vllm/utils/flashinfer.py | 33 +++++ 5 files changed, 265 insertions(+) create mode 100644 tests/kernels/quantization/test_flashinfer_mxfp8_trtllm.py diff --git a/docs/features/quantization/modelopt.md b/docs/features/quantization/modelopt.md index 7850dc75ccca..139497ababf7 100644 --- a/docs/features/quantization/modelopt.md +++ b/docs/features/quantization/modelopt.md @@ -33,6 +33,11 @@ following `quantization.quant_algo` values: [Engine Arguments](../../configuration/engine_args.md) page and shown by `vllm serve --help=KernelConfig`. +!!! note + For models quantized to MXFP8 with BF16 activations on SM100-family GPUs, + use `--linear-backend flashinfer_trtllm` to select FlashInfer's TensorRT-LLM + GEMM backend. + ## Quantizing HuggingFace Models with PTQ You can quantize HuggingFace models using the example scripts provided in the Model Optimizer repository. The primary script for LLM PTQ is typically found within the `examples/llm_ptq` directory. diff --git a/tests/kernels/quantization/test_flashinfer_mxfp8_trtllm.py b/tests/kernels/quantization/test_flashinfer_mxfp8_trtllm.py new file mode 100644 index 000000000000..315e35b782a0 --- /dev/null +++ b/tests/kernels/quantization/test_flashinfer_mxfp8_trtllm.py @@ -0,0 +1,123 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch +import torch.nn.functional as F + +from vllm.model_executor.kernels.linear import ( + FlashInferTrtllmMxfp8LinearKernel, + Mxfp8LinearLayerConfig, +) +from vllm.platforms import current_platform +from vllm.utils import flashinfer as vllm_flashinfer +from vllm.utils.flashinfer import has_flashinfer + +if not ( + current_platform.is_cuda() + and current_platform.is_device_capability_family(100) + and has_flashinfer() +): + pytest.skip( + reason="FlashInfer TRTLLM MXFP8 requires an SM100-family GPU", + allow_module_level=True, + ) + + +def _make_layer(weight: torch.Tensor) -> torch.nn.Module: + from flashinfer import SfLayout, mxfp8_quantize + + weight_mxfp8, weight_scale = mxfp8_quantize( + weight, + sf_swizzle_layout=SfLayout.layout_linear, + ) + layer = torch.nn.Module() + layer.weight = torch.nn.Parameter(weight_mxfp8, requires_grad=False) + layer.weight_scale = torch.nn.Parameter( + weight_scale.view(weight.shape[0], weight.shape[1] // 32), + requires_grad=False, + ) + return layer + + +@pytest.mark.parametrize("shape", [(1, 130, 256), (7, 256, 512), (128, 130, 768)]) +@torch.inference_mode() +def test_flashinfer_trtllm_mxfp8_linear_numerics( + shape: tuple[int, int, int], +) -> None: + torch.manual_seed(0) + m, n, k = shape + x = torch.randn((m, k), dtype=torch.bfloat16, device="cuda") + weight = torch.randn((n, k), dtype=torch.bfloat16, device="cuda") + layer = _make_layer(weight) + kernel = FlashInferTrtllmMxfp8LinearKernel(Mxfp8LinearLayerConfig()) + kernel.process_weights_after_loading(layer) + + output = kernel.apply_weights(layer, x) + reference = torch.mm(x, weight.t()) + similarity = F.cosine_similarity( + output.float().flatten(), reference.float().flatten(), dim=0 + ) + + assert output.shape == (m, n) + assert output.is_contiguous() + assert similarity.item() > 0.98 + + +@torch.inference_mode() +def test_flashinfer_trtllm_mxfp8_custom_ops() -> None: + x = torch.randn((7, 512), dtype=torch.bfloat16, device="cuda") + weight = torch.randn((256, 512), dtype=torch.bfloat16, device="cuda") + layer = _make_layer(weight) + kernel = FlashInferTrtllmMxfp8LinearKernel(Mxfp8LinearLayerConfig()) + kernel.process_weights_after_loading(layer) + + torch.library.opcheck( + torch.ops.vllm.flashinfer_mxfp8_quantize_8x4.default, + (x,), + ) + x_mxfp8, x_scale = vllm_flashinfer.flashinfer_mxfp8_quantize_8x4(x) + # SchemaCheckMode compares inputs with allclose, which CUDA does not + # implement for float8. The numerical tests above guard input mutation. + torch.library.opcheck( + torch.ops.vllm.mm_mxfp8.default, + ( + x_mxfp8, + layer.weight.t(), + x_scale, + layer.weight_scale, + torch.bfloat16, + "trtllm", + True, + ), + test_utils=( + "test_autograd_registration", + "test_faketensor", + "test_aot_dispatch_dynamic", + ), + ) + + +@torch.inference_mode() +def test_flashinfer_trtllm_mxfp8_linear_cuda_graph() -> None: + torch.manual_seed(0) + m, n, k = 7, 130, 512 + weight = torch.randn((n, k), dtype=torch.bfloat16, device="cuda") + layer = _make_layer(weight) + kernel = FlashInferTrtllmMxfp8LinearKernel(Mxfp8LinearLayerConfig()) + kernel.process_weights_after_loading(layer) + + static_x = torch.randn((m, k), dtype=torch.bfloat16, device="cuda") + kernel.apply_weights(layer, static_x) + torch.accelerator.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_output = kernel.apply_weights(layer, static_x) + + new_x = torch.randn_like(static_x) + static_x.copy_(new_x) + graph.replay() + eager_output = kernel.apply_weights(layer, new_x) + + torch.testing.assert_close(graph_output, eager_output, rtol=0, atol=0) diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index b94c8f7d58d0..16d3098df3b5 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -115,6 +115,7 @@ from vllm.model_executor.kernels.linear.mxfp8.flashinfer import ( FlashInferCutedslMxfp8LinearKernel, FlashInferCutlassMxfp8LinearKernel, + FlashInferTrtllmMxfp8LinearKernel, ) from vllm.model_executor.kernels.linear.mxfp8.humming import ( HummingMxfp8LinearKernel, @@ -267,6 +268,7 @@ def _get_linear_backend() -> str: FlashInferCutedslMxfp8LinearKernel, }, "flashinfer_trtllm": { + FlashInferTrtllmMxfp8LinearKernel, FlashInferTrtllmNvFp4LinearKernel, }, "flashinfer_cudnn": { @@ -507,6 +509,7 @@ def _resolve_backend_kernels( B12xMxfp8LinearKernel, EmulationMxfp8LinearKernel, HummingMxfp8LinearKernel, + FlashInferTrtllmMxfp8LinearKernel, ], PlatformEnum.ROCM: [ # Native CDNA4 (gfx950) MX linear; is_supported() gates to gfx95x and @@ -1216,6 +1219,7 @@ def register_linear_kernel( "MarlinMxFp4LinearKernel", "FlashInferCutedslMxfp8LinearKernel", "FlashInferCutlassMxfp8LinearKernel", + "FlashInferTrtllmMxfp8LinearKernel", "MarlinMxfp8LinearKernel", "XPUMxFp8LinearKernel", "EmulationMxfp8LinearKernel", diff --git a/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py b/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py index 3150d93f99e1..6ea9138de92a 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py +++ b/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py @@ -177,3 +177,103 @@ def apply_weights( output_shape = (*input_shape[:-1], N) return output.view(output_shape) + + +class FlashInferTrtllmMxfp8LinearKernel(Mxfp8LinearKernel): + """MXFP8 W8A8 GEMM via FlashInfer's TensorRT-LLM wrapper.""" + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not ( + current_platform.is_cuda() + and current_platform.is_device_capability_family(100) + ): + return False, "requires SM100-family GPU" + if not has_flashinfer(): + return False, "requires FlashInfer" + return True, None + + @classmethod + def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]: + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + from flashinfer import shuffle_matrix_a, shuffle_matrix_sf_a + + if hasattr(layer, "_mxfp8_trtllm_output_size") and layer.weight_scale.ndim == 1: + return + + weight = layer.weight.data # [N, K] + N, K = weight.shape + if K % 256 != 0: + raise ValueError( + f"FlashInfer TRTLLM MXFP8 requires K to be divisible by 256, got K={K}." + ) + + scale_k = K // MXFP8_BLOCK_SIZE + weight_scale = layer.weight_scale.data[:N, :scale_k].contiguous() + padded_n = ((N + 127) // 128) * 128 + if padded_n != N: + padded_weight = weight.new_zeros((padded_n, K)) + padded_weight[:N] = weight + weight = padded_weight + + padded_scale = weight_scale.new_zeros((padded_n, scale_k)) + padded_scale[:N] = weight_scale + weight_scale = padded_scale + else: + weight = weight.contiguous() + + layer.weight = Parameter( + shuffle_matrix_a(weight, 128).reshape(padded_n, K), + requires_grad=False, + ) + layer.weight_scale = Parameter( + shuffle_matrix_sf_a( + weight_scale, + 128, + num_elts_per_sf=MXFP8_BLOCK_SIZE, + ).reshape(-1), + requires_grad=False, + ) + layer._mxfp8_trtllm_output_size = N + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + assert x.dtype == torch.bfloat16, ( + f"FlashInfer TRTLLM MXFP8 requires bfloat16 activations, got {x.dtype}." + ) + + weight = layer.weight # shuffled [padded N, K] + weight_scale = layer.weight_scale + _, K = weight.shape + output_size = layer._mxfp8_trtllm_output_size + input_shape = x.shape + input_2d = x.view(-1, K) + + input_mxfp8, input_scale = vllm_flashinfer.flashinfer_mxfp8_quantize_8x4( + input_2d + ) + output = vllm_flashinfer.mm_mxfp8( + input_mxfp8, + weight.t(), + input_scale, + weight_scale, + out_dtype=x.dtype, + backend="trtllm", + use_8x4_sf_layout=True, + ) + if output.shape[-1] != output_size: + output = output[:, :output_size].contiguous() + + if bias is not None: + output = output + bias + + output_shape = (*input_shape[:-1], output_size) + return output.view(output_shape) diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index 4f7d2eecd780..f77b669cc266 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -726,6 +726,36 @@ def flashinfer_nvfp4_quantize_fake( rounded_m, rounded_n, dtype=torch.uint8, device=a.device ) + @torch.library.custom_op( + "vllm::flashinfer_mxfp8_quantize_8x4", + mutates_args=[], + device_types="cuda", + ) + def flashinfer_mxfp8_quantize_8x4( + a: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + from flashinfer import SfLayout + from flashinfer import mxfp8_quantize as mxfp8_quantize_ + + return mxfp8_quantize_( + a, + backend="cuda", + sf_swizzle_layout=SfLayout.layout_8x4, + ) + + @torch.library.register_fake( + "vllm::flashinfer_mxfp8_quantize_8x4", + ) + def flashinfer_mxfp8_quantize_8x4_fake( + a: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + m, k = a.shape + scale_size = cdiv(m, 8) * 8 * cdiv(k // 32, 4) * 4 + return ( + torch.empty(m, k, dtype=torch.float8_e4m3fn, device=a.device), + torch.empty(scale_size, dtype=torch.uint8, device=a.device), + ) + @torch.library.custom_op( "vllm::mm_mxfp8", mutates_args=[], @@ -738,6 +768,7 @@ def mm_mxfp8( B_scale: torch.Tensor, out_dtype: torch.dtype, backend: str = "cutlass", + use_8x4_sf_layout: bool = False, ) -> torch.Tensor: from flashinfer import mm_mxfp8 as mm_mxfp8_ @@ -749,6 +780,7 @@ def mm_mxfp8( out=None, out_dtype=out_dtype, backend=backend, + use_8x4_sf_layout=use_8x4_sf_layout, ) @torch.library.register_fake( @@ -761,6 +793,7 @@ def mm_mxfp8_fake( B_scale: torch.Tensor, out_dtype: torch.dtype, backend: str = "cutlass", + use_8x4_sf_layout: bool = False, ) -> torch.Tensor: # A is [m, k], B is [k, n] -> output [m, n] return torch.empty(A.shape[0], B.shape[1], dtype=out_dtype, device=A.device) From cb09dd7488d2ad13dbf610aa2a580af2e7b050b5 Mon Sep 17 00:00:00 2001 From: Yiqin <61896954+Yiqin-17@users.noreply.github.com> Date: Thu, 20 Aug 2026 23:04:29 +0800 Subject: [PATCH 212/839] [Core][Multimodal] Skip redundant placeholder scan when token match succeeds (#52925) Co-authored-by: shenyiqin --- tests/multimodal/test_processing.py | 506 ++++++++++++++++------- vllm/model_executor/models/gemma3_mm.py | 34 ++ vllm/model_executor/models/gemma3n_mm.py | 34 ++ vllm/multimodal/processing/processor.py | 140 +++++-- 4 files changed, 539 insertions(+), 175 deletions(-) diff --git a/tests/multimodal/test_processing.py b/tests/multimodal/test_processing.py index fa2f76730d21..601fe9b8a819 100644 --- a/tests/multimodal/test_processing.py +++ b/tests/multimodal/test_processing.py @@ -17,6 +17,7 @@ PromptInsertion, PromptReplacement, _apply_matches, + _apply_token_matches_with_placeholders, apply_text_matches, apply_token_matches, find_mm_placeholders, @@ -567,171 +568,309 @@ def test_find_update_text( assert new_prompt == expected -@pytest.mark.parametrize( - ("prompt", "target_by_key", "repl_by_key", "expected_by_update_type_mm_count"), # noqa: E501 - [ - # Tokenized test cases of `test_find_update_text` - # using the vocab of llava-hf/llava-v1.6-mistral-7b-hf - ( - [1, 9833, 28747, 32000, 9833, 28747, 32000, 32000, 918], - { - # We use `` before `Image:` to test matches that - # occur out of order - "pattern_1": [32000], - "pattern_2": [9833, 28747], - "pattern_3": [918], +FIND_UPDATE_TOKENS_TEST_CASES = [ + # Tokenized test cases of `test_find_update_text` + # using the vocab of llava-hf/llava-v1.6-mistral-7b-hf + ( + [1, 9833, 28747, 32000, 9833, 28747, 32000, 32000, 918], + { + # We use `` before `Image:` to test matches that + # occur out of order + "pattern_1": [32000], + "pattern_2": [9833, 28747], + "pattern_3": [918], + }, + { + # Test whether target is confused with replacement + "pattern_1": [32000, 32000], + # Test empty replacement + "pattern_2": [], + # Test dynamic replacement (beyond the form of `unit * count`) + "pattern_3": [1550, 918, 1550], + }, + { + PromptInsertion: { + 0: [1, 9833, 28747, 32000, 9833, 28747, 32000, 32000, 918], + 1: [ + 1, + 9833, + 28747, + 32000, + 32000, + 32000, + 9833, + 28747, + 32000, + 32000, + 918, + 1550, + 918, + 1550, + ], # noqa: E501 + 2: [ + 1, + 9833, + 28747, + 32000, + 32000, + 32000, + 32000, + 32000, + 9833, + 28747, + 32000, + 32000, + 918, + 1550, + 918, + 1550, + 1550, + 918, + 1550, + ], # noqa: E501 }, - { - # Test whether target is confused with replacement - "pattern_1": [32000, 32000], - # Test empty replacement - "pattern_2": [], - # Test dynamic replacement (beyond the form of `unit * count`) - "pattern_3": [1550, 918, 1550], + PromptReplacement: { + 0: [1, 9833, 28747, 32000, 9833, 28747, 32000, 32000, 918], + 1: [1, 32000, 32000, 9833, 28747, 32000, 32000, 1550, 918, 1550], # noqa: E501 + 2: [1, 32000, 32000, 32000, 32000, 32000, 1550, 918, 1550], }, - { - PromptInsertion: { - 0: [1, 9833, 28747, 32000, 9833, 28747, 32000, 32000, 918], - 1: [ - 1, - 9833, - 28747, - 32000, - 32000, - 32000, - 9833, - 28747, - 32000, - 32000, - 918, - 1550, - 918, - 1550, - ], # noqa: E501 - 2: [ - 1, - 9833, - 28747, - 32000, - 32000, - 32000, - 32000, - 32000, - 9833, - 28747, - 32000, - 32000, - 918, - 1550, - 918, - 1550, - 1550, - 918, - 1550, - ], # noqa: E501 - }, - PromptReplacement: { - 0: [1, 9833, 28747, 32000, 9833, 28747, 32000, 32000, 918], - 1: [1, 32000, 32000, 9833, 28747, 32000, 32000, 1550, 918, 1550], # noqa: E501 - 2: [1, 32000, 32000, 32000, 32000, 32000, 1550, 918, 1550], - }, + }, + ), + # Test index targets + ( + [], + { + "pattern_1": PromptIndexTargets.start(), + "pattern_2": PromptIndexTargets.prefix([32000]), + "pattern_3": PromptIndexTargets.end(), + }, + { + "pattern_1": [-1], + "pattern_2": [-2], + "pattern_3": [-3], + }, + { + PromptInsertion: { + 0: [], + 1: [-1, -3], + 2: [-1, -1, -3, -3], }, - ), - # Test index targets - ( - [], - { - "pattern_1": PromptIndexTargets.start(), - "pattern_2": PromptIndexTargets.prefix([32000]), - "pattern_3": PromptIndexTargets.end(), + PromptReplacement: { + 0: [], + 1: [-1, -3], + 2: [-1, -1, -3, -3], }, - { - "pattern_1": [-1], - "pattern_2": [-2], - "pattern_3": [-3], + }, + ), + ( + [32000], + { + "pattern_1": PromptIndexTargets.start(), + "pattern_2": PromptIndexTargets.prefix([32000]), + "pattern_3": PromptIndexTargets.end(), + }, + { + "pattern_1": [-1], + "pattern_2": [-2], + "pattern_3": [-3], + }, + { + PromptInsertion: { + 0: [32000], + 1: [-1, 32000, -2, -3], + 2: [-1, -1, 32000, -2, -2, -3, -3], }, - { - PromptInsertion: { - 0: [], - 1: [-1, -3], - 2: [-1, -1, -3, -3], - }, - PromptReplacement: { - 0: [], - 1: [-1, -3], - 2: [-1, -1, -3, -3], - }, + PromptReplacement: { + 0: [32000], + 1: [-1, 32000, -2, -3], + 2: [-1, -1, 32000, -2, -2, -3, -3], }, - ), - ( - [32000], - { - "pattern_1": PromptIndexTargets.start(), - "pattern_2": PromptIndexTargets.prefix([32000]), - "pattern_3": PromptIndexTargets.end(), + }, + ), + # Test different replacement per item + ( + [32000, 32000, 32000], + { + "pattern_1": [32000], + }, + { + "pattern_1": lambda idx: [-(idx + 1)], + }, + { + PromptInsertion: { + 0: [32000, 32000, 32000], + 1: [32000, -1, 32000, 32000], + 2: [32000, -1, -2, 32000, 32000], }, - { - "pattern_1": [-1], - "pattern_2": [-2], - "pattern_3": [-3], + PromptReplacement: { + 0: [32000, 32000, 32000], + 1: [-1, 32000, 32000], + 2: [-1, -2, 32000], }, - { - PromptInsertion: { - 0: [32000], - 1: [-1, 32000, -2, -3], - 2: [-1, -1, 32000, -2, -2, -3, -3], - }, - PromptReplacement: { - 0: [32000], - 1: [-1, 32000, -2, -3], - 2: [-1, -1, 32000, -2, -2, -3, -3], - }, + }, + ), + ( + [32000, 32000, 32000], + { + "pattern_1": PromptIndexTargets.prefix([32000]), + }, + { + "pattern_1": lambda idx: [-(idx + 1)], + }, + { + PromptInsertion: { + 0: [32000, 32000, 32000], + 1: [32000, -1, 32000, 32000], + 2: [32000, -1, -2, 32000, 32000], }, - ), - # Test different replacement per item - ( - [32000, 32000, 32000], - { - "pattern_1": [32000], + PromptReplacement: { + 0: [32000, 32000, 32000], + 1: [32000, -1, 32000, 32000], + 2: [32000, -1, -2, 32000, 32000], }, - { - "pattern_1": lambda idx: [-(idx + 1)], + }, + ), +] + + +def _placeholder(modality, item_idx, start_idx, tokens): + return PlaceholderFeaturesInfo( + modality=modality, + item_idx=item_idx, + start_idx=start_idx, + tokens=tokens, + is_embed=None, + ) + + +FIND_UPDATE_TOKENS_PLACEHOLDER_EXPECTED = [ + { + PromptInsertion: { + 0: {}, + 1: { + "pattern_1": [_placeholder("pattern_1", 0, 4, [32000, 32000])], + "pattern_3": [_placeholder("pattern_3", 0, 11, [1550, 918, 1550])], }, - { - PromptInsertion: { - 0: [32000, 32000, 32000], - 1: [32000, -1, 32000, 32000], - 2: [32000, -1, -2, 32000, 32000], - }, - PromptReplacement: { - 0: [32000, 32000, 32000], - 1: [-1, 32000, 32000], - 2: [-1, -2, 32000], - }, + 2: { + "pattern_1": [ + _placeholder("pattern_1", 0, 4, [32000, 32000]), + _placeholder("pattern_1", 1, 6, [32000, 32000]), + ], + "pattern_3": [ + _placeholder("pattern_3", 0, 13, [1550, 918, 1550]), + _placeholder("pattern_3", 1, 16, [1550, 918, 1550]), + ], }, - ), - ( - [32000, 32000, 32000], - { - "pattern_1": PromptIndexTargets.prefix([32000]), + }, + PromptReplacement: { + 0: {}, + 1: { + "pattern_1": [_placeholder("pattern_1", 0, 1, [32000, 32000])], + "pattern_3": [_placeholder("pattern_3", 0, 7, [1550, 918, 1550])], }, - { - "pattern_1": lambda idx: [-(idx + 1)], + 2: {}, + }, + }, + { + PromptInsertion: {0: {}, 1: {}, 2: {}}, + PromptReplacement: {0: {}, 1: {}, 2: {}}, + }, + { + PromptInsertion: { + 0: {}, + 1: { + "pattern_1": [_placeholder("pattern_1", 0, 0, [-1])], + "pattern_2": [_placeholder("pattern_2", 0, 2, [-2])], + "pattern_3": [_placeholder("pattern_3", 0, 3, [-3])], }, - { - PromptInsertion: { - 0: [32000, 32000, 32000], - 1: [32000, -1, 32000, 32000], - 2: [32000, -1, -2, 32000, 32000], - }, - PromptReplacement: { - 0: [32000, 32000, 32000], - 1: [32000, -1, 32000, 32000], - 2: [32000, -1, -2, 32000, 32000], - }, + 2: { + "pattern_1": [ + _placeholder("pattern_1", 0, 0, [-1]), + _placeholder("pattern_1", 1, 1, [-1]), + ], + "pattern_2": [ + _placeholder("pattern_2", 0, 3, [-2]), + _placeholder("pattern_2", 1, 4, [-2]), + ], + "pattern_3": [ + _placeholder("pattern_3", 0, 5, [-3]), + _placeholder("pattern_3", 1, 6, [-3]), + ], }, - ), - ], + }, + PromptReplacement: { + 0: {}, + 1: { + "pattern_1": [_placeholder("pattern_1", 0, 0, [-1])], + "pattern_2": [_placeholder("pattern_2", 0, 2, [-2])], + "pattern_3": [_placeholder("pattern_3", 0, 3, [-3])], + }, + 2: { + "pattern_1": [ + _placeholder("pattern_1", 0, 0, [-1]), + _placeholder("pattern_1", 1, 1, [-1]), + ], + "pattern_2": [ + _placeholder("pattern_2", 0, 3, [-2]), + _placeholder("pattern_2", 1, 4, [-2]), + ], + "pattern_3": [ + _placeholder("pattern_3", 0, 5, [-3]), + _placeholder("pattern_3", 1, 6, [-3]), + ], + }, + }, + }, + { + PromptInsertion: { + 0: {}, + 1: {"pattern_1": [_placeholder("pattern_1", 0, 1, [-1])]}, + 2: { + "pattern_1": [ + _placeholder("pattern_1", 0, 1, [-1]), + _placeholder("pattern_1", 1, 2, [-2]), + ] + }, + }, + PromptReplacement: { + 0: {}, + 1: {"pattern_1": [_placeholder("pattern_1", 0, 0, [-1])]}, + 2: { + "pattern_1": [ + _placeholder("pattern_1", 0, 0, [-1]), + _placeholder("pattern_1", 1, 1, [-2]), + ] + }, + }, + }, + { + PromptInsertion: { + 0: {}, + 1: {"pattern_1": [_placeholder("pattern_1", 0, 1, [-1])]}, + 2: { + "pattern_1": [ + _placeholder("pattern_1", 0, 1, [-1]), + _placeholder("pattern_1", 1, 2, [-2]), + ] + }, + }, + PromptReplacement: { + 0: {}, + 1: {"pattern_1": [_placeholder("pattern_1", 0, 1, [-1])]}, + 2: { + "pattern_1": [ + _placeholder("pattern_1", 0, 1, [-1]), + _placeholder("pattern_1", 1, 2, [-2]), + ] + }, + }, + }, +] + + +@pytest.mark.parametrize( + ("prompt", "target_by_key", "repl_by_key", "expected_by_update_type_mm_count"), # noqa: E501 + FIND_UPDATE_TOKENS_TEST_CASES, ) def test_find_update_tokens( prompt, @@ -769,6 +908,73 @@ def test_find_update_tokens( assert new_prompt == expected +@pytest.mark.parametrize( + ( + "prompt", + "target_by_key", + "repl_by_key", + "expected_by_update_type_mm_count", + "expected_placeholders_by_update_type_mm_count", + ), + [ + (*case, placeholder_expected) + for case, placeholder_expected in zip( + FIND_UPDATE_TOKENS_TEST_CASES, + FIND_UPDATE_TOKENS_PLACEHOLDER_EXPECTED, + strict=True, + ) + ], +) +def test_apply_token_matches_with_placeholders( + prompt, + target_by_key, + repl_by_key, + expected_by_update_type_mm_count, + expected_placeholders_by_update_type_mm_count, +): + for update_type, expected_by_mm_count in expected_by_update_type_mm_count.items(): + for mm_count, expected in expected_by_mm_count.items(): + mm_prompt_updates = { + key: [ + [update_type(key, target, repl_by_key[key]).resolve(i)] + for i in range(mm_count) + ] + for key, target in target_by_key.items() + } + + new_prompt, result, placeholders = _apply_token_matches_with_placeholders( + prompt, + mm_prompt_updates, + tokenizer=None, + ) + + if any( + update_idx is None + for update_idxs in result.values() + for update_idx in update_idxs + ): + continue + + expected_placeholders = expected_placeholders_by_update_type_mm_count[ + update_type + ][mm_count] + + # Only displayed on error + print("update_type:", update_type) + print("mm_count:", mm_count) + print("mm_prompt_updates:", mm_prompt_updates) + print("new_prompt:", new_prompt) + print("result:", result) + print("placeholders:", placeholders) + + assert new_prompt == expected + assert { + modality: ph_list + for modality, ph_list in placeholders.items() + if ph_list + } == expected_placeholders + + @pytest.mark.parametrize( "repl_by_key", [ diff --git a/vllm/model_executor/models/gemma3_mm.py b/vllm/model_executor/models/gemma3_mm.py index 0d9f8f14188f..5551e0d8b9d3 100644 --- a/vllm/model_executor/models/gemma3_mm.py +++ b/vllm/model_executor/models/gemma3_mm.py @@ -373,6 +373,40 @@ def _apply_token_matches( return token_ids, res + def _apply_token_matches_with_placeholders( + self, + token_ids: list[int], + mm_prompt_updates: MultiModalPromptUpdates, + ) -> tuple[ + list[int], + MultiModalPromptUpdatesApplyResult, + Mapping[str, list[PlaceholderFeaturesInfo]], + ]: + new_token_ids, match_result = self._apply_token_matches( + token_ids, + mm_prompt_updates, + ) + + placeholders: dict[str, list[PlaceholderFeaturesInfo]] = { + modality: [] for modality in mm_prompt_updates + } + + if all( + all(update_idx is not None for update_idx in update_idxs) + for update_idxs in match_result.values() + ): + placeholders = dict( + self._find_mm_placeholders( + new_token_ids, + self._matched_updates_from_result( + mm_prompt_updates, + match_result, + ), + ) + ) + + return new_token_ids, match_result, placeholders + def _find_mm_placeholders( self, new_token_ids: list[int], diff --git a/vllm/model_executor/models/gemma3n_mm.py b/vllm/model_executor/models/gemma3n_mm.py index 2f5688853d62..3a01f1457aed 100644 --- a/vllm/model_executor/models/gemma3n_mm.py +++ b/vllm/model_executor/models/gemma3n_mm.py @@ -391,6 +391,40 @@ def _apply_token_matches( return token_ids, res + def _apply_token_matches_with_placeholders( + self, + token_ids: list[int], + mm_prompt_updates: MultiModalPromptUpdates, + ) -> tuple[ + list[int], + MultiModalPromptUpdatesApplyResult, + Mapping[str, list[PlaceholderFeaturesInfo]], + ]: + new_token_ids, match_result = self._apply_token_matches( + token_ids, + mm_prompt_updates, + ) + + placeholders: dict[str, list[PlaceholderFeaturesInfo]] = { + modality: [] for modality in mm_prompt_updates + } + + if all( + all(update_idx is not None for update_idx in update_idxs) + for update_idxs in match_result.values() + ): + placeholders = dict( + self._find_mm_placeholders( + new_token_ids, + self._matched_updates_from_result( + mm_prompt_updates, + match_result, + ), + ) + ) + + return new_token_ids, match_result, placeholders + def _find_mm_placeholders( self, new_token_ids: list[int], diff --git a/vllm/multimodal/processing/processor.py b/vllm/multimodal/processing/processor.py index d35856e04f03..6dc5bdd35063 100644 --- a/vllm/multimodal/processing/processor.py +++ b/vllm/multimodal/processing/processor.py @@ -958,6 +958,64 @@ def apply_token_matches( return flatten_2d_lists(token_id_seqs), result +def _apply_token_matches_with_placeholders( + token_ids: list[int], + mm_prompt_updates: "MultiModalPromptUpdates", + tokenizer: TokenizerLike | None, +) -> tuple[ + list[int], + "MultiModalPromptUpdatesApplyResult", + Mapping[str, list[PlaceholderFeaturesInfo]], +]: + matched_updates, result = _plan_prompt_updates( + token_ids, + mm_prompt_updates, + tokenizer, + ) + placeholders: dict[str, list[PlaceholderFeaturesInfo]] = { + modality: [] for modality in mm_prompt_updates + } + + new_token_ids = list[int]() + prev_end_idx = 0 + for matched_update in matched_updates: + update = matched_update.update + match = matched_update.match + matched_content = update.content.full + + if update.mode == UpdateMode.INSERT: + end_idx_to_insert = match.end_idx + elif update.mode == UpdateMode.REPLACE: + end_idx_to_insert = match.start_idx + else: + assert_never(update.mode) + + new_token_ids.extend(token_ids[prev_end_idx:end_idx_to_insert]) + start_idx = len(new_token_ids) + + tokens = _seq2tokens(tokenizer, matched_content) + if tokens: + content_is_embed = update.content.is_embed + if content_is_embed is not None: + content_is_embed = content_is_embed(tokenizer, matched_content) + + placeholders[update.modality].append( + PlaceholderFeaturesInfo( + modality=update.modality, + item_idx=update.item_idx, + start_idx=start_idx, + tokens=tokens, + is_embed=content_is_embed, + ) + ) + new_token_ids.extend(tokens) + + prev_end_idx = match.end_idx + + new_token_ids.extend(token_ids[prev_end_idx:]) + return new_token_ids, result, placeholders + + def apply_text_matches( prompt: str, mm_prompt_updates: "MultiModalPromptUpdates", @@ -1671,6 +1729,22 @@ def _apply_token_matches( tokenizer = self.info.get_tokenizer() return apply_token_matches(prompt, mm_prompt_updates, tokenizer) + def _apply_token_matches_with_placeholders( + self, + token_ids: list[int], + mm_prompt_updates: MultiModalPromptUpdates, + ) -> tuple[ + list[int], + MultiModalPromptUpdatesApplyResult, + Mapping[str, list[PlaceholderFeaturesInfo]], + ]: + tokenizer = self.info.get_tokenizer() + return _apply_token_matches_with_placeholders( + token_ids, + mm_prompt_updates, + tokenizer, + ) + def _apply_text_matches( self, prompt: str, @@ -1689,6 +1763,25 @@ def _apply_text_matches_as_segmented_tokens( prompt, mm_prompt_updates, tokenizer ) + def _matched_updates_from_result( + self, + mm_prompt_updates: MultiModalPromptUpdates, + match_result: MultiModalPromptUpdatesApplyResult, + ) -> dict[str, list[Sequence[ResolvedPromptUpdate]]]: + matched_updates = defaultdict[str, list[Sequence[ResolvedPromptUpdate]]](list) + for modality, update_idxs in match_result.items(): + for item_idx, update_idx in enumerate(update_idxs): + assert update_idx is not None, ( + "Failed to apply prompt replacement for " + f"mm_items[{modality!r}][{item_idx}]" + ) + + matched_updates[modality].append( + [mm_prompt_updates[modality][item_idx][update_idx]] + ) + + return dict(matched_updates) + def _apply_prompt_updates( self, token_ids: list[int], @@ -1697,11 +1790,24 @@ def _apply_prompt_updates( """Apply multi-modal prompt updates to token IDs.""" tokenizer = self.info.get_tokenizer() - new_token_ids, match_result = self._apply_token_matches( - token_ids, - mm_prompt_updates, + new_token_ids, match_result, placeholders = ( + self._apply_token_matches_with_placeholders( + token_ids, + mm_prompt_updates, + ) ) + if all( + all(update_idx is not None for update_idx in update_idxs) + for update_idxs in match_result.values() + ): + placeholders = { + modality: modality_placeholders + for modality, modality_placeholders in placeholders.items() + if modality_placeholders + } + return new_token_ids, placeholders + # If the search text does not represent a special token, # it may have different token IDs in the prompt, because # the tokens may go across the boundaries of the search text. @@ -1712,32 +1818,16 @@ def _apply_prompt_updates( # Since it is inefficient to search for all possible tokenizations # of the search text in the prompt, we instead perform string-based # updates on the decoded token IDs, then encode them back. - if not all( - all(update_idx is not None for update_idx in update_idxs) - for update_idxs in match_result.values() - ): - new_text, match_result = self._apply_text_matches( - _seq2text(tokenizer, token_ids, use_cache=False), - mm_prompt_updates, - ) - - new_token_ids = _seq2tokens(tokenizer, new_text, use_cache=False) - - matched_updates = defaultdict[str, list[Sequence[ResolvedPromptUpdate]]](list) - for modality, update_idxs in match_result.items(): - for item_idx, update_idx in enumerate(update_idxs): - assert update_idx is not None, ( - "Failed to apply prompt replacement for " - f"mm_items[{modality!r}][{item_idx}]" - ) + new_text, match_result = self._apply_text_matches( + _seq2text(tokenizer, token_ids, use_cache=False), + mm_prompt_updates, + ) - matched_updates[modality].append( - [mm_prompt_updates[modality][item_idx][update_idx]] - ) + new_token_ids = _seq2tokens(tokenizer, new_text, use_cache=False) placeholders = self._find_mm_placeholders( new_token_ids, - dict(matched_updates), + self._matched_updates_from_result(mm_prompt_updates, match_result), ) return new_token_ids, placeholders From 01af92e175407231b1433b0aef01a1b9c983d955 Mon Sep 17 00:00:00 2001 From: Zupeng Wang Date: Thu, 20 Aug 2026 23:13:43 +0800 Subject: [PATCH 213/839] [Feature][Model Runner V2] Support extract_hidden_states speculation (#49811) Signed-off-by: Zupeng Wang <71580390+zupengwang@users.noreply.github.com> Signed-off-by: Misha Goin Co-authored-by: OpenAI Co-authored-by: Misha Goin Co-authored-by: OpenAI Codex --- tests/test_config.py | 15 ++ ...st_gpu_extract_hidden_states_speculator.py | 132 +++++++++++++++ vllm/config/vllm.py | 1 + vllm/v1/worker/gpu/model_runner.py | 7 +- vllm/v1/worker/gpu/spec_decode/__init__.py | 8 +- .../gpu/spec_decode/extract_hidden_states.py | 153 ++++++++++++++++++ 6 files changed, 314 insertions(+), 2 deletions(-) create mode 100644 tests/v1/worker/test_gpu_extract_hidden_states_speculator.py create mode 100644 vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py diff --git a/tests/test_config.py b/tests/test_config.py index 7ec18b3e400e..ee3e0c96b58d 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -5,6 +5,7 @@ import os from dataclasses import MISSING, Field, asdict, dataclass, field from types import SimpleNamespace +from typing import cast from unittest.mock import patch import pydantic @@ -263,6 +264,20 @@ def test_dsa_models_select_matching_mtp(model_type, expected_architecture): assert hf_config.architectures == [expected_architecture] +def test_v2_model_runner_supports_extract_hidden_states(): + config = VllmConfig() + config.speculative_config = cast( + SpeculativeConfig, + SimpleNamespace( + method="extract_hidden_states", + parallel_drafting=False, + enable_adaptive_verification=False, + ), + ) + + assert config._get_v2_model_runner_unsupported_features() == [] + + @pytest.mark.parametrize( ("use_v2_model_runner", "expected_capture_sizes"), [ diff --git a/tests/v1/worker/test_gpu_extract_hidden_states_speculator.py b/tests/v1/worker/test_gpu_extract_hidden_states_speculator.py new file mode 100644 index 000000000000..d1aa0798feb2 --- /dev/null +++ b/tests/v1/worker/test_gpu_extract_hidden_states_speculator.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from contextlib import nullcontext +from types import SimpleNamespace +from typing import Any, cast + +import pytest +import torch + +from vllm.v1.worker.gpu.spec_decode import extract_hidden_states as spec_module +from vllm.v1.worker.gpu.spec_decode import init_speculator +from vllm.v1.worker.gpu.spec_decode.extract_hidden_states import ( + ExtractHiddenStatesSpeculator, +) + + +class _RecordingModel(torch.nn.Module): + def forward(self, *, hidden_states: torch.Tensor) -> None: + self.hidden_states = hidden_states.clone() + + +def test_init_requires_greedy_draft_sampling(): + vllm_config = cast( + Any, + SimpleNamespace( + speculative_config=SimpleNamespace(draft_sample_method="probabilistic") + ), + ) + + with pytest.raises(ValueError, match="only supports draft_sample_method='greedy'"): + ExtractHiddenStatesSpeculator(vllm_config, torch.device("cpu")) + + +def test_init_speculator_dispatches_extract_hidden_states(monkeypatch): + vllm_config = cast( + Any, + SimpleNamespace( + speculative_config=SimpleNamespace(method="extract_hidden_states") + ), + ) + device = torch.device("cpu") + + def fake_speculator(config, target_device): + return config, target_device + + monkeypatch.setattr(spec_module, "ExtractHiddenStatesSpeculator", fake_speculator) + + assert init_speculator(vllm_config, device) == (vllm_config, device) + + +def test_propose_caches_hidden_states_and_returns_sampled_tokens(monkeypatch): + contexts = [] + + def fake_set_forward_context(*args, **kwargs): + contexts.append((args, kwargs)) + return nullcontext() + + monkeypatch.setattr(spec_module, "set_forward_context", fake_set_forward_context) + + layer_name = "cache_only_layers.2" + speculator = object.__new__(ExtractHiddenStatesSpeculator) + speculator.vllm_config = cast(Any, SimpleNamespace()) + speculator.num_hidden_states = 2 + speculator.hidden_states = torch.zeros(4, 2, 3) + speculator.draft_attn_layer_names = {layer_name} + speculator.model = _RecordingModel() + + input_batch = cast( + Any, + SimpleNamespace( + idx_mapping=torch.tensor([2, 0], dtype=torch.int32), + is_padding=torch.zeros(4, dtype=torch.bool), + ), + ) + aux_hidden_states = [ + torch.full((4, 3), 1.0), + torch.full((4, 3), 2.0), + ] + attn_metadata = {layer_name: object(), "target_layer": object()} + slot_mappings = { + layer_name: torch.arange(4), + "target_layer": torch.arange(4), + } + last_sampled = torch.tensor([[10], [11], [12]], dtype=torch.int64) + + draft_tokens = speculator.propose( + input_batch=input_batch, + attn_metadata=attn_metadata, + slot_mappings=slot_mappings, + last_hidden_states=torch.empty(0), + aux_hidden_states=aux_hidden_states, + num_sampled=torch.empty(0), + num_rejected=torch.empty(0), + last_sampled=last_sampled, + next_prefill_tokens=torch.empty(0), + temperature=torch.empty(0), + seeds=torch.empty(0), + ) + + expected_hidden_states = torch.stack(aux_hidden_states, dim=1) + assert torch.equal(speculator.model.hidden_states, expected_hidden_states) + assert torch.equal(draft_tokens, torch.tensor([[12], [10]])) + + assert len(contexts) == 1 + args, kwargs = contexts[0] + assert args[0] == {layer_name: attn_metadata[layer_name]} + assert kwargs["num_tokens"] == 4 + assert set(kwargs["slot_mapping"]) == {layer_name} + assert torch.equal(kwargs["slot_mapping"][layer_name], slot_mappings[layer_name]) + + +def test_propose_requires_aux_hidden_states(): + speculator = object.__new__(ExtractHiddenStatesSpeculator) + speculator.num_hidden_states = 2 + input_batch = cast( + Any, SimpleNamespace(idx_mapping=torch.tensor([0], dtype=torch.int32)) + ) + + with pytest.raises(ValueError, match="aux_hidden_states are required"): + speculator.propose( + input_batch=input_batch, + attn_metadata={}, + slot_mappings={}, + last_hidden_states=torch.empty(0), + aux_hidden_states=None, + num_sampled=torch.empty(0), + num_rejected=torch.empty(0), + last_sampled=torch.tensor([[10]]), + next_prefill_tokens=torch.empty(0), + temperature=torch.empty(0), + seeds=torch.empty(0), + ) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index ca009b91840a..5146a362254b 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -2429,6 +2429,7 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]: "mtp", "dflash", "dspark", + "extract_hidden_states", ): unsupported.append(f"speculative method '{speculative_config.method}'") diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 66b3c8db4a1e..93a8af4f062e 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -247,7 +247,12 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): if self.is_last_pp_rank: self.speculator = init_speculator(self.vllm_config, self.device) - if self.speculative_config.method in ("eagle3", "dflash", "dspark"): + if self.speculative_config.method in ( + "eagle3", + "dflash", + "dspark", + "extract_hidden_states", + ): # Drafting may require auxiliary hidden states from target model outputs self.use_aux_hidden_state_outputs = True if self.use_pp: diff --git a/vllm/v1/worker/gpu/spec_decode/__init__.py b/vllm/v1/worker/gpu/spec_decode/__init__.py index 4229696f255c..2f75109893d0 100644 --- a/vllm/v1/worker/gpu/spec_decode/__init__.py +++ b/vllm/v1/worker/gpu/spec_decode/__init__.py @@ -8,7 +8,13 @@ def init_speculator(vllm_config: VllmConfig, device: torch.device): speculative_config = vllm_config.speculative_config assert speculative_config is not None - if speculative_config.method == "dflash": + if speculative_config.method == "extract_hidden_states": + from vllm.v1.worker.gpu.spec_decode.extract_hidden_states import ( + ExtractHiddenStatesSpeculator, + ) + + return ExtractHiddenStatesSpeculator(vllm_config, device) + elif speculative_config.method == "dflash": from vllm.v1.worker.gpu.spec_decode.dflash.speculator import ( DFlashSpeculator, ) diff --git a/vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py b/vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py new file mode 100644 index 000000000000..efcdec80182e --- /dev/null +++ b/vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py @@ -0,0 +1,153 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import Any + +import torch +import torch.nn as nn + +from vllm.compilation.backends import set_model_tag +from vllm.config import VllmConfig +from vllm.config.compilation import CUDAGraphMode +from vllm.forward_context import set_forward_context +from vllm.model_executor.model_loader import get_model +from vllm.v1.worker.gpu.input_batch import InputBatch +from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator + + +class ExtractHiddenStatesSpeculator(DraftModelSpeculator): + """Cache target hidden states while returning always-accepted draft tokens.""" + + def __init__(self, vllm_config: VllmConfig, device: torch.device): + assert vllm_config.speculative_config is not None + if vllm_config.speculative_config.draft_sample_method != "greedy": + raise ValueError( + "extract_hidden_states only supports draft_sample_method='greedy'" + ) + super().__init__(vllm_config, device) + + if self.num_speculative_steps != 1: + raise ValueError( + "extract_hidden_states requires num_speculative_tokens to be 1" + ) + if self.speculative_config.disable_padded_drafter_batch: + raise ValueError( + "disable_padded_drafter_batch is not supported with " + "extract_hidden_states method" + ) + + self.supports_mm_inputs = False + layer_ids = getattr( + self.draft_model_config.hf_config, + "eagle_aux_hidden_state_layer_ids", + None, + ) + if not layer_ids: + raise ValueError( + "eagle_aux_hidden_state_layer_ids must be set in the draft " + "model config for extract_hidden_states method" + ) + + self.num_hidden_states = len(layer_ids) + assert isinstance(self.dtype, torch.dtype) + self.hidden_states = torch.zeros( + self.max_num_tokens, + self.num_hidden_states, + self.vllm_config.model_config.get_hidden_size(), + dtype=self.dtype, + device=device, + ) + + def load_draft_model( + self, + target_model: nn.Module, + target_attn_layer_names: set[str], + ) -> nn.Module: + del target_model, target_attn_layer_names + with set_model_tag("extract_hidden_states"): + return get_model( + vllm_config=self.vllm_config, + model_config=self.draft_model_config, + ) + + def load_model(self, target_model: nn.Module) -> None: + super().load_model(target_model) + if len(self.draft_attn_layer_names) != 1: + raise ValueError( + "ExtractHiddenStatesModel should have exactly one attention " + f"layer, found {len(self.draft_attn_layer_names)}" + ) + + def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None: + del cudagraph_mode + + def capture(self) -> None: + return None + + @torch.inference_mode() + def propose( + self, + input_batch: InputBatch, + attn_metadata: dict[str, Any], + slot_mappings: dict[str, torch.Tensor], + last_hidden_states: torch.Tensor, + aux_hidden_states: list[torch.Tensor] | None, + num_sampled: torch.Tensor, + num_rejected: torch.Tensor, + last_sampled: torch.Tensor, + next_prefill_tokens: torch.Tensor, + temperature: torch.Tensor, + seeds: torch.Tensor, + num_tokens_across_dp: torch.Tensor | None = None, + dummy_run: bool = False, + skip_attn_for_dummy_run: bool = False, + mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, + is_profile: bool = False, + ) -> torch.Tensor: + del ( + last_hidden_states, + num_sampled, + num_rejected, + next_prefill_tokens, + temperature, + seeds, + dummy_run, + mm_inputs, + is_profile, + ) + + draft_tokens = last_sampled[input_batch.idx_mapping, :1] + if skip_attn_for_dummy_run: + return draft_tokens + if aux_hidden_states is None: + raise ValueError( + "aux_hidden_states are required when using extract_hidden_states" + ) + if len(aux_hidden_states) != self.num_hidden_states: + raise ValueError( + f"Expected {self.num_hidden_states} auxiliary hidden states, " + f"got {len(aux_hidden_states)}" + ) + + stacked_hidden_states = torch.stack(aux_hidden_states, dim=1) + num_tokens = stacked_hidden_states.shape[0] + self.hidden_states[:num_tokens].copy_(stacked_hidden_states) + + draft_attn_metadata = { + name: attn_metadata[name] for name in self.draft_attn_layer_names + } + draft_slot_mappings = { + name: slot_mappings[name][:num_tokens] + for name in self.draft_attn_layer_names + } + with set_forward_context( + draft_attn_metadata, + self.vllm_config, + num_tokens=num_tokens, + num_tokens_across_dp=num_tokens_across_dp, + cudagraph_runtime_mode=CUDAGraphMode.NONE, + slot_mapping=draft_slot_mappings, + is_padding=input_batch.is_padding[:num_tokens], + ): + self.model(hidden_states=self.hidden_states[:num_tokens]) + + return draft_tokens From 2dd17225c6a579c8ed0041943c636d74684fcbb5 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Thu, 20 Aug 2026 16:21:23 +0100 Subject: [PATCH 214/839] Fix Transformers modelling backend `RMSNormFuser.fuse` performance (#52766) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- .../transformers/fusers/test_rms_norm.py | 90 +++++++++++++++++++ .../models/transformers/fusers/mla.py | 7 +- .../models/transformers/fusers/rms_norm.py | 70 +++++++++++++-- 3 files changed, 157 insertions(+), 10 deletions(-) diff --git a/tests/models/transformers/fusers/test_rms_norm.py b/tests/models/transformers/fusers/test_rms_norm.py index 711f9b8b2217..ac259b66b48e 100644 --- a/tests/models/transformers/fusers/test_rms_norm.py +++ b/tests/models/transformers/fusers/test_rms_norm.py @@ -60,6 +60,54 @@ def forward(self, x): return self.weight * self._rms(x) +class NamedEpsRMSNorm(nn.Module): + """Holds eps under `attr`, with no `variance_epsilon` to find instead.""" + + attr = "eps" + + def __init__(self, hidden: int = 16, eps: float = 1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden)) + setattr(self, self.attr, eps) + + def forward(self, x): + eps = getattr(self, self.attr) + return self.weight * (x * torch.rsqrt(x.pow(2).mean(-1, True) + eps)) + + +class ToleranceRMSNorm(NamedEpsRMSNorm): + """Names eps something other than `eps`, so only its value identifies it.""" + + attr = "tolerance" + + +class LiteralEpsRMSNorm(RMSNorm): + """eps is a literal in the source, and an unrelated attribute happens to + hold the same value: matching on the value alone would bind to it.""" + + def __init__(self, hidden: int = 16, eps: float = 1e-5): + super().__init__(hidden, eps) + self.epsilon_lookalike = 1e-4 + + def forward(self, x): + return self.weight * (x * torch.rsqrt(x.pow(2).mean(-1, True) + 1e-4)) + + +class AmbiguousEpsRMSNorm(nn.Module): + """Two attributes hold the eps value, and the forward reads the second, so + their order cannot pick it.""" + + def __init__(self, hidden: int = 16, eps: float = 1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden)) + self.decoy_eps = eps + self.variance_epsilon = eps + + def forward(self, x): + var = x.pow(2).mean(-1, True) + return self.weight * (x * torch.rsqrt(var + self.variance_epsilon)) + + class NotAnRMSNorm(RMSNorm): """Mean-subtracting LayerNorm-like math -> not an RMSNorm.""" @@ -207,6 +255,48 @@ def test_eps_is_derived_per_instance(default_vllm_config): assert built.variance_epsilon == eps +@pytest.mark.parametrize("cls", [NamedEpsRMSNorm, ToleranceRMSNorm]) +def test_eps_attr_is_found_by_value_not_name(cls, default_vllm_config): + """The eps attribute is identified by holding the traced value, so a norm + stays per-instance correct whatever it names it.""" + with torch.device("meta"): + for eps in (1e-5, 1e-6): + module = cls(16, eps=eps) + fuser = get_fuser(module) + assert fuser.eps_attr == cls.attr + built = fuser.fuse(module, "norm", default_vllm_config) + assert built.variance_epsilon == eps + + +def test_literal_eps_is_not_mistaken_for_an_attribute(default_vllm_config, caplog): + """A literal eps is recognised as coming from no attribute, even when one + holds the same value, and is taken from the traced source instead.""" + logger = "vllm.model_executor.models.transformers.fusers.rms_norm" + with caplog.at_level("DEBUG", logger=logger), torch.device("meta"): + module = LiteralEpsRMSNorm() + fuser = get_fuser(module) + built = fuser.fuse(module, "norm", default_vllm_config) + assert fuser.eps_attr is None + assert built.variance_epsilon == 1e-4 + assert "does not hold its eps" in caplog.text + + +def test_ambiguous_eps_attrs_are_disambiguated(default_vllm_config): + """When several attributes hold the eps value, the one the forward actually + reads is identified, and they are all left as they were found.""" + with torch.device("meta"): + module = AmbiguousEpsRMSNorm(16, eps=1e-6) + before = dict(vars(module)) + fuser = get_fuser(module) + assert fuser.eps_attr == "variance_epsilon" + assert vars(module) == before + + other = AmbiguousEpsRMSNorm(16, eps=1e-5) + other.decoy_eps = 1e-3 + built = fuser.fuse(other, "norm", default_vllm_config) + assert built.variance_epsilon == 1e-5 + + def test_fused_norm_is_gather_capable(default_vllm_config): """Every weighted fused norm is emitted gather-capable, so a norm on a head-sharded projection (OLMoE-style) self-corrects at runtime with no QKV-specific diff --git a/vllm/model_executor/models/transformers/fusers/mla.py b/vllm/model_executor/models/transformers/fusers/mla.py index 8e17bff49024..db29ded91104 100644 --- a/vllm/model_executor/models/transformers/fusers/mla.py +++ b/vllm/model_executor/models/transformers/fusers/mla.py @@ -22,7 +22,6 @@ replace_expr, returned_linear, single_self_call, - trace, upstream_linear, ) from vllm.model_executor.models.transformers.utils import ( @@ -57,8 +56,10 @@ def _norm_size(norm: nn.Module) -> int: def _is_rms_norm(module: nn.Module) -> bool: """Whether `module` computes an RMSNorm, verified by `RMSNormFuser`'s matcher.""" - graph = trace(module) - return graph is not None and RMSNormFuser.match(graph, module) is not None + # Go via `get_fuser` because it caches + from vllm.model_executor.models.transformers.fuser import get_fuser + + return isinstance(get_fuser(module), RMSNormFuser) def _top_level_index(funcdef: ast.FunctionDef, node: ast.AST) -> int: diff --git a/vllm/model_executor/models/transformers/fusers/rms_norm.py b/vllm/model_executor/models/transformers/fusers/rms_norm.py index 7241f8539c19..47d20085a2c7 100644 --- a/vllm/model_executor/models/transformers/fusers/rms_norm.py +++ b/vllm/model_executor/models/transformers/fusers/rms_norm.py @@ -15,6 +15,7 @@ ) from vllm.distributed.parallel_state import model_parallel_is_initialized from vllm.distributed.utils import split_tensor_along_last_dim +from vllm.logger import init_logger from vllm.model_executor.models.transformers.fusers.base import BaseFuser from vllm.model_executor.models.transformers.fx_utils import ( find_node, @@ -32,6 +33,8 @@ if TYPE_CHECKING: from vllm.config import VllmConfig +logger = init_logger(__name__) + def _is_squared(node: object, x: fx.Node) -> bool: """`x**2`, `x.square()` or `x * x`, through any dtype casts.""" @@ -122,6 +125,10 @@ class RMSNormFuser(BaseFuser): """Gemma-style `(1 + weight)` scaling (weight initialised at zero).""" source_cls: str """Class name of the norm this was matched from (for logging).""" + eps_attr: str | None = None + """Attribute holding `eps`, read per instance in `fuse`.""" + eps: float | None = None + """`eps` itself, when it is not held in an attribute (see `_eps_source`).""" def info(self, name: str) -> str: norm = "GemmaRMSNorm" if self.zero_centered else "RMSNorm" @@ -140,7 +147,13 @@ def match(cls, graph: fx.Graph, module: nn.Module) -> "RMSNormFuser | None": if rms_norm is not None and rms_norm.args and peel(rms_norm.args[0]) is x: if _has_trailing_compute(graph, rms_norm): return None - return cls(zero_centered=False, source_cls=type(module).__name__) + eps_attr, eps = cls._eps_source(graph, module) + return cls( + zero_centered=False, + source_cls=type(module).__name__, + eps_attr=eps_attr, + eps=eps, + ) # Handle explicit `x * rsqrt(mean(x**2, -1) + eps)` pattern. # The rsqrt over the mean-square variance is the spine of the norm. rsqrt = None @@ -169,7 +182,51 @@ def match(cls, graph: fx.Graph, module: nn.Module) -> "RMSNormFuser | None": # The norm must be the last compute in forward, or it is not a pure norm. if _has_trailing_compute(graph, tail): return None - return cls(zero_centered=zero_centered, source_cls=type(module).__name__) + eps_attr, eps = cls._eps_source(graph, module) + return cls( + zero_centered=zero_centered, + source_cls=type(module).__name__, + eps_attr=eps_attr, + eps=eps, + ) + + @classmethod + def _eps_source( + cls, graph: fx.Graph, module: nn.Module + ) -> tuple[str | None, float | None]: + """Where `fuse` should read `eps` from, resolved once per class.""" + eps = cls._eps_from_graph(graph) + if eps is None: + return None, None + # Whatever supplied the constant must still equal it. + candidates = { + name: value + for name, value in vars(module).items() + if isinstance(value, float) and value == eps + } + # Use unique markers and retrace to verify exactly which attribute is eps. + markers = {float(-index - 1): name for index, name in enumerate(candidates)} + marked = None + if markers: + try: + for marker, name in markers.items(): + setattr(module, name, marker) + if (remarked := trace(module)) is not None: + marked = cls._eps_from_graph(remarked) + finally: + for name, value in candidates.items(): + setattr(module, name, value) + if (name := markers.get(marked)) is not None: + return name, None + logger.debug_once( + "%s does not hold its eps (%s) in an attribute. Every instance in this " + "model will use the value traced from this instance. If this is not " + "desired, consider storing and reading eps using attribute of %s.", + type(module).__name__, + eps, + type(module).__name__, + ) + return None, eps @staticmethod def _eps_from_graph(graph: fx.Graph) -> float | None: @@ -180,7 +237,7 @@ def _eps_from_graph(graph: fx.Graph) -> float | None: if fused is not None and fused.args and peel(fused.args[0]) is x: args, kwargs = fused.args, fused.kwargs eps = args[3] if len(args) > 3 else kwargs.get("eps") - return eps if isinstance(eps, (int, float)) else None + return float(eps) if isinstance(eps, (int, float)) else None for node in graph.nodes: if is_op(node, "rsqrt") and (eps := _variance_eps(node, x)) is not None: return eps @@ -196,10 +253,9 @@ def fuse( weight = getattr(module, "weight", None) has_weight = weight is not None hidden_size = weight.size(0) if has_weight else 0 - graph = trace(module) - eps = self._eps_from_graph(graph) if graph is not None else None - if eps is None: - # If eps not in graph, match torch behaviour. + eps = getattr(module, self.eps_attr, None) if self.eps_attr else self.eps + if not isinstance(eps, (int, float)): + # If eps was not detected, match torch behaviour. dtype = weight.dtype if has_weight else vllm_config.model_config.dtype eps = torch.finfo(dtype).eps if self.zero_centered: From 1fe3a1571ac67581478a11743e55a306de1d136f Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Thu, 20 Aug 2026 16:40:11 +0100 Subject: [PATCH 215/839] Reduce `AutoWeightsLoader` kwargs (#53106) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- .../models/inkling/test_moe_weight_layout.py | 2 +- tests/models/test_dspark_mla.py | 8 +- tests/models/test_utils.py | 104 +++++------------- tests/quantization/test_modelopt.py | 4 +- vllm/lora/worker_manager.py | 9 +- vllm/model_executor/model_loader/utils.py | 2 +- vllm/model_executor/models/adapters.py | 13 ++- vllm/model_executor/models/aimv2.py | 15 +-- vllm/model_executor/models/arcee.py | 7 +- vllm/model_executor/models/aria.py | 22 ++-- vllm/model_executor/models/bagel.py | 12 +- vllm/model_executor/models/bert.py | 13 ++- vllm/model_executor/models/blip.py | 8 +- vllm/model_executor/models/clip.py | 13 ++- vllm/model_executor/models/cohere2_moe.py | 5 +- vllm/model_executor/models/cohere_asr.py | 19 ++-- vllm/model_executor/models/cohere_eagle.py | 14 +-- vllm/model_executor/models/colbert.py | 9 +- vllm/model_executor/models/colqwen3_5.py | 6 +- vllm/model_executor/models/commandr.py | 5 +- vllm/model_executor/models/deepseek_eagle.py | 5 +- vllm/model_executor/models/deepseek_eagle3.py | 15 ++- vllm/model_executor/models/ernie45_moe.py | 9 +- vllm/model_executor/models/exaone4_5.py | 16 ++- vllm/model_executor/models/exaone_moe.py | 4 +- vllm/model_executor/models/fireredasr2.py | 9 +- vllm/model_executor/models/fireredlid.py | 22 ++-- vllm/model_executor/models/funaudiochat.py | 6 +- vllm/model_executor/models/gemma3n.py | 17 +-- vllm/model_executor/models/gemma4.py | 23 ++-- vllm/model_executor/models/glm4.py | 27 +++-- vllm/model_executor/models/glmasr.py | 9 +- vllm/model_executor/models/gpt2.py | 13 ++- vllm/model_executor/models/gpt_j.py | 8 +- vllm/model_executor/models/gpt_neox.py | 11 +- .../models/idefics2_vision_model.py | 16 +-- vllm/model_executor/models/interfaces.py | 5 +- vllm/model_executor/models/interns1_pro.py | 17 +-- .../model_executor/models/interns2_preview.py | 14 ++- vllm/model_executor/models/internvl.py | 45 +++++--- vllm/model_executor/models/jina.py | 6 +- vllm/model_executor/models/keye.py | 5 +- vllm/model_executor/models/kimi_audio.py | 23 ++-- vllm/model_executor/models/lfm2_siglip2.py | 9 +- vllm/model_executor/models/llama.py | 3 +- vllm/model_executor/models/llama4_eagle.py | 6 +- vllm/model_executor/models/llama_eagle.py | 5 +- vllm/model_executor/models/llama_eagle3.py | 19 ++-- .../models/longcat_flash_ngram.py | 12 +- vllm/model_executor/models/mimo.py | 10 +- vllm/model_executor/models/mimo_v2_omni.py | 3 +- vllm/model_executor/models/minicpmo.py | 14 ++- vllm/model_executor/models/minicpmv.py | 20 ++-- vllm/model_executor/models/minicpmv4_6.py | 3 +- vllm/model_executor/models/minimax_m2.py | 20 ++-- vllm/model_executor/models/modernbert.py | 6 +- vllm/model_executor/models/moss_audio.py | 6 +- vllm/model_executor/models/nemotron_h.py | 4 +- vllm/model_executor/models/nemotron_vl.py | 13 ++- vllm/model_executor/models/paddleocr_vl.py | 24 ++-- vllm/model_executor/models/phi4mm.py | 3 +- .../models/qwen2_5_omni_thinker.py | 4 +- vllm/model_executor/models/qwen3_5.py | 17 ++- vllm/model_executor/models/qwen3_asr.py | 7 +- .../models/qwen3_asr_forced_aligner.py | 7 +- vllm/model_executor/models/qwen3_dflash.py | 19 ++-- vllm/model_executor/models/qwen3_dspark.py | 22 ++-- vllm/model_executor/models/qwen3_eagle3.py | 19 ++-- vllm/model_executor/models/qwen3_next.py | 7 +- .../models/qwen3_omni_moe_thinker.py | 7 +- vllm/model_executor/models/roberta.py | 8 +- vllm/model_executor/models/siglip.py | 22 ++-- vllm/model_executor/models/skyworkr1v.py | 44 +++++--- .../models/transformers/base.py | 6 - .../models/transformers/legacy.py | 32 +++--- vllm/model_executor/models/utils.py | 72 ++++++------ vllm/model_executor/models/whisper.py | 3 +- vllm/models/deepseek_v4/amd/model.py | 7 +- vllm/models/deepseek_v4/nvidia/model.py | 3 +- vllm/models/deepseek_v4/xpu/model.py | 3 +- vllm/models/dots3_note/nvidia/multimodal.py | 17 ++- vllm/models/inkling/amd/model.py | 5 +- vllm/models/inkling/nvidia/model.py | 5 +- vllm/models/kimi_k3/nvidia/dspark_mla.py | 16 +-- 84 files changed, 578 insertions(+), 569 deletions(-) diff --git a/tests/models/inkling/test_moe_weight_layout.py b/tests/models/inkling/test_moe_weight_layout.py index e15a573b8459..d9f1ceac47db 100644 --- a/tests/models/inkling/test_moe_weight_layout.py +++ b/tests/models/inkling/test_moe_weight_layout.py @@ -104,7 +104,7 @@ def test_inkling_mapper_maps_modelopt_exclusions() -> None: ) quant_config.apply_vllm_mapper( - _TmlForCausalLMBase.hf_to_vllm_mapper.get_unstacked_mapper() + _TmlForCausalLMBase.hf_to_vllm_mapper.get_rename_mapper() ) assert quant_config.is_layer_excluded("model.layers.2.mlp.experts") diff --git a/tests/models/test_dspark_mla.py b/tests/models/test_dspark_mla.py index 3188f0fadd3f..01ee0b0533e5 100644 --- a/tests/models/test_dspark_mla.py +++ b/tests/models/test_dspark_mla.py @@ -55,11 +55,9 @@ def test_dspark_mla_checkpoint_weight_mapping(checkpoint_name, runtime_name, sha def test_dspark_mla_shares_frozen_target_weights_and_skips_training_head(): assert not K3DSparkForCausalLM.has_own_embed_tokens assert not K3DSparkForCausalLM.has_own_lm_head - assert set(K3DSparkForCausalLM.checkpoint_skip_substrs) == { - "confidence_head", - "embed_tokens", - "lm_head", - } + mapper = K3DSparkForCausalLM.hf_to_vllm_mapper + for name in ("confidence_head.weight", "embed_tokens.weight", "lm_head.weight"): + assert mapper._map_name(name) is None @pytest.mark.cpu_test diff --git a/tests/models/test_utils.py b/tests/models/test_utils.py index 2b32a63ac6e1..02903504c684 100644 --- a/tests/models/test_utils.py +++ b/tests/models/test_utils.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest +import regex as re import torch from vllm.model_executor.layers.vocab_parallel_embedding import ( @@ -10,6 +11,7 @@ ) from vllm.model_executor.models.utils import ( AutoWeightsLoader, + WeightsMapper, _merge_multimodal_embeddings, ) from vllm.platforms import current_platform @@ -93,82 +95,6 @@ def weight_generator(): assert new_mod.nested_mod.bn.num_batches_tracked.item() == 1 -@pytest.mark.cpu_test -def test_module_skip_prefix(): - """Ensure the auto weight loader can skip prefix.""" - mod = ModuleWithNestedBatchNorm() - # Run some data through the module with batchnorm - mod(torch.Tensor([[1, 2], [3, 4]])) - - # Try to load the weights to a new instance - def weight_generator(): - # weights needed to be filtered out - redundant_weights = { - "prefix.bn.weight": torch.Tensor([1, 2]), - "prefix.bn.bias": torch.Tensor([3, 4]), - } - yield from (mod.state_dict() | redundant_weights).items() - - new_mod = ModuleWithNestedBatchNorm() - - assert not torch.all( - new_mod.nested_mod.bn.running_mean == mod.nested_mod.bn.running_mean - ) - assert not torch.all( - new_mod.nested_mod.bn.running_var == mod.nested_mod.bn.running_var - ) - assert new_mod.nested_mod.bn.num_batches_tracked.item() == 0 - - loader = AutoWeightsLoader(new_mod, skip_prefixes=["prefix."]) - loader.load_weights(weight_generator()) - - # Ensure the stats are updated - assert torch.all( - new_mod.nested_mod.bn.running_mean == mod.nested_mod.bn.running_mean - ) - assert torch.all(new_mod.nested_mod.bn.running_var == mod.nested_mod.bn.running_var) - assert new_mod.nested_mod.bn.num_batches_tracked.item() == 1 - - -@pytest.mark.cpu_test -def test_module_skip_substr(): - """Ensure the auto weight loader can skip prefix.""" - mod = ModuleWithNestedBatchNorm() - # Run some data through the module with batchnorm - mod(torch.Tensor([[1, 2], [3, 4]])) - - # Try to load the weights to a new instance - def weight_generator(): - # weights needed to be filtered out - redundant_weights = { - "nested_mod.0.substr.weight": torch.Tensor([1, 2]), - "nested_mod.0.substr.bias": torch.Tensor([3, 4]), - "nested_mod.substr.weight": torch.Tensor([1, 2]), - "nested_mod.substr.bias": torch.Tensor([3, 4]), - } - yield from (mod.state_dict() | redundant_weights).items() - - new_mod = ModuleWithNestedBatchNorm() - - assert not torch.all( - new_mod.nested_mod.bn.running_mean == mod.nested_mod.bn.running_mean - ) - assert not torch.all( - new_mod.nested_mod.bn.running_var == mod.nested_mod.bn.running_var - ) - assert new_mod.nested_mod.bn.num_batches_tracked.item() == 0 - - loader = AutoWeightsLoader(new_mod, skip_substrs=["substr."]) - loader.load_weights(weight_generator()) - - # Ensure the stats are updated - assert torch.all( - new_mod.nested_mod.bn.running_mean == mod.nested_mod.bn.running_mean - ) - assert torch.all(new_mod.nested_mod.bn.running_var == mod.nested_mod.bn.running_var) - assert new_mod.nested_mod.bn.num_batches_tracked.item() == 1 - - VOCAB_SIZE = 16 HIDDEN_SIZE = 2 @@ -265,3 +191,29 @@ def test_merge_multimodal_embeddings_no_sync(): _merge_multimodal_embeddings( inputs_embeds, multimodal_embeddings, is_multimodal ) + + +@pytest.mark.cpu_test +def test_get_rename_mapper_keeps_only_renames(): + """`None` means "do not load", which is meaningless to the consumers of + this mapper (LoRA name parsing, quantization config layer lists), and + applying it would silently shrink their lists.""" + mapper = WeightsMapper( + orig_to_new_regex={re.compile(r"^drop_regex\."): None}, + orig_to_new_substr={"drop_substr": None, "keep_substr": "kept"}, + orig_to_new_stacked={".q_proj": (".qkv_proj", "q")}, + orig_to_new_prefix={"drop_prefix.": None, "keep_prefix.": "kept."}, + orig_to_new_suffix={".drop_suffix": None}, + ) + renames = mapper.get_rename_mapper() + + assert renames.orig_to_new_regex == {} + assert renames.orig_to_new_substr == {"keep_substr": "kept"} + assert renames.orig_to_new_stacked == {} + assert renames.orig_to_new_prefix == {"keep_prefix.": "kept."} + assert renames.orig_to_new_suffix == {} + + # Names the full mapper drops now survive unchanged. + for name in ("drop_regex.w", "drop_substr.w", "drop_prefix.w", "w.drop_suffix"): + assert mapper._map_name(name) is None + assert renames._map_name(name) == name diff --git a/tests/quantization/test_modelopt.py b/tests/quantization/test_modelopt.py index 4124085c23e3..41cedca9d3a6 100644 --- a/tests/quantization/test_modelopt.py +++ b/tests/quantization/test_modelopt.py @@ -214,9 +214,9 @@ def test_modelopt_mixed_precision_composes_gemma4_mappers(): ) config.apply_vllm_mapper( - Gemma4ForConditionalGeneration.hf_to_vllm_mapper.get_unstacked_mapper() + Gemma4ForConditionalGeneration.hf_to_vllm_mapper.get_rename_mapper() ) - config.apply_vllm_mapper(Gemma4ForCausalLM.hf_to_vllm_mapper.get_unstacked_mapper()) + config.apply_vllm_mapper(Gemma4ForCausalLM.hf_to_vllm_mapper.get_rename_mapper()) expected_prefix = "language_model.model.layers.0.moe.experts" assert set(config.quantized_layers) == { diff --git a/vllm/lora/worker_manager.py b/vllm/lora/worker_manager.py index 105a99807257..4c00c6531bc8 100644 --- a/vllm/lora/worker_manager.py +++ b/vllm/lora/worker_manager.py @@ -128,14 +128,13 @@ def _load_adapter(self, lora_request: LoRARequest) -> LoRAModel: # loading weights, throwing an exception if validation fails. peft_helper.validate_legal(self.lora_config) - # For some models like Qwen2VL, we need to use hf_to_vllm_mapper - # to ensure correct loading of lora weights. Drop the QKV/MLP fusion - # substr maps so constituent names (e.g. `q_proj`) survive for the - # LoRA manager to pack, while keeping genuine renames/prefixes. + # For some models like Qwen2VL, we need to use hf_to_vllm_mapper to ensure + # correct loading of lora weights. We only need to know about renames for + # this, so we use get_rename_mapper() to ignore stacking and deletions. model = self._adapter_manager.model hf_to_vllm_mapper = getattr(model, "hf_to_vllm_mapper", None) if hf_to_vllm_mapper is not None: - hf_to_vllm_mapper = hf_to_vllm_mapper.get_unstacked_mapper() + hf_to_vllm_mapper = hf_to_vllm_mapper.get_rename_mapper() # Get model-defined prefixes to skip during LoRA loading. lora_skip_prefixes = getattr(model, "lora_skip_prefixes", None) diff --git a/vllm/model_executor/model_loader/utils.py b/vllm/model_executor/model_loader/utils.py index d28051a54892..8fbd31ab4e02 100644 --- a/vllm/model_executor/model_loader/utils.py +++ b/vllm/model_executor/model_loader/utils.py @@ -281,6 +281,6 @@ def configure_quant_config( # pass mappings by reference to quant_config if hf_to_vllm_mapper is not None: - quant_config.apply_vllm_mapper(hf_to_vllm_mapper.get_unstacked_mapper()) + quant_config.apply_vllm_mapper(hf_to_vllm_mapper.get_rename_mapper()) if packed_mapping is not None: quant_config.packed_modules_mapping = packed_mapping diff --git a/vllm/model_executor/models/adapters.py b/vllm/model_executor/models/adapters.py index 8f269fcabc4b..ea75ff8cc0b1 100644 --- a/vllm/model_executor/models/adapters.py +++ b/vllm/model_executor/models/adapters.py @@ -236,7 +236,8 @@ def _load_pooling_model_weights( def default_load_weights(weights): loader = AutoWeightsLoader(self) - return loader.load_weights(weights) + mapper = getattr(self, "hf_to_vllm_mapper", None) + return loader.load_weights(weights, mapper=mapper) load_weights = getattr(super(), "load_weights", default_load_weights) return load_weights(mapped_weights) @@ -584,10 +585,11 @@ def load_weights_using_from_2_way_softmax( ) loaded_weights.add(score_weight_name) - lm_head_name = "lm_head.weight" + lm_head_name: str | None = "lm_head.weight" if hf_to_vllm_mapper := getattr(model, "hf_to_vllm_mapper", None): lm_head_name = hf_to_vllm_mapper._map_name(lm_head_name) - loaded_weights.discard(lm_head_name) + if lm_head_name is not None: + loaded_weights.discard(lm_head_name) return loaded_weights @@ -649,10 +651,11 @@ def load_weights_no_post_processing(model, weights: Iterable[tuple[str, torch.Te ) loaded_weights.add(score_weight_name) - lm_head_name = "lm_head.weight" + lm_head_name: str | None = "lm_head.weight" if hf_to_vllm_mapper := getattr(model, "hf_to_vllm_mapper", None): lm_head_name = hf_to_vllm_mapper._map_name(lm_head_name) - loaded_weights.discard(lm_head_name) + if lm_head_name is not None: + loaded_weights.discard(lm_head_name) return loaded_weights diff --git a/vllm/model_executor/models/aimv2.py b/vllm/model_executor/models/aimv2.py index eaff05c4f0d6..3461b5713734 100644 --- a/vllm/model_executor/models/aimv2.py +++ b/vllm/model_executor/models/aimv2.py @@ -219,6 +219,11 @@ def __init__( require_post_norm=require_post_norm, prefix=f"{prefix}.trunk", ) + # post_trunk_norm is optional (absent for clip-skip backbones). + if self.trunk.post_trunk_norm is None: + self.hf_to_vllm_mapper = self.hf_to_vllm_mapper | WeightsMapper( + orig_to_new_prefix={"trunk.post_trunk_norm.": None} + ) def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: x = self.preprocessor(pixel_values) @@ -227,13 +232,5 @@ def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: return x def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - # post_trunk_norm is optional (absent for clip-skip backbones). - skip_prefixes=( - ["trunk.post_trunk_norm."] - if self.trunk.post_trunk_norm is None - else None - ), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/arcee.py b/vllm/model_executor/models/arcee.py index 81cf3c5e0f64..32ad19a8e73b 100644 --- a/vllm/model_executor/models/arcee.py +++ b/vllm/model_executor/models/arcee.py @@ -285,7 +285,10 @@ class ArceeForCausalLM( ".q_proj": (".qkv_proj", "q"), ".k_proj": (".qkv_proj", "k"), ".v_proj": (".qkv_proj", "v"), - } + }, + orig_to_new_substr={ + "gate_proj": None, + }, ) # Map fused module names to their submodule components # (for quantization and LoRA) @@ -356,7 +359,7 @@ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load weights into the model (delegates to inner model and handles tied embeddings).""" - loader = AutoWeightsLoader(self, skip_substrs=["gate_proj"]) + loader = AutoWeightsLoader(self) # AutoWeightLoader handles weight name remapping, including fusing # separate q_proj, k_proj, v_proj into qkv_proj return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/aria.py b/vllm/model_executor/models/aria.py index f35a32fab57c..57eb62d3e26a 100644 --- a/vllm/model_executor/models/aria.py +++ b/vllm/model_executor/models/aria.py @@ -13,18 +13,13 @@ from vllm.config.multimodal import BaseDummyOptions, ImageDummyOptions from vllm.inputs import MultiModalDataDict from vllm.model_executor.layers.activation import get_act_fn -from vllm.model_executor.layers.fused_moe import ( - FusedMoEFactory, -) +from vllm.model_executor.layers.fused_moe import FusedMoEFactory from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.multimodal.inputs import ( - MultiModalFieldConfig, - MultiModalKwargsItems, -) +from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems from vllm.multimodal.parse import MultiModalDataItems from vllm.multimodal.processing import ( BaseDummyInputsBuilder, @@ -42,11 +37,7 @@ ) from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsQuant from .llama import LlamaDecoderLayer, LlamaMLP, LlamaModel -from .utils import ( - AutoWeightsLoader, - WeightsMapper, - maybe_prefix, -) +from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix class AriaImagePixelInputs(TensorSchema): @@ -88,16 +79,17 @@ def __init__( self.post_layernorm = nn.Identity() hf_to_vllm_mapper = WeightsMapper( + # NOTE: post_layernorm is not used in Aria. + orig_to_new_substr={"post_layernorm": None}, orig_to_new_stacked={ ".q_proj": (".qkv_proj", "q"), ".k_proj": (".qkv_proj", "k"), ".v_proj": (".qkv_proj", "v"), - } + }, ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - # NOTE: post_layernorm is not used in Aria. - loader = AutoWeightsLoader(self, skip_substrs=["post_layernorm"]) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/bagel.py b/vllm/model_executor/models/bagel.py index e592dd3728c9..28c8b64cd170 100644 --- a/vllm/model_executor/models/bagel.py +++ b/vllm/model_executor/models/bagel.py @@ -338,14 +338,9 @@ class BagelForConditionalGeneration( The image generation part is not supported in vLLM. """ - # Weight mapping from HF to vLLM + # pos_embed is handled by the PositionEmbedding module hf_to_vllm_mapper = WeightsMapper( - orig_to_new_prefix={ - "language_model.": "language_model.", - "vit_model.": "vit_model.", - "connector.": "connector.", - "vit_pos_embed.": "vit_pos_embed.", - } + orig_to_new_prefix={"vit_pos_embed.pos_embed": None} ) @classmethod @@ -580,6 +575,5 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: filtered_weights.append((name, tensor)) - # Skip vit_pos_embed.pos_embed as it's handled by PositionEmbedding module - loader = AutoWeightsLoader(self, skip_prefixes=["vit_pos_embed.pos_embed"]) + loader = AutoWeightsLoader(self) return loader.load_weights(filtered_weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/bert.py b/vllm/model_executor/models/bert.py index 673011305459..cdc7870f3ec5 100644 --- a/vllm/model_executor/models/bert.py +++ b/vllm/model_executor/models/bert.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Iterable, Set +from dataclasses import replace import torch from torch import nn @@ -382,6 +383,7 @@ class BertModel(nn.Module, SupportsQuant): ".self.key": (".self.qkv_proj", "k"), ".self.value": (".self.qkv_proj", "v"), }, + orig_to_new_prefix={"pooler.": None}, ) def __init__( @@ -416,13 +418,16 @@ def forward( return self.encoder(hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self, skip_prefixes=["pooler."]) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) class BertPoolingModel(BertModel): is_pooling_model = True + # Unlike `BertModel`, this model has a pooler to load weights into. + hf_to_vllm_mapper = replace(BertModel.hf_to_vllm_mapper, orig_to_new_prefix={}) + def __init__( self, *, @@ -488,11 +493,13 @@ def forward( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): weights_list = list(weights) + orig_to_new_prefix: dict[str, str | None] = {"lm_head.": None} has_model_prefix = any(name.startswith("model.") for name, _ in weights_list) if not has_model_prefix: - mapper = WeightsMapper(orig_to_new_prefix={"": "model."}) + orig_to_new_prefix[""] = "model." + mapper = WeightsMapper(orig_to_new_prefix=orig_to_new_prefix) - loader = AutoWeightsLoader(self, skip_prefixes=["lm_head."]) + loader = AutoWeightsLoader(self) return loader.load_weights(weights_list, mapper=mapper) def _build_model(self, vllm_config: VllmConfig, prefix: str = "") -> BertModel: diff --git a/vllm/model_executor/models/blip.py b/vllm/model_executor/models/blip.py index aecb99716137..aefd17003d2f 100644 --- a/vllm/model_executor/models/blip.py +++ b/vllm/model_executor/models/blip.py @@ -313,6 +313,9 @@ def __init__( ) else: self.post_layernorm = None + self.hf_to_vllm_mapper = self.hf_to_vllm_mapper | WeightsMapper( + orig_to_new_prefix={"post_layernorm.": None} + ) def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: hidden_states = self.embeddings(pixel_values) @@ -324,10 +327,7 @@ def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: return self.post_layernorm(hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - skip_prefixes: list[str] = [] - if self.post_layernorm is None: - skip_prefixes.append("post_layernorm.") - loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes) + loader = AutoWeightsLoader(self) # omit layers when num_hidden_layers_override is set def _filter(ws): diff --git a/vllm/model_executor/models/clip.py b/vllm/model_executor/models/clip.py index c37e2e62146e..2e0190350538 100644 --- a/vllm/model_executor/models/clip.py +++ b/vllm/model_executor/models/clip.py @@ -669,6 +669,9 @@ def __init__( self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) else: self.post_layernorm = None + self.hf_to_vllm_mapper = self.hf_to_vllm_mapper | WeightsMapper( + orig_to_new_prefix={"post_layernorm.": None} + ) @property def dtype(self): @@ -707,10 +710,7 @@ def forward( return encoder_outputs def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - skip_prefixes: list[str] = [] - if self.post_layernorm is None: - skip_prefixes.append("post_layernorm.") - loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes) + loader = AutoWeightsLoader(self) # Drop layers beyond num_hidden_layers_override. def _filter(ws): @@ -775,6 +775,8 @@ def device(self): class CLIPEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant): is_pooling_model = True + hf_to_vllm_mapper = WeightsMapper(orig_to_new_substr={".position_ids": None}) + packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]} @classmethod @@ -981,8 +983,7 @@ def forward( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): loader = AutoWeightsLoader( self, - skip_substrs=[".position_ids"], ignore_unexpected_prefixes=["logit_scale."], ) - return loader.load_weights(weights) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/cohere2_moe.py b/vllm/model_executor/models/cohere2_moe.py index 9af1a3439b90..d657b78a5da1 100644 --- a/vllm/model_executor/models/cohere2_moe.py +++ b/vllm/model_executor/models/cohere2_moe.py @@ -474,7 +474,8 @@ class Cohere2MoeForCausalLM(nn.Module, SupportsPP, SupportsQuant): ".mlp.up_proj": (".mlp.gate_up_proj", 1), ".shared_experts.gate_proj": (".shared_experts.gate_up_proj", 0), ".shared_experts.up_proj": (".shared_experts.gate_up_proj", 1), - } + }, + orig_to_new_prefix={"lm_head.": None}, ) packed_modules_mapping = { "qkv_proj": [ @@ -529,5 +530,5 @@ def compute_logits( return self.logits_processor(self.model.embed_tokens, hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self, skip_prefixes=["lm_head."]) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/cohere_asr.py b/vllm/model_executor/models/cohere_asr.py index 0f9bd86dd09a..7ffdd01cd602 100644 --- a/vllm/model_executor/models/cohere_asr.py +++ b/vllm/model_executor/models/cohere_asr.py @@ -2014,7 +2014,15 @@ class CohereAsrForConditionalGeneration( } hf_to_vllm_mapper = WeightsMapper( - orig_to_new_substr={".fc1.": ".mlp.fc1.", ".fc2.": ".mlp.fc2."} + orig_to_new_substr={ + ".fc1.": ".mlp.fc1.", + ".fc2.": ".mlp.fc2.", + "model.conv.batch_norm.num_batches_tracked": None, + }, + orig_to_new_prefix={ + "model.preprocessor.featurizer.fb": None, + "model.preprocessor.featurizer.window": None, + }, ) supports_transcription_only = True @@ -2273,14 +2281,7 @@ def transform(inputs): return name, loaded_weight - loader = AutoWeightsLoader( - self, - skip_prefixes=[ - "model.preprocessor.featurizer.fb", - "model.preprocessor.featurizer.window", - ], - skip_substrs=["model.conv.batch_norm.num_batches_tracked"], - ) + loader = AutoWeightsLoader(self) return loader.load_weights( map(transform, weights), mapper=self.hf_to_vllm_mapper diff --git a/vllm/model_executor/models/cohere_eagle.py b/vllm/model_executor/models/cohere_eagle.py index 519d1164b0dd..a8b540bbcf5c 100644 --- a/vllm/model_executor/models/cohere_eagle.py +++ b/vllm/model_executor/models/cohere_eagle.py @@ -22,6 +22,7 @@ from .utils import ( AutoWeightsLoader, + WeightsMapper, get_draft_quant_config, maybe_prefix, process_eagle_weight, @@ -145,6 +146,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): # use tied embeddings so these weights are absent from the draft file. self.has_own_embed_tokens = False self.has_own_lm_head = False + if self.config.tie_word_embeddings: + self.hf_to_vllm_mapper = self.hf_to_vllm_mapper | WeightsMapper( + orig_to_new_prefix={"model.embed_tokens.": None} + ) target_layer_num = vllm_config.model_config.get_num_layers( vllm_config.parallel_config ) @@ -181,14 +186,7 @@ def _track_and_forward(inputs): process_eagle_weight(self, name) return name, weight - loader = AutoWeightsLoader( - self, - skip_prefixes=( - ["lm_head.", "model.embed_tokens."] - if self.config.tie_word_embeddings - else None - ), - ) + loader = AutoWeightsLoader(self) loaded_weight_names = loader.load_weights( map(_track_and_forward, weights), mapper=self.hf_to_vllm_mapper diff --git a/vllm/model_executor/models/colbert.py b/vllm/model_executor/models/colbert.py index cc5483fd7b3f..124cfb5fdbec 100644 --- a/vllm/model_executor/models/colbert.py +++ b/vllm/model_executor/models/colbert.py @@ -358,11 +358,12 @@ def forward( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): other_weights, colbert_loaded = self._load_colbert_weights(weights) - mapper = WeightsMapper(orig_to_new_prefix={"roberta.": "model."}) - - # Skip HF pooler weights (model.pooler.*) as they not used in ColBERT - loader = AutoWeightsLoader(self, skip_prefixes=["model.pooler."]) + # HF pooler weights are not used in ColBERT + mapper = WeightsMapper( + orig_to_new_prefix={"roberta.": "model.", "model.pooler.": None} + ) + loader = AutoWeightsLoader(self) loaded = loader.load_weights(other_weights, mapper=mapper) return loaded | colbert_loaded diff --git a/vllm/model_executor/models/colqwen3_5.py b/vllm/model_executor/models/colqwen3_5.py index 760a54f485f7..68ec62383533 100644 --- a/vllm/model_executor/models/colqwen3_5.py +++ b/vllm/model_executor/models/colqwen3_5.py @@ -140,6 +140,7 @@ class ColQwen3_5Model( hf_to_vllm_mapper = WeightsMapper( orig_to_new_prefix={ "language_model.": "language_model.model.", + "mtp.": None, } ) @@ -238,10 +239,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: else: model_weights.append((name, weight)) - loader = AutoWeightsLoader( - self, - skip_prefixes=["mtp."], - ) + loader = AutoWeightsLoader(self) loaded = loader.load_weights(model_weights, mapper=self.hf_to_vllm_mapper) for name, weight in proj_weights: diff --git a/vllm/model_executor/models/commandr.py b/vllm/model_executor/models/commandr.py index ca2543055db4..87b832f91e01 100644 --- a/vllm/model_executor/models/commandr.py +++ b/vllm/model_executor/models/commandr.py @@ -415,6 +415,7 @@ class CohereForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant): # (e.g. "*.weight_quantizer._double_scale"); drop them before loading. # See #41925. orig_to_new_substr={"_quantizer.": None}, + orig_to_new_prefix={"lm_head": None}, ) packed_modules_mapping = { "qkv_proj": ["q_proj", "k_proj", "v_proj"], @@ -472,7 +473,5 @@ def compute_logits( return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, skip_prefixes=["lm_head", "rotary_emb.inv_freq"] - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/deepseek_eagle.py b/vllm/model_executor/models/deepseek_eagle.py index 91e158fa5a20..5c9c78e7846d 100644 --- a/vllm/model_executor/models/deepseek_eagle.py +++ b/vllm/model_executor/models/deepseek_eagle.py @@ -253,8 +253,5 @@ def transform(inputs): process_eagle_weight(self, name) return name, loaded_weight - loader = AutoWeightsLoader( - self, - skip_prefixes=None, - ) + loader = AutoWeightsLoader(self) loader.load_weights(map(transform, weights)) diff --git a/vllm/model_executor/models/deepseek_eagle3.py b/vllm/model_executor/models/deepseek_eagle3.py index 52f476adbce5..63673093439d 100644 --- a/vllm/model_executor/models/deepseek_eagle3.py +++ b/vllm/model_executor/models/deepseek_eagle3.py @@ -384,18 +384,17 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): model_weights[name] = loaded_weight process_eagle_weight(self, name) - skip_substrs = [] + orig_to_new_substr: dict[str, str | None] = {} if not includes_draft_id_mapping: - skip_substrs.append("draft_id_to_target_id") + orig_to_new_substr["draft_id_to_target_id"] = None if not includes_embed_tokens: - skip_substrs.append("embed_tokens") + orig_to_new_substr["embed_tokens"] = None - loader = AutoWeightsLoader( - self, - skip_prefixes=None, - skip_substrs=skip_substrs, + loader = AutoWeightsLoader(self) + loader.load_weights( + model_weights.items(), + mapper=WeightsMapper(orig_to_new_substr=orig_to_new_substr), ) - loader.load_weights(model_weights.items()) # Aliases for compatibility diff --git a/vllm/model_executor/models/ernie45_moe.py b/vllm/model_executor/models/ernie45_moe.py index c7f26b6dd4a2..746d12c596f1 100644 --- a/vllm/model_executor/models/ernie45_moe.py +++ b/vllm/model_executor/models/ernie45_moe.py @@ -403,7 +403,8 @@ class Ernie4_5_MoeModel(nn.Module): ".mlp.up_proj": (".mlp.gate_up_proj", 1), ".shared_experts.gate_proj": (".shared_experts.gate_up_proj", 0), ".shared_experts.up_proj": (".shared_experts.gate_up_proj", 1), - } + }, + orig_to_new_substr={"mtp": None}, ) def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): @@ -497,11 +498,7 @@ def _preprocess( yield name, loaded_weight def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_substrs=["mtp"], - ignore_unexpected_suffixes=[".bias", "_bias"], - ) + loader = AutoWeightsLoader(self) return loader.load_weights( self._preprocess(weights), mapper=self.hf_to_vllm_mapper ) diff --git a/vllm/model_executor/models/exaone4_5.py b/vllm/model_executor/models/exaone4_5.py index 58ad3d4c61a3..676faa9d4f3a 100644 --- a/vllm/model_executor/models/exaone4_5.py +++ b/vllm/model_executor/models/exaone4_5.py @@ -52,7 +52,12 @@ from .qwen2_vl import Qwen2VLDummyInputsBuilder as Exaone4_5_DummyInputsBuilder from .qwen2_vl import Qwen2VLMultiModalProcessor as Exaone4_5_MultiModalProcessor -from .utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix +from .utils import ( + AutoWeightsLoader, + WeightsMapper, + init_vllm_registered_model, + maybe_prefix, +) logger = init_logger(__name__) @@ -317,6 +322,10 @@ def get_hf_processor(self, **kwargs: object) -> Exaone4_5_Processor: dummy_inputs=Exaone4_5_DummyInputsBuilder, ) class Exaone4_5_ForConditionalGeneration(Qwen2_5_VLForConditionalGeneration): + hf_to_vllm_mapper = Qwen2_5_VLForConditionalGeneration.hf_to_vllm_mapper | ( + WeightsMapper(orig_to_new_prefix={"mtp.": None}) + ) + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): nn.Module.__init__(self) @@ -353,10 +362,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["mtp."]), - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) @classmethod diff --git a/vllm/model_executor/models/exaone_moe.py b/vllm/model_executor/models/exaone_moe.py index 229a741dcaaf..6fbeb03af8db 100644 --- a/vllm/model_executor/models/exaone_moe.py +++ b/vllm/model_executor/models/exaone_moe.py @@ -471,7 +471,8 @@ class ExaoneMoeForCausalLM(nn.Module, SupportsLoRA, SupportsPP): ".mlp.up_proj": (".mlp.gate_up_proj", 1), ".shared_experts.gate_proj": (".shared_experts.gate_up_proj", 0), ".shared_experts.up_proj": (".shared_experts.gate_up_proj", 1), - } + }, + orig_to_new_prefix={"mtp.": None}, ) packed_modules_mapping = { "qkv_proj": [ @@ -560,7 +561,6 @@ def compute_logits( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader( self, - skip_prefixes=["mtp."], # Skip loading extra parameters for GPTQ/modelopt models. ignore_unexpected_suffixes=[ ".bias", diff --git a/vllm/model_executor/models/fireredasr2.py b/vllm/model_executor/models/fireredasr2.py index 231417d83367..7dedda3c3c41 100644 --- a/vllm/model_executor/models/fireredasr2.py +++ b/vllm/model_executor/models/fireredasr2.py @@ -332,7 +332,10 @@ class FireRedASR2ForConditionalGeneration( "net.0": "pre_layer_norm", "net.1": "linear_expand", "net.4": "linear_project", - } + }, + orig_to_new_prefix={ + "model.encoder.audio_encoder.positional_encoding.pe": None, + }, ) supports_transcription_only = True @@ -477,8 +480,6 @@ def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, skip_prefixes=["model.encoder.audio_encoder.positional_encoding.pe"] - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/fireredlid.py b/vllm/model_executor/models/fireredlid.py index 09faf75da1ff..9f1786327b24 100644 --- a/vllm/model_executor/models/fireredlid.py +++ b/vllm/model_executor/models/fireredlid.py @@ -589,7 +589,15 @@ class FireRedLIDForConditionalGeneration( "net.0": "pre_layer_norm", "net.1": "linear_expand", "net.4": "linear_project", - } + }, + orig_to_new_prefix={ + # Position encoding buffers are rebuilt at init, and the output + # projection is tied to the embedding. + "model.encoder.positional_encoding.pe": None, + "model.decoder.positional_encoding.pe": None, + "model.decoder.tgt_word_prj.weight": None, + "proj_out.": None, + }, ) def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): @@ -778,15 +786,5 @@ def post_process_output(cls, text: str) -> str: return text.strip() def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=[ - # Position encoding buffers are rebuilt at init - "model.encoder.positional_encoding.pe", - "model.decoder.positional_encoding.pe", - # Tied output projection (shared with embedding) - "model.decoder.tgt_word_prj.weight", - "proj_out.", - ], - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/funaudiochat.py b/vllm/model_executor/models/funaudiochat.py index b336b9801e2b..10563f7efc46 100644 --- a/vllm/model_executor/models/funaudiochat.py +++ b/vllm/model_executor/models/funaudiochat.py @@ -760,6 +760,8 @@ def get_replacement_funaudiochat(item_idx: int): dummy_inputs=FunAudioChatDummyInputsBuilder, ) class FunAudioChatForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP): + hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"audio_invert_tower.": None}) + @classmethod def get_placeholder_str(cls, modality: str, i: int) -> str | None: if modality.startswith("audio"): @@ -969,5 +971,5 @@ def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: return self.language_model.compute_logits(hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self, skip_prefixes=["audio_invert_tower."]) - return loader.load_weights(weights) + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/gemma3n.py b/vllm/model_executor/models/gemma3n.py index 4b06fd418f3e..6e73ad8baa64 100644 --- a/vllm/model_executor/models/gemma3n.py +++ b/vllm/model_executor/models/gemma3n.py @@ -1056,6 +1056,14 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: class Gemma3nForCausalLM(nn.Module): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_substr={ + "embed_audio.": None, + "embed_vision.": None, + "audio_tower.": None, + "vision_tower.": None, + } + ) packed_modules_mapping = { "qkv_proj": [ "q_proj", @@ -1112,10 +1120,5 @@ def compute_logits( return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_substrs=( - ["embed_audio.", "embed_vision.", "audio_tower.", "vision_tower."] - ), - ) - return loader.load_weights(weights) + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/gemma4.py b/vllm/model_executor/models/gemma4.py index e8bb72ac3703..1e3b25d88ead 100644 --- a/vllm/model_executor/models/gemma4.py +++ b/vllm/model_executor/models/gemma4.py @@ -1693,13 +1693,16 @@ def _weight_iterator(): yield name, weight - # Skip multimodal weights — handled by the multimodal wrapper. - skip = [ - "audio_tower.", - "vision_tower.", - "embed_audio.", - "embed_vision.", - ] - - loader = AutoWeightsLoader(self, skip_substrs=skip) - return loader.load_weights(_weight_iterator()) + # Drop multimodal weights, which are handled by the multimodal wrapper. + # `_weight_iterator` already applies this model's renames by hand, so + # `hf_to_vllm_mapper` is deliberately not passed here. + mapper = WeightsMapper( + orig_to_new_substr={ + "audio_tower.": None, + "vision_tower.": None, + "embed_audio.": None, + "embed_vision.": None, + } + ) + loader = AutoWeightsLoader(self) + return loader.load_weights(_weight_iterator(), mapper=mapper) diff --git a/vllm/model_executor/models/glm4.py b/vllm/model_executor/models/glm4.py index 5d95175ab03e..f3f65bf512fa 100644 --- a/vllm/model_executor/models/glm4.py +++ b/vllm/model_executor/models/glm4.py @@ -45,11 +45,7 @@ from .interfaces import SupportsLoRA, SupportsPP from .llama import LlamaMLP as Glm4MLP from .llama import LlamaModel -from .utils import ( - AutoWeightsLoader, - PPMissingLayer, - maybe_prefix, -) +from .utils import AutoWeightsLoader, PPMissingLayer, WeightsMapper, maybe_prefix class Glm4Attention(nn.Module): @@ -269,6 +265,16 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.model.make_empty_intermediate_tensors ) + # Drop the speculative (MTP) layers, which are loaded by the draft + # model instead. They are appended after the main decoder layers. + num_nextn_layers = getattr(config, "num_nextn_predict_layers", 0) + self.hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + f"model.layers.{config.num_hidden_layers + i}.": None + for i in range(num_nextn_layers) + } + ) + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.model.embed_input_ids(input_ids) @@ -292,12 +298,5 @@ def compute_logits( return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - # Skip the speculative (MTP) layers, which are loaded by the - # draft model instead. - num_nextn_layers = getattr(self.config, "num_nextn_predict_layers", 0) - skip_prefixes = [ - f"model.layers.{self.config.num_hidden_layers + i}." - for i in range(num_nextn_layers) - ] - loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes) - return loader.load_weights(weights) + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/glmasr.py b/vllm/model_executor/models/glmasr.py index 454687f469f7..df1b4e0d4b60 100644 --- a/vllm/model_executor/models/glmasr.py +++ b/vllm/model_executor/models/glmasr.py @@ -916,6 +916,10 @@ class GlmAsrForConditionalGeneration( ): supported_languages = ISO639_1_SUPPORTED_LANGS + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={"audio_tower.embed_positions": None} + ) + packed_modules_mapping = { "qkv_proj": ["q_proj", "k_proj", "v_proj"], "gate_up_proj": ["gate_proj", "up_proj"], @@ -1083,9 +1087,8 @@ def compute_logits( return self.language_model.compute_logits(hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - skip_prefixes = ["audio_tower.embed_positions"] - loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes) - return loader.load_weights(weights) + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) @classmethod def _get_audio_token(cls, model_config: ModelConfig) -> str: diff --git a/vllm/model_executor/models/gpt2.py b/vllm/model_executor/models/gpt2.py index 01dc119f8507..07801cd098b1 100644 --- a/vllm/model_executor/models/gpt2.py +++ b/vllm/model_executor/models/gpt2.py @@ -52,6 +52,7 @@ from .interfaces import SupportsCrossEncoding, SupportsPP from .utils import ( AutoWeightsLoader, + WeightsMapper, make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, @@ -182,6 +183,11 @@ def forward( @support_torch_compile class GPT2Model(nn.Module): + # Drop attention mask buffers; NOTE: "c_attn.bias" must not be dropped. + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_substr={".attn.bias": None, ".attn.masked_bias": None} + ) + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() @@ -253,11 +259,10 @@ def _transpose_conv1d( yield name, loaded_weight def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - # Skip attention mask buffers; NOTE: "c_attn.bias" must not be skipped. - loader = AutoWeightsLoader( - self, skip_substrs=[".attn.bias", ".attn.masked_bias"] + loader = AutoWeightsLoader(self) + return loader.load_weights( + self._transpose_conv1d(weights), mapper=self.hf_to_vllm_mapper ) - return loader.load_weights(self._transpose_conv1d(weights)) class GPT2LMHeadModel(nn.Module, SupportsPP): diff --git a/vllm/model_executor/models/gpt_j.py b/vllm/model_executor/models/gpt_j.py index 44dec8734574..8c74e6ad6fa9 100644 --- a/vllm/model_executor/models/gpt_j.py +++ b/vllm/model_executor/models/gpt_j.py @@ -243,7 +243,11 @@ class GPTJForCausalLM(nn.Module, SupportsPP): ".q_proj": (".qkv_proj", "q"), ".k_proj": (".qkv_proj", "k"), ".v_proj": (".qkv_proj", "v"), - } + }, + orig_to_new_substr={ + "attn.bias": None, + "attn.masked_bias": None, + }, ) def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): @@ -291,5 +295,5 @@ def compute_logits( return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self, skip_substrs=["attn.bias", "attn.masked_bias"]) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/gpt_neox.py b/vllm/model_executor/models/gpt_neox.py index 8af7eb67923e..e360580b9389 100644 --- a/vllm/model_executor/models/gpt_neox.py +++ b/vllm/model_executor/models/gpt_neox.py @@ -48,6 +48,7 @@ from .interfaces import SupportsPP from .utils import ( AutoWeightsLoader, + WeightsMapper, make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, @@ -196,6 +197,10 @@ def forward( @support_torch_compile class GPTNeoXModel(nn.Module): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_substr={"attention.bias": None, "attention.masked_bias": None} + ) + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() @@ -262,10 +267,10 @@ def _repack_qkv( yield name, loaded_weight def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, skip_substrs=["attention.bias", "attention.masked_bias"] + loader = AutoWeightsLoader(self) + return loader.load_weights( + self._repack_qkv(weights), mapper=self.hf_to_vllm_mapper ) - return loader.load_weights(self._repack_qkv(weights)) class GPTNeoXForCausalLM(nn.Module, SupportsPP): diff --git a/vllm/model_executor/models/idefics2_vision_model.py b/vllm/model_executor/models/idefics2_vision_model.py index 72a26993a537..e3f59138d086 100644 --- a/vllm/model_executor/models/idefics2_vision_model.py +++ b/vllm/model_executor/models/idefics2_vision_model.py @@ -19,7 +19,6 @@ """PyTorch Idefics2 model.""" from collections.abc import Iterable -from typing import ClassVar import torch from torch import nn @@ -360,7 +359,7 @@ def forward( class Idefics2VisionTransformer(nn.Module): - hf_to_vllm_mapper: ClassVar[WeightsMapper] = WeightsMapper( + hf_to_vllm_mapper = WeightsMapper( orig_to_new_stacked={ ".q_proj": (".qkv_proj", "q"), ".k_proj": (".qkv_proj", "k"), @@ -408,6 +407,13 @@ def __init__( if require_post_norm else nn.Identity() ) + # head is a pooling header absent from this model. + orig_to_new_prefix: dict[str, str | None] = {"head.": None} + if not require_post_norm: + orig_to_new_prefix["post_layernorm."] = None + self.hf_to_vllm_mapper = self.hf_to_vllm_mapper | WeightsMapper( + orig_to_new_prefix=orig_to_new_prefix + ) def get_input_embeddings(self): return self.embeddings @@ -469,11 +475,7 @@ def forward( return last_hidden_state def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - # head is a pooling header absent from this model. - skip_prefixes = ["head."] - if not self.require_post_norm: - skip_prefixes.append("post_layernorm.") - loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes) + loader = AutoWeightsLoader(self) layer_count = len(self.encoder.layers) diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 555051237336..c9d8b47ec19f 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -1150,7 +1150,7 @@ class SupportsLateInteraction(Protocol): class SupportsQuant: """The interface required for all models that support quantization.""" - hf_to_vllm_mapper: ClassVar["WeightsMapper | None"] = None + hf_to_vllm_mapper: "WeightsMapper | None" = None packed_modules_mapping: ClassVar[dict[str, list[str]]] quant_config: QuantizationConfig | None = None @@ -1184,8 +1184,7 @@ def _maybe_apply_model_mapping(self): if self.quant_config is None: return if (hf_to_vllm_mapper := self.hf_to_vllm_mapper) is not None: - unstacked_mapper = hf_to_vllm_mapper.get_unstacked_mapper() - self.quant_config.apply_vllm_mapper(unstacked_mapper) + self.quant_config.apply_vllm_mapper(hf_to_vllm_mapper.get_rename_mapper()) if packed_modules_mapping := getattr(self, "packed_modules_mapping", None): self.quant_config.packed_modules_mapping.update(packed_modules_mapping) diff --git a/vllm/model_executor/models/interns1_pro.py b/vllm/model_executor/models/interns1_pro.py index e85ad87cdc7c..ee52fbd03ecd 100644 --- a/vllm/model_executor/models/interns1_pro.py +++ b/vllm/model_executor/models/interns1_pro.py @@ -616,17 +616,18 @@ def get_frope_params_map(self) -> str: def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): """load weights""" - skip_prefixes = ["model.time_series."] + orig_to_new_prefix: dict[str, str | None] = { + "model.visual.": "visual.", + "lm_head.": "language_model.lm_head.", + "model.language_model.": "language_model.model.", + "model.time_series.": None, + } if self.visual is None: - skip_prefixes.append("visual.") + orig_to_new_prefix["visual."] = None # FIXME(Isotr0py): See if we can avoid tighing FoPE to PP layers weights_mapper = WeightsMapper( - orig_to_new_prefix={ - "model.visual.": "visual.", - "lm_head.": "language_model.lm_head.", - "model.language_model.": "language_model.model.", - }, + orig_to_new_prefix=orig_to_new_prefix, orig_to_new_suffix=self.get_frope_params_map(), ) - loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=weights_mapper) diff --git a/vllm/model_executor/models/interns2_preview.py b/vllm/model_executor/models/interns2_preview.py index 6efc98aabc1f..7131ef33bd00 100644 --- a/vllm/model_executor/models/interns2_preview.py +++ b/vllm/model_executor/models/interns2_preview.py @@ -13,7 +13,7 @@ Qwen3VLMultiModalProcessor, Qwen3VLProcessingInfo, ) -from .utils import AutoWeightsLoader +from .utils import AutoWeightsLoader, WeightsMapper class InternS2PreviewProcessingInfo(Qwen3VLProcessingInfo): @@ -30,9 +30,13 @@ def get_hf_processor(self, **kwargs: object) -> AutoProcessor: dummy_inputs=Qwen3VLDummyInputsBuilder, ) class InternS2PreviewForConditionalGeneration(Qwen3_5MoeForConditionalGeneration): - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=["mtp.", "model.time_series.", "time_series."], + # `mtp.` is already dropped by `Qwen3_5ForConditionalGeneration`. + hf_to_vllm_mapper = Qwen3_5MoeForConditionalGeneration.hf_to_vllm_mapper | ( + WeightsMapper( + orig_to_new_prefix={"model.time_series.": None, "time_series.": None} ) + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/internvl.py b/vllm/model_executor/models/internvl.py index cf602750403f..d360e5018283 100644 --- a/vllm/model_executor/models/internvl.py +++ b/vllm/model_executor/models/internvl.py @@ -59,7 +59,12 @@ SupportsMultiModal, SupportsPP, ) -from .utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix +from .utils import ( + AutoWeightsLoader, + WeightsMapper, + init_vllm_registered_model, + maybe_prefix, +) class InternVLImagePixelInputs(TensorSchema): @@ -554,6 +559,25 @@ class InternVLChatModel( ): supports_encoder_tp_data = True + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix=dict.fromkeys( + [ + "action_embed", + "temporal_embed", + "track_embed", + "track_embed_decoder", + "box_token", + "cg_criterion", + "cg_model", + "loc_encoder", + "loc_decoder", + "sam", + "temporal_token", + "track_token", + ] + ) + ) + @classmethod def get_placeholder_str(cls, modality: str, i: int) -> str | None: if modality.startswith("image"): @@ -865,23 +889,8 @@ def compute_logits( return self.language_model.compute_logits(hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - # unused modules appear in OpenGVLab/InternVideo2_5_Chat_8B - skip_prefixes = [ - "action_embed", - "temporal_embed", - "track_embed", - "track_embed_decoder", - "box_token", - "cg_criterion", - "cg_model", - "loc_encoder", - "loc_decoder", - "sam", - "temporal_token", - "track_token", - ] - loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes) - return loader.load_weights(weights) + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) def get_mm_mapping(self) -> MultiModelKeys: """ diff --git a/vllm/model_executor/models/jina.py b/vllm/model_executor/models/jina.py index 0595df8e3a28..456a51b3f396 100644 --- a/vllm/model_executor/models/jina.py +++ b/vllm/model_executor/models/jina.py @@ -43,6 +43,8 @@ class JinaForRanking(nn.Module, SupportsLateInteraction): is_pooling_model = True + hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"lm_head.": None}) + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() config = vllm_config.model_config.hf_config @@ -87,8 +89,8 @@ def forward( return hidden_states def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self, skip_prefixes=(["lm_head."])) - return loader.load_weights(weights) + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) class JinaForRankingPool(StepPool): diff --git a/vllm/model_executor/models/keye.py b/vllm/model_executor/models/keye.py index 2fee80b271c3..d64d2ebc64e0 100644 --- a/vllm/model_executor/models/keye.py +++ b/vllm/model_executor/models/keye.py @@ -721,7 +721,8 @@ class KeyeSiglipVisionModel(nn.Module): ".q_proj": (".qkv_proj", "q"), ".k_proj": (".qkv_proj", "k"), ".v_proj": (".qkv_proj", "v"), - } + }, + orig_to_new_prefix={"vision_model.head.": None}, ) def __init__( @@ -782,7 +783,7 @@ def forward( ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self, skip_prefixes=["vision_model.head."]) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/kimi_audio.py b/vllm/model_executor/models/kimi_audio.py index 900894612ce2..c999a694ad4e 100644 --- a/vllm/model_executor/models/kimi_audio.py +++ b/vllm/model_executor/models/kimi_audio.py @@ -380,6 +380,12 @@ class KimiAudioForConditionalGeneration( "model.embed_tokens.": "language_model.model.embed_tokens.", "model.norm.": "language_model.model.norm.", "lm_head.": "language_model.lm_head.", + # MIMO/TTS weights and any `model.` residue no rule above + # claimed: this model only does ASR (speech-to-text). + "model.": None, + "mimo_layers.": None, + "mimo_output.": None, + "mimo_norm.": None, }, orig_to_new_substr={ ".fc1.": ".mlp.fc1.", @@ -573,21 +579,8 @@ def compute_logits( return self.language_model.compute_logits(hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - """Load weights, skipping MIMO layers (TTS-only) for ASR.""" - # Filter out MIMO/TTS weights since we only do ASR (speech-to-text) - skipped_patterns = [ - # Audio tower - "model.", - # MIMO/TTS - "mimo_layers.", - "mimo_output.", - "mimo_norm.", - ] - - # Load main model weights (LLM + projector) with mapper - loader = AutoWeightsLoader(self, skip_prefixes=skipped_patterns) - loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) - return loaded + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) @classmethod def get_speech_to_text_config( diff --git a/vllm/model_executor/models/lfm2_siglip2.py b/vllm/model_executor/models/lfm2_siglip2.py index 31e7c0f14438..54fc1c5d2ff0 100644 --- a/vllm/model_executor/models/lfm2_siglip2.py +++ b/vllm/model_executor/models/lfm2_siglip2.py @@ -484,6 +484,10 @@ def __init__( require_post_norm=require_post_norm, prefix=maybe_prefix(prefix, "vision_model"), ) + if self.vision_model.post_layernorm is None: + self.hf_to_vllm_mapper = self.hf_to_vllm_mapper | WeightsMapper( + orig_to_new_prefix={"vision_model.post_layernorm.": None} + ) def forward( self, @@ -510,10 +514,7 @@ def forward( ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - skip_prefixes = [] - if self.vision_model.post_layernorm is None: - skip_prefixes.append("vision_model.post_layernorm.") - loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes) + loader = AutoWeightsLoader(self) # Drop layers omitted by num_hidden_layers_override. layer_count = len(self.vision_model.encoder.layers) diff --git a/vllm/model_executor/models/llama.py b/vllm/model_executor/models/llama.py index e3c5272224d4..f136cea23b1a 100644 --- a/vllm/model_executor/models/llama.py +++ b/vllm/model_executor/models/llama.py @@ -26,7 +26,6 @@ from collections.abc import Iterable from itertools import islice -from typing import ClassVar import torch from torch import nn @@ -343,7 +342,7 @@ def get_quant_config(self, vllm_config: VllmConfig) -> QuantizationConfig | None }, ) class LlamaModel(nn.Module, EagleModelMixin): - hf_to_vllm_mapper: ClassVar[WeightsMapper] = WeightsMapper( + hf_to_vllm_mapper = WeightsMapper( orig_to_new_stacked={ # weight_name: (param_name, shard_id) ".q_proj": (".qkv_proj", "q"), diff --git a/vllm/model_executor/models/llama4_eagle.py b/vllm/model_executor/models/llama4_eagle.py index e41b0c5dd283..ba9a8e325460 100644 --- a/vllm/model_executor/models/llama4_eagle.py +++ b/vllm/model_executor/models/llama4_eagle.py @@ -206,9 +206,5 @@ def transform(inputs): process_eagle_weight(self, name) return name, weight - loader = AutoWeightsLoader( - self, - # lm_head is tied with target model (Llama4ForCausalLM) - skip_prefixes=([]), - ) + loader = AutoWeightsLoader(self) loader.load_weights(map(transform, weights)) diff --git a/vllm/model_executor/models/llama_eagle.py b/vllm/model_executor/models/llama_eagle.py index 5d13b29c0cbf..50706bb32615 100644 --- a/vllm/model_executor/models/llama_eagle.py +++ b/vllm/model_executor/models/llama_eagle.py @@ -175,8 +175,5 @@ def transform(inputs): process_eagle_weight(self, name) return name, loaded_weight - loader = AutoWeightsLoader( - self, - skip_prefixes=None, - ) + loader = AutoWeightsLoader(self) loader.load_weights(map(transform, weights)) diff --git a/vllm/model_executor/models/llama_eagle3.py b/vllm/model_executor/models/llama_eagle3.py index 549e8b7bf63a..93c94e759949 100644 --- a/vllm/model_executor/models/llama_eagle3.py +++ b/vllm/model_executor/models/llama_eagle3.py @@ -415,18 +415,17 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): "Please provide mask_hidden in the weights." ) - skip_substrs = ["mask_hidden"] + orig_to_new_substr = {"mask_hidden": None} if not includes_draft_id_mapping: - skip_substrs.append("draft_id_to_target_id") + orig_to_new_substr["draft_id_to_target_id"] = None if not includes_embed_tokens: - skip_substrs.append("embed_tokens") + orig_to_new_substr["embed_tokens"] = None if not self.model.use_aux_hidden_state: - skip_substrs.append("fc.") + orig_to_new_substr["fc."] = None if not self.model.norm_before_fc: - skip_substrs.append("input_norm.") - loader = AutoWeightsLoader( - self, - skip_prefixes=None, - skip_substrs=skip_substrs, + orig_to_new_substr["input_norm."] = None + loader = AutoWeightsLoader(self) + loader.load_weights( + model_weights.items(), + mapper=WeightsMapper(orig_to_new_substr=orig_to_new_substr), ) - loader.load_weights(model_weights.items()) diff --git a/vllm/model_executor/models/longcat_flash_ngram.py b/vllm/model_executor/models/longcat_flash_ngram.py index 41f913cdfec1..e01672bb14c0 100644 --- a/vllm/model_executor/models/longcat_flash_ngram.py +++ b/vllm/model_executor/models/longcat_flash_ngram.py @@ -31,7 +31,7 @@ from .interfaces import SupportsLoRA, SupportsPP from .longcat_flash import FlashConfig, FlashModel -from .utils import AutoWeightsLoader, PPMissingLayer, maybe_prefix +from .utils import AutoWeightsLoader, PPMissingLayer, WeightsMapper, maybe_prefix def uses_ngram_embedding(config: FlashConfig) -> bool: @@ -204,6 +204,9 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: class LongcatFlashNgramForCausalLM(nn.Module, SupportsLoRA, SupportsPP): """LongCat-Flash-Lite for causal LM (MRV2-only, n-gram embedding).""" + # MTP weights are not part of this model. + hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"model.mtp.": None}) + packed_modules_mapping = { "qkv_proj": ["q_proj", "k_proj", "v_proj"], "gate_up_proj": ["gate_proj", "up_proj"], @@ -266,10 +269,9 @@ def get_expert_mapping(self): def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # AutoWeightsLoader routes ``model.*`` to FlashNgramModel.load_weights - # (which handles the ngram split) and ``lm_head.*`` to the head. MTP - # weights are not part of this model. - loader = AutoWeightsLoader(self, skip_prefixes=["model.mtp."]) - return loader.load_weights(weights) + # (which handles the ngram split) and ``lm_head.*`` to the head. + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) class LongcatNgramModelState(DefaultModelState): diff --git a/vllm/model_executor/models/mimo.py b/vllm/model_executor/models/mimo.py index 7be514e990a5..6de0c0e5967d 100644 --- a/vllm/model_executor/models/mimo.py +++ b/vllm/model_executor/models/mimo.py @@ -41,7 +41,7 @@ from vllm.model_executor.models.qwen2 import Qwen2ForCausalLM, Qwen2Model from vllm.sequence import IntermediateTensors -from .utils import AutoWeightsLoader, PPMissingLayer, maybe_prefix +from .utils import AutoWeightsLoader, PPMissingLayer, WeightsMapper, maybe_prefix logger = init_logger(__name__) @@ -87,6 +87,9 @@ def forward( class MiMoForCausalLM(Qwen2ForCausalLM, nn.Module): + # MTP layers are loaded by the draft model, not the main model. + hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"model.mtp_layers.": None}) + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): nn.Module.__init__(self) config = vllm_config.model_config.hf_config @@ -119,9 +122,8 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - # MTP layers are loaded by the draft model, not the main model. - loader = AutoWeightsLoader(self, skip_prefixes=["model.mtp_layers."]) - return loader.load_weights(weights) + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) def compute_logits( self, diff --git a/vllm/model_executor/models/mimo_v2_omni.py b/vllm/model_executor/models/mimo_v2_omni.py index bd2bdbf9018b..28205f9da29d 100644 --- a/vllm/model_executor/models/mimo_v2_omni.py +++ b/vllm/model_executor/models/mimo_v2_omni.py @@ -1189,6 +1189,7 @@ class MiMoV2OmniForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsQ # mapping for original checkpoint "lm_head.": "language_model.lm_head.", "model.": "language_model.model.", + "audio_tokenizer.": None, } ) @@ -1490,6 +1491,6 @@ def compute_logits( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: audio_loaded: set[str] = set() - loader = AutoWeightsLoader(self, skip_prefixes=["audio_tokenizer."]) + loader = AutoWeightsLoader(self) auto_loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) return audio_loaded | auto_loaded diff --git a/vllm/model_executor/models/minicpmo.py b/vllm/model_executor/models/minicpmo.py index bd8547420c6d..dd6f8d2f56b9 100644 --- a/vllm/model_executor/models/minicpmo.py +++ b/vllm/model_executor/models/minicpmo.py @@ -71,7 +71,12 @@ MiniCPMVProcessingInfo, _minicpmv_field_config, ) -from .utils import AutoWeightsLoader, cast_overflow_tensors, maybe_prefix +from .utils import ( + AutoWeightsLoader, + WeightsMapper, + cast_overflow_tensors, + maybe_prefix, +) CPU_DEVICE = torch.device("cpu") @@ -672,6 +677,9 @@ def forward( class MiniCPMOBaseModel: """Base mixin class for MiniCPM-O models with audio support.""" + # Unlike the vision-only MiniCPM-V models, audio weights are loaded here. + hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"tts": None}) + packed_modules_mapping = { "qkv_proj": [ "q_proj", @@ -718,8 +726,8 @@ def init_audio_module(self, *, vllm_config: VllmConfig, prefix: str = ""): return model def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self, skip_prefixes=["tts"]) - loaded = loader.load_weights(weights) + loader = AutoWeightsLoader(self) + loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) self._ensure_resampler_device() return loaded diff --git a/vllm/model_executor/models/minicpmv.py b/vllm/model_executor/models/minicpmv.py index 5d8eecfaaaa9..d7ab703f38c8 100644 --- a/vllm/model_executor/models/minicpmv.py +++ b/vllm/model_executor/models/minicpmv.py @@ -109,7 +109,7 @@ SupportsMultiModal, SupportsPP, ) -from .utils import AutoWeightsLoader, flatten_bn, maybe_prefix +from .utils import AutoWeightsLoader, WeightsMapper, flatten_bn, maybe_prefix # For profile run _MAX_FRAMES_PER_VIDEO = 16 @@ -1567,6 +1567,10 @@ def get_vision_hidden_states(self, data: MiniCPMVImagePixelInputs) -> torch.Tens class MiniCPMV2_6(MiniCPMVBaseModel, SupportsLoRA): + hf_to_vllm_mapper = WeightsMapper( + # The vision-only models have no audio tower or TTS head to load weights into. + orig_to_new_prefix={"apm.": None, "audio": None, "tts": None} + ) packed_modules_mapping = { "qkv_proj": [ "q_proj", @@ -1660,13 +1664,14 @@ def get_vision_hidden_states(self, data: MiniCPMVImagePixelInputs) -> torch.Tens return self.resampler(vision_embedding, tgt_sizes) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self, skip_prefixes=["apm.", "audio", "tts"]) - loaded = loader.load_weights(weights) + loader = AutoWeightsLoader(self) + loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) self._ensure_resampler_device() return loaded class MiniCPMV4_0(MiniCPMVBaseModel, SupportsLoRA): + hf_to_vllm_mapper = MiniCPMV2_6.hf_to_vllm_mapper packed_modules_mapping = { "qkv_proj": [ "q_proj", @@ -1758,13 +1763,14 @@ def get_vision_hidden_states(self, data: MiniCPMVImagePixelInputs) -> torch.Tens return self.resampler(vision_embedding, tgt_sizes) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self, skip_prefixes=["apm.", "audio", "tts"]) - loaded = loader.load_weights(weights) + loader = AutoWeightsLoader(self) + loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) self._ensure_resampler_device() return loaded class MiniCPMV4_5(MiniCPMVBaseModel, SupportsLoRA): + hf_to_vllm_mapper = MiniCPMV2_6.hf_to_vllm_mapper packed_modules_mapping = { "qkv_proj": [ "q_proj", @@ -1861,8 +1867,8 @@ def get_vision_hidden_states(self, data: MiniCPMVImagePixelInputs) -> torch.Tens return self.resampler(vision_embedding, tgt_sizes, all_temporal_ids) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self, skip_prefixes=["apm.", "audio", "tts"]) - loaded = loader.load_weights(weights) + loader = AutoWeightsLoader(self) + loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) self._ensure_resampler_device() return loaded diff --git a/vllm/model_executor/models/minicpmv4_6.py b/vllm/model_executor/models/minicpmv4_6.py index 69124a061c73..d0948715a9d1 100644 --- a/vllm/model_executor/models/minicpmv4_6.py +++ b/vllm/model_executor/models/minicpmv4_6.py @@ -958,6 +958,7 @@ class MiniCPMV4_6ForConditionalGeneration( "model.merger.": "merger.", "model.language_model.": "language_model.model.", "lm_head.": "language_model.lm_head.", + "mtp.": None, } ) @@ -1281,7 +1282,7 @@ def load_weights( self, weights: Iterable[tuple[str, torch.Tensor]], ) -> set[str]: - loader = AutoWeightsLoader(self, skip_prefixes=["mtp."]) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) def get_mm_mapping(self) -> MultiModelKeys: diff --git a/vllm/model_executor/models/minimax_m2.py b/vllm/model_executor/models/minimax_m2.py index b94d6fef0ac3..139529e55931 100644 --- a/vllm/model_executor/models/minimax_m2.py +++ b/vllm/model_executor/models/minimax_m2.py @@ -371,6 +371,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ["hidden_states", "residual"], config.hidden_size ) + # Drop spec-decode (MTP) layers; they are appended after the main + # decoder layers and have no destination in the main model. + if num_mtp := getattr(config, "num_mtp_modules", 0): + base = config.num_hidden_layers + self.hf_to_vllm_mapper = self.hf_to_vllm_mapper | WeightsMapper( + orig_to_new_prefix={f"layers.{base + i}.": None for i in range(num_mtp)} + ) + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) @@ -413,17 +421,7 @@ def forward( return hidden_states def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - # Skip spec-decode (MTP) layers; they are appended after the main - # decoder layers and have no destination in the main model. - skip_prefixes = None - num_mtp = getattr(self.config, "num_mtp_modules", 0) - if num_mtp: - base = self.config.num_hidden_layers - skip_prefixes = [f"layers.{base + i}." for i in range(num_mtp)] - loader = AutoWeightsLoader( - self, - skip_prefixes=skip_prefixes, - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/modernbert.py b/vllm/model_executor/models/modernbert.py index d182fa071594..dca5c9ec4559 100644 --- a/vllm/model_executor/models/modernbert.py +++ b/vllm/model_executor/models/modernbert.py @@ -434,6 +434,8 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: class ModernBertForTokenClassification(nn.Module): is_pooling_model = True + hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"drop": None}) + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() config = vllm_config.model_config.hf_config @@ -456,8 +458,8 @@ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.model.embed_input_ids(input_ids) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): - loader = AutoWeightsLoader(self, skip_prefixes=["drop"]) - loaded_params = loader.load_weights(weights) + loader = AutoWeightsLoader(self) + loaded_params = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) return loaded_params def forward( diff --git a/vllm/model_executor/models/moss_audio.py b/vllm/model_executor/models/moss_audio.py index bfcebe7f9210..324b3f55ae8a 100644 --- a/vllm/model_executor/models/moss_audio.py +++ b/vllm/model_executor/models/moss_audio.py @@ -1450,6 +1450,7 @@ class MossAudioModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA): "language_model.embed_tokens.": "language_model.model.embed_tokens.", "language_model.layers.": "language_model.model.layers.", "language_model.norm.": "language_model.model.norm.", + "audio_encoder.embed_positions": None, }, orig_to_new_stacked={ ".gate_proj": (".gate_up_proj", 0), @@ -1865,8 +1866,5 @@ def compute_logits( return self.language_model.compute_logits(hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=["audio_encoder.embed_positions"], - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index 6c7b253063a8..b62e46d44af4 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -708,7 +708,7 @@ class NemotronHForCausalLM( is_non_gated_moe: bool = True hf_to_vllm_mapper = WeightsMapper( - orig_to_new_prefix={"backbone": "model"}, + orig_to_new_prefix={"backbone": "model", "mtp": None}, orig_to_new_substr={"A_log": "A", "embeddings": "embed_tokens"}, orig_to_new_stacked={ ".q_proj": (".qkv_proj", "q"), @@ -887,5 +887,5 @@ def compute_logits( return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self, skip_prefixes=["mtp"]) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/nemotron_vl.py b/vllm/model_executor/models/nemotron_vl.py index 734968819b9a..818d8191971c 100644 --- a/vllm/model_executor/models/nemotron_vl.py +++ b/vllm/model_executor/models/nemotron_vl.py @@ -89,6 +89,12 @@ def get_hf_processor(self, **kwargs: object) -> LlamaNemotronNanoVLProcessor: dummy_inputs=BaseInternVLDummyInputsBuilder[NemotronVLProcessingInfo], ) class LlamaNemotronVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA): + # Ignore registered buffers, see + # https://huggingface.co/nvidia/C-RADIOv2-H/blob/main/input_conditioner.py#L28 + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_substr={"norm_mean": None, "norm_std": None} + ) + @classmethod def get_placeholder_str(cls, modality: str, i: int) -> str | None: if modality.startswith("image"): @@ -377,11 +383,8 @@ def compute_logits( return self.language_model.compute_logits(hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - ## Ignore registered_buffers - ## see https://huggingface.co/nvidia/C-RADIOv2-H/blob/main/input_conditioner.py#L28 # noqa: E501 - skip_substrs = ["norm_mean", "norm_std"] - loader = AutoWeightsLoader(self, skip_substrs=skip_substrs) - return loader.load_weights(weights) + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) def get_mm_mapping(self) -> MultiModelKeys: """ diff --git a/vllm/model_executor/models/paddleocr_vl.py b/vllm/model_executor/models/paddleocr_vl.py index 5f5302e7b78d..e9ecc332bdf8 100644 --- a/vllm/model_executor/models/paddleocr_vl.py +++ b/vllm/model_executor/models/paddleocr_vl.py @@ -887,7 +887,16 @@ class SiglipVisionModel(nn.Module): ".q_proj": (".qkv_proj", "q"), ".k_proj": (".qkv_proj", "k"), ".v_proj": (".qkv_proj", "v"), - } + }, + # The SigLIP attention pooling head and packing pos embedding are + # present in the checkpoint but absent from this vision tower. + orig_to_new_substr={ + "head.attention": None, + "head.layernorm": None, + "head.mlp": None, + "head.probe": None, + "packing_position_embedding": None, + }, ) def __init__( @@ -934,18 +943,7 @@ def forward( ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - # Skip the SigLIP attention pooling head and packing pos embedding - # present in the checkpoint but absent from this vision tower. - loader = AutoWeightsLoader( - self, - skip_substrs=[ - "head.attention", - "head.layernorm", - "head.mlp", - "head.probe", - "packing_position_embedding", - ], - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/phi4mm.py b/vllm/model_executor/models/phi4mm.py index 359cc1613bea..cc15ed31462c 100644 --- a/vllm/model_executor/models/phi4mm.py +++ b/vllm/model_executor/models/phi4mm.py @@ -1016,6 +1016,7 @@ class Phi4MMForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal): hf_to_vllm_mapper = WeightsMapper( orig_to_new_substr={ "base_layer.": "", + "lora": None, }, orig_to_new_prefix={ "model.embed_tokens_extend.audio_embed.audio_projection.vision.": "embed_tokens_extend.audio_projection_for_vision.", # noqa: E501 @@ -1258,7 +1259,7 @@ def compute_logits( return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> None: - loader = AutoWeightsLoader(self, skip_substrs=["lora"]) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) def get_mm_mapping(self) -> MultiModelKeys: diff --git a/vllm/model_executor/models/qwen2_5_omni_thinker.py b/vllm/model_executor/models/qwen2_5_omni_thinker.py index 149ab7bc15da..13ad6173b067 100644 --- a/vllm/model_executor/models/qwen2_5_omni_thinker.py +++ b/vllm/model_executor/models/qwen2_5_omni_thinker.py @@ -1112,6 +1112,8 @@ class Qwen2_5OmniThinkerForConditionalGeneration( "thinker.lm_head.": "language_model.lm_head.", "thinker.model.": "language_model.model.", "thinker.": "", + "talker.": None, + "token2wav.": None, } ) packed_modules_mapping = { @@ -1591,7 +1593,7 @@ def compute_logits( return self.language_model.compute_logits(hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self, skip_prefixes=["talker.", "token2wav."]) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) def get_mm_mapping(self) -> MultiModelKeys: diff --git a/vllm/model_executor/models/qwen3_5.py b/vllm/model_executor/models/qwen3_5.py index e8cc041ce423..2624102192ea 100644 --- a/vllm/model_executor/models/qwen3_5.py +++ b/vllm/model_executor/models/qwen3_5.py @@ -319,7 +319,7 @@ class Qwen3_5ForCausalLMBase( # `model.language_model.` prefix inherited from the VL training stack. # Strip it so both prefixed and clean checkpoints load correctly. hf_to_vllm_mapper = WeightsMapper( - orig_to_new_prefix={"model.language_model.": "model."}, + orig_to_new_prefix={"model.language_model.": "model.", "mtp.": None}, ) def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): @@ -430,10 +430,7 @@ def compute_logits( return self.logits_processor(self.lm_head, hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=["mtp."], - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) def get_mrope_input_positions( @@ -470,6 +467,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): class Qwen3_5ForConditionalGeneration(Qwen3VLForConditionalGeneration, IsHybrid): supports_multimodal_pruning = True + hf_to_vllm_mapper = ( + Qwen3VLForConditionalGeneration.hf_to_vllm_mapper + | WeightsMapper(orig_to_new_prefix={"mtp.": None}) + ) + packed_modules_mapping = Qwen3VLForConditionalGeneration.packed_modules_mapping | { "in_proj_qkvz": ["in_proj_qkv", "in_proj_z"], "in_proj_ba": ["in_proj_b", "in_proj_a"], @@ -590,10 +592,7 @@ def forward( return hidden_states def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=["mtp."], - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) @classmethod diff --git a/vllm/model_executor/models/qwen3_asr.py b/vllm/model_executor/models/qwen3_asr.py index 328fcdbdac27..301c08cee06e 100644 --- a/vllm/model_executor/models/qwen3_asr.py +++ b/vllm/model_executor/models/qwen3_asr.py @@ -372,6 +372,8 @@ class Qwen3ASRForConditionalGeneration( "model.language_model.": "language_model.model.", "model.multi_modal_projector.linear_1.": "audio_tower.proj1.", "model.multi_modal_projector.linear_2.": "audio_tower.proj2.", + "talker.": None, + "code2wav.": None, } ) @@ -545,10 +547,7 @@ def compute_logits( return self.language_model.compute_logits(hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=["talker.", "code2wav."], - ) + loader = AutoWeightsLoader(self) loaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) return loaded_weights diff --git a/vllm/model_executor/models/qwen3_asr_forced_aligner.py b/vllm/model_executor/models/qwen3_asr_forced_aligner.py index 56c57f477da1..2cbf4c111106 100644 --- a/vllm/model_executor/models/qwen3_asr_forced_aligner.py +++ b/vllm/model_executor/models/qwen3_asr_forced_aligner.py @@ -59,6 +59,8 @@ class Qwen3ASRForcedAlignerForTokenClassification( "thinker.lm_head.": "classifier.", "thinker.model.": "language_model.model.", "thinker.": "", + "talker.": None, + "code2wav.": None, } ) @@ -113,8 +115,5 @@ def forward( return self.classifier(hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=["talker.", "code2wav."], - ) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index 35c04331b001..7ce397d24796 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -835,21 +835,18 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): model_weights["model.mask_embedding"] = mask_embedding self.model.has_separate_mask_embedding = True - skip_substrs = [] + orig_to_new_substr = {} if not includes_draft_id_mapping: - skip_substrs.append("draft_id_to_target_id") + orig_to_new_substr["draft_id_to_target_id"] = None if not includes_embed_tokens: - skip_substrs.append("embed_tokens") + orig_to_new_substr["embed_tokens"] = None if not self.model.use_aux_hidden_state: - skip_substrs.append("fc.") + orig_to_new_substr["fc."] = None if not self.model.has_separate_mask_embedding: - skip_substrs.append("mask_embedding") - loader = AutoWeightsLoader( - self, - skip_prefixes=None, - skip_substrs=skip_substrs, - ) - loader.load_weights(model_weights.items()) + orig_to_new_substr["mask_embedding"] = None + mapper = WeightsMapper(orig_to_new_substr=orig_to_new_substr) + loader = AutoWeightsLoader(self) + loader.load_weights(model_weights.items(), mapper=mapper) self.model._build_fused_kv_buffers() def _read_mask_embedding(self) -> torch.Tensor | None: diff --git a/vllm/model_executor/models/qwen3_dspark.py b/vllm/model_executor/models/qwen3_dspark.py index d15aa45461cf..250d54d9ca88 100644 --- a/vllm/model_executor/models/qwen3_dspark.py +++ b/vllm/model_executor/models/qwen3_dspark.py @@ -30,7 +30,12 @@ ) from .qwen3_dflash import DFlashQwen3ForCausalLM, DFlashQwen3Model -from .utils import AutoWeightsLoader, maybe_prefix, process_eagle_weight +from .utils import ( + AutoWeightsLoader, + WeightsMapper, + maybe_prefix, + process_eagle_weight, +) logger = init_logger(__name__) @@ -275,16 +280,17 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): # mask_embedding is an unused placeholder param; DSpark masks via the vocab row. # embed_tokens / lm_head are optional; when omitted they are shared from # the target by load_dspark_model, so skip the unloaded params here. - skip_substrs = ["mask_embedding"] + orig_to_new_substr = {"mask_embedding": None} if not includes_embed_tokens: - skip_substrs.append("embed_tokens") + orig_to_new_substr["embed_tokens"] = None if not includes_lm_head: - skip_substrs.append("lm_head") + orig_to_new_substr["lm_head"] = None if not includes_draft_id_mapping: - skip_substrs.append("draft_id_to_target_id") + orig_to_new_substr["draft_id_to_target_id"] = None if self.model.confidence_head is None or not includes_confidence_head: self.model.confidence_head = None - skip_substrs.append("confidence_head") - loader = AutoWeightsLoader(self, skip_substrs=skip_substrs) - loader.load_weights(model_weights.items()) + orig_to_new_substr["confidence_head"] = None + mapper = WeightsMapper(orig_to_new_substr=orig_to_new_substr) + loader = AutoWeightsLoader(self) + loader.load_weights(model_weights.items(), mapper=mapper) self.model._build_fused_kv_buffers() diff --git a/vllm/model_executor/models/qwen3_eagle3.py b/vllm/model_executor/models/qwen3_eagle3.py index 3255c65e2736..5b4e5a7a57cd 100644 --- a/vllm/model_executor/models/qwen3_eagle3.py +++ b/vllm/model_executor/models/qwen3_eagle3.py @@ -425,18 +425,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): "Please provide mask_hidden in the weights." ) - skip_substrs = ["mask_hidden"] + orig_to_new_substr = {"mask_hidden": None} if not includes_draft_id_mapping: - skip_substrs.append("draft_id_to_target_id") + orig_to_new_substr["draft_id_to_target_id"] = None if not includes_embed_tokens: - skip_substrs.append("embed_tokens") + orig_to_new_substr["embed_tokens"] = None if not self.model.use_aux_hidden_state: - skip_substrs.append("fc.") + orig_to_new_substr["fc."] = None if not self.model.norm_before_fc: - skip_substrs.append("input_norm.") - loader = AutoWeightsLoader( - self, - skip_prefixes=None, - skip_substrs=skip_substrs, - ) - loader.load_weights(model_weights.items()) + orig_to_new_substr["input_norm."] = None + mapper = WeightsMapper(orig_to_new_substr=orig_to_new_substr) + loader = AutoWeightsLoader(self) + loader.load_weights(model_weights.items(), mapper=mapper) diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py index 9511fcee9e7f..51fd036f4e39 100644 --- a/vllm/model_executor/models/qwen3_next.py +++ b/vllm/model_executor/models/qwen3_next.py @@ -727,6 +727,9 @@ class Qwen3NextForCausalLM( IsHybrid, SupportsEagle3, ): + # MTP weights are loaded by the draft model, not this one. + hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"mtp.": None}) + packed_modules_mapping = { "qkv_proj": [ "q_proj", @@ -833,5 +836,5 @@ def compute_logits( return self.logits_processor(self.lm_head, hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self, skip_prefixes=["mtp."]) - return loader.load_weights(weights) + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/qwen3_omni_moe_thinker.py b/vllm/model_executor/models/qwen3_omni_moe_thinker.py index 0df29a7e713e..814bce5c8794 100755 --- a/vllm/model_executor/models/qwen3_omni_moe_thinker.py +++ b/vllm/model_executor/models/qwen3_omni_moe_thinker.py @@ -1587,6 +1587,8 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( "thinker.lm_head.": "language_model.lm_head.", "thinker.model.": "language_model.model.", "thinker.": "", + "talker.": None, + "code2wav.": None, } ) @@ -1943,10 +1945,7 @@ def compute_logits( return self.language_model.compute_logits(hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=["talker.", "code2wav."], - ) + loader = AutoWeightsLoader(self) loaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) return loaded_weights diff --git a/vllm/model_executor/models/roberta.py b/vllm/model_executor/models/roberta.py index 948a939b9d58..aca929755f10 100644 --- a/vllm/model_executor/models/roberta.py +++ b/vllm/model_executor/models/roberta.py @@ -147,19 +147,21 @@ def _build_model( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): weights_list = list(weights) + orig_to_new_prefix: dict[str, str | None] = {"lm_head.": None} has_roberta_prefix = any( name.startswith("roberta.") for name, _ in weights_list ) if has_roberta_prefix: # For models with the `roberta.` prefix e.g. # `FacebookAI/roberta-base` - mapper = WeightsMapper(orig_to_new_prefix={"roberta.": "model."}) + orig_to_new_prefix["roberta."] = "model." else: # For models without the `roberta.` prefix e.g. # `sentence-transformers/stsb-roberta-base-v2` - mapper = WeightsMapper(orig_to_new_prefix={"": "model."}) + orig_to_new_prefix[""] = "model." + mapper = WeightsMapper(orig_to_new_prefix=orig_to_new_prefix) - loader = AutoWeightsLoader(self, skip_prefixes=["lm_head."]) + loader = AutoWeightsLoader(self) return loader.load_weights(weights_list, mapper=mapper) diff --git a/vllm/model_executor/models/siglip.py b/vllm/model_executor/models/siglip.py index 6d924a286501..823cb8d73767 100644 --- a/vllm/model_executor/models/siglip.py +++ b/vllm/model_executor/models/siglip.py @@ -738,6 +738,16 @@ def __init__( ) self.last_hs_proc = partial(self.maybe_layer_norm_and_apply_head) + drops = dict[str, None]() + if self.post_layernorm is None: + drops["post_layernorm."] = None + if self.head is None: + drops["head."] = None + if drops: + self.hf_to_vllm_mapper = self.hf_to_vllm_mapper | WeightsMapper( + orig_to_new_prefix=drops + ) + @property def dtype(self): return next(self.parameters()).dtype @@ -797,12 +807,7 @@ def maybe_layer_norm_and_apply_head( return encoder_outputs def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - skip_prefixes = [] - if self.post_layernorm is None: - skip_prefixes.append("post_layernorm.") - if self.head is None: - skip_prefixes.append("head.") - loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes) + loader = AutoWeightsLoader(self) layer_count = len(self.encoder.layers) @@ -912,6 +917,8 @@ def forward( class SiglipEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant): is_pooling_model = True + hf_to_vllm_mapper = WeightsMapper(orig_to_new_substr={".position_ids": None}) + packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]} @classmethod @@ -1154,8 +1161,7 @@ def forward( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): loader = AutoWeightsLoader( self, - skip_substrs=[".position_ids"], ignore_unexpected_prefixes=["logit_scale.", "logit_bias."], ) - return loader.load_weights(weights) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/skyworkr1v.py b/vllm/model_executor/models/skyworkr1v.py index 255827eff341..c44460aa0c72 100644 --- a/vllm/model_executor/models/skyworkr1v.py +++ b/vllm/model_executor/models/skyworkr1v.py @@ -35,7 +35,12 @@ BaseInternVLMultiModalProcessor, BaseInternVLProcessingInfo, ) -from .utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix +from .utils import ( + AutoWeightsLoader, + WeightsMapper, + init_vllm_registered_model, + maybe_prefix, +) class SkyworkR1VImagePixelInputs(TensorSchema): @@ -120,6 +125,25 @@ def get_hf_processor(self, **kwargs: object) -> InternVLProcessor: dummy_inputs=BaseInternVLDummyInputsBuilder, ) class SkyworkR1VChatModel(nn.Module, SupportsMultiModal, SupportsPP): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix=dict.fromkeys( + [ + "action_embed", + "temporal_embed", + "track_embed", + "track_embed_decoder", + "box_token", + "cg_criterion", + "cg_model", + "loc_encoder", + "loc_decoder", + "sam", + "temporal_token", + "track_token", + ] + ) + ) + @classmethod def get_placeholder_str(cls, modality: str, i: int) -> str | None: if modality.startswith("image"): @@ -381,19 +405,5 @@ def compute_logits( return self.language_model.compute_logits(hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - skip_prefixes = [ - "action_embed", - "temporal_embed", - "track_embed", - "track_embed_decoder", - "box_token", - "cg_criterion", - "cg_model", - "loc_encoder", - "loc_decoder", - "sam", - "temporal_token", - "track_token", - ] - loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes) - return loader.load_weights(weights) + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py index 1d0bbb71a69f..1d8fa5b5d7ec 100644 --- a/vllm/model_executor/models/transformers/base.py +++ b/vllm/model_executor/models/transformers/base.py @@ -122,10 +122,6 @@ def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""): self.tp_group = get_tp_group() # Attrs for weight loading (see self.load_weights) - self.skip_prefixes: list[str] = [] - """Skip loading weights whose qualname starts with these prefixes.""" - self.skip_substrs: list[str] = [] - """Skip loading weights whose qualname contains these substrings.""" self.ignore_unexpected_prefixes: list[str] = [] """Ignore unexpected weights whose qualname starts with these prefixes.""" self.ignore_unexpected_suffixes: list[str] = [] @@ -731,8 +727,6 @@ def forward( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader( self, - skip_prefixes=self.skip_prefixes, - skip_substrs=self.skip_substrs, ignore_unexpected_prefixes=self.ignore_unexpected_prefixes, ignore_unexpected_suffixes=self.ignore_unexpected_suffixes, ) diff --git a/vllm/model_executor/models/transformers/legacy.py b/vllm/model_executor/models/transformers/legacy.py index 49c5e9dcf68a..9c8a5c45b6f1 100644 --- a/vllm/model_executor/models/transformers/legacy.py +++ b/vllm/model_executor/models/transformers/legacy.py @@ -30,26 +30,24 @@ class LegacyMixin: def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""): super().__init__(vllm_config=vllm_config, prefix=prefix) - # Skip unsupported/unwanted output embeddings layers - self.skip_prefixes.extend( - [ - "model.lm_head.", - "model.predictions.", - "model.qa_outputs.", - "model.embeddings_project.", - "model.discriminator_predictions.", - ] + # Drop unsupported/unwanted output embeddings layers. + self.hf_to_vllm_mapper.orig_to_new_prefix.update( + { + "model.lm_head.": None, + "model.predictions.": None, + "model.qa_outputs.": None, + "model.embeddings_project.": None, + "model.discriminator_predictions.": None, + } ) - # Some encoder models have the position_ids buffer in the checkpoint. - # vLLM will always pass position_ids as an argument, so we skip loading - # the buffer if it exists - self.skip_substrs.append("position_ids") + # Some encoder models have the position_ids buffer in the checkpoint. vLLM will + # always pass position_ids as an argument, so we drop the buffer if it exists. + self.hf_to_vllm_mapper.orig_to_new_substr["position_ids"] = None - # Some encoder models have the bias of the final classifier layer - # in the checkpoint. vLLM does not use this bias, so we skip loading - # it if it exists - self.skip_substrs.append("score.bias") + # Some encoder models have the bias of the final classifier layer in the + # checkpoint. vLLM does not use this bias, so we drop it if it exists. + self.hf_to_vllm_mapper.orig_to_new_substr["score.bias"] = None # roberta-like models an extra padding in positions. # FIXME(Isotr0py): This is quite hacky for roberta edge case, diff --git a/vllm/model_executor/models/utils.py b/vllm/model_executor/models/utils.py index 2f1d6f3b5c17..9c6c7f6bb2cb 100644 --- a/vllm/model_executor/models/utils.py +++ b/vllm/model_executor/models/utils.py @@ -160,14 +160,26 @@ def apply_dict(self, values: dict[str, Any]) -> dict[str, Any]: if (out_name := self._map_name(name)) is not None } - def get_unstacked_mapper(self) -> "WeightsMapper": - """Mapper variant that drops stacked maps, keeping all genuine renames/prefixes. - - Consumers that reference the checkpoint's *unstacked* module names (LoRA name - parsing and the quantization config's layer lists) need the constituent names - (e.g. `q_proj`) to survive rather than being rewritten to the stacked vLLM name - (`qkv_proj`).""" - return replace(self, orig_to_new_stacked={}) + def get_rename_mapper(self) -> "WeightsMapper": + """Mapper variant keeping only the renames. + + This is what consumers that *name* modules rather than load them need: + LoRA name parsing and the quantization config's layer lists. + + Stacked maps are dropped so that constituent names (e.g. `q_proj`) survive + rather than being rewritten to the stacked vLLM name (`qkv_proj`). Mappings to + `None` are dropped because "do not load this weight" is meaningless to such a + consumer, and applying it would silently shrink a quantization config's ignore + list or make LoRA name parsing fail.""" + remove_none = lambda d: {k: v for k, v in d.items() if v is not None} + return replace( + self, + orig_to_new_regex=remove_none(self.orig_to_new_regex), + orig_to_new_substr=remove_none(self.orig_to_new_substr), + orig_to_new_stacked={}, + orig_to_new_prefix=remove_none(self.orig_to_new_prefix), + orig_to_new_suffix=remove_none(self.orig_to_new_suffix), + ) def _get_tied_embedding_params(module: nn.Module) -> dict[str, str]: @@ -201,32 +213,28 @@ class AutoWeightsLoader: """ # Models trained using early version ColossalAI or quantized by - # GPTQModel may include these tensors in checkpoint. Skip them. - ROTARY_EMBEDS_UNUSED_WEIGHTS = [ - "rotary_pos_emb.inv_freq", - "rotary_emb.inv_freq", - "rotary_emb.cos_cached", - "rotary_emb.sin_cached", - ] + # GPTQModel may include these tensors in checkpoint. Drop them. + REMOVE_UNUSED_ROTARY_EMBEDS_MAPPER = WeightsMapper( + orig_to_new_substr={ + "rotary_pos_emb.inv_freq": None, + "rotary_emb.inv_freq": None, + "rotary_emb.cos_cached": None, + "rotary_emb.sin_cached": None, + } + ) def __init__( self, module: nn.Module, *, - skip_prefixes: list[str] | None = None, - skip_substrs: list[str] | None = None, ignore_unexpected_prefixes: list[str] | None = None, ignore_unexpected_suffixes: list[str] | None = None, ) -> None: super().__init__() self.module = module - self.skip_prefixes = skip_prefixes or [] - self.skip_substrs = skip_substrs or [] self.ignore_unexpected_prefixes = ignore_unexpected_prefixes or [] self.ignore_unexpected_suffixes = ignore_unexpected_suffixes or [] - # update default skip_substrs - self.skip_substrs += self.ROTARY_EMBEDS_UNUSED_WEIGHTS # Weight tying makes two qualnames point at the same `nn.Parameter` # (e.g. `lm_head.weight` and `model.embed_tokens.weight`). Loading both @@ -265,11 +273,7 @@ def _get_qualname(self, prefix: str, rest: str) -> str: return ".".join((prefix, rest)) def _can_skip(self, qualname: str) -> bool: - return ( - qualname in self.aliased_params - or any(qualname.startswith(p) for p in self.skip_prefixes) - or any(substr in qualname for substr in self.skip_substrs) - ) + return qualname in self.aliased_params def _can_ignore_unexpected(self, qualname: str) -> bool: iup = (qualname.startswith(p) for p in self.ignore_unexpected_prefixes) @@ -375,11 +379,6 @@ def _load_module( prefix = self._get_qualname(base_prefix, child_prefix) if child_prefix in child_modules: - if self._can_skip(prefix + "."): - logger.debug("Skipping module %s", prefix) - - continue - yield from self._load_module( prefix, child_modules[child_prefix], child_weights ) @@ -393,9 +392,7 @@ def _load_module( prefix, child_params[child_prefix], child_weights ) else: - can_skip_module = self._can_skip(prefix + ".") - can_skip_param = self._can_skip(prefix) - if can_skip_module or can_skip_param: + if self._can_skip(prefix): logger.debug("Skipping missing %s", prefix) continue @@ -429,18 +426,19 @@ def load_weights( # Ignore unexpected biases (typically from GPTQ models) self.ignore_unexpected_suffixes.append(".bias") + mapper = mapper or WeightsMapper() # Many models store quant_config in the base model instead of the causal model. # We look at the causal model's direct children for this reason. modules = (self.module, *self.module.children()) iterator = (m.quant_config for m in modules if hasattr(m, "quant_config")) if quant_config := next(iterator, None): # Get mappings and ignore prefixes for KV cache quantization scales - mapper = mapper or WeightsMapper() mapper |= quant_config.get_cache_scale_mapper() ignore_unexpected_suffixes = quant_config._ignore_unexpected_suffixes self.ignore_unexpected_suffixes.extend(ignore_unexpected_suffixes) - if mapper is not None: - weights = mapper.apply(weights) + mapper |= self.REMOVE_UNUSED_ROTARY_EMBEDS_MAPPER + + weights = mapper.apply(weights) weights = self._filter_skipped(weights) autoloaded_weights = set(self._load_module("", self.module, weights)) diff --git a/vllm/model_executor/models/whisper.py b/vllm/model_executor/models/whisper.py index 30d1795b058c..82d20bf0f72a 100644 --- a/vllm/model_executor/models/whisper.py +++ b/vllm/model_executor/models/whisper.py @@ -838,6 +838,7 @@ class WhisperForConditionalGeneration( ".encoder_attn.k_proj": (".encoder_attn.kv_proj", 0), ".encoder_attn.v_proj": (".encoder_attn.kv_proj", 1), }, + orig_to_new_prefix={"proj_out.": None}, ) # Whisper only supports audio-conditioned generation. @@ -1048,7 +1049,7 @@ def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self, skip_prefixes=["proj_out."]) + loader = AutoWeightsLoader(self) # add fake zeros bias for k_proj to state_dict weights = _create_fake_bias_for_k_proj(weights, ".k_proj.weight") diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index f28df04d4b8f..2a82915f6e58 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -947,11 +947,12 @@ def _make_deepseek_v4_weights_mapper( # When shared experts are fused into the routed MXFP4 grouped GEMM, the # shared_experts tensors are redirected to routed expert slots ; leave # their names untouched here. - substr_map = ( + orig_to_new_substr: dict[str, str | None] = ( {} if fuse_shared_experts else {".shared_experts.w2": ".shared_experts.down_proj"} ) + orig_to_new_substr["mtp."] = None return WeightsMapper( orig_to_new_prefix={ "layers.": "model.layers.", @@ -967,7 +968,7 @@ def _make_deepseek_v4_weights_mapper( ".ffn.gate.bias": ".ffn.gate.e_score_correction_bias", ".input_scale": ".input_scale_2", }, - orig_to_new_substr=substr_map, + orig_to_new_substr=orig_to_new_substr, ) @@ -1038,7 +1039,7 @@ def get_mtp_target_hidden_states(self) -> torch.Tensor | None: return getattr(self.model, "_mtp_hidden_buffer", None) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self, skip_substrs=["mtp."]) + loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) def process_weights_after_loading(self) -> None: diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 922d1c449871..27fded04079f 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -1392,6 +1392,7 @@ def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper: }, orig_to_new_substr={ ".shared_experts.w2": ".shared_experts.down_proj", + "mtp.": None, }, ) @@ -1517,7 +1518,7 @@ def get_mtp_target_hidden_states(self) -> torch.Tensor | None: return getattr(self.model, "_mtp_hidden_buffer", None) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self, skip_substrs=["mtp."]) + loader = AutoWeightsLoader(self) loaded_params = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) self.process_weights_after_loading() return loaded_params diff --git a/vllm/models/deepseek_v4/xpu/model.py b/vllm/models/deepseek_v4/xpu/model.py index 2c8aabaaf2e6..ddf8ae175b22 100644 --- a/vllm/models/deepseek_v4/xpu/model.py +++ b/vllm/models/deepseek_v4/xpu/model.py @@ -1299,6 +1299,7 @@ def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper: }, orig_to_new_substr={ ".shared_experts.w2": ".shared_experts.down_proj", + "mtp.": None, }, ) @@ -1364,7 +1365,7 @@ def get_mtp_target_hidden_states(self) -> torch.Tensor | None: return getattr(self.model, "_mtp_hidden_buffer", None) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self, skip_substrs=["mtp."]) + loader = AutoWeightsLoader(self) loaded_params = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) self.model.finalize_mega_moe_weights() return loaded_params diff --git a/vllm/models/dots3_note/nvidia/multimodal.py b/vllm/models/dots3_note/nvidia/multimodal.py index cd4bfeae6a7a..00f50918222d 100644 --- a/vllm/models/dots3_note/nvidia/multimodal.py +++ b/vllm/models/dots3_note/nvidia/multimodal.py @@ -141,6 +141,16 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: self.language_model.make_empty_intermediate_tensors ) + orig_to_new_prefix = dict[str, None]() + if self.visual is None: + orig_to_new_prefix["visual."] = None + if self.audio_tower is None: + orig_to_new_prefix["audio_tower."] = None + if orig_to_new_prefix: + self.hf_to_vllm_mapper = self.hf_to_vllm_mapper | WeightsMapper( + orig_to_new_prefix=orig_to_new_prefix + ) + def _process_image_input( self, pixel_values: torch.Tensor, @@ -283,12 +293,7 @@ def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: return self.language_model.compute_logits(hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - skip_prefixes = [] - if self.visual is None: - skip_prefixes.append("visual.") - if self.audio_tower is None: - skip_prefixes.append("audio_tower.") - return AutoWeightsLoader(self, skip_prefixes=skip_prefixes).load_weights( + return AutoWeightsLoader(self).load_weights( weights, mapper=self.hf_to_vllm_mapper, ) diff --git a/vllm/models/inkling/amd/model.py b/vllm/models/inkling/amd/model.py index e37a54938d8c..711a33855cfc 100644 --- a/vllm/models/inkling/amd/model.py +++ b/vllm/models/inkling/amd/model.py @@ -656,8 +656,9 @@ def _iter_loadable_weights() -> Iterable[tuple[str, torch.Tensor]]: yield name, weight - loader = AutoWeightsLoader(module, skip_prefixes=["model.mtp."]) - loaded |= loader.load_weights(_iter_loadable_weights()) + mapper = WeightsMapper(orig_to_new_prefix={"model.mtp.": None}) + loader = AutoWeightsLoader(module) + loaded |= loader.load_weights(_iter_loadable_weights(), mapper=mapper) # Post-load MoE fixups (default input scales, zeroed EP-padding experts). for moe_name, moe in moe_modules.items(): diff --git a/vllm/models/inkling/nvidia/model.py b/vllm/models/inkling/nvidia/model.py index a7df6243e79d..aeb460d64a21 100644 --- a/vllm/models/inkling/nvidia/model.py +++ b/vllm/models/inkling/nvidia/model.py @@ -693,8 +693,9 @@ def _iter_loadable_weights() -> Iterable[tuple[str, torch.Tensor]]: # The release checkpoint also carries auxiliary prediction-head weights; # they are not part of the causal LM served by this implementation. - loader = AutoWeightsLoader(module, skip_prefixes=["model.mtp."]) - loaded |= loader.load_weights(_iter_loadable_weights()) + mapper = WeightsMapper(orig_to_new_prefix={"model.mtp.": None}) + loader = AutoWeightsLoader(module) + loaded |= loader.load_weights(_iter_loadable_weights(), mapper=mapper) # Post-load MoE fixups (default input scales, zeroed EP-padding experts). for moe_name, moe in moe_modules.items(): diff --git a/vllm/models/kimi_k3/nvidia/dspark_mla.py b/vllm/models/kimi_k3/nvidia/dspark_mla.py index 81715e02c921..5008015abdc8 100644 --- a/vllm/models/kimi_k3/nvidia/dspark_mla.py +++ b/vllm/models/kimi_k3/nvidia/dspark_mla.py @@ -399,9 +399,14 @@ class K3DSparkForCausalLM(nn.Module): has_own_embed_tokens = False has_own_lm_head = False draft_id_to_target_id = None - checkpoint_skip_substrs = ("confidence_head", "embed_tokens", "lm_head") - hf_to_vllm_mapper = WeightsMapper( + # confidence_head is training-only. The frozen target embedding and LM + # head are shared after this draft-specific checkpoint is loaded. + orig_to_new_substr={ + "confidence_head": None, + "embed_tokens": None, + "lm_head": None, + }, orig_to_new_prefix={"": "model."}, orig_to_new_stacked={ ".gate_proj": (".gate_up_proj", 0), @@ -477,12 +482,7 @@ def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor: return self.model.markov_head.bias(markov_embed, self.logits_processor) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - # confidence_head is training-only. The frozen target embedding and LM - # head are shared after this draft-specific checkpoint is loaded. - loader = AutoWeightsLoader( - self, - skip_substrs=list(self.checkpoint_skip_substrs), - ) + loader = AutoWeightsLoader(self) # read: 1. all weights. 2. context kv weights weights = _duplicate_context_kv_weights(weights, len(self.model.layers)) loaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) From ae256289956945c750d5b3cd13848dc734501a6d Mon Sep 17 00:00:00 2001 From: Dennis Yeh Date: Fri, 21 Aug 2026 00:21:39 +0800 Subject: [PATCH 216/839] upgrade tpu-inference to v0.27.0 (#53088) --- requirements/tpu.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements/tpu.txt b/requirements/tpu.txt index bc25cd3c23ab..ab8fa029d82f 100644 --- a/requirements/tpu.txt +++ b/requirements/tpu.txt @@ -12,4 +12,4 @@ ray[data] setuptools==78.1.0 setuptools-rust>=1.9.0 nixl==0.3.0 -tpu-inference==0.26.0 +tpu-inference==0.27.0 From 00f7f258282ad3ea400a760ba3bc4679ef4e40fe Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Fri, 21 Aug 2026 01:26:47 +0800 Subject: [PATCH 217/839] [Misc] Don't allow language-model-only used with encoder CG together (#53127) Signed-off-by: Isotr0py --- vllm/config/vllm.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 5146a362254b..4fcb3ca17d81 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1147,6 +1147,19 @@ def __post_init__(self): "connectors (PD disaggregation, KV cache offload)." ) + if ( + self.model_config is not None + and self.model_config.multimodal_config is not None + and self.model_config.multimodal_config.language_model_only + and self.compilation_config.cudagraph_mm_encoder + ): + raise ValueError( + "--language-model-only is incompatible with " + "cudagraph_mm_encoder=True, since it disables all multimodal " + "inputs and the multimodal encoder is never run. Please " + "disable one of them." + ) + self._verify_sampling_replay_config() self._verify_trace_replay_config() From 7c8b68b9ce30c45ea17063cc040b2473e152fa7d Mon Sep 17 00:00:00 2001 From: kyleliang-nv Date: Thu, 20 Aug 2026 10:45:28 -0700 Subject: [PATCH 218/839] [Bugfix][MiniMax-M3] Keep FP8 query allocation stable across CUDA graph replay (#51203) Signed-off-by: Kyle Liang --- ...st_minimax_m3_msa_cutlass_sparse_decode.py | 82 ++++++++++++++++++- vllm/models/minimax_m3/nvidia/model.py | 20 ++--- 2 files changed, 91 insertions(+), 11 deletions(-) diff --git a/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py b/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py index e2b59dc33462..f6ea80cca020 100644 --- a/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py +++ b/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py @@ -9,12 +9,14 @@ import torch from vllm import _custom_ops as ops -from vllm.config import AttentionConfig +from vllm.config import AttentionConfig, CUDAGraphMode +from vllm.forward_context import ForwardContext, override_forward_context from vllm.models.minimax_m3.common.ops.sparse_attn import ( minimax_m3_sparse_attn_decode, ) from vllm.models.minimax_m3.common.sparse_attention import ( MiniMaxM3SparseBackend, + MiniMaxM3SparseMetadata, MiniMaxM3SparseMetadataBuilder, MiniMaxM3SparseTritonImpl, select_main_backend_and_impl_cls, @@ -22,6 +24,7 @@ from vllm.models.minimax_m3.nvidia import ( sparse_attention_msa as sparse_attention_msa_module, ) +from vllm.models.minimax_m3.nvidia.model import MiniMaxM3SparseAttention from vllm.models.minimax_m3.nvidia.msa_cutlass_sparse_decode import ( MSACutlassDecodePlanCache, msa_cutlass_sparse_decode, @@ -316,6 +319,83 @@ def fake_build_plan(**kwargs): assert built_query_lens == [1, 2] +def test_query_fp8_stays_valid_when_cutlass_plan_appears_on_replay( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import vllm.envs as envs + from vllm.compilation.breakable_cudagraph import ( + BreakableCUDAGraphCapture, + eager_break_during_capture, + ) + + monkeypatch.setenv("VLLM_USE_BREAKABLE_CUDAGRAPH", "1") + envs.disable_envs_cache() + + num_tokens = 16 + layer_name = "model.layers.0.self_attn.attn" + seq_lens = torch.full((num_tokens,), 257, dtype=torch.int32, device="cuda") + block_table = torch.zeros(num_tokens, 3, dtype=torch.int32, device="cuda") + slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device="cuda") + decode = MiniMaxM3SparseMSADecodeMetadata( + seq_lens=seq_lens, + block_table=block_table, + decode_query_len=1, + msa_cutlass=None, + ) + metadata = MiniMaxM3SparseMetadata( + seq_lens=seq_lens, + max_seq_len=257, + slot_mapping=slot_mapping, + num_actual_tokens=num_tokens, + num_decodes=num_tokens, + num_decode_tokens=num_tokens, + num_prefills=0, + num_prefill_tokens=0, + decode=decode, + ) + forward_context = ForwardContext( + no_compile_layers={}, + attn_metadata={layer_name: metadata}, + slot_mapping={}, + cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE, + ) + + impl = object.__new__(MiniMaxM3SparseMSAImpl) + impl.use_cutlass_decode = True + attention = SimpleNamespace(impl=impl, q_size=HEAD_DIM) + qkv = torch.empty(num_tokens, HEAD_DIM, device="cuda") + observed_ptrs: list[int] = [] + + @eager_break_during_capture + def run_attention(query_fp8: torch.Tensor | None) -> None: + if impl.should_use_msa_decode(layer_name): + assert query_fp8 is not None + observed_ptrs.append(query_fp8.data_ptr()) + + stream = torch.cuda.Stream() + + with torch.cuda.stream(stream), override_forward_context(forward_context): + capture = BreakableCUDAGraphCapture() + with capture: + query_fp8 = MiniMaxM3SparseAttention._allocate_query_fp8(attention, qkv) + if query_fp8 is not None: + query_fp8.zero_() + run_attention(query_fp8) + qkv.zero_() + + assert capture.num_graphs == 2 + assert capture.num_eager_breaks == 1 + assert observed_ptrs == [] + + decode.msa_cutlass = object() # type: ignore[assignment] + for _ in range(3): + capture.replay() + stream.synchronize() + + assert len(observed_ptrs) == 3 + assert len(set(observed_ptrs)) == 1 + + def _make_topk( seq_lens: list[int], num_kv_heads: int, diff --git a/vllm/models/minimax_m3/nvidia/model.py b/vllm/models/minimax_m3/nvidia/model.py index a75eca2b630e..99e6d1f7aeef 100644 --- a/vllm/models/minimax_m3/nvidia/model.py +++ b/vllm/models/minimax_m3/nvidia/model.py @@ -564,6 +564,15 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), ) + def _allocate_query_fp8(self, qkv: torch.Tensor) -> torch.Tensor | None: + if not getattr(self.impl, "use_cutlass_decode", False): + return None + return torch.empty( + (qkv.shape[0], self.q_size), + dtype=torch.float8_e4m3fn, + device=qkv.device, + ) + def forward( self, positions: torch.Tensor, @@ -598,16 +607,7 @@ def forward( main_slot_mapping = fwd_slot_mapping[self.layer_name] index_slot_mapping = fwd_slot_mapping[self.indexer.index_cache.prefix] q = qkv.new_empty((num_tokens, self.q_size)) - use_msa_decode = self.impl.should_use_msa_decode(self.layer_name) - query_fp8 = ( - torch.empty( - (num_tokens, self.q_size), - dtype=torch.float8_e4m3fn, - device=qkv.device, - ) - if use_msa_decode - else None - ) + query_fp8 = self._allocate_query_fp8(qkv) # index_q matches the index-K cache dtype (e4m3 for the fp8 score path); # the fused kernel emits fp8 directly when this buffer is e4m3. index_q = qkv.new_empty( From d56bbf3995b820864a94971264cd5d2e5265261a Mon Sep 17 00:00:00 2001 From: Thanh Phan <69446444+thanhpt1110@users.noreply.github.com> Date: Fri, 21 Aug 2026 01:08:50 +0700 Subject: [PATCH 219/839] [Bugfix] Support MistralCommonBackend tokenizers in structured output (#52720) Signed-off-by: Thanh Phan --- .../test_mistral_common_tokenizer.py | 102 ++++++++++++++++++ vllm/v1/structured_output/__init__.py | 5 +- vllm/v1/structured_output/utils.py | 13 +++ 3 files changed, 118 insertions(+), 2 deletions(-) create mode 100644 tests/v1/structured_output/test_mistral_common_tokenizer.py diff --git a/tests/v1/structured_output/test_mistral_common_tokenizer.py b/tests/v1/structured_output/test_mistral_common_tokenizer.py new file mode 100644 index 000000000000..74d7cb085cc5 --- /dev/null +++ b/tests/v1/structured_output/test_mistral_common_tokenizer.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Structured output with a `MistralCommonBackend` tokenizer.""" + +import pytest +from transformers import AutoTokenizer, MistralCommonBackend + +from vllm.config import StructuredOutputsConfig, VllmConfig +from vllm.tokenizers import get_tokenizer +from vllm.tokenizers.mistral import MistralTokenizer +from vllm.v1.structured_output.backend_outlines import OutlinesBackend +from vllm.v1.structured_output.backend_types import StructuredOutputOptions +from vllm.v1.structured_output.backend_xgrammar import XgrammarBackend +from vllm.v1.structured_output.utils import ( + _reduced_vocabulary, + maybe_wrap_mistral_common_tokenizer, +) + +TOKENIZER = "mistralai/Mistral-Nemo-Instruct-2407" +JSON_SCHEMA = ( + '{"type": "object", "properties": {"x": {"type": "integer"}}, ' + '"required": ["x"], "additionalProperties": false}' +) +VALID_DOCUMENT = '{"x": 1}' + +BACKENDS = [("xgrammar", XgrammarBackend), ("outlines", OutlinesBackend)] + + +@pytest.fixture(scope="module") +def loaded_tokenizer(): + tokenizer = get_tokenizer(tokenizer_name=TOKENIZER, tokenizer_mode="hf") + assert isinstance(tokenizer, MistralCommonBackend) + return tokenizer + + +@pytest.fixture(scope="module") +def wrapped_tokenizer(loaded_tokenizer): + return maybe_wrap_mistral_common_tokenizer(loaded_tokenizer) + + +def _encode(tokenizer, text): + special = set(tokenizer.all_special_ids) + return [tid for tid in tokenizer.encode(text) if tid not in special] + + +def _backend(cls, name, tokenizer): + vllm_config = VllmConfig( + structured_outputs_config=StructuredOutputsConfig(backend=name) + ) + return cls(vllm_config, tokenizer=tokenizer, vocab_size=tokenizer.vocab_size) + + +def test_wrapping_produces_a_mistral_tokenizer(loaded_tokenizer, wrapped_tokenizer): + assert isinstance(wrapped_tokenizer, MistralTokenizer) + assert len(wrapped_tokenizer.vocab) == loaded_tokenizer.vocab_size + + +def test_other_tokenizers_are_left_alone(): + tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2") + assert maybe_wrap_mistral_common_tokenizer(tokenizer) is tokenizer + + +def test_wrapping_fixes_the_outlines_vocabulary(loaded_tokenizer, wrapped_tokenizer): + unwrapped_vocab = _reduced_vocabulary(loaded_tokenizer) + wrapped_vocab = _reduced_vocabulary(wrapped_tokenizer) + + assert unwrapped_vocab.keys() == wrapped_vocab.keys() + + mismapped = { + token for token, ids in unwrapped_vocab.items() if ids != wrapped_vocab[token] + } + assert mismapped + + assert len({tuple(wrapped_vocab[t]) for t in mismapped}) == len(mismapped) + assert len({tuple(unwrapped_vocab[t]) for t in mismapped}) < len(mismapped) + + +@pytest.mark.parametrize("name,cls", BACKENDS) +def test_backend_accepts_a_valid_document(wrapped_tokenizer, name, cls): + backend = _backend(cls, name, wrapped_tokenizer) + grammar = backend.compile_grammar(StructuredOutputOptions.JSON, JSON_SCHEMA) + + for token in _encode(wrapped_tokenizer, VALID_DOCUMENT): + assert grammar.accept_tokens("req", [token]) + + +@pytest.mark.parametrize("name,cls", BACKENDS) +def test_backend_rejects_an_invalid_token(wrapped_tokenizer, name, cls): + backend = _backend(cls, name, wrapped_tokenizer) + grammar = backend.compile_grammar(StructuredOutputOptions.JSON, JSON_SCHEMA) + + assert not grammar.accept_tokens("req", [_encode(wrapped_tokenizer, "]")[-1]]) + + +def test_xgrammar_vocab_size_matches_the_tokenizer(wrapped_tokenizer): + backend = _backend(XgrammarBackend, "xgrammar", wrapped_tokenizer) + assert backend.vocab_size == len(wrapped_tokenizer.vocab) + + +def test_xgrammar_rejects_the_unwrapped_tokenizer(loaded_tokenizer): + with pytest.raises(ValueError, match="Unsupported tokenizer type"): + _backend(XgrammarBackend, "xgrammar", loaded_tokenizer) diff --git a/vllm/v1/structured_output/__init__.py b/vllm/v1/structured_output/__init__.py index 6a2409c8a5a9..2fe1399d4d32 100644 --- a/vllm/v1/structured_output/__init__.py +++ b/vllm/v1/structured_output/__init__.py @@ -17,6 +17,7 @@ StructuredOutputGrammar, ) from vllm.v1.structured_output.backend_xgrammar import XgrammarBackend +from vllm.v1.structured_output.utils import maybe_wrap_mistral_common_tokenizer if TYPE_CHECKING: import numpy as np @@ -75,8 +76,8 @@ def __init__(self, vllm_config: VllmConfig): # of CPUs. max_workers = max(1, (multiprocessing.cpu_count() + 1) // 2) self.executor = ThreadPoolExecutor(max_workers=max_workers) - self.tokenizer = cached_tokenizer_from_config( - model_config=self.vllm_config.model_config + self.tokenizer = maybe_wrap_mistral_common_tokenizer( + cached_tokenizer_from_config(model_config=self.vllm_config.model_config) ) reasoning_parser_plugin = ( self.vllm_config.structured_outputs_config.reasoning_parser_plugin diff --git a/vllm/v1/structured_output/utils.py b/vllm/v1/structured_output/utils.py index 0629a6d2e0f6..0633c55fe687 100644 --- a/vllm/v1/structured_output/utils.py +++ b/vllm/v1/structured_output/utils.py @@ -14,6 +14,7 @@ import regex as re import torch from cachetools import LRUCache +from transformers import MistralCommonBackend import vllm.envs as envs from vllm.logger import init_logger @@ -305,6 +306,18 @@ def get_outlines_cache(): re_replacement_seq = re.compile(r"^.{0,6}�+.{0,6}$") +def maybe_wrap_mistral_common_tokenizer(tokenizer: TokenizerLike) -> TokenizerLike: + """The grammar backends cannot consume a `MistralCommonBackend` directly.""" + if not isinstance(tokenizer, MistralCommonBackend): + return tokenizer + + # Deferred: `vllm.tokenizers.mistral` pulls in a large dependency tree, and + # this module is imported by `vllm.v1.worker.gpu_model_runner`. + from vllm.tokenizers.mistral import MistralTokenizer + + return MistralTokenizer(tokenizer) + + def _reduced_vocabulary(tokenizer: TokenizerLike) -> dict[bytes, list[int]]: """Create a map from vocabulary tokens to lists of equivalent token ids. From 3b829cf176aa4ef357b7edd61eb9352419a8adbd Mon Sep 17 00:00:00 2001 From: stefankoncarevic Date: Thu, 20 Aug 2026 20:40:33 +0200 Subject: [PATCH 220/839] [CI/Build][ROCm] Keep the CUDA-only kernel tests out of the ROCm run (#53113) Signed-off-by: Stefan Koncarevic --- tests/kernels/test_fp32_router_gemm.py | 7 +++++-- tests/kernels/test_kimi_k3_gemm_rs.py | 6 +++++- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/tests/kernels/test_fp32_router_gemm.py b/tests/kernels/test_fp32_router_gemm.py index 571dfa077d32..6384cae7cdd6 100644 --- a/tests/kernels/test_fp32_router_gemm.py +++ b/tests/kernels/test_fp32_router_gemm.py @@ -16,6 +16,7 @@ import torch from vllm._custom_ops import fp32_router_gemm +from vllm.platforms import current_platform # (hidden_size, num_experts) SHAPES = [(3072, 256), (6144, 128), (6144, 256)] @@ -26,8 +27,10 @@ def _requires_sm90(): - if not torch.cuda.is_available(): - pytest.skip("CUDA not available") + # ROCm reports a CUDA-like device capability (gfx950 -> (9, 5)), which would + # pass the SM90 check below for a kernel that is only built for CUDA. + if not current_platform.is_cuda(): + pytest.skip("fp32_router_gemm is built for CUDA only") major, minor = torch.cuda.get_device_capability() if major * 10 + minor < 90: pytest.skip(f"fp32_router_gemm requires SM90+, got SM{major}{minor}") diff --git a/tests/kernels/test_kimi_k3_gemm_rs.py b/tests/kernels/test_kimi_k3_gemm_rs.py index 5ca54d2ca4de..d4d6347875a0 100644 --- a/tests/kernels/test_kimi_k3_gemm_rs.py +++ b/tests/kernels/test_kimi_k3_gemm_rs.py @@ -12,7 +12,6 @@ init_distributed_environment, initialize_model_parallel, ) -from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs import GemmRS from vllm.platforms import current_platform from vllm.utils.network_utils import get_open_port from vllm.utils.system_utils import update_environment_variables @@ -70,6 +69,11 @@ def _assert_valid_rows_close( def _worker(local_rank: int, world_size: int, master_port: int) -> None: + # This module pulls in cute_dsl, which is unavailable off CUDA, so importing + # it at module level would fail collection. Import it where it is used, as + # `kda.py` does. + from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs import GemmRS + device = torch.device("cuda", local_rank) torch.accelerator.set_device_index(device) update_environment_variables( From 4f6885fffc931ce030d55a97ffb03bbeb0934d10 Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Fri, 21 Aug 2026 03:07:02 +0800 Subject: [PATCH 221/839] [DSV4][Kernel] Fuse shared experts into MegaMoE (#53040) Signed-off-by: Canlin Guo --- tests/models/test_deepseek_v4_mega_moe.py | 232 +++++++++++++ vllm/envs.py | 7 + vllm/models/deepseek_v4/nvidia/model.py | 316 +++++++++++++++--- .../deepseek_v4/nvidia/ops/prepare_megamoe.py | 51 +++ vllm/models/kimi_k3/nvidia/model.py | 7 +- 5 files changed, 569 insertions(+), 44 deletions(-) diff --git a/tests/models/test_deepseek_v4_mega_moe.py b/tests/models/test_deepseek_v4_mega_moe.py index ee1e765db6cf..f3c70dc88a0d 100644 --- a/tests/models/test_deepseek_v4_mega_moe.py +++ b/tests/models/test_deepseek_v4_mega_moe.py @@ -13,6 +13,7 @@ from vllm.models.deepseek_v4.nvidia.model import ( DeepseekV4ForCausalLM, DeepseekV4MegaMoEExperts, + DeepseekV4MoE, make_deepseek_v4_expert_params_mapping, ) from vllm.models.deepseek_v4.nvidia.mtp import DeepSeekV4MTP @@ -168,6 +169,149 @@ def test_deepseek_v4_mega_moe_weight_loader_uses_ep_expert_ownership(): assert torch.count_nonzero(experts.w13_weight[1]) == 0 +def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights(monkeypatch): + class FakeDeepGemm: + transformed_dims: list[tuple[int, int]] = [] + scale_inputs: list[tuple[int, ...]] = [] + + @staticmethod + def get_symm_buffer_for_mega_moe(*args, num_shared_experts=0, **kwargs): + return None + + @staticmethod + def get_block_m_for_mega_moe(*args, **kwargs): + return 128 + + @staticmethod + def fp8_fp4_mega_moe( + y, + l1_weights, + l2_weights, + sym_buffer, + shared_l1_weights=None, + shared_l2_weights=None, + **kwargs, + ): + return None + + @classmethod + def transform_sf_into_required_layout(cls, sf, mn, k, *args, **kwargs): + cls.scale_inputs.append(tuple(sf.shape)) + return torch.empty((sf.shape[0], mn, k // 128), dtype=torch.int32) + + @classmethod + def transform_weights_for_mega_moe(cls, l1_weights, l2_weights): + cls.transformed_dims.append((l1_weights[0].dim(), l2_weights[0].dim())) + if l1_weights[0].dim() == 2: + return (l1_weights[0].clone(), l1_weights[1]), l2_weights + return l1_weights, l2_weights + + vllm_config = SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_batched_tokens=4), + compilation_config=SimpleNamespace(static_forward_context={}), + ) + experts = DeepseekV4MegaMoEExperts( + vllm_config, + num_experts=2, + num_local_experts=1, + experts_start_idx=0, + top_k=1, + hidden_size=128, + intermediate_size=128, + num_shared_experts=1, + ) + experts._check_runtime_supported = lambda: None + + def fp8_parameter(*shape): + return torch.nn.Parameter( + torch.empty(*shape, dtype=torch.float8_e4m3fn), requires_grad=False + ) + + def scale_parameter(*shape, dtype=torch.int32): + return torch.nn.Parameter(torch.ones(*shape, dtype=dtype), requires_grad=False) + + shared_experts = SimpleNamespace( + gate_up_proj=SimpleNamespace( + weight=fp8_parameter(256, 128), + weight_block_size=(128, 128), + weight_scale_inv=scale_parameter(2, 1, dtype=torch.float8_e8m0fnu), + ), + down_proj=SimpleNamespace( + weight=fp8_parameter(128, 128), + weight_block_size=(128, 128), + weight_scale_inv=scale_parameter(1, 1, dtype=torch.float8_e8m0fnu), + ), + ) + monkeypatch.setattr("vllm.utils.deep_gemm._import_deep_gemm", lambda: FakeDeepGemm) + + original_gate_up_ptr = shared_experts.gate_up_proj.weight.data_ptr() + experts.finalize_weights(shared_experts) + + assert FakeDeepGemm.transformed_dims == [(3, 3), (2, 2)] + assert FakeDeepGemm.scale_inputs[-2:] == [(1, 256, 4), (1, 128, 4)] + assert experts.has_fused_shared_experts + assert shared_experts.gate_up_proj.weight.data_ptr() != original_gate_up_ptr + assert ( + experts._transformed_shared_l1_weights[0].data_ptr() + == shared_experts.gate_up_proj.weight.data_ptr() + ) + assert ( + experts._transformed_shared_l2_weights[0].data_ptr() + == shared_experts.down_proj.weight.data_ptr() + ) + + +@pytest.mark.parametrize("fused", [False, True]) +def test_deepseek_v4_mega_moe_does_not_double_add_fused_shared_expert( + monkeypatch, fused +): + class FakeGate(torch.nn.Module): + tid2eid = None + e_score_correction_bias = None + + def forward(self, hidden_states): + return torch.empty(hidden_states.shape[0], 2), None + + class FakeExperts(torch.nn.Module): + has_fused_shared_experts = fused + + def forward(self, hidden_states, *args, **kwargs): + return torch.ones_like(hidden_states) + + class FakeSharedExperts(torch.nn.Module): + calls = 0 + + def forward(self, hidden_states): + self.calls += 1 + return torch.full_like(hidden_states, 2) + + moe = DeepseekV4MoE.__new__(DeepseekV4MoE) + torch.nn.Module.__init__(moe) + moe.use_mega_moe = True + moe.gate = FakeGate() + moe.experts = FakeExperts() + moe.shared_experts = FakeSharedExperts() + moe.scoring_func = "sqrtsoftplus" + moe.n_activated_experts = 1 + moe.renormalize = True + moe.hash_indices_dtype = torch.int64 + moe.routed_scaling_factor = 1.0 + moe.swiglu_limit = 10.0 + monkeypatch.setattr( + "vllm.models.deepseek_v4.nvidia.model.fused_topk_bias", + lambda **kwargs: ( + torch.ones(kwargs["hidden_states"].shape[0], 1), + torch.zeros(kwargs["hidden_states"].shape[0], 1, dtype=torch.int64), + ), + ) + + output = moe(torch.zeros(2, 128)) + + expected = 1 if fused else 3 + assert torch.all(output == expected) + assert moe.shared_experts.calls == (0 if fused else 1) + + @pytest.mark.skipif( not torch.cuda.is_available(), reason="DeepSeek V4 MegaMoE fused input staging requires CUDA.", @@ -246,6 +390,94 @@ def test_deepseek_v4_mega_moe_fused_input_staging_is_bitwise_exact(): ) +@pytest.mark.parametrize("shared_block_m", [8, 32, 96, 128, 192]) +def test_deepseek_v4_mega_moe_stages_shared_scale_tma_layout(shared_block_m): + from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8 + + device = torch.device("cuda") + num_tokens = shared_block_m + 7 + hidden_size = 256 + top_k = 8 + generator = torch.Generator(device=device) + generator.manual_seed(shared_block_m) + hidden_states = torch.randn( + num_tokens, + hidden_size, + device=device, + dtype=torch.bfloat16, + generator=generator, + ) + topk_ids = torch.randint( + 0, + 256, + (num_tokens, top_k), + device=device, + dtype=torch.int32, + generator=generator, + ) + topk_weights = torch.randn( + num_tokens, + top_k, + device=device, + dtype=torch.float32, + generator=generator, + ) + + ref_x, ref_x_sf = per_token_cast_to_fp8( + hidden_states, + use_ue8m0=True, + gran_k=32, + use_packed_ue8m0=True, + ) + aligned_block_m = ((shared_block_m + 127) // 128) * 128 + num_shared_rows = ((num_tokens + shared_block_m - 1) // shared_block_m) * ( + aligned_block_m + ) + ref_shared_x_sf = torch.zeros( + num_shared_rows, + hidden_size // 128, + dtype=torch.int32, + device=device, + ) + for token_id in range(num_tokens): + m_in_block = token_id % shared_block_m + transposed_m = ( + (m_in_block // 128) * 128 + (m_in_block % 32) * 4 + (m_in_block % 128) // 32 + ) + shared_row = token_id // shared_block_m * aligned_block_m + transposed_m + ref_shared_x_sf[shared_row].copy_(ref_x_sf[token_id]) + + fused_x = torch.empty_like(ref_x) + fused_x_sf = torch.empty_like(ref_x_sf) + fused_shared_storage = torch.full( + (hidden_size // 128, num_shared_rows), + -1, + dtype=torch.int32, + device=device, + ) + fused_shared_x_sf = fused_shared_storage.t() + fused_topk_idx = torch.empty_like(topk_ids, dtype=torch.int64) + fused_topk_weights = torch.empty_like(topk_weights) + + prepare_megamoe_inputs( + hidden_states, + topk_weights, + topk_ids, + fused_x, + fused_x_sf, + fused_topk_idx, + fused_topk_weights, + shared_x_sf=fused_shared_x_sf, + shared_block_m=shared_block_m, + ) + torch.accelerator.synchronize() + + populated = ref_shared_x_sf != 0 + assert torch.equal(fused_x.view(torch.uint8), ref_x.view(torch.uint8)) + assert torch.equal(fused_x_sf, ref_x_sf) + assert torch.equal(fused_shared_x_sf[populated], ref_shared_x_sf[populated]) + + def test_deepseek_v4_pwal_hook_finalizes_mega_moe_and_mhc_broadcast(): """The loader invokes the model-level PWAL hook for every load format, so it must finalize megamoe + mhc broadcast weights to cover dummy diff --git a/vllm/envs.py b/vllm/envs.py index a928461da210..90398ee4e0ee 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -288,6 +288,7 @@ VLLM_GC_DEBUG: str = "" VLLM_DEBUG_WORKSPACE: bool = False VLLM_DISABLE_SHARED_EXPERTS_STREAM: bool = False + VLLM_DISABLE_DSV4_MEGAMOE_SHARED_EXPERT_FUSION: bool = False VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD: int = 256 VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD: int = 1024 VLLM_COMPILE_CACHE_SAVE_FORMAT: Literal["binary", "unpacked"] = "binary" @@ -1993,6 +1994,12 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_DISABLE_SHARED_EXPERTS_STREAM": lambda: bool( int(os.getenv("VLLM_DISABLE_SHARED_EXPERTS_STREAM", "0")) ), + # Emergency rollback for the DeepSeek-V4 NVIDIA MegaMoE path. By default, + # DeepGEMM computes replicated FP8 shared experts in the same persistent + # SM100 kernel as the routed FP4 experts. + "VLLM_DISABLE_DSV4_MEGAMOE_SHARED_EXPERT_FUSION": lambda: bool( + int(os.getenv("VLLM_DISABLE_DSV4_MEGAMOE_SHARED_EXPERT_FUSION", "0")) + ), # Limits when we run shared_experts in a separate stream. # We found out that for large batch sizes, the separate stream # execution is not beneficial (most likely because of the input clone) diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 27fded04079f..be8214f21e9f 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import typing from collections.abc import Callable, Iterable +from inspect import signature from itertools import islice import regex as re @@ -18,6 +19,7 @@ ) from vllm.distributed.eplb.eplb_state import EplbLayerState from vllm.forward_context import get_forward_context, is_forward_context_available +from vllm.logger import init_logger from vllm.model_executor.kernels.mhc.tilelang import ( hc_head_fused_kernel_tilelang, mhc_fused_post_pre_tilelang, @@ -84,6 +86,8 @@ from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.worker.ubatching import dbo_current_ubatch_id +logger = init_logger(__name__) + class DeepseekV4MLP(nn.Module): def __init__( @@ -165,7 +169,7 @@ def make_deepseek_v4_expert_params_mapping( class DeepseekV4MegaMoEExperts(nn.Module): - _symm_buffer_cache: dict[tuple[int, int, int, int, int, int, int], object] = {} + _symm_buffer_cache: dict[tuple[int, int, int, int, int, int, int, int], object] = {} def __init__( self, @@ -177,6 +181,7 @@ def __init__( top_k: int, hidden_size: int, intermediate_size: int, + num_shared_experts: int = 0, prefix: str = "", num_logical_experts: int | None = None, ): @@ -190,6 +195,7 @@ def __init__( self.top_k = top_k self.hidden_size = hidden_size self.intermediate_size = intermediate_size + self.num_shared_experts = num_shared_experts self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens self.num_logical_experts = ( @@ -247,6 +253,12 @@ def __init__( self._transformed_l1_weights: tuple[torch.Tensor, torch.Tensor] | None = None self._transformed_l2_weights: tuple[torch.Tensor, torch.Tensor] | None = None + self._transformed_shared_l1_weights: ( + tuple[torch.Tensor, torch.Tensor] | None + ) = None + self._transformed_shared_l2_weights: ( + tuple[torch.Tensor, torch.Tensor] | None + ) = None # Register in the static forward context so the custom-op wrapper # can look up this module by name from within a torch.compile graph. @@ -321,45 +333,220 @@ def _check_runtime_supported(self) -> None: "to be multiples of 128." ) - def finalize_weights(self) -> None: - if self._transformed_l1_weights is not None: - return + @staticmethod + def _deep_gemm_supports_shared_experts(deep_gemm) -> bool: + """Check the Python API before touching a symmetric-memory group. - self._check_runtime_supported() - from vllm.utils.deep_gemm import _import_deep_gemm + This also gives users of an older precompiled vLLM wheel a safe serial + fallback instead of failing halfway through multi-rank buffer setup. + """ + try: + buffer_params = signature(deep_gemm.get_symm_buffer_for_mega_moe).parameters + kernel_params = signature(deep_gemm.fp8_fp4_mega_moe).parameters + except (TypeError, ValueError): + return False + return ( + hasattr(deep_gemm, "get_block_m_for_mega_moe") + and hasattr(deep_gemm, "transform_weights_for_mega_moe") + and "num_shared_experts" in buffer_params + and "shared_l1_weights" in kernel_params + and "shared_l2_weights" in kernel_params + ) - deep_gemm = _import_deep_gemm() + def _finalize_shared_expert_weights( + self, deep_gemm, shared_experts: DeepseekV4MLP + ) -> None: + gate_up = shared_experts.gate_up_proj + down = shared_experts.down_proj + gate_up_weight = gate_up.weight.data + gate_up_scale = gate_up.weight_scale_inv.data + down_weight = down.weight.data + down_scale = down.weight_scale_inv.data + + # MegaMoE's shared FP8 MMA consumes a 1x32 scale for every weight row, + # while the checkpoint uses coarser block-FP8 scales (usually + # 128x128). Build a dedicated, numerically equivalent scale view before + # the generic linear post-load hook replaces the raw checkpoint scales + # with its 128x128 DeepGEMM layout. + checkpoint_scale_dtypes = (torch.float8_e8m0fnu, torch.uint8) + if ( + gate_up_scale.dtype in checkpoint_scale_dtypes + and down_scale.dtype in checkpoint_scale_dtypes + ): + gate_up_scale = self._prepare_shared_expert_scale( + deep_gemm, + gate_up, + gate_up_scale, + gate_up_weight.shape[0], + gate_up_weight.shape[1], + ) + down_scale = self._prepare_shared_expert_scale( + deep_gemm, + down, + down_scale, + down_weight.shape[0], + down_weight.shape[1], + ) + + if gate_up_scale is None or down_scale is None: + self.num_shared_experts = 0 + return - w13_scale = deep_gemm.transform_sf_into_required_layout( - self._ue8m0_uint8_to_float(self.w13_weight_scale.data).contiguous(), - 2 * self.intermediate_size, + shared_intermediate_size = self.intermediate_size * self.num_shared_experts + expected_gate_up_shape = ( + 2 * shared_intermediate_size, self.hidden_size, - (1, 32), - self.num_local_experts, ) - w2_scale = deep_gemm.transform_sf_into_required_layout( - self._ue8m0_uint8_to_float(self.w2_weight_scale.data).contiguous(), - self.hidden_size, - self.intermediate_size, - (1, 32), - self.num_local_experts, + expected_down_shape = (self.hidden_size, shared_intermediate_size) + if ( + gate_up_weight.dtype != torch.float8_e4m3fn + or down_weight.dtype != torch.float8_e4m3fn + or gate_up_scale.dtype != torch.int32 + or down_scale.dtype != torch.int32 + or tuple(gate_up_weight.shape) != expected_gate_up_shape + or tuple(down_weight.shape) != expected_down_shape + ): + logger.warning( + "Disabling native MegaMoE shared-expert fusion for %s: expected " + "replicated block-FP8 weights with gate_up=%s, down=%s, and " + "DeepGEMM int32 scales; got gate_up=%s/%s/%s and down=%s/%s/%s.", + self.prefix, + expected_gate_up_shape, + expected_down_shape, + tuple(gate_up_weight.shape), + gate_up_weight.dtype, + gate_up_scale.dtype, + tuple(down_weight.shape), + down_weight.dtype, + down_scale.dtype, + ) + self.num_shared_experts = 0 + return + + transformed_l1, transformed_l2 = deep_gemm.transform_weights_for_mega_moe( + (gate_up_weight, gate_up_scale), + (down_weight, down_scale), + ) + # L1 interleaving allocates a full copy. Re-home the loader Parameter on + # that storage so the original 2*intermediate*hidden FP8 tensor can be + # released instead of adding roughly 0.7 GiB per rank on DSV4-Flash. + # The generic linear post-load hook may still repack the serial scales, + # but this shared MLP is never called after native fusion is enabled. + gate_up.weight.data = transformed_l1[0] + self._transformed_shared_l1_weights = ( + gate_up.weight.data, + transformed_l1[1], ) - self._transformed_l1_weights, self._transformed_l2_weights = ( - deep_gemm.transform_weights_for_mega_moe( - (self.w13_weight.data.view(torch.int8).contiguous(), w13_scale), - (self.w2_weight.data.view(torch.int8).contiguous(), w2_scale), + self._transformed_shared_l2_weights = transformed_l2 + + def _prepare_shared_expert_scale( + self, + deep_gemm, + linear: nn.Module, + scale: torch.Tensor, + mn: int, + k: int, + ) -> torch.Tensor | None: + block_size = getattr(linear, "weight_block_size", None) + if block_size is None or len(block_size) != 2: + logger.warning( + "Disabling native MegaMoE shared-expert fusion for %s: " + "shared FP8 weight block size is unavailable.", + self.prefix, + ) + return None + + block_m, block_k = block_size + expected_shape = ( + (mn + block_m - 1) // block_m, + (k + block_k - 1) // block_k, + ) + if block_k % 32 != 0 or tuple(scale.shape) != expected_shape: + logger.warning( + "Disabling native MegaMoE shared-expert fusion for %s: " + "cannot convert shared scale shape %s with block size %s " + "to MegaMoE's 1x32 layout for weight (%d, %d).", + self.prefix, + tuple(scale.shape), + tuple(block_size), + mn, + k, ) + return None + + scale_fp32 = self._ue8m0_uint8_to_float(scale.view(torch.uint8)) + scale_1x32 = ( + scale_fp32.repeat_interleave(block_m, dim=0) + .repeat_interleave(block_k // 32, dim=1)[:mn, : k // 32] + .contiguous() ) - # Drop the original loader-side parameters: the MegaMoE kernels only - # consume the transformed views above. transform_weights_for_mega_moe - # allocates a fresh tensor for the L1 weight (see _interleave_l1_weights) - # and fresh SF tensors for L1/L2; the L2 weight is the only tensor that - # aliases the original storage, and _transformed_l2_weights still holds - # it, so the storage stays live after we drop the Parameter. - self.w13_weight = None - self.w13_weight_scale = None - self.w2_weight = None - self.w2_weight_scale = None + # The grouped API is used with a singleton dimension to request the + # MN-major, TMA-aligned packed-UE8M0 strides, then squeezed back to the + # 2D layout required for a shared expert. + return deep_gemm.transform_sf_into_required_layout( + scale_1x32.unsqueeze(0), + mn, + k, + (1, 32), + 1, + ).squeeze(0) + + def finalize_weights(self, shared_experts: DeepseekV4MLP | None = None) -> None: + from vllm.utils.deep_gemm import _import_deep_gemm + + deep_gemm = _import_deep_gemm() + + if self._transformed_l1_weights is None: + self._check_runtime_supported() + w13_scale = deep_gemm.transform_sf_into_required_layout( + self._ue8m0_uint8_to_float(self.w13_weight_scale.data).contiguous(), + 2 * self.intermediate_size, + self.hidden_size, + (1, 32), + self.num_local_experts, + ) + w2_scale = deep_gemm.transform_sf_into_required_layout( + self._ue8m0_uint8_to_float(self.w2_weight_scale.data).contiguous(), + self.hidden_size, + self.intermediate_size, + (1, 32), + self.num_local_experts, + ) + self._transformed_l1_weights, self._transformed_l2_weights = ( + deep_gemm.transform_weights_for_mega_moe( + (self.w13_weight.data.view(torch.int8).contiguous(), w13_scale), + (self.w2_weight.data.view(torch.int8).contiguous(), w2_scale), + ) + ) + # Drop the original loader-side parameters: the MegaMoE kernels only + # consume the transformed views above. transform_weights_for_mega_moe + # allocates a fresh tensor for the L1 weight (see + # _interleave_l1_weights) and fresh SF tensors for L1/L2; the L2 + # weight is the only tensor that aliases the original storage, and + # _transformed_l2_weights still holds it, so the storage stays live + # after we drop the Parameter. + self.w13_weight = None + self.w13_weight_scale = None + self.w2_weight = None + self.w2_weight_scale = None + + if shared_experts is None or self.num_shared_experts == 0: + return + if self._transformed_shared_l1_weights is not None: + return + if not self._deep_gemm_supports_shared_experts(deep_gemm): + logger.warning_once( + "Disabling native MegaMoE shared-expert fusion because the " + "installed DeepGEMM Python API is older than the vLLM " + "source. Rebuild the vendored _deep_gemm_C extension to enable it.", + ) + self.num_shared_experts = 0 + return + self._finalize_shared_expert_weights(deep_gemm, shared_experts) + + @property + def has_fused_shared_experts(self) -> bool: + return self._transformed_shared_l1_weights is not None def get_symm_buffer(self): from vllm.utils.deep_gemm import _import_deep_gemm @@ -376,6 +563,7 @@ def get_symm_buffer(self): self.top_k, self.hidden_size, self.intermediate_size, + self.num_shared_experts if self.has_fused_shared_experts else 0, ) symm_buffer = self._symm_buffer_cache.get(key) if symm_buffer is None: @@ -386,6 +574,9 @@ def get_symm_buffer(self): self.top_k, self.hidden_size, self.intermediate_size, + num_shared_experts=( + self.num_shared_experts if self.has_fused_shared_experts else 0 + ), ) self._symm_buffer_cache[key] = symm_buffer return symm_buffer @@ -492,6 +683,19 @@ def forward( else None, ) + shared_x_sf = None + shared_block_m = None + if self.has_fused_shared_experts: + shared_x_sf = symm_buffer.shared_l1_acts_sf + shared_block_m = deep_gemm.get_block_m_for_mega_moe( + get_ep_group().world_size, + self.num_experts, + symm_buffer.num_max_tokens_per_rank, + num_tokens, + self.top_k, + "fp8xfp4", + ) + prepare_megamoe_inputs( hidden_states, topk_weights, @@ -501,18 +705,32 @@ def forward( symm_buffer.topk_idx[:num_tokens], symm_buffer.topk_weights[:num_tokens], is_padding=is_padding, + shared_x_sf=shared_x_sf, + shared_block_m=shared_block_m, ) assert self._transformed_l1_weights is not None assert self._transformed_l2_weights is not None - deep_gemm.fp8_fp4_mega_moe( - y, - self._transformed_l1_weights, - self._transformed_l2_weights, - symm_buffer, - activation_clamp=activation_clamp, - fast_math=fast_math, - ) + if self.has_fused_shared_experts: + deep_gemm.fp8_fp4_mega_moe( + y, + self._transformed_l1_weights, + self._transformed_l2_weights, + symm_buffer, + shared_l1_weights=self._transformed_shared_l1_weights, + shared_l2_weights=self._transformed_shared_l2_weights, + activation_clamp=activation_clamp, + fast_math=fast_math, + ) + else: + deep_gemm.fp8_fp4_mega_moe( + y, + self._transformed_l1_weights, + self._transformed_l2_weights, + symm_buffer, + activation_clamp=activation_clamp, + fast_math=fast_math, + ) return y @@ -645,6 +863,16 @@ def _init_mega_moe_experts( self.experts_start_idx = self.physical_expert_start self.experts_end_idx = self.physical_expert_end + # Native DeepGEMM fusion requires each EP rank to own the complete + # shared MLP. Sequence parallel replicates those weights while sharding + # tokens. TP=1 is also naturally replicated. With PP+TP the shared MLP + # remains tensor-sharded, so retain the serial path. + fuse_shared_experts = bool( + self.shared_experts is not None + and not envs.VLLM_DISABLE_DSV4_MEGAMOE_SHARED_EXPERT_FUSION + and (self.use_sequence_parallel or self.tp_size == 1) + ) + self.experts = DeepseekV4MegaMoEExperts( vllm_config, num_experts=self.n_physical_experts, @@ -654,6 +882,7 @@ def _init_mega_moe_experts( top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, + num_shared_experts=(self.n_shared_experts if fuse_shared_experts else 0), prefix=f"{prefix}.experts", ) @@ -739,7 +968,10 @@ def forward( activation_clamp=activation_clamp, ) - if self.shared_experts is not None: + if ( + self.shared_experts is not None + and not self.experts.has_fused_shared_experts + ): shared_output = self.shared_experts(hidden_states) final_hidden_states += shared_output @@ -759,7 +991,7 @@ def _forward_fused_moe( def finalize_mega_moe_weights(self) -> None: if self.use_mega_moe: - self.experts.finalize_weights() + self.experts.finalize_weights(self.shared_experts) def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]: diff --git a/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py b/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py index dac86be6edb1..d254dad9d6a1 100644 --- a/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py +++ b/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py @@ -17,6 +17,7 @@ def _prepare_megamoe_inputs_kernel( hidden_states, x_fp8, x_sf, + shared_x_sf, topk_ids, topk_weights, is_padding, @@ -28,6 +29,8 @@ def _prepare_megamoe_inputs_kernel( x_stride_k: tl.constexpr, x_sf_stride_m: tl.constexpr, x_sf_stride_k: tl.constexpr, + shared_x_sf_stride_m: tl.constexpr, + shared_x_sf_stride_k: tl.constexpr, topk_ids_stride_m: tl.constexpr, topk_ids_stride_k: tl.constexpr, topk_weights_stride_m: tl.constexpr, @@ -42,6 +45,7 @@ def _prepare_megamoe_inputs_kernel( BLOCK_K: tl.constexpr, GROUP_K: tl.constexpr, BLOCK_TOPK: tl.constexpr, + SHARED_BLOCK_M: tl.constexpr, ) -> None: token_id = tl.program_id(0) k_block_id = tl.program_id(1) @@ -84,6 +88,25 @@ def _prepare_megamoe_inputs_kernel( packed_scale, ) + # DeepGEMM's SM100 shared-expert TMA loads require the activation scales + # in an MN-major layout whose row permutation depends on the MegaMoE + # scheduler's runtime BLOCK_M. Write that view while the packed UE8M0 scale + # is already resident, avoiding another kernel and temporary tensor. + if shared_x_sf is not None: + m_block_id = token_id // SHARED_BLOCK_M + m_in_block = token_id % SHARED_BLOCK_M + aligned_block_m: tl.constexpr = triton.cdiv(SHARED_BLOCK_M, 128) * 128 + transposed_m = ( + (m_in_block // 128) * 128 + (m_in_block % 32) * 4 + (m_in_block % 128) // 32 + ) + shared_row = m_block_id * aligned_block_m + transposed_m + tl.store( + shared_x_sf + + shared_row * shared_x_sf_stride_m + + k_block_id * shared_x_sf_stride_k, + packed_scale, + ) + if k_block_id == 0: topk_offsets = tl.arange(0, BLOCK_TOPK) topk_mask = topk_offsets < top_k @@ -131,6 +154,8 @@ def prepare_megamoe_inputs( topk_idx_out: torch.Tensor, topk_weights_out: torch.Tensor, is_padding: torch.Tensor | None = None, + shared_x_sf: torch.Tensor | None = None, + shared_block_m: int | None = None, ) -> None: num_tokens, hidden_size = hidden_states.shape if num_tokens == 0: @@ -146,6 +171,28 @@ def prepare_megamoe_inputs( "DeepSeek V4 MegaMoE input staging requires topk_weights and " "topk_ids to have the same shape." ) + if (shared_x_sf is None) != (shared_block_m is None): + raise ValueError( + "DeepSeek V4 MegaMoE shared input staging requires both " + "shared_x_sf and shared_block_m." + ) + if shared_x_sf is not None: + assert shared_block_m is not None + if shared_block_m <= 0: + raise ValueError("MegaMoE shared_block_m must be positive.") + expected_sf_k = hidden_size // 128 + if shared_x_sf.ndim != 2 or shared_x_sf.shape[1] != expected_sf_k: + raise ValueError( + "MegaMoE shared_x_sf must have shape " + f"(*, {expected_sf_k}), got {tuple(shared_x_sf.shape)}." + ) + aligned_block_m = triton.cdiv(shared_block_m, 128) * 128 + required_rows = triton.cdiv(num_tokens, shared_block_m) * aligned_block_m + if shared_x_sf.shape[0] < required_rows: + raise ValueError( + "MegaMoE shared_x_sf has insufficient rows: requires " + f"{required_rows}, got {shared_x_sf.shape[0]}." + ) block_k = 128 grid = (num_tokens, triton.cdiv(hidden_size, block_k)) @@ -155,6 +202,7 @@ def prepare_megamoe_inputs( hidden_states, x_fp8, x_sf, + shared_x_sf, topk_ids, topk_weights, is_padding, @@ -166,6 +214,8 @@ def prepare_megamoe_inputs( x_fp8.stride(1), x_sf.stride(0), x_sf.stride(1), + shared_x_sf.stride(0) if shared_x_sf is not None else 0, + shared_x_sf.stride(1) if shared_x_sf is not None else 0, topk_ids.stride(0), topk_ids.stride(1), topk_weights.stride(0), @@ -180,5 +230,6 @@ def prepare_megamoe_inputs( BLOCK_K=block_k, GROUP_K=32, BLOCK_TOPK=block_topk, + SHARED_BLOCK_M=shared_block_m or 1, num_warps=4, ) diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index 472272a35741..d2ceafd5147b 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -93,7 +93,10 @@ sp_reduce_scatter, sp_shard, ) -from vllm.models.deepseek_v4.nvidia.model import DeepseekV4MegaMoEExperts +from vllm.models.deepseek_v4.nvidia.model import ( + DeepseekV4MegaMoEExperts, + DeepseekV4MLP, +) from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention from vllm.models.kimi_k3.nvidia.latent_moe_runner import ( @@ -354,7 +357,7 @@ def synchronize_first_launch(self) -> None: torch.distributed.barrier(group=ep_group.cpu_group) self._synchronized_ep_groups.add(key) - def finalize_weights(self) -> None: + def finalize_weights(self, shared_experts: DeepseekV4MLP | None = None) -> None: if self._transformed_l1_weights is not None: return From bfb6c134997aace3e801c9ae3251728bd5312003 Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Thu, 20 Aug 2026 16:26:24 -0400 Subject: [PATCH 222/839] [Bugfix][MoE] Tune FlashInfer experts to scheduler token limit (#52989) Signed-off-by: mgoin Co-authored-by: OpenAI Codex --- tests/kernels/moe/test_flashinfer_moe.py | 2 ++ .../layers/fused_moe/experts/flashinfer_cutlass_moe.py | 2 ++ .../layers/fused_moe/experts/trtllm_bf16_moe.py | 1 + .../layers/fused_moe/experts/trtllm_lora_moe.py | 2 ++ .../layers/fused_moe/experts/trtllm_mxint4_moe.py | 2 ++ .../layers/quantization/utils/flashinfer_mxint4_moe.py | 4 +++- 6 files changed, 12 insertions(+), 1 deletion(-) diff --git a/tests/kernels/moe/test_flashinfer_moe.py b/tests/kernels/moe/test_flashinfer_moe.py index a98844d5f8b5..ece11c371c6e 100644 --- a/tests/kernels/moe/test_flashinfer_moe.py +++ b/tests/kernels/moe/test_flashinfer_moe.py @@ -72,6 +72,7 @@ def test_flashinfer_swigluoai_params_are_forwarded(activation, monkeypatch): moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), in_dtype=torch.bfloat16, routing_method=RoutingMethodType.TopK, + max_num_tokens=16_384, ) quant_config = FusedMoEQuantConfig.make( gemm1_alpha=1.702, @@ -110,6 +111,7 @@ def fake_flashinfer_cutlass_fused_moe(**kwargs): assert experts._supports_activation(activation) assert call_args["activation_type"] == ActivationType.Swiglu + assert call_args["tune_max_num_tokens"] == 16_384 for name, value in ( ("swiglu_alpha", 1.702), ("swiglu_beta", 1.0), diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py index 1614e41b3303..543968b67159 100644 --- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py @@ -13,6 +13,7 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( TopKWeightAndReduceNoOP, ) +from vllm.model_executor.layers.fused_moe.utils import fi_moe_largest_bucket from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( activation_to_flashinfer_type, ) @@ -388,6 +389,7 @@ def apply( use_deepseek_fp8_block_scale=self.use_deepseek_fp8_block_scale, use_mxfp8_act_scaling=use_mxfp8_act_scaling, use_w4_group_scaling=use_w4_group_scaling, + tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config), ) def moe_sum(self, input: torch.Tensor, output: torch.Tensor) -> None: diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py index 4372634be8d1..031c8f6cf9da 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py @@ -219,6 +219,7 @@ def apply( weight_layout=WeightLayout.BlockMajorK, do_finalize=True, activation_type=activation_to_flashinfer_int(activation), + tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config), ) # FlashInfer's BF16 routed wrapper does not expose an output= argument. output.copy_(result[0] if isinstance(result, list) else result) diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py index 1a947756ce1c..2ecda740c0b1 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py @@ -39,6 +39,7 @@ TopKWeightAndReduceNoOP, ) from vllm.model_executor.layers.fused_moe.utils import ( + fi_moe_largest_bucket, trtllm_moe_pack_topk_ids_weights, ) from vllm.platforms import current_platform @@ -549,6 +550,7 @@ def invoke_routed_moe( routing_method_type=self.routing_method_type, do_finalize=do_finalize, output=output if do_finalize else None, + tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config), ) if not do_finalize: # [gemm2_output, expert_weights, expanded_idx, gemm1_activation] diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py index a9e725f2965c..12b0dce12ecd 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py @@ -11,6 +11,7 @@ FusedMoEQuantConfig, RoutingMethodType, ) +from vllm.model_executor.layers.fused_moe.utils import fi_moe_largest_bucket from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, kInt4Static32, @@ -169,6 +170,7 @@ def apply( e_score_correction_bias=e_score_correction_bias, routing_method_type=self.routing_method, routing_replay_out=routing_replay_out, + tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config), ) self._maybe_dispatch_routing_replay( routing_replay_out, num_tokens=hidden_states.shape[0] diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_mxint4_moe.py b/vllm/model_executor/layers/quantization/utils/flashinfer_mxint4_moe.py index 1b5320615f4a..af15b206f047 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_mxint4_moe.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_mxint4_moe.py @@ -191,6 +191,7 @@ def flashinfer_trtllm_mxint4_moe( e_score_correction_bias: torch.Tensor | None = None, routing_method_type: int | None = None, routing_replay_out: torch.Tensor | None = None, + tune_max_num_tokens: int = 8192, ) -> torch.Tensor: """ Apply FlashInfer TensorRT-LLM MxInt4 MoE kernel. @@ -211,6 +212,7 @@ def flashinfer_trtllm_mxint4_moe( topk_group: Top-k within groups (default: None -> 0) e_score_correction_bias: Optional routing bias. dtype: bfloat16 routing_method_type: FlashInfer RoutingMethodType enum value + tune_max_num_tokens: Maximum token count covered by autotuning. Returns: Output tensor from MoE layer. dtype: same as x (bfloat16) @@ -262,7 +264,7 @@ def flashinfer_trtllm_mxint4_moe( enable_pdl=None, do_finalize=True, output=None, - tune_max_num_tokens=8192, + tune_max_num_tokens=tune_max_num_tokens, routing_replay_out=routing_replay_out, ) if isinstance(out, (tuple, list)): From 54ba80d9611d3a9cce7ba348e8719e24b5d9211c Mon Sep 17 00:00:00 2001 From: Andrey Talman Date: Thu, 20 Aug 2026 18:19:40 -0400 Subject: [PATCH 223/839] [CI][Docker] Pin manylinux2_28-builder:cuda13.0 to the release/2.13 image (#52994) Signed-off-by: Andrey Talman --- .buildkite/image_build/image_build_torch_nightly.sh | 2 +- .buildkite/release-pipeline.yaml | 6 +++--- docker/Dockerfile | 2 +- docker/versions.json | 2 +- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/.buildkite/image_build/image_build_torch_nightly.sh b/.buildkite/image_build/image_build_torch_nightly.sh index ff60378102dd..8aa313d0e6d4 100755 --- a/.buildkite/image_build/image_build_torch_nightly.sh +++ b/.buildkite/image_build/image_build_torch_nightly.sh @@ -48,7 +48,7 @@ echo "Image not found, proceeding with build..." # --- CUDA 13.0 for nightly builds --- # Nightly CI uses CUDA 13.0 while regular CI stays on CUDA 12.9 NIGHTLY_CUDA_VERSION="13.0.2" -NIGHTLY_BUILD_BASE_IMAGE="pytorch/manylinux2_28-builder:cuda13.0" +NIGHTLY_BUILD_BASE_IMAGE="pytorch/manylinux2_28-builder:cuda13.0-78e737ad29420ffc4800e677c51e2a852caf8359" NIGHTLY_FINAL_BASE_IMAGE="nvidia/cuda:${NIGHTLY_CUDA_VERSION}-base-ubuntu22.04" echo "--- :docker: Building torch nightly image (CUDA ${NIGHTLY_CUDA_VERSION})" diff --git a/.buildkite/release-pipeline.yaml b/.buildkite/release-pipeline.yaml index 1dd8f0146c31..15e50d3a75fc 100644 --- a/.buildkite/release-pipeline.yaml +++ b/.buildkite/release-pipeline.yaml @@ -47,7 +47,7 @@ steps: agents: queue: cpu_queue_release commands: - - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda13.0 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ." + - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda13.0-78e737ad29420ffc4800e677c51e2a852caf8359 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ." - "mkdir artifacts" - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'" - "bash .buildkite/scripts/upload-nightly-wheels.sh" @@ -234,7 +234,7 @@ steps: --build-arg CUDA_VERSION=13.0.2 \ --build-arg torch_cuda_arch_list="${CUDA_ARCH_X86}" \ --build-arg INSTALL_KV_CONNECTORS=true \ - --build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda13.0 \ + --build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda13.0-78e737ad29420ffc4800e677c51e2a852caf8359 \ --target vllm-openai \ --progress plain \ -f docker/Dockerfile . @@ -334,7 +334,7 @@ steps: --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \ --build-arg torch_cuda_arch_list="${CUDA_ARCH_X86}" \ --build-arg INSTALL_KV_CONNECTORS=true \ - --build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda13.0 \ + --build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda13.0-78e737ad29420ffc4800e677c51e2a852caf8359 \ --target vllm-openai \ --progress plain \ -f docker/Dockerfile . diff --git a/docker/Dockerfile b/docker/Dockerfile index 550cc6f4dca7..543f4b12813c 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -38,7 +38,7 @@ ARG NCCL_VERSION=2.30.7 # docker build --build-arg BUILD_BASE_IMAGE=registry.acme.org/mirror/pytorch/manylinux2_28-builder:cuda13.0 # Build wheels against the same glibc floor as PyTorch's published wheels. -ARG BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda13.0 +ARG BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda13.0-78e737ad29420ffc4800e677c51e2a852caf8359 # Using cuda base image with minimal dependencies necessary for JIT compilation (FlashInfer, DeepGEMM, EP kernels) ARG FINAL_BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-base-ubuntu${UBUNTU_VERSION} diff --git a/docker/versions.json b/docker/versions.json index 7ac2bada5e44..936114523375 100644 --- a/docker/versions.json +++ b/docker/versions.json @@ -14,7 +14,7 @@ "default": "2.30.7" }, "BUILD_BASE_IMAGE": { - "default": "pytorch/manylinux2_28-builder:cuda13.0" + "default": "pytorch/manylinux2_28-builder:cuda13.0-78e737ad29420ffc4800e677c51e2a852caf8359" }, "FINAL_BASE_IMAGE": { "default": "nvidia/cuda:13.0.3-base-ubuntu24.04" From 2f41c894e320a07a46a360ee72149dd84d6f4bec Mon Sep 17 00:00:00 2001 From: MKQuantum <121593006+MKQuantum@users.noreply.github.com> Date: Thu, 20 Aug 2026 18:38:53 -0400 Subject: [PATCH 224/839] Fix seed loss when batch contains unseeded requests (#51866) Signed-off-by: MKQuantum <121593006+MKQuantum@users.noreply.github.com> --- vllm/v1/sample/ops/topk_topp_sampler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py index e6e0e054c06c..03716cb3c135 100644 --- a/vllm/v1/sample/ops/topk_topp_sampler.py +++ b/vllm/v1/sample/ops/topk_topp_sampler.py @@ -200,7 +200,7 @@ def forward_cpu( elif self.logprobs_mode == "processed_logprobs": logits_to_return = logits.log_softmax(dim=-1, dtype=torch.float32) - if len(generators) != logits.shape[0] and not self.use_fp64_gumbel: + if not generators and not self.use_fp64_gumbel: return compiled_random_sample(logits), logits_to_return probs = logits.softmax(dim=-1, dtype=torch.float32) From 7cfb97e33791a348cd5d7b622cca521d82d8399f Mon Sep 17 00:00:00 2001 From: Matthew Kotila <7692737+matthewkotila@users.noreply.github.com> Date: Thu, 20 Aug 2026 15:42:49 -0700 Subject: [PATCH 225/839] [Frontend][Core][Spec Decode] Per-request acceptance stats in OpenAI API responses (#48915) Signed-off-by: Matthew Kotila <7692737+matthewkotila@users.noreply.github.com> Co-authored-by: Reed Meyerson <31574681+reed-meyerson@users.noreply.github.com> Co-authored-by: Claude Opus 4.8 (1M context) --- docs/features/per_request_metrics.md | 9 ++ docs/features/speculative_decoding/README.md | 1 + .../acceptance_metrics.md | 100 ++++++++++++ .../engine-core-client/src/protocol/output.rs | 6 + .../engine-core-client/src/tests/client.rs | 1 + .../completion/test_completion_error.py | 110 ++++++++++++- tests/test_config.py | 14 ++ tests/v1/core/test_scheduler.py | 122 ++++++++++++-- tests/v1/core/utils.py | 5 + .../v1/spec_decode/test_request_acceptance.py | 150 ++++++++++++++++++ vllm/config/observability.py | 12 ++ vllm/config/vllm.py | 9 ++ vllm/engine/arg_utils.py | 8 + vllm/entrypoints/generate/base/serving.py | 31 +++- .../openai/chat_completion/protocol.py | 6 +- .../openai/chat_completion/serving.py | 53 ++++--- .../entrypoints/openai/completion/protocol.py | 6 +- vllm/entrypoints/openai/completion/serving.py | 64 ++++---- vllm/entrypoints/openai/engine/protocol.py | 26 ++- vllm/outputs.py | 8 +- vllm/v1/core/sched/scheduler.py | 32 +++- vllm/v1/engine/__init__.py | 10 +- vllm/v1/engine/output_processor.py | 8 + vllm/v1/metrics/stats.py | 76 +++++++++ vllm/v1/request.py | 7 +- 25 files changed, 798 insertions(+), 76 deletions(-) create mode 100644 docs/features/speculative_decoding/acceptance_metrics.md create mode 100644 tests/v1/spec_decode/test_request_acceptance.py diff --git a/docs/features/per_request_metrics.md b/docs/features/per_request_metrics.md index 9bc64d2b86dc..25211ecfc39b 100644 --- a/docs/features/per_request_metrics.md +++ b/docs/features/per_request_metrics.md @@ -125,3 +125,12 @@ a single prompt's generation. The `metrics` response field provides per-request values for a single request. The `/metrics` Prometheus endpoint exposes server-level histograms (e.g. `vllm:time_to_first_token_seconds`) that aggregate across all requests. + +## Speculative Decoding Acceptance + +When speculative decoding is enabled, per-request acceptance metrics +(mean acceptance length and the accepted-draft-length distribution) can be +returned via `--per-request-spec-decode-metrics`. They share this `metrics` +object as `metrics.speculative_decoding`, and — like the timing fields — are +reported only for single-sequence (`n == 1`) requests. See +[Per-Request Acceptance Metrics](speculative_decoding/acceptance_metrics.md). diff --git a/docs/features/speculative_decoding/README.md b/docs/features/speculative_decoding/README.md index 7858fff2d4a1..f4417c7fdefe 100644 --- a/docs/features/speculative_decoding/README.md +++ b/docs/features/speculative_decoding/README.md @@ -19,6 +19,7 @@ vLLM supports a variety of methods of speculative decoding. Model-based methods - [Custom Proposer Backend (Experimental)](#custom-proposer-backend-experimental) - [Dynamic Speculative Decoding](dynamic_speculative_decoding.md) - [Adaptive Verification](adaptive_verification.md) +- [Per-Request Acceptance Metrics](acceptance_metrics.md) ## Method Selection at a Glance diff --git a/docs/features/speculative_decoding/acceptance_metrics.md b/docs/features/speculative_decoding/acceptance_metrics.md new file mode 100644 index 000000000000..9e31f920edcf --- /dev/null +++ b/docs/features/speculative_decoding/acceptance_metrics.md @@ -0,0 +1,100 @@ +# Per-Request Acceptance Metrics + +When speculative decoding is enabled, vLLM can report per-request acceptance +metrics in the response, under `metrics.speculative_decoding`. This lets a +client compute the mean acceptance length and the accepted-draft-length +distribution for an individual request, as a complement to the server-aggregated +spec-decode metrics exposed at `/metrics`. + +!!! warning "Experimental" + `metrics.speculative_decoding` is experimental and its shape may change in a + future release. Pin to a vLLM version if you depend on it. + +## Enabling + +Start the server with `--per-request-spec-decode-metrics` set to `summary` or +`detailed` (default `none`): + +```bash +vllm serve \ + --speculative-config '{"method": "ngram", "num_speculative_tokens": 3, "prompt_lookup_min": 1, "prompt_lookup_max": 3}' \ + --per-request-spec-decode-metrics summary +``` + +| Level | Behavior | +| --- | --- | +| `none` (default) | No collection; responses are unchanged. | +| `summary` | Acceptance metrics per request. | +| `detailed` | `summary` plus ordered per-step arrays. | + +Collection is gated at the source: with `none`, nothing is accumulated. + +## Response Format + +Acceptance metrics share the top-level `metrics` object with the timing +[per-request metrics](../per_request_metrics.md) — `metrics.speculative_decoding` +sits alongside the timing fields. Like timing, they describe a single generation +stream, so they are reported only for single-sequence requests and are `null` +for `n > 1`. + +A `summary` response's `metrics` looks like: + +```json +{ + "choices": [ ... ], + "usage": { ... }, + "metrics": { + "speculative_decoding": { + "mean_acceptance_length": 1.2325581395348837, + "draft_acceptance_rate": 0.07751937984496124, + "acceptance_histogram": [39, 1, 0, 3], + "num_spec_steps": 43, + "num_accepted_draft_tokens": 10, + "num_draft_tokens": 129, + "num_spec_tokens": 3 + } + } +} +``` + +| Field | Description | +| --- | --- | +| `mean_acceptance_length` | Mean tokens emitted per verification step, including the bonus token: `1 + num_accepted_draft_tokens / num_spec_steps`. Ranges from `1.0` (nothing accepted) to `num_spec_tokens + 1`. | +| `draft_acceptance_rate` | Fraction of proposed draft tokens accepted: `num_accepted_draft_tokens / num_draft_tokens`. | +| `acceptance_histogram` | Dense list of length `num_spec_tokens + 1`; index `j` is the number of steps that accepted exactly `j` draft tokens. Excludes the always-accepted bonus token. | +| `num_spec_steps` | Number of verification steps for this request (the sum of the histogram). | +| `num_accepted_draft_tokens` | Total accepted draft tokens, excluding bonus tokens. | +| `num_draft_tokens` | Total proposed draft tokens, after subtracting drafts invalidated by structured-output constraints. | +| `num_spec_tokens` | Configured `num_speculative_tokens` (`k`), i.e. the maximum draft length per step. | + +With `detailed`, two ordered arrays are added, one entry per verification step: + +| Field | Description | +| --- | --- | +| `per_step_accepted` | Accepted draft count at each step. | +| `per_step_drafted` | Proposed draft count at each step. Records the effective proposal length per step, so variable-length drafting (e.g. adaptive speculation) is represented without a schema change. | + +`metrics.speculative_decoding` is present whenever `--per-request-spec-decode-metrics` +is `summary`/`detailed`, speculative decoding is enabled, and `n == 1` (with an +all-zero histogram if the request drafted nothing). It is `null` otherwise. + +## Streaming + +In streaming responses, `metrics` (including `speculative_decoding`) rides the +final usage chunk, which is only emitted when usage reporting is enabled — set +`stream_options.include_usage: true` or start the server with +`--enable-force-include-usage`. + +## Relationship to Prometheus metrics + +The per-request fields are the individual-request counterpart of the +server-aggregated spec-decode counters at `/metrics`. Summed across the +single-sequence requests that report them, they reconcile with the aggregate +counters (which also count `n > 1` requests, so the totals match only for +all-`n == 1` workloads): + +| Per-request field (summed) | Prometheus counter | +| --- | --- | +| `num_spec_steps` | `vllm:spec_decode_num_drafts_total` | +| `num_draft_tokens` | `vllm:spec_decode_num_draft_tokens_total` | +| `num_accepted_draft_tokens` | `vllm:spec_decode_num_accepted_tokens_total` | diff --git a/rust/src/engine-core-client/src/protocol/output.rs b/rust/src/engine-core-client/src/protocol/output.rs index a9695bd1cf11..cc7541eae1b2 100644 --- a/rust/src/engine-core-client/src/protocol/output.rs +++ b/rust/src/engine-core-client/src/protocol/output.rs @@ -125,6 +125,11 @@ pub struct EngineCoreOutput { pub mm_cache_miss_hashes: Option>, #[serde(default)] pub new_sampling_mask: Option, + /// Per-request speculative-decoding acceptance metrics, set on the final + /// output when `--per-request-spec-decode-metrics` is enabled. Opaque here; + /// the Rust frontend does not yet surface it in responses. + #[serde(default)] + pub spec_decode_metrics: Option, } impl EngineCoreOutput { @@ -443,6 +448,7 @@ mod tests { num_nans_in_logits: 0, mm_cache_miss_hashes: None, new_sampling_mask: None, + spec_decode_metrics: None, }, ], scheduler_stats: None, diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index fa5d8e0db855..6f4e53ff1b41 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -2743,6 +2743,7 @@ fn python_msgpack_fixtures_match_rust_encoding() { num_nans_in_logits: 0, mm_cache_miss_hashes: None, new_sampling_mask: None, + spec_decode_metrics: None, }, ], scheduler_stats: None, diff --git a/tests/entrypoints/openai/completion/test_completion_error.py b/tests/entrypoints/openai/completion/test_completion_error.py index 79fd6206d610..56a2ba66f1ad 100644 --- a/tests/entrypoints/openai/completion/test_completion_error.py +++ b/tests/entrypoints/openai/completion/test_completion_error.py @@ -23,7 +23,7 @@ from vllm.renderers.online_renderer import OnlineRenderer from vllm.tokenizers.registry import cached_tokenizer_from_config from vllm.v1.engine.async_llm import AsyncLLM -from vllm.v1.metrics.stats import RequestStateStats +from vllm.v1.metrics.stats import RequestSpecDecodeMetrics, RequestStateStats MODEL_NAME = "openai-community/gpt2" MODEL_NAME_SHORT = "gpt2" @@ -196,6 +196,114 @@ def test_completion_per_request_metrics_suppressed_for_multiple_prompts(): assert response.metrics is None +def _spec_decode_metrics() -> RequestSpecDecodeMetrics: + # Two verify steps: accept 3 drafts, then 1 -> histogram [0, 1, 0, 1]. + m = RequestSpecDecodeMetrics.new(num_spec_tokens=3) + m.observe(num_draft_tokens=3, num_accepted=3) + m.observe(num_draft_tokens=3, num_accepted=1) + return m + + +def _make_spec_decode_request_output( + num_seqs: int = 1, with_metrics: bool = True +) -> RequestOutput: + outputs = [ + CompletionOutput( + index=i, + text="Hello", + token_ids=[100, 101], + cumulative_logprob=None, + logprobs=None, + finish_reason="stop", + spec_decode_metrics=_spec_decode_metrics() if with_metrics else None, + ) + for i in range(num_seqs) + ] + return RequestOutput( + request_id="test-id", + prompt="Test prompt", + prompt_token_ids=[1, 2, 3], + prompt_logprobs=None, + outputs=outputs, + finished=True, + metrics=None, + ) + + +def _completion_response(serving, request, request_output): + return serving.request_output_to_completion_response( + [request_output], + request, + "cmpl-test-id", + 0, + MODEL_NAME, + None, + RequestResponseMetadata(request_id="cmpl-test-id"), + ) + + +def test_completion_spec_decode_metrics_present_for_single_sequence(): + # Timing off, but the sequence carries acceptance metrics -> the metrics + # object is created just to hold metrics.speculative_decoding. + serving = _build_minimal_metrics_serving_completion( + enable_per_request_metrics=False + ) + response = _completion_response( + serving, + CompletionRequest(model=MODEL_NAME, prompt="Test prompt", max_tokens=10), + _make_spec_decode_request_output(num_seqs=1), + ) + assert response.metrics is not None + assert response.metrics.time_to_first_token_ms is None # timing not requested + spec = response.metrics.speculative_decoding + assert spec is not None + assert spec.acceptance_histogram == [0, 1, 0, 1] # dense, index j + assert spec.num_spec_steps == 2 + assert spec.num_spec_tokens == 3 + assert spec.mean_acceptance_length == pytest.approx(3.0) # 1 + (3 + 1) / 2 + + +def test_completion_spec_decode_metrics_suppressed_for_n_gt_1(): + # Per-request metrics can't be attributed to one of the n sequences. + serving = _build_minimal_metrics_serving_completion( + enable_per_request_metrics=False + ) + response = _completion_response( + serving, + CompletionRequest(model=MODEL_NAME, prompt="Test prompt", n=2, max_tokens=10), + _make_spec_decode_request_output(num_seqs=2), + ) + assert response.metrics is None + + +def test_completion_spec_decode_metrics_absent_when_not_collected(): + # Flag off -> the sequence carries no acceptance metrics -> no metrics object. + serving = _build_minimal_metrics_serving_completion( + enable_per_request_metrics=False + ) + response = _completion_response( + serving, + CompletionRequest(model=MODEL_NAME, prompt="Test prompt", max_tokens=10), + _make_spec_decode_request_output(num_seqs=1, with_metrics=False), + ) + assert response.metrics is None + + +def test_completion_metrics_carries_both_timing_and_spec_decode(): + serving = _build_minimal_metrics_serving_completion(enable_per_request_metrics=True) + request_output = _make_spec_decode_request_output(num_seqs=1) + request_output.metrics = _PER_REQUEST_STATS # timing source + response = _completion_response( + serving, + CompletionRequest(model=MODEL_NAME, prompt="Test prompt", max_tokens=10), + request_output, + ) + assert response.metrics is not None + assert response.metrics.time_to_first_token_ms == pytest.approx(500.0) + assert response.metrics.speculative_decoding is not None + assert response.metrics.speculative_decoding.num_spec_steps == 2 + + @pytest.mark.asyncio async def test_completion_error_non_stream(): """test finish_reason='error' returns 500 InternalServerError (non-streaming)""" diff --git a/tests/test_config.py b/tests/test_config.py index ee3e0c96b58d..0b36258039fa 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -20,6 +20,7 @@ CompilationConfig, KernelConfig, ModelConfig, + ObservabilityConfig, ParallelConfig, PoolerConfig, SchedulerConfig, @@ -85,6 +86,19 @@ def test_kda_recoverssm_derivation_is_revalidated(): VllmConfig.validate_mamba_cached_kernel(config) +def test_per_request_spec_decode_metrics_requires_spec_decode(): + # The flag only makes sense with speculative decoding configured; enabling + # it without --speculative-config should fail fast rather than silently + # produce no metrics. + for level in ("summary", "detailed"): + with pytest.raises(ValueError, match="speculative"): + VllmConfig( + observability_config=ObservabilityConfig( + per_request_spec_decode_metrics=level + ) + ) + + def test_compile_config_repr_succeeds(): # setup: VllmBackend mutates the config object config = VllmConfig() diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index c0932d975c18..920823baeb8d 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -1306,29 +1306,42 @@ def test_draft_slots_budgeted_per_scheduled_request(tmp_path, monkeypatch): # Note - these test cases mirror some of those in test_rejection_sampler.py @pytest.mark.parametrize( - "spec_tokens,output_tokens,expected", + "spec_tokens,output_tokens,expected,expected_per_req", [ - ([[1, 2, 3]], [[1, 2, 3, 4]], (1, 3, 3, [1, 1, 1])), # perfect match - ([[1, 2, 3]], [[1, 5]], (1, 3, 1, [1, 0, 0])), # early mismatch - ([[1, 2], [3]], [[1, 2, 5], [3, 4]], (2, 3, 3, [2, 1])), # multiple sequences - ([[1]], [[1, 2]], (1, 1, 1, [1])), # single token sequence - ([[]], [[5]], (0, 0, 0, [0])), # empty sequence + ([[1, 2, 3]], [[1, 2, 3, 4]], (1, 3, 3, [1, 1, 1]), [[0, 0, 0, 1]]), # perfect + ([[1, 2, 3]], [[1, 5]], (1, 3, 1, [1, 0, 0]), [[0, 1, 0, 0]]), # early mismatch + ( + [[1, 2], [3]], + [[1, 2, 5], [3, 4]], + (2, 3, 3, [2, 1]), + [[0, 0, 1], [0, 1, 0]], + ), # multiple sequences + ([[1]], [[1, 2]], (1, 1, 1, [1]), [[0, 1]]), # single token sequence + ([[]], [[5]], (0, 0, 0, [0]), [[0, 0]]), # empty sequence -> empty accumulator ( [[1, 2, 3], [4, 5, 6]], [[1, 2, 7], [4, 8]], (2, 6, 3, [2, 1, 0]), + [[0, 0, 1, 0], [0, 1, 0, 0]], ), # multiple mismatches ], ) -def test_schedule_spec_decoding_stats(spec_tokens, output_tokens, expected): +def test_schedule_spec_decoding_stats( + spec_tokens, output_tokens, expected, expected_per_req +): """Test scheduling behavior with speculative decoding. This test verifies that: 1. Speculated tokens get scheduled correctly - 2. Spec decoding stats properly count number of draft and accepted tokens + 2. The aggregate SpecDecodingStats count draft and accepted tokens + 3. The per-request accumulator (enabled via per_request_spec_decode_metrics) + buckets the same acceptance by accepted draft count (j) """ num_spec_tokens = max(1, max(len(t) for t in spec_tokens)) - scheduler = create_scheduler(num_speculative_tokens=num_spec_tokens) + scheduler = create_scheduler( + num_speculative_tokens=num_spec_tokens, + per_request_spec_decode_metrics="summary", + ) requests = create_requests(num_requests=len(spec_tokens), num_tokens=1) req_ids = [] req_to_index = {} @@ -1414,6 +1427,97 @@ def test_schedule_spec_decoding_stats(spec_tokens, output_tokens, expected): assert stats.num_accepted_tokens == expected[2] assert stats.num_accepted_tokens_per_pos == expected[3] + # Per-request accumulator: the same acceptance, bucketed by accepted draft + # count (j) on each request rather than summed across the batch. The + # accumulator is created eagerly on add_request, so every request has one + # (an empty histogram when it drafted nothing). + for i, req_id in enumerate(req_ids): + payload = scheduler.requests[req_id].spec_decode_metrics.to_dict() + assert payload["acceptance_histogram"] == expected_per_req[i] + assert payload["num_draft_tokens"] == len(spec_tokens[i]) + assert "per_step_accepted" not in payload # summary level + + +def _run_spec_verify_steps(scheduler, rounds, num_invalid_per_round=None): + """Drive prefill + one draft/verify step per round for a single request. + + ``rounds`` is a list of ``(spec_token_ids, output_token_ids)`` -- one verify + step each. ``num_invalid_per_round`` optionally injects the grammar-invalid + draft count structured-output decoding would set on the verify output. + Returns the request. + """ + [req] = create_requests(num_requests=1, num_tokens=1) + scheduler.add_request(req) + rid = req.request_id + req_to_index = {rid: 0} + + def _mk_output(sampled): + return ModelRunnerOutput( + req_ids=[rid], + req_id_to_index=req_to_index, + sampled_token_ids=sampled, + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ) + + # Prefill: sample one token. + scheduler.update_from_output(scheduler.schedule(), _mk_output([[0]])) + for i, (spec, out) in enumerate(rounds): + scheduler.update_draft_token_ids(DraftTokenIds([rid], [spec])) + output = scheduler.schedule() + if num_invalid_per_round is not None and num_invalid_per_round[i]: + output.num_invalid_spec_tokens = {rid: num_invalid_per_round[i]} + scheduler.update_from_output(output, _mk_output([out])) + return req + + +def test_per_request_spec_decode_detailed_records_per_step(): + scheduler = create_scheduler( + num_speculative_tokens=3, + per_request_spec_decode_metrics="detailed", + ) + # Three verify steps for one request; num_accepted = len(output) - 1, so the + # outputs below accept 3, 1, then 0 drafts across the steps. + req = _run_spec_verify_steps( + scheduler, + [ + ([1, 2, 3], [1, 2, 3, 4]), # accept 3 + ([5, 6, 7], [5, 8]), # accept 1 + ([9, 10, 11], [12]), # accept 0 + ], + ) + payload = scheduler.requests[req.request_id].spec_decode_metrics.to_dict() + assert payload["per_step_accepted"] == [3, 1, 0] + assert payload["per_step_drafted"] == [3, 3, 3] + assert payload["num_spec_steps"] == 3 + + +def test_per_request_spec_decode_subtracts_invalid_drafts(): + # Grammar-invalidated drafts (num_invalid_spec_tokens, set by structured + # output) are excluded from the proposed count, mirroring the aggregate. + scheduler = create_scheduler( + num_speculative_tokens=3, + per_request_spec_decode_metrics="summary", + ) + # One verify step: 3 drafted, 1 grammar-invalid, output accepts 2. + req = _run_spec_verify_steps( + scheduler, + [([1, 2, 3], [1, 2, 5])], + num_invalid_per_round=[1], + ) + payload = scheduler.requests[req.request_id].spec_decode_metrics.to_dict() + assert payload["num_draft_tokens"] == 2 # 3 drafted - 1 invalid + assert payload["num_accepted_draft_tokens"] == 2 # len([1,2,5]) - 1 + assert payload["acceptance_histogram"] == [0, 0, 1, 0] + + +def test_per_request_spec_decode_acceptance_disabled_by_default(): + scheduler = create_scheduler(num_speculative_tokens=3) + assert scheduler.spec_decode_metrics_level == "none" + req = _run_spec_verify_steps(scheduler, [([1, 2, 3], [1, 2, 3, 4])]) + assert scheduler.requests[req.request_id].spec_decode_metrics is None + def test_spec_decoding_stats_empty_output(): """Test that spec decoding stats handle empty output tokens gracefully. diff --git a/tests/v1/core/utils.py b/tests/v1/core/utils.py index 0baa68af2589..96352324e968 100644 --- a/tests/v1/core/utils.py +++ b/tests/v1/core/utils.py @@ -10,6 +10,7 @@ ECTransferConfig, KVTransferConfig, ModelConfig, + ObservabilityConfig, ParallelConfig, SchedulerConfig, SpeculativeConfig, @@ -71,6 +72,7 @@ def create_scheduler( ec_role: str | None = None, use_v2_model_runner: bool | None = None, kv_cache_spec: KVCacheSpec | None = None, + per_request_spec_decode_metrics: str = "none", ) -> Scheduler | AsyncScheduler: """Create scheduler under test. @@ -175,6 +177,9 @@ def create_scheduler( kv_transfer_config=kv_transfer_config, speculative_config=speculative_config, ec_transfer_config=ec_transfer_config, + observability_config=ObservabilityConfig( + per_request_spec_decode_metrics=per_request_spec_decode_metrics, + ), ) if kv_cache_spec is None: kv_cache_spec = FullAttentionSpec( diff --git a/tests/v1/spec_decode/test_request_acceptance.py b/tests/v1/spec_decode/test_request_acceptance.py new file mode 100644 index 000000000000..1a00308abdf8 --- /dev/null +++ b/tests/v1/spec_decode/test_request_acceptance.py @@ -0,0 +1,150 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for per-request speculative-decoding metrics accumulation. + +These cover the pure histogram/per-step math (no GPU / no model): the engine-core +accumulator ``RequestSpecDecodeMetrics`` and its ``to_dict`` payload surfaced in the +response as ``metrics.speculative_decoding``. +""" + +import msgspec +import pytest + +from vllm.outputs import CompletionOutput +from vllm.v1.engine import EngineCoreOutput +from vllm.v1.metrics.stats import RequestSpecDecodeMetrics + + +def _metrics(pairs, num_spec_tokens=3, detailed=False): + s = RequestSpecDecodeMetrics.new(num_spec_tokens) + for k, j in pairs: + s.observe(num_draft_tokens=k, num_accepted=j, detailed=detailed) + return s + + +def test_new_allocates_dense_histogram_of_k_plus_one(): + s = RequestSpecDecodeMetrics.new(num_spec_tokens=3) + assert s.num_spec_tokens == 3 + assert s.histogram == [0, 0, 0, 0] + assert s.num_draft_tokens == 0 + assert s.per_step_accepted == [] + + +def test_observe_buckets_by_accepted_draft_count(): + s = _metrics([(3, 0), (3, 3), (3, 2), (3, 3), (3, 1)]) + # j=0 ->1, j=1 ->1, j=2 ->1, j=3 ->2 + assert s.histogram == [1, 1, 1, 2] + assert s.num_draft_tokens == 15 + # summary level does not record the ordered per-step arrays + assert s.per_step_accepted == [] + assert s.per_step_drafted == [] + + +def test_observe_detailed_records_ordered_per_step_arrays(): + # Distinct step count (3), max draft length (k=4), and per-step drafted + # counts (4, 2, 4) so no accidental "everything is 3" pattern is implied. + s = _metrics([(4, 3), (2, 2), (4, 0)], num_spec_tokens=4, detailed=True) + assert s.per_step_accepted == [3, 2, 0] + assert s.per_step_drafted == [4, 2, 4] + # histogram (indexed by accepted j, length k+1=5) is still maintained + assert s.histogram == [1, 0, 1, 1, 0] + + +def test_to_dict_summary_omits_per_step_arrays(): + d = _metrics([(3, 0), (3, 3), (3, 2), (3, 3), (3, 1)]).to_dict() + assert d == { + "mean_acceptance_length": pytest.approx(1 + 9 / 5), # j+1 + "draft_acceptance_rate": pytest.approx(9 / 15), + "acceptance_histogram": [1, 1, 1, 2], # dense, index j = step count + "num_spec_steps": 5, + "num_accepted_draft_tokens": 9, + "num_draft_tokens": 15, + "num_spec_tokens": 3, + } + + +def test_to_dict_detailed_appends_per_step_arrays(): + d = _metrics([(3, 3), (3, 2), (3, 0)], detailed=True).to_dict() + assert d["per_step_accepted"] == [3, 2, 0] + assert d["per_step_drafted"] == [3, 3, 3] + # summary fields still present in detailed mode + assert d["num_spec_steps"] == 3 + assert d["num_accepted_draft_tokens"] == 5 + + +def test_to_dict_histogram_is_dense_list_indexed_by_j(): + # length k+1, index j holds the step count that accepted exactly j drafts + d = _metrics([(3, 0), (3, 0), (3, 3)]).to_dict() + assert d["acceptance_histogram"] == [2, 0, 0, 1] + + +def test_all_rejected_gives_mean_one_and_rate_zero(): + d = _metrics([(2, 0), (2, 0), (2, 0), (2, 0)], num_spec_tokens=2).to_dict() + assert d["num_spec_steps"] == 4 + assert d["num_accepted_draft_tokens"] == 0 + assert d["acceptance_histogram"] == [4, 0, 0] + assert d["mean_acceptance_length"] == pytest.approx(1.0) + assert d["draft_acceptance_rate"] == pytest.approx(0.0) + + +def test_empty_metrics_do_not_divide_by_zero(): + d = RequestSpecDecodeMetrics.new(3).to_dict() + assert d["num_spec_steps"] == 0 + assert d["num_draft_tokens"] == 0 + assert d["draft_acceptance_rate"] == 0.0 + assert d["mean_acceptance_length"] == 1.0 + assert "per_step_accepted" not in d + + +def test_observe_records_proposed_and_accepted_independently(): + # observe() takes proposed and accepted as independent inputs: the histogram + # is keyed by accepted, num_draft_tokens sums the proposed as given. (The + # grammar-invalidated-draft subtraction happens in the scheduler before + # observe() -- see test_per_request_spec_decode_subtracts_invalid_drafts.) + s = RequestSpecDecodeMetrics.new(num_spec_tokens=3) + s.observe(num_draft_tokens=2, num_accepted=1) + s.observe(num_draft_tokens=3, num_accepted=1) + d = s.to_dict() + assert d["acceptance_histogram"] == [0, 2, 0, 0] # both steps accepted 1 + assert d["num_draft_tokens"] == 5 # proposed summed independently: 2 + 3 + + +def test_engine_core_output_round_trips_spec_decode_metrics(): + # The accumulator rides EngineCoreOutput (msgspec, array_like) to the + # frontend; verify it serializes (incl. per-step arrays) and is omitted + # when absent. ``new_token_ids`` is EngineCoreOutput's required "tokens + # generated this step" field -- a dummy value here since we only exercise + # spec_decode_metrics. + metrics = _metrics([(3, 0), (3, 3), (3, 2)], detailed=True) + out = EngineCoreOutput( + request_id="r1", new_token_ids=[1, 2], spec_decode_metrics=metrics + ) + decoder = msgspec.msgpack.Decoder(EngineCoreOutput) + decoded = decoder.decode(msgspec.msgpack.encode(out)) + assert decoded.spec_decode_metrics.histogram == [1, 0, 1, 1] + assert decoded.spec_decode_metrics.num_draft_tokens == 9 + assert decoded.spec_decode_metrics.per_step_accepted == [0, 3, 2] + + without = EngineCoreOutput(request_id="r2", new_token_ids=[1]) + decoded_without = decoder.decode(msgspec.msgpack.encode(without)) + assert decoded_without.spec_decode_metrics is None + + +def _completion_output(**kwargs): + return CompletionOutput( + index=0, + text="", + token_ids=[], + cumulative_logprob=None, + logprobs=None, + **kwargs, + ) + + +def test_completion_output_carries_spec_decode_metrics(): + # Metrics are per output sequence, so they ride the CompletionOutput + # (=> choices[i]), not the request-level RequestOutput. + metrics = _metrics([(3, 3), (3, 1)]) + out = _completion_output(spec_decode_metrics=metrics) + assert out.spec_decode_metrics is metrics + assert _completion_output().spec_decode_metrics is None diff --git a/vllm/config/observability.py b/vllm/config/observability.py index 093ed2f684d1..0394245f3cf2 100644 --- a/vllm/config/observability.py +++ b/vllm/config/observability.py @@ -45,6 +45,18 @@ def show_hidden_metrics(self) -> bool: Note that collecting detailed timing information for each request can be expensive.""" + per_request_spec_decode_metrics: Literal["none", "summary", "detailed"] = "none" + """Include per-request speculative-decoding acceptance metrics in the + response under `metrics.speculative_decoding`. `none` disables; `summary` adds mean + acceptance length, draft acceptance rate, and the step-by-draft-length + histogram; `detailed` additionally records the ordered per-step + accepted/proposed arrays (one entry per verify step). Only reported for + single-sequence requests (`n == 1`), mirroring the timing metrics. No effect + unless speculative decoding is enabled. Independent of `--disable-log-stats`. + This is the per-request response-body counterpart of the aggregate + `vllm:spec_decode_*` Prometheus metrics. The response field is experimental + and its shape may change in a future release.""" + kv_cache_metrics: bool = False """Enable KV cache residency metrics (lifetime, idle time, reuse gaps). Uses sampling to minimize overhead. diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 4fcb3ca17d81..5ad162ae5f5c 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1331,6 +1331,15 @@ def __post_init__(self): ) self.model_config.disable_cascade_attn = True + if ( + self.observability_config.per_request_spec_decode_metrics != "none" + and self.speculative_config is None + ): + raise ValueError( + "--per-request-spec-decode-metrics requires speculative decoding " + "to be enabled (via --speculative-config)." + ) + if ( self.model_config is not None and self.model_config.multimodal_config is not None diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 035595d5a8df..293e67be5d24 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -656,6 +656,9 @@ class EngineArgs: collect_detailed_traces: list[DetailedTraceModules] | None = ( ObservabilityConfig.collect_detailed_traces ) + per_request_spec_decode_metrics: Literal["none", "summary", "detailed"] = ( + ObservabilityConfig.per_request_spec_decode_metrics + ) kv_cache_metrics: bool = ObservabilityConfig.kv_cache_metrics kv_cache_metrics_sample: float = get_field( ObservabilityConfig, "kv_cache_metrics_sample" @@ -1476,6 +1479,10 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: "--collect-detailed-traces", **observability_kwargs["collect_detailed_traces"], ) + observability_group.add_argument( + "--per-request-spec-decode-metrics", + **observability_kwargs["per_request_spec_decode_metrics"], + ) observability_group.add_argument( "--kv-cache-metrics", **observability_kwargs["kv_cache_metrics"] ) @@ -1934,6 +1941,7 @@ def create_observability_config(self) -> ObservabilityConfig: show_hidden_metrics_for_version=self.show_hidden_metrics_for_version, otlp_traces_endpoint=self.otlp_traces_endpoint, collect_detailed_traces=self.collect_detailed_traces, + per_request_spec_decode_metrics=self.per_request_spec_decode_metrics, kv_cache_metrics=self.kv_cache_metrics, kv_cache_metrics_sample=self.kv_cache_metrics_sample, cudagraph_metrics=self.cudagraph_metrics, diff --git a/vllm/entrypoints/generate/base/serving.py b/vllm/entrypoints/generate/base/serving.py index 4f27b3bc350d..d6ae0a209065 100644 --- a/vllm/entrypoints/generate/base/serving.py +++ b/vllm/entrypoints/generate/base/serving.py @@ -5,7 +5,7 @@ from collections.abc import Awaitable, Mapping from dataclasses import dataclass, field from http import HTTPStatus -from typing import ClassVar, Generic, TypeVar +from typing import TYPE_CHECKING, ClassVar, Generic, TypeVar from fastapi import Request from pydantic import ConfigDict @@ -18,7 +18,8 @@ from vllm.entrypoints.openai.engine.protocol import ( ErrorResponse, GenerationError, - PerRequestTimingMetrics, + PerRequestMetrics, + SpeculativeDecodingMetrics, ) from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.entrypoints.openai.responses.protocol import ResponsesRequest @@ -37,6 +38,9 @@ ) from vllm.v1.metrics.stats import RequestStateStats +if TYPE_CHECKING: + from vllm.outputs import RequestOutput + logger = init_logger(__name__) RequestT = TypeVar("RequestT", bound=AnyRequest) @@ -48,7 +52,7 @@ def build_per_request_timing_metrics( metrics: RequestStateStats | None, num_generation_tokens: int, -) -> PerRequestTimingMetrics: +) -> PerRequestMetrics: """Build per-request timing metrics from ``RequestStateStats``. ``generation_time_ms`` is the decode interval only (first output token to @@ -60,7 +64,7 @@ def build_per_request_timing_metrics( unavailable. """ if metrics is None: - return PerRequestTimingMetrics() + return PerRequestMetrics() queued_ts = metrics.queued_ts scheduled_ts = metrics.scheduled_ts @@ -91,7 +95,7 @@ def build_per_request_timing_metrics( if inference_time_ms > 0: tokens_per_second = num_generation_tokens / inference_time_ms * 1000 - return PerRequestTimingMetrics( + return PerRequestMetrics( time_to_first_token_ms=time_to_first_token_ms, generation_time_ms=generation_time_ms, queue_time_ms=queue_time_ms, @@ -100,6 +104,23 @@ def build_per_request_timing_metrics( ) +def build_spec_decoding_metrics( + final_res: "RequestOutput | None", +) -> SpeculativeDecodingMetrics | None: + """Build per-request spec-decode acceptance metrics from the single output + sequence, or ``None`` when unavailable (metrics disabled, or no sequence + yet). + + Only meaningful for single-sequence requests; callers suppress it for n>1. + """ + if final_res is None or not final_res.outputs: + return None + metrics = final_res.outputs[0].spec_decode_metrics + if metrics is None: + return None + return SpeculativeDecodingMetrics(**metrics.to_dict()) + + @dataclass(kw_only=True) class ServeContext(Generic[RequestT]): request: RequestT diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index e8bb4e567826..f3a3b398d6d9 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -29,7 +29,7 @@ FunctionCall, FunctionDefinition, OpenAIBaseModel, - PerRequestTimingMetrics, + PerRequestMetrics, StopParam, StreamOptions, ToolCall, @@ -146,7 +146,7 @@ class ChatCompletionResponse(OpenAIBaseModel): ec_transfer_params: dict[str, Any] | None = Field( default=None, description="ECTransfer parameters." ) - metrics: PerRequestTimingMetrics | None = None + metrics: PerRequestMetrics | None = None class ChatCompletionResponseStreamChoice(OpenAIBaseModel): @@ -178,7 +178,7 @@ class ChatCompletionStreamResponse(OpenAIBaseModel): # Rendered prompt text from chat templating (only set when # ``return_prompt_text=True`` on the request); only sent on the first chunk. prompt_text: str | None = None - metrics: PerRequestTimingMetrics | None = None + metrics: PerRequestMetrics | None = None class ChatCompletionToolsParam(OpenAIBaseModel): diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index 28ec1e8a1a59..13932aac2d78 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -20,6 +20,7 @@ GenerateBaseServing, GenerationError, build_per_request_timing_metrics, + build_spec_decoding_metrics, clamp_prompt_logprobs, format_token_id_placeholder, ) @@ -40,7 +41,7 @@ DeltaMessage, ErrorResponse, FunctionCall, - PerRequestTimingMetrics, + PerRequestMetrics, PromptTokenUsageInfo, RequestResponseMetadata, ToolCall, @@ -819,16 +820,21 @@ async def chat_completion_stream_generator( # only emitted when usage reporting is enabled (i.e. # ``stream_options.include_usage=true`` or # ``--enable-force-include-usage``). - stream_per_request_metrics: PerRequestTimingMetrics | None = None - if ( - self.enable_per_request_metrics - # See note in chat_completion_full_generator: suppress for n>1. - and (request.n or 1) == 1 - ): - last_metrics = last_res.metrics if last_res is not None else None - stream_per_request_metrics = build_per_request_timing_metrics( - last_metrics, completion_tokens - ) + stream_per_request_metrics: PerRequestMetrics | None = None + # See note in chat_completion_full_generator: suppress for n>1. + if (request.n or 1) == 1: + if self.enable_per_request_metrics: + last_metrics = ( + last_res.metrics if last_res is not None else None + ) + stream_per_request_metrics = build_per_request_timing_metrics( + last_metrics, completion_tokens + ) + spec_stats = build_spec_decoding_metrics(last_res) + if spec_stats is not None: + if stream_per_request_metrics is None: + stream_per_request_metrics = PerRequestMetrics() + stream_per_request_metrics.speculative_decoding = spec_stats final_usage_chunk = ChatCompletionStreamResponse( id=request_id, @@ -1118,17 +1124,20 @@ async def chat_completion_full_generator( request_metadata.final_usage_info = usage - per_request_metrics: PerRequestTimingMetrics | None = None - if ( - self.enable_per_request_metrics - # Timing metrics describe a single generation stream. For n>1 the - # returned stats belong to only one of the n sequences, so they - # cannot be accurately attributed to the request; suppress instead. - and (request.n or 1) == 1 - ): - per_request_metrics = build_per_request_timing_metrics( - final_res.metrics, num_generated_tokens - ) + per_request_metrics: PerRequestMetrics | None = None + # Per-request metrics (timing + spec-decode acceptance) describe a single + # generation stream. For n>1 the stats belong to only one of the n + # sequences, so they cannot be attributed to the request; suppress. + if (request.n or 1) == 1: + if self.enable_per_request_metrics: + per_request_metrics = build_per_request_timing_metrics( + final_res.metrics, num_generated_tokens + ) + spec_stats = build_spec_decoding_metrics(final_res) + if spec_stats is not None: + if per_request_metrics is None: + per_request_metrics = PerRequestMetrics() + per_request_metrics.speculative_decoding = spec_stats # ``final_res.prompt`` is the rendered chat-templated prompt text prompt_text = final_res.prompt if request.return_prompt_text else None diff --git a/vllm/entrypoints/openai/completion/protocol.py b/vllm/entrypoints/openai/completion/protocol.py index 2fc6ce1cd9cf..4b9ec1827e9d 100644 --- a/vllm/entrypoints/openai/completion/protocol.py +++ b/vllm/entrypoints/openai/completion/protocol.py @@ -13,7 +13,7 @@ from vllm.entrypoints.openai.engine.protocol import ( AnyResponseFormat, OpenAIBaseModel, - PerRequestTimingMetrics, + PerRequestMetrics, StopParam, StreamOptions, UsageInfo, @@ -660,7 +660,7 @@ class CompletionResponse(OpenAIBaseModel): ec_transfer_params: dict[str, Any] | None = Field( default=None, description="ECTransfer parameters." ) - metrics: PerRequestTimingMetrics | None = None + metrics: PerRequestMetrics | None = None class CompletionResponseStreamChoice(OpenAIBaseModel): @@ -692,4 +692,4 @@ class CompletionStreamResponse(OpenAIBaseModel): # Set only on the final chunk of a stream to mirror non-streaming responses # without the per-chunk serialization overhead. system_fingerprint: str | None = None - metrics: PerRequestTimingMetrics | None = None + metrics: PerRequestMetrics | None = None diff --git a/vllm/entrypoints/openai/completion/serving.py b/vllm/entrypoints/openai/completion/serving.py index 37f55f4116c2..93a2b7d3aa6a 100644 --- a/vllm/entrypoints/openai/completion/serving.py +++ b/vllm/entrypoints/openai/completion/serving.py @@ -14,6 +14,7 @@ GenerateBaseServing, GenerationError, build_per_request_timing_metrics, + build_spec_decoding_metrics, clamp_prompt_logprobs, format_token_id_placeholder, ) @@ -27,7 +28,7 @@ ) from vllm.entrypoints.openai.engine.protocol import ( ErrorResponse, - PerRequestTimingMetrics, + PerRequestMetrics, PromptTokenUsageInfo, RequestResponseMetadata, UsageInfo, @@ -460,18 +461,22 @@ async def completion_stream_generator( # only emitted when usage reporting is enabled (i.e. # ``stream_options.include_usage=true`` or # ``--enable-force-include-usage``). - stream_per_request_metrics: PerRequestTimingMetrics | None = None - if ( - self.enable_per_request_metrics - # See note in request_output_to_completion_response: suppress - # when not attributable to one stream (multi-prompt or n>1). - and num_prompts == 1 - and (request.n or 1) == 1 - ): - last_metrics = last_res.metrics if last_res is not None else None - stream_per_request_metrics = build_per_request_timing_metrics( - last_metrics, total_completion_tokens - ) + stream_per_request_metrics: PerRequestMetrics | None = None + # See note in request_output_to_completion_response: suppress when + # not attributable to one stream (multi-prompt or n>1). + if num_prompts == 1 and (request.n or 1) == 1: + if self.enable_per_request_metrics: + last_metrics = ( + last_res.metrics if last_res is not None else None + ) + stream_per_request_metrics = build_per_request_timing_metrics( + last_metrics, total_completion_tokens + ) + spec_stats = build_spec_decoding_metrics(last_res) + if spec_stats is not None: + if stream_per_request_metrics is None: + stream_per_request_metrics = PerRequestMetrics() + stream_per_request_metrics.speculative_decoding = spec_stats final_usage_chunk = CompletionStreamResponse( id=request_id, @@ -612,21 +617,24 @@ def request_output_to_completion_response( request_metadata.final_usage_info = usage - per_request_metrics: PerRequestTimingMetrics | None = None - if ( - self.enable_per_request_metrics - # Metrics describe a single generation stream, so suppress them when - # they cannot be attributed to one: multiple prompts (timestamps - # span prompts) or n>1 (stats belong to one of the n sequences). - and len(final_res_batch) == 1 - and (request.n or 1) == 1 - ): - last_metrics = ( - last_final_res.metrics if last_final_res is not None else None - ) - per_request_metrics = build_per_request_timing_metrics( - last_metrics, num_generated_tokens - ) + per_request_metrics: PerRequestMetrics | None = None + # Per-request metrics (timing + spec-decode acceptance) describe a single + # generation stream, so suppress them when they cannot be attributed to + # one: multiple prompts (timestamps span prompts) or n>1 (stats belong to + # one of the n sequences). + if len(final_res_batch) == 1 and (request.n or 1) == 1: + if self.enable_per_request_metrics: + last_metrics = ( + last_final_res.metrics if last_final_res is not None else None + ) + per_request_metrics = build_per_request_timing_metrics( + last_metrics, num_generated_tokens + ) + spec_stats = build_spec_decoding_metrics(last_final_res) + if spec_stats is not None: + if per_request_metrics is None: + per_request_metrics = PerRequestMetrics() + per_request_metrics.speculative_decoding = spec_stats if final_res_batch: kv_transfer_params = final_res_batch[0].kv_transfer_params diff --git a/vllm/entrypoints/openai/engine/protocol.py b/vllm/entrypoints/openai/engine/protocol.py index c32be4459c90..9635ece46b03 100644 --- a/vllm/entrypoints/openai/engine/protocol.py +++ b/vllm/entrypoints/openai/engine/protocol.py @@ -129,12 +129,36 @@ class UsageInfo(OpenAIBaseModel): completion_tokens_details: CompletionTokenUsageInfo | None = None -class PerRequestTimingMetrics(OpenAIBaseModel): +class SpeculativeDecodingMetrics(OpenAIBaseModel): + """Per-request speculative-decoding acceptance metrics. + + Experimental, subject to change. Only populated for single-sequence requests + (`n == 1`); `null` for `n > 1`, mirroring the timing metrics. + """ + + mean_acceptance_length: float + draft_acceptance_rate: float + # Dense histogram: index j holds the number of verify steps that accepted + # exactly j draft tokens (length num_spec_tokens + 1). Excludes the + # always-accepted bonus token. + acceptance_histogram: list[int] + num_spec_steps: int + num_accepted_draft_tokens: int + num_draft_tokens: int + num_spec_tokens: int + # Ordered per-verify-step arrays; populated only at the `detailed` level. + per_step_accepted: list[int] | None = None + per_step_drafted: list[int] | None = None + + +class PerRequestMetrics(OpenAIBaseModel): time_to_first_token_ms: float | None = None generation_time_ms: float | None = None queue_time_ms: float | None = None mean_itl_ms: float | None = None tokens_per_second: float | None = None + # Experimental, subject to change. + speculative_decoding: SpeculativeDecodingMetrics | None = None class RequestResponseMetadata(BaseModel): diff --git a/vllm/outputs.py b/vllm/outputs.py index 29584e0e34cc..84b9fccdc88d 100644 --- a/vllm/outputs.py +++ b/vllm/outputs.py @@ -13,7 +13,7 @@ from vllm.logger import init_logger from vllm.logprobs import PromptLogprobs, SampleLogprobs from vllm.lora.request import LoRARequest -from vllm.v1.metrics.stats import RequestStateStats +from vllm.v1.metrics.stats import RequestSpecDecodeMetrics, RequestStateStats logger = init_logger(__name__) @@ -48,6 +48,11 @@ class CompletionOutput: to stop, None if the completion finished for some other reason including encountering the EOS token. lora_request: The LoRA request that was used to generate the output. + spec_decode_metrics: Per-sequence speculative-decoding acceptance metrics, + populated on finish when speculative decoding ran and + ``--per-request-spec-decode-metrics`` is enabled; None otherwise. + Surfaced in the response as ``metrics.speculative_decoding`` for + single-sequence (``n == 1``) requests. """ index: int @@ -60,6 +65,7 @@ class CompletionOutput: stop_reason: int | str | None = None lora_request: LoRARequest | None = None sampling_mask: SamplingMask | None = None + spec_decode_metrics: RequestSpecDecodeMetrics | None = None def finished(self) -> bool: return self.finish_reason is not None diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 2115ca673496..26674154c648 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -55,7 +55,11 @@ from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.metrics.perf import ModelMetrics, PerfStats -from vllm.v1.metrics.stats import PrefixCacheStats, SchedulerStats +from vllm.v1.metrics.stats import ( + PrefixCacheStats, + RequestSpecDecodeMetrics, + SchedulerStats, +) from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput from vllm.v1.request import Request, RequestStatus, StreamingUpdate from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup @@ -88,6 +92,9 @@ def __init__( self.parallel_config = vllm_config.parallel_config self.log_stats = log_stats self.observability_config = vllm_config.observability_config + self.spec_decode_metrics_level = ( + self.observability_config.per_request_spec_decode_metrics + ) self.kv_metrics_collector: KVCacheMetricsCollector | None = None if self.observability_config.kv_cache_metrics: self.kv_metrics_collector = KVCacheMetricsCollector( @@ -1857,6 +1864,20 @@ def update_from_output( num_invalid_spec_tokens=scheduler_output.num_invalid_spec_tokens, request_id=req_id, ) + if request.spec_decode_metrics is not None: + # Exclude grammar-invalidated drafts from the proposed + # count, mirroring make_spec_decoding_stats; the accepted + # bucket (j) is unaffected. + adj_draft_tokens = num_draft_tokens + if scheduler_output.num_invalid_spec_tokens: + adj_draft_tokens -= ( + scheduler_output.num_invalid_spec_tokens.get(req_id, 0) + ) + request.spec_decode_metrics.observe( + num_draft_tokens=adj_draft_tokens, + num_accepted=num_accepted, + detailed=self.spec_decode_metrics_level == "detailed", + ) # Free encoder inputs only after the step has actually executed. if request.has_encoder_inputs: @@ -2020,6 +2041,11 @@ def update_from_output( stop_reason=request.stop_reason, events=request.take_events(), prefill_stats=prefill_stats, + spec_decode_metrics=( + request.spec_decode_metrics + if finish_reason is not None + else None + ), kv_transfer_params=kv_transfer_params, ec_transfer_params=ec_transfer_params, trace_headers=request.trace_headers, @@ -2321,6 +2347,10 @@ def add_request(self, request: Request) -> None: request.streaming_queue = deque() self._enqueue_waiting_request(request) self.requests[request.request_id] = request + if self.spec_decode_metrics_level != "none": + request.spec_decode_metrics = RequestSpecDecodeMetrics.new( + self.num_spec_tokens + ) if self.connector is not None: self.connector.on_new_request(request) if self.log_stats: diff --git a/vllm/v1/engine/__init__.py b/vllm/v1/engine/__init__.py index 0a7f440ef799..5ae9ee0cac83 100644 --- a/vllm/v1/engine/__init__.py +++ b/vllm/v1/engine/__init__.py @@ -16,7 +16,11 @@ from vllm.multimodal.inputs import MultiModalFeatureSpec from vllm.pooling_params import PoolingParams from vllm.sampling_params import SamplingParams -from vllm.v1.metrics.stats import PrefillStats, SchedulerStats +from vllm.v1.metrics.stats import ( + PrefillStats, + RequestSpecDecodeMetrics, + SchedulerStats, +) from vllm.v1.outputs import LogprobsLists, LogprobsTensors, SamplingMaskLists from vllm.v1.serial_utils import UtilityResult @@ -225,6 +229,10 @@ class EngineCoreOutput( new_sampling_mask: SamplingMaskLists | None = None + # Per-request spec-decode acceptance; attached only on the final output. + # Appended last so `array_like` positional serialization stays compatible. + spec_decode_metrics: RequestSpecDecodeMetrics | None = None + @property def finished(self) -> bool: return self.finish_reason is not None diff --git a/vllm/v1/engine/output_processor.py b/vllm/v1/engine/output_processor.py index 99f60d5d5df3..055e3af574c1 100644 --- a/vllm/v1/engine/output_processor.py +++ b/vllm/v1/engine/output_processor.py @@ -35,6 +35,7 @@ from vllm.v1.metrics.stats import ( IterationStats, LoRARequestStates, + RequestSpecDecodeMetrics, RequestStateStats, SchedulerStats, ) @@ -175,6 +176,9 @@ def __init__( self.queue = queue self.num_cached_tokens = 0 self.num_cache_creation_tokens = 0 + # Per-sequence spec-decode accumulator; arrives once (on finish) via + # EngineCoreOutput, then attached to this sequence's CompletionOutput. + self.spec_decode_metrics: RequestSpecDecodeMetrics | None = None self.stats = RequestStateStats(arrival_time=arrival_time) if log_stats else None @@ -429,6 +433,7 @@ def _new_completion_output( cumulative_logprob=self.logprobs_processor.cumulative_logprob, finish_reason=str(finish_reason) if finished else None, stop_reason=stop_reason if finished else None, + spec_decode_metrics=self.spec_decode_metrics if finished else None, ) def _new_pooling_output(self, pooling_output: torch.Tensor) -> PoolingOutput: @@ -658,6 +663,9 @@ def process_outputs( ) req_state.is_prefilling = False + if engine_core_output.spec_decode_metrics is not None: + req_state.spec_decode_metrics = engine_core_output.spec_decode_metrics + if pooling_output is None: assert req_state.detokenizer is not None assert req_state.logprobs_processor is not None diff --git a/vllm/v1/metrics/stats.py b/vllm/v1/metrics/stats.py index 3956f7e44137..3dbc5206ca94 100644 --- a/vllm/v1/metrics/stats.py +++ b/vllm/v1/metrics/stats.py @@ -297,6 +297,82 @@ def finalize(self, num_cached_tokens: int) -> None: ) +@dataclass +class RequestSpecDecodeMetrics: + """Per-output-sequence speculative-decoding statistics accumulator. + + Accumulates, over one sequence's verify steps, a histogram of accepted + draft-token counts (``j``, draft-only) and the total number of proposed + draft tokens. When ``detailed`` is requested it also records the ordered + per-step accepted/proposed sequences (``summary`` omits them). Tracked per + engine ``Request`` (one per sampled sequence, so ``n > 1`` yields one per + child), surfaced via ``EngineCoreOutput`` and the response + ``metrics.speculative_decoding`` for single-sequence requests (see + ``to_dict``). + + Fields: + num_spec_tokens: Configured ``num_speculative_tokens`` (the max ``k``); + also the histogram's upper bound. + histogram: Dense counts indexed by accepted draft tokens ``j`` + (length ``num_spec_tokens + 1``). + num_draft_tokens: Total proposed draft tokens, after the + grammar-invalidated (``num_invalid_spec_tokens``) adjustment. + per_step_accepted: Ordered accepted-draft count per verify step + (``detailed`` only; empty otherwise). + per_step_drafted: Ordered proposed-draft count per verify step + (``detailed`` only; empty otherwise). + """ + + num_spec_tokens: int + histogram: list[int] = field(default_factory=list) + num_draft_tokens: int = 0 + per_step_accepted: list[int] = field(default_factory=list) + per_step_drafted: list[int] = field(default_factory=list) + + @classmethod + def new(cls, num_spec_tokens: int) -> "RequestSpecDecodeMetrics": + return cls( + num_spec_tokens=num_spec_tokens, + histogram=[0] * (num_spec_tokens + 1), + ) + + def observe( + self, num_draft_tokens: int, num_accepted: int, detailed: bool = False + ) -> None: + self.histogram[num_accepted] += 1 + self.num_draft_tokens += num_draft_tokens + if detailed: + self.per_step_accepted.append(num_accepted) + self.per_step_drafted.append(num_draft_tokens) + + def to_dict(self) -> dict[str, Any]: + """Payload matching ``SpeculativeDecodingMetrics`` for the response. + + ``acceptance_histogram`` is a dense list indexed by accepted draft count + ``j`` (length ``num_spec_tokens + 1``). ``mean_acceptance_length`` + includes the bonus token (``j + 1``); ``draft_acceptance_rate`` is + draft-only, full precision. Per-step arrays are included only when + populated (``detailed`` level). + """ + num_spec_steps = sum(self.histogram) + num_accepted = sum(j * count for j, count in enumerate(self.histogram)) + mean_al = 1.0 + num_accepted / num_spec_steps if num_spec_steps else 1.0 + rate = num_accepted / self.num_draft_tokens if self.num_draft_tokens else 0.0 + result: dict[str, Any] = { + "mean_acceptance_length": mean_al, + "draft_acceptance_rate": rate, + "acceptance_histogram": list(self.histogram), + "num_spec_steps": num_spec_steps, + "num_accepted_draft_tokens": num_accepted, + "num_draft_tokens": self.num_draft_tokens, + "num_spec_tokens": self.num_spec_tokens, + } + if self.per_step_accepted: + result["per_step_accepted"] = self.per_step_accepted + result["per_step_drafted"] = self.per_step_drafted + return result + + @dataclass class PromptTokenStats: """Breakdown of prompt tokens by source. diff --git a/vllm/v1/request.py b/vllm/v1/request.py index 44de9ad1f7c8..8b453a09069e 100644 --- a/vllm/v1/request.py +++ b/vllm/v1/request.py @@ -20,7 +20,7 @@ EngineCoreRequest, FinishReason, ) -from vllm.v1.metrics.stats import PrefillStats +from vllm.v1.metrics.stats import PrefillStats, RequestSpecDecodeMetrics from vllm.v1.structured_output.request import StructuredOutputRequest from vllm.v1.utils import ConstantList @@ -211,6 +211,11 @@ def __init__( self.prefill_stats: PrefillStats | None = PrefillStats() + # Per-request speculative-decoding acceptance accumulator. Populated by + # the scheduler when --per-request-spec-decode-metrics is set (eagerly on + # add_request, then observed each verify step); stays None otherwise. + self.spec_decode_metrics: RequestSpecDecodeMetrics | None = None + self.block_hashes: list[BlockHash] = [] # Store the block hasher without binding self to avoid creating a # reference cycle (Request -> partial -> Request) that prevents From 0a5a55136fa0552a19816b211c4a57f1a341d830 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Thu, 20 Aug 2026 16:17:02 -0700 Subject: [PATCH 226/839] [CI][Docker] Pin remaining manylinux builder images (#53172) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Co-authored-by: OpenAI Codex --- .buildkite/image_build/image_build_arm64.sh | 2 +- .buildkite/release-pipeline.yaml | 18 +++++++++--------- .../scripts/hardware_ci/run-gh200-test.sh | 2 +- .../installation/gpu.cuda.inc.md | 4 ++-- 4 files changed, 13 insertions(+), 13 deletions(-) diff --git a/.buildkite/image_build/image_build_arm64.sh b/.buildkite/image_build/image_build_arm64.sh index c003f1f03d47..eb56f02b89e4 100755 --- a/.buildkite/image_build/image_build_arm64.sh +++ b/.buildkite/image_build/image_build_arm64.sh @@ -26,7 +26,7 @@ else --platform linux/arm64 \ --build-arg max_jobs=16 \ --build-arg nvcc_threads=4 \ - --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0 \ + --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0-78e737ad29420ffc4800e677c51e2a852caf8359 \ --build-arg torch_cuda_arch_list="9.0 10.0 11.0 12.0" \ --build-arg USE_SCCACHE=1 \ --build-arg buildkite_commit="$BUILDKITE_COMMIT" \ diff --git a/.buildkite/release-pipeline.yaml b/.buildkite/release-pipeline.yaml index 15e50d3a75fc..1d421d664a50 100644 --- a/.buildkite/release-pipeline.yaml +++ b/.buildkite/release-pipeline.yaml @@ -33,7 +33,7 @@ steps: agents: queue: arm64_cpu_queue_release commands: - - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ." + - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0-78e737ad29420ffc4800e677c51e2a852caf8359 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ." - "mkdir artifacts" - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'" - "bash .buildkite/scripts/upload-nightly-wheels.sh" @@ -71,7 +71,7 @@ steps: agents: queue: arm64_cpu_queue_release commands: - - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda12.9 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ." + - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda12.9-78e737ad29420ffc4800e677c51e2a852caf8359 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ." - "mkdir artifacts" - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'" - "bash .buildkite/scripts/upload-nightly-wheels.sh" @@ -142,7 +142,7 @@ steps: agents: queue: cpu_queue_release commands: - - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda12.9 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ." + - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda12.9-78e737ad29420ffc4800e677c51e2a852caf8359 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ." - "mkdir artifacts" - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'" - "bash .buildkite/scripts/upload-nightly-wheels.sh" @@ -260,7 +260,7 @@ steps: --build-arg CUDA_VERSION=13.0.2 \ --build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64}" \ --build-arg INSTALL_KV_CONNECTORS=true \ - --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0 \ + --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0-78e737ad29420ffc4800e677c51e2a852caf8359 \ --target vllm-openai \ --progress plain \ -f docker/Dockerfile . @@ -281,7 +281,7 @@ steps: --build-arg USE_SCCACHE=1 \ --build-arg GIT_REPO_CHECK=1 \ --build-arg CUDA_VERSION=12.9.1 \ - --build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda12.9 \ + --build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda12.9-78e737ad29420ffc4800e677c51e2a852caf8359 \ --build-arg torch_cuda_arch_list="${CUDA_ARCH_X86_CU129}" \ --build-arg INSTALL_KV_CONNECTORS=true \ --target vllm-openai \ @@ -307,7 +307,7 @@ steps: --build-arg USE_SCCACHE=1 \ --build-arg GIT_REPO_CHECK=1 \ --build-arg CUDA_VERSION=12.9.1 \ - --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda12.9 \ + --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda12.9-78e737ad29420ffc4800e677c51e2a852caf8359 \ --build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64_CU129}" \ --build-arg INSTALL_KV_CONNECTORS=true \ --target vllm-openai \ @@ -361,7 +361,7 @@ steps: --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \ --build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64}" \ --build-arg INSTALL_KV_CONNECTORS=true \ - --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0 \ + --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0-78e737ad29420ffc4800e677c51e2a852caf8359 \ --target vllm-openai \ --progress plain \ -f docker/Dockerfile . @@ -382,7 +382,7 @@ steps: --build-arg USE_SCCACHE=1 \ --build-arg GIT_REPO_CHECK=1 \ --build-arg CUDA_VERSION=12.9.1 \ - --build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda12.9 \ + --build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda12.9-78e737ad29420ffc4800e677c51e2a852caf8359 \ --build-arg UBUNTU_VERSION=24.04 \ --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \ --build-arg torch_cuda_arch_list="${CUDA_ARCH_X86_CU129}" \ @@ -409,7 +409,7 @@ steps: --build-arg USE_SCCACHE=1 \ --build-arg GIT_REPO_CHECK=1 \ --build-arg CUDA_VERSION=12.9.1 \ - --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda12.9 \ + --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda12.9-78e737ad29420ffc4800e677c51e2a852caf8359 \ --build-arg UBUNTU_VERSION=24.04 \ --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \ --build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64_CU129}" \ diff --git a/.buildkite/scripts/hardware_ci/run-gh200-test.sh b/.buildkite/scripts/hardware_ci/run-gh200-test.sh index ed82e914909f..dca25b177a93 100644 --- a/.buildkite/scripts/hardware_ci/run-gh200-test.sh +++ b/.buildkite/scripts/hardware_ci/run-gh200-test.sh @@ -15,7 +15,7 @@ DOCKER_BUILDKIT=1 docker build . \ -t gh200-test \ --build-arg max_jobs=66 \ --build-arg nvcc_threads=2 \ - --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0 \ + --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0-78e737ad29420ffc4800e677c51e2a852caf8359 \ --build-arg torch_cuda_arch_list="9.0+PTX" # Setup cleanup diff --git a/docs/getting_started/installation/gpu.cuda.inc.md b/docs/getting_started/installation/gpu.cuda.inc.md index a417bf4af7e1..555c8b2edf94 100644 --- a/docs/getting_started/installation/gpu.cuda.inc.md +++ b/docs/getting_started/installation/gpu.cuda.inc.md @@ -389,7 +389,7 @@ A docker container can be built for aarch64 systems such as the Nvidia Grace-Hop -t vllm/vllm-gh200-openai:latest \ --build-arg max_jobs=66 \ --build-arg nvcc_threads=2 \ - --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0 \ + --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0-78e737ad29420ffc4800e677c51e2a852caf8359 \ --build-arg torch_cuda_arch_list="9.0 10.0+PTX" ``` @@ -400,7 +400,7 @@ For (G)B300, we recommend using CUDA 13, as shown in the following command. ```bash DOCKER_BUILDKIT=1 docker build \ --build-arg CUDA_VERSION=13.0.2 \ - --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0 \ + --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0-78e737ad29420ffc4800e677c51e2a852caf8359 \ --build-arg max_jobs=256 \ --build-arg nvcc_threads=2 \ --build-arg torch_cuda_arch_list='9.0 10.0+PTX' \ From f32b17b6d616605eb5ed498777f150242bd1bc30 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Fri, 21 Aug 2026 10:31:25 +1000 Subject: [PATCH 227/839] [Rust Frontend] Support `--generation-config vllm` (#53044) Signed-off-by: Bugen Zhao --- rust/src/chat/src/backend/hf.rs | 8 +- rust/src/chat/src/backend/mod.rs | 4 +- rust/src/chat/src/lib.rs | 1 + rust/src/cmd/src/cli.rs | 10 +- rust/src/cmd/src/cli/tests.rs | 32 +++- rust/src/cmd/src/cli/unsupported.rs | 13 +- .../examples/external_engine_openai_qwen.rs | 1 + rust/src/server/src/config.rs | 6 +- rust/src/server/src/lib.rs | 5 +- rust/src/server/src/render.rs | 1 + rust/src/text/src/backend/hf/mod.rs | 141 +++++++++++++++--- rust/src/text/src/backend/mod.rs | 44 ++++++ rust/src/text/src/lib.rs | 4 +- rust/src/text/src/lower.rs | 15 +- 14 files changed, 244 insertions(+), 41 deletions(-) diff --git a/rust/src/chat/src/backend/hf.rs b/rust/src/chat/src/backend/hf.rs index 20d2b81cc178..0be33f91a8ae 100644 --- a/rust/src/chat/src/backend/hf.rs +++ b/rust/src/chat/src/backend/hf.rs @@ -129,8 +129,11 @@ pub(super) async fn load_model_backends( options: LoadModelBackendsOptions, ) -> Result { let files = ResolvedModelFiles::new(model_id).await?; - let text_backend = - HfTextBackend::from_resolved_model_files(files.clone(), model_id.to_string())?; + let text_backend = HfTextBackend::from_resolved_model_files( + files.clone(), + model_id.to_string(), + options.generation_config, + )?; let tokenizer = text_backend.tokenizer(); let text_backend: DynTextBackend = Arc::new(text_backend); @@ -227,6 +230,7 @@ mod tests { resolved_files(config_json, tokenizer_config_json), "test-model".to_string(), LoadModelBackendsOptions { + generation_config: Default::default(), renderer, language_model_only: false, chat_template_content_format: Default::default(), diff --git a/rust/src/chat/src/backend/mod.rs b/rust/src/chat/src/backend/mod.rs index 0a0f757d654f..77b9114a5852 100644 --- a/rust/src/chat/src/backend/mod.rs +++ b/rust/src/chat/src/backend/mod.rs @@ -5,7 +5,7 @@ use std::collections::HashMap; use std::sync::Arc; use serde_json::Value; -use vllm_text::{DynTextBackend, TextBackend}; +use vllm_text::{DynTextBackend, GenerationConfigMode, TextBackend}; use crate::error::Result; use crate::multimodal::{MmLimitPerPrompt, MultimodalModelInfo}; @@ -61,6 +61,8 @@ pub type DynChatTextBackend = Arc; /// Frontend-side chat backend loading options. #[derive(Debug, Clone, PartialEq, Eq, Default)] pub struct LoadModelBackendsOptions { + /// Which generation-config sampling defaults to inherit. + pub generation_config: GenerationConfigMode, /// Which chat renderer implementation to use. pub renderer: RendererSelection, /// Disable frontend-side multimodal preprocessing and render the model as diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs index cf22722dd7f1..c6b91679e027 100644 --- a/rust/src/chat/src/lib.rs +++ b/rust/src/chat/src/lib.rs @@ -44,6 +44,7 @@ pub use request::{ pub use stream::{ChatEventStream, ChatEventStreamTrait, CollectedAssistantMessage}; pub use vllm_engine_core_client::protocol::multimodal::MmFeatures; pub use vllm_llm::FinishReason; +pub use vllm_text::GenerationConfigMode; mod backend; mod error; diff --git a/rust/src/cmd/src/cli.rs b/rust/src/cmd/src/cli.rs index 64b6dc9800cf..a061c8eff077 100644 --- a/rust/src/cmd/src/cli.rs +++ b/rust/src/cmd/src/cli.rs @@ -22,8 +22,8 @@ use serde_json::Value; use serde_with::{DefaultOnNull, OneOrMany, serde_as}; use thiserror_ext::AsReport as _; use uuid::Uuid; -use vllm_chat::ReasoningParserFactory; use vllm_chat::multimodal::MmLimitPerPrompt; +use vllm_chat::{GenerationConfigMode, ReasoningParserFactory}; use vllm_engine_core_client::TransportMode; use vllm_managed_engine::ManagedEngineConfig; use vllm_managed_engine::cli::{ManagedEngineArgs, repartition_managed_engine_args}; @@ -186,6 +186,12 @@ pub struct SharedRuntimeArgs { /// public model ID. pub model: String, + /// The source of generation-config sampling defaults. `"auto"` loads the + /// model's defaults, while `"vllm"` uses vLLM's neutral defaults. + #[arg(long, default_value_t)] + #[serde(default)] + pub generation_config: GenerationConfigMode, + /// Maximum time to wait for the expected engines to register on the /// frontend transport. #[arg( @@ -482,6 +488,7 @@ impl SharedRuntimeArgs { None => CoordinatorMode::None, }, model: self.model, + generation_config: self.generation_config, served_model_name: self.served_model_name, listener_mode: HttpListenerMode::InheritedFd { fd: listen_fd }, tool_call_parser: self.tool_call_parser, @@ -535,6 +542,7 @@ impl SharedRuntimeArgs { }, coordinator_mode: CoordinatorMode::MaybeInProc, model: self.model, + generation_config: self.generation_config, served_model_name: self.served_model_name, listener_mode, tool_call_parser: self.tool_call_parser, diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs index 80436ec423e6..89979ec6d094 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -3,7 +3,9 @@ use expect_test::expect; use vllm_engine_core_client::TransportMode; -use vllm_server::{Config, HttpListenerMode, ParserSelection, RendererSelection}; +use vllm_server::{ + Config, GenerationConfigMode, HttpListenerMode, ParserSelection, RendererSelection, +}; use super::{BenchCommand, Cli, Command}; @@ -88,6 +90,7 @@ fn serve_args_forward_python_flags_with_separator() { uds: None, runtime: SharedRuntimeArgs { model: "Qwen/Qwen3-0.6B", + generation_config: Auto, engine_ready_timeout_secs: 600, tool_call_parser: Auto, reasoning_parser: Auto, @@ -732,6 +735,24 @@ fn serve_args_reject_unsupported_flag_arg() { "#]].assert_eq(&error.to_string()); } +#[test] +fn serve_args_reject_custom_generation_config_source() { + let error = Cli::try_parse_from([ + "vllm-rs", + "serve", + "Qwen/Qwen3-0.6B", + "--generation-config", + "/tmp/custom-config", + ]) + .unwrap_err(); + + expect![[r#" + error: invalid value '/tmp/custom-config' for '--generation-config ': generation config source `/tmp/custom-config` is not implemented yet (expected one of: auto, vllm) + + For more information, try '--help'. + "#]].assert_eq(&error.to_string()); +} + #[test] fn serve_args_reject_unsupported_no_flag_alias() { let error = Cli::try_parse_from([ @@ -788,6 +809,7 @@ fn frontend_args_accept_json() { data_parallel_size: None, runtime: SharedRuntimeArgs { model: "Qwen/Qwen3-0.6B", + generation_config: Auto, engine_ready_timeout_secs: 600, tool_call_parser: None, reasoning_parser: None, @@ -858,6 +880,7 @@ fn frontend_args_json_applies_defaults() { panic!("expected frontend args"); }; assert_eq!(args.runtime.model, "Qwen/Qwen3-0.6B"); + assert_eq!(args.runtime.generation_config, GenerationConfigMode::Auto); assert_eq!(args.runtime.engine_ready_timeout_secs, 600); assert_eq!(args.runtime.tool_call_parser, ParserSelection::None); assert_eq!(args.runtime.reasoning_parser, ParserSelection::None); @@ -901,7 +924,7 @@ fn frontend_args_json_accepts_supported_non_default_fields() { "--output-address", "ipc:///tmp/output.sock", "--args-json", - r#"{"model_tag":"Qwen/Qwen3-0.6B","engine_ready_timeout_secs":42,"tool_call_parser":"hermes","reasoning_parser":"qwen3_thinking","tokenizer_mode":"deepseek_v32","language_model_only":true,"max_logprobs":-1,"shutdown_timeout":3}"#, + r#"{"model_tag":"Qwen/Qwen3-0.6B","generation_config":"vllm","engine_ready_timeout_secs":42,"tool_call_parser":"hermes","reasoning_parser":"qwen3_thinking","tokenizer_mode":"deepseek_v32","language_model_only":true,"max_logprobs":-1,"shutdown_timeout":3}"#, ]) .unwrap(); @@ -909,6 +932,7 @@ fn frontend_args_json_accepts_supported_non_default_fields() { panic!("expected frontend args"); }; assert_eq!(args.runtime.engine_ready_timeout_secs, 42); + assert_eq!(args.runtime.generation_config, GenerationConfigMode::Vllm); assert_eq!( args.runtime.tool_call_parser, ParserSelection::Explicit("hermes".to_string()) @@ -1379,6 +1403,7 @@ fn serve_args_accept_handshake_aliases() { uds: None, runtime: SharedRuntimeArgs { model: "Qwen/Qwen3-0.6B", + generation_config: Auto, engine_ready_timeout_secs: 600, tool_call_parser: Auto, reasoning_parser: Auto, @@ -1523,6 +1548,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() { }, coordinator_mode: MaybeInProc, model: "Qwen/Qwen3-0.6B", + generation_config: Auto, served_model_name: [], listener_mode: BindTcp { host: "127.0.0.1", @@ -1608,6 +1634,7 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() { }, coordinator_mode: MaybeInProc, model: "Qwen/Qwen3-0.6B", + generation_config: Auto, served_model_name: [], listener_mode: BindTcp { host: "127.0.0.1", @@ -1715,6 +1742,7 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present address: "tcp://127.0.0.1:7000", }, model: "Qwen/Qwen3-0.6B", + generation_config: Auto, served_model_name: [], listener_mode: InheritedFd { fd: 3, diff --git a/rust/src/cmd/src/cli/unsupported.rs b/rust/src/cmd/src/cli/unsupported.rs index a48e61dbf7a6..d2372ed56686 100644 --- a/rust/src/cmd/src/cli/unsupported.rs +++ b/rust/src/cmd/src/cli/unsupported.rs @@ -240,15 +240,12 @@ pub struct EngineUnsupportedArgs { #[arg(long)] pub hf_overrides: Option, - /// The folder path to the generation config. Defaults to `"auto"`, the - /// generation config will be loaded from model path. If set to `"vllm"`, no - /// generation config is loaded, vLLM defaults will be used. If set to a - /// folder path, the generation config will be loaded from the specified - /// folder path. If `max_new_tokens` is specified in generation config, - /// then it sets a server-wide limit on the number of output tokens for - /// all requests. + /// Overrides or sets generation config. e.g. `{"temperature": 0.5}`. If + /// used with `--generation-config auto`, the override parameters will be + /// merged with the default config from the model. If used with + /// `--generation-config vllm`, only the override parameters are used. #[arg(long)] - pub generation_config: Option, + pub override_generation_config: Option, /// IOProcessor plugin name to load at model startup #[arg(long)] diff --git a/rust/src/server/examples/external_engine_openai_qwen.rs b/rust/src/server/examples/external_engine_openai_qwen.rs index 387776263059..5c99fc92ebbc 100644 --- a/rust/src/server/examples/external_engine_openai_qwen.rs +++ b/rust/src/server/examples/external_engine_openai_qwen.rs @@ -60,6 +60,7 @@ async fn main() -> Result<()> { }, coordinator_mode: CoordinatorMode::MaybeInProc, model: args.model, + generation_config: Default::default(), served_model_name: vec![], listener_mode: HttpListenerMode::BindTcp { host: "127.0.0.1".to_string(), diff --git a/rust/src/server/src/config.rs b/rust/src/server/src/config.rs index a33d7e25fed8..eb690d4ea0bd 100644 --- a/rust/src/server/src/config.rs +++ b/rust/src/server/src/config.rs @@ -11,7 +11,9 @@ use educe::Educe; use serde::Serialize; use serde_json::Value; use vllm_chat::multimodal::MmLimitPerPrompt; -use vllm_chat::{ChatTemplateContentFormatOption, ParserSelection, RendererSelection}; +use vllm_chat::{ + ChatTemplateContentFormatOption, GenerationConfigMode, ParserSelection, RendererSelection, +}; use vllm_engine_core_client::{CoordinatorMode as EngineCoreCoordinatorMode, TransportMode}; /// Default keep-alive idle timeout (seconds); also the head-read bound @@ -165,6 +167,8 @@ pub struct Config { pub coordinator_mode: CoordinatorMode, /// Backend model identifier used for engine-core loading. pub model: String, + /// Which generation-config sampling defaults to inherit. + pub generation_config: GenerationConfigMode, /// Model name(s) exposed to clients via the OpenAI API. When non-empty, /// the first entry is used as the primary ID in responses and all entries /// are accepted in requests. When empty, falls back to `model`. diff --git a/rust/src/server/src/lib.rs b/rust/src/server/src/lib.rs index f9f5c423506f..e8138b2494aa 100644 --- a/rust/src/server/src/lib.rs +++ b/rust/src/server/src/lib.rs @@ -45,7 +45,9 @@ use tonic_health::server::health_reporter; use tower::ServiceExt as _; use tracing::{info, trace, warn}; use vllm_chat::{ChatLlm, LoadModelBackendsOptions, load_model_backends}; -pub use vllm_chat::{ChatTemplateContentFormatOption, ParserSelection, RendererSelection}; +pub use vllm_chat::{ + ChatTemplateContentFormatOption, GenerationConfigMode, ParserSelection, RendererSelection, +}; use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig}; use vllm_llm::Llm; use vllm_text::TextLlm; @@ -96,6 +98,7 @@ async fn build_state(config: &Config) -> Result> { let loaded = load_model_backends( &config.model, LoadModelBackendsOptions { + generation_config: config.generation_config, renderer: config.renderer, language_model_only: config.language_model_only, chat_template: config.chat_template.clone(), diff --git a/rust/src/server/src/render.rs b/rust/src/server/src/render.rs index ee8098f8e93b..35c821bc6eee 100644 --- a/rust/src/server/src/render.rs +++ b/rust/src/server/src/render.rs @@ -56,6 +56,7 @@ async fn build_state(config: &RenderConfig) -> Result> { let loaded = load_model_backends( &config.model, LoadModelBackendsOptions { + generation_config: Default::default(), renderer: config.renderer, language_model_only: true, chat_template: config.chat_template.clone(), diff --git a/rust/src/text/src/backend/hf/mod.rs b/rust/src/text/src/backend/hf/mod.rs index 4fc7a18753a4..ae768767d99a 100644 --- a/rust/src/text/src/backend/hf/mod.rs +++ b/rust/src/text/src/backend/hf/mod.rs @@ -18,7 +18,7 @@ pub use self::config::{ load_tokenizer_config, }; pub use self::model_files::{ResolvedModelFiles, TokenizerSource}; -use crate::backend::{SamplingHints, TextBackend}; +use crate::backend::{GenerationConfigMode, SamplingHints, TextBackend}; use crate::error::Result; fn load_tokenizer(tokenizer: &TokenizerSource) -> Result { @@ -41,6 +41,7 @@ pub struct HfTextBackend { /// Generation-config for sampling defaults that may be inherited when the /// user does not explicitly override them. generation_config: GenerationConfig, + generation_config_mode: GenerationConfigMode, /// Model vocabulary size from the selected text config. model_vocab_size: usize, /// Model config (`config.json`). @@ -48,14 +49,12 @@ pub struct HfTextBackend { } impl HfTextBackend { - /// Load the text backend with the given model id. - pub async fn from_model(model_id: &str) -> Result { - let files = ResolvedModelFiles::new(model_id).await?; - Self::from_resolved_model_files(files, model_id.to_string()) - } - /// Load the text backend from resolved Hugging Face model files. - pub fn from_resolved_model_files(files: ResolvedModelFiles, model_id: String) -> Result { + pub fn from_resolved_model_files( + files: ResolvedModelFiles, + model_id: String, + generation_config_mode: GenerationConfigMode, + ) -> Result { let tokenizer_config = load_tokenizer_config(files.tokenizer_config_path.as_deref())?; let tokenizer = load_tokenizer(&files.tokenizer)?; let model_config = load_model_config(files.config_path.as_deref())?; @@ -80,6 +79,7 @@ impl HfTextBackend { primary_eos_token_id, extra_eos_token_ids, generation_config, + generation_config_mode, model_vocab_size, model_config, }) @@ -146,16 +146,38 @@ impl TextBackend for HfTextBackend { } fn sampling_hints(&self) -> Result { - Ok(SamplingHints { - primary_eos_token_id: self.primary_eos_token_id, - extra_eos_token_ids: self.extra_eos_token_ids.clone(), - default_temperature: self.generation_config.temperature, - default_top_p: self.generation_config.top_p, - default_top_k: self.generation_config.top_k, - default_min_p: self.generation_config.min_p, - default_repetition_penalty: self.generation_config.repetition_penalty, - default_max_tokens: self.generation_config.max_new_tokens, - }) + Ok(build_sampling_hints( + self.generation_config_mode, + &self.generation_config, + self.primary_eos_token_id, + self.extra_eos_token_ids.clone(), + )) + } +} + +fn build_sampling_hints( + mode: GenerationConfigMode, + generation_config: &GenerationConfig, + primary_eos_token_id: Option, + extra_eos_token_ids: BTreeSet, +) -> SamplingHints { + let sampling_config = match mode { + // Auto inherits sampling defaults from the model's generation config. + GenerationConfigMode::Auto => Some(generation_config), + // Vllm skips model-provided sampling defaults. `lower_sampling_params` + // applies vLLM fallback values; separately resolved EOS tokens remain. + GenerationConfigMode::Vllm => None, + }; + + SamplingHints { + primary_eos_token_id, + extra_eos_token_ids, + default_temperature: sampling_config.and_then(|config| config.temperature), + default_top_p: sampling_config.and_then(|config| config.top_p), + default_top_k: sampling_config.and_then(|config| config.top_k), + default_min_p: sampling_config.and_then(|config| config.min_p), + default_repetition_penalty: sampling_config.and_then(|config| config.repetition_penalty), + default_max_tokens: sampling_config.and_then(|config| config.max_new_tokens), } } @@ -163,7 +185,10 @@ impl TextBackend for HfTextBackend { mod tests { use std::collections::BTreeSet; - use super::{GenerationConfig, HfTokenizerConfig, ModelConfig, resolve_eos_token_ids}; + use super::{ + GenerationConfig, GenerationConfigMode, HfTokenizerConfig, ModelConfig, + build_sampling_hints, resolve_eos_token_ids, + }; use vllm_tokenizer::Tokenizer; struct FakeTokenizer; @@ -235,4 +260,82 @@ mod tests { assert_eq!(primary, Some(2)); assert_eq!(extra, BTreeSet::from([200006, 200010])); } + + #[test] + fn vllm_generation_config_keeps_eos_and_uses_neutral_sampling_defaults() { + let generation_config: GenerationConfig = serde_json::from_str( + r#"{ + "eos_token_id": [2, 3], + "temperature": 0.6, + "top_p": 0.95, + "top_k": 20, + "min_p": 0.1, + "repetition_penalty": 1.1, + "max_new_tokens": 4096 + }"#, + ) + .unwrap(); + let extra_eos_token_ids = BTreeSet::from([3]); + + let hints = ( + build_sampling_hints( + GenerationConfigMode::Auto, + &generation_config, + Some(2), + extra_eos_token_ids.clone(), + ), + build_sampling_hints( + GenerationConfigMode::Vllm, + &generation_config, + Some(2), + extra_eos_token_ids, + ), + ); + + expect_test::expect![[r#" + ( + SamplingHints { + primary_eos_token_id: Some( + 2, + ), + extra_eos_token_ids: { + 3, + }, + default_temperature: Some( + 0.6, + ), + default_top_p: Some( + 0.95, + ), + default_top_k: Some( + 20, + ), + default_min_p: Some( + 0.1, + ), + default_repetition_penalty: Some( + 1.1, + ), + default_max_tokens: Some( + 4096, + ), + }, + SamplingHints { + primary_eos_token_id: Some( + 2, + ), + extra_eos_token_ids: { + 3, + }, + default_temperature: None, + default_top_p: None, + default_top_k: None, + default_min_p: None, + default_repetition_penalty: None, + default_max_tokens: None, + }, + ) + "#]] + .assert_debug_eq(&hints); + } } diff --git a/rust/src/text/src/backend/mod.rs b/rust/src/text/src/backend/mod.rs index 98b709f11818..3bda9212a820 100644 --- a/rust/src/text/src/backend/mod.rs +++ b/rust/src/text/src/backend/mod.rs @@ -3,12 +3,56 @@ pub mod hf; +use std::fmt; +use std::str::FromStr; use std::sync::Arc; +use serde_with::{DeserializeFromStr, SerializeDisplay}; use vllm_tokenizer::DynTokenizer; use crate::error::Result; +/// Select which sampling defaults to inherit from generation config. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, DeserializeFromStr, SerializeDisplay)] +pub enum GenerationConfigMode { + /// Inherit sampling defaults from the model's generation config. + #[default] + Auto, + /// Use vLLM's neutral sampling defaults. + Vllm, +} + +impl GenerationConfigMode { + pub const AUTO_LITERAL: &str = "auto"; + pub const VLLM_LITERAL: &str = "vllm"; +} + +impl FromStr for GenerationConfigMode { + type Err = String; + + fn from_str(value: &str) -> std::result::Result { + if value.eq_ignore_ascii_case(Self::AUTO_LITERAL) { + Ok(Self::Auto) + } else if value.eq_ignore_ascii_case(Self::VLLM_LITERAL) { + Ok(Self::Vllm) + } else { + Err(format!( + "generation config source `{value}` is not implemented yet \ + (expected one of: auto, vllm)" + )) + } + } +} + +impl fmt::Display for GenerationConfigMode { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Auto => f.write_str(Self::AUTO_LITERAL), + Self::Vllm => f.write_str(Self::VLLM_LITERAL), + } + } +} + /// Tokenizer/model-derived defaults used to enrich text-generation requests /// before they are lowered into engine-core. #[derive(Debug, Clone, Default, PartialEq)] diff --git a/rust/src/text/src/lib.rs b/rust/src/text/src/lib.rs index 53edf7e30918..f49038b27fed 100644 --- a/rust/src/text/src/lib.rs +++ b/rust/src/text/src/lib.rs @@ -9,7 +9,9 @@ use std::mem::take; -pub use backend::{DynTextBackend, SamplingHints, SamplingLimits, TextBackend}; +pub use backend::{ + DynTextBackend, GenerationConfigMode, SamplingHints, SamplingLimits, TextBackend, +}; pub use error::{Error, LogprobsError, Result, SamplingParamsError, TokenIdsError}; use futures::Stream; pub use lower::{ diff --git a/rust/src/text/src/lower.rs b/rust/src/text/src/lower.rs index 55be99cc4e57..0e1fbd26bcbe 100644 --- a/rust/src/text/src/lower.rs +++ b/rust/src/text/src/lower.rs @@ -322,8 +322,8 @@ mod tests { use vllm_tokenizer::test_utils::TestTokenizer; use super::*; - use crate::backend::hf::HfTextBackend; - use crate::backend::{SamplingHints, TextBackend as _}; + use crate::backend::hf::{HfTextBackend, ResolvedModelFiles}; + use crate::backend::{GenerationConfigMode, SamplingHints, TextBackend as _}; use crate::error::{LogprobsError, SamplingParamsError, TokenIdsError}; use crate::request::{Prompt, TextRequest}; @@ -785,9 +785,14 @@ mod tests { #[tokio::test] #[file_serial(hf_qwen3)] async fn lower_text_request_uses_real_qwen_generation_defaults() { - let backend = HfTextBackend::from_model("Qwen/Qwen3-0.6B") - .await - .expect("load qwen tokenizer and generation config"); + let model_id = "Qwen/Qwen3-0.6B"; + let files = ResolvedModelFiles::new(model_id).await.expect("resolve qwen model files"); + let backend = HfTextBackend::from_resolved_model_files( + files, + model_id.to_string(), + GenerationConfigMode::Auto, + ) + .expect("load qwen tokenizer and generation config"); let hints = backend.sampling_hints().expect("collect sampling hints"); expect_test::expect![[r#" From d29f7f5c9294be8e489dac34d45a939b95a06336 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Thu, 20 Aug 2026 17:47:08 -0700 Subject: [PATCH 228/839] [Bugfix] Load untied Gemma LM head weights (#53170) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> --- .../models/language/generation/test_gemma.py | 45 +++++++++++++++++++ vllm/model_executor/models/gemma.py | 15 ++++++- 2 files changed, 58 insertions(+), 2 deletions(-) diff --git a/tests/models/language/generation/test_gemma.py b/tests/models/language/generation/test_gemma.py index 246b893be315..2375fc2ab7b2 100644 --- a/tests/models/language/generation/test_gemma.py +++ b/tests/models/language/generation/test_gemma.py @@ -1,11 +1,56 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace +from typing import cast + import numpy as np import pytest +import torch + +from vllm.config import VllmConfig +from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding +from vllm.model_executor.models import gemma MODELS = ["google/gemma-2b", "google/gemma-2-2b", "google/gemma-3-4b-it"] +@pytest.mark.cpu_test +@pytest.mark.usefixtures("dist_init") +def test_checkpoint_lm_head_can_override_tied_config(monkeypatch) -> None: + """A physical LM head must load after checkpoint-driven untying.""" + + class StubGemmaModel(torch.nn.Module): + def __init__(self, *, vllm_config, prefix): + super().__init__() + self.embed_tokens = VocabParallelEmbedding(4, 2) + self.make_empty_intermediate_tensors = None + + monkeypatch.setattr(gemma, "GemmaModel", StubGemmaModel) + config = SimpleNamespace( + vocab_size=4, + hidden_size=2, + tie_word_embeddings=False, + ) + vllm_config = SimpleNamespace( + model_config=SimpleNamespace(hf_config=config), + quant_config=None, + ) + model = gemma.GemmaForCausalLM(vllm_config=cast(VllmConfig, vllm_config)) + embedding_weight = torch.full((4, 2), 1.0) + lm_head_weight = torch.full((4, 2), 2.0) + + loaded = model.load_weights( + [ + ("model.embed_tokens.weight", embedding_weight), + ("lm_head.weight", lm_head_weight), + ] + ) + + assert loaded == {"model.embed_tokens.weight", "lm_head.weight"} + assert torch.equal(model.model.embed_tokens.weight[:4], embedding_weight) + assert torch.equal(model.lm_head.weight[:4], lm_head_weight) + + @pytest.mark.parametrize("model", MODELS) def test_dummy_loader(vllm_runner, monkeypatch, model: str) -> None: with monkeypatch.context() as m: diff --git a/vllm/model_executor/models/gemma.py b/vllm/model_executor/models/gemma.py index 0f1282fd969c..7058290865d9 100644 --- a/vllm/model_executor/models/gemma.py +++ b/vllm/model_executor/models/gemma.py @@ -41,7 +41,10 @@ from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.rotary_embedding import get_rope -from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) from vllm.sequence import IntermediateTensors from .interfaces import SupportsLoRA, SupportsPP, SupportsQuant @@ -351,6 +354,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.model = GemmaModel( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + if config.tie_word_embeddings: + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) self.logits_processor = LogitsProcessor(config.vocab_size) self.make_empty_intermediate_tensors = ( self.model.make_empty_intermediate_tensors @@ -375,7 +386,7 @@ def compute_logits( self, hidden_states: torch.Tensor, ) -> torch.Tensor | None: - logits = self.logits_processor(self.model.embed_tokens, hidden_states) + logits = self.logits_processor(self.lm_head, hidden_states) return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: From 91a893de64722019ea2faf852e06cabe143b3490 Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Thu, 20 Aug 2026 21:42:08 -0400 Subject: [PATCH 229/839] [Bugfix][Spec Decode] Scope DSpark backend inheritance to DeepSeek V4 (#52809) Signed-off-by: mgoin Co-authored-by: Codex Co-authored-by: OpenAI Codex --- .../v1/worker/gpu/spec_decode/dspark/utils.py | 44 ++++++++++++++----- 1 file changed, 34 insertions(+), 10 deletions(-) diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py index 78ae392132ce..a33fa259c3e7 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py @@ -3,13 +3,32 @@ import torch.nn as nn -from vllm.config import VllmConfig, replace +from vllm.config import ModelConfig, VllmConfig, replace from vllm.distributed.parallel_state import get_pp_group -from vllm.model_executor.model_loader import get_model -from vllm.v1.worker.gpu.spec_decode.eagle.utils import ( - _should_share, - get_target_lm_head, -) +from vllm.logger import init_logger +from vllm.v1.attention.backends.registry import AttentionBackendEnum + +logger = init_logger(__name__) + + +def _resolve_dspark_attention_backend( + draft_model_config: ModelConfig, + draft_backend: AttentionBackendEnum | None, + target_backend: AttentionBackendEnum | None, +) -> AttentionBackendEnum | None: + if draft_backend is not None: + return draft_backend + # DeepSeek-V4 draft layers share the target's KV-cache layout. Other + # DSpark architectures may use a different attention kind. + if draft_model_config.hf_config.model_type == "deepseek_v4": + if target_backend is not None: + logger.info_once( + "Using the target model's %s attention backend for the " + "DeepSeek-V4 DSpark drafter.", + target_backend.name, + ) + return target_backend + return None def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module: @@ -18,13 +37,18 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo draft_model_config = speculative_config.draft_model_config from vllm.compilation.backends import set_model_tag + from vllm.model_executor.model_loader import get_model from vllm.model_executor.models.qwen3_dflash import dflash_has_any_non_causal from vllm.model_executor.models.utils import get_draft_quant_config + from vllm.v1.worker.gpu.spec_decode.eagle.utils import ( + _should_share, + get_target_lm_head, + ) - # None re-runs backend auto-selection for the draft, which can pick a - # different attention class than the target; fall back to the target's. - draft_attention_backend = ( - speculative_config.attention_backend or vllm_config.attention_config.backend + draft_attention_backend = _resolve_dspark_attention_backend( + draft_model_config, + speculative_config.attention_backend, + vllm_config.attention_config.backend, ) draft_vllm_config = replace( From 83c5d59209d2a3629cafaa1ceca0893530442835 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Fri, 21 Aug 2026 11:43:51 +1000 Subject: [PATCH 230/839] [Rust Frontend] Replace external `protoc` with pure Rust lib `protox` (#52892) Signed-off-by: Bugen Zhao Co-authored-by: OpenAI Codex Co-authored-by: Nick Hill --- .buildkite/ci_config_rocm.yaml | 1 - .buildkite/scripts/build-macos-wheel.sh | 5 - .buildkite/scripts/ci-bake-rocm.sh | 4 +- .../scripts/run-rust-frontend-cargo-ci.sh | 30 ---- docker/Dockerfile | 9 +- docker/Dockerfile.cpu | 7 +- docker/Dockerfile.rocm | 18 +- docker/Dockerfile.rocm_gfx1250 | 10 +- docker/Dockerfile.xpu | 7 +- rust/Cargo.lock | 155 +++++++++++++++++- rust/Cargo.toml | 1 + rust/src/server/Cargo.toml | 1 + rust/src/server/build.rs | 21 ++- tools/install_protoc.sh | 37 ----- 14 files changed, 181 insertions(+), 125 deletions(-) delete mode 100755 tools/install_protoc.sh diff --git a/.buildkite/ci_config_rocm.yaml b/.buildkite/ci_config_rocm.yaml index cc7c2ad27f64..c2da292a838a 100644 --- a/.buildkite/ci_config_rocm.yaml +++ b/.buildkite/ci_config_rocm.yaml @@ -25,7 +25,6 @@ run_all_patterns: - "rust/" - "rust-toolchain.toml" - "tools/build_rust.py" - - "tools/install_protoc.sh" - "tools/install_torchcodec_rocm.sh" - "tests/vllm_test_utils/" - "vllm/envs.py" diff --git a/.buildkite/scripts/build-macos-wheel.sh b/.buildkite/scripts/build-macos-wheel.sh index ac0e4eb1d76e..cb95d92ff14f 100755 --- a/.buildkite/scripts/build-macos-wheel.sh +++ b/.buildkite/scripts/build-macos-wheel.sh @@ -10,11 +10,6 @@ set -euo pipefail # The macmini queue uses persistent checkouts, so refresh tags for setuptools-scm. git fetch --tags --force origin -# The Rust frontend build needs protoc. -if ! command -v protoc >/dev/null 2>&1; then - brew install protobuf -fi - # upload-nightly-wheels.sh expects exactly one wheel. rm -rf artifacts/dist mkdir -p artifacts/dist diff --git a/.buildkite/scripts/ci-bake-rocm.sh b/.buildkite/scripts/ci-bake-rocm.sh index 614987f6d148..ef38b8be5824 100644 --- a/.buildkite/scripts/ci-bake-rocm.sh +++ b/.buildkite/scripts/ci-bake-rocm.sh @@ -15,7 +15,7 @@ set -euo pipefail DEFAULT_REPO_SLUG="vllm-project/vllm" DEFAULT_CI_HCL_SOURCE="docker/ci-rocm.hcl" -DEFAULT_CI_BASE_CONTENT_FILES=".dockerignore requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt tools/install_torchcodec_rocm.sh tools/install_protoc.sh rust-toolchain.toml tests/vllm_test_utils" +DEFAULT_CI_BASE_CONTENT_FILES=".dockerignore requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt tools/install_torchcodec_rocm.sh rust-toolchain.toml tests/vllm_test_utils" DEFAULT_CI_BASE_DOCKERFILE="docker/Dockerfile.rocm" DEFAULT_CI_BASE_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust-toolchain-input rust-toolchain build_nixl lmcache_source build_lmcache build_rocshmem build_deepep mori_base ci_base" DEFAULT_CI_BASE_METADATA_VERSION="3" @@ -23,7 +23,7 @@ DEFAULT_CI_BASE_METADATA_VERSION="3" # local-source stages rather than unreachable remote-fetch alternatives. DEFAULT_ROCM_CSRC_CONTENT_FILES=".dockerignore requirements/common.txt requirements/rocm.txt pyproject.toml setup.py CMakeLists.txt cmake csrc vllm/envs.py vllm/__init__.py tools/build_rust.py" DEFAULT_ROCM_CSRC_DOCKERFILE_STAGES="base fetch_vllm_0 fetch_vllm build_vllm_dependencies rocm-triton-kernels csrc-build" -DEFAULT_ROCM_RUST_CONTENT_FILES=".dockerignore .git_archival.txt pyproject.toml requirements/build/rust.txt rust/Cargo.lock rust/Cargo.toml rust/proto rust/src rust-toolchain.toml tools/build_rust.py tools/install_protoc.sh build_rust.sh" +DEFAULT_ROCM_RUST_CONTENT_FILES=".dockerignore .git_archival.txt pyproject.toml requirements/build/rust.txt rust/Cargo.lock rust/Cargo.toml rust/proto rust/src rust-toolchain.toml tools/build_rust.py build_rust.sh" DEFAULT_ROCM_RUST_DOCKERFILE_STAGES="base fetch_vllm_0 fetch_vllm vllm-version rust_toolchain_input_0 rust-toolchain-input rust_input_0 rust-input rust-toolchain rust-build" # Docker's 128-character tag limit minus the longest cache prefix # ("csrc-rocm-branch-" and "rust-rocm-branch-", both 17 characters). diff --git a/.buildkite/scripts/run-rust-frontend-cargo-ci.sh b/.buildkite/scripts/run-rust-frontend-cargo-ci.sh index 215650a07fbb..62a70fd66450 100755 --- a/.buildkite/scripts/run-rust-frontend-cargo-ci.sh +++ b/.buildkite/scripts/run-rust-frontend-cargo-ci.sh @@ -21,7 +21,6 @@ export CARGO_HOME="${CARGO_HOME:-$HOME/.cargo}" export RUSTUP_HOME="${RUSTUP_HOME:-$HOME/.rustup}" export PATH="$CARGO_HOME/bin:$PATH" -PROTOC_VERSION="${PROTOC_VERSION:-31.1}" CARGO_BINSTALL_VERSION="${CARGO_BINSTALL_VERSION:-1.20.1}" UV_VERSION="${UV_VERSION:-0.11.28}" PYO3_PYTHON_VERSION="${PYO3_PYTHON_VERSION:-3.12}" @@ -34,34 +33,6 @@ log_section() { echo "--- $*" } -install_protoc() { - local arch - case "$(uname -m)" in - x86_64) - arch="x86_64" - ;; - aarch64|arm64) - arch="aarch_64" - ;; - *) - echo "Unsupported protoc architecture: $(uname -m)" >&2 - return 1 - ;; - esac - - local url="https://github.com/protocolbuffers/protobuf/releases/download/v${PROTOC_VERSION}/protoc-${PROTOC_VERSION}-linux-${arch}.zip" - local tmp_dir - tmp_dir="$(mktemp -d)" - - log_section "Installing protoc ${PROTOC_VERSION}" - curl -L --proto '=https' --tlsv1.2 -sSf "$url" -o "$tmp_dir/protoc.zip" - mkdir -p "$CARGO_HOME/bin" - unzip -q "$tmp_dir/protoc.zip" bin/protoc 'include/*' -d "$CARGO_HOME" - chmod +x "$CARGO_HOME/bin/protoc" - rm -rf "$tmp_dir" - protoc --version -} - rust_toolchain() { awk -F '"' '/channel[[:space:]]*=/ { print $2; exit }' rust-toolchain.toml } @@ -186,7 +157,6 @@ run_tests() { --no-fail-fast } -install_protoc install_rust_toolchain case "$MODE" in diff --git a/docker/Dockerfile b/docker/Dockerfile index 543f4b12813c..17ee9be31648 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -243,21 +243,18 @@ ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list} #################### RUST BUILD IMAGE #################### # Build the Rust frontend (`vllm-rs`) in a dedicated stage so the main wheel -# build stage doesn't need the rust toolchain, protoc, or the rust source. +# build stage doesn't need the rust toolchain or the rust source. # This stage reuses the Python environment from base and runs in parallel with # csrc-build/extensions-build. FROM base AS rust-build ARG USE_SCCACHE ARG SCCACHE_ENDPOINT -# Install native tools needed only for Rust/protoc builds. -RUN dnf install -y --setopt=install_weak_deps=False make unzip \ +# Install native tools needed only for Rust builds. +RUN dnf install -y --setopt=install_weak_deps=False make \ && dnf clean all \ && rm -rf /var/cache/dnf -COPY tools/install_protoc.sh /tmp/install_protoc.sh -RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh - WORKDIR /workspace COPY requirements/build/rust.txt requirements/build/rust.txt diff --git a/docker/Dockerfile.cpu b/docker/Dockerfile.cpu index 43fab4bccb16..744336d926d2 100644 --- a/docker/Dockerfile.cpu +++ b/docker/Dockerfile.cpu @@ -127,19 +127,16 @@ RUN echo 'ulimit -c 0' >> ~/.bashrc ######################### RUST BUILD IMAGE ######################### # Build the Rust frontend (`vllm-rs`) in a dedicated stage so the wheel build -# stage doesn't need the rust toolchain or protoc. This stage runs in parallel +# stage doesn't need the rust toolchain. This stage runs in parallel # with the main vllm-build stage. FROM base AS rust-build ENV DEBIAN_FRONTEND=noninteractive RUN apt-get update -y \ && apt-get install -y --no-install-recommends \ - build-essential unzip python3 python3-pip \ + build-essential python3 python3-pip \ && rm -rf /var/lib/apt/lists/* -COPY tools/install_protoc.sh /tmp/install_protoc.sh -RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh - WORKDIR /workspace COPY requirements/build/rust.txt requirements/build/rust.txt diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 0b02b9ec2153..5bfd99bc3a09 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -138,13 +138,12 @@ RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ > /vllm-version.txt # ----------------------- -# Rust/protoc toolchain inputs +# Rust toolchain inputs # # Keep these separate from Rust source inputs so Rust source-only changes do not # invalidate the long-lived ci_base image that carries the toolchain. FROM scratch AS rust_toolchain_input_0 COPY rust-toolchain.toml /rust-toolchain-input/vllm/rust-toolchain.toml -COPY tools/install_protoc.sh /rust-toolchain-input/vllm/tools/install_protoc.sh FROM base AS rust_toolchain_input_1 ARG VLLM_REPO @@ -155,7 +154,6 @@ RUN git clone --no-checkout --filter=blob:none ${VLLM_REPO} /rust-toolchain-inpu && git sparse-checkout init --no-cone \ && git sparse-checkout set \ rust-toolchain.toml \ - tools/install_protoc.sh \ && git checkout FETCH_HEAD FROM rust_toolchain_input_${REMOTE_VLLM} AS rust-toolchain-input @@ -195,21 +193,17 @@ RUN git clone --no-checkout --filter=blob:none ${VLLM_REPO} /rust-input/vllm \ FROM rust_input_${REMOTE_VLLM} AS rust-input # ----------------------- -# Rust/protoc toolchain +# Rust toolchain FROM base AS rust-toolchain ENV CARGO_HOME=/root/.cargo ENV RUSTUP_HOME=/root/.rustup ENV PATH=${CARGO_HOME}/bin:${PATH} -# protoc is used by tonic-build/prost-build. RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \ - ca-certificates curl unzip \ + ca-certificates curl \ && rm -rf /var/lib/apt/lists/* -COPY --from=rust-toolchain-input /rust-toolchain-input/vllm/tools/install_protoc.sh /tmp/install_protoc.sh -RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh - COPY --from=rust-toolchain-input /rust-toolchain-input/vllm/rust-toolchain.toml /tmp/rust-toolchain.toml RUN TOOLCHAIN="$(grep '^channel' /tmp/rust-toolchain.toml | sed 's/.*= *"\(.*\)"/\1/')" \ && if ! command -v rustup >/dev/null 2>&1; then \ @@ -224,7 +218,7 @@ RUN TOOLCHAIN="$(grep '^channel' /tmp/rust-toolchain.toml | sed 's/.*= *"\(.*\)" # ----------------------- # Rust build stage # Builds the `vllm-rs` frontend in a dedicated stage so the wheel build stages -# don't need the rust toolchain or protoc. +# don't need the rust toolchain. FROM rust-toolchain AS rust-build ARG COMMON_WORKDIR ARG USE_SCCACHE @@ -780,12 +774,10 @@ ENV CARGO_HOME=/root/.cargo ENV RUSTUP_HOME=/root/.rustup ENV PATH=${CARGO_HOME}/bin:${PATH} -# Bake the Rust/protoc toolchain into ci_base so Rust build/test steps do not +# Bake the Rust toolchain into ci_base so Rust build/test steps do not # need to fetch it on each per-commit image build. COPY --from=rust-toolchain /root/.cargo /root/.cargo COPY --from=rust-toolchain /root/.rustup /root/.rustup -COPY --from=rust-toolchain /usr/local/bin/protoc /usr/local/bin/protoc -COPY --from=rust-toolchain /usr/local/include/google /usr/local/include/google # Update rdma-core to support latest rocshmem. ARG DEEPEP_NIC diff --git a/docker/Dockerfile.rocm_gfx1250 b/docker/Dockerfile.rocm_gfx1250 index 6e88eb9fed9e..e1f2c75d97dd 100644 --- a/docker/Dockerfile.rocm_gfx1250 +++ b/docker/Dockerfile.rocm_gfx1250 @@ -141,18 +141,10 @@ RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ # ----------------------- # Rust build stage # Builds the `vllm-rs` frontend in a dedicated stage so the wheel build stages -# don't need the rust toolchain or protoc. +# don't need the rust toolchain. FROM fetch_vllm AS rust-build ARG COMMON_WORKDIR -# protoc is used by tonic-build/prost-build. -RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \ - ca-certificates curl unzip \ - && rm -rf /var/lib/apt/lists/* - -COPY tools/install_protoc.sh /tmp/install_protoc.sh -RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh - # Cap cargo parallelism to avoid exhausting the AMD CI host's open-file limit # (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise). ENV CARGO_BUILD_JOBS=4 diff --git a/docker/Dockerfile.xpu b/docker/Dockerfile.xpu index 0811d64dff26..12ebfa36fe0d 100644 --- a/docker/Dockerfile.xpu +++ b/docker/Dockerfile.xpu @@ -1,17 +1,14 @@ ######################### RUST BUILD IMAGE ######################### # Build the Rust frontend (`vllm-rs`) in a dedicated stage so the main image -# doesn't need the rust toolchain or protoc. Runs in parallel with vllm-base. +# doesn't need the rust toolchain. Runs in parallel with vllm-base. FROM ubuntu:22.04 AS rust-build ENV DEBIAN_FRONTEND=noninteractive RUN apt-get update -y \ && apt-get install -y --no-install-recommends \ - ca-certificates curl git build-essential unzip python3 python3-pip \ + ca-certificates curl git build-essential python3 python3-pip \ && rm -rf /var/lib/apt/lists/* -COPY tools/install_protoc.sh /tmp/install_protoc.sh -RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh - WORKDIR /workspace COPY requirements/build/rust.txt requirements/build/rust.txt diff --git a/rust/Cargo.lock b/rust/Cargo.lock index a8608e34f2fd..de6c97654b75 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -385,6 +385,12 @@ version = "1.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06" +[[package]] +name = "beef" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a8241f3ebb85c056b509d4327ad0358fbbba6ffb340bf388f26350aeda225b1" + [[package]] name = "bit-set" version = "0.5.3" @@ -661,7 +667,7 @@ dependencies = [ "encode_unicode", "libc", "once_cell", - "unicode-width", + "unicode-width 0.2.2", "windows-sys 0.61.2", ] @@ -1499,7 +1505,7 @@ version = "0.2.24" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cfe4fbac503b8d1f88e6676011885f34b7174f46e59956bba534ba83abded4df" dependencies = [ - "unicode-width", + "unicode-width 0.2.2", ] [[package]] @@ -2005,7 +2011,7 @@ checksum = "25470f23803092da7d239834776d653104d551bc4d7eacaf31e6837854b8e9eb" dependencies = [ "console", "portable-atomic", - "unicode-width", + "unicode-width 0.2.2", "unit-prefix", "web-time", ] @@ -2264,6 +2270,72 @@ version = "0.4.29" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" +[[package]] +name = "logos" +version = "0.15.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff472f899b4ec2d99161c51f60ff7075eeb3097069a36050d8037a6325eb8154" +dependencies = [ + "logos-derive 0.15.1", +] + +[[package]] +name = "logos" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eb2c55a318a87600ea870ff8c2012148b44bf18b74fad48d0f835c38c7d07c5f" +dependencies = [ + "logos-derive 0.16.1", +] + +[[package]] +name = "logos-codegen" +version = "0.15.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "192a3a2b90b0c05b27a0b2c43eecdb7c415e29243acc3f89cc8247a5b693045c" +dependencies = [ + "beef", + "fnv", + "lazy_static", + "proc-macro2", + "quote", + "regex-syntax", + "rustc_version", + "syn 2.0.117", +] + +[[package]] +name = "logos-codegen" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "58b3ffaa284e1350d017a57d04ada118c4583cf260c8fb01e0fe28a2e9cf8970" +dependencies = [ + "fnv", + "proc-macro2", + "quote", + "regex-automata", + "regex-syntax", + "syn 2.0.117", +] + +[[package]] +name = "logos-derive" +version = "0.15.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "605d9697bcd5ef3a42d38efc51541aa3d6a4a25f7ab6d1ed0da5ac632a26b470" +dependencies = [ + "logos-codegen 0.15.1", +] + +[[package]] +name = "logos-derive" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52d3a9855747c17eaf4383823f135220716ab49bea5fbea7dd42cc9a92f8aa31" +dependencies = [ + "logos-codegen 0.16.1", +] + [[package]] name = "macro_rules_attribute" version = "0.2.2" @@ -2375,6 +2447,28 @@ version = "0.3.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "38d1115007560874e373613744c6fba374c17688327a71c1476d1a5954cc857b" +[[package]] +name = "miette" +version = "7.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5f98efec8807c63c752b5bd61f862c165c115b0a35685bdcfd9238c7aeb592b7" +dependencies = [ + "cfg-if", + "miette-derive", + "unicode-width 0.1.14", +] + +[[package]] +name = "miette-derive" +version = "7.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db5b29714e950dbb20d5e6f74f9dcec4edbcc1067bb7f8ed198c097b8c1a818b" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "mimalloc" version = "0.1.52" @@ -3116,6 +3210,18 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "prost-reflect" +version = "0.16.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "01b80ea363c31af2de2b92e3c07ed1156628f7838c4afb4df75ee78a37fedbd1" +dependencies = [ + "logos 0.16.1", + "miette", + "prost", + "prost-types", +] + [[package]] name = "prost-types" version = "0.14.3" @@ -3125,6 +3231,33 @@ dependencies = [ "prost", ] +[[package]] +name = "protox" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4f25a07a73c6717f0b9bbbd685918f5df9815f7efba450b83d9c9dea41f0e3a1" +dependencies = [ + "bytes", + "miette", + "prost", + "prost-reflect", + "prost-types", + "protox-parse", + "thiserror 2.0.18", +] + +[[package]] +name = "protox-parse" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "072eee358134396a4643dff81cfff1c255c9fbd3fb296be14bdb6a26f9156366" +dependencies = [ + "logos 0.15.1", + "miette", + "prost-types", + "thiserror 2.0.18", +] + [[package]] name = "pulldown-cmark" version = "0.13.3" @@ -3623,6 +3756,15 @@ version = "2.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "357703d41365b4b27c590e3ed91eabb1b663f07c4c084095e60cbed4362dff0d" +[[package]] +name = "rustc_version" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfcb3a22ef46e85b45de6ee7e79d063319ebb6594faafcf1c225ea92ab6e9b92" +dependencies = [ + "semver", +] + [[package]] name = "rustfft" version = "6.4.1" @@ -5284,6 +5426,12 @@ version = "1.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "da36089a805484bcccfffe0739803392c8298778a2d2f09febf76fac5ad9025b" +[[package]] +name = "unicode-width" +version = "0.1.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7dd6e30e90baa6f72411720665d41d89b9a3d039dc45b8faea1ddd07f617f6af" + [[package]] name = "unicode-width" version = "0.2.2" @@ -5733,6 +5881,7 @@ dependencies = [ "openssl", "prost", "prost-types", + "protox", "rmp-serde", "rmpv", "serde", diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 5c4b47d5817e..3e0b7b4365f2 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -75,6 +75,7 @@ prometheus-client = "0.24.0" prometheus-client-derive-encode = "0.5.0" prost = "0.14.3" prost-types = "0.14.3" +protox = "0.9.1" pyo3 = "0.28.3" pythonize = "0.28.0" rand = "0.9.2" diff --git a/rust/src/server/Cargo.toml b/rust/src/server/Cargo.toml index f6e6f08b530f..cc0d278f2774 100644 --- a/rust/src/server/Cargo.toml +++ b/rust/src/server/Cargo.toml @@ -53,6 +53,7 @@ vllm-metrics.workspace = true vllm-text.workspace = true [build-dependencies] +protox.workspace = true tonic-prost-build.workspace = true [dev-dependencies] diff --git a/rust/src/server/build.rs b/rust/src/server/build.rs index 585c3c70b990..1a2ba8a6b095 100644 --- a/rust/src/server/build.rs +++ b/rust/src/server/build.rs @@ -3,19 +3,22 @@ fn main() -> Result<(), Box> { let manifest_dir = std::env::var("CARGO_MANIFEST_DIR").unwrap(); - let proto_dir = format!("{manifest_dir}/../../proto"); + let proto_dir = std::path::PathBuf::from(format!("{manifest_dir}/../../proto")); + let protos = [ + proto_dir.join("control.proto"), + proto_dir.join("inference.proto"), + ]; + + for proto in &protos { + println!("cargo:rerun-if-changed={}", proto.display()); + } + + let file_descriptor_set = protox::compile(&protos, [&proto_dir])?; tonic_prost_build::configure() .build_server(true) .build_client(true) - .protoc_arg("--experimental_allow_proto3_optional") // be compatible with old compilers - .compile_protos( - &[ - format!("{proto_dir}/control.proto"), - format!("{proto_dir}/inference.proto"), - ], - &[proto_dir], - )?; + .compile_fds(file_descriptor_set)?; Ok(()) } diff --git a/tools/install_protoc.sh b/tools/install_protoc.sh deleted file mode 100755 index a995fdb6f6a2..000000000000 --- a/tools/install_protoc.sh +++ /dev/null @@ -1,37 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -# Install a pinned protoc binary from upstream GitHub releases. -# -# Distro protobuf-compiler packages vary widely in version (e.g. -# AlmaLinux/RHEL 8 ships protoc 3.5, predating the -# --experimental_allow_proto3_optional flag the rust frontend's build.rs -# passes), so we pin the protoc version here instead. -# -# Override the version via the PROTOC_VERSION env var. -# Requires: curl, unzip, root privileges. - -if [[ $(id -u) -ne 0 ]]; then - echo "Must be run as root" >&2 - exit 1 -fi - -VERSION="${PROTOC_VERSION:-34.2}" - -ARCH="$(uname -m)" -case "${ARCH}" in - # protoc release archives use "aarch_64" (with an underscore), not - # "aarch64". Don't "fix" this. - aarch64|arm64) URL_ARCH="aarch_64" ;; - x86_64|amd64) URL_ARCH="x86_64" ;; - *) echo "Unsupported arch for protoc binary: ${ARCH}" >&2; exit 1 ;; -esac - -URL="https://github.com/protocolbuffers/protobuf/releases/download/v${VERSION}/protoc-${VERSION}-linux-${URL_ARCH}.zip" -TMPDIR="$(mktemp -d)" -trap 'rm -rf "${TMPDIR}"' EXIT - -echo "Downloading: ${URL}" -curl -fsSL -o "${TMPDIR}/protoc.zip" "${URL}" -unzip -q -o "${TMPDIR}/protoc.zip" -d /usr/local -echo "Installed $(protoc --version)" From c6e19b3be24338759a443e03c8325d76da9ee202 Mon Sep 17 00:00:00 2001 From: Chauncey Date: Fri, 21 Aug 2026 10:01:06 +0800 Subject: [PATCH 231/839] [Bugfix][Structured Output] Avoid spurious FSM errors after speculative reasoning end (#53046) Signed-off-by: chaunceyjiang --- tests/v1/spec_decode/test_mtp_structured_output.py | 3 ++- vllm/v1/structured_output/__init__.py | 7 ++++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/v1/spec_decode/test_mtp_structured_output.py b/tests/v1/spec_decode/test_mtp_structured_output.py index 8bd599733a24..1fce9037d49e 100644 --- a/tests/v1/spec_decode/test_mtp_structured_output.py +++ b/tests/v1/spec_decode/test_mtp_structured_output.py @@ -189,7 +189,7 @@ def is_reasoning_end_streaming(self, input_ids, delta_ids): @pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) -def test_bitmask_post_reasoning_end_drafts_skip_grammar_advance(backend): +def test_bitmask_post_reasoning_end_drafts_skip_grammar_advance(backend, caplog): """Post-marker drafts predate the bitmask and may be grammar-invalid; grammar_bitmask must skip the grammar advance instead of asserting. """ @@ -235,6 +235,7 @@ def is_reasoning_end_streaming(self, input_ids, delta_ids): assert not (bitmask[2] == -1).all() # Grammar must not have advanced through the unvalidated draft. assert not grammar.is_terminated() + assert "Failed to advance FSM" not in caplog.text @pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) diff --git a/vllm/v1/structured_output/__init__.py b/vllm/v1/structured_output/__init__.py index 2fe1399d4d32..d0ec080d84bf 100644 --- a/vllm/v1/structured_output/__init__.py +++ b/vllm/v1/structured_output/__init__.py @@ -326,7 +326,12 @@ def grammar_bitmask( advance_grammar = False post_reasoning_end_in_window = True if advance_grammar and not grammar.is_terminated(): - accepted = grammar.accept_tokens(req_id, [token]) + if post_reasoning_end_in_window: + accepted = bool(grammar.validate_tokens([token])) + if accepted: + accepted = grammar.accept_tokens(req_id, [token]) + else: + accepted = grammar.accept_tokens(req_id, [token]) if accepted: state_advancements += 1 elif not post_reasoning_end_in_window: From df3b3422b495a33fd10d0b2f06c051144d6bb2c2 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Fri, 21 Aug 2026 12:13:54 +1000 Subject: [PATCH 232/839] [Rust Frontend] Add HY3 unified parser and local XGrammar structural-tag builder (#53054) Signed-off-by: Bugen Zhao --- rust/src/chat/src/lib.rs | 2 +- rust/src/chat/src/parser/reasoning/mod.rs | 3 + rust/src/chat/src/parser/tool/mod.rs | 5 +- rust/src/chat/src/parser/unified.rs | 6 +- rust/src/chat/tests/roundtrip.rs | 67 ++++- rust/src/parser/src/reasoning/delimited.rs | 22 +- rust/src/parser/src/reasoning/hy_v3.rs | 51 ++++ rust/src/parser/src/reasoning/mod.rs | 2 + rust/src/parser/src/tool/hy_v3.rs | 204 +++++++++----- .../parser/src/tool/hy_v3/structural_tag.rs | 261 ++++++++++++++++++ rust/src/parser/src/tool/mod.rs | 2 +- rust/src/parser/src/unified/hy_v3.rs | 207 ++++++++++++++ rust/src/parser/src/unified/mod.rs | 2 + rust/src/text/src/backend/hf/config.rs | 39 ++- rust/src/tokenizer/src/hf.rs | 33 +++ rust/src/tokenizer/src/lib.rs | 9 + rust/src/tokenizer/src/test_utils.rs | 20 +- 17 files changed, 849 insertions(+), 86 deletions(-) create mode 100644 rust/src/parser/src/reasoning/hy_v3.rs create mode 100644 rust/src/parser/src/tool/hy_v3/structural_tag.rs create mode 100644 rust/src/parser/src/unified/hy_v3.rs diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs index c6b91679e027..53caf31c3a04 100644 --- a/rust/src/chat/src/lib.rs +++ b/rust/src/chat/src/lib.rs @@ -415,6 +415,6 @@ mod tests { ) .unwrap_err(); - expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, gemma4, glm45, inkling, kimi, kimi_k2, kimi_k3, minimax_m2, minimax_m3, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string()); + expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, gemma4, glm45, hy_v3, inkling, kimi, kimi_k2, kimi_k3, minimax_m2, minimax_m3, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string()); } } diff --git a/rust/src/chat/src/parser/reasoning/mod.rs b/rust/src/chat/src/parser/reasoning/mod.rs index f85c4a930d56..c8e1c2df3955 100644 --- a/rust/src/chat/src/parser/reasoning/mod.rs +++ b/rust/src/chat/src/parser/reasoning/mod.rs @@ -25,6 +25,7 @@ pub mod names { pub const GEMMA4: &str = "gemma4"; pub const INKLING: &str = "inkling"; pub const GLM45: &str = "glm45"; + pub const HY_V3: &str = "hy_v3"; pub const KIMI: &str = "kimi"; pub const KIMI_K2: &str = "kimi_k2"; pub const KIMI_K3: &str = "kimi_k3"; @@ -67,6 +68,7 @@ impl ReasoningParserFactory { .register_unified_dummy(names::GEMMA4) .register_unified_dummy(names::INKLING) .register_parser::(names::GLM45) + .register_unified_dummy(names::HY_V3) .register_parser::(names::KIMI) .register_parser::(names::KIMI_K2) .register_unified_dummy(names::KIMI_K3) @@ -91,6 +93,7 @@ impl ReasoningParserFactory { .register_pattern("glm-4.7", names::GLM45) .register_pattern("glm-4.6", names::GLM45) .register_pattern("glm-4.5", names::GLM45) + .register_pattern("hy3", names::HY_V3) .register_pattern("kimi-k2", names::KIMI_K2) .register_pattern("kimi", names::KIMI) // step3p5 patterns must precede `step3`: substring matching would diff --git a/rust/src/chat/src/parser/tool/mod.rs b/rust/src/chat/src/parser/tool/mod.rs index cc8b160d937e..580ede001929 100644 --- a/rust/src/chat/src/parser/tool/mod.rs +++ b/rust/src/chat/src/parser/tool/mod.rs @@ -7,7 +7,7 @@ use std::sync::{Arc, LazyLock}; pub use vllm_parser::tool::{ DeepSeekV3ToolParser, DeepSeekV4ToolParser, DeepSeekV31ToolParser, DeepSeekV32ToolParser, - Glm45MoeToolParser, Glm47MoeToolParser, Granite4ToolParser, HermesToolParser, HyV3ToolParser, + Glm45MoeToolParser, Glm47MoeToolParser, Granite4ToolParser, HermesToolParser, Internlm2ToolParser, KimiK2ToolParser, Llama3JsonToolParser, MinimaxM2ToolParser, MinimaxM3ToolParser, MistralToolParser, Phi4MiniJsonToolParser, Qwen3CoderToolParser, Qwen3XmlToolParser, SeedOssToolParser, ToolParser, ToolParserError, @@ -76,7 +76,7 @@ impl ToolParserFactory { .register_unified_dummy(names::INKLING) .register_parser::(names::GRANITE4) .register_parser::(names::HERMES) - .register_parser::(names::HY_V3) + .register_unified_dummy(names::HY_V3) .register_parser::(names::INTERNLM) .register_parser::(names::KIMI_K2) .register_unified_dummy(names::KIMI_K3) @@ -100,7 +100,6 @@ impl ToolParserFactory { .register_pattern("qwen3", names::QWEN3_XML) .register_pattern("hermes", names::HERMES) .register_pattern("hy3", names::HY_V3) - .register_pattern("hy_v3", names::HY_V3) // Narrow to `internlm2` substring so it matches `internlm2-chat-7b` // and `internlm2_5-7b-chat` but NOT `internlm-chat-7b` (InternLM v1, // routes to Llama), `internlm3-*` (also Llama-architecture per diff --git a/rust/src/chat/src/parser/unified.rs b/rust/src/chat/src/parser/unified.rs index 6246733db281..ccd0fff637bb 100644 --- a/rust/src/chat/src/parser/unified.rs +++ b/rust/src/chat/src/parser/unified.rs @@ -6,7 +6,8 @@ use std::sync::LazyLock; pub use vllm_parser::unified::{ - Gemma4UnifiedParser, InklingUnifiedParser, KimiK3UnifiedParser, UnifiedParser, + Gemma4UnifiedParser, HyV3UnifiedParser, InklingUnifiedParser, KimiK3UnifiedParser, + UnifiedParser, }; use vllm_tokenizer::DynTokenizer; @@ -16,6 +17,7 @@ use crate::request::ChatTool; /// Canonical public names for registered unified parsers. pub mod names { pub const GEMMA4: &str = "gemma4"; + pub const HY_V3: &str = "hy_v3"; pub const INKLING: &str = "inkling"; pub const KIMI_K3: &str = "kimi_k3"; } @@ -41,12 +43,14 @@ impl UnifiedParserFactory { let mut factory = Self::default(); factory.register_parser::(names::GEMMA4); + factory.register_parser::(names::HY_V3); factory.register_parser::(names::INKLING); factory.register_parser::(names::KIMI_K3); factory .register_pattern("gemma-4", names::GEMMA4) .register_pattern("gemma4", names::GEMMA4) + .register_pattern("hy3", names::HY_V3) .register_pattern("inkling", names::INKLING) .register_pattern("kimi-k3", names::KIMI_K3) .register_pattern("kimi_k3", names::KIMI_K3); diff --git a/rust/src/chat/tests/roundtrip.rs b/rust/src/chat/tests/roundtrip.rs index 049dc6a40308..e9bc14a5e45e 100644 --- a/rust/src/chat/tests/roundtrip.rs +++ b/rust/src/chat/tests/roundtrip.rs @@ -55,6 +55,12 @@ enum ThinkingBehavior { Toggleable { default: bool }, /// The chat template always behaves as `value` for this fixture. Always { value: bool }, + /// The chat template selects thinking mode through `reasoning_effort`. + ReasoningEffort { + default: bool, + enabled: &'static str, + disabled: &'static str, + }, } impl ThinkingBehavior { @@ -62,6 +68,7 @@ impl ThinkingBehavior { match self { Self::Toggleable { default } => default, Self::Always { value } => value, + Self::ReasoningEffort { default, .. } => default, } } @@ -76,6 +83,33 @@ impl ThinkingBehavior { Some(value), // explicitly request the supported thinking behavior None, // use default template behavior ], + Self::ReasoningEffort { .. } => vec![ + Some(true), // explicitly enable thinking + Some(false), // explicitly disable thinking + None, // use default template behavior + ], + } + } + + fn apply(self, request: &mut ChatRequest, thinking: Option) { + let Some(thinking) = thinking else { + return; + }; + + match self { + Self::Toggleable { .. } | Self::Always { .. } => { + for key in ["thinking", "enable_thinking"] { + request.chat_options.template_kwargs.insert(key.to_string(), thinking.into()); + } + } + Self::ReasoningEffort { + enabled, disabled, .. + } => { + request.chat_options.template_kwargs.insert( + "reasoning_effort".to_string(), + if thinking { enabled } else { disabled }.into(), + ); + } } } } @@ -242,6 +276,23 @@ impl RoundtripCase { } } + /// HY3 suffixed reasoning/tool markers discovered from tokenizer added vocab. + fn hy_v3() -> Self { + Self { + model_id: "tencent/Hy3", + assistant_stop_suffix: "<｜hy_eos:opensource｜>", + tool_call_parser: ParserSelection::Auto, + reasoning_parser: ParserSelection::Auto, + thinking_behavior: ThinkingBehavior::ReasoningEffort { + default: false, + enabled: "high", + disabled: "no_think", + }, + json_fmt: compact_json_fmt(), + sort_json_keys: false, + } + } + /// SeedOSS with `` / `` reasoning tags. fn seed_oss() -> Self { Self { @@ -345,6 +396,9 @@ roundtrip_tests! { kimi_k25 => [tool_call_mix], // Kimi K2.5 strips reasoning in history // K3 drops plain-assistant reasoning in history; tool-call turns keep it. kimi_k3 => [tool_call_mix], + // HY3's final plain-assistant history omits EOS; tool-call history keeps it + // and exercises both stages of the tokenizer-derived unified parser. + hy_v3 => [tool_call_mix], gpt_oss => [tool_call_mix], // Harmony strips reasoning in history if there's no tool call inkling => [reasoning_and_content, tool_call_mix], } @@ -370,6 +424,7 @@ async fn run_roundtrip_reasoning_and_content_inner( vec![ChatMessage::text(ChatRole::User, "What is 2 + 2?")], Vec::new(), thinking, + case.thinking_behavior, ); let expected_reasoning = "Need compute 2 + 2 directly."; let expected_text = "The answer is 4."; @@ -417,6 +472,7 @@ async fn run_roundtrip_tool_call_mix( )], test_tools(), Some(true), // always enable thinking in this fixture + case.thinking_behavior, ); let expected_reasoning = "Need call the weather and add tools."; let expected_text = "I will call the tools."; @@ -838,6 +894,7 @@ fn roundtrip_request( messages: Vec, tools: Vec, thinking: Option, + thinking_behavior: ThinkingBehavior, ) -> ChatRequest { let tool_context = vllm_chat::ResolvedToolContext::new(&messages, tools, None, true) .expect("tool context should resolve"); @@ -848,13 +905,9 @@ fn roundtrip_request( ..ChatRequest::for_test() }; - // Explicitly enable or disable thinking so that rendering and parsing the reasoning block is - // exercised or skipped in the roundtrip. If unspecified, use the default template behavior. - if let Some(thinking) = thinking { - for key in ["thinking", "enable_thinking"] { - request.chat_options.template_kwargs.insert(key.to_string(), thinking.into()); - } - } + // Explicitly enable or disable thinking using the controls understood by + // this model's template. If unspecified, use the template default. + thinking_behavior.apply(&mut request, thinking); request } diff --git a/rust/src/parser/src/reasoning/delimited.rs b/rust/src/parser/src/reasoning/delimited.rs index 51cb4eaf9173..78213d291a48 100644 --- a/rust/src/parser/src/reasoning/delimited.rs +++ b/rust/src/parser/src/reasoning/delimited.rs @@ -35,25 +35,29 @@ impl DelimitedReasoningParser { /// start or end delimiter, that prompt boundary always wins. pub(crate) fn new( tokenizer: DynTokenizer, - start_token: &'static str, - end_token: &'static str, + start_token: impl Into, + end_token: impl Into, default_in_reasoning: bool, ) -> Result { + let start_token = start_token.into(); + let end_token = end_token.into(); let start_token_id = - tokenizer.token_to_id(start_token).ok_or_else(|| ReasoningError::MissingToken { - token: start_token.to_string(), - })?; + tokenizer + .token_to_id(&start_token) + .ok_or_else(|| ReasoningError::MissingToken { + token: start_token.clone(), + })?; let end_token_id = - tokenizer.token_to_id(end_token).ok_or_else(|| ReasoningError::MissingToken { - token: end_token.to_string(), + tokenizer.token_to_id(&end_token).ok_or_else(|| ReasoningError::MissingToken { + token: end_token.clone(), })?; Ok(Self { tokenizer, current_in_reasoning: default_in_reasoning, buffer: String::new(), - start_token: start_token.to_string(), - end_token: end_token.to_string(), + start_token, + end_token, start_token_id, end_token_id, default_in_reasoning, diff --git a/rust/src/parser/src/reasoning/hy_v3.rs b/rust/src/parser/src/reasoning/hy_v3.rs new file mode 100644 index 000000000000..f25f4a22e06c --- /dev/null +++ b/rust/src/parser/src/reasoning/hy_v3.rs @@ -0,0 +1,51 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use vllm_tokenizer::DynTokenizer; + +use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningError, ReasoningParser, Result}; + +/// Internal HY3 reasoning stage used by the unified HY3 parser. +pub(crate) struct HyV3ReasoningParser { + inner: DelimitedReasoningParser, +} + +impl HyV3ReasoningParser { + /// Create a HY3 reasoning parser for the tokenizer-specific marker suffix. + pub(crate) fn new(tokenizer: DynTokenizer, suffix: &str) -> Result { + Ok(Self { + inner: DelimitedReasoningParser::new( + tokenizer, + format!(""), + format!(""), + false, + )?, + }) + } +} + +impl ReasoningParser for HyV3ReasoningParser { + // Suffix discovery belongs to `HyV3UnifiedParser` so its reasoning and + // tool delimiters always use the same tokenizer-derived value. + fn create(_tokenizer: DynTokenizer) -> Result> + where + Self: Sized + 'static, + { + Err(ReasoningError::DummyUnifiedParser { + name: "hy_v3".to_string(), + }) + } + + fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> { + self.inner.initialize(prompt_token_ids); + Ok(()) + } + + fn push(&mut self, delta: &str) -> Result { + Ok(self.inner.push(delta)) + } + + fn finish(&mut self) -> Result { + Ok(self.inner.finish()) + } +} diff --git a/rust/src/parser/src/reasoning/mod.rs b/rust/src/parser/src/reasoning/mod.rs index e8dda6631a21..7588fe8f589c 100644 --- a/rust/src/parser/src/reasoning/mod.rs +++ b/rust/src/parser/src/reasoning/mod.rs @@ -20,6 +20,7 @@ mod cohere_cmd; mod deepseek_r1; mod delimited; +mod hy_v3; mod kimi; mod minimax_m3; mod qwen3; @@ -32,6 +33,7 @@ use vllm_tokenizer::DynTokenizer; pub use self::cohere_cmd::CohereCmdReasoningParser; pub use self::deepseek_r1::DeepSeekR1ReasoningParser; pub(crate) use self::delimited::{DelimitedReasoningParser, last_reasoning_boundary}; +pub(crate) use self::hy_v3::HyV3ReasoningParser; pub use self::kimi::KimiReasoningParser; pub use self::minimax_m3::MiniMaxM3ReasoningParser; pub use self::qwen3::Qwen3ReasoningParser; diff --git a/rust/src/parser/src/tool/hy_v3.rs b/rust/src/parser/src/tool/hy_v3.rs index a1426d6a3bd2..e289ab66ff05 100644 --- a/rust/src/parser/src/tool/hy_v3.rs +++ b/rust/src/parser/src/tool/hy_v3.rs @@ -1,26 +1,65 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project +mod structural_tag; + +use std::sync::Arc; + use winnow::ascii::multispace0 as ws0; use winnow::combinator::{alt, delimited, eof, repeat, seq, terminated}; use winnow::prelude::*; use winnow::stream::Partial; use winnow::token::{literal, rest, take_until}; +use self::structural_tag::HyV3StructuralTagBuilder; use super::parameters::ToolSchemas; use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker}; -use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput}; +use super::{Result, ToolCallDelta, ToolParser, ToolParserError, ToolParserOutput}; use crate::tool::{StructuralTagBuilder, Tool}; -const TOOL_CALLS_START: &str = ""; -const TOOL_CALLS_END: &str = ""; -const TOOL_CALL_START: &str = ""; -const TOOL_CALL_END: &str = ""; -const TOOL_SEP: &str = ""; -const ARG_KEY_START: &str = ""; -const ARG_KEY_END: &str = ""; -const ARG_VALUE_START: &str = ""; -const ARG_VALUE_END: &str = ""; +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct HyV3ToolMarkers { + tool_calls_start: String, + tool_calls_end: String, + tool_call_start: String, + tool_call_end: String, + tool_sep: String, + arg_key_start: String, + arg_key_end: String, + arg_value_start: String, + arg_value_end: String, +} + +impl HyV3ToolMarkers { + pub(crate) fn new(suffix: &str) -> Self { + Self { + tool_calls_start: format!(""), + tool_calls_end: format!(""), + tool_call_start: format!(""), + tool_call_end: format!(""), + tool_sep: format!(""), + arg_key_start: format!(""), + arg_key_end: format!(""), + arg_value_start: format!(""), + arg_value_end: format!(""), + } + } + + pub(crate) fn iter(&self) -> impl Iterator { + [ + self.tool_calls_start.as_str(), + self.tool_calls_end.as_str(), + self.tool_call_start.as_str(), + self.tool_call_end.as_str(), + self.tool_sep.as_str(), + self.arg_key_start.as_str(), + self.arg_key_end.as_str(), + self.arg_value_start.as_str(), + self.arg_value_end.as_str(), + ] + .into_iter() + } +} type HyV3Input<'i> = Partial<&'i str>; @@ -60,21 +99,26 @@ enum HyV3Event { /// Arguments are emitted only after a full `` block is parsed. /// HY3 marker tokens are added-vocabulary tokens rather than tokenizer special /// tokens, so the default `preserve_special_tokens() == false` is sufficient. -pub struct HyV3ToolParser { +pub(crate) struct HyV3ToolParser { buffer: String, mode: HyV3Mode, emitted_tool_count: usize, tool_parameters: ToolSchemas, + markers: Arc, + structural_tag_builder: HyV3StructuralTagBuilder, } impl HyV3ToolParser { /// Create a HY3 tool parser. - fn new(tools: &[Tool]) -> Self { + pub(crate) fn new(tools: &[Tool], suffix: &str) -> Self { + let markers = Arc::new(HyV3ToolMarkers::new(suffix)); Self { buffer: String::new(), mode: HyV3Mode::Text, emitted_tool_count: 0, tool_parameters: ToolSchemas::from_tools(tools), + markers: Arc::clone(&markers), + structural_tag_builder: HyV3StructuralTagBuilder::new(markers), } } @@ -109,22 +153,26 @@ impl HyV3ToolParser { } impl ToolParser for HyV3ToolParser { - fn create(tools: &[Tool]) -> Result> + // Suffix discovery belongs to `HyV3UnifiedParser` so its reasoning and + // tool delimiters always use the same tokenizer-derived value. + fn create(_tools: &[Tool]) -> Result> where Self: Sized + 'static, { - Ok(Box::new(Self::new(tools))) + Err(ToolParserError::DummyUnifiedParser { + name: "hy_v3".to_string(), + }) } fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> { - Some(xgrammar_structural_tag::Model::HyV3.builder()) + Some(&self.structural_tag_builder) } fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> { self.buffer.push_str(chunk); while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| { - parse_next_hy_v3_event(input, &mut self.mode) + parse_next_hy_v3_event(input, &mut self.mode, &self.markers) })? { self.apply_event(event, output)?; self.buffer.drain(..consumed_len); @@ -155,62 +203,83 @@ impl ToolParser for HyV3ToolParser { fn parse_next_hy_v3_event( input: &mut HyV3Input<'_>, mode: &mut HyV3Mode, + markers: &HyV3ToolMarkers, ) -> ModalResult { match mode { - HyV3Mode::Text => parse_text_event(input), + HyV3Mode::Text => parse_text_event(input, markers), HyV3Mode::ToolBlock { tool_call_end_scan } => { - parse_tool_block_event(input, tool_call_end_scan) + parse_tool_block_event(input, tool_call_end_scan, markers) } HyV3Mode::Done => ignored_rest_event(input), } } /// Parse a text-mode HY3 event. -fn parse_text_event(input: &mut HyV3Input<'_>) -> ModalResult { - alt((tool_block_start_event, safe_text_event)).parse_next(input) +fn parse_text_event( + input: &mut HyV3Input<'_>, + markers: &HyV3ToolMarkers, +) -> ModalResult { + alt(( + |input: &mut HyV3Input<'_>| tool_block_start_event(input, markers), + |input: &mut HyV3Input<'_>| safe_text_event(input, markers), + )) + .parse_next(input) } /// Parse a HY3 tool-block start marker. -fn tool_block_start_event(input: &mut HyV3Input<'_>) -> ModalResult { - literal(TOOL_CALLS_START).value(HyV3Event::ToolBlockStart).parse_next(input) +fn tool_block_start_event( + input: &mut HyV3Input<'_>, + markers: &HyV3ToolMarkers, +) -> ModalResult { + literal(markers.tool_calls_start.as_str()) + .value(HyV3Event::ToolBlockStart) + .parse_next(input) } /// Parse a safe text run before the next HY3 marker. -fn safe_text_event(input: &mut HyV3Input<'_>) -> ModalResult { - safe_text_len(input, TOOL_CALLS_START).map(|len| HyV3Event::Text { len }) +fn safe_text_event(input: &mut HyV3Input<'_>, markers: &HyV3ToolMarkers) -> ModalResult { + safe_text_len(input, &markers.tool_calls_start).map(|len| HyV3Event::Text { len }) } /// Parse one event inside a HY3 tool block. fn parse_tool_block_event( input: &mut HyV3Input<'_>, tool_call_end_scan: &mut MarkerScanState, + markers: &HyV3ToolMarkers, ) -> ModalResult { - alt((tool_block_end_event, |input: &mut HyV3Input<'_>| { - tool_call_event(input, tool_call_end_scan) - })) + alt(( + |input: &mut HyV3Input<'_>| tool_block_end_event(input, markers), + |input: &mut HyV3Input<'_>| tool_call_event(input, tool_call_end_scan, markers), + )) .parse_next(input) } /// Parse a HY3 tool-block end marker. -fn tool_block_end_event(input: &mut HyV3Input<'_>) -> ModalResult { - (ws0, literal(TOOL_CALLS_END)).value(HyV3Event::ToolBlockEnd).parse_next(input) +fn tool_block_end_event( + input: &mut HyV3Input<'_>, + markers: &HyV3ToolMarkers, +) -> ModalResult { + (ws0, literal(markers.tool_calls_end.as_str())) + .value(HyV3Event::ToolBlockEnd) + .parse_next(input) } /// Parse a complete HY3 tool-call block. fn tool_call_event( input: &mut HyV3Input<'_>, tool_call_end_scan: &mut MarkerScanState, + markers: &HyV3ToolMarkers, ) -> ModalResult { let (name, body) = seq!( _: ws0, - _: literal(TOOL_CALL_START), - take_until(0.., TOOL_SEP), - _: literal(TOOL_SEP), - take_until_marker(TOOL_CALL_END, tool_call_end_scan), - _: literal(TOOL_CALL_END), + _: literal(markers.tool_call_start.as_str()), + take_until(0.., markers.tool_sep.as_str()), + _: literal(markers.tool_sep.as_str()), + take_until_marker(markers.tool_call_end.as_str(), tool_call_end_scan), + _: literal(markers.tool_call_end.as_str()), ) .parse_next(input)?; - let raw_params = parse_tool_call_params(body)?; + let raw_params = parse_tool_call_params(body, markers)?; Ok(HyV3Event::ToolCall { name: name.trim().to_string(), @@ -219,21 +288,32 @@ fn tool_call_event( } /// Parse all parameter blocks inside a complete HY3 tool call. -fn parse_tool_call_params(tool_call_body: &str) -> ModalResult> { +fn parse_tool_call_params( + tool_call_body: &str, + markers: &HyV3ToolMarkers, +) -> ModalResult> { let mut input = tool_call_body; - delimited(ws0, repeat(0.., terminated(parameter, ws0)), eof).parse_next(&mut input) + delimited( + ws0, + repeat( + 0.., + terminated(|input: &mut &str| parameter(input, markers), ws0), + ), + eof, + ) + .parse_next(&mut input) } /// Parse a HY3 argument key/value block. -fn parameter(input: &mut &str) -> ModalResult<(String, String)> { +fn parameter(input: &mut &str, markers: &HyV3ToolMarkers) -> ModalResult<(String, String)> { let (name, value) = seq!( - _: literal(ARG_KEY_START), - take_until(0.., ARG_KEY_END), - _: literal(ARG_KEY_END), + _: literal(markers.arg_key_start.as_str()), + take_until(0.., markers.arg_key_end.as_str()), + _: literal(markers.arg_key_end.as_str()), _: ws0, - _: literal(ARG_VALUE_START), - take_until(0.., ARG_VALUE_END), - _: literal(ARG_VALUE_END), + _: literal(markers.arg_value_start.as_str()), + take_until(0.., markers.arg_value_end.as_str()), + _: literal(markers.arg_value_end.as_str()), ) .parse_next(input)?; @@ -274,14 +354,14 @@ mod tests { #[test] fn hy_v3_does_not_preserve_special_tokens() { - let parser = HyV3ToolParser::new(&test_tools()); + let parser = HyV3ToolParser::new(&test_tools(), ""); assert!(!parser.preserve_special_tokens()); } #[test] fn hy_v3_parse_complete_without_tool_call_keeps_text() { - let mut parser = HyV3ToolParser::new(&test_tools()); + let mut parser = HyV3ToolParser::new(&test_tools(), ""); let output = parser.parse_complete("This is a plain response.").unwrap(); assert_eq!(output.normal_text(), "This is a plain response."); @@ -290,7 +370,7 @@ mod tests { #[test] fn hy_v3_parse_complete_extracts_zero_arg_inline_tool_call() { - let mut parser = HyV3ToolParser::new(&test_tools()); + let mut parser = HyV3ToolParser::new(&test_tools(), ""); let output = parser .parse_complete( "get_current_date", @@ -305,7 +385,7 @@ mod tests { #[test] fn hy_v3_parse_complete_extracts_zero_arg_newline_tool_call() { - let mut parser = HyV3ToolParser::new(&test_tools()); + let mut parser = HyV3ToolParser::new(&test_tools(), ""); let output = parser .parse_complete( "\nget_current_date\n\n", @@ -318,7 +398,7 @@ mod tests { #[test] fn hy_v3_parse_complete_extracts_arguments_on_same_line() { - let mut parser = HyV3ToolParser::new(&test_tools()); + let mut parser = HyV3ToolParser::new(&test_tools(), ""); let output = parser .parse_complete( "get_weathercityBeijingdate2026-03-30", @@ -333,7 +413,7 @@ mod tests { #[test] fn hy_v3_parse_complete_extracts_arguments_with_newlines() { - let mut parser = HyV3ToolParser::new(&test_tools()); + let mut parser = HyV3ToolParser::new(&test_tools(), ""); let output = parser .parse_complete(&build_tool_calls(&[build_tool_call( "get_weather", @@ -349,7 +429,7 @@ mod tests { #[test] fn hy_v3_parse_complete_preserves_prefix_and_ignores_trailing_text() { - let mut parser = HyV3ToolParser::new(&test_tools()); + let mut parser = HyV3ToolParser::new(&test_tools(), ""); let output = parser .parse_complete(&format!( "Checking.{} trailing text", @@ -363,7 +443,7 @@ mod tests { #[test] fn hy_v3_parse_complete_extracts_multiple_tool_calls_in_one_block() { - let mut parser = HyV3ToolParser::new(&test_tools()); + let mut parser = HyV3ToolParser::new(&test_tools(), ""); let output = parser .parse_complete(&build_tool_calls(&[ build_tool_call( @@ -406,7 +486,7 @@ mod tests { #[test] fn hy_v3_parse_complete_converts_schema_types() { - let mut parser = HyV3ToolParser::new(&test_tools()); + let mut parser = HyV3ToolParser::new(&test_tools(), ""); let output = parser .parse_complete(&build_tool_calls(&[build_tool_call( "convert", @@ -432,7 +512,7 @@ mod tests { #[test] fn hy_v3_streaming_without_tool_call_emits_text_incrementally() { - let mut parser = HyV3ToolParser::new(&test_tools()); + let mut parser = HyV3ToolParser::new(&test_tools(), ""); let mut output = ToolParserOutput::default(); output.append(parser.parse_chunk("This is ").unwrap()); @@ -446,7 +526,7 @@ mod tests { #[test] fn hy_v3_streaming_extracts_zero_arg_tool_call() { - let mut parser = HyV3ToolParser::new(&test_tools()); + let mut parser = HyV3ToolParser::new(&test_tools(), ""); let chunks = [ "", "\n", @@ -465,7 +545,7 @@ mod tests { #[test] fn hy_v3_streaming_extracts_arguments() { - let mut parser = HyV3ToolParser::new(&test_tools()); + let mut parser = HyV3ToolParser::new(&test_tools(), ""); let chunks = [ "", "\n", @@ -491,7 +571,7 @@ mod tests { #[test] fn hy_v3_streaming_preserves_prefix_text() { - let mut parser = HyV3ToolParser::new(&test_tools()); + let mut parser = HyV3ToolParser::new(&test_tools(), ""); let chunks = [ "Checking.", "", @@ -521,7 +601,7 @@ mod tests { ), ]); let chunks = split_by_chars(&input, 9); - let mut parser = HyV3ToolParser::new(&test_tools()); + let mut parser = HyV3ToolParser::new(&test_tools(), ""); let output = collect_stream(&mut parser, &chunks); @@ -537,7 +617,7 @@ mod tests { build_tool_calls(&[build_tool_call("get_weather", &[("city", "Beijing")])]) ); let chunks = split_by_chars(&input, 5); - let mut parser = HyV3ToolParser::new(&test_tools()); + let mut parser = HyV3ToolParser::new(&test_tools(), ""); let output = collect_stream(&mut parser, &chunks); @@ -548,7 +628,7 @@ mod tests { #[test] fn hy_v3_streaming_does_not_emit_incomplete_tool_call() { - let mut parser = HyV3ToolParser::new(&test_tools()); + let mut parser = HyV3ToolParser::new(&test_tools(), ""); let mut output = ToolParserOutput::default(); parser @@ -564,7 +644,7 @@ mod tests { #[test] fn hy_v3_finish_fails_incomplete_tool_call() { - let mut parser = HyV3ToolParser::new(&test_tools()); + let mut parser = HyV3ToolParser::new(&test_tools(), ""); parser.parse_chunk("get_weather").unwrap(); let error = parser.finish().unwrap_err(); @@ -575,7 +655,7 @@ mod tests { #[test] fn hy_v3_malformed_tool_call_fails_fast() { - let mut parser = HyV3ToolParser::new(&test_tools()); + let mut parser = HyV3ToolParser::new(&test_tools(), ""); let error = parser .parse_complete( "get_weathercityBeijing", diff --git a/rust/src/parser/src/tool/hy_v3/structural_tag.rs b/rust/src/parser/src/tool/hy_v3/structural_tag.rs new file mode 100644 index 000000000000..ef182da267b8 --- /dev/null +++ b/rust/src/parser/src/tool/hy_v3/structural_tag.rs @@ -0,0 +1,261 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Structural-tag grammar for HY3 XML-style tool calls. + +use std::collections::HashSet; +use std::sync::Arc; + +use serde_json::Value; +use xgrammar_structural_tag::Result; +use xgrammar_structural_tag::builders::{StructuralTagBuilder, StructuralTagContext}; +use xgrammar_structural_tag::format::{Format, StructuralTag, TagFormat}; +use xgrammar_structural_tag::tool::{BuilderToolChoice, FunctionToolParam}; + +use super::HyV3ToolMarkers; + +/// HY3 structural-tag builder using tokenizer-specific structural markers. +#[derive(Debug, Clone)] +pub(super) struct HyV3StructuralTagBuilder { + markers: Arc, +} + +impl HyV3StructuralTagBuilder { + pub(super) fn new(markers: Arc) -> Self { + Self { markers } + } + + fn argument_pair(&self, key: &str) -> Format { + let excludes = self.markers.iter().collect::>(); + Format::sequence(vec![ + Format::const_string(&self.markers.arg_key_start), + Format::const_string(key), + Format::const_string(&self.markers.arg_key_end), + Format::const_string("\n"), + Format::const_string(&self.markers.arg_value_start), + Format::any_text_excluding(&excludes), + Format::const_string(&self.markers.arg_value_end), + Format::const_string("\n"), + ]) + } + + fn tool_call(&self, tool: &FunctionToolParam) -> TagFormat { + let (required_keys, optional_keys) = argument_keys(tool.function.parameters.as_ref()); + let mut elements = + required_keys.into_iter().map(|key| self.argument_pair(key)).collect::>(); + + if !optional_keys.is_empty() { + let mut pairs = + optional_keys.into_iter().map(|key| self.argument_pair(key)).collect::>(); + let optional = if pairs.len() == 1 { + pairs.pop().unwrap() + } else { + Format::or(pairs) + }; + elements.push(Format::star(optional)); + } + + let content = if elements.is_empty() { + Format::any_text() + } else { + Format::sequence(elements) + }; + TagFormat::new( + format!( + "{}{}{}\n", + self.markers.tool_call_start, tool.function.name, self.markers.tool_sep + ), + content, + self.markers.tool_call_end.clone(), + ) + } + + fn tool_calls(&self, tools: &[FunctionToolParam], choice: BuilderToolChoice) -> Format { + let mut calls = tools.iter().map(|tool| self.tool_call(tool)).collect::>(); + let begin = format!("{}\n", self.markers.tool_calls_start); + let end = format!("\n{}", self.markers.tool_calls_end); + + match choice { + BuilderToolChoice::Auto if calls.is_empty() => Format::any_text(), + BuilderToolChoice::Auto => { + let outer = TagFormat::new( + begin, + Format::tags_with_separator(calls, "\n", true, false), + end, + ); + Format::triggered_tags(&[&self.markers.tool_calls_start], vec![outer]) + } + BuilderToolChoice::Forced => Format::sequence(vec![ + Format::const_string(begin), + Format::Tag(calls.pop().unwrap()), + Format::const_string(end), + ]), + BuilderToolChoice::Required => Format::sequence(vec![ + Format::const_string(begin), + Format::tags_with_separator(calls, "\n", true, false), + Format::const_string(end), + ]), + } + } +} + +impl StructuralTagBuilder for HyV3StructuralTagBuilder { + fn build(&self, ctx: StructuralTagContext<'_>) -> Result { + Ok(StructuralTag::new( + self.tool_calls(ctx.function_tools, ctx.tool_choice), + )) + } +} + +/// Split argument keys into required and optional declaration-order groups. +fn argument_keys(parameters: Option<&Value>) -> (Vec<&str>, Vec<&str>) { + let Some(parameters) = parameters.and_then(Value::as_object) else { + return (Vec::new(), Vec::new()); + }; + let properties = parameters.get("properties").and_then(Value::as_object); + let required = parameters + .get("required") + .and_then(Value::as_array) + .into_iter() + .flatten() + .filter_map(Value::as_str) + .collect::>(); + let required_set = required.iter().copied().collect::>(); + + let mut required_keys = properties + .into_iter() + .flat_map(|properties| properties.keys()) + .map(String::as_str) + .filter(|key| required_set.contains(key)) + .collect::>(); + required_keys.extend( + required + .into_iter() + .filter(|key| properties.is_none_or(|properties| !properties.contains_key(*key))), + ); + let optional_keys = properties + .into_iter() + .flat_map(|properties| properties.keys()) + .map(String::as_str) + .filter(|key| !required_set.contains(key)) + .collect(); + (required_keys, optional_keys) +} + +#[cfg(test)] +mod tests { + use expect_test::expect; + use serde_json::{Value, json}; + use std::sync::Arc; + use xgrammar_structural_tag::builders::StructuralTagOptions; + use xgrammar_structural_tag::{ + FunctionDefinition, FunctionToolParam, ToolChoice, ToolParam, build_structural_tag, + }; + + use super::HyV3StructuralTagBuilder; + use crate::tool::HyV3ToolMarkers; + + fn tool(name: &str, parameters: Value) -> ToolParam { + ToolParam::Function(FunctionToolParam::new( + FunctionDefinition::new(name).with_parameters(parameters), + )) + } + + fn build( + suffix: &str, + tools: &[ToolParam], + choice: ToolChoice, + ) -> xgrammar_structural_tag::format::StructuralTag { + build_structural_tag( + HyV3StructuralTagBuilder::new(Arc::new(HyV3ToolMarkers::new(suffix))), + tools, + choice, + StructuralTagOptions::default().with_reasoning(false), + ) + .unwrap() + } + + #[test] + fn required_uses_suffixed_hy3_skeleton_and_bounded_values() { + let tag = build( + ":opensource", + &[tool( + "get_weather", + json!({ + "type": "object", + "properties": { + "city": { "type": "string" }, + "days": { "type": "integer" } + }, + "required": ["city"] + }), + )], + ToolChoice::required(), + ); + + expect![[r#"{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"const_string","value":"\n"},{"type":"tags_with_separator","tags":[{"begin":"get_weather\n","content":{"type":"sequence","elements":[{"type":"sequence","elements":[{"type":"const_string","value":""},{"type":"const_string","value":"city"},{"type":"const_string","value":""},{"type":"const_string","value":"\n"},{"type":"const_string","value":""},{"type":"any_text","excludes":["","","","","","","","",""]},{"type":"const_string","value":""},{"type":"const_string","value":"\n"}]},{"type":"star","content":{"type":"sequence","elements":[{"type":"const_string","value":""},{"type":"const_string","value":"days"},{"type":"const_string","value":""},{"type":"const_string","value":"\n"},{"type":"const_string","value":""},{"type":"any_text","excludes":["","","","","","","","",""]},{"type":"const_string","value":""},{"type":"const_string","value":"\n"}]}}]},"end":""}],"separator":"\n","at_least_one":true,"stop_after_first":false},{"type":"const_string","value":"\n"}]}}"#]].assert_eq(&tag.to_json_string().unwrap()); + } + + #[test] + fn optional_only_schema_keeps_declared_key_alternatives() { + let tag = build( + "", + &[tool( + "lookup", + json!({ + "type": "object", + "properties": { + "query": { "type": "string" }, + "limit": { "type": "integer" } + } + }), + )], + ToolChoice::required(), + ); + let value = serde_json::to_value(tag).unwrap(); + + expect![[r#"{"type":"sequence","elements":[{"type":"star","content":{"type":"or","elements":[{"type":"sequence","elements":[{"type":"const_string","value":""},{"type":"const_string","value":"query"},{"type":"const_string","value":""},{"type":"const_string","value":"\n"},{"type":"const_string","value":""},{"type":"any_text","excludes":["","","","","","","","",""]},{"type":"const_string","value":""},{"type":"const_string","value":"\n"}]},{"type":"sequence","elements":[{"type":"const_string","value":""},{"type":"const_string","value":"limit"},{"type":"const_string","value":""},{"type":"const_string","value":"\n"},{"type":"const_string","value":""},{"type":"any_text","excludes":["","","","","","","","",""]},{"type":"const_string","value":""},{"type":"const_string","value":"\n"}]}]}}]}"#]].assert_eq( + &serde_json::to_string( + &value["format"]["elements"][1]["tags"][0]["content"], + ) + .unwrap(), + ); + } + + #[test] + fn auto_and_forced_preserve_tool_choice_shape() { + let tools = [ + tool("search", json!({ "type": "object" })), + tool("lookup", json!({ "type": "object" })), + ]; + let auto = build("", &tools, ToolChoice::auto()).to_json_string().unwrap(); + let forced = build("", &tools, ToolChoice::function("lookup")).to_json_string().unwrap(); + + expect![[r#"{"type":"structural_tag","format":{"type":"triggered_tags","triggers":[""],"tags":[{"begin":"\n","content":{"type":"tags_with_separator","tags":[{"begin":"search\n","content":{"type":"any_text","excludes":[]},"end":""},{"begin":"lookup\n","content":{"type":"any_text","excludes":[]},"end":""}],"separator":"\n","at_least_one":true,"stop_after_first":false},"end":"\n"}],"at_least_one":false,"stop_after_first":false,"excludes":[]}}"#]].assert_eq(&auto); + expect![[r#"{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"const_string","value":"\n"},{"type":"tag","begin":"lookup\n","content":{"type":"any_text","excludes":[]},"end":""},{"type":"const_string","value":"\n"}]}}"#]].assert_eq(&forced); + } + + #[test] + fn missing_required_property_is_still_emitted() { + let tag = build( + "", + &[tool( + "search", + json!({ + "type": "object", + "properties": { "query": { "type": "string" } }, + "required": ["query", "tenant"] + }), + )], + ToolChoice::required(), + ); + let value = serde_json::to_value(tag).unwrap(); + + expect![[r#"{"type":"sequence","elements":[{"type":"sequence","elements":[{"type":"const_string","value":""},{"type":"const_string","value":"query"},{"type":"const_string","value":""},{"type":"const_string","value":"\n"},{"type":"const_string","value":""},{"type":"any_text","excludes":["","","","","","","","",""]},{"type":"const_string","value":""},{"type":"const_string","value":"\n"}]},{"type":"sequence","elements":[{"type":"const_string","value":""},{"type":"const_string","value":"tenant"},{"type":"const_string","value":""},{"type":"const_string","value":"\n"},{"type":"const_string","value":""},{"type":"any_text","excludes":["","","","","","","","",""]},{"type":"const_string","value":""},{"type":"const_string","value":"\n"}]}]}"#]].assert_eq( + &serde_json::to_string( + &value["format"]["elements"][1]["tags"][0]["content"], + ) + .unwrap(), + ); + } +} diff --git a/rust/src/parser/src/tool/mod.rs b/rust/src/parser/src/tool/mod.rs index e25eae882f9b..940072bdaca5 100644 --- a/rust/src/parser/src/tool/mod.rs +++ b/rust/src/parser/src/tool/mod.rs @@ -24,7 +24,7 @@ pub use deepseek_dsml::{DeepSeekV4ToolParser, DeepSeekV32ToolParser}; pub use deepseek_json::{DeepSeekV3ToolParser, DeepSeekV31ToolParser}; pub use error::{Result, ToolParserError}; pub use glm_xml::{Glm45MoeToolParser, Glm47MoeToolParser}; -pub use hy_v3::HyV3ToolParser; +pub(crate) use hy_v3::{HyV3ToolMarkers, HyV3ToolParser}; pub use json::{ Granite4ToolParser, HermesToolParser, Internlm2ToolParser, Llama3JsonToolParser, MistralToolParser, Phi4MiniJsonToolParser, Qwen3XmlToolParser, diff --git a/rust/src/parser/src/unified/hy_v3.rs b/rust/src/parser/src/unified/hy_v3.rs new file mode 100644 index 000000000000..ede3b51dd2e9 --- /dev/null +++ b/rust/src/parser/src/unified/hy_v3.rs @@ -0,0 +1,207 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use vllm_tokenizer::{DynTokenizer, Tokenizer}; + +use super::{CombinedParser, Result, UnifiedParser, UnifiedParserOutput, token_id}; +use crate::reasoning::HyV3ReasoningParser; +use crate::tool::{HyV3ToolMarkers, HyV3ToolParser, StructuralTagBuilder, Tool}; + +const HY_V3_MARKER_STEMS: &[&str] = &[ + "think", + "tool_calls", + "tool_call", + "tool_sep", + "arg_key", + "arg_value", +]; + +/// Unified reasoning and tool parser for HY3 output. +pub struct HyV3UnifiedParser { + inner: CombinedParser, +} + +impl HyV3UnifiedParser { + /// Create a HY3 parser using the suffix encoded in tokenizer added tokens. + pub fn new(tools: &[Tool], tokenizer: DynTokenizer) -> Result { + let suffix = detect_token_suffix(tokenizer.as_ref()); + let markers = HyV3ToolMarkers::new(&suffix); + for marker in markers.iter() { + token_id(tokenizer.as_ref(), marker)?; + } + + let reasoning = HyV3ReasoningParser::new(tokenizer, &suffix)?; + let tool = HyV3ToolParser::new(tools, &suffix); + Ok(Self { + inner: CombinedParser::new(Some(Box::new(reasoning)), Some(Box::new(tool))), + }) + } +} + +impl UnifiedParser for HyV3UnifiedParser { + fn create(tools: &[Tool], tokenizer: DynTokenizer) -> Result> + where + Self: Sized + 'static, + { + Self::new(tools, tokenizer).map(|parser| Box::new(parser) as Box) + } + + fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> { + self.inner.initialize(prompt_token_ids) + } + + fn preserve_special_tokens(&self) -> bool { + self.inner.preserve_special_tokens() + } + + fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> { + self.inner.structural_tag_builder() + } + + fn tool_call_id(&self, tool_index: usize) -> Option<&str> { + self.inner.tool_call_id(tool_index) + } + + fn parse_into(&mut self, delta: &str, output: &mut UnifiedParserOutput) -> Result<()> { + self.inner.parse_into(delta, output) + } + + fn finish(&mut self) -> Result { + self.inner.finish() + } + + fn reset(&mut self) -> String { + self.inner.reset() + } +} + +/// Detect the HY3 structural-token suffix from tokenizer added vocabulary. +fn detect_token_suffix(tokenizer: &dyn Tokenizer) -> String { + tokenizer + .added_vocab() + .iter() + .filter_map(|(token, id)| marker_suffix(token).map(|suffix| (*id, suffix))) + .min_by_key(|(id, _)| *id) + .map(|(_, suffix)| suffix.to_string()) + .unwrap_or_default() +} + +/// Extract the suffix from one opening or closing HY3 structural token. +fn marker_suffix(token: &str) -> Option<&str> { + let body = token.strip_prefix('<')?.strip_suffix('>')?; + let body = body.strip_prefix('/').unwrap_or(body); + + HY_V3_MARKER_STEMS.iter().find_map(|stem| { + let suffix = body.strip_prefix(stem)?; + (suffix.is_empty() || suffix.starts_with(':')).then_some(suffix) + }) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use serde_json::json; + use vllm_tokenizer::{Tokenizer, test_utils::TestTokenizer}; + use xgrammar_structural_tag::builders::StructuralTagOptions; + use xgrammar_structural_tag::{ + FunctionDefinition, FunctionToolParam, ToolChoice, ToolParam, build_structural_tag, + }; + + use super::{HyV3UnifiedParser, UnifiedParser}; + use crate::tool::Tool; + use crate::unified::{UnifiedParserEvent, UnifiedParserOutput}; + + fn tokenizer() -> TestTokenizer { + [ + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + ] + .into_iter() + .enumerate() + .fold(TestTokenizer::new(), |tokenizer, (index, token)| { + tokenizer.with_regular_token(token, 1000 + index as u32) + }) + } + + fn tools() -> Vec { + vec![Tool { + name: "get_weather".to_string(), + description: None, + parameters: json!({ + "type": "object", + "properties": { "city": { "type": "string" } }, + }), + strict: None, + }] + } + + #[test] + fn parses_suffixed_reasoning_and_tool_call_as_one_stream() { + let tokenizer = Arc::new(tokenizer()); + let think_start_id = tokenizer.token_to_id("").unwrap(); + let mut parser = HyV3UnifiedParser::new(&tools(), tokenizer).unwrap(); + parser.initialize(&[think_start_id]).unwrap(); + + let chunks = [ + "reasoninganswerget_weather", + "city", + "Beijing", + "", + ]; + let mut output = UnifiedParserOutput::default(); + for chunk in chunks { + parser.parse_into(chunk, &mut output).unwrap(); + } + output.append(parser.finish().unwrap()); + + assert_eq!( + output.events, + vec![ + UnifiedParserEvent::Reasoning("reasoning".to_string()), + UnifiedParserEvent::Text("answer".to_string()), + UnifiedParserEvent::ToolCall(crate::tool::ToolCallDelta { + tool_index: 0, + name: Some("get_weather".to_string()), + arguments: r#"{"city":"Beijing"}"#.to_string(), + }), + ] + ); + } + + #[test] + fn structural_tag_uses_tokenizer_detected_suffix() { + let parser = HyV3UnifiedParser::new(&tools(), Arc::new(tokenizer())).unwrap(); + let structural_tools = [ToolParam::Function(FunctionToolParam::new( + FunctionDefinition::new("get_weather").with_parameters(json!({ + "type": "object", + "properties": { "city": { "type": "string" } }, + "required": ["city"] + })), + ))]; + + let tag = build_structural_tag( + parser.structural_tag_builder().unwrap(), + &structural_tools, + ToolChoice::required(), + StructuralTagOptions::default().with_reasoning(false), + ) + .unwrap() + .to_json_string() + .unwrap(); + + assert!(tag.contains("")); + assert!(tag.contains("")); + assert!(!tag.contains("glm_xml")); + } +} diff --git a/rust/src/parser/src/unified/mod.rs b/rust/src/parser/src/unified/mod.rs index ae5524759df3..9deac1853f09 100644 --- a/rust/src/parser/src/unified/mod.rs +++ b/rust/src/parser/src/unified/mod.rs @@ -5,11 +5,13 @@ mod combined; mod gemma4; +mod hy_v3; mod inkling; mod kimi_k3; pub use combined::CombinedParser; pub use gemma4::Gemma4UnifiedParser; +pub use hy_v3::HyV3UnifiedParser; pub use inkling::InklingUnifiedParser; pub use kimi_k3::{KimiK3StructuralTagBuilder, KimiK3UnifiedParser}; use thiserror::Error; diff --git a/rust/src/text/src/backend/hf/config.rs b/rust/src/text/src/backend/hf/config.rs index 2a69632c872b..7f393014b79a 100644 --- a/rust/src/text/src/backend/hf/config.rs +++ b/rust/src/text/src/backend/hf/config.rs @@ -111,12 +111,32 @@ pub(super) struct GenerationConfig { pub eos_token_id: Option, pub temperature: Option, pub top_p: Option, + #[serde(deserialize_with = "deserialize_top_k")] pub top_k: Option, pub min_p: Option, pub repetition_penalty: Option, pub max_new_tokens: Option, } +/// Deserialize vLLM-compatible `top_k` values from generation configs. +/// +/// Both `-1` and `0` disable top-k sampling and become `None`; positive values +/// are preserved. +fn deserialize_top_k<'de, D>(deserializer: D) -> std::result::Result, D::Error> +where + D: serde::Deserializer<'de>, +{ + match Option::::deserialize(deserializer)? { + None | Some(-1) | Some(0) => Ok(None), + Some(value) if value > 0 => { + u32::try_from(value).map(Some).map_err(serde::de::Error::custom) + } + Some(value) => Err(serde::de::Error::custom(format!( + "top_k must be -1, 0, or a positive integer, got {value}" + ))), + } +} + /// HF generation configs allow either one EOS id or a list of EOS ids. #[derive(Debug, Clone, Deserialize)] #[serde(untagged)] @@ -294,7 +314,24 @@ where #[cfg(test)] mod tests { - use super::ModelConfig; + use super::{GenerationConfig, ModelConfig}; + + #[test] + fn generation_config_normalizes_vllm_top_k_values() { + for (input, expected) in [ + (r#"{"top_k":null}"#, None), + (r#"{"top_k":-1}"#, None), + (r#"{"top_k":0}"#, None), + (r#"{"top_k":20}"#, Some(20)), + ] { + let config: GenerationConfig = serde_json::from_str(input).unwrap(); + assert_eq!(config.top_k, expected, "input={input}"); + } + + for input in [r#"{"top_k":-2}"#, r#"{"top_k":4294967296}"#] { + assert!(serde_json::from_str::(input).is_err()); + } + } #[test] fn model_config_detects_moe_from_named_expert_fields() { diff --git a/rust/src/tokenizer/src/hf.rs b/rust/src/tokenizer/src/hf.rs index eb527294a5b0..d936ac80c6b7 100644 --- a/rust/src/tokenizer/src/hf.rs +++ b/rust/src/tokenizer/src/hf.rs @@ -149,10 +149,20 @@ fn encode_fastokens_ordinary( pub struct HuggingFaceTokenizer { backend: Backend, special_token_ids: Arc<[u32]>, + added_vocab: Box<[(String, u32)]>, } impl HuggingFaceTokenizer { fn from_hf_backend(tokenizer: HfTokenizer) -> Self { + let added_vocab = { + let mut vocab: Vec<_> = tokenizer + .get_added_tokens_decoder() + .iter() + .map(|(&id, token)| (token.content.clone(), id)) + .collect(); + vocab.sort_unstable_by_key(|(_, id)| *id); + vocab.into_boxed_slice() + }; let special_token_ids = { let mut ids: Vec = tokenizer .get_added_tokens_decoder() @@ -167,10 +177,21 @@ impl HuggingFaceTokenizer { Self { backend: Backend::Hf(Box::new(tokenizer)), special_token_ids, + added_vocab, } } fn from_fastokens_backend(tokenizer: FastokensTokenizer) -> Self { + let added_vocab = { + let mut vocab: Vec<_> = tokenizer + .added_tokens() + .into_iter() + .flat_map(|added_tokens| added_tokens.iter()) + .map(|token| (token.content.to_string(), token.id)) + .collect(); + vocab.sort_unstable_by_key(|(_, id)| *id); + vocab.into_boxed_slice() + }; let special_token_ids = { let mut ids: Vec = tokenizer .added_tokens() @@ -192,6 +213,7 @@ impl HuggingFaceTokenizer { Self { backend, special_token_ids, + added_vocab, } } @@ -292,6 +314,10 @@ impl Tokenizer for HuggingFaceTokenizer { } } + fn added_vocab(&self) -> &[(String, u32)] { + &self.added_vocab + } + fn is_special_id(&self, token_id: u32) -> bool { self.special_token_ids.binary_search(&token_id).is_ok() } @@ -477,6 +503,13 @@ mod tests { tokenizer.encode(SPECIAL_TOKEN, false).unwrap(), vec![tokenizer.token_to_id(SPECIAL_TOKEN).unwrap()] ); + assert_eq!( + tokenizer.added_vocab(), + vec![ + (REGULAR_TOKEN.to_string(), 256), + (SPECIAL_TOKEN.to_string(), 257), + ] + ); for text in [ "", diff --git a/rust/src/tokenizer/src/lib.rs b/rust/src/tokenizer/src/lib.rs index 0f8c7dc16e59..2bc1cfca3d7b 100644 --- a/rust/src/tokenizer/src/lib.rs +++ b/rust/src/tokenizer/src/lib.rs @@ -39,6 +39,15 @@ pub trait Tokenizer: Send + Sync { /// Convert one token ID into the tokenizer's raw token string. fn id_to_token(&self, id: u32) -> Option; + /// Borrow tokenizer added-vocabulary entries as `(token, id)` pairs. + /// + /// Backends that do not expose the distinction between model vocabulary + /// and added vocabulary return an empty list. + // TODO: add support to all tokenizer backends + fn added_vocab(&self) -> &[(String, u32)] { + &[] + } + /// Return the vocabulary size. Backends that cannot report it fall back to /// `usize::MAX`, an effectively unbounded value used only by test stubs. fn vocab_size(&self) -> usize { diff --git a/rust/src/tokenizer/src/test_utils.rs b/rust/src/tokenizer/src/test_utils.rs index 6fdb6e027211..924f70c98a72 100644 --- a/rust/src/tokenizer/src/test_utils.rs +++ b/rust/src/tokenizer/src/test_utils.rs @@ -69,6 +69,7 @@ struct TestToken { pub struct TestTokenizer { token_to_id: BTreeMap, id_to_token: BTreeMap, + added_vocab: Vec<(String, u32)>, unknown_decode: UnknownDecode, vocab_size: Option, bos_token_id: Option, @@ -86,6 +87,7 @@ impl TestTokenizer { Self { token_to_id: BTreeMap::new(), id_to_token: BTreeMap::new(), + added_vocab: Vec::new(), unknown_decode: UnknownDecode::Error, vocab_size: None, bos_token_id: None, @@ -154,9 +156,21 @@ impl TestTokenizer { if self.token_to_id.insert(token.clone(), id).is_some() { panic!("configured test token text {token:?} was registered more than once"); } - if self.id_to_token.insert(id, TestToken { text: token, kind }).is_some() { + if self + .id_to_token + .insert( + id, + TestToken { + text: token.clone(), + kind, + }, + ) + .is_some() + { panic!("configured test token id {id} was registered more than once"); } + self.added_vocab.push((token, id)); + self.added_vocab.sort_unstable_by_key(|(_, id)| *id); } fn byte_to_token(id: u32) -> Option { @@ -254,6 +268,10 @@ impl Tokenizer for TestTokenizer { .or_else(|| Self::byte_to_token(id)) } + fn added_vocab(&self) -> &[(String, u32)] { + &self.added_vocab + } + fn vocab_size(&self) -> usize { self.vocab_size.unwrap_or_else(|| self.inferred_vocab_size()) } From 5df31ea52d7d175dd920fce61a19017941b18b4b Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Thu, 20 Aug 2026 22:44:55 -0400 Subject: [PATCH 233/839] [Spec Decode] Enable adaptive verification on DSv4 + sm90 (#52795) Signed-off-by: Lucas Wilkinson Co-authored-by: OpenAI Codex --- vllm/v1/attention/backends/mla/indexer.py | 32 ++++++++++++++++++++--- 1 file changed, 28 insertions(+), 4 deletions(-) diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 1177d8b27afe..2a21445f5da1 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -157,8 +157,11 @@ def supports_pcp(cls) -> bool: @classmethod def supports_device_cpu_query_lens_mismatch(cls) -> bool: # Only the varlen paged MQA logits kernel takes per-request query - # lengths from device tensors; otherwise the indexer needs uniform ones. - return _supports_varlen_paged_mqa_logits() + # lengths from device tensors natively. Hopper can instead flatten each + # query into a single-token row using device-built metadata. + return _supports_varlen_paged_mqa_logits() or ( + _supports_flattened_device_query_lens() + ) @staticmethod def get_name() -> str: @@ -493,6 +496,14 @@ def _supports_varlen_paged_mqa_logits() -> bool: ) +def _supports_flattened_device_query_lens() -> bool: + return ( + current_platform.is_cuda() + and current_platform.is_device_capability_family(90) + and has_deep_gemm() + ) + + def _supports_native_decode(next_n: int) -> bool: """Whether decode can pass `next_n` Q rows per request to the kernel instead of flattening to one single-token row per query, which re-reads @@ -507,6 +518,16 @@ def _supports_native_decode(next_n: int) -> bool: return next_n in (1, 2) +def _use_flattening(vllm_config: VllmConfig) -> bool: + speculative_config = vllm_config.speculative_config + next_n = 1 + vllm_config.num_speculative_tokens + return not _supports_native_decode(next_n) or ( + speculative_config is not None + and speculative_config.enable_adaptive_verification + and _supports_flattened_device_query_lens() + ) + + class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): # The indexer opts out of the shared reorder-threshold vote (see __init__), # so this is None; its own split uses self.decode_threshold. @@ -519,7 +540,7 @@ def get_cudagraph_support( vllm_config: VllmConfig, kv_cache_spec: KVCacheSpec, ) -> AttentionCGSupport: - if _supports_varlen_paged_mqa_logits(): + if _supports_varlen_paged_mqa_logits() or _use_flattening(vllm_config): return AttentionCGSupport.ALWAYS return AttentionCGSupport.UNIFORM_BATCH @@ -553,7 +574,7 @@ def __init__(self, *args, block_table_width: int, **kwargs) -> None: next_n = self.num_speculative_tokens + 1 self.decode_threshold = next_n self.reorder_batch_threshold = None - self.use_flattening = not _supports_native_decode(next_n) + self.use_flattening = _use_flattening(self.vllm_config) self.supports_varlen = _supports_varlen_paged_mqa_logits() logger.info_once( "DSA indexer decode path: use_flattening=%s supports_varlen=%s " @@ -674,11 +695,14 @@ def _prepare_decode_tensors( max_decode_len: int, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, bool]: """Prepare native or per-token flattened decode tensors.""" + spec_config = self.vllm_config.speculative_config + adaptive = bool(spec_config and spec_config.enable_adaptive_verification) min_decode_len = int(decode_lens_cpu.min().item()) if not use_native: assert self.decode_seq_lens_buffer.dim() == 1 if ( not self.supports_varlen + and (num_decodes == 1 or not adaptive) and min_decode_len == max_decode_len and num_decodes * max_decode_len == num_decode_tokens ): From 2adf4b9e2a915bd30876ab0c5dba2243f54fe5b7 Mon Sep 17 00:00:00 2001 From: Reid <61492567+reidliu41@users.noreply.github.com> Date: Fri, 21 Aug 2026 10:47:07 +0800 Subject: [PATCH 234/839] [Rust Frontend] Fix Kimi K3 reasoning_effort="none" handling (#53043) Co-authored-by: Bugen Zhao Signed-off-by: reidliu41 --- .../src/chat/src/renderer/kimi_k3/encoding.rs | 25 ++++++++++---- rust/src/chat/src/renderer/kimi_k3/tests.rs | 34 +++++++++++++++++-- 2 files changed, 50 insertions(+), 9 deletions(-) diff --git a/rust/src/chat/src/renderer/kimi_k3/encoding.rs b/rust/src/chat/src/renderer/kimi_k3/encoding.rs index 3ff9242ff4bf..392b6b642d29 100644 --- a/rust/src/chat/src/renderer/kimi_k3/encoding.rs +++ b/rust/src/chat/src/renderer/kimi_k3/encoding.rs @@ -236,11 +236,15 @@ fn thinking_enabled(request: &ChatRequest) -> Result { if let Some(enable_thinking) = request.parse_template_bool("enable_thinking")? { return Ok(enable_thinking); } + if let Some(reasoning_effort) = request.chat_options.reasoning_effort { + return Ok(reasoning_effort != crate::request::ReasoningEffort::None); + } Ok(request .chat_options - .reasoning_effort - .map(|effort| effort != crate::request::ReasoningEffort::None) - .unwrap_or(true)) + .template_kwargs + .get("reasoning_effort") + .and_then(Value::as_str) + != Some("none")) } fn thinking_effort(request: &ChatRequest) -> Result { @@ -251,13 +255,22 @@ fn thinking_effort(request: &ChatRequest) -> Result { )) })? } else if let Some(effort) = request.chat_options.reasoning_effort { - effort.as_str() + if effort == crate::request::ReasoningEffort::None { + DEFAULT_THINKING_EFFORT + } else { + effort.as_str() + } } else if let Some(value) = request.chat_options.template_kwargs.get("reasoning_effort") { - value.as_str().ok_or_else(|| { + let effort = value.as_str().ok_or_else(|| { Error::ChatTemplate(format!( "template kwarg `reasoning_effort` must be a string, got {value}" )) - })? + })?; + if effort == "none" { + DEFAULT_THINKING_EFFORT + } else { + effort + } } else { DEFAULT_THINKING_EFFORT }; diff --git a/rust/src/chat/src/renderer/kimi_k3/tests.rs b/rust/src/chat/src/renderer/kimi_k3/tests.rs index 94a3ad4cfef2..1d9f984dbbd3 100644 --- a/rust/src/chat/src/renderer/kimi_k3/tests.rs +++ b/rust/src/chat/src/renderer/kimi_k3/tests.rs @@ -329,16 +329,44 @@ fn native_k3_kwargs_take_precedence() { #[test] fn standard_none_disables_thinking() { + let mut request = crate::request::ChatRequest::for_test(); + request + .chat_options + .template_kwargs + .insert("reasoning_effort".to_string(), json!("none")); + + let rendered = render_request(&request); + + assert!(!rendered.contains("type=\"thinking-effort\"")); + assert!(rendered.ends_with("<|open|>response<|sep|>")); +} + +#[test] +fn native_thinking_true_overrides_standard_none() { let mut request = crate::request::ChatRequest::for_test(); request.chat_options.template_kwargs.extend([ - ("enable_thinking".to_string(), json!(false)), + ("thinking".to_string(), json!(true)), ("reasoning_effort".to_string(), json!("none")), ]); let rendered = render_request(&request); - assert!(!rendered.contains("type=\"thinking-effort\"")); - assert!(rendered.ends_with("<|open|>response<|sep|>")); + assert!(rendered.contains("thinking_effort=max")); + assert!(rendered.ends_with("<|open|>think<|sep|>")); +} + +#[test] +fn enable_thinking_true_overrides_standard_none() { + let mut request = crate::request::ChatRequest::for_test(); + request.chat_options.template_kwargs.extend([ + ("enable_thinking".to_string(), json!(true)), + ("reasoning_effort".to_string(), json!("none")), + ]); + + let rendered = render_request(&request); + + assert!(rendered.contains("thinking_effort=max")); + assert!(rendered.ends_with("<|open|>think<|sep|>")); } #[test] From b00f475f09b7d54e811efc006cab2562d3b23893 Mon Sep 17 00:00:00 2001 From: Cyrus Leung Date: Fri, 21 Aug 2026 11:42:53 +0800 Subject: [PATCH 235/839] [MM] Remove text components from ProcessorInputs (#53093) Signed-off-by: DarkLight1337 --- docs/contributing/model/multimodal.md | 7 - .../processing/test_audioflamingo3.py | 2 +- .../multimodal/processing/test_common.py | 68 +--------- .../multimodal/processing/test_llava_next.py | 19 ++- .../processing/test_llava_onevision.py | 20 ++- .../multimodal/processing/test_moss_audio.py | 4 +- .../multimodal/processing/test_openvla.py | 5 - .../processing/test_tensor_schema.py | 6 +- vllm/model_executor/models/audioflamingo3.py | 2 - vllm/model_executor/models/bagel.py | 9 -- vllm/model_executor/models/blip2.py | 2 - vllm/model_executor/models/chameleon.py | 2 - vllm/model_executor/models/cheers.py | 12 +- vllm/model_executor/models/clip.py | 37 +---- vllm/model_executor/models/cohere2_vision.py | 2 - vllm/model_executor/models/cohere_asr.py | 6 +- vllm/model_executor/models/colmodernvbert.py | 11 -- vllm/model_executor/models/colpali.py | 14 +- vllm/model_executor/models/deepseek_ocr.py | 1 - vllm/model_executor/models/deepseek_ocr2.py | 1 - vllm/model_executor/models/deepseek_vl2.py | 2 - vllm/model_executor/models/ernie45_vl.py | 3 +- vllm/model_executor/models/fireredasr2.py | 2 - vllm/model_executor/models/fireredlid.py | 6 +- vllm/model_executor/models/funasr.py | 2 - vllm/model_executor/models/funaudiochat.py | 12 +- vllm/model_executor/models/gemma3_mm.py | 2 - vllm/model_executor/models/gemma3n_mm.py | 2 - vllm/model_executor/models/gemma4_mm.py | 22 --- vllm/model_executor/models/glm4_1v.py | 4 - vllm/model_executor/models/glm4v.py | 9 -- vllm/model_executor/models/glmasr.py | 2 - vllm/model_executor/models/granite_speech.py | 2 - vllm/model_executor/models/hunyuan_vision.py | 3 +- .../models/hyperclovax_vision.py | 10 -- .../models/hyperclovax_vision_v2.py | 32 ++--- vllm/model_executor/models/idefics3.py | 2 - vllm/model_executor/models/interns1.py | 5 +- vllm/model_executor/models/internvl.py | 7 +- vllm/model_executor/models/jina_vl.py | 3 +- vllm/model_executor/models/kanana_v.py | 1 - vllm/model_executor/models/keye.py | 3 +- vllm/model_executor/models/keye_vl1_5.py | 3 +- vllm/model_executor/models/kimi_audio.py | 3 +- vllm/model_executor/models/kimi_k25.py | 3 +- vllm/model_executor/models/lfm2_vl.py | 2 - vllm/model_executor/models/lightonocr.py | 2 - vllm/model_executor/models/llava.py | 2 - vllm/model_executor/models/llava_onevision.py | 21 --- .../model_executor/models/llava_onevision2.py | 8 +- vllm/model_executor/models/midashenglm.py | 2 - vllm/model_executor/models/mimo_v2_omni.py | 3 +- vllm/model_executor/models/minicpmo.py | 7 +- vllm/model_executor/models/minicpmv.py | 26 +--- vllm/model_executor/models/minicpmv4_6.py | 2 - vllm/model_executor/models/mistral3.py | 2 - vllm/model_executor/models/mllama4.py | 2 - vllm/model_executor/models/molmo.py | 3 +- vllm/model_executor/models/molmo2.py | 5 +- vllm/model_executor/models/moondream3.py | 15 +-- vllm/model_executor/models/moss_audio.py | 42 +++++- .../models/moss_transcribe_diarize.py | 17 +-- vllm/model_executor/models/muse_glimmer.py | 13 -- .../model_executor/models/nano_nemotron_vl.py | 10 +- vllm/model_executor/models/nemotron_parse.py | 9 +- vllm/model_executor/models/opencua.py | 10 -- vllm/model_executor/models/openvla.py | 9 -- vllm/model_executor/models/ovis.py | 2 - vllm/model_executor/models/ovis2_5.py | 2 - vllm/model_executor/models/paddleocr_vl.py | 3 +- vllm/model_executor/models/paligemma.py | 2 - vllm/model_executor/models/phi3v.py | 10 +- vllm/model_executor/models/phi4mm.py | 5 +- vllm/model_executor/models/phi4siglip.py | 16 --- vllm/model_executor/models/pixtral.py | 10 +- .../models/qwen2_5_omni_thinker.py | 39 +----- vllm/model_executor/models/qwen2_5_vl.py | 3 +- vllm/model_executor/models/qwen2_audio.py | 2 - .../models/qwen3_omni_moe_thinker.py | 5 +- vllm/model_executor/models/qwen3_vl.py | 3 - vllm/model_executor/models/siglip.py | 37 +---- .../models/transformers/multimodal.py | 79 +++-------- vllm/model_executor/models/ultravox.py | 7 - vllm/model_executor/models/voxtral.py | 10 +- vllm/model_executor/models/whisper.py | 14 +- vllm/models/dots3_note/common/processor.py | 12 +- vllm/models/inkling/common/mm_preprocess.py | 18 +-- vllm/models/kimi_k3/common/mm_preprocess.py | 12 +- .../models/minimax_m3/common/mm_preprocess.py | 3 +- vllm/multimodal/processing/dummy_inputs.py | 16 ++- vllm/multimodal/processing/inputs.py | 3 +- vllm/multimodal/processing/processor.py | 126 ++++-------------- 92 files changed, 220 insertions(+), 806 deletions(-) diff --git a/docs/contributing/model/multimodal.md b/docs/contributing/model/multimodal.md index 33d89db75d3c..3d16e0e91fbc 100644 --- a/docs/contributing/model/multimodal.md +++ b/docs/contributing/model/multimodal.md @@ -456,13 +456,11 @@ return a schema of the tensors outputted by the HF processor that are related to prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: processed_outputs = super()._call_hf_processor( prompt=prompt, mm_data=mm_data, mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, ) pixel_values = processed_outputs.get("pixel_values") @@ -479,11 +477,6 @@ return a schema of the tensors outputted by the HF processor that are related to return processed_outputs ``` - !!! note - The `_call_hf_processor` method specifies both `mm_kwargs` and `tok_kwargs` for - processing. `mm_kwargs` is used to both initialize and call the huggingface - processor, whereas `tok_kwargs` is only used to call the huggingface processor. - Since `pixel_values` is now a list with one tensor per image, we can override [_get_mm_fields_config][vllm.multimodal.processing.BaseMultiModalProcessor._get_mm_fields_config] as follows: diff --git a/tests/models/multimodal/processing/test_audioflamingo3.py b/tests/models/multimodal/processing/test_audioflamingo3.py index ae0da09e1fb0..c7dffa90de29 100644 --- a/tests/models/multimodal/processing/test_audioflamingo3.py +++ b/tests/models/multimodal/processing/test_audioflamingo3.py @@ -107,7 +107,7 @@ def test_audio_chunk_counting(mock_ctx): mm_data = {"audio": [audio_1, audio_2]} prompt = "<|user|>Listen.<|end|>" - processed = processor._call_hf_processor(prompt, mm_data, {}, {}) + processed = processor._call_hf_processor(prompt, mm_data, {}) chunk_counts = processed["chunk_counts"] diff --git a/tests/models/multimodal/processing/test_common.py b/tests/models/multimodal/processing/test_common.py index 0754892e4e09..ce79ce243583 100644 --- a/tests/models/multimodal/processing/test_common.py +++ b/tests/models/multimodal/processing/test_common.py @@ -145,10 +145,10 @@ def get_transformers_backend_model_ids_to_test(): ) -def get_text_token_prompts( +def get_token_prompt( processor: BaseMultiModalProcessor, mm_data: MultiModalDataDict, -): +) -> list[int]: dummy_inputs = processor.dummy_inputs tokenizer: TokenizerLike = processor.info.get_tokenizer() model_config = processor.info.ctx.model_config @@ -178,21 +178,10 @@ def get_text_token_prompts( mm_options={}, ) - text_prompt: str | None - token_prompt: list[int] - if isinstance(inputs.prompt, list): - text_prompt = None - token_prompt = inputs.prompt - elif isinstance(inputs.prompt, str): - text_prompt = inputs.prompt - token_prompt = tokenizer.encode( - text_prompt, - **processor.info.get_default_tok_params().get_encode_kwargs(), - ) - else: + if not isinstance(inputs.prompt, list): raise TypeError(type(inputs.prompt)) - return text_prompt, token_prompt + return inputs.prompt def random_vision_chunk( @@ -358,7 +347,7 @@ def _test_processing_correctness_one( ): model_type = model_config.hf_config.model_type - text_prompt, token_prompt = get_text_token_prompts(baseline_processor, mm_data) + token_prompt = get_token_prompt(baseline_processor, mm_data) mm_items = baseline_processor.info.parse_mm_data(mm_data) ignore_mm_keys = _IGNORE_MM_KEYS.get(model_type, set[str]()) @@ -381,55 +370,10 @@ def _test_processing_correctness_one( msg=( f"Failed ({batch_idx=}, {hit_rate=}, " f"{num_batches=}, {simplify_rate=}, " - f"{text_prompt=}, {token_prompt=}, {mm_data=})" + f"{token_prompt=}, {mm_data=})" ), ) - if text_prompt is not None: - baseline_text_result = baseline_processor( - text_prompt, - mm_items=mm_items, - hf_processor_mm_kwargs={}, - ) - cached_text_result = cached_processor( - text_prompt, - mm_items=mm_items, - hf_processor_mm_kwargs={}, - ) - - _assert_inputs_equal( - baseline_text_result, - cached_text_result, - ignore_mm_keys=ignore_mm_keys, - msg=( - f"Failed ({batch_idx=}, {hit_rate=}, " - f"{num_batches=}, {simplify_rate=}, " - f"{text_prompt=}, {token_prompt=}, {mm_data=})" - ), - ) - - _assert_inputs_equal( - baseline_text_result, - baseline_tokenized_result, - ignore_mm_keys=ignore_mm_keys, - msg=( - f"Failed ({batch_idx=}, {hit_rate=}, " - f"{num_batches=}, {simplify_rate=}, " - f"{text_prompt=}, {token_prompt=}, {mm_data=})" - ), - ) - - _assert_inputs_equal( - cached_text_result, - cached_tokenized_result, - ignore_mm_keys=ignore_mm_keys, - msg=( - f"Failed ({batch_idx=}, {hit_rate=}, " - f"{num_batches=}, {simplify_rate=}, " - f"{text_prompt=}, {token_prompt=}, {mm_data=})" - ), - ) - @pytest.mark.parametrize("model_id", get_model_ids_to_test()) @pytest.mark.parametrize("hit_rate", [0.3, 0.5, 1.0]) diff --git a/tests/models/multimodal/processing/test_llava_next.py b/tests/models/multimodal/processing/test_llava_next.py index b72c1bfd8ece..d7eaa35dd8d9 100644 --- a/tests/models/multimodal/processing/test_llava_next.py +++ b/tests/models/multimodal/processing/test_llava_next.py @@ -11,6 +11,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.parse import ImageSize from vllm.multimodal.processing import BaseMultiModalProcessor +from vllm.tokenizers.hf import maybe_make_thread_pool from ...utils import build_model_context @@ -144,7 +145,14 @@ def test_processor_prompt_replacements_regression(model_id, num_imgs): mm_processor_kwargs=None, limit_mm_per_prompt={"image": num_imgs}, ) - processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) + + # Avoid tokenizer already borrowed error + maybe_make_thread_pool(ctx.tokenizer) + + processor = MULTIMODAL_REGISTRY.create_processor( + ctx.model_config, + tokenizer=ctx.tokenizer, + ) image_ratios = [ (171, 152), @@ -177,7 +185,14 @@ def test_processor_prompt_replacements_all(model_id, num_imgs): mm_processor_kwargs=None, limit_mm_per_prompt={"image": num_imgs}, ) - processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) + + # Avoid tokenizer already borrowed error + maybe_make_thread_pool(ctx.tokenizer) + + processor = MULTIMODAL_REGISTRY.create_processor( + ctx.model_config, + tokenizer=ctx.tokenizer, + ) seen_aspect_ratios = set[float]() image_sizes = list[ImageSize]() diff --git a/tests/models/multimodal/processing/test_llava_onevision.py b/tests/models/multimodal/processing/test_llava_onevision.py index 2bac464e78f4..4e8a8103a592 100644 --- a/tests/models/multimodal/processing/test_llava_onevision.py +++ b/tests/models/multimodal/processing/test_llava_onevision.py @@ -11,6 +11,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.parse import ImageSize from vllm.multimodal.processing import BaseMultiModalProcessor +from vllm.tokenizers.hf import maybe_make_thread_pool from ...utils import build_model_context @@ -42,7 +43,15 @@ def test_processor_max_tokens(model_id): mm_processor_kwargs=None, limit_mm_per_prompt={"image": 1}, ) - processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) + + # Avoid tokenizer already borrowed error + maybe_make_thread_pool(ctx.tokenizer) + + processor = MULTIMODAL_REGISTRY.create_processor( + ctx.model_config, + tokenizer=ctx.tokenizer, + ) + info = processor.info seen_aspect_ratios = set[float]() @@ -142,7 +151,14 @@ def test_processor_prompt_replacements_regression(model_id, num_imgs): mm_processor_kwargs=None, limit_mm_per_prompt={"image": num_imgs}, ) - processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) + + # Avoid tokenizer already borrowed error + maybe_make_thread_pool(ctx.tokenizer) + + processor = MULTIMODAL_REGISTRY.create_processor( + ctx.model_config, + tokenizer=ctx.tokenizer, + ) image_ratios = [ (171, 152), diff --git a/tests/models/multimodal/processing/test_moss_audio.py b/tests/models/multimodal/processing/test_moss_audio.py index 8a18d90f560c..87e7079a810b 100644 --- a/tests/models/multimodal/processing/test_moss_audio.py +++ b/tests/models/multimodal/processing/test_moss_audio.py @@ -35,12 +35,10 @@ class _Tokenizer: - def encode(self, text, add_special_tokens=False): - del add_special_tokens + def encode(self, text, **kwargs): return [ord(char) for char in text] def decode(self, token_ids, **kwargs): - del kwargs return "".join(chr(token_id) for token_id in token_ids) def batch_decode(self, batch_token_ids, **kwargs): diff --git a/tests/models/multimodal/processing/test_openvla.py b/tests/models/multimodal/processing/test_openvla.py index b9ed02000d9a..992608f893cc 100644 --- a/tests/models/multimodal/processing/test_openvla.py +++ b/tests/models/multimodal/processing/test_openvla.py @@ -185,11 +185,6 @@ def test_openvla_prompt_update_inserts_image_tokens_after_bos() -> None: image = Image.new("RGB", (640, 480), color=(255, 255, 255)) mm_items = MultiModalDataItems({"image": ImageProcessorItems([image])}) - assert ( - processor._hf_processor_applies_updates("In: test\nOut:", mm_items, {}, {}) - is False - ) - prompt_update = processor._get_prompt_updates(mm_items, {}, {})[0] resolved = prompt_update.resolve(0) content = resolved.content diff --git a/tests/models/multimodal/processing/test_tensor_schema.py b/tests/models/multimodal/processing/test_tensor_schema.py index 12c5071978ff..87cf65b8f978 100644 --- a/tests/models/multimodal/processing/test_tensor_schema.py +++ b/tests/models/multimodal/processing/test_tensor_schema.py @@ -37,7 +37,7 @@ from ....utils import create_new_process_for_each_test from ...registry import HF_EXAMPLE_MODELS from ...utils import dummy_hf_overrides -from .test_common import get_model_ids_to_test, get_text_token_prompts +from .test_common import get_model_ids_to_test, get_token_prompt ImageInput = list[Image.Image] VideoInput: TypeAlias = ( @@ -107,10 +107,10 @@ def create_batched_mm_kwargs( } # video metadata will be added back to the resized video data here. - text_prompt, token_prompt = get_text_token_prompts(processor, resized_mm_data) + token_prompt = get_token_prompt(processor, resized_mm_data) mm_kwargs = processor( - prompt=token_prompt if text_prompt is None else text_prompt, + prompt=token_prompt, mm_items=processor.info.parse_mm_data(resized_mm_data), hf_processor_mm_kwargs=processor_inputs.hf_processor_mm_kwargs, )["mm_kwargs"].require_data() diff --git a/vllm/model_executor/models/audioflamingo3.py b/vllm/model_executor/models/audioflamingo3.py index d8099c311a59..c5583b669b2c 100644 --- a/vllm/model_executor/models/audioflamingo3.py +++ b/vllm/model_executor/models/audioflamingo3.py @@ -382,7 +382,6 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, Any], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: processor_mm_data = dict(mm_data) audios = processor_mm_data.pop("audios", None) @@ -393,7 +392,6 @@ def _call_hf_processor( prompt=prompt, mm_data=processor_mm_data, mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, ) if "input_features_mask" in outputs: diff --git a/vllm/model_executor/models/bagel.py b/vllm/model_executor/models/bagel.py index 28c8b64cd170..9603dc043cc7 100644 --- a/vllm/model_executor/models/bagel.py +++ b/vllm/model_executor/models/bagel.py @@ -273,15 +273,6 @@ def get_dummy_mm_data( class BagelMultiModalProcessor(BaseMultiModalProcessor[BagelProcessingInfo]): """Multimodal processor for BAGEL model.""" - def _hf_processor_applies_updates( - self, - prompt_text: str, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - ) -> bool: - return False - def _get_prompt_updates( self, mm_items: MultiModalDataItems, diff --git a/vllm/model_executor/models/blip2.py b/vllm/model_executor/models/blip2.py index baa1a16a59bb..2f11b1dee076 100644 --- a/vllm/model_executor/models/blip2.py +++ b/vllm/model_executor/models/blip2.py @@ -472,7 +472,6 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: if not mm_data: # HF processor always adds placeholders even when there's no image @@ -484,7 +483,6 @@ def _call_hf_processor( prompt=prompt, mm_data=mm_data, mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, ) def _get_mm_fields_config( diff --git a/vllm/model_executor/models/chameleon.py b/vllm/model_executor/models/chameleon.py index 517697669c8e..af147b7f1a59 100644 --- a/vllm/model_executor/models/chameleon.py +++ b/vllm/model_executor/models/chameleon.py @@ -140,7 +140,6 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: if not mm_data: prompt_ids = self.info.get_tokenizer().encode(prompt) @@ -151,7 +150,6 @@ def _call_hf_processor( prompt=prompt, mm_data=mm_data, mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, ) def _apply_hf_processor_tokens_only( diff --git a/vllm/model_executor/models/cheers.py b/vllm/model_executor/models/cheers.py index 5f74c6771e4e..b9e88e10558a 100644 --- a/vllm/model_executor/models/cheers.py +++ b/vllm/model_executor/models/cheers.py @@ -497,18 +497,8 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: - return super()._call_hf_processor(prompt, mm_data, mm_kwargs, tok_kwargs) - - def _hf_processor_applies_updates( - self, - prompt_text: str, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - ) -> bool: - return False + return super()._call_hf_processor(prompt, mm_data, mm_kwargs) def _get_prompt_updates( self, diff --git a/vllm/model_executor/models/clip.py b/vllm/model_executor/models/clip.py index 2e0190350538..4294a203df28 100644 --- a/vllm/model_executor/models/clip.py +++ b/vllm/model_executor/models/clip.py @@ -208,40 +208,17 @@ def apply( timing_ctx: TimingContext, ) -> MultiModalInput: if inputs.mm_data_items: - if isinstance(inputs.prompt, str): - if len(inputs.prompt) > 0: - raise ValueError( - "CLIP accepts text-only or image-only inputs, not both! " - "You must pass an image with an empty text prompt." - ) + special_tokens = self.info.get_tokenizer().all_special_ids + if all(tok in special_tokens for tok in inputs.prompt): + inputs.prompt = [] else: - special_tokens = self.info.get_tokenizer().all_special_ids - if all(tok in special_tokens for tok in inputs.prompt): - inputs.prompt = [] - else: - raise ValueError( - "CLIP accepts text-only or image-only inputs, not both! " - "You must pass an image with an empty token prompt." - ) - - # For multi-modal data, the prompt after processing should - # only contain the dummy image tokens - inputs.tokenization_kwargs = { - **inputs.tokenization_kwargs, - "add_special_tokens": False, - } + raise ValueError( + "CLIP accepts text-only or image-only inputs, not both! " + "You must pass an image with an empty token prompt." + ) return super().apply(inputs, timing_ctx) - def _hf_processor_applies_updates( - self, - prompt_text: str, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - ) -> bool: - return False - def _get_mm_fields_config( self, hf_inputs: BatchFeature, diff --git a/vllm/model_executor/models/cohere2_vision.py b/vllm/model_executor/models/cohere2_vision.py index b6def10aab3b..147d462b23c6 100644 --- a/vllm/model_executor/models/cohere2_vision.py +++ b/vllm/model_executor/models/cohere2_vision.py @@ -227,13 +227,11 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: processed_outputs = super()._call_hf_processor( prompt, mm_data, mm_kwargs, - tok_kwargs, ) # Ensure num_patches is available for proper tensor splitting diff --git a/vllm/model_executor/models/cohere_asr.py b/vllm/model_executor/models/cohere_asr.py index 7ffdd01cd602..9cbf9bb205c5 100644 --- a/vllm/model_executor/models/cohere_asr.py +++ b/vllm/model_executor/models/cohere_asr.py @@ -1936,9 +1936,9 @@ def pad_dummy_encoder_prompt(self) -> bool: def create_encoder_prompt( self, - prompt: str | list[int], + prompt: list[int], mm_items: MultiModalDataItems, - ) -> str | list[int]: + ) -> list[int]: return [0] def _call_hf_processor( @@ -1946,7 +1946,6 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ): if mm_data: feature_extractor = self.info.get_feature_extractor(**mm_kwargs) @@ -1959,7 +1958,6 @@ def _call_hf_processor( prompt=prompt, mm_data=mm_data, mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, ) if "labels" in processed_outputs: processed_outputs["input_ids"] = processed_outputs.pop("labels") diff --git a/vllm/model_executor/models/colmodernvbert.py b/vllm/model_executor/models/colmodernvbert.py index 798b87b77b1d..ffa10452c529 100644 --- a/vllm/model_executor/models/colmodernvbert.py +++ b/vllm/model_executor/models/colmodernvbert.py @@ -159,14 +159,12 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: tokenizer = self.info.get_tokenizer() assert isinstance(tokenizer, HfTokenizer) text_encoding = tokenizer( prompt, return_tensors="pt", - **tok_kwargs, ) result = BatchFeature(data=dict(text_encoding)) @@ -187,15 +185,6 @@ def _call_hf_processor( return result - def _hf_processor_applies_updates( - self, - prompt_text: str, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - ) -> bool: - return False - def _get_mm_fields_config( self, hf_inputs: BatchFeature, diff --git a/vllm/model_executor/models/colpali.py b/vllm/model_executor/models/colpali.py index 7948da01bab7..aaf3b307a1f5 100644 --- a/vllm/model_executor/models/colpali.py +++ b/vllm/model_executor/models/colpali.py @@ -64,7 +64,6 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: if mm_data: # The ColPali tokenizer_config.json ships with a small default @@ -72,13 +71,12 @@ def _call_hf_processor( # by PaliGemmaProcessor, causing a token-count mismatch. # vLLM enforces its own max_model_len, so we disable HF # truncation to keep all image + text tokens intact. - tok_kwargs = dict(tok_kwargs, truncation=False) - return super()._call_hf_processor( - prompt=prompt, - mm_data=mm_data, - mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, - ) + return self.info.ctx.call_hf_processor( + self.info.get_hf_processor(**mm_kwargs), + dict(text=prompt, **mm_data), + dict(**mm_kwargs, truncation=False), + ) + return super()._call_hf_processor(prompt, mm_data, mm_kwargs) @default_pooling_type(seq_pooling_type="CLS", tok_pooling_type="ALL") diff --git a/vllm/model_executor/models/deepseek_ocr.py b/vllm/model_executor/models/deepseek_ocr.py index 4bb5598db3a6..5365b031f9f3 100644 --- a/vllm/model_executor/models/deepseek_ocr.py +++ b/vllm/model_executor/models/deepseek_ocr.py @@ -292,7 +292,6 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: if mm_data: processed_outputs = self.info.ctx.call_hf_processor( diff --git a/vllm/model_executor/models/deepseek_ocr2.py b/vllm/model_executor/models/deepseek_ocr2.py index aa6e9b1aec43..d19af4713b6f 100644 --- a/vllm/model_executor/models/deepseek_ocr2.py +++ b/vllm/model_executor/models/deepseek_ocr2.py @@ -163,7 +163,6 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: if mm_data: processed_outputs = self.info.ctx.call_hf_processor( diff --git a/vllm/model_executor/models/deepseek_vl2.py b/vllm/model_executor/models/deepseek_vl2.py index 9ca05b6c77d1..ac83b6028ad4 100644 --- a/vllm/model_executor/models/deepseek_vl2.py +++ b/vllm/model_executor/models/deepseek_vl2.py @@ -241,7 +241,6 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: if not mm_data: tokenizer = self.info.get_tokenizer() @@ -252,7 +251,6 @@ def _call_hf_processor( prompt=prompt, mm_data=mm_data, mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, ) processed_outputs["num_patches"] = ( diff --git a/vllm/model_executor/models/ernie45_vl.py b/vllm/model_executor/models/ernie45_vl.py index 8c4b150e30d5..4b495fae7aaa 100644 --- a/vllm/model_executor/models/ernie45_vl.py +++ b/vllm/model_executor/models/ernie45_vl.py @@ -1115,7 +1115,6 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: # when the prompt is not empty but the multimodal data is empty, # directly invoke the tokenizer. @@ -1151,7 +1150,7 @@ def _call_hf_processor( processor_output = self.info.ctx.call_hf_processor( hf_processor, dict(text=[prompt], images=mm_data["images"], videos=mm_data["videos"]), - dict(**mm_kwargs, **tok_kwargs), + mm_kwargs, ) # Divide the processor_output into two modalities: image and video. diff --git a/vllm/model_executor/models/fireredasr2.py b/vllm/model_executor/models/fireredasr2.py index 7dedda3c3c41..670e44cb8642 100644 --- a/vllm/model_executor/models/fireredasr2.py +++ b/vllm/model_executor/models/fireredasr2.py @@ -233,7 +233,6 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: if mm_data: feature_extractor = self.info.get_feature_extractor(**mm_kwargs) @@ -246,7 +245,6 @@ def _call_hf_processor( prompt=prompt, mm_data=mm_data, mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, ) if "labels" in processed_outputs: processed_outputs["input_ids"] = processed_outputs.pop("labels") diff --git a/vllm/model_executor/models/fireredlid.py b/vllm/model_executor/models/fireredlid.py index 9f1786327b24..32efe3dc0eac 100644 --- a/vllm/model_executor/models/fireredlid.py +++ b/vllm/model_executor/models/fireredlid.py @@ -445,9 +445,9 @@ class FireRedLIDMultiModalProcessor( ): def create_encoder_prompt( self, - prompt: str | list[int], + prompt: list[int], mm_items: MultiModalDataItems, - ) -> str | list[int]: + ) -> list[int]: # Dummy encoder prompt for profiling (encoder only processes audio). return [0] @@ -456,7 +456,6 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: if mm_data: feature_extractor = self.info.get_feature_extractor(**mm_kwargs) @@ -469,7 +468,6 @@ def _call_hf_processor( prompt=prompt, mm_data=mm_data, mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, ) if "labels" in processed_outputs: processed_outputs["input_ids"] = processed_outputs.pop("labels") diff --git a/vllm/model_executor/models/funasr.py b/vllm/model_executor/models/funasr.py index 69f63d770baa..205cb5c0e629 100644 --- a/vllm/model_executor/models/funasr.py +++ b/vllm/model_executor/models/funasr.py @@ -752,7 +752,6 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: if mm_data: feature_extractor = self.info.get_feature_extractor(**mm_kwargs) @@ -765,7 +764,6 @@ def _call_hf_processor( prompt=prompt, mm_data=mm_data, mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, ) if "labels" in processed_outputs: processed_outputs["input_ids"] = processed_outputs.pop("labels") diff --git a/vllm/model_executor/models/funaudiochat.py b/vllm/model_executor/models/funaudiochat.py index 10563f7efc46..255e691df3d4 100644 --- a/vllm/model_executor/models/funaudiochat.py +++ b/vllm/model_executor/models/funaudiochat.py @@ -617,10 +617,9 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: tokenizer = self.info.get_tokenizer() - input_ids = torch.tensor([tokenizer.encode(prompt, **tok_kwargs)]) + input_ids = torch.tensor([tokenizer.encode(prompt)]) audios = mm_data.get("audios", []) if not audios: @@ -679,15 +678,6 @@ def _call_hf_processor( return BatchFeature({"input_ids": input_ids, **mm_inputs}) - def _hf_processor_applies_updates( - self, - prompt_text: str, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - ) -> bool: - return False - def _get_mm_fields_config( self, hf_inputs: BatchFeature, diff --git a/vllm/model_executor/models/gemma3_mm.py b/vllm/model_executor/models/gemma3_mm.py index 5551e0d8b9d3..31d3e82575ce 100644 --- a/vllm/model_executor/models/gemma3_mm.py +++ b/vllm/model_executor/models/gemma3_mm.py @@ -266,13 +266,11 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: processed_outputs = super()._call_hf_processor( prompt, mm_data, mm_kwargs, - tok_kwargs, ) # HF processor pops the `num_crops` kwarg, which is needed by vLLM diff --git a/vllm/model_executor/models/gemma3n_mm.py b/vllm/model_executor/models/gemma3n_mm.py index 3a01f1457aed..787642fa1eb1 100644 --- a/vllm/model_executor/models/gemma3n_mm.py +++ b/vllm/model_executor/models/gemma3n_mm.py @@ -262,7 +262,6 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: # HF Transformers audio processor no longer accepts `audios` key. # We pop `audios` and replace it with `audio` key to suppress @@ -273,7 +272,6 @@ def _call_hf_processor( prompt, mm_data, mm_kwargs, - tok_kwargs, ) if "input_features" in processed_outputs: diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py index aa9669aeb157..e993b1922f5e 100644 --- a/vllm/model_executor/models/gemma4_mm.py +++ b/vllm/model_executor/models/gemma4_mm.py @@ -560,31 +560,11 @@ def _get_dummy_videos( class Gemma4MultiModalProcessor(BaseMultiModalProcessor[Gemma4ProcessingInfo]): - def _apply_hf_processor_text_only( - self, - prompt_text: str, - tokenization_kwargs: Mapping[str, object], - ) -> list[int]: - # Bypass the HF processor and tokenize directly. The HF - # processor expands multimodal placeholders (<|video|>, etc.) - # via get_text_with_replacements, which raises StopIteration - # when the prompt contains placeholders without matching data. - # The text-only path only needs token IDs, so the tokenizer - # alone is sufficient. - processor = self.info.get_hf_processor() - text_inputs = processor.tokenizer([prompt_text], **tokenization_kwargs) - input_ids = text_inputs["input_ids"] - if not isinstance(input_ids, list): - input_ids = input_ids.tolist() - (prompt_ids,) = input_ids - return prompt_ids - def _call_hf_processor( self, prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: merged_kwargs = self.info.ctx.get_merged_mm_kwargs(mm_kwargs) val, is_top_level_max_soft_tokens = _get_max_soft_tokens(merged_kwargs) @@ -640,7 +620,6 @@ def _call_hf_processor( prompt=dummy_prompt, mm_data={"images": frames}, mm_kwargs=video_mm_kwargs, - tok_kwargs=tok_kwargs, ) # Remap HF key name @@ -741,7 +720,6 @@ def _call_hf_processor( prompt, mm_data, patched_mm_kwargs, - tok_kwargs, ) # HF uses 'image_position_ids'; vLLM uses 'pixel_position_ids'. diff --git a/vllm/model_executor/models/glm4_1v.py b/vllm/model_executor/models/glm4_1v.py index 679c90728bd3..c9ab921a8f30 100644 --- a/vllm/model_executor/models/glm4_1v.py +++ b/vllm/model_executor/models/glm4_1v.py @@ -1581,7 +1581,6 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: mm_data = dict(mm_data) if not mm_data: @@ -1602,7 +1601,6 @@ def _call_hf_processor( prompt=prompt, mm_data=prepared_data, mm_kwargs=prepared_kwargs, - tok_kwargs=tok_kwargs, ) if ( @@ -1631,7 +1629,6 @@ def _call_hf_processor( prompt="<|begin_of_video|><|video|><|end_of_video|>", mm_data=video_mm_data, mm_kwargs=video_mm_kwargs, - tok_kwargs=tok_kwargs, ) input_ids = video_outputs.pop("input_ids") if swap_video_frame_tokens: @@ -1659,7 +1656,6 @@ def _call_hf_processor( prompt=prompt, mm_data=mm_data, mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, ) if swap_video_frame_tokens: input_ids = processed_outputs["input_ids"] diff --git a/vllm/model_executor/models/glm4v.py b/vllm/model_executor/models/glm4v.py index 2e3a301579d4..a8512f49fa48 100644 --- a/vllm/model_executor/models/glm4v.py +++ b/vllm/model_executor/models/glm4v.py @@ -466,15 +466,6 @@ def get_dummy_mm_data( class GLM4VMultiModalProcessor(BaseMultiModalProcessor[GLM4VProcessingInfo]): - def _hf_processor_applies_updates( - self, - prompt_text: str, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - ) -> bool: - return False - def _get_mm_fields_config( self, hf_inputs: BatchFeature, diff --git a/vllm/model_executor/models/glmasr.py b/vllm/model_executor/models/glmasr.py index df1b4e0d4b60..d90bbb5634da 100644 --- a/vllm/model_executor/models/glmasr.py +++ b/vllm/model_executor/models/glmasr.py @@ -755,7 +755,6 @@ def _call_hf_processor( prompt: str, mm_data: dict[str, object], mm_kwargs: Mapping[str, Any], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: # Normalize input: handle deprecated key and list conversion. if "audios" in mm_data: @@ -782,7 +781,6 @@ def _call_hf_processor( prompt=prompt, mm_data=mm_data, mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, ) # Postprocess: rename mask and add chunk counts diff --git a/vllm/model_executor/models/granite_speech.py b/vllm/model_executor/models/granite_speech.py index 0e101212744e..626750bd960d 100644 --- a/vllm/model_executor/models/granite_speech.py +++ b/vllm/model_executor/models/granite_speech.py @@ -185,7 +185,6 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: mm_data = dict(mm_data) audios = mm_data.pop("audios", []) @@ -198,7 +197,6 @@ def _call_hf_processor( prompt=prompt, mm_data=mm_data, mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, ) if "audio" in mm_data: diff --git a/vllm/model_executor/models/hunyuan_vision.py b/vllm/model_executor/models/hunyuan_vision.py index 252556e86c43..4454ce47c1a6 100644 --- a/vllm/model_executor/models/hunyuan_vision.py +++ b/vllm/model_executor/models/hunyuan_vision.py @@ -736,7 +736,6 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: hf_processor = self.info.get_hf_processor(**mm_kwargs) # HunYuanVLProcessor requires image placeholders wrapped with start/end tokens. @@ -751,7 +750,7 @@ def _call_hf_processor( return self.info.ctx.call_hf_processor( hf_processor, dict(text=prompt, **mm_data), - dict(**mm_kwargs, **tok_kwargs), + mm_kwargs, ) def _get_prompt_updates( diff --git a/vllm/model_executor/models/hyperclovax_vision.py b/vllm/model_executor/models/hyperclovax_vision.py index 647103aa44e1..e469c0202254 100644 --- a/vllm/model_executor/models/hyperclovax_vision.py +++ b/vllm/model_executor/models/hyperclovax_vision.py @@ -195,7 +195,6 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: for video_idx, video_arr in enumerate(mm_data.get("videos", [])): if video_arr.dtype != np.uint8: @@ -263,15 +262,6 @@ def _call_hf_processor( return processed_outputs - def _hf_processor_applies_updates( - self, - prompt_text: str, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - ) -> bool: - return False - def _get_prompt_updates( self, mm_items: MultiModalDataItems, diff --git a/vllm/model_executor/models/hyperclovax_vision_v2.py b/vllm/model_executor/models/hyperclovax_vision_v2.py index 6cfeac67a527..e50faf885529 100644 --- a/vllm/model_executor/models/hyperclovax_vision_v2.py +++ b/vllm/model_executor/models/hyperclovax_vision_v2.py @@ -209,11 +209,16 @@ def get_dummy_processor_inputs( ) dummy_mm_items = self.info.parse_mm_data(dummy_mm_data, validate=False) + tokenizer = self.info.get_tokenizer() + prompt = tokenizer.encode( + prompt_text, + **self.info.default_tok_params.get_encode_kwargs(), + ) + return ProcessorInputs( - prompt=prompt_text, + prompt=prompt, mm_data_items=dummy_mm_items, hf_processor_mm_kwargs=mm_processor_kwargs or {}, - tokenization_kwargs={"truncation": False}, ) def get_dummy_mm_data( @@ -261,7 +266,6 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: images = mm_data.get("images") videos = mm_data.get("videos") @@ -271,8 +275,7 @@ def _call_hf_processor( # Build data dict for HF processor (images/videos only) # NOTE: We pass the prompt as-is without token normalization. - # Token expansion is handled by vLLM via _get_prompt_updates since - # _hf_processor_applies_updates returns False. + # Token expansion is handled by vLLM via _get_prompt_updates. data: dict[str, object] = dict( text=prompt, images=images, @@ -282,28 +285,11 @@ def _call_hf_processor( processed_outputs = self.info.ctx.call_hf_processor( hf_processor=hf_processor, data=data, - kwargs=dict(**mm_kwargs, **tok_kwargs), + kwargs=mm_kwargs, ) return processed_outputs - def _hf_processor_applies_updates( - self, - prompt_text: str, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - ) -> bool: - # Match BaseMultiModalProcessor behavior: - # - raw multimodal inputs: HF processor applies updates - # - embedding inputs: vLLM applies updates - return super()._hf_processor_applies_updates( - prompt_text, - mm_items, - hf_processor_mm_kwargs, - tokenization_kwargs, - ) - def _get_prompt_updates( self, mm_items: MultiModalDataItems, diff --git a/vllm/model_executor/models/idefics3.py b/vllm/model_executor/models/idefics3.py index 7b3e552f0cdd..ab6c64f46f19 100644 --- a/vllm/model_executor/models/idefics3.py +++ b/vllm/model_executor/models/idefics3.py @@ -244,7 +244,6 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: # Text-only input not supported in composite processor if not (images := mm_data.get("images", [])): @@ -260,7 +259,6 @@ def _call_hf_processor( prompt, mm_data, mm_kwargs, - tok_kwargs, ) mm_items = self.info.parse_mm_data({"image": images}, validate=False) diff --git a/vllm/model_executor/models/interns1.py b/vllm/model_executor/models/interns1.py index 6dbc358e089e..bb3ba6a55620 100644 --- a/vllm/model_executor/models/interns1.py +++ b/vllm/model_executor/models/interns1.py @@ -337,7 +337,6 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: mm_data = dict(mm_data) videos = mm_data.pop("videos", []) @@ -364,7 +363,6 @@ def _call_hf_processor( prompt=hf_processor.image_token, mm_data={"images": image}, mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, ) image_pixel_values.append(processed_outputs.pop("pixel_values")) @@ -387,7 +385,6 @@ def _call_hf_processor( prompt=hf_processor.video_token, mm_data={"videos": video}, mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, ) video_pixel_values.append(processed_outputs.pop("pixel_values")) @@ -406,7 +403,7 @@ def _call_hf_processor( prompt = re.sub("", hf_processor.image_token, prompt) prompt = re.sub("", hf_processor.video_token, prompt) - text_outputs = tokenizer(prompt, **tok_kwargs, return_tensors="pt") + text_outputs = tokenizer(prompt, return_tensors="pt") return BatchFeature({**text_outputs, **image_outputs, **video_outputs}) diff --git a/vllm/model_executor/models/internvl.py b/vllm/model_executor/models/internvl.py index d360e5018283..f818ce9dfdb5 100644 --- a/vllm/model_executor/models/internvl.py +++ b/vllm/model_executor/models/internvl.py @@ -225,13 +225,11 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: processed_outputs = super()._call_hf_processor( prompt=prompt, mm_data=mm_data, mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, ) hf_processor = self.info.get_hf_processor(**mm_kwargs) @@ -462,11 +460,8 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: - processed_outputs = super()._call_hf_processor( - prompt, mm_data, mm_kwargs, tok_kwargs - ) + processed_outputs = super()._call_hf_processor(prompt, mm_data, mm_kwargs) hf_processor = self.info.get_hf_processor(**mm_kwargs) if (video_token_id := hf_processor.ctx_video_token_id) is not None: diff --git a/vllm/model_executor/models/jina_vl.py b/vllm/model_executor/models/jina_vl.py index 6970f74a2768..c1cdf207e329 100644 --- a/vllm/model_executor/models/jina_vl.py +++ b/vllm/model_executor/models/jina_vl.py @@ -59,7 +59,6 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: # NOTE: We should reverse the order of the mm_data because the # query prompt is placed after the document prompt in the score @@ -67,7 +66,7 @@ def _call_hf_processor( # stored in the opposite order (query first, then document). for _, value in mm_data.items(): value.reverse() - return super()._call_hf_processor(prompt, mm_data, mm_kwargs, tok_kwargs) + return super()._call_hf_processor(prompt, mm_data, mm_kwargs) @MULTIMODAL_REGISTRY.register_processor( diff --git a/vllm/model_executor/models/kanana_v.py b/vllm/model_executor/models/kanana_v.py index 7a1421b515d7..88999149977e 100644 --- a/vllm/model_executor/models/kanana_v.py +++ b/vllm/model_executor/models/kanana_v.py @@ -466,7 +466,6 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: """Run the underlying HF processor on text and image data.""" # Text-only input is handled as a special case here. diff --git a/vllm/model_executor/models/keye.py b/vllm/model_executor/models/keye.py index d64d2ebc64e0..dd894e60b108 100644 --- a/vllm/model_executor/models/keye.py +++ b/vllm/model_executor/models/keye.py @@ -1167,11 +1167,10 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: # Override to use the text path instead of token path to use the # video-specific logic in processing_keye.py - return super()._call_hf_processor(prompt, mm_data, mm_kwargs, tok_kwargs) + return super()._call_hf_processor(prompt, mm_data, mm_kwargs) def _get_prompt_updates( self, diff --git a/vllm/model_executor/models/keye_vl1_5.py b/vllm/model_executor/models/keye_vl1_5.py index a4dcd729db3b..80a94373cb4a 100644 --- a/vllm/model_executor/models/keye_vl1_5.py +++ b/vllm/model_executor/models/keye_vl1_5.py @@ -383,11 +383,10 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: # Override to use the text path instead of token path to use the # video-specific logic in processing_keye.py - return super()._call_hf_processor(prompt, mm_data, mm_kwargs, tok_kwargs) + return super()._call_hf_processor(prompt, mm_data, mm_kwargs) def _get_prompt_updates( self, diff --git a/vllm/model_executor/models/kimi_audio.py b/vllm/model_executor/models/kimi_audio.py index c999a694ad4e..60534658ae6d 100644 --- a/vllm/model_executor/models/kimi_audio.py +++ b/vllm/model_executor/models/kimi_audio.py @@ -231,7 +231,6 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: """Call the HuggingFace processor.""" # Convert mm_data format: {'audios': [...]} -> {'audio': ...} @@ -256,7 +255,7 @@ def _call_hf_processor( return self.info.ctx.call_hf_processor( self.info.get_hf_processor(**mm_kwargs), dict(text=prompt, **mm_data), - dict(**mm_kwargs, **tok_kwargs), + mm_kwargs, ) def _get_mm_fields_config( diff --git a/vllm/model_executor/models/kimi_k25.py b/vllm/model_executor/models/kimi_k25.py index b5c1c925d1d0..3054fcc5d89f 100644 --- a/vllm/model_executor/models/kimi_k25.py +++ b/vllm/model_executor/models/kimi_k25.py @@ -263,11 +263,10 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: # Override to use the text path instead of token path because vision chunk # is not considered - return super()._call_hf_processor(prompt, mm_data, mm_kwargs, tok_kwargs) + return super()._call_hf_processor(prompt, mm_data, mm_kwargs) def _get_prompt_updates( self, diff --git a/vllm/model_executor/models/lfm2_vl.py b/vllm/model_executor/models/lfm2_vl.py index d247fef4d7e3..85286712feb6 100644 --- a/vllm/model_executor/models/lfm2_vl.py +++ b/vllm/model_executor/models/lfm2_vl.py @@ -398,7 +398,6 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: # Text-only input not supported in composite processor if not (images := mm_data.get("images", [])): @@ -412,7 +411,6 @@ def _call_hf_processor( prompt, mm_data, mm_kwargs, - tok_kwargs, ) mm_items = self.info.parse_mm_data({"image": images}, validate=False) diff --git a/vllm/model_executor/models/lightonocr.py b/vllm/model_executor/models/lightonocr.py index c1ee640f63a5..83133e400f84 100644 --- a/vllm/model_executor/models/lightonocr.py +++ b/vllm/model_executor/models/lightonocr.py @@ -44,13 +44,11 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: processed_outputs = super()._call_hf_processor( prompt=prompt, mm_data=mm_data, mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, ) # NOTE: LightOnOCR does not use break/end tokens, so we remove them here. diff --git a/vllm/model_executor/models/llava.py b/vllm/model_executor/models/llava.py index 1e850a7efc6e..9ef36e182446 100644 --- a/vllm/model_executor/models/llava.py +++ b/vllm/model_executor/models/llava.py @@ -326,13 +326,11 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: processed_outputs = super()._call_hf_processor( prompt=prompt, mm_data=mm_data, mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, ) pixel_values = processed_outputs.get("pixel_values") diff --git a/vllm/model_executor/models/llava_onevision.py b/vllm/model_executor/models/llava_onevision.py index 1beec4207d53..7b8896f9e0d1 100644 --- a/vllm/model_executor/models/llava_onevision.py +++ b/vllm/model_executor/models/llava_onevision.py @@ -326,7 +326,6 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: mm_data = dict(mm_data) videos = mm_data.pop("videos", []) @@ -337,7 +336,6 @@ def _call_hf_processor( prompt=prompt, mm_data=mm_data, mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, ) # LLaVA-OneVision processor doesn't support multiple videos @@ -352,7 +350,6 @@ def _call_hf_processor( prompt=prompt, mm_data={}, mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, ) images = mm_data.pop("images", []) @@ -362,7 +359,6 @@ def _call_hf_processor( prompt=image_token * len(images), mm_data={"images": images}, mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, ) image_outputs = { k: v @@ -378,7 +374,6 @@ def _call_hf_processor( prompt=video_token, mm_data={"videos": video}, mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, ) pixel_values_videos.append(item_outputs["pixel_values_videos"][0]) @@ -392,22 +387,6 @@ def _call_hf_processor( ) return BatchFeature(combined_outputs) - def _hf_processor_applies_updates( - self, - prompt_text: str, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - ) -> bool: - base_result = super()._hf_processor_applies_updates( - prompt_text=prompt_text, - mm_items=mm_items, - hf_processor_mm_kwargs=hf_processor_mm_kwargs, - tokenization_kwargs=tokenization_kwargs, - ) - - return base_result and mm_items.get_count("video", strict=False) == 0 - def _get_prompt_updates( self, mm_items: MultiModalDataItems, diff --git a/vllm/model_executor/models/llava_onevision2.py b/vllm/model_executor/models/llava_onevision2.py index f094a168fb21..be51b1f0f5f3 100644 --- a/vllm/model_executor/models/llava_onevision2.py +++ b/vllm/model_executor/models/llava_onevision2.py @@ -1552,16 +1552,13 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: # The wrapped OV2 processor is a bare custom class without the standard # ProcessorMixin ``_merge_kwargs`` machinery, so vLLM's default path # fails; overriding this method routes the base class to call us # directly. hf_processor = self.info.get_hf_processor(**mm_kwargs) - merged_kwargs = self.info.ctx.get_merged_mm_kwargs( - dict(**mm_kwargs, **tok_kwargs) - ) + merged_kwargs = self.info.ctx.get_merged_mm_kwargs(mm_kwargs) merged_kwargs.setdefault("return_tensors", "pt") call_kwargs = { k: v @@ -1638,7 +1635,6 @@ def _call_hf_processor( prompt=prompt, mm_data=mm_data, mm_kwargs={**mm_kwargs, "video_backend": "codec"}, - tok_kwargs=tok_kwargs, ) data = dict(output) return BatchFeature( @@ -1736,7 +1732,6 @@ def _call_hf_processor( prompt=new_prompt, mm_data=merged_mm_data, mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, ) data = dict(output) @@ -1803,7 +1798,6 @@ def _gather_rows(tensor, rows): prompt=prompt, mm_data=mm_data, mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, ) def _rename_codec_outputs_to_video( diff --git a/vllm/model_executor/models/midashenglm.py b/vllm/model_executor/models/midashenglm.py index 5ecc92e4d04b..8ddfd05f7294 100644 --- a/vllm/model_executor/models/midashenglm.py +++ b/vllm/model_executor/models/midashenglm.py @@ -588,7 +588,6 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, Any], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: audios = mm_data.pop("audios", []) @@ -622,7 +621,6 @@ def _call_hf_processor( prompt=prompt, mm_data=mm_data, mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, ) def _get_mm_fields_config( diff --git a/vllm/model_executor/models/mimo_v2_omni.py b/vllm/model_executor/models/mimo_v2_omni.py index 28205f9da29d..8f78409581f8 100644 --- a/vllm/model_executor/models/mimo_v2_omni.py +++ b/vllm/model_executor/models/mimo_v2_omni.py @@ -917,7 +917,6 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: """Convert numpy video arrays to (TCHW, timestamps) tuples for MiMo. Also remap 'audios' → 'audio' since MiMoOmniProcessor.__call__ uses @@ -1000,7 +999,7 @@ def _call_hf_processor( mm_data = {**mm_data, "videos": converted} - return super()._call_hf_processor(prompt, mm_data, mm_kwargs, tok_kwargs) + return super()._call_hf_processor(prompt, mm_data, mm_kwargs) def _get_prompt_updates( self, diff --git a/vllm/model_executor/models/minicpmo.py b/vllm/model_executor/models/minicpmo.py index dd6f8d2f56b9..3dd3f8fd3b0c 100644 --- a/vllm/model_executor/models/minicpmo.py +++ b/vllm/model_executor/models/minicpmo.py @@ -436,7 +436,6 @@ def process_audios( self, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> Mapping[str, NestedTensors]: if (audios := mm_data.get("audios")) is None: return {} @@ -453,7 +452,6 @@ def process_audios( prompts=[self.info.audio_pattern] * len(parsed_audios), mm_data={"audios": [[audio] for audio in parsed_audios]}, mm_kwargs={**mm_kwargs, "chunk_input": True}, - tok_kwargs=tok_kwargs, out_keys={"audio_features", "audio_feature_lens"}, ) @@ -486,11 +484,10 @@ def process_mm_inputs( self, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> Mapping[str, NestedTensors]: return { - **super().process_mm_inputs(mm_data, mm_kwargs, tok_kwargs), - **self.process_audios(mm_data, mm_kwargs, tok_kwargs), + **super().process_mm_inputs(mm_data, mm_kwargs), + **self.process_audios(mm_data, mm_kwargs), } def _get_prompt_updates( diff --git a/vllm/model_executor/models/minicpmv.py b/vllm/model_executor/models/minicpmv.py index d7ab703f38c8..968490d4e9de 100644 --- a/vllm/model_executor/models/minicpmv.py +++ b/vllm/model_executor/models/minicpmv.py @@ -869,7 +869,6 @@ def process_images( self, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> Mapping[str, NestedTensors]: if (images := mm_data.get("images")) is None: return {} @@ -886,7 +885,6 @@ def process_images( prompts=[self.info.image_pattern] * len(parsed_images), mm_data={"images": [[image] for image in parsed_images]}, mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, out_keys={"pixel_values", "image_sizes", "tgt_sizes"}, ) @@ -896,7 +894,6 @@ def process_videos( self, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> Mapping[str, NestedTensors]: if (videos := mm_data.get("videos")) is None: return {} @@ -918,7 +915,6 @@ def process_videos( **mm_kwargs, "max_slice_nums": self.info.get_video_max_slice_num(), }, - tok_kwargs=tok_kwargs, out_keys={"pixel_values", "image_sizes", "tgt_sizes"}, ) @@ -930,11 +926,10 @@ def process_mm_inputs( self, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> Mapping[str, NestedTensors]: return { - **self.process_images(mm_data, mm_kwargs, tok_kwargs), - **self.process_videos(mm_data, mm_kwargs, tok_kwargs), + **self.process_images(mm_data, mm_kwargs), + **self.process_videos(mm_data, mm_kwargs), } def _apply_prompt_updates( @@ -993,7 +988,6 @@ def _base_call_hf_processor( prompts: list[str], mm_data: Mapping[str, Sequence[object]], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], *, out_keys: set[str], ) -> dict[str, NestedTensors]: @@ -1003,7 +997,6 @@ def _base_call_hf_processor( prompt=prompts, # type: ignore mm_data=mm_data, mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, ) else: inputs = defaultdict[str, list[torch.Tensor]](list) @@ -1013,7 +1006,6 @@ def _base_call_hf_processor( prompt=prompt, mm_data={k: v[i] for k, v in mm_data.items()}, mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, ) for k, v in inputs_one.items(): @@ -1027,12 +1019,11 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: tokenizer = self.info.get_tokenizer() - input_ids = torch.tensor([tokenizer.encode(prompt, **tok_kwargs)]) - mm_inputs = self.process_mm_inputs(mm_data, mm_kwargs, tok_kwargs) + input_ids = torch.tensor([tokenizer.encode(prompt)]) + mm_inputs = self.process_mm_inputs(mm_data, mm_kwargs) return BatchFeature( { @@ -1041,15 +1032,6 @@ def _call_hf_processor( } ) - def _hf_processor_applies_updates( - self, - prompt_text: str, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - ) -> bool: - return False - def _get_prompt_updates( self, mm_items: MultiModalDataItems, diff --git a/vllm/model_executor/models/minicpmv4_6.py b/vllm/model_executor/models/minicpmv4_6.py index d0948715a9d1..501be460504f 100644 --- a/vllm/model_executor/models/minicpmv4_6.py +++ b/vllm/model_executor/models/minicpmv4_6.py @@ -154,7 +154,6 @@ def process_images( self, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> Mapping[str, NestedTensors]: if (images := mm_data.get("images")) is None: return {} @@ -221,7 +220,6 @@ def process_videos( self, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> Mapping[str, NestedTensors]: if (videos := mm_data.get("videos")) is None: return {} diff --git a/vllm/model_executor/models/mistral3.py b/vllm/model_executor/models/mistral3.py index 025ce564083c..22adf1c484af 100644 --- a/vllm/model_executor/models/mistral3.py +++ b/vllm/model_executor/models/mistral3.py @@ -238,13 +238,11 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: processed_outputs = super()._call_hf_processor( prompt=prompt, mm_data=mm_data, mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, ) pixel_values = processed_outputs.get("pixel_values") diff --git a/vllm/model_executor/models/mllama4.py b/vllm/model_executor/models/mllama4.py index fe5ff17a3b46..04748465e08a 100644 --- a/vllm/model_executor/models/mllama4.py +++ b/vllm/model_executor/models/mllama4.py @@ -594,13 +594,11 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: processed_outputs = super()._call_hf_processor( prompt=prompt, mm_data=mm_data, mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, ) processor = self.info.get_hf_processor(**mm_kwargs) diff --git a/vllm/model_executor/models/molmo.py b/vllm/model_executor/models/molmo.py index 3fccb7a86cd2..2b963e339a92 100644 --- a/vllm/model_executor/models/molmo.py +++ b/vllm/model_executor/models/molmo.py @@ -1128,13 +1128,12 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: hf_processor = self.info.get_hf_processor(**mm_kwargs) processed_outputs = self.info.ctx.call_hf_processor( hf_processor.process, dict(text=prompt, **mm_data), - dict(**mm_kwargs, **tok_kwargs), + mm_kwargs, ) tokenizer = hf_processor.tokenizer diff --git a/vllm/model_executor/models/molmo2.py b/vllm/model_executor/models/molmo2.py index 42f5ebfdc5f6..861e31928bdc 100644 --- a/vllm/model_executor/models/molmo2.py +++ b/vllm/model_executor/models/molmo2.py @@ -1948,7 +1948,6 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: mm_data = dict(mm_data) @@ -2008,7 +2007,7 @@ def patched_call(text=None, images=None, videos=None, **kwargs) -> BatchFeature: video_outputs = self.info.ctx.call_hf_processor( patched_call, dict(text=VIDEO_PROMPT, **video_mm_data), - dict(**video_mm_kwargs, **tok_kwargs), + video_mm_kwargs, ) input_ids = video_outputs.pop("input_ids") @@ -2058,7 +2057,7 @@ def patched_call(text=None, images=None, videos=None, **kwargs) -> BatchFeature: processed_outputs = self.info.ctx.call_hf_processor( patched_call, dict(text=prompt, **mm_data), - dict(**mm_kwargs, **tok_kwargs), + mm_kwargs, ) if (images := mm_data.get("images")) is not None: diff --git a/vllm/model_executor/models/moondream3.py b/vllm/model_executor/models/moondream3.py index d5f3e6b195fb..0464150bd49e 100644 --- a/vllm/model_executor/models/moondream3.py +++ b/vllm/model_executor/models/moondream3.py @@ -947,11 +947,10 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: # Moondream3's processor handles images directly rather than exposing a # separate `image_processor`, so keep the cache path on text+MM calls. - return super()._call_hf_processor(prompt, mm_data, mm_kwargs, tok_kwargs) + return super()._call_hf_processor(prompt, mm_data, mm_kwargs) @cached_property def bos_image_placeholder_tokens(self) -> list[int]: @@ -977,18 +976,6 @@ def _get_mm_fields_config( "tilings": MultiModalFieldConfig.batched("image", keep_on_cpu=True), } - def _hf_processor_applies_updates( - self, - prompt_text: str, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - ) -> bool: - # Moondream3 HF processor does NOT expand placeholder tokens. - # vLLM expands BOS + so the whole HF image prefix is marked - # bidirectional by the multimodal prefix-LM mask. - return False - def _get_prompt_updates( self, mm_items: MultiModalDataItems, diff --git a/vllm/model_executor/models/moss_audio.py b/vllm/model_executor/models/moss_audio.py index 324b3f55ae8a..ceed6245a1ca 100644 --- a/vllm/model_executor/models/moss_audio.py +++ b/vllm/model_executor/models/moss_audio.py @@ -49,6 +49,7 @@ BaseDummyInputsBuilder, BaseMultiModalProcessor, BaseProcessingInfo, + InputProcessingContext, PromptReplacement, PromptUpdate, PromptUpdateDetails, @@ -102,6 +103,34 @@ "enable_time_marker", "mel_config", } +MOSS_AUDIO_PLACEHOLDER_TOKENS = ( + MOSS_AUDIO_BOS_TOKEN, + MOSS_AUDIO_TOKEN, + MOSS_AUDIO_EOS_TOKEN, +) + + +def _ensure_moss_audio_placeholder_tokens(tokenizer: object) -> None: + """Register the audio placeholder tokens if the tokenizer lacks them. + + The MOSS-Audio hub tokenizer does not define the audio placeholder + tokens (the reference HF processor patches them in at runtime). Left + as plain text, they merge with adjacent tokens under BPE, so prompts + no longer contain a stable token sequence for prompt updates to match. + """ + convert_tokens_to_ids = getattr(tokenizer, "convert_tokens_to_ids", None) + add_tokens = getattr(tokenizer, "add_tokens", None) + if convert_tokens_to_ids is None or add_tokens is None: + return + + unk_token_id = getattr(tokenizer, "unk_token_id", None) + missing_tokens = [ + token + for token in MOSS_AUDIO_PLACEHOLDER_TOKENS + if convert_tokens_to_ids(token) in (None, unk_token_id) + ] + if missing_tokens: + add_tokens(missing_tokens, special_tokens=True) class MossAudioAudioInputs(TensorSchema): @@ -1157,6 +1186,11 @@ def batch_decode(self, *args: object, **kwargs: object) -> list[str]: class MossAudioProcessingInfo(BaseProcessingInfo): + def __init__(self, ctx: InputProcessingContext) -> None: + super().__init__(ctx) + if ctx.tokenizer is not None: + _ensure_moss_audio_placeholder_tokens(ctx.tokenizer) + def get_hf_config(self) -> MossAudioConfig: config = self.ctx.get_hf_config() if isinstance(config, MossAudioConfig): @@ -1306,7 +1340,6 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: mm_data = dict(mm_data) audios = mm_data.pop("audios", []) @@ -1314,15 +1347,10 @@ def _call_hf_processor( mm_data["audio"] = audios mm_kwargs = dict(mm_kwargs) processor_kwargs = _filter_moss_audio_processor_config(mm_kwargs) - tok_kwargs = { - key: value - for key, value in tok_kwargs.items() - if key not in MOSS_AUDIO_PROCESSOR_CONFIG_KEYS - } return self.info.ctx.call_hf_processor( self.info.get_hf_processor(**processor_kwargs), dict(text=prompt, **mm_data), - dict(**tok_kwargs), + {}, ) def _get_mm_fields_config( diff --git a/vllm/model_executor/models/moss_transcribe_diarize.py b/vllm/model_executor/models/moss_transcribe_diarize.py index a862c27e6f25..421901d0ee8f 100644 --- a/vllm/model_executor/models/moss_transcribe_diarize.py +++ b/vllm/model_executor/models/moss_transcribe_diarize.py @@ -466,33 +466,20 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: tokenizer = self.info.get_tokenizer() audios = _get_audios_from_mm_data(mm_data) if not audios: - input_ids = tokenizer.encode( - prompt, - add_special_tokens=tok_kwargs.get("add_special_tokens", False), - ) + input_ids = tokenizer.encode(prompt, add_special_tokens=False) return BatchFeature({"input_ids": [input_ids]}, tensor_type="pt") processed = self.info.ctx.call_hf_processor( self.info.get_hf_processor(**mm_kwargs), dict(text=prompt, audio=audios), - dict(**mm_kwargs, **tok_kwargs), + mm_kwargs, ) return _add_vllm_audio_metadata(processed, len(audios)) - def _hf_processor_applies_updates( - self, - prompt_text: str, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - ) -> bool: - return mm_items.get_count("audio", strict=False) > 0 - def _get_mm_fields_config( self, hf_inputs: BatchFeature, diff --git a/vllm/model_executor/models/muse_glimmer.py b/vllm/model_executor/models/muse_glimmer.py index 83ec17a75d56..fcdfb82333cd 100644 --- a/vllm/model_executor/models/muse_glimmer.py +++ b/vllm/model_executor/models/muse_glimmer.py @@ -262,23 +262,11 @@ def get_dummy_mm_data( class MuseGlimmerMultiModalProcessor( BaseMultiModalProcessor[MuseGlimmerProcessingInfo] ): - def _apply_hf_processor_text_only( - self, - prompt_text: str, - tokenization_kwargs: Mapping[str, object], - ) -> list[int]: - tokenizer = self.info.get_tokenizer() - return tokenizer.encode( - prompt_text, - **{"add_special_tokens": False, **tokenization_kwargs}, - ) - def _call_hf_processor( self, prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: processor = self.info.get_hf_processor(**mm_kwargs) tokenizer = processor.tokenizer @@ -366,7 +354,6 @@ def _call_hf_processor( video_pixel_values=video_pixels, video_feature_sizes=torch.tensor(video_sizes), ) - del tok_kwargs return BatchFeature(data=data, tensor_type=None) def _get_mm_fields_config( diff --git a/vllm/model_executor/models/nano_nemotron_vl.py b/vllm/model_executor/models/nano_nemotron_vl.py index 7ecf8677bb33..679495e3ecf0 100644 --- a/vllm/model_executor/models/nano_nemotron_vl.py +++ b/vllm/model_executor/models/nano_nemotron_vl.py @@ -366,14 +366,13 @@ def _call_hf_processor( prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: """ Bypass `call_hf_processor_mm_only` by no-op overriding`_call_hf_processor`, so it chooses this path: `type(self)._call_hf_processor != BaseMultiModalProcessor._call_hf_processor` """ - return super()._call_hf_processor(prompt, mm_data, mm_kwargs, tok_kwargs) + return super()._call_hf_processor(prompt, mm_data, mm_kwargs) def _get_image_fields_config(self, hf_inputs: BatchFeature): if self.info.is_dynamic_tiler: @@ -714,10 +713,8 @@ def apply( if not audio_items: return super().apply(inputs, timing_ctx) - prompt = inputs.prompt tokenizer = self.info.get_tokenizer() - if not isinstance(prompt, str): - prompt = tokenizer.decode(prompt, skip_special_tokens=False) + prompt = tokenizer.decode(inputs.prompt, skip_special_tokens=False) # Inject AUDIO_CONTEXT only after