From 4c4b6f7a9764bac8bf9f2a0bfedf852d8e59c98e Mon Sep 17 00:00:00 2001 From: Daniel Mescheder Date: Tue, 3 Feb 2026 06:51:10 +0100 Subject: [PATCH 001/810] [Frontend] Add sampling parameters to Responses API (#32609) Signed-off-by: Daniel Mescheder Co-authored-by: Daniel Mescheder --- .../openai/responses/test_sampling_params.py | 113 ++++++++++++++++++ .../openai/responses/test_simple.py | 25 +++- vllm/entrypoints/openai/responses/protocol.py | 29 ++++- 3 files changed, 165 insertions(+), 2 deletions(-) create mode 100644 tests/entrypoints/openai/responses/test_sampling_params.py diff --git a/tests/entrypoints/openai/responses/test_sampling_params.py b/tests/entrypoints/openai/responses/test_sampling_params.py new file mode 100644 index 00000000000..b8d1aa66404 --- /dev/null +++ b/tests/entrypoints/openai/responses/test_sampling_params.py @@ -0,0 +1,113 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Unit tests for ResponsesRequest.to_sampling_params() parameter mapping.""" + +import pytest + +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest + + +class TestResponsesRequestSamplingParams: + """Test that ResponsesRequest correctly maps parameters to SamplingParams.""" + + def test_basic_sampling_params(self): + """Test basic sampling parameters are correctly mapped.""" + request = ResponsesRequest( + model="test-model", + input="test input", + temperature=0.8, + top_p=0.95, + top_k=50, + max_output_tokens=100, + ) + + sampling_params = request.to_sampling_params(default_max_tokens=1000) + + assert sampling_params.temperature == 0.8 + assert sampling_params.top_p == 0.95 + assert sampling_params.top_k == 50 + assert sampling_params.max_tokens == 100 + + def test_extra_sampling_params(self): + """Test extra sampling parameters are correctly mapped.""" + request = ResponsesRequest( + model="test-model", + input="test input", + repetition_penalty=1.2, + seed=42, + stop=["END", "STOP"], + ignore_eos=True, + vllm_xargs={"custom": "value"}, + ) + + sampling_params = request.to_sampling_params(default_max_tokens=1000) + + assert sampling_params.repetition_penalty == 1.2 + assert sampling_params.seed == 42 + assert sampling_params.stop == ["END", "STOP"] + assert sampling_params.ignore_eos is True + assert sampling_params.extra_args == {"custom": "value"} + + def test_stop_string_conversion(self): + """Test that single stop string is converted to list.""" + request = ResponsesRequest( + model="test-model", + input="test input", + stop="STOP", + ) + + sampling_params = request.to_sampling_params(default_max_tokens=1000) + + assert sampling_params.stop == ["STOP"] + + def test_default_values(self): + """Test default values for optional parameters.""" + request = ResponsesRequest( + model="test-model", + input="test input", + ) + + sampling_params = request.to_sampling_params(default_max_tokens=1000) + + assert sampling_params.repetition_penalty == 1.0 # None → 1.0 + assert sampling_params.stop == [] # Empty list + assert sampling_params.extra_args == {} # Empty dict + + def test_seed_bounds_validation(self): + """Test that seed values outside torch.long bounds are rejected.""" + import torch + from pydantic import ValidationError + + # Test seed below minimum + with pytest.raises(ValidationError) as exc_info: + ResponsesRequest( + model="test-model", + input="test input", + seed=torch.iinfo(torch.long).min - 1, + ) + assert "greater_than_equal" in str(exc_info.value).lower() + + # Test seed above maximum + with pytest.raises(ValidationError) as exc_info: + ResponsesRequest( + model="test-model", + input="test input", + seed=torch.iinfo(torch.long).max + 1, + ) + assert "less_than_equal" in str(exc_info.value).lower() + + # Test valid seed at boundaries + request_min = ResponsesRequest( + model="test-model", + input="test input", + seed=torch.iinfo(torch.long).min, + ) + assert request_min.seed == torch.iinfo(torch.long).min + + request_max = ResponsesRequest( + model="test-model", + input="test input", + seed=torch.iinfo(torch.long).max, + ) + assert request_max.seed == torch.iinfo(torch.long).max diff --git a/tests/entrypoints/openai/responses/test_simple.py b/tests/entrypoints/openai/responses/test_simple.py index 30423788bf7..8f07b02a308 100644 --- a/tests/entrypoints/openai/responses/test_simple.py +++ b/tests/entrypoints/openai/responses/test_simple.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project - import pytest import pytest_asyncio from openai import OpenAI @@ -147,3 +146,27 @@ async def test_max_tokens(client: OpenAI, model_name: str): assert response is not None assert response.status == "incomplete" assert response.incomplete_details.reason == "max_output_tokens" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model_name", [MODEL_NAME]) +async def test_extra_sampling_params(client: OpenAI, model_name: str): + """Test that extra sampling parameters are accepted and work.""" + # Test with multiple sampling parameters - just verify they're accepted + response = await client.responses.create( + model=model_name, + input="Write a short sentence", + max_output_tokens=50, + temperature=0.7, + top_p=0.9, + extra_body={ + "top_k": 40, + "repetition_penalty": 1.2, + "seed": 42, + }, + ) + + # Verify request succeeded and parameters were accepted + assert response.status in ["completed", "incomplete"] + assert len(response.output) > 0 + assert response.output[0].content[0].text # Has text output diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index 81abebee291..9a471852ba2 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -6,6 +6,7 @@ import time from typing import Any, Literal, TypeAlias +import torch from openai.types.responses import ( ResponseCodeInterpreterCallCodeDeltaEvent, ResponseCodeInterpreterCallCodeDoneEvent, @@ -77,6 +78,8 @@ from vllm.utils import random_uuid logger = init_logger(__name__) +_LONG_INFO = torch.iinfo(torch.long) + class InputTokensDetails(OpenAIBaseModel): cached_tokens: int @@ -230,6 +233,18 @@ class ResponsesRequest(OpenAIBaseModel): # this cannot be used in conjunction with previous_response_id # TODO: consider supporting non harmony messages as well previous_input_messages: list[OpenAIHarmonyMessage | dict] | None = None + + repetition_penalty: float | None = None + seed: int | None = Field(None, ge=_LONG_INFO.min, le=_LONG_INFO.max) + stop: str | list[str] | None = [] + ignore_eos: bool = False + vllm_xargs: dict[str, str | int | float | list[str | int | float]] | None = Field( + default=None, + description=( + "Additional request parameters with (list of) string or " + "numeric values, used by custom extensions." + ), + ) # --8<-- [end:responses-extra-params] def build_chat_params( @@ -297,6 +312,10 @@ class ResponsesRequest(OpenAIBaseModel): top_k = default_sampling_params.get( "top_k", self._DEFAULT_SAMPLING_PARAMS["top_k"] ) + + if (repetition_penalty := self.repetition_penalty) is None: + repetition_penalty = default_sampling_params.get("repetition_penalty", 1.0) + stop_token_ids = default_sampling_params.get("stop_token_ids") # Structured output @@ -313,7 +332,10 @@ class ResponsesRequest(OpenAIBaseModel): elif response_format.type == "json_object": raise NotImplementedError("json_object is not supported") - # TODO: add more parameters + stop = self.stop if self.stop else [] + if isinstance(stop, str): + stop = [stop] + return SamplingParams.from_optional( temperature=temperature, top_p=top_p, @@ -321,11 +343,16 @@ class ResponsesRequest(OpenAIBaseModel): max_tokens=max_tokens, logprobs=self.top_logprobs if self.is_include_output_logprobs() else None, stop_token_ids=stop_token_ids, + stop=stop, + repetition_penalty=repetition_penalty, + seed=self.seed, + ignore_eos=self.ignore_eos, output_kind=( RequestOutputKind.DELTA if self.stream else RequestOutputKind.FINAL_ONLY ), structured_outputs=structured_outputs, logit_bias=self.logit_bias, + extra_args=self.vllm_xargs or {}, skip_clone=True, # Created fresh per request, safe to skip clone skip_special_tokens=self.skip_special_tokens, include_stop_str_in_output=self.include_stop_str_in_output, From f1cb9b554492bd198ea63cf2233f8599aa850723 Mon Sep 17 00:00:00 2001 From: Shengliang Xu <106840466+shengliangxu@users.noreply.github.com> Date: Mon, 2 Feb 2026 22:31:27 -0800 Subject: [PATCH 002/810] Fix quantized Falcon-H1 model loading issues (#32728) Signed-off-by: Shengliang Xu Co-authored-by: Cyrus Leung --- vllm/model_executor/models/falcon_h1.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/models/falcon_h1.py b/vllm/model_executor/models/falcon_h1.py index 3d4d253c390..fba2e216e3f 100644 --- a/vllm/model_executor/models/falcon_h1.py +++ b/vllm/model_executor/models/falcon_h1.py @@ -35,7 +35,10 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) -from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) from vllm.sequence import IntermediateTensors from vllm.transformers_utils.config import set_default_rope_theta @@ -278,6 +281,7 @@ class FalconH1AttentionDecoderLayer(nn.Module): self.scaling, num_kv_heads=self.num_kv_heads, cache_config=cache_config, + quant_config=quant_config, prefix=f"{prefix}.attn", ) self.key_multiplier = config.key_multiplier @@ -360,7 +364,9 @@ class FalconH1ParallelHybrid(nn.Module): self.attention_in_multiplier = config.attention_in_multiplier self.attn_out_multiplier = config.attention_out_multiplier - self.feed_forward = FalconH1MLP(config, prefix=f"{prefix}.feed_forward") + self.feed_forward = FalconH1MLP( + config, quant_config=quant_config, prefix=f"{prefix}.feed_forward" + ) self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.pre_ff_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -647,6 +653,12 @@ class FalconH1ForCausalLM( if "mamba" in name: name = name.replace("mamba", "mamba.mamba") + if "scale" in name: + # Remapping the name of kv-scale. + name = maybe_remap_kv_scale_name(name, params_dict) + if name is None: + continue + for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: continue From a0a984ac2e4503de1a76f55ece65ac0847678503 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9D=A8=E6=9C=B1=20=C2=B7=20Kiki?= Date: Tue, 3 Feb 2026 14:32:39 +0800 Subject: [PATCH 003/810] [CI/Build] Remove hardcoded America/Los_Angeles timezone from Dockerfiles (#33553) Signed-off-by: carlory Co-authored-by: Claude Opus 4.5 --- docker/Dockerfile | 12 +++--------- docker/Dockerfile.nightly_torch | 8 ++------ 2 files changed, 5 insertions(+), 15 deletions(-) diff --git a/docker/Dockerfile b/docker/Dockerfile index 5a1929e5024..72299119735 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -97,9 +97,7 @@ ARG PYTHON_VERSION ENV DEBIAN_FRONTEND=noninteractive # Install system dependencies including build tools -RUN echo 'tzdata tzdata/Areas select America' | debconf-set-selections \ - && echo 'tzdata tzdata/Zones/America select Los_Angeles' | debconf-set-selections \ - && apt-get update -y \ +RUN apt-get update -y \ && apt-get install -y --no-install-recommends \ ccache \ software-properties-common \ @@ -502,9 +500,7 @@ RUN PYTHON_VERSION_STR=$(echo ${PYTHON_VERSION} | sed 's/\.//g') && \ echo "export PYTHON_VERSION_STR=${PYTHON_VERSION_STR}" >> /etc/environment # Install Python and system dependencies -RUN echo 'tzdata tzdata/Areas select America' | debconf-set-selections \ - && echo 'tzdata tzdata/Zones/America select Los_Angeles' | debconf-set-selections \ - && apt-get update -y \ +RUN apt-get update -y \ && apt-get install -y --no-install-recommends \ software-properties-common \ curl \ @@ -713,9 +709,7 @@ ENV UV_INDEX_STRATEGY="unsafe-best-match" # Use copy mode to avoid hardlink failures with Docker cache mounts ENV UV_LINK_MODE=copy -RUN echo 'tzdata tzdata/Areas select America' | debconf-set-selections \ - && echo 'tzdata tzdata/Zones/America select Los_Angeles' | debconf-set-selections \ - && apt-get update -y \ +RUN apt-get update -y \ && apt-get install -y git # We can specify the standard or nightly build of PyTorch diff --git a/docker/Dockerfile.nightly_torch b/docker/Dockerfile.nightly_torch index 87a9144952e..5c17d2a3aea 100644 --- a/docker/Dockerfile.nightly_torch +++ b/docker/Dockerfile.nightly_torch @@ -20,9 +20,7 @@ ARG PYTHON_VERSION=3.12 ARG TARGETPLATFORM ENV DEBIAN_FRONTEND=noninteractive # Install Python and other dependencies -RUN echo 'tzdata tzdata/Areas select America' | debconf-set-selections \ - && echo 'tzdata tzdata/Zones/America select Los_Angeles' | debconf-set-selections \ - && apt-get update -y \ +RUN apt-get update -y \ && apt-get install -y ccache software-properties-common git curl sudo \ && for i in 1 2 3; do \ add-apt-repository -y ppa:deadsnakes/ppa && break || \ @@ -172,9 +170,7 @@ RUN PYTHON_VERSION_STR=$(echo ${PYTHON_VERSION} | sed 's/\.//g') && \ echo "export PYTHON_VERSION_STR=${PYTHON_VERSION_STR}" >> /etc/environment # Install Python and other dependencies -RUN echo 'tzdata tzdata/Areas select America' | debconf-set-selections \ - && echo 'tzdata tzdata/Zones/America select Los_Angeles' | debconf-set-selections \ - && apt-get update -y \ +RUN apt-get update -y \ && apt-get install -y ccache software-properties-common git curl wget sudo vim python3-pip \ && apt-get install -y ffmpeg libsm6 libxext6 libgl1 \ && for i in 1 2 3; do \ From bf001da4bfb53854927b68055a12efd05d494786 Mon Sep 17 00:00:00 2001 From: Chauncey Date: Tue, 3 Feb 2026 14:46:05 +0800 Subject: [PATCH 004/810] [Bugfix] Interleaved thinking keeps compatibility with reasoning_content (#33635) Signed-off-by: chaunceyjiang Co-authored-by: Koushik Dutta --- vllm/entrypoints/chat_utils.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/vllm/entrypoints/chat_utils.py b/vllm/entrypoints/chat_utils.py index 50b664bdba8..c77c18a5887 100644 --- a/vllm/entrypoints/chat_utils.py +++ b/vllm/entrypoints/chat_utils.py @@ -1463,6 +1463,9 @@ def _parse_chat_message_content( # Include reasoning if present for interleaved thinking. if reasoning is not None: result_msg["reasoning"] = cast(str, reasoning) + result_msg["reasoning_content"] = cast( + str, reasoning + ) # keep compatibility elif role == "tool": parsed_msg = _ToolParser(message) if "tool_call_id" in parsed_msg: From e10604480bb8177d563253194667ee9c1590e31a Mon Sep 17 00:00:00 2001 From: Kunshang Ji Date: Tue, 3 Feb 2026 14:46:10 +0800 Subject: [PATCH 005/810] [XPU][1/N] Deprecate ipex and switch to vllm-xpu-kernels for xpu platform (#33379) Signed-off-by: Kunshang Ji --- .../scripts/hardware_ci/run-xpu-test.sh | 3 +- docker/Dockerfile.xpu | 14 +- requirements/xpu.txt | 4 +- vllm/_ipex_ops.py | 445 ++++-------------- vllm/config/model.py | 1 - vllm/model_executor/layers/activation.py | 82 ++-- vllm/model_executor/layers/layernorm.py | 19 +- vllm/model_executor/layers/linear.py | 2 - .../layers/quantization/__init__.py | 4 - .../model_executor/layers/quantization/fp8.py | 33 +- .../model_executor/layers/quantization/inc.py | 26 +- .../layers/quantization/ipex_quant.py | 403 ---------------- .../layers/rotary_embedding/base.py | 7 +- vllm/platforms/__init__.py | 2 - vllm/platforms/xpu.py | 14 +- vllm/v1/attention/backends/fa_utils.py | 5 +- vllm/v1/attention/backends/registry.py | 1 - vllm/v1/worker/xpu_worker.py | 12 +- 18 files changed, 150 insertions(+), 927 deletions(-) delete mode 100644 vllm/model_executor/layers/quantization/ipex_quant.py diff --git a/.buildkite/scripts/hardware_ci/run-xpu-test.sh b/.buildkite/scripts/hardware_ci/run-xpu-test.sh index ab8dfb0360a..36775152f1e 100644 --- a/.buildkite/scripts/hardware_ci/run-xpu-test.sh +++ b/.buildkite/scripts/hardware_ci/run-xpu-test.sh @@ -38,10 +38,9 @@ docker run \ python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 -O3 -cc.cudagraph_mode=NONE python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend ray python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend mp - python3 examples/offline_inference/basic/generate.py --model Intel/Qwen2.5-0.5B-W4A16-G128-AutoRound-LLMC-TEST-ONLY --enforce-eager python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --attention-backend=TRITON_ATTN cd tests - pytest -v -s v1/core + pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py pytest -v -s v1/engine pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py diff --git a/docker/Dockerfile.xpu b/docker/Dockerfile.xpu index f63ce2c5037..04051827ba4 100644 --- a/docker/Dockerfile.xpu +++ b/docker/Dockerfile.xpu @@ -1,8 +1,8 @@ -FROM intel/deep-learning-essentials:2025.2.2-0-devel-ubuntu24.04 AS vllm-base +FROM intel/deep-learning-essentials:2025.3.2-0-devel-ubuntu24.04 AS vllm-base RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \ echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list && \ - add-apt-repository -y ppa:kobuk-team/intel-graphics-staging + add-apt-repository -y ppa:kobuk-team/intel-graphics RUN apt clean && apt-get update -y && \ apt-get install -y --no-install-recommends --fix-missing \ @@ -25,10 +25,13 @@ RUN apt clean && apt-get update -y && \ RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.12 1 RUN update-alternatives --install /usr/bin/python python /usr/bin/python3.12 1 -RUN apt install -y libze1 libze-dev libze-intel-gpu1 intel-opencl-icd libze-intel-gpu-raytracing intel-ocloc +RUN apt update && apt upgrade -y && \ + apt install -y libze1 libze-dev libze-intel-gpu1 intel-opencl-icd libze-intel-gpu-raytracing intel-ocloc && \ + apt install -y intel-oneapi-compiler-dpcpp-cpp-2025.3 + # This oneccl contains the BMG support which is not the case for default version of oneapi 2025.2. -ARG ONECCL_INSTALLER="intel-oneccl-2021.15.7.6_offline.sh" +ARG ONECCL_INSTALLER="intel-oneccl-2021.15.7.8_offline.sh" RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.7/${ONECCL_INSTALLER}" && \ bash "${ONECCL_INSTALLER}" -a --silent --eula accept && \ rm "${ONECCL_INSTALLER}" && \ @@ -85,6 +88,9 @@ RUN python3 -m pip install -e tests/vllm_test_utils ENV NIXL_VERSION=0.7.0 RUN python3 /workspace/vllm/tools/install_nixl_from_source_ubuntu.py +# FIX triton +RUN --mount=type=cache,target=/root/.cache/pip pip uninstall triton triton-xpu -y && pip install triton-xpu==3.6.0 --extra-index-url=https://download.pytorch.org/whl/xpu + # PyJWT-2.7.0 will influence some wheel behaviors, remove its dist-info to avoid conflicts RUN rm /usr/lib/python3/dist-packages/PyJWT-2.7.0.dist-info/ -rf diff --git a/requirements/xpu.txt b/requirements/xpu.txt index c1dc4195b52..6fde5b8f916 100644 --- a/requirements/xpu.txt +++ b/requirements/xpu.txt @@ -11,8 +11,8 @@ jinja2>=3.1.6 datasets # for benchmark scripts numba == 0.61.2 # Required for N-gram speculative decoding --extra-index-url=https://download.pytorch.org/whl/xpu -torch==2.9.0+xpu +torch==2.10.0+xpu torchaudio torchvision -intel-extension-for-pytorch @ https://intel-extension-for-pytorch.s3.us-east-1.amazonaws.com/ipex_dev/xpu/intel_extension_for_pytorch-2.9.10.post0%2Bxpu-cp312-cp312-linux_x86_64.whl +vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.0/vllm_xpu_kernels-0.1.0-cp312-cp312-linux_x86_64.whl \ No newline at end of file diff --git a/vllm/_ipex_ops.py b/vllm/_ipex_ops.py index 239f5376eb4..22133eaef08 100644 --- a/vllm/_ipex_ops.py +++ b/vllm/_ipex_ops.py @@ -1,273 +1,59 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import TYPE_CHECKING import torch +from vllm_xpu_kernels.flash_attn_interface import flash_attn_varlen_func from vllm.logger import init_logger -from vllm.platforms import current_platform logger = init_logger(__name__) -try: - import intel_extension_for_pytorch as ipex -except ImportError as e: - logger.debug("Import error msg: %s", e.msg) +if TYPE_CHECKING: + + def register_fake(fn): + return lambda name: fn +else: + try: + from torch.library import register_fake + except ImportError: + from torch.library import impl_abstract as register_fake + +if hasattr(torch.ops._xpu_C, "fp8_gemm_w8a16"): + + @register_fake("_xpu_C::fp8_gemm_w8a16") + def _fp8_gemm_w8a16_fake( + input: torch.Tensor, + q_weight: torch.Tensor, + weight_scale: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + input_2d = input.view(-1, input.shape[-1]) + M = input_2d.size(0) + N = q_weight.size(1) + return torch.empty((M, N), dtype=input.dtype, device=input.device) + + +if hasattr(torch.ops._xpu_C, "int4_gemm_w4a16"): + + @register_fake("_xpu_C::int4_gemm_w4a16") + def _int4_gemm_w4a16_fake( + input: torch.Tensor, + q_weight: torch.Tensor, + bias: torch.Tensor | None, + weight_scale: torch.Tensor, + qzeros: torch.Tensor, + group_size: int, + group_idx: torch.Tensor | None = None, + ) -> torch.Tensor: + input_2d = input.view(-1, input.shape[-1]) + M = input_2d.size(0) + N = q_weight.size(1) + return torch.empty((M, N), dtype=input.dtype, device=input.device) class ipex_ops: - @staticmethod - def _reshape_activation_tensor( - x: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - num = x.size(0) - d = x.size(1) // 2 - x = x.reshape(num, 2, d) - x1, x2 = torch.chunk(x, chunks=2, dim=1) - x1 = x1.reshape(num, d) - x2 = x2.reshape(num, d) - return x1, x2 - - @staticmethod - def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None: - ipex.llm.functional.silu_and_mul(x, out) - - @staticmethod - def gelu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None: - ipex.llm.functional.gelu_and_mul(x, out) - - @staticmethod - def gelu_tanh_and_mul(out: torch.Tensor, x: torch.Tensor) -> None: - ipex.llm.functional.gelu_and_mul(x, out) - - @staticmethod - def gelu_fast(x: torch.Tensor) -> torch.Tensor: - return torch.nn.functional.gelu(x) - - @staticmethod - def gelu_new(x: torch.Tensor) -> torch.Tensor: - return torch.nn.functional.gelu(x) - - @staticmethod - def gelu_quick(out: torch.Tensor, x: torch.Tensor) -> None: - ipex.llm.functional.gelu_quick(x, out) - - @staticmethod - def paged_attention_v1( - out: torch.Tensor, - query: torch.Tensor, - key_cache: torch.Tensor, - value_cache: torch.Tensor, - num_kv_heads: int, - scale: float, - block_tables: torch.Tensor, - context_lens: torch.Tensor, - block_size: int, - max_context_len: int, - alibi_slopes: torch.Tensor | None, - kv_cache_dtype: str, - k_scale: float, - v_scale: float, - tp_rank: int = 0, - blocksparse_local_blocks: int = 0, - blocksparse_vert_stride: int = 0, - blocksparse_block_size: int = 64, - blocksparse_head_sliding_step: int = 0, - ) -> None: - assert kv_cache_dtype == "auto" - num_heads = out.size(1) - num_queries_per_tokens = num_heads // num_kv_heads - ipex.llm.modules.PagedAttention.single_query_kv_attention( - out, - query.contiguous(), - key_cache.view_as(value_cache), - value_cache, - num_queries_per_tokens, - scale, - block_tables, - context_lens, - block_size, - max_context_len, - alibi_slopes, - ) - - @staticmethod - def paged_attention_v2( - out: torch.Tensor, - exp_sum: torch.Tensor, - max_logits: torch.Tensor, - tmp_out: torch.Tensor, - query: torch.Tensor, - key_cache: torch.Tensor, - value_cache: torch.Tensor, - num_kv_heads: int, - scale: float, - block_tables: torch.Tensor, - context_lens: torch.Tensor, - block_size: int, - max_context_len: int, - alibi_slopes: torch.Tensor | None, - kv_cache_dtype: str, - k_scale: float, - v_scale: float, - tp_rank: int = 0, - blocksparse_local_blocks: int = 0, - blocksparse_vert_stride: int = 0, - blocksparse_block_size: int = 64, - blocksparse_head_sliding_step: int = 0, - ) -> None: - assert kv_cache_dtype == "auto" - num_heads = out.size(1) - num_queries_per_tokens = num_heads // num_kv_heads - ipex.llm.modules.PagedAttention.single_query_kv_attention( - out, - query.contiguous(), - key_cache.view_as(value_cache), - value_cache, - num_queries_per_tokens, - scale, - block_tables, - context_lens, - block_size, - max_context_len, - alibi_slopes, - ) - - @staticmethod - def rotary_embedding( - positions: torch.Tensor, # [batch_size, seq_len] - query: torch.Tensor, # [batch_size, seq_len, num_heads*head_size] - key: torch.Tensor, # [batch_size, seq_len, num_kv_heads*head_size] - head_size: int, - cos_sin_cache: torch.Tensor, # [cos_sin_dim, rot_dim] - is_neox: bool, - ) -> None: - rot_dim = cos_sin_cache.size(1) - ipex.llm.functional.rotary_embedding_batched( - positions, query, key, head_size, cos_sin_cache, is_neox, rot_dim - ) - - @staticmethod - def rms_norm( - input: torch.Tensor, weight: torch.Tensor, epsilon: float - ) -> torch.Tensor: - out = torch.empty_like(input) - torch.ops.torch_ipex.rms_norm_vllm(out, input.contiguous(), weight, epsilon) - return out - - @staticmethod - def fused_add_rms_norm( - input: torch.Tensor, - residual: torch.Tensor, - weight: torch.Tensor, - epsilon: float, - ) -> None: - torch.ops.torch_ipex.fused_add_rms_norm_vllm(input, residual, weight, epsilon) - - @staticmethod - def varlen_attention( - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - out: torch.Tensor, - seqlen_q: torch.Tensor, - seqlen_k: torch.Tensor, - alibi_slopes: torch.Tensor | None, - max_seqlen_q: int, - max_seqlen_k: int, - pdropout: float, - softmax_scale: float, - zero_tensors: bool, - is_causal: bool, - return_softmax: bool, - gen_: torch.Generator, - window_size_left: float, - window_size_right: float, - logits_soft_cap: float, - ) -> None: - if ipex.__version__.endswith("cpu"): - if logits_soft_cap != 0.0: - raise ValueError("IPEX CPU does not support logits_soft_cap") - assert alibi_slopes is None - assert window_size_left < 0 and window_size_right < 0 - ipex.llm.functional.varlen_attention( - query.contiguous(), - key.contiguous(), - value.contiguous(), - out, - seqlen_q.int(), - seqlen_k.int(), - max_seqlen_q, - max_seqlen_k, - pdropout, - softmax_scale, - zero_tensors, - is_causal, - return_softmax, - gen_, - ) - else: # XPU build - ipex.llm.functional.varlen_attention( - query.contiguous(), - key.contiguous(), - value.contiguous(), - out, - seqlen_q.int(), - seqlen_k.int(), - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - pdropout, - softmax_scale, - zero_tensors, - is_causal, - return_softmax, - gen_, - window_size_left, - window_size_right, - logits_soft_cap, - ) - - @staticmethod - def reshape_and_cache( - key: torch.Tensor, - value: torch.Tensor, - key_cache: torch.Tensor, - value_cache: torch.Tensor, - slot_mapping: torch.Tensor, - kv_cache_dtype: str, - k_scale: float, - v_scale: float, - ) -> None: - assert kv_cache_dtype == "auto" - ipex.llm.modules.PagedAttention.reshape_and_cache( - key, value, key_cache, value_cache, slot_mapping - ) - - @staticmethod - def reshape_and_cache_flash( - key: torch.Tensor, - value: torch.Tensor, - key_cache: torch.Tensor, - value_cache: torch.Tensor, - slot_mapping: torch.Tensor, - kv_cache_dtype: str, - k_scale: torch.Tensor | None = None, - v_scale: torch.Tensor | None = None, - k_scale_float: float = 1.0, - v_scale_float: float = 1.0, - ) -> None: - ipex.llm.modules.PagedAttention.reshape_and_cache_flash( - key, - value, - key_cache, - value_cache, - slot_mapping, - kv_cache_dtype, - k_scale_float, - v_scale_float, - ) - @staticmethod def flash_attn_varlen_func( q: torch.Tensor, @@ -295,8 +81,21 @@ class ipex_ops: k_descale=None, v_descale=None, num_splits=0, + return_softmax_lse: bool | None = False, s_aux: torch.Tensor | None = None, ): + assert cu_seqlens_k is not None or seqused_k is not None, ( + "cu_seqlens_k or seqused_k must be provided" + ) + assert cu_seqlens_k is None or seqused_k is None, ( + "cu_seqlens_k and seqused_k cannot be provided at the same time" + ) + assert block_table is None or seqused_k is not None, ( + "when enable block_table, seqused_k is needed" + ) + assert block_table is not None or cu_seqlens_k is not None, ( + "when block_table is disabled, cu_seqlens_k is needed" + ) if out is None: out = torch.empty(q.shape, dtype=q.dtype, device=q.device) real_window_size: tuple[int, int] @@ -304,56 +103,31 @@ class ipex_ops: real_window_size = (-1, -1) else: assert len(window_size) == 2 - real_window_size = (window_size[0], window_size[1]) + real_window_size = (window_size[0], window_size[1]) # noqa: F841 + # In encode attention, v maybe not contiguous and current + # kernel can't handle it if block_table is None: - assert cu_seqlens_k is not None, ( - "cu_seqlens_k can't be None when calling varlen_attention." - ) - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - ipex_ops.varlen_attention( - q.contiguous(), - k.contiguous(), - v.contiguous(), - out, - cu_seqlens_q, - cu_seqlens_k, - None, - max_seqlen_q, - max_seqlen_k, - 0.0, - softmax_scale, - False, - causal, - False, - None, - real_window_size[0], - real_window_size[1], - -1, - ) - return out - else: - return ipex.llm.modules.PagedAttention.flash_attn_varlen_func( - out, - q.contiguous(), - k, - v, - cu_seqlens_q, - seqused_k, - max_seqlen_q, - max_seqlen_k, - softmax_scale, - causal, - block_table, - alibi_slopes, - sink=s_aux, - softcap=softcap, - window_size_left=real_window_size[0], - window_size_right=real_window_size[1], - k_scale=1.0, - v_scale=1.0, - ) + v = v.contiguous() + return flash_attn_varlen_func( + out=out, + q=q.contiguous(), + k=k, + v=v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + seqused_k=seqused_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + softmax_scale=softmax_scale, + causal=causal, + block_table=block_table, + s_aux=s_aux, + window_size=real_window_size, + # alibi_slopes = alibi_slopes, + # softcap=softcap, + return_softmax_lse=return_softmax_lse, + ) @staticmethod def get_scheduler_metadata( @@ -382,64 +156,3 @@ class ipex_ops: "get_scheduler_metadata is not implemented for ipex_ops, returning None." ) return None - - @staticmethod - def swap_blocks( - src: torch.Tensor, dst: torch.Tensor, block_mapping: torch.Tensor - ) -> None: - torch.xpu.swap_blocks(src, dst, block_mapping) # type: ignore - - @staticmethod - def scaled_fp8_quant( - input: torch.Tensor, - scale: torch.Tensor | None = None, - num_token_padding: int | None = None, - scale_ub: torch.Tensor | None = None, - use_per_token_if_dynamic: bool = False, - output: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Quantize input tensor to FP8 and return quantized tensor and scale. - - This function is designed for both static and dynamic quantization: - If you provide the scale, it will use static scaling and if you omit - it, the scale will be determined dynamically. Currently, XPU platform - only supports dynamic quantization. The function also allows optional - padding of the output tensors for downstream kernels that will benefit - from padding. - - Args: - input: The input tensor to be quantized to FP8 - scale: Optional scaling factor for the FP8 quantization - scale_ub: Optional upper bound for scaling factor in dynamic - per token case - num_token_padding: If specified, pad the first dimension - of the output to at least this value. - use_per_token_if_dynamic: Whether to do per_tensor or per_token - in the dynamic quantization case. - - Returns: - tuple[torch.Tensor, torch.Tensor]: The output tensor in FP8 and - scaling factor. - """ - # This code assumes batch_dim and num_tokens are flattened - assert input.ndim == 2 - shape: tuple[int, int] | torch.Size = input.shape - out_dtype: torch.dtype = current_platform.fp8_dtype() - if num_token_padding: - shape = (max(num_token_padding, input.shape[0]), shape[1]) - if output is None: - output = torch.empty(shape, device=input.device, dtype=out_dtype) - else: - assert num_token_padding is None, ( - "padding not supported if output passed in" - ) - assert output.dtype == out_dtype - assert scale is None, "only dynamic fp8 quantization supported on XPU" - assert not use_per_token_if_dynamic, ( - "per token dynamic fp8 quantization not supported on XPU" - ) - scale = torch.zeros(1, device=input.device, dtype=torch.float32) - torch.ops.torch_ipex.dynamic_scaled_fp8_quant(output, input, scale) - - return output, scale diff --git a/vllm/config/model.py b/vllm/config/model.py index 563f8ac56e0..48ff44ac9fd 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -877,7 +877,6 @@ class ModelConfig: overrides = [ "gptq_marlin", "awq_marlin", - "ipex", "inc", "moe_wna16", "modelopt", diff --git a/vllm/model_executor/layers/activation.py b/vllm/model_executor/layers/activation.py index b53a37a3176..3e00d21d5a1 100644 --- a/vllm/model_executor/layers/activation.py +++ b/vllm/model_executor/layers/activation.py @@ -129,12 +129,8 @@ class SiluAndMul(CustomOp): def __init__(self, *, compile_native: bool = True): super().__init__(compile_native=compile_native) - if current_platform.is_cuda_alike(): + if current_platform.is_cuda_alike() or current_platform.is_xpu(): self.op = torch.ops._C.silu_and_mul - elif current_platform.is_xpu(): - from vllm._ipex_ops import ipex_ops - - self.op = ipex_ops.silu_and_mul elif current_platform.is_cpu(): self._forward_method = self.forward_native @@ -152,11 +148,7 @@ class SiluAndMul(CustomOp): return out def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: - d = x.shape[-1] // 2 - output_shape = x.shape[:-1] + (d,) - out = torch.empty(output_shape, dtype=x.dtype, device=x.device) - self.op(out, x) - return out + return self.forward_cuda(x) # --8<-- [start:mul_and_silu] @@ -175,12 +167,8 @@ class MulAndSilu(CustomOp): def __init__(self): super().__init__() - if current_platform.is_cuda_alike(): + if current_platform.is_cuda_alike() or current_platform.is_xpu(): self.op = torch.ops._C.mul_and_silu - elif current_platform.is_xpu(): - from vllm._ipex_ops import ipex_ops - - self.op = ipex_ops.silu_and_mul elif current_platform.is_cpu(): self._forward_method = self.forward_native @@ -196,8 +184,8 @@ class MulAndSilu(CustomOp): self.op(out, x) return out - # TODO implement forward_xpu for MulAndSilu - # def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: + def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: + return self.forward_cuda(x) # --8<-- [start:gelu_and_mul_sparse] @@ -278,7 +266,11 @@ class GeluAndMul(CustomOp): self.approximate = approximate if approximate not in ("none", "tanh"): raise ValueError(f"Unknown approximate mode: {approximate}") - if current_platform.is_cuda_alike() or current_platform.is_cpu(): + if ( + current_platform.is_cuda_alike() + or current_platform.is_cpu() + or current_platform.is_xpu() + ): if approximate == "none": self.op = torch.ops._C.gelu_and_mul elif approximate == "tanh": @@ -289,13 +281,6 @@ class GeluAndMul(CustomOp): "with torch.compile. For native implementation, fallback to 'none' " "approximation. The custom kernel implementation is unaffected." ) - elif current_platform.is_xpu(): - from vllm._ipex_ops import ipex_ops - - if approximate == "none": - self.op = ipex_ops.gelu_and_mul - else: - self.op = ipex_ops.gelu_tanh_and_mul def forward_native(self, x: torch.Tensor) -> torch.Tensor: """PyTorch-native implementation equivalent to forward().""" @@ -314,11 +299,7 @@ class GeluAndMul(CustomOp): return out def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: - d = x.shape[-1] // 2 - output_shape = x.shape[:-1] + (d,) - out = torch.empty(output_shape, dtype=x.dtype, device=x.device) - self.op(out, x) - return out + return self.forward_cuda(x) def extra_repr(self) -> str: return f"approximate={repr(self.approximate)}" @@ -401,12 +382,12 @@ class NewGELU(CustomOp): def __init__(self): super().__init__() - if current_platform.is_cuda_alike() or current_platform.is_cpu(): + if ( + current_platform.is_cuda_alike() + or current_platform.is_cpu() + or current_platform.is_xpu() + ): self.op = torch.ops._C.gelu_new - elif current_platform.is_xpu(): - from vllm._ipex_ops import ipex_ops - - self.op = ipex_ops.gelu_new def forward_native(self, x: torch.Tensor) -> torch.Tensor: """PyTorch-native implementation equivalent to forward().""" @@ -419,7 +400,7 @@ class NewGELU(CustomOp): return out def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: - return self.op(x) + return self.forward_cuda(x) # --8<-- [start:gelu_fast] @@ -429,12 +410,12 @@ class FastGELU(CustomOp): def __init__(self): super().__init__() - if current_platform.is_cuda_alike() or current_platform.is_cpu(): + if ( + current_platform.is_cuda_alike() + or current_platform.is_cpu() + or current_platform.is_xpu() + ): self.op = torch.ops._C.gelu_fast - elif current_platform.is_xpu(): - from vllm._ipex_ops import ipex_ops - - self.op = ipex_ops.gelu_fast def forward_native(self, x: torch.Tensor) -> torch.Tensor: """PyTorch-native implementation equivalent to forward().""" @@ -446,7 +427,7 @@ class FastGELU(CustomOp): return out def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: - return self.op(x) + return self.forward_cuda(x) # --8<-- [start:quick_gelu] @@ -457,12 +438,12 @@ class QuickGELU(CustomOp): def __init__(self): super().__init__() - if current_platform.is_cuda_alike() or current_platform.is_cpu(): + if ( + current_platform.is_cuda_alike() + or current_platform.is_cpu() + or current_platform.is_xpu() + ): self.op = torch.ops._C.gelu_quick - elif current_platform.is_xpu(): - from vllm._ipex_ops import ipex_ops - - self.op = ipex_ops.gelu_quick def forward_native(self, x: torch.Tensor) -> torch.Tensor: """PyTorch-native implementation equivalent to forward().""" @@ -474,12 +455,7 @@ class QuickGELU(CustomOp): return out def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: - out = torch.empty_like(x) - self.op(out, x) - return out - - # TODO implement forward_xpu for QuickGELU - # def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: + return self.forward_cuda(x) # --8<-- [start:relu2] diff --git a/vllm/model_executor/layers/layernorm.py b/vllm/model_executor/layers/layernorm.py index 2db8ce2bdf0..3b669c55965 100644 --- a/vllm/model_executor/layers/layernorm.py +++ b/vllm/model_executor/layers/layernorm.py @@ -231,24 +231,7 @@ class RMSNorm(CustomOp): x: torch.Tensor, residual: torch.Tensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - if self.variance_size_override is not None: - return self.forward_native(x, residual) - - from vllm._ipex_ops import ipex_ops as ops - - if residual is not None: - ops.fused_add_rms_norm( - x, - residual, - self.weight.data, - self.variance_epsilon, - ) - return x, residual - return ops.rms_norm( - x, - self.weight.data, - self.variance_epsilon, - ) + return self.forward_cuda(x, residual) def extra_repr(self) -> str: s = f"hidden_size={self.weight.data.size(0)}" diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index 61d86cea48c..bbd7267fdf7 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -60,8 +60,6 @@ WEIGHT_LOADER_V2_SUPPORTED = [ "ModelOptFp8LinearMethod", "ModelOptFp8PcPtLinearMethod", "ModelOptFp8PbWoLinearMethod", - "IPEXAWQLinearMethod", - "IPEXGPTQLinearMethod", "QuarkLinearMethod", "ModelOptNvFp4LinearMethod", "PetitNvFp4LinearMethod", diff --git a/vllm/model_executor/layers/quantization/__init__.py b/vllm/model_executor/layers/quantization/__init__.py index cc0fdfa8e21..82de32af347 100644 --- a/vllm/model_executor/layers/quantization/__init__.py +++ b/vllm/model_executor/layers/quantization/__init__.py @@ -24,7 +24,6 @@ QuantizationMethods = Literal[ "compressed-tensors", "bitsandbytes", "experts_int8", - "ipex", "quark", "moe_wna16", "torchao", @@ -41,7 +40,6 @@ DEPRECATED_QUANTIZATION_METHODS = [ "fbgemm_fp8", "fp_quant", "experts_int8", - "ipex", "petit_nvfp4", ] @@ -121,7 +119,6 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: from .gptq import GPTQConfig from .gptq_marlin import GPTQMarlinConfig from .inc import INCConfig - from .ipex_quant import IPEXConfig from .modelopt import ModelOptFp8Config, ModelOptNvFp4Config from .moe_wna16 import MoeWNA16Config from .mxfp4 import Mxfp4Config @@ -144,7 +141,6 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: "bitsandbytes": BitsAndBytesConfig, "ptpc_fp8": PTPCFp8Config, "experts_int8": ExpertsInt8Config, - "ipex": IPEXConfig, "quark": QuarkConfig, "moe_wna16": MoeWNA16Config, "torchao": TorchAOConfig, diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index a8467b5f07f..9b7d654335d 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -184,39 +184,10 @@ class Fp8Config(QuantizationConfig): def get_xpu_quant_method( self, layer: torch.nn.Module, prefix: str ) -> "QuantizeMethodBase | None": - from vllm.model_executor.layers.quantization.ipex_quant import ( - XPUFp8LinearMethod, - XPUFp8MoEMethod, + raise NotImplementedError( + "FP8 quantization is not supported during xpu kernel migration." ) - fp8_config = Fp8Config( - is_checkpoint_fp8_serialized=self.is_checkpoint_fp8_serialized, - activation_scheme=self.activation_scheme, - ignored_layers=self.ignored_layers, - weight_block_size=self.weight_block_size, - ) - - if isinstance(layer, LinearBase): - if is_layer_skipped( - prefix=prefix, - ignored_layers=self.ignored_layers, - fused_mapping=self.packed_modules_mapping, - ): - return UnquantizedLinearMethod() - return XPUFp8LinearMethod(fp8_config) - elif isinstance(layer, FusedMoE): - if is_layer_skipped( - prefix=prefix, - ignored_layers=self.ignored_layers, - fused_mapping=self.packed_modules_mapping, - ): - return UnquantizedFusedMoEMethod(layer.moe_config) - - return XPUFp8MoEMethod(fp8_config, layer) - elif isinstance(layer, Attention): - return Fp8KVCacheMethod(self) - return None - def get_quant_method( self, layer: torch.nn.Module, prefix: str ) -> "QuantizeMethodBase | None": diff --git a/vllm/model_executor/layers/quantization/inc.py b/vllm/model_executor/layers/quantization/inc.py index f68fd957867..359f24688ce 100644 --- a/vllm/model_executor/layers/quantization/inc.py +++ b/vllm/model_executor/layers/quantization/inc.py @@ -38,7 +38,6 @@ class INCConfig(QuantizationConfig): "awq", "awq:marlin", "marlin", - "ipex", } def __init__( @@ -410,31 +409,10 @@ class INCConfig(QuantizationConfig): return UnquantizedLinearMethod() else: return None - from vllm.model_executor.layers.quantization.ipex_quant import ( - IPEXAWQLinearMethod, - IPEXConfig, - IPEXGPTQLinearMethod, + raise NotImplementedError( + "INC quantization is not supported during xpu kernel migration." ) - if isinstance(layer, (LinearBase, ParallelLMHead)): - if "awq" in self.packing_format: - config = IPEXConfig( - method="awq", weight_bits=weight_bits, group_size=group_size - ) - return IPEXAWQLinearMethod(config) - elif "gptq" in self.packing_format: - config = IPEXConfig( - method="gptq", weight_bits=weight_bits, group_size=group_size - ) - return IPEXGPTQLinearMethod(config) - else: - raise ValueError( - f"ipex backend only supports awq " - f"and gptq format,but got {self.packing_format}" - ) - else: - return None - def get_quant_method(self, layer: torch.nn.Module, prefix: str): if prefix and self.extra_config: for layer_name in self.extra_config: diff --git a/vllm/model_executor/layers/quantization/ipex_quant.py b/vllm/model_executor/layers/quantization/ipex_quant.py deleted file mode 100644 index f957b3991ee..00000000000 --- a/vllm/model_executor/layers/quantization/ipex_quant.py +++ /dev/null @@ -1,403 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from typing import Any - -import torch -from packaging import version -from torch.nn import Module - -from vllm._ipex_ops import ipex_ops as ops -from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig -from vllm.model_executor.layers.linear import ( - LinearBase, - LinearMethodBase, - UnquantizedLinearMethod, -) -from vllm.model_executor.layers.quantization import ( - QuantizationConfig, - QuantizationMethods, -) -from vllm.model_executor.layers.quantization.awq import AWQLinearMethod -from vllm.model_executor.layers.quantization.fp8 import ( - Fp8Config, - Fp8LinearMethod, - Fp8OnlineMoEMethod, -) -from vllm.model_executor.layers.quantization.gptq import GPTQLinearMethod -from vllm.model_executor.layers.quantization.utils.quant_utils import is_layer_skipped -from vllm.model_executor.utils import replace_parameter -from vllm.platforms import current_platform - -MIN_IPEX_VERSION = "2.6.0" - - -class IPEXConfig(QuantizationConfig): - """INT8 quantization config class using IPEX for the CPU/XPU backend, - including AWQ, GPTQ. - """ - - IPEX_QUANT_METHOD_MAP = { - "awq": 1, - "gptq": 0, - } - - def __init__( - self, - method: str, - weight_bits: int, - group_size: int, - modules_to_not_convert: list[str] | None = None, - desc_act: bool | None = None, - lm_head_quantized: bool | None = None, - is_sym: bool | None = None, - ) -> None: - super().__init__() - self.method = method - self.weight_bits = weight_bits - self.group_size = group_size - self.modules_to_not_convert = modules_to_not_convert or [] - self.desc_act = desc_act - self.lm_head_quantized = lm_head_quantized - self.is_sym = is_sym - self.pack_factor = 32 // self.weight_bits - - if self.weight_bits not in [4]: - raise ValueError( - f"IPEX quantization supports weight bits [4], " - f"but got {self.weight_bits}." - ) - - if self.method not in ["awq", "gptq"]: - raise ValueError( - f"IPEX quantization supports [awq, gptq], but got {self.method}." - ) - - def __repr__(self) -> str: - return ( - f"IPEXConfig(method={self.method}," - f"weight_bits={self.weight_bits}, " - f"group_size={self.group_size})" - ) - - @classmethod - def get_name(cls) -> QuantizationMethods: - return "ipex" - - @classmethod - def get_supported_act_dtypes(cls) -> list[torch.dtype]: - return [torch.bfloat16, torch.float16] - - @classmethod - def get_min_capability(cls) -> int: - return -1 - - @staticmethod - def get_config_filenames() -> list[str]: - return [ - "quant_config.json", - "quantize_config.json", - ] - - @classmethod - def from_config(cls, config: dict[str, Any]) -> "IPEXConfig": - method = cls.get_from_keys(config, ["quant_method"]).lower() - if method == "awq": - weight_bits = cls.get_from_keys(config, ["w_bit", "bits"]) - group_size = cls.get_from_keys(config, ["q_group_size", "group_size"]) - modules_to_not_convert = cls.get_from_keys_or( - config, ["modules_to_not_convert"], None - ) - is_sym = not cls.get_from_keys_or(config, ["zero_point"], default=False) - return cls( - method, - weight_bits, - group_size, - modules_to_not_convert, - False, - False, - is_sym, - ) - # otherwise for gptq - weight_bits = cls.get_from_keys(config, ["bits"]) - group_size = cls.get_from_keys(config, ["group_size"]) - lm_head_quantized = cls.get_from_keys_or(config, ["lm_head"], default=False) - desc_act = cls.get_from_keys_or(config, ["desc_act"], default=False) - is_sym = cls.get_from_keys_or(config, ["sym"], default=True) - return cls( - method, weight_bits, group_size, [], desc_act, lm_head_quantized, is_sym - ) - - @classmethod - def override_quantization_method( - cls, hf_quant_cfg, user_quant - ) -> QuantizationMethods | None: - if not current_platform.is_xpu(): - return None - - quant_method = hf_quant_cfg.get("quant_method", "").lower() - - if quant_method in ["awq", "gptq"]: - return cls.get_name() - - return None - - def get_quant_method( - self, layer: torch.nn.Module, prefix: str - ) -> "LinearMethodBase | None": - if isinstance(layer, LinearBase): - if self.method == "awq": - if is_layer_skipped( - prefix, - self.modules_to_not_convert, - self.packed_modules_mapping, - skip_with_substr=True, - ): - return UnquantizedLinearMethod() - return IPEXAWQLinearMethod(self) - if self.method == "gptq": - return IPEXGPTQLinearMethod(self) - return None - - -class IPEXGPTQLinearMethod(GPTQLinearMethod): - """GPTQ linear method using IPEX for the CPU/XPU backend.""" - - def __init__(self, quant_config: IPEXConfig): - self.quant_config = quant_config # type: ignore - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - bias = layer.bias if not layer.skip_bias_add else None - - try: - import intel_extension_for_pytorch as ipex - - if version.parse(ipex.__version__) < version.parse(MIN_IPEX_VERSION): - raise ImportError( - "intel_extension_for_pytorch version is " - "wrong. Please install " - f"intel_extension_for_pytorch>={MIN_IPEX_VERSION}." - ) - except ImportError as err: - raise ImportError( - "Please install " - f"intel_extension_for_pytorch>={MIN_IPEX_VERSION} via " - f"`pip install intel_extension_for_pytorch>={MIN_IPEX_VERSION}`" - " to use IPEX-AWQ linear method." - ) from err - # Using the compute dtype (lowp_mode) as INT8 to leverage instructions - # with better performance. - lowp_mode = ipex.quantization.WoqLowpMode.INT8 - # The weight will be de-packed from INT4 to INT8. - weight_dtype = ipex.quantization.WoqWeightDtype.INT4 - # The float activation will be quantized (dynamic, per-token) to INT8. - act_quant_mode = ipex.quantization.WoqActQuantMode.PER_BATCH_IC_BLOCK - - assert isinstance(self.quant_config, IPEXConfig) - qconfig = ipex.quantization.get_weight_only_quant_qconfig_mapping( - weight_dtype=weight_dtype, - lowp_mode=lowp_mode, - act_quant_mode=act_quant_mode, - group_size=self.quant_config.group_size, - ) - layer.ipex_output_size = layer.qweight.shape[-1] - g_idx = layer.g_idx if self.quant_config.desc_act else None - layer.ipex_qlinear = ( - ipex.llm.quantization.woq_linear.IPEXWeightOnlyQuantizedLinear.from_weight( - layer.qweight, - layer.scales, - layer.qzeros, - layer.qweight.size(0), - layer.ipex_output_size, - qconfig=qconfig, - g_idx=g_idx, - bias=bias, - group_size=self.quant_config.group_size, - quant_method=IPEXConfig.IPEX_QUANT_METHOD_MAP["gptq"], - weight_qscheme="sym" if self.quant_config.is_sym else "asym", - ) - ) - - def apply( - self, - layer: torch.nn.Module, - x: torch.Tensor, - bias: torch.Tensor | None = None, - ) -> torch.Tensor: - reshaped_x = x.reshape(-1, x.shape[-1]) - out = layer.ipex_qlinear(reshaped_x) - return out.reshape(x.shape[:-1] + (layer.ipex_output_size,)) - - -class IPEXAWQLinearMethod(AWQLinearMethod): - """AWQ linear method using IPEX for the CPU/XPU backend.""" - - def __init__(self, quant_config: IPEXConfig): - self.quant_config = quant_config # type: ignore - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - super().process_weights_after_loading(layer=layer) - - bias = layer.bias if not layer.skip_bias_add else None - - try: - import intel_extension_for_pytorch as ipex - - if version.parse(ipex.__version__) < version.parse(MIN_IPEX_VERSION): - raise ImportError( - "intel_extension_for_pytorch version is " - "wrong. Please install " - f"intel_extension_for_pytorch>={MIN_IPEX_VERSION}." - ) - except ImportError as err: - raise ImportError( - "Please install " - f"intel_extension_for_pytorch>={MIN_IPEX_VERSION} via " - f"`pip install intel_extension_for_pytorch>={MIN_IPEX_VERSION}`" - " to use IPEX-AWQ linear method." - ) from err - - # Using the compute dtype (lowp_mode) as INT8 to leverage instructions - # with better performance. - lowp_mode = ipex.quantization.WoqLowpMode.INT8 - # The weight will be de-packed from INT4 to INT8. - weight_dtype = ipex.quantization.WoqWeightDtype.INT4 - # The float activation will be quantized (dynamic, per-token) to INT8. - act_quant_mode = ipex.quantization.WoqActQuantMode.PER_BATCH - - assert isinstance(self.quant_config, IPEXConfig) - qconfig = ipex.quantization.get_weight_only_quant_qconfig_mapping( - weight_dtype=weight_dtype, - lowp_mode=lowp_mode, - act_quant_mode=act_quant_mode, - group_size=self.quant_config.group_size, - ) - - layer.ipex_output_size = layer.qweight.size(1) * self.quant_config.pack_factor - layer.ipex_qlinear = ( - ipex.llm.quantization.woq_linear.IPEXWeightOnlyQuantizedLinear.from_weight( - layer.qweight, - layer.scales, - layer.qzeros, - layer.qweight.size(0), - layer.ipex_output_size, - qconfig=qconfig, - bias=bias, - group_size=self.quant_config.group_size, - quant_method=IPEXConfig.IPEX_QUANT_METHOD_MAP["awq"], # type: ignore - weight_qscheme="sym" if self.quant_config.is_sym else "asym", - ) - ) - - def apply( - self, - layer: torch.nn.Module, - x: torch.Tensor, - bias: torch.Tensor | None = None, - ) -> torch.Tensor: - reshaped_x = x.reshape(-1, x.shape[-1]) - out = layer.ipex_qlinear(reshaped_x) - return out.reshape(x.shape[:-1] + (layer.ipex_output_size,)) - - -class XPUFp8LinearMethod(Fp8LinearMethod): - def __init__(self, quant_config: Fp8Config): - super().__init__(quant_config) - - def process_weights_after_loading(self, layer: Module) -> None: - if getattr(layer, "_already_called_process_weights_after_loading", False): - return - # If checkpoint not serialized fp8, quantize the weights. - if not self.quant_config.is_checkpoint_fp8_serialized: - qweight, weight_scale = ops.scaled_fp8_quant(layer.weight, scale=None) - # Update the layer with the new values. - replace_parameter(layer, "weight", qweight.data) - replace_parameter(layer, "weight_scale", weight_scale.data) - layer.input_scale = None - - def apply( - self, - layer: torch.nn.Module, - x: torch.Tensor, - bias: torch.Tensor | None = None, - ) -> torch.Tensor: - weight = layer.weight.data - weight_scale = layer.weight_scale.data - output = torch.ops.torch_ipex.fp8_gemm_w8a16( - x, weight, True, weight_scale, bias - ) - return output - - -class XPUFp8MoEMethod(Fp8OnlineMoEMethod): - def __init__(self, quant_config: Fp8Config, layer: torch.nn.Module): - super().__init__(quant_config, layer) - self.quant_config = quant_config - - def process_weights_after_loading(self, layer: Module) -> None: - if getattr(layer, "_already_called_process_weights_after_loading", False): - return - if not self.quant_config.is_checkpoint_fp8_serialized: - fp8_dtype = current_platform.fp8_dtype() - w13_weight = torch.empty_like(layer.w13_weight.data, dtype=fp8_dtype) - w2_weight = torch.empty_like(layer.w2_weight.data, dtype=fp8_dtype) - - # Re-initialize w13_scale because we directly quantize - # merged w13 weights and generate a single scaling factor. - layer.w13_weight_scale = torch.nn.Parameter( - torch.ones( - layer.local_num_experts, - dtype=torch.float32, - device=w13_weight.device, - ), - requires_grad=False, - ) - for expert in range(layer.local_num_experts): - w13_weight[expert, :, :], layer.w13_weight_scale[expert] = ( - ops.scaled_fp8_quant(layer.w13_weight.data[expert, :, :]) - ) - w2_weight[expert, :, :], layer.w2_weight_scale[expert] = ( - ops.scaled_fp8_quant(layer.w2_weight.data[expert, :, :]) - ) - replace_parameter(layer, "w13_weight", w13_weight) - replace_parameter(layer, "w2_weight", w2_weight) - - import intel_extension_for_pytorch as ipex - - ep_rank_start = self.moe.ep_rank * self.moe.num_local_experts - layer.ipex_fusion = ipex.llm.modules.GatedMLPMOE( - layer.w13_weight, - layer.w2_weight, - w1_scale_inv=layer.w13_weight_scale, - w2_scale_inv=layer.w2_weight_scale, - a1_scale_inv=layer.w13_input_scale, - a2_scale_inv=layer.w2_input_scale, - use_prepack=True, - experts_start_id=ep_rank_start, - ) - - def get_fused_moe_quant_config( - self, layer: torch.nn.Module - ) -> FusedMoEQuantConfig | None: - return None - - @property - def is_monolithic(self) -> bool: - return True - - def apply_monolithic( - self, - layer: torch.nn.Module, - x: torch.Tensor, - router_logits: torch.Tensor, - ) -> torch.Tensor: - return layer.ipex_fusion( - x, - layer.use_grouped_topk, - layer.top_k, - router_logits, - layer.renormalize, - layer.topk_group, - layer.num_expert_group, - custom_routing_function=layer.custom_routing_function, - ) diff --git a/vllm/model_executor/layers/rotary_embedding/base.py b/vllm/model_executor/layers/rotary_embedding/base.py index d63367af5fe..ffc6f67daef 100644 --- a/vllm/model_executor/layers/rotary_embedding/base.py +++ b/vllm/model_executor/layers/rotary_embedding/base.py @@ -232,17 +232,14 @@ class RotaryEmbedding(RotaryEmbeddingBase): query: torch.Tensor, key: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor | None]: - from vllm._ipex_ops import ipex_ops as ops - self._match_cos_sin_cache_dtype(query) # ops.rotary_embedding() is an in-place operation # that updates the query and key tensors. if key is None: - # XPU kernel doesn't support key=None so fall back to native impl - # TODO(sarckk): add support for optional key in - # ipex.llm.functional.rotary_embedding_batched return self.forward_native(positions, query, key) else: + from vllm import _custom_ops as ops + ops.rotary_embedding( positions, query, diff --git a/vllm/platforms/__init__.py b/vllm/platforms/__init__.py index a0e5af1aba5..758409ae1a3 100644 --- a/vllm/platforms/__init__.py +++ b/vllm/platforms/__init__.py @@ -132,8 +132,6 @@ def xpu_platform_plugin() -> str | None: is_xpu = False logger.debug("Checking if XPU platform is available.") try: - # installed IPEX if the machine has XPUs. - import intel_extension_for_pytorch # noqa: F401 import torch if supports_xccl(): diff --git a/vllm/platforms/xpu.py b/vllm/platforms/xpu.py index 439d21cb88e..6e299f30ee6 100644 --- a/vllm/platforms/xpu.py +++ b/vllm/platforms/xpu.py @@ -7,6 +7,11 @@ from typing import TYPE_CHECKING import torch +# import custom ops, trigger op registration +import vllm_xpu_kernels._C # noqa +import vllm_xpu_kernels._moe_C # noqa +import vllm_xpu_kernels._xpu_C # noqa + from vllm.logger import init_logger from vllm.v1.attention.backends.registry import AttentionBackendEnum @@ -55,6 +60,9 @@ class XPUPlatform(Platform): dtype = attn_selector_config.dtype if attn_selector_config.use_sparse: raise NotImplementedError("Sparse Attention is not supported on XPU.") + if attn_selector_config.use_mla: + logger.info_once("Using Triton MLA backend on V1 engine.") + return AttentionBackendEnum.TRITON_MLA.get_path() if selected_backend == AttentionBackendEnum.TRITON_ATTN: logger.info_once("Using Triton backend.") return AttentionBackendEnum.TRITON_ATTN.get_path() @@ -78,9 +86,9 @@ class XPUPlatform(Platform): @classmethod def get_supported_vit_attn_backends(cls) -> list["AttentionBackendEnum"]: - # XPU only supports FLASH_ATTN for vision attention. return [ AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA, ] @classmethod @@ -145,7 +153,7 @@ class XPUPlatform(Platform): def check_and_update_config(cls, vllm_config: VllmConfig) -> None: cache_config = vllm_config.cache_config model_config = vllm_config.model_config - # in V1(or with ipex chunked prefill) block_size is 64 + # in V1(or with chunked prefill) block_size is 64 if cache_config and cache_config.block_size is None: cache_config.block_size = 64 @@ -206,7 +214,7 @@ class XPUPlatform(Platform): @classmethod def fp8_dtype(cls) -> torch.dtype: - return torch.float8_e5m2 + return torch.float8_e4m3fn @classmethod def is_data_center_gpu(cls) -> bool: diff --git a/vllm/v1/attention/backends/fa_utils.py b/vllm/v1/attention/backends/fa_utils.py index 988cf7c2711..281d188557f 100644 --- a/vllm/v1/attention/backends/fa_utils.py +++ b/vllm/v1/attention/backends/fa_utils.py @@ -16,12 +16,13 @@ if current_platform.is_cuda(): ) elif current_platform.is_xpu(): + from vllm import _custom_ops as ops + + reshape_and_cache_flash = ops.reshape_and_cache_flash from vllm._ipex_ops import ipex_ops - reshape_and_cache_flash = ipex_ops.reshape_and_cache_flash flash_attn_varlen_func = ipex_ops.flash_attn_varlen_func # type: ignore[assignment] get_scheduler_metadata = ipex_ops.get_scheduler_metadata # type: ignore[assignment] - elif current_platform.is_rocm(): try: from flash_attn import flash_attn_varlen_func # type: ignore[no-redef] diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index bd45702fa58..2a80bbd94a1 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -69,7 +69,6 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): "vllm.v1.attention.backends.mla.flashmla_sparse.FlashMLASparseBackend" ) FLASH_ATTN_MLA = "vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend" - IPEX = "vllm.v1.attention.backends.ipex.IpexAttentionBackend" NO_ATTENTION = "vllm.v1.attention.backends.no_attention.NoAttentionBackend" FLEX_ATTENTION = "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend" TREE_ATTN = "vllm.v1.attention.backends.tree_attn.TreeAttentionBackend" diff --git a/vllm/v1/worker/xpu_worker.py b/vllm/v1/worker/xpu_worker.py index f1bdd5da3be..6e45a107ca1 100644 --- a/vllm/v1/worker/xpu_worker.py +++ b/vllm/v1/worker/xpu_worker.py @@ -1,10 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import gc import os from typing import Any import torch -import torch.distributed from vllm.config import VllmConfig from vllm.logger import init_logger @@ -85,7 +85,14 @@ class XPUWorker(Worker): current_platform.dist_backend, ) + # Set random seed. + set_random_seed(self.model_config.seed) + + # Now take memory snapshot after NCCL is initialized + gc.collect() torch.xpu.empty_cache() + + # take current memory snapshot self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device) self.requested_memory = request_memory(init_snapshot, self.cache_config) logger.debug("worker init memory snapshot: %r", self.init_snapshot) @@ -93,9 +100,6 @@ class XPUWorker(Worker): "worker requested memory: %sGiB", format_gib(self.requested_memory) ) - # Set random seed. - set_random_seed(self.model_config.seed) - # Initialize workspace manager num_ubatches = 2 if self.vllm_config.parallel_config.enable_dbo else 1 init_workspace_manager(self.device, num_ubatches) From ef248ff740200c91791ba952b3458a5d5a016d26 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9D=A8=E6=9C=B1=20=C2=B7=20Kiki?= Date: Tue, 3 Feb 2026 14:58:44 +0800 Subject: [PATCH 006/810] [Misc] Remove deprecated profiler environment variables (#33536) Signed-off-by: carlory Co-authored-by: Claude Opus 4.5 --- vllm/config/profiler.py | 81 ----------------------------------------- vllm/envs.py | 14 ------- 2 files changed, 95 deletions(-) diff --git a/vllm/config/profiler.py b/vllm/config/profiler.py index e3cedf33e3c..425f3fb6bcd 100644 --- a/vllm/config/profiler.py +++ b/vllm/config/profiler.py @@ -8,7 +8,6 @@ from pydantic import Field, model_validator from pydantic.dataclasses import dataclass from typing_extensions import Self -import vllm.envs as envs from vllm.config.utils import config from vllm.logger import init_logger from vllm.utils.hashing import safe_hash @@ -102,88 +101,8 @@ class ProfilerConfig: hash_str = safe_hash(str(factors).encode(), usedforsecurity=False).hexdigest() return hash_str - def _get_from_env_if_set(self, field_name: str, env_var_name: str) -> None: - """Get field from env var if set, with deprecation warning.""" - - if envs.is_set(env_var_name): - value = getattr(envs, env_var_name) - logger.warning_once( - "Using %s environment variable is deprecated and will be removed in " - "v0.15.0 or v1.0.0, whichever is soonest. Please use " - "--profiler-config.%s command line argument or " - "ProfilerConfig(%s=...) config field instead.", - env_var_name, - field_name, - field_name, - ) - return value - return None - - def _set_from_env_if_set( - self, - field_name: str, - env_var_name: str, - to_bool: bool = True, - to_int: bool = False, - ) -> None: - """Set field from env var if set, with deprecation warning.""" - value = self._get_from_env_if_set(field_name, env_var_name) - if value is not None: - if to_bool: - value = value == "1" - if to_int: - value = int(value) - setattr(self, field_name, value) - @model_validator(mode="after") def _validate_profiler_config(self) -> Self: - maybe_use_cuda_profiler = self._get_from_env_if_set( - "profiler", "VLLM_TORCH_CUDA_PROFILE" - ) - if maybe_use_cuda_profiler is not None: - self.profiler = "cuda" if maybe_use_cuda_profiler == "1" else None - else: - self._set_from_env_if_set( - "torch_profiler_dir", "VLLM_TORCH_PROFILER_DIR", to_bool=False - ) - if self.torch_profiler_dir: - self.profiler = "torch" - self._set_from_env_if_set( - "torch_profiler_record_shapes", - "VLLM_TORCH_PROFILER_RECORD_SHAPES", - ) - self._set_from_env_if_set( - "torch_profiler_with_memory", - "VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY", - ) - self._set_from_env_if_set( - "torch_profiler_with_stack", - "VLLM_TORCH_PROFILER_WITH_STACK", - ) - self._set_from_env_if_set( - "torch_profiler_with_flops", - "VLLM_TORCH_PROFILER_WITH_FLOPS", - ) - self._set_from_env_if_set( - "ignore_frontend", - "VLLM_TORCH_PROFILER_DISABLE_ASYNC_LLM", - ) - self._set_from_env_if_set( - "torch_profiler_use_gzip", - "VLLM_TORCH_PROFILER_USE_GZIP", - ) - self._set_from_env_if_set( - "torch_profiler_dump_cuda_time_total", - "VLLM_TORCH_PROFILER_DUMP_CUDA_TIME_TOTAL", - ) - - self._set_from_env_if_set( - "delay_iterations", "VLLM_PROFILER_DELAY_ITERS", to_bool=False, to_int=True - ) - self._set_from_env_if_set( - "max_iterations", "VLLM_PROFILER_MAX_ITERS", to_bool=False, to_int=True - ) - has_delay_or_limit = self.delay_iterations > 0 or self.max_iterations > 0 if self.profiler == "torch" and has_delay_or_limit and not self.ignore_frontend: logger.warning_once( diff --git a/vllm/envs.py b/vllm/envs.py index 741a2163c91..5fe65fa750a 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -88,20 +88,6 @@ if TYPE_CHECKING: VLLM_PLUGINS: list[str] | None = None VLLM_LORA_RESOLVER_CACHE_DIR: str | None = None VLLM_LORA_RESOLVER_HF_REPO_LIST: str | None = None - # Deprecated env variables for profiling, kept for backward compatibility - # See also vllm/config/profiler.py and `--profiler-config` argument - VLLM_TORCH_CUDA_PROFILE: str | None = None - VLLM_TORCH_PROFILER_DIR: str | None = None - VLLM_TORCH_PROFILER_RECORD_SHAPES: str | None = None - VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY: str | None = None - VLLM_TORCH_PROFILER_DISABLE_ASYNC_LLM: str | None = None - VLLM_TORCH_PROFILER_WITH_STACK: str | None = None - VLLM_TORCH_PROFILER_WITH_FLOPS: str | None = None - VLLM_TORCH_PROFILER_USE_GZIP: str | None = None - VLLM_TORCH_PROFILER_DUMP_CUDA_TIME_TOTAL: str | None = None - VLLM_PROFILER_DELAY_ITERS: str | None = None - VLLM_PROFILER_MAX_ITERS: str | None = None - # End of deprecated env variables for profiling VLLM_USE_AOT_COMPILE: bool = False VLLM_USE_BYTECODE_HOOK: bool = False VLLM_FORCE_AOT_LOAD: bool = False From 61397891ce00c6e28ca9918fab11be1b9e925a20 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Mon, 2 Feb 2026 23:00:00 -0800 Subject: [PATCH 007/810] [Minor] Some code simplification in `scheduler.py` (#33597) Signed-off-by: Nick Hill --- vllm/v1/core/sched/scheduler.py | 48 +++++++++++++-------------------- 1 file changed, 19 insertions(+), 29 deletions(-) diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 84e6ae1f1cd..a4c692b3d60 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -440,17 +440,13 @@ class Scheduler(SchedulerInterface): ) self.running.remove(preempted_req) if preempted_req in scheduled_running_reqs: + preempted_req_id = preempted_req.request_id scheduled_running_reqs.remove(preempted_req) - token_budget += num_scheduled_tokens[ - preempted_req.request_id - ] - req_to_new_blocks.pop(preempted_req.request_id) - num_scheduled_tokens.pop(preempted_req.request_id) - scheduled_spec_decode_tokens.pop( - preempted_req.request_id, None - ) + token_budget += num_scheduled_tokens.pop(preempted_req_id) + req_to_new_blocks.pop(preempted_req_id) + scheduled_spec_decode_tokens.pop(preempted_req_id, None) preempted_encoder_inputs = scheduled_encoder_inputs.pop( - preempted_req.request_id, None + preempted_req_id, None ) if preempted_encoder_inputs: # Restore encoder compute budget if the preempted @@ -476,8 +472,9 @@ class Scheduler(SchedulerInterface): # Schedule the request. scheduled_running_reqs.append(request) - req_to_new_blocks[request.request_id] = new_blocks - num_scheduled_tokens[request.request_id] = num_new_tokens + request_id = request.request_id + req_to_new_blocks[request_id] = new_blocks + num_scheduled_tokens[request_id] = num_new_tokens token_budget -= num_new_tokens req_index += 1 @@ -492,18 +489,14 @@ class Scheduler(SchedulerInterface): if num_scheduled_spec_tokens > 0: # Trim spec_token_ids list to num_scheduled_spec_tokens. del request.spec_token_ids[num_scheduled_spec_tokens:] - scheduled_spec_decode_tokens[request.request_id] = ( - request.spec_token_ids - ) + scheduled_spec_decode_tokens[request_id] = request.spec_token_ids # New spec tokens will be set in `update_draft_token_ids` before the # next step when applicable. request.spec_token_ids = [] # Encoder-related. if encoder_inputs_to_schedule: - scheduled_encoder_inputs[request.request_id] = ( - encoder_inputs_to_schedule - ) + scheduled_encoder_inputs[request_id] = encoder_inputs_to_schedule # Allocate the encoder cache. for i in encoder_inputs_to_schedule: self.encoder_cache_manager.allocate(request, i) @@ -535,6 +528,7 @@ class Scheduler(SchedulerInterface): break request = self.waiting.peek_request() + request_id = request.request_id # KVTransfer: skip request if still waiting for remote kvs. if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS: @@ -549,7 +543,7 @@ class Scheduler(SchedulerInterface): else: logger.debug( "%s is still in WAITING_FOR_REMOTE_KVS state.", - request.request_id, + request_id, ) self.waiting.pop_request() skipped_waiting_requests.prepend_request(request) @@ -729,7 +723,7 @@ class Scheduler(SchedulerInterface): if self.connector is not None: self.connector.update_state_after_alloc( request, - self.kv_cache_manager.get_blocks(request.request_id), + self.kv_cache_manager.get_blocks(request_id), num_external_computed_tokens, ) @@ -759,10 +753,10 @@ class Scheduler(SchedulerInterface): if self.lora_config and request.lora_request: scheduled_loras.add(request.lora_request.lora_int_id) - req_to_new_blocks[request.request_id] = ( - self.kv_cache_manager.get_blocks(request.request_id) + req_to_new_blocks[request_id] = self.kv_cache_manager.get_blocks( + request_id ) - num_scheduled_tokens[request.request_id] = num_new_tokens + num_scheduled_tokens[request_id] = num_new_tokens token_budget -= num_new_tokens request.status = RequestStatus.RUNNING request.num_computed_tokens = num_computed_tokens @@ -771,9 +765,7 @@ class Scheduler(SchedulerInterface): request.num_cached_tokens = num_computed_tokens # Encoder-related. if encoder_inputs_to_schedule: - scheduled_encoder_inputs[request.request_id] = ( - encoder_inputs_to_schedule - ) + scheduled_encoder_inputs[request_id] = encoder_inputs_to_schedule # Allocate the encoder cache. for i in encoder_inputs_to_schedule: self.encoder_cache_manager.allocate(request, i) @@ -806,11 +798,9 @@ class Scheduler(SchedulerInterface): num_common_prefix_blocks = [0] * len(self.kv_cache_config.kv_cache_groups) with record_function_or_nullcontext("schedule: get_num_common_prefix_blocks"): if self.running: - any_request = self.running[0] + any_request_id = self.running[0].request_id num_common_prefix_blocks = ( - self.kv_cache_manager.get_num_common_prefix_blocks( - any_request.request_id - ) + self.kv_cache_manager.get_num_common_prefix_blocks(any_request_id) ) # Construct the scheduler output. From b95cc5014dc7b260e5c70ae33d1b30c54d11306d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9D=A8=E6=9C=B1=20=C2=B7=20Kiki?= Date: Tue, 3 Feb 2026 15:01:59 +0800 Subject: [PATCH 008/810] [Misc] Remove deprecated VLLM_ALL2ALL_BACKEND environment variable (#33535) Signed-off-by: carlory Co-authored-by: Claude Opus 4.5 --- .../deepseek_v2_lite_ep_eplb.sh | 2 +- .../moe/modular_kernel_tools/common.py | 5 +-- vllm/config/parallel.py | 9 ----- vllm/envs.py | 33 ------------------- 4 files changed, 2 insertions(+), 47 deletions(-) diff --git a/.buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh b/.buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh index 8106f50f18f..463969cbc2a 100644 --- a/.buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh +++ b/.buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh @@ -43,7 +43,6 @@ trap cleanup EXIT for BACK in "${BACKENDS[@]}"; do VLLM_DEEP_GEMM_WARMUP=skip \ - VLLM_ALL2ALL_BACKEND=$BACK \ vllm serve "$MODEL" \ --enforce-eager \ --tensor-parallel-size 2 \ @@ -52,6 +51,7 @@ for BACK in "${BACKENDS[@]}"; do --enable-eplb \ --trust-remote-code \ --max-model-len 2048 \ + --all2all-backend $BACK \ --port $PORT & SERVER_PID=$! wait_for_server $PORT diff --git a/tests/kernels/moe/modular_kernel_tools/common.py b/tests/kernels/moe/modular_kernel_tools/common.py index 4ee18e3428e..327cd44f612 100644 --- a/tests/kernels/moe/modular_kernel_tools/common.py +++ b/tests/kernels/moe/modular_kernel_tools/common.py @@ -150,10 +150,7 @@ class Config: "VLLM_USE_DEEP_GEMM": str(int(self.needs_deep_gemm())), } - backend = self.all2all_backend() - vllm_config.parallel_config.all2all_backend = backend - if backend is not None: - env_dict.update({"VLLM_ALL2ALL_BACKEND": backend}) + vllm_config.parallel_config.all2all_backend = self.all2all_backend() if self.fused_moe_chunk_size is not None: env_dict.update( diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 4bc12b986be..fa1aa03121b 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -549,15 +549,6 @@ class ParallelConfig: return hash_factors(factors) def __post_init__(self) -> None: - # Set all2all_backend from env var if not specified, with deprecation warning - if envs.is_set("VLLM_ALL2ALL_BACKEND"): - logger.warning_once( - "VLLM_ALL2ALL_BACKEND environment variable is deprecated and " - "will be removed in v0.15.0. Please use the " - "--all2all-backend command-line argument instead." - ) - self.all2all_backend = envs.VLLM_ALL2ALL_BACKEND - # Continue with the rest of the initialization self.world_size = ( self.pipeline_parallel_size diff --git a/vllm/envs.py b/vllm/envs.py index 5fe65fa750a..f9aaa4f380c 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -171,15 +171,6 @@ if TYPE_CHECKING: VLLM_NIXL_SIDE_CHANNEL_HOST: str = "localhost" VLLM_NIXL_SIDE_CHANNEL_PORT: int = 5600 VLLM_MOONCAKE_BOOTSTRAP_PORT: int = 8998 - VLLM_ALL2ALL_BACKEND: Literal[ - "naive", - "pplx", - "deepep_high_throughput", - "deepep_low_latency", - "mori", - "allgather_reducescatter", - "flashinfer_all2allv", - ] = "allgather_reducescatter" VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE: int = 163840 VLLM_TOOL_PARSE_REGEX_TIMEOUT_SECONDS: int = 1 VLLM_SLEEP_WHEN_IDLE: bool = False @@ -1292,30 +1283,6 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_MOONCAKE_BOOTSTRAP_PORT": lambda: int( os.getenv("VLLM_MOONCAKE_BOOTSTRAP_PORT", "8998") ), - # [DEPRECATED - will be removed in v0.15.0] all2all backend for vllm's - # expert parallel communication. Use --all2all-backend CLI argument instead. - # Available options: - # - "naive": naive all2all implementation using broadcasts - # - "allgather_reducescatter": all2all implementation based on allgather and - # reducescatter - # - "pplx": use pplx kernels - # - "deepep_high_throughput", use deepep high-throughput kernels - # - "deepep_low_latency", use deepep low-latency kernels - # - "mori", use MoRI kernels - # - "flashinfer_all2allv", use flashinfer alltoallv kernels for mnnvl - "VLLM_ALL2ALL_BACKEND": env_with_choices( - "VLLM_ALL2ALL_BACKEND", - None, - [ - "naive", - "pplx", - "deepep_high_throughput", - "deepep_low_latency", - "mori", - "allgather_reducescatter", - "flashinfer_all2allv", - ], - ), # Flashinfer MoE backend for vLLM's fused Mixture-of-Experts support. # Both require compute capability 10.0 or above. # Available options: From fd9c83d0e05e6a4214be7dabf3b2bd64a9696ed8 Mon Sep 17 00:00:00 2001 From: Richard Zou Date: Mon, 2 Feb 2026 23:16:55 -0800 Subject: [PATCH 009/810] [torch.compile] Document the workaround to standalone_compile failing (#33571) Signed-off-by: Richard Zou --- docs/design/debug_vllm_compile.md | 9 +++++++++ vllm/compilation/compiler_interface.py | 19 +++++++++++++++++++ 2 files changed, 28 insertions(+) diff --git a/docs/design/debug_vllm_compile.md b/docs/design/debug_vllm_compile.md index bc5f46022e9..262782243e7 100644 --- a/docs/design/debug_vllm_compile.md +++ b/docs/design/debug_vllm_compile.md @@ -282,6 +282,15 @@ If vLLM's compile cache is wrong, this usually means that a factor is missing. Please see [this example](https://github.com/vllm-project/vllm/blob/18b39828d90413d05d770dfd2e2f48304f4ca0eb/vllm/config/model.py#L310) of how vLLM computes part of the cache key. +vLLM's compilation cache requires that the code being compiled ends up being serializable. +If this is not the case, then it will error out on save. Usually the fixes are to either: + +- rewrite the non-serializable pieces (perhaps difficult because it's difficult to + tell right now what is serializable and what isn't) +- file a bug report +- ignore the error by setting `VLLM_DISABLE_COMPILE_CACHE=1` (note that this will + make warm server starts a lot slower). + ## Debugging CUDAGraphs CUDAGraphs is a feature that allows one to: diff --git a/vllm/compilation/compiler_interface.py b/vllm/compilation/compiler_interface.py index 331063ff145..60650353971 100644 --- a/vllm/compilation/compiler_interface.py +++ b/vllm/compilation/compiler_interface.py @@ -273,7 +273,26 @@ class InductorStandaloneAdaptor(CompilerInterface): assert key is not None path = os.path.join(self.cache_dir, key) + def is_saveable_2_10(compiled_artifact): + # can just use compiled_artifact.is_saveable in 2.11 + if compiled_artifact._artifacts is None: + return False + _, cache_info = compiled_artifact._artifacts + return len(cache_info.aot_autograd_artifacts) == 1 + if is_compile_cache_enabled(compiler_config): + if not is_saveable_2_10(compiled_graph): + raise RuntimeError( + "The compiled artifact is not serializable. This usually means " + "that the model code has something that is not serializable " + "by torch.compile in it. You can fix this by either " + "figuring out what is not serializable and rewriting it, " + "filing a bug report, " + "or suppressing this error by " + "disabling vLLM's compilation cache via " + "VLLM_DISABLE_COMPILE_CACHE=1 " + "(this will greatly increase vLLM server warm start times)." + ) compiled_graph.save(path=path, format=self.save_format) compilation_counter.num_compiled_artifacts_saved += 1 return compiled_graph, (key, path) From 32e84fa1ff4d371d52042657a05d825c475cba3a Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Tue, 3 Feb 2026 15:49:17 +0800 Subject: [PATCH 010/810] [CI/Build] Investigate torchrun distributed tests hanging issue (#33650) Signed-off-by: Isotr0py --- tests/distributed/test_torchrun_example.py | 3 +++ tests/distributed/test_torchrun_example_moe.py | 3 +++ 2 files changed, 6 insertions(+) diff --git a/tests/distributed/test_torchrun_example.py b/tests/distributed/test_torchrun_example.py index f415409d7b3..35400951966 100644 --- a/tests/distributed/test_torchrun_example.py +++ b/tests/distributed/test_torchrun_example.py @@ -32,6 +32,9 @@ llm = LLM( gpu_memory_utilization=random.uniform(0.7, 0.9), swap_space=random.randint(1, 4), seed=0, + # FIXME(Isotr0py): async scheduling causes deadlock + # on torchrun with PP, need to investigate further. + async_scheduling=False, ) outputs = llm.generate(prompts, sampling_params) diff --git a/tests/distributed/test_torchrun_example_moe.py b/tests/distributed/test_torchrun_example_moe.py index 1aa7f179357..25f55a968c1 100644 --- a/tests/distributed/test_torchrun_example_moe.py +++ b/tests/distributed/test_torchrun_example_moe.py @@ -39,6 +39,9 @@ llm = LLM( gpu_memory_utilization=random.uniform(0.7, 0.9), swap_space=random.randint(1, 4), seed=0, + # FIXME(Isotr0py): async scheduling causes deadlock + # on torchrun with PP, need to investigate further. + async_scheduling=False, ) outputs = llm.generate(prompts, sampling_params) From dad2d6a590207cb8938fb915602794665b8e9326 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lucas=20H=C3=A4nke=20de=20Cansino?= Date: Tue, 3 Feb 2026 09:35:58 +0100 Subject: [PATCH 011/810] [Bugfix][Model] Fix DeepSeek-OCR-2 chat template to include BOS token (#33642) Signed-off-by: l4b4r4b4b4 --- vllm/transformers_utils/configs/deepseek_vl2.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/vllm/transformers_utils/configs/deepseek_vl2.py b/vllm/transformers_utils/configs/deepseek_vl2.py index 8b02a4ddd4b..05067c04cf4 100644 --- a/vllm/transformers_utils/configs/deepseek_vl2.py +++ b/vllm/transformers_utils/configs/deepseek_vl2.py @@ -119,8 +119,9 @@ class DeepseekVLV2Config(PretrainedConfig): self.candidate_resolutions = candidate_resolutions self.vocab_size = self.text_config.vocab_size - # update model_type for OCR model - if "DeepseekOCRForCausalLM" in ( - self.architectures or kwargs.get("architectures", []) - ): + # update model_type for OCR models + architectures = self.architectures or kwargs.get("architectures", []) + if "DeepseekOCRForCausalLM" in architectures: self.model_type = "deepseek_ocr" + elif "DeepseekOCR2ForCausalLM" in architectures: + self.model_type = "deepseek_ocr2" From 83449a5ff04f70a20c24e8e6fc719881b29e10ac Mon Sep 17 00:00:00 2001 From: Cyrus Leung Date: Tue, 3 Feb 2026 18:29:18 +0800 Subject: [PATCH 012/810] [Refactor] Clean up pooling serial utils (#33665) Signed-off-by: DarkLight1337 --- .../embed/embedding_requests_base64_online.py | 8 +- .../embed/embedding_requests_bytes_online.py | 9 +- .../entrypoints/pooling/embed/test_online.py | 16 +- .../pooling/pooling/test_online.py | 14 +- tests/utils_/test_serial_utils.py | 8 +- vllm/entrypoints/pooling/embed/serving.py | 184 +++++++++------- vllm/entrypoints/pooling/pooling/serving.py | 188 ++++++++++------- vllm/entrypoints/pooling/utils.py | 124 +++++++++++ vllm/utils/serial_utils.py | 198 ++++-------------- 9 files changed, 417 insertions(+), 332 deletions(-) create mode 100644 vllm/entrypoints/pooling/utils.py diff --git a/examples/pooling/embed/embedding_requests_base64_online.py b/examples/pooling/embed/embedding_requests_base64_online.py index 88c961370c8..e85af4b858a 100644 --- a/examples/pooling/embed/embedding_requests_base64_online.py +++ b/examples/pooling/embed/embedding_requests_base64_online.py @@ -12,11 +12,7 @@ import base64 import requests import torch -from vllm.utils.serial_utils import ( - EMBED_DTYPE_TO_TORCH_DTYPE, - ENDIANNESS, - binary2tensor, -) +from vllm.utils.serial_utils import EMBED_DTYPES, ENDIANNESS, binary2tensor def post_http_request(prompt: dict, api_url: str) -> requests.Response: @@ -45,7 +41,7 @@ def main(args): ] * 2 # The OpenAI client does not support the embed_dtype and endianness parameters. - for embed_dtype in EMBED_DTYPE_TO_TORCH_DTYPE: + for embed_dtype in EMBED_DTYPES: for endianness in ENDIANNESS: prompt = { "model": model, diff --git a/examples/pooling/embed/embedding_requests_bytes_online.py b/examples/pooling/embed/embedding_requests_bytes_online.py index 6a45beb0bca..fa2fe853f16 100644 --- a/examples/pooling/embed/embedding_requests_bytes_online.py +++ b/examples/pooling/embed/embedding_requests_bytes_online.py @@ -12,13 +12,12 @@ import json import requests import torch -from vllm.utils.serial_utils import ( - EMBED_DTYPE_TO_TORCH_DTYPE, - ENDIANNESS, +from vllm.entrypoints.pooling.utils import ( MetadataItem, build_metadata_items, decode_pooling_output, ) +from vllm.utils.serial_utils import EMBED_DTYPES, ENDIANNESS def post_http_request(prompt: dict, api_url: str) -> requests.Response: @@ -51,7 +50,7 @@ def main(args): # The OpenAI client does not support the bytes encoding_format. # The OpenAI client does not support the embed_dtype and endianness parameters. - for embed_dtype in EMBED_DTYPE_TO_TORCH_DTYPE: + for embed_dtype in EMBED_DTYPES: for endianness in ENDIANNESS: prompt = { "model": model, @@ -74,7 +73,7 @@ def main(args): # The vllm server always sorts the returned embeddings in the order of input. So # returning metadata is not necessary. You can set encoding_format to bytes_only # to let the server not return metadata. - for embed_dtype in EMBED_DTYPE_TO_TORCH_DTYPE: + for embed_dtype in EMBED_DTYPES: for endianness in ENDIANNESS: prompt = { "model": model, diff --git a/tests/entrypoints/pooling/embed/test_online.py b/tests/entrypoints/pooling/embed/test_online.py index 8f3f8a85054..092a5c008ed 100644 --- a/tests/entrypoints/pooling/embed/test_online.py +++ b/tests/entrypoints/pooling/embed/test_online.py @@ -17,16 +17,14 @@ from tests.models.utils import check_embeddings_close from tests.utils import RemoteOpenAIServer from vllm.entrypoints.pooling.embed.protocol import EmbeddingResponse from vllm.entrypoints.pooling.pooling.protocol import PoolingResponse -from vllm.platforms import current_platform -from vllm.tokenizers import get_tokenizer -from vllm.utils.serial_utils import ( - EMBED_DTYPE_TO_TORCH_DTYPE, - ENDIANNESS, +from vllm.entrypoints.pooling.utils import ( MetadataItem, - binary2tensor, build_metadata_items, decode_pooling_output, ) +from vllm.platforms import current_platform +from vllm.tokenizers import get_tokenizer +from vllm.utils.serial_utils import EMBED_DTYPES, ENDIANNESS, binary2tensor MODEL_NAME = "intfloat/multilingual-e5-small" DUMMY_CHAT_TEMPLATE = """{% for message in messages %}{{message['role'] + ': ' + message['content'] + '\\n'}}{% endfor %}""" # noqa: E501 @@ -535,7 +533,7 @@ async def test_base64_embed_dtype_and_endianness( ) float_data = [d.embedding for d in responses_float.data] - for embed_dtype in EMBED_DTYPE_TO_TORCH_DTYPE: + for embed_dtype in EMBED_DTYPES: for endianness in ENDIANNESS: responses_base64 = requests.post( server.url_for("/v1/embeddings"), @@ -574,7 +572,7 @@ async def test_bytes_embed_dtype_and_endianness( ) float_data = [d.embedding for d in responses_float.data] - for embed_dtype in list(EMBED_DTYPE_TO_TORCH_DTYPE.keys()): + for embed_dtype in EMBED_DTYPES: for endianness in ENDIANNESS: responses_bytes = requests.post( server.url_for("/v1/embeddings"), @@ -618,7 +616,7 @@ async def test_bytes_only_embed_dtype_and_endianness( float_data = [d.embedding for d in responses_float.data] embedding_size = len(float_data[0]) - for embed_dtype in list(EMBED_DTYPE_TO_TORCH_DTYPE.keys()): + for embed_dtype in EMBED_DTYPES: for endianness in ENDIANNESS: responses_bytes = requests.post( server.url_for("/v1/embeddings"), diff --git a/tests/entrypoints/pooling/pooling/test_online.py b/tests/entrypoints/pooling/pooling/test_online.py index 0ca841b4a5b..c6a62c19688 100644 --- a/tests/entrypoints/pooling/pooling/test_online.py +++ b/tests/entrypoints/pooling/pooling/test_online.py @@ -12,15 +12,13 @@ import torch from tests.models.utils import check_embeddings_close from tests.utils import RemoteOpenAIServer from vllm.entrypoints.pooling.pooling.protocol import PoolingResponse -from vllm.tokenizers import get_tokenizer -from vllm.utils.serial_utils import ( - EMBED_DTYPE_TO_TORCH_DTYPE, - ENDIANNESS, +from vllm.entrypoints.pooling.utils import ( MetadataItem, - binary2tensor, build_metadata_items, decode_pooling_output, ) +from vllm.tokenizers import get_tokenizer +from vllm.utils.serial_utils import EMBED_DTYPES, ENDIANNESS, binary2tensor MODEL_NAME = "internlm/internlm2-1_8b-reward" DUMMY_CHAT_TEMPLATE = """{% for message in messages %}{{message['role'] + ': ' + message['content'] + '\\n'}}{% endfor %}""" # noqa: E501 @@ -342,7 +340,7 @@ async def test_base64_embed_dtype_and_endianness( responses_float = PoolingResponse.model_validate(float_response.json()) float_data = [np.array(d.data).squeeze(-1).tolist() for d in responses_float.data] - for embed_dtype in EMBED_DTYPE_TO_TORCH_DTYPE: + for embed_dtype in EMBED_DTYPES: for endianness in ENDIANNESS: responses_base64 = requests.post( url, @@ -389,7 +387,7 @@ async def test_bytes_embed_dtype_and_endianness( responses_float = PoolingResponse.model_validate(float_response.json()) float_data = [np.array(d.data).squeeze(-1).tolist() for d in responses_float.data] - for embed_dtype in list(EMBED_DTYPE_TO_TORCH_DTYPE.keys()): + for embed_dtype in EMBED_DTYPES: for endianness in ENDIANNESS: responses_bytes = requests.post( url, @@ -438,7 +436,7 @@ async def test_bytes_only_embed_dtype_and_endianness( float_data = [np.array(d.data).squeeze(-1).tolist() for d in responses_float.data] n_tokens = responses_float.usage.prompt_tokens // len(input_texts) - for embed_dtype in list(EMBED_DTYPE_TO_TORCH_DTYPE.keys()): + for embed_dtype in EMBED_DTYPES: for endianness in ENDIANNESS: responses_bytes = requests.post( url, diff --git a/tests/utils_/test_serial_utils.py b/tests/utils_/test_serial_utils.py index 51b2e4de026..42e466709cb 100644 --- a/tests/utils_/test_serial_utils.py +++ b/tests/utils_/test_serial_utils.py @@ -5,17 +5,19 @@ import torch from tests.models.utils import check_embeddings_close from vllm.utils.serial_utils import ( - EMBED_DTYPE_TO_TORCH_DTYPE, + EMBED_DTYPES, ENDIANNESS, + EmbedDType, + Endianness, binary2tensor, tensor2binary, ) @pytest.mark.parametrize("endianness", ENDIANNESS) -@pytest.mark.parametrize("embed_dtype", EMBED_DTYPE_TO_TORCH_DTYPE.keys()) +@pytest.mark.parametrize("embed_dtype", EMBED_DTYPES.keys()) @torch.inference_mode() -def test_encode_and_decode(embed_dtype: str, endianness: str): +def test_encode_and_decode(embed_dtype: EmbedDType, endianness: Endianness): for i in range(10): tensor = torch.rand(2, 3, 5, 7, 11, 13, device="cpu", dtype=torch.float32) shape = tensor.shape diff --git a/vllm/entrypoints/pooling/embed/serving.py b/vllm/entrypoints/pooling/embed/serving.py index 7c9e840ead4..a535801351c 100644 --- a/vllm/entrypoints/pooling/embed/serving.py +++ b/vllm/entrypoints/pooling/embed/serving.py @@ -1,8 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import json -from collections.abc import AsyncGenerator, Mapping -from typing import Any, Final, TypeAlias +from collections.abc import AsyncGenerator, Callable, Mapping +from functools import partial +from typing import Any, Final, Literal, TypeAlias, cast import torch from fastapi import Request @@ -22,16 +23,18 @@ from vllm.entrypoints.pooling.embed.protocol import ( EmbeddingResponse, EmbeddingResponseData, ) +from vllm.entrypoints.pooling.utils import ( + encode_pooling_bytes, + encode_pooling_output_base64, + encode_pooling_output_float, +) from vllm.inputs.data import EmbedsPrompt, TokensPrompt from vllm.logger import init_logger from vllm.outputs import PoolingOutput, PoolingRequestOutput from vllm.pooling_params import PoolingParams from vllm.utils.async_utils import merge_async_iterators from vllm.utils.collection_utils import chunk_list -from vllm.utils.serial_utils import ( - encode_pooling_bytes, - encode_pooling_output, -) +from vllm.utils.serial_utils import EmbedDType, Endianness logger = init_logger(__name__) @@ -113,79 +116,120 @@ class OpenAIServingEmbedding(OpenAIServing): logger.exception("Error in preprocessing prompt inputs") return self.create_error_response(str(e)) + def request_output_to_embed_json_response( + self, + final_res_batch: list[PoolingRequestOutput], + request_id: str, + created_time: int, + model_name: str, + encoding_format: Literal["float", "base64"], + embed_dtype: EmbedDType, + endianness: Endianness, + ) -> EmbeddingResponse: + encode_fn = cast( + Callable[[PoolingRequestOutput], list[float] | str], + ( + encode_pooling_output_float + if encoding_format == "float" + else partial( + encode_pooling_output_base64, + embed_dtype=embed_dtype, + endianness=endianness, + ) + ), + ) + + items: list[EmbeddingResponseData] = [] + num_prompt_tokens = 0 + + for idx, final_res in enumerate(final_res_batch): + item = EmbeddingResponseData( + index=idx, + embedding=encode_fn(final_res), + ) + prompt_token_ids = final_res.prompt_token_ids + + items.append(item) + num_prompt_tokens += len(prompt_token_ids) + + usage = UsageInfo( + prompt_tokens=num_prompt_tokens, + total_tokens=num_prompt_tokens, + ) + + return EmbeddingResponse( + id=request_id, + created=created_time, + model=model_name, + data=items, + usage=usage, + ) + + def request_output_to_embed_bytes_response( + self, + final_res_batch: list[PoolingRequestOutput], + request_id: str, + created_time: int, + model_name: str, + encoding_format: Literal["bytes", "bytes_only"], + embed_dtype: EmbedDType, + endianness: Endianness, + ) -> EmbeddingBytesResponse: + content, items, usage = encode_pooling_bytes( + pooling_outputs=final_res_batch, + embed_dtype=embed_dtype, + endianness=endianness, + ) + + headers = ( + None + if encoding_format == "bytes_only" + else { + "metadata": json.dumps( + { + "id": request_id, + "created": created_time, + "model": model_name, + "data": items, + "usage": usage, + } + ) + } + ) + + return EmbeddingBytesResponse(content=content, headers=headers) + def _build_response( self, ctx: EmbeddingServeContext, ) -> EmbeddingResponse | EmbeddingBytesResponse | ErrorResponse: - final_res_batch_checked = ctx.final_res_batch - encoding_format = ctx.request.encoding_format embed_dtype = ctx.request.embed_dtype endianness = ctx.request.endianness - def encode_float_base64(): - items: list[EmbeddingResponseData] = [] - num_prompt_tokens = 0 - - for idx, final_res in enumerate(final_res_batch_checked): - item = EmbeddingResponseData( - index=idx, - embedding=encode_pooling_output( - final_res, - encoding_format=encoding_format, - embed_dtype=embed_dtype, - endianness=endianness, - ), - ) - prompt_token_ids = final_res.prompt_token_ids - - items.append(item) - num_prompt_tokens += len(prompt_token_ids) - - usage = UsageInfo( - prompt_tokens=num_prompt_tokens, - total_tokens=num_prompt_tokens, - ) - - return EmbeddingResponse( - id=ctx.request_id, - created=ctx.created_time, - model=ctx.model_name, - data=items, - usage=usage, - ) - - def encode_bytes(bytes_only: bool) -> EmbeddingBytesResponse: - content, items, usage = encode_pooling_bytes( - pooling_outputs=final_res_batch_checked, - embed_dtype=embed_dtype, - endianness=endianness, - ) - - headers = ( - None - if bytes_only - else { - "metadata": json.dumps( - { - "id": ctx.request_id, - "created": ctx.created_time, - "model": ctx.model_name, - "data": items, - "usage": usage, - } - ) - } - ) - - return EmbeddingBytesResponse(content=content, headers=headers) - if encoding_format == "float" or encoding_format == "base64": - return encode_float_base64() - elif encoding_format == "bytes" or encoding_format == "bytes_only": - return encode_bytes(bytes_only=encoding_format == "bytes_only") - else: - assert_never(encoding_format) + return self.request_output_to_embed_json_response( + ctx.final_res_batch, + ctx.request_id, + ctx.created_time, + ctx.model_name, + encoding_format, + embed_dtype, + endianness, + ) + + if encoding_format == "bytes" or encoding_format == "bytes_only": + return self.request_output_to_embed_bytes_response( + ctx.final_res_batch, + ctx.request_id, + ctx.created_time, + ctx.model_name, + encoding_format, + embed_dtype, + endianness, + ) + + assert_never(encoding_format) def _get_max_position_embeddings(self) -> int: """Get the model's effective maximum sequence length for chunking.""" diff --git a/vllm/entrypoints/pooling/pooling/serving.py b/vllm/entrypoints/pooling/pooling/serving.py index 4efc7572b27..423474ca958 100644 --- a/vllm/entrypoints/pooling/pooling/serving.py +++ b/vllm/entrypoints/pooling/pooling/serving.py @@ -4,8 +4,9 @@ import asyncio import json import time -from collections.abc import AsyncGenerator, Sequence -from typing import Any, Final, cast +from collections.abc import AsyncGenerator, Callable, Sequence +from functools import partial +from typing import Any, Final, Literal, cast import jinja2 from fastapi import Request @@ -27,17 +28,16 @@ from vllm.entrypoints.pooling.pooling.protocol import ( PoolingResponse, PoolingResponseData, ) +from vllm.entrypoints.pooling.utils import ( + encode_pooling_bytes, + encode_pooling_output_base64, + encode_pooling_output_float, +) from vllm.logger import init_logger from vllm.outputs import PoolingRequestOutput from vllm.tasks import PoolingTask, SupportedTask from vllm.utils.async_utils import merge_async_iterators -from vllm.utils.serial_utils import ( - EmbedDType, - EncodingFormat, - Endianness, - encode_pooling_bytes, - encode_pooling_output, -) +from vllm.utils.serial_utils import EmbedDType, EncodingFormat, Endianness logger = init_logger(__name__) @@ -256,6 +256,89 @@ class OpenAIServingPooling(OpenAIServing): return response + def request_output_to_pooling_json_response( + self, + final_res_batch: list[PoolingRequestOutput], + request_id: str, + created_time: int, + model_name: str, + encoding_format: Literal["float", "base64"], + embed_dtype: EmbedDType, + endianness: Endianness, + ) -> PoolingResponse: + encode_fn = cast( + Callable[[PoolingRequestOutput], list[float] | str], + ( + encode_pooling_output_float + if encoding_format == "float" + else partial( + encode_pooling_output_base64, + embed_dtype=embed_dtype, + endianness=endianness, + ) + ), + ) + + items: list[PoolingResponseData] = [] + num_prompt_tokens = 0 + + for idx, final_res in enumerate(final_res_batch): + item = PoolingResponseData( + index=idx, + data=encode_fn(final_res), + ) + prompt_token_ids = final_res.prompt_token_ids + + items.append(item) + num_prompt_tokens += len(prompt_token_ids) + + usage = UsageInfo( + prompt_tokens=num_prompt_tokens, + total_tokens=num_prompt_tokens, + ) + + return PoolingResponse( + id=request_id, + created=created_time, + model=model_name, + data=items, + usage=usage, + ) + + def request_output_to_pooling_bytes_response( + self, + final_res_batch: list[PoolingRequestOutput], + request_id: str, + created_time: int, + model_name: str, + encoding_format: Literal["bytes", "bytes_only"], + embed_dtype: EmbedDType, + endianness: Endianness, + ) -> PoolingBytesResponse: + content, items, usage = encode_pooling_bytes( + pooling_outputs=final_res_batch, + embed_dtype=embed_dtype, + endianness=endianness, + ) + + headers = ( + None + if encoding_format == "bytes_only" + else { + "metadata": json.dumps( + { + "id": request_id, + "created": created_time, + "model": model_name, + "data": items, + "usage": usage, + } + ) + } + ) + + return PoolingBytesResponse(content=content, headers=headers) + def request_output_to_pooling_response( self, final_res_batch: list[PoolingRequestOutput], @@ -266,69 +349,26 @@ class OpenAIServingPooling(OpenAIServing): embed_dtype: EmbedDType, endianness: Endianness, ) -> PoolingResponse | PoolingBytesResponse: - def encode_float_base64(): - items: list[PoolingResponseData] = [] - num_prompt_tokens = 0 - - for idx, final_res in enumerate(final_res_batch): - item = PoolingResponseData( - index=idx, - data=encode_pooling_output( - final_res, - encoding_format=encoding_format, - embed_dtype=embed_dtype, - endianness=endianness, - ), - ) - prompt_token_ids = final_res.prompt_token_ids - - items.append(item) - num_prompt_tokens += len(prompt_token_ids) - - usage = UsageInfo( - prompt_tokens=num_prompt_tokens, - total_tokens=num_prompt_tokens, - ) - - return PoolingResponse( - id=request_id, - created=created_time, - model=model_name, - data=items, - usage=usage, - ) - - def encode_bytes(bytes_only: bool) -> PoolingBytesResponse: - content, items, usage = encode_pooling_bytes( - pooling_outputs=final_res_batch, - embed_dtype=embed_dtype, - endianness=endianness, - ) - - headers = ( - None - if bytes_only - else { - "metadata": json.dumps( - { - "id": request_id, - "created": created_time, - "model": model_name, - "data": items, - "usage": usage, - } - ) - } - ) - - return PoolingBytesResponse( - content=content, - headers=headers, - ) - if encoding_format == "float" or encoding_format == "base64": - return encode_float_base64() - elif encoding_format == "bytes" or encoding_format == "bytes_only": - return encode_bytes(bytes_only=encoding_format == "bytes_only") - else: - assert_never(encoding_format) + return self.request_output_to_pooling_json_response( + final_res_batch, + request_id, + created_time, + model_name, + encoding_format, + embed_dtype, + endianness, + ) + + if encoding_format == "bytes" or encoding_format == "bytes_only": + return self.request_output_to_pooling_bytes_response( + final_res_batch, + request_id, + created_time, + model_name, + encoding_format, + embed_dtype, + endianness, + ) + + assert_never(encoding_format) diff --git a/vllm/entrypoints/pooling/utils.py b/vllm/entrypoints/pooling/utils.py new file mode 100644 index 00000000000..dd2f3c874fc --- /dev/null +++ b/vllm/entrypoints/pooling/utils.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import math +from dataclasses import dataclass +from typing import Any + +import pybase64 +import torch + +from vllm.outputs import PoolingRequestOutput +from vllm.utils.serial_utils import ( + EMBED_DTYPES, + EmbedDType, + Endianness, + binary2tensor, + tensor2binary, +) + + +@dataclass +class MetadataItem: + index: int + embed_dtype: EmbedDType + endianness: Endianness + start: int + end: int + shape: tuple[int, ...] + + +def build_metadata_items( + embed_dtype: EmbedDType, + endianness: Endianness, + shape: tuple[int, ...], + n_request: int, +) -> list[MetadataItem]: + n_bytes = EMBED_DTYPES[embed_dtype].nbytes + size = math.prod(shape) + + return [ + MetadataItem( + index=i, + embed_dtype=embed_dtype, + endianness=endianness, + start=i * size * n_bytes, + end=(i + 1) * size * n_bytes, + shape=shape, + ) + for i in range(n_request) + ] + + +def encode_pooling_output_float(output: PoolingRequestOutput) -> list[float]: + return output.outputs.data.tolist() + + +def encode_pooling_output_binary( + output: PoolingRequestOutput, + embed_dtype: EmbedDType, + endianness: Endianness, +) -> bytes: + return tensor2binary(output.outputs.data, embed_dtype, endianness) + + +def encode_pooling_output_base64( + output: PoolingRequestOutput, + embed_dtype: EmbedDType, + endianness: Endianness, +) -> str: + embedding_bytes = tensor2binary(output.outputs.data, embed_dtype, endianness) + return pybase64.b64encode(embedding_bytes).decode("utf-8") + + +def encode_pooling_bytes( + pooling_outputs: list[PoolingRequestOutput], + embed_dtype: EmbedDType, + endianness: Endianness, +) -> tuple[list[bytes], list[dict[str, Any]], dict[str, Any]]: + num_prompt_tokens = 0 + items: list[dict[str, Any]] = [] + body: list[bytes] = [] + offset = 0 + for idx, output in enumerate(pooling_outputs): + binary = tensor2binary( + tensor=output.outputs.data, + embed_dtype=embed_dtype, + endianness=endianness, + ) + size = len(binary) + + # Dictionary form of MetadataItem + item = dict( + index=idx, + embed_dtype=embed_dtype, + endianness=endianness, + start=offset, + end=offset + size, + shape=output.outputs.data.shape, + ) + + body.append(binary) + items.append(item) + prompt_token_ids = output.prompt_token_ids + num_prompt_tokens += len(prompt_token_ids) + offset += size + + # Dictionary form of UsageInfo + usage = dict( + prompt_tokens=num_prompt_tokens, + total_tokens=num_prompt_tokens, + ) + + return body, items, usage + + +def decode_pooling_output(items: list[MetadataItem], body: bytes) -> list[torch.Tensor]: + return [ + binary2tensor( + body[item.start : item.end], + item.shape, + item.embed_dtype, + item.endianness, + ) + for item in sorted(items, key=lambda x: x.index) + ] diff --git a/vllm/utils/serial_utils.py b/vllm/utils/serial_utils.py index 07db5eaf74c..596a7193510 100644 --- a/vllm/utils/serial_utils.py +++ b/vllm/utils/serial_utils.py @@ -1,69 +1,49 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import base64 import io -import math import sys +from collections.abc import Mapping from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Literal +from typing import Literal, get_args import numpy as np +import numpy.typing as npt +import pybase64 import torch -from typing_extensions import assert_never - -if TYPE_CHECKING: - from vllm import PoolingRequestOutput -else: - PoolingRequestOutput = Any sys_byteorder = sys.byteorder -EMBED_DTYPE_TO_TORCH_DTYPE = { - "float32": torch.float32, - "float16": torch.float16, - "bfloat16": torch.bfloat16, - # I'm not sure if other platforms' CPUs support the fp8 data format. - # EMBED_DTYPE only uses the fp8 data representation, - # does not use fp8 computation, and only occurs on the CPU. - # Apologize for any possible break. - "fp8_e4m3": torch.float8_e4m3fn, - "fp8_e5m2": torch.float8_e5m2, -} +@dataclass(frozen=True) +class DTypeInfo: + torch_dtype: torch.dtype -EMBED_DTYPE_TO_N_BYTES = { - "float32": 4, - "float16": 2, - "bfloat16": 2, - "fp8_e4m3": 1, - "fp8_e5m2": 1, -} + torch_view_dtype: torch.dtype + numpy_view_dtype: npt.DTypeLike + @property + def nbytes(self) -> int: + return self.torch_dtype.itemsize -EMBED_DTYPE_TO_TORCH_DTYPE_VIEW = { - "float32": torch.float32, - "float16": torch.float16, - # numpy does not support bfloat16 and fp8 - "bfloat16": torch.float16, - "fp8_e4m3": torch.uint8, - "fp8_e5m2": torch.uint8, -} - -EMBED_DTYPE_TO_NUMPY_DTYPE_VIEW = { - "float32": np.float32, - "float16": np.float16, - # numpy does not support bfloat16 and fp8 - "bfloat16": np.float16, - "fp8_e4m3": np.uint8, - "fp8_e5m2": np.uint8, -} - -ENDIANNESS = ["native", "big", "little"] EmbedDType = Literal["float32", "float16", "bfloat16", "fp8_e4m3", "fp8_e5m2"] Endianness = Literal["native", "big", "little"] EncodingFormat = Literal["float", "base64", "bytes", "bytes_only"] +# I'm not sure if other platforms' CPUs support the fp8 data format. +# EMBED_DTYPE only uses the fp8 data representation, +# does not use fp8 computation, and only occurs on the CPU. +# Apologize for any possible break. +# NOTE: numpy does not support bfloat16 and fp8 +EMBED_DTYPES: Mapping[EmbedDType, DTypeInfo] = { + "float32": DTypeInfo(torch.float32, torch.float32, np.float32), + "float16": DTypeInfo(torch.float16, torch.float16, np.float16), + "bfloat16": DTypeInfo(torch.bfloat16, torch.float16, np.float16), + "fp8_e4m3": DTypeInfo(torch.float8_e4m3fn, torch.uint8, np.uint8), + "fp8_e5m2": DTypeInfo(torch.float8_e5m2, torch.uint8, np.uint8), +} +ENDIANNESS: tuple[Endianness, ...] = get_args(Endianness) + def tensor2base64(x: torch.Tensor) -> str: with io.BytesIO() as buf: @@ -71,21 +51,26 @@ def tensor2base64(x: torch.Tensor) -> str: buf.seek(0) binary_data = buf.read() - return base64.b64encode(binary_data).decode("utf-8") + return pybase64.b64encode(binary_data).decode("utf-8") def tensor2binary( - tensor: torch.Tensor, embed_dtype: EmbedDType, endianness: Endianness + tensor: torch.Tensor, + embed_dtype: EmbedDType, + endianness: Endianness, ) -> bytes: assert isinstance(tensor, torch.Tensor) - assert embed_dtype in EMBED_DTYPE_TO_TORCH_DTYPE + assert embed_dtype in EMBED_DTYPES assert endianness in ENDIANNESS - torch_dtype = EMBED_DTYPE_TO_TORCH_DTYPE[embed_dtype] - torch_view_dtype = EMBED_DTYPE_TO_TORCH_DTYPE_VIEW[embed_dtype] + dtype_info = EMBED_DTYPES[embed_dtype] np_array = ( - tensor.to(torch_dtype).flatten().contiguous().view(torch_view_dtype).numpy() + tensor.to(dtype_info.torch_dtype) + .flatten() + .contiguous() + .view(dtype_info.torch_view_dtype) + .numpy() ) if endianness != "native" and endianness != sys_byteorder: @@ -100,115 +85,14 @@ def binary2tensor( embed_dtype: EmbedDType, endianness: Endianness, ) -> torch.Tensor: - assert embed_dtype in EMBED_DTYPE_TO_TORCH_DTYPE - assert embed_dtype in EMBED_DTYPE_TO_NUMPY_DTYPE_VIEW + assert embed_dtype in EMBED_DTYPES assert endianness in ENDIANNESS - torch_dtype = EMBED_DTYPE_TO_TORCH_DTYPE[embed_dtype] - np_dtype = EMBED_DTYPE_TO_NUMPY_DTYPE_VIEW[embed_dtype] + dtype_info = EMBED_DTYPES[embed_dtype] - np_array = np.frombuffer(binary, dtype=np_dtype).reshape(shape) + np_array = np.frombuffer(binary, dtype=dtype_info.numpy_view_dtype).reshape(shape) if endianness != "native" and endianness != sys_byteorder: np_array = np_array.byteswap() - return torch.from_numpy(np_array).view(torch_dtype) - - -def encode_pooling_output( - output: PoolingRequestOutput, - encoding_format: EncodingFormat, - embed_dtype: EmbedDType, - endianness: Endianness, -) -> list[float] | str | bytes: - if encoding_format == "float": - return output.outputs.data.tolist() - elif encoding_format == "base64": - embedding_bytes = tensor2binary(output.outputs.data, embed_dtype, endianness) - return base64.b64encode(embedding_bytes).decode("utf-8") - elif encoding_format == "bytes" or encoding_format == "bytes_only": - return tensor2binary(output.outputs.data, embed_dtype, endianness) - assert_never(encoding_format) - - -@dataclass -class MetadataItem: - index: int - embed_dtype: EmbedDType - endianness: Endianness - start: int - end: int - shape: tuple[int, ...] - - -def build_metadata_items( - embed_dtype: EmbedDType, - endianness: Endianness, - shape: tuple[int, ...], - n_request: int, -): - n_bytes = EMBED_DTYPE_TO_N_BYTES[embed_dtype] - size = math.prod(shape) - items = [ - MetadataItem( - index=i, - embed_dtype=embed_dtype, - endianness=endianness, - start=i * size * n_bytes, - end=(i + 1) * size * n_bytes, - shape=shape, - ) - for i in range(n_request) - ] - - return items - - -def encode_pooling_bytes( - pooling_outputs: list[PoolingRequestOutput], - embed_dtype: EmbedDType, - endianness: Endianness, -): - num_prompt_tokens = 0 - items: list[dict[str, MetadataItem]] = [] - body = [] - offset = 0 - for idx, output in enumerate(pooling_outputs): - binary = tensor2binary( - tensor=output.outputs.data, - embed_dtype=embed_dtype, - endianness=endianness, - ) - size = len(binary) - - item = { - "index": idx, - "embed_dtype": embed_dtype, - "endianness": endianness, - "start": offset, - "end": offset + size, - "shape": output.outputs.data.shape, - } - - body.append(binary) - items.append(item) - prompt_token_ids = output.prompt_token_ids - num_prompt_tokens += len(prompt_token_ids) - offset += size - - usage = { - "prompt_tokens": num_prompt_tokens, - "total_tokens": num_prompt_tokens, - } - return body, items, usage - - -def decode_pooling_output(items: list[MetadataItem], body: bytes) -> list[torch.Tensor]: - items.sort(key=lambda x: x.index) - - tensor_list: list[torch.Tensor] = [] - for item in items: - binary = body[item.start : item.end] - tensor = binary2tensor(binary, item.shape, item.embed_dtype, item.endianness) - tensor_list.append(tensor) - return tensor_list + return torch.from_numpy(np_array).view(dtype_info.torch_dtype) From e346e2d056a66bb84287e4fea049bde9a37bd72b Mon Sep 17 00:00:00 2001 From: Michael Goin Date: Tue, 3 Feb 2026 05:37:15 -0500 Subject: [PATCH 013/810] [Bugfix] Disable RoutingMethodType.[Renormalize,RenormalizeNaive] TRTLLM per-tensor FP8 MoE (#33620) Signed-off-by: mgoin --- .../layers/fused_moe/flashinfer_trtllm_moe.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py index a066535c51e..43e02d51043 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py @@ -72,8 +72,10 @@ def _supports_routing_method( # NOTE(dbari): as above, potentially allow others here. return routing_method in [ RoutingMethodType.Llama4, - RoutingMethodType.Renormalize, - RoutingMethodType.RenormalizeNaive, + # NOTE(mgoin): Disabled to investigate accuracy issues. + # See https://github.com/vllm-project/vllm/issues/33532 + # RoutingMethodType.Renormalize, + # RoutingMethodType.RenormalizeNaive, ] else: raise ValueError("Unsupported quantization scheme.") From 52683ccbe194688b5c2a1a8ff6b6d9a060a2b2e7 Mon Sep 17 00:00:00 2001 From: Cyrus Leung Date: Tue, 3 Feb 2026 19:13:16 +0800 Subject: [PATCH 014/810] [Misc] Update default image format of `encode_base64` (#33656) Signed-off-by: DarkLight1337 --- vllm/multimodal/media/image.py | 17 +++-------------- vllm/multimodal/utils.py | 5 +---- 2 files changed, 4 insertions(+), 18 deletions(-) diff --git a/vllm/multimodal/media/image.py b/vllm/multimodal/media/image.py index 977a6700736..260ebadd4a3 100644 --- a/vllm/multimodal/media/image.py +++ b/vllm/multimodal/media/image.py @@ -8,13 +8,11 @@ import pybase64 import torch from PIL import Image -from vllm.logger import init_logger +from vllm.utils.serial_utils import tensor2base64 from ..image import convert_image_mode, rgba_to_rgb from .base import MediaIO, MediaWithBytes -logger = init_logger(__file__) - class ImageMediaIO(MediaIO[Image.Image]): def __init__(self, image_mode: str = "RGB", **kwargs) -> None: @@ -77,17 +75,8 @@ class ImageMediaIO(MediaIO[Image.Image]): self, media: Image.Image, *, - image_format: str | None = None, + image_format: str = "PNG", ) -> str: - if image_format is None: - logger.warning_once( - "The default format of `ImageMediaIO.encode_base64` will be changed " - 'from "JPEG" to "PNG" in v0.15 to avoid lossy compression. ' - "To continue using the old default, " - 'pass `format="JPEG"` explicitly to silence this warning.' - ) - image_format = "JPEG" - image = media with BytesIO() as buffer: @@ -121,4 +110,4 @@ class ImageEmbeddingMediaIO(MediaIO[torch.Tensor]): return tensor.to_dense() def encode_base64(self, media: torch.Tensor) -> str: - return pybase64.b64encode(media.numpy()).decode("utf-8") + return tensor2base64(media) diff --git a/vllm/multimodal/utils.py b/vllm/multimodal/utils.py index 2f8c343ca0e..cd116b9b8bc 100644 --- a/vllm/multimodal/utils.py +++ b/vllm/multimodal/utils.py @@ -11,7 +11,6 @@ import numpy as np import numpy.typing as npt from PIL import Image -from vllm.logger import init_logger from vllm.utils.import_utils import LazyLoader from .inputs import ( @@ -27,8 +26,6 @@ if TYPE_CHECKING: else: torch = LazyLoader("torch", globals(), "torch") -logger = init_logger(__name__) - def __getattr__(name: str): if name == "MEDIA_CONNECTOR_REGISTRY": @@ -74,7 +71,7 @@ def encode_image_base64( image: Image.Image, *, image_mode: str = "RGB", - format: str | None = None, + format: str = "PNG", ) -> str: """ Encode a pillow image to base64 format. From ceab70c89d2b1f5eeaeb4582eb927b16dacb7671 Mon Sep 17 00:00:00 2001 From: Song Zhixin Date: Tue, 3 Feb 2026 19:33:56 +0800 Subject: [PATCH 015/810] [Bugfix] fix qwen3-asr response error (#33644) Signed-off-by: jesse Signed-off-by: Cyrus Leung Co-authored-by: Cyrus Leung --- vllm/model_executor/models/qwen3_asr.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/vllm/model_executor/models/qwen3_asr.py b/vllm/model_executor/models/qwen3_asr.py index b27f710dbe1..e63e03e23e4 100644 --- a/vllm/model_executor/models/qwen3_asr.py +++ b/vllm/model_executor/models/qwen3_asr.py @@ -125,6 +125,13 @@ class Qwen3ASRProcessingInfo(BaseProcessingInfo): def get_supported_mm_limits(self) -> Mapping[str, int | None]: return {"audio": None} + def get_data_parser(self) -> MultiModalDataParser: + feature_extractor = self.get_feature_extractor() + return Qwen3ASRMultiModalDataParser( + target_sr=feature_extractor.sampling_rate, + expected_hidden_size=self._get_expected_hidden_size(), + ) + class Qwen3ASRDummyInputsBuilder(BaseDummyInputsBuilder[Qwen3ASRProcessingInfo]): def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: @@ -194,12 +201,6 @@ class Qwen3ASRMultiModalDataParser(MultiModalDataParser): class Qwen3ASRMultiModalProcessor( Qwen3OmniMoeThinkerMultiModalProcessor, ): - def _get_data_parser(self) -> MultiModalDataParser: - feature_extractor = self.info.get_feature_extractor() - return Qwen3ASRMultiModalDataParser( - target_sr=feature_extractor.sampling_rate, - ) - def _get_mm_fields_config( self, hf_inputs: BatchFeature, From f6af34626d37f63ecb128e1f775ebcbbc1d0e5bf Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Tue, 3 Feb 2026 12:07:24 +0000 Subject: [PATCH 016/810] Fix offline test for Transformers v5 (#33682) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- tests/entrypoints/offline_mode/test_offline_mode.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/entrypoints/offline_mode/test_offline_mode.py b/tests/entrypoints/offline_mode/test_offline_mode.py index 539ff89abe9..ad7a3960dd1 100644 --- a/tests/entrypoints/offline_mode/test_offline_mode.py +++ b/tests/entrypoints/offline_mode/test_offline_mode.py @@ -109,8 +109,15 @@ def _re_import_modules(): if k.startswith("transformers") and not k.startswith("transformers_modules") ] + # These modules are aliased in Transformers v5 and so cannot be reloaded directly + aliased_modules = ["tokenization_utils", "tokenization_utils_fast"] + reload_exception = None for module_name in hf_hub_module_names + transformers_module_names: + if any(module_name.endswith(f".{alias}") for alias in aliased_modules): + # Remove from sys.modules so they are re-aliased on next import + del sys.modules[module_name] + continue try: importlib.reload(sys.modules[module_name]) except Exception as e: From be8168ff889aa8981d4e8a158fc1b4d0a4deb18b Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Tue, 3 Feb 2026 12:36:53 +0000 Subject: [PATCH 017/810] Fix Gemma3 GGUF for Transformers v5 (#33683) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- vllm/transformers_utils/gguf_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/transformers_utils/gguf_utils.py b/vllm/transformers_utils/gguf_utils.py index 81d7733160d..3faa5ee60e9 100644 --- a/vllm/transformers_utils/gguf_utils.py +++ b/vllm/transformers_utils/gguf_utils.py @@ -250,7 +250,7 @@ def maybe_patch_hf_config_from_gguf( text_config = hf_config.get_text_config() is_gemma3 = hf_config.model_type in ("gemma3", "gemma3_text") if vision_config is not None and is_gemma3: - new_hf_config = Gemma3Config.from_text_vision_configs( + new_hf_config = Gemma3Config( text_config=text_config, vision_config=vision_config, architectures=["Gemma3ForConditionalGeneration"], From a3acfa10719a931111caccd08ef19f1551b2fe1e Mon Sep 17 00:00:00 2001 From: zxy <46674730+CUHKSZzxy@users.noreply.github.com> Date: Tue, 3 Feb 2026 21:49:45 +0800 Subject: [PATCH 018/810] [Models] Intern-S1-Pro (#33636) Signed-off-by: zxy Signed-off-by: Isotr0py Co-authored-by: Isotr0py --- docs/models/supported_models.md | 1 + examples/offline_inference/vision_language.py | 35 + tests/models/registry.py | 6 + .../layers/rotary_embedding/__init__.py | 23 + .../layers/rotary_embedding/base.py | 21 +- .../layers/rotary_embedding/fope.py | 199 ++++++ vllm/model_executor/models/interns1_pro.py | 633 ++++++++++++++++++ vllm/model_executor/models/qwen3_moe.py | 10 +- vllm/model_executor/models/qwen3_vl.py | 6 +- vllm/model_executor/models/qwen3_vl_moe.py | 15 +- vllm/model_executor/models/registry.py | 4 + 11 files changed, 942 insertions(+), 11 deletions(-) create mode 100644 vllm/model_executor/layers/rotary_embedding/fope.py create mode 100644 vllm/model_executor/models/interns1_pro.py diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 45465d9c42e..a96abd891fb 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -689,6 +689,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen | `Idefics3ForConditionalGeneration` | Idefics3 | T + I | `HuggingFaceM4/Idefics3-8B-Llama3`, etc. | ✅︎ | | | `IsaacForConditionalGeneration` | Isaac | T + I+ | `PerceptronAI/Isaac-0.1` | ✅︎ | ✅︎ | | `InternS1ForConditionalGeneration` | Intern-S1 | T + IE+ + VE+ | `internlm/Intern-S1`, `internlm/Intern-S1-mini`, etc. | ✅︎ | ✅︎ | +| `InternS1ProForConditionalGeneration` | Intern-S1-Pro | T + IE+ + VE+ | `internlm/Intern-S1-Pro`, etc. | ✅︎ | ✅︎ | | `InternVLChatModel` | InternVL 3.5, InternVL 3.0, InternVideo 2.5, InternVL 2.5, Mono-InternVL, InternVL 2.0 | T + IE+ + (VE+) | `OpenGVLab/InternVL3_5-14B`, `OpenGVLab/InternVL3-9B`, `OpenGVLab/InternVideo2_5_Chat_8B`, `OpenGVLab/InternVL2_5-4B`, `OpenGVLab/Mono-InternVL-2B`, `OpenGVLab/InternVL2-4B`, etc. | ✅︎ | ✅︎ | | `InternVLForConditionalGeneration` | InternVL 3.0 (HF format) | T + IE+ + VE+ | `OpenGVLab/InternVL3-1B-hf`, etc. | ✅︎ | ✅︎ | | `KananaVForConditionalGeneration` | Kanana-V | T + I+ | `kakaocorp/kanana-1.5-v-3b-instruct`, etc. | | ✅︎ | diff --git a/examples/offline_inference/vision_language.py b/examples/offline_inference/vision_language.py index dd442d9e3f7..d0122b31840 100755 --- a/examples/offline_inference/vision_language.py +++ b/examples/offline_inference/vision_language.py @@ -842,6 +842,40 @@ def run_interns1(questions: list[str], modality: str) -> ModelRequestData: ) +# Intern-S1-Pro +def run_interns1_pro(questions: list[str], modality: str) -> ModelRequestData: + model_name = "internlm/Intern-S1-Pro" + + engine_args = EngineArgs( + model=model_name, + trust_remote_code=True, + max_model_len=8192, + max_num_seqs=2, + limit_mm_per_prompt={modality: 1}, + enforce_eager=True, + tensor_parallel_size=4, + ) + + if modality == "image": + placeholder = "<|vision_start|><|image_pad|><|vision_end|>" + elif modality == "video": + placeholder = "<|vision_start|><|video_pad|><|vision_end|>" + + tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) + messages = [ + [{"role": "user", "content": f"{placeholder}\n{question}"}] + for question in questions + ] + prompts = tokenizer.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True + ) + + return ModelRequestData( + engine_args=engine_args, + prompts=prompts, + ) + + # InternVL def run_internvl(questions: list[str], modality: str) -> ModelRequestData: model_name = "OpenGVLab/InternVL3-2B" @@ -2130,6 +2164,7 @@ model_example_map = { "hyperclovax_seed_vision": run_hyperclovax_seed_vision, "idefics3": run_idefics3, "interns1": run_interns1, + "interns1_pro": run_interns1_pro, "internvl_chat": run_internvl, "kanana_v": run_kanana_v, "keye_vl": run_keye_vl, diff --git a/tests/models/registry.py b/tests/models/registry.py index 0e3d0d3128c..c38637c1c67 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -755,6 +755,12 @@ _MULTIMODAL_EXAMPLE_MODELS = { "InternS1ForConditionalGeneration": _HfExamplesInfo( "internlm/Intern-S1", trust_remote_code=True ), + "InternS1ProForConditionalGeneration": _HfExamplesInfo( + "internlm/Intern-S1-Pro", + trust_remote_code=True, + min_transformers_version="5.0.0", + is_available_online=False, + ), "InternVLChatModel": _HfExamplesInfo( "OpenGVLab/InternVL2-1B", extras={ diff --git a/vllm/model_executor/layers/rotary_embedding/__init__.py b/vllm/model_executor/layers/rotary_embedding/__init__.py index 127d84555c5..9ad7c9cdafd 100644 --- a/vllm/model_executor/layers/rotary_embedding/__init__.py +++ b/vllm/model_executor/layers/rotary_embedding/__init__.py @@ -11,6 +11,7 @@ from .deepseek_scaling_rope import DeepseekScalingRotaryEmbedding from .dual_chunk_rope import DualChunkRotaryEmbedding from .dynamic_ntk_alpha_rope import DynamicNTKAlphaRotaryEmbedding from .dynamic_ntk_scaling_rope import DynamicNTKScalingRotaryEmbedding +from .fope import FourierRotaryEmbedding from .linear_scaling_rope import LinearScalingRotaryEmbedding from .llama3_rope import Llama3RotaryEmbedding from .llama4_vision_rope import Llama4VisionRotaryEmbedding @@ -102,6 +103,28 @@ def get_rope( mrope_section=rope_parameters["mrope_section"], mrope_interleaved=rope_parameters.get("mrope_interleaved", False), ) + elif "use_fope" in rope_parameters and rope_parameters["use_fope"]: + extra_kwargs = { + k: v + for k, v in rope_parameters.items() + if k + in ( + "num_key_value_heads", + "num_inv_freq", + "fope_sep_head", + "fope_init_factor", + ) + } + extra_kwargs["init_cache"] = False + rotary_emb = FourierRotaryEmbedding( + head_size, + rotary_dim, + max_position, + base, + is_neox_style, + dtype, + **extra_kwargs, + ) else: rotary_emb = RotaryEmbedding( head_size, diff --git a/vllm/model_executor/layers/rotary_embedding/base.py b/vllm/model_executor/layers/rotary_embedding/base.py index ffc6f67daef..2147e00d2db 100644 --- a/vllm/model_executor/layers/rotary_embedding/base.py +++ b/vllm/model_executor/layers/rotary_embedding/base.py @@ -25,6 +25,7 @@ class RotaryEmbeddingBase(CustomOp): base: float, is_neox_style: bool, dtype: torch.dtype, + init_cache: bool = True, ) -> None: super().__init__() self.head_size = head_size @@ -46,11 +47,12 @@ class RotaryEmbeddingBase(CustomOp): if not hasattr(self, "use_flashinfer"): self.use_flashinfer = False - cache = self._compute_cos_sin_cache() - if not self.use_flashinfer: - cache = cache.to(dtype) - self.cos_sin_cache: torch.Tensor - self.register_buffer("cos_sin_cache", cache, persistent=False) + if init_cache: + cache = self._compute_cos_sin_cache() + if not self.use_flashinfer: + cache = cache.to(dtype) + self.cos_sin_cache: torch.Tensor + self.register_buffer("cos_sin_cache", cache, persistent=False) self.is_rocm_triton_rotary_embed_enabled = ( rocm_aiter_ops.is_triton_rotary_embed_enabled() ) @@ -108,9 +110,16 @@ class RotaryEmbedding(RotaryEmbeddingBase): base: float, is_neox_style: bool, dtype: torch.dtype, + init_cache: bool = True, ) -> None: super().__init__( - head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype + head_size=head_size, + rotary_dim=rotary_dim, + max_position_embeddings=max_position_embeddings, + base=base, + is_neox_style=is_neox_style, + dtype=dtype, + init_cache=init_cache, ) @staticmethod diff --git a/vllm/model_executor/layers/rotary_embedding/fope.py b/vllm/model_executor/layers/rotary_embedding/fope.py new file mode 100644 index 00000000000..4c8a7bcbfa1 --- /dev/null +++ b/vllm/model_executor/layers/rotary_embedding/fope.py @@ -0,0 +1,199 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch +import torch.nn.functional as F +from torch import nn + +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) + +from .base import RotaryEmbedding +from .common import rotate_neox + + +class FourierRotaryEmbedding(RotaryEmbedding): + def __init__( + self, + head_size: int, + rotary_dim: int, + max_position_embeddings: int, + base: float, + is_neox_style: bool, + dtype: torch.dtype, + init_cache: bool, + # extra parameters for FoPE + num_key_value_heads: int, + num_inv_freq: int, + fope_sep_head: bool, + fope_init_factor: float, + ): + # fope related parameters + self.num_key_value_heads = num_key_value_heads + self.num_inv_freq = num_inv_freq + self.fope_sep_head = fope_sep_head + self.fope_init_factor = fope_init_factor + + super().__init__( + head_size=head_size, + rotary_dim=rotary_dim, + max_position_embeddings=max_position_embeddings, + base=base, + is_neox_style=is_neox_style, + dtype=dtype, + init_cache=init_cache, + ) + + # setup buffers and parameters + self.inv_freq: torch.Tensor + self.register_buffer( + "inv_freq", self._compute_inv_freq(self.base), persistent=False + ) + + self.input_dim = self.inv_freq.shape[-1] + self.output_dim = self.inv_freq.shape[-1] + self.cos_coef = nn.Parameter( + torch.empty(num_key_value_heads, self.input_dim, self.output_dim), + requires_grad=False, + ) + self.sin_coef = nn.Parameter( + torch.empty(num_key_value_heads, self.input_dim, self.output_dim), + requires_grad=False, + ) + self.sin_coef.weight_loader = self.weight_loader + self.cos_coef.weight_loader = self.weight_loader + + self.cos_sin_cache: torch.Tensor + cache = self._compute_cos_sin_cache().to(dtype) + self.register_buffer("cos_sin_cache", cache, persistent=False) + + # update cache in the first forward, where sin/cos_coef weights are ready + self.update_cache = True + + def _compute_inv_freq(self, base: float) -> torch.Tensor: + """Compute the inverse frequency.""" + inv_freq = 1.0 / ( + base + ** ( + torch.arange(0, self.rotary_dim, 2, dtype=torch.float) / self.rotary_dim + ) + ) + + inv_freq_idx_selected = torch.ones_like(inv_freq, dtype=torch.bool) + if self.num_inv_freq is not None: + inv_freq_idx_selected[self.num_inv_freq :] = False + else: + inv_freq_idx_selected = inv_freq > ( + 2.0 * torch.pi / self.max_position_embeddings + ) + + inv_freq = inv_freq[inv_freq_idx_selected] + return inv_freq + + def _compute_cos_sin_cache(self) -> torch.Tensor: + """Compute the cos and sin cache.""" + device = self.inv_freq.device + t = torch.arange(self.max_position_embeddings, dtype=torch.float, device=device) + + freqs = torch.einsum("j,i -> ji", t, self.inv_freq) + if self.fope_sep_head: + pos_cos = freqs.cos().unsqueeze(0).expand(self.num_key_value_heads, -1, -1) + pos_sin = freqs.sin().unsqueeze(0).expand(self.num_key_value_heads, -1, -1) + else: + pos_cos = freqs.cos() + pos_sin = freqs.sin() + + if self.fope_sep_head: + sin = torch.einsum("htD, hDd -> thd", pos_sin, self.sin_coef.float()) + cos = torch.einsum("htD, hDd -> thd", pos_cos, self.cos_coef.float()) + else: + sin = torch.einsum("tD, Dd -> td", pos_sin, self.sin_coef.float()) + cos = torch.einsum("tD, Dd -> td", pos_cos, self.cos_coef.float()) + + sin = F.pad( + input=sin, + pad=(0, self.head_size // 2 - sin.size(-1)), + mode="constant", + value=1, + ) + cos = F.pad( + input=cos, + pad=(0, self.head_size // 2 - cos.size(-1)), + mode="constant", + value=1, + ) + + sin = torch.cat((sin, sin), dim=-1) + cos = torch.cat((cos, cos), dim=-1) + + # cache: (max_position_embeddings, num_kv_heads, kv_size * 2) + cache = torch.cat((cos, sin), dim=-1) + return cache + + def forward_native( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor | None = None, + offsets: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + # update cos/sin cache in the first forward + if self.update_cache: + cache = self._compute_cos_sin_cache().to(self.dtype) + self.cos_sin_cache.copy_(cache) + self.update_cache = False + + positions = positions.flatten() + cos_sin = self.cos_sin_cache.index_select(0, positions) + cos, sin = cos_sin.chunk(2, dim=-1) + + # apply rotary embedding + # query: (seq_len, num_heads, head_size) + # key: (seq_len, num_kv_heads, head_size) + query = query.unflatten(-1, (-1, self.head_size)) + assert key is not None, "Key tensor is required for FoPE." + key = key.unflatten(-1, (-1, self.head_size)) + + assert query.dim() == key.dim() == 3, ( + "Expected query key (seq_len, heads, head_dim)" + ) + assert cos.dim() <= 3 and sin.dim() <= 3 + + need_reshape = False + if cos.dim() == 3: + # for fope + need_reshape = True + query_shape = query.shape + key_shape = key.shape + cos = cos.flatten(0, 1) + sin = sin.flatten(0, 1) + seq_len = cos.size(0) + query = query.view(seq_len, -1, query.size(-1)) + key = key.view(seq_len, -1, key.size(-1)) + + # native implementation of apply rope for neox style + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + query = (query * cos) + (rotate_neox(query) * sin) + key = (key * cos) + (rotate_neox(key) * sin) + + if need_reshape: + query = query.view(query_shape) + key = key.view(key_shape) + + return query, key + + def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor): + """load fope weights""" + world_size = get_tensor_model_parallel_world_size() + rank = get_tensor_model_parallel_rank() + num_key_value_heads = loaded_weight.size(0) + + if num_key_value_heads < world_size: + n_replicate = world_size // num_key_value_heads + world_size = num_key_value_heads + rank = rank // n_replicate + + loaded_weight = loaded_weight.chunk(world_size, dim=0)[rank] + param.data.copy_(loaded_weight) diff --git a/vllm/model_executor/models/interns1_pro.py b/vllm/model_executor/models/interns1_pro.py new file mode 100644 index 00000000000..60c92cddab3 --- /dev/null +++ b/vllm/model_executor/models/interns1_pro.py @@ -0,0 +1,633 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Copyright 2025 The vLLM team. +# Copyright 2025 The Qwen Team. +# Copyright 2025 The HuggingFace Inc. team. +# All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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 InternS1Pro model compatible with HuggingFace weights.""" + +import functools +from collections.abc import Iterable +from typing import Any + +import torch +from torch import nn +from transformers import AutoProcessor, PretrainedConfig + +from vllm.attention.layer import Attention +from vllm.config import CacheConfig, VllmConfig +from vllm.distributed import ( + get_ep_group, + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_gather, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.activation import SiluAndMul +from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.fused_moe.config import RoutingMethodType +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +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 ( + ParallelLMHead, +) +from vllm.model_executor.models.utils import sequence_parallel_chunk +from vllm.multimodal import MULTIMODAL_REGISTRY + +from .interfaces import MixtureOfExperts +from .qwen3_moe import ( + Qwen3MoeForCausalLM, +) +from .qwen3_vl import ( + Qwen3_VisionTransformer, + Qwen3VLDummyInputsBuilder, + Qwen3VLForConditionalGeneration, + Qwen3VLMultiModalProcessor, + Qwen3VLProcessingInfo, +) +from .qwen3_vl_moe import Qwen3MoeLLMModel +from .utils import ( + AutoWeightsLoader, + WeightsMapper, + extract_layer_index, + maybe_prefix, +) + +logger = init_logger(__name__) + + +class InternS1ProProcessingInfo(Qwen3VLProcessingInfo): + def get_hf_config(self): + return self.ctx.get_hf_config() + + def get_hf_processor(self, **kwargs: object) -> AutoProcessor: + return AutoProcessor.from_pretrained( + self.ctx.model_config.model, + trust_remote_code=True, + **kwargs, + ) + + +class InternS1ProMoeMLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: QuantizationConfig | None = None, + reduce_results: bool = True, + prefix: str = "", + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=reduce_results, + prefix=f"{prefix}.down_proj", + ) + if hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {hidden_act}. Only silu is supported for now." + ) + self.act_fn = SiluAndMul() + + def forward(self, x): + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class InternS1ProMoeSparseMoeBlock(nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + prefix: str = "", + ): + super().__init__() + + config = vllm_config.model_config.hf_text_config + parallel_config = vllm_config.parallel_config + quant_config = vllm_config.quant_config + + self.tp_size = get_tensor_model_parallel_world_size() + + self.ep_group = get_ep_group().device_group + self.ep_rank = get_ep_group().rank_in_group + self.ep_size = self.ep_group.size() + self.n_routed_experts = config.num_experts + + self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe + + if self.tp_size > config.num_experts: + raise ValueError( + f"Tensor parallel size {self.tp_size} is greater than " + f"the number of experts {config.num_experts}." + ) + + # Load balancing settings. + eplb_config = vllm_config.parallel_config.eplb_config + self.enable_eplb = parallel_config.enable_eplb + + self.n_logical_experts = self.n_routed_experts + self.n_redundant_experts = eplb_config.num_redundant_experts + self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts + self.n_local_physical_experts = self.n_physical_experts // self.ep_size + + self.physical_expert_start = self.ep_rank * self.n_local_physical_experts + self.physical_expert_end = ( + self.physical_expert_start + self.n_local_physical_experts + ) + + # For custom routing function + self.n_groups = getattr(config, "router_n_groups", -1) + + self.experts = FusedMoE( + num_experts=self.n_routed_experts, + top_k=config.num_experts_per_tok, + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + reduce_results=True, + renormalize=config.norm_topk_prob, + quant_config=quant_config, + prefix=f"{prefix}.experts", + enable_eplb=self.enable_eplb, + num_redundant_experts=self.n_redundant_experts, + is_sequence_parallel=self.is_sequence_parallel, + routing_method_type=RoutingMethodType.Renormalize, + custom_routing_function=self._custom_routing_function, + ) + + self.gate = ReplicatedLinear( + config.hidden_size, + config.num_experts, + bias=False, + prefix=f"{prefix}.gate", + ) + + @staticmethod + @functools.lru_cache + def get_group_offsets(n_groups: int, group_size: int, device: str): + group_offsets = (torch.arange(n_groups, device=device) * group_size).view( + 1, -1, 1 + ) # [1, n_groups, 1] + return group_offsets + + # TODO: zhouxinyu, use vllm routing functions + def _custom_routing_function( + self, + hidden_states: torch.Tensor, + gating_output: torch.Tensor, + topk: int, + renormalize: bool, + ) -> torch.Tensor: + routing_weights = torch.softmax(gating_output, dim=-1, dtype=torch.float32) + + if self.n_groups > 0: + assert routing_weights.shape[-1] % self.n_groups == 0, ( + f"{routing_weights.shape[-1]} cannot be divided by {self.n_groups}" + ) + per_group_top_k = topk // self.n_groups + group_size = routing_weights.shape[-1] // self.n_groups + group_offsets = self.get_group_offsets( + self.n_groups, group_size, routing_weights.device + ) + routing_weights = routing_weights.unflatten(-1, (self.n_groups, group_size)) + topk_weights, topk_ids = torch.topk( + routing_weights, per_group_top_k, dim=-1 + ) + topk_ids = (topk_ids + group_offsets).flatten(-2, -1) + topk_weights = topk_weights.flatten(-2, -1) + else: + topk_weights, topk_ids = torch.topk(routing_weights, topk, dim=-1) + + if renormalize: + topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) + + return topk_weights, topk_ids + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + assert hidden_states.dim() <= 2, ( + "InternS1ProMoeSparseMoeBlock only supports 1D or 2D inputs" + ) + is_input_1d = hidden_states.dim() == 1 + num_tokens, hidden_dim = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_dim) + + if self.is_sequence_parallel: + hidden_states = sequence_parallel_chunk(hidden_states) + + # router_logits: (num_tokens, n_experts) + router_logits, _ = self.gate(hidden_states) + final_hidden_states = self.experts( + hidden_states=hidden_states, router_logits=router_logits + ) + + if self.is_sequence_parallel: + final_hidden_states = tensor_model_parallel_all_gather( + final_hidden_states, 0 + ) + final_hidden_states = final_hidden_states[:num_tokens] + + # return to 1d if input is 1d + return final_hidden_states.squeeze(0) if is_input_1d else final_hidden_states + + +class InternS1ProMoeAttention(nn.Module): + def __init__( + self, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + rope_parameters: dict[str, Any], + max_position_embeddings: int = 32768, + head_dim: int | None = None, + rms_norm_eps: float = 1e-06, + qkv_bias: bool = False, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + dual_chunk_attention_config: dict[str, Any] | None = None, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + self.head_dim = head_dim or (hidden_size // self.total_num_heads) + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.max_position_embeddings = max_position_embeddings + self.dual_chunk_attention_config = dual_chunk_attention_config + + self.qkv_proj = QKVParallelLinear( + hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=qkv_bias, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + rope_parameters["num_key_value_heads"] = self.num_kv_heads + self.rotary_emb = get_rope( + self.head_dim, + max_position=max_position_embeddings, + rope_parameters=rope_parameters, + dual_chunk_attention_config=dual_chunk_attention_config, + ) + + self.attn = Attention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.attn", + **{ + "layer_idx": extract_layer_index(prefix), + "dual_chunk_attention_config": dual_chunk_attention_config, + } + if dual_chunk_attention_config + else {}, + ) + + self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + # Add qk-norm + q_by_head = q.view(*q.shape[:-1], q.shape[-1] // self.head_dim, self.head_dim) + q_by_head = self.q_norm(q_by_head) + q = q_by_head.view(q.shape) + + k_by_head = k.view(*k.shape[:-1], k.shape[-1] // self.head_dim, self.head_dim) + k_by_head = self.k_norm(k_by_head) + k = k_by_head.view(k.shape) + q, k = self.rotary_emb.forward_native(positions, q, k) + attn_output = self.attn(q, k, v) + output, _ = self.o_proj(attn_output) + return output + + +class InternS1ProMoeDecoderLayer(nn.Module): + def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + + config = vllm_config.model_config.hf_text_config + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + + self.hidden_size = config.hidden_size + max_position_embeddings = getattr(config, "max_position_embeddings", 32768) + dual_chunk_attention_config = getattr( + config, "dual_chunk_attention_config", None + ) + + # update rope related parameters + rope_scaling = config.rope_scaling + fope_keys = {"fope_init_factor", "fope_sep_head", "num_inv_freq"} + use_fope = any(rope_scaling.get(key) is not None for key in fope_keys) + fope_init_factor = rope_scaling.get("fope_init_factor", None) + fope_sep_head = rope_scaling.get("fope_sep_head", None) + num_inv_freq = rope_scaling.get("num_inv_freq", None) + + config.rope_parameters["use_fope"] = use_fope + config.rope_parameters["fope_init_factor"] = fope_init_factor + config.rope_parameters["fope_sep_head"] = fope_sep_head + config.rope_parameters["num_inv_freq"] = num_inv_freq + + assert use_fope, "should use FOPE for InternS1Pro model" + self.self_attn = InternS1ProMoeAttention( + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=config.num_key_value_heads, + rope_parameters=config.rope_parameters, + max_position_embeddings=max_position_embeddings, + rms_norm_eps=config.rms_norm_eps, + qkv_bias=getattr(config, "attention_bias", False), + head_dim=getattr(config, "head_dim", None), + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + dual_chunk_attention_config=dual_chunk_attention_config, + ) + + # `mlp_only_layers` in the config. + layer_idx = extract_layer_index(prefix) + mlp_only_layers = ( + [] if not hasattr(config, "mlp_only_layers") else config.mlp_only_layers + ) + if (layer_idx not in mlp_only_layers) and ( + config.num_experts > 0 and (layer_idx + 1) % config.decoder_sparse_step == 0 + ): + self.mlp = InternS1ProMoeSparseMoeBlock( + vllm_config=vllm_config, prefix=f"{prefix}.mlp" + ) + else: + self.mlp = InternS1ProMoeMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + # Self Attention + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + ) + + # Fully Connected + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + +class InternS1ProMoeLLMModel(Qwen3MoeLLMModel): + def __init__( + self, + *, + vllm_config: VllmConfig, + prefix: str = "", + decoder_layer_type: type[torch.nn.Module] = InternS1ProMoeDecoderLayer, + ): + super().__init__( + vllm_config=vllm_config, + prefix=prefix, + decoder_layer_type=decoder_layer_type, + ) + + +class InternS1ProMoeLLMForCausalLM(Qwen3MoeForCausalLM): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.config = vllm_config.model_config.hf_config.text_config + self.quant_config = vllm_config.quant_config + self.model = InternS1ProMoeLLMModel( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + self.lm_head = ParallelLMHead( + self.config.vocab_size, + self.config.hidden_size, + quant_config=self.quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + if self.config.tie_word_embeddings: + self.lm_head.weight = self.model.embed_tokens.weight + self.logits_processor = LogitsProcessor(self.config.vocab_size) + self.make_empty_intermediate_tensors = ( + self.model.make_empty_intermediate_tensors + ) + + +class Qwen3VLMoeMixtureOfExperts(MixtureOfExperts): + def update_physical_experts_metadata( + self, + num_physical_experts: int, + num_local_physical_experts: int, + ) -> None: + assert self.num_local_physical_experts == num_local_physical_experts + self.num_physical_experts = num_physical_experts + self.num_local_physical_experts = num_local_physical_experts + self.num_redundant_experts = num_physical_experts - self.num_logical_experts + for layer in self.language_model.model.layers: + if isinstance(layer.mlp, InternS1ProMoeSparseMoeBlock): + moe = layer.mlp + moe.n_local_physical_experts = num_local_physical_experts + moe.n_physical_experts = num_physical_experts + moe.n_redundant_experts = self.num_redundant_experts + moe.experts.update_expert_map() + + def set_moe_parameters(self): + self.expert_weights = [] + + self.moe_layers = [] + example_moe = None + for layer in self.language_model.model.layers: + if hasattr(layer, "mlp") and isinstance( + layer.mlp, InternS1ProMoeSparseMoeBlock + ): + example_moe = layer.mlp + self.moe_layers.append(layer.mlp.experts) + + if example_moe is None: + raise RuntimeError("No InternS1ProMoe layer found in the language_model.") + + # Set MoE hyperparameters + self.num_moe_layers = len(self.moe_layers) + self.num_expert_groups = 1 + self.num_shared_experts = 0 + self.num_logical_experts = example_moe.n_logical_experts + self.num_physical_experts = example_moe.n_physical_experts + self.num_local_physical_experts = example_moe.n_local_physical_experts + self.num_routed_experts = example_moe.n_routed_experts + self.num_redundant_experts = example_moe.n_redundant_experts + + +@MULTIMODAL_REGISTRY.register_processor( + Qwen3VLMultiModalProcessor, + info=InternS1ProProcessingInfo, + dummy_inputs=Qwen3VLDummyInputsBuilder, +) +class InternS1ProForConditionalGeneration( + Qwen3VLForConditionalGeneration, Qwen3VLMoeMixtureOfExperts +): + is_3d_moe_weight: bool = True + packed_modules_mapping = { + "qkv_proj": [ + "q_proj", + "k_proj", + "v_proj", + ], + } + + # To ensure correct weight loading and mapping. + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "model.visual.": "visual.", + "lm_head.": "language_model.lm_head.", + "model.language_model.": "language_model.model.", + }, + orig_to_new_suffix={ + # Handle FOPE rotary embeddings + ".rotary_emb.sin_coef": ".layers.0.self_attn.rotary_emb.sin_coef", + ".rotary_emb.cos_coef": ".layers.0.self_attn.rotary_emb.cos_coef", + }, + ) + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config: PretrainedConfig = vllm_config.model_config.hf_config + multimodal_config = vllm_config.model_config.multimodal_config + + self.config = config + self.multimodal_config = multimodal_config + self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" + self.video_pruning_rate = multimodal_config.video_pruning_rate + self.is_multimodal_pruning_enabled = ( + multimodal_config.is_multimodal_pruning_enabled() + ) + + if not multimodal_config.get_limit_per_prompt( + "image" + ) and not multimodal_config.get_limit_per_prompt("video"): + self.visual = None + else: + self.visual = Qwen3_VisionTransformer( + config.vision_config, + norm_eps=getattr(config, "rms_norm_eps", 1e-6), + multimodal_config=multimodal_config, + prefix=maybe_prefix(prefix, "visual"), + ) + + self.language_model = InternS1ProMoeLLMForCausalLM( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "language_model") + ) + # Whether to include the gate_up_proj mapping is determined by + # the language model. + self.packed_modules_mapping = ( + self.packed_modules_mapping | self.language_model.packed_modules_mapping + ) + + self.make_empty_intermediate_tensors = ( + self.language_model.make_empty_intermediate_tensors + ) + + self.use_deepstack = hasattr(config.vision_config, "deepstack_visual_indexes") + self.deepstack_num_level = ( + len(config.vision_config.deepstack_visual_indexes) + if self.use_deepstack + else 0 + ) + self.visual_dim = config.vision_config.out_hidden_size + self.multiscale_dim = self.visual_dim * self.deepstack_num_level + + # Set MoE hyperparameters + self.set_moe_parameters() + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + """load weights""" + skip_prefixes = ["model.time_series."] + if self.visual is None: + skip_prefixes.append("visual.") + loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/qwen3_moe.py b/vllm/model_executor/models/qwen3_moe.py index 2f95f4141c0..45aa58ab2d9 100644 --- a/vllm/model_executor/models/qwen3_moe.py +++ b/vllm/model_executor/models/qwen3_moe.py @@ -428,7 +428,13 @@ class Qwen3MoeDecoderLayer(nn.Module): @support_torch_compile class Qwen3MoeModel(nn.Module): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + def __init__( + self, + *, + vllm_config: VllmConfig, + prefix: str = "", + decoder_layer_type: type[torch.nn.Module] = Qwen3MoeDecoderLayer, + ): super().__init__() config = vllm_config.model_config.hf_text_config @@ -449,7 +455,7 @@ class Qwen3MoeModel(nn.Module): ) self.start_layer, self.end_layer, self.layers = make_layers( config.num_hidden_layers, - lambda prefix: Qwen3MoeDecoderLayer(vllm_config=vllm_config, prefix=prefix), + lambda prefix: decoder_layer_type(vllm_config=vllm_config, prefix=prefix), prefix=f"{prefix}.layers", ) self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) diff --git a/vllm/model_executor/models/qwen3_vl.py b/vllm/model_executor/models/qwen3_vl.py index 97754833953..102d846090c 100644 --- a/vllm/model_executor/models/qwen3_vl.py +++ b/vllm/model_executor/models/qwen3_vl.py @@ -325,7 +325,11 @@ class Qwen3_VisionTransformer(nn.Module): self.spatial_merge_size = vision_config.spatial_merge_size self.spatial_merge_unit = self.spatial_merge_size**2 self.temporal_patch_size = vision_config.temporal_patch_size - self.deepstack_visual_indexes = vision_config.deepstack_visual_indexes + self.deepstack_visual_indexes = ( + vision_config.deepstack_visual_indexes + if hasattr(vision_config, "deepstack_visual_indexes") + else [] + ) self.num_grid_per_side = int(self.num_position_embeddings**0.5) # NOTE: This is used for creating empty tensor for all_gather for diff --git a/vllm/model_executor/models/qwen3_vl_moe.py b/vllm/model_executor/models/qwen3_vl_moe.py index b39a3d297fe..af8536e3f19 100644 --- a/vllm/model_executor/models/qwen3_vl_moe.py +++ b/vllm/model_executor/models/qwen3_vl_moe.py @@ -48,6 +48,7 @@ from vllm.sequence import IntermediateTensors from .interfaces import MixtureOfExperts from .qwen3_moe import ( + Qwen3MoeDecoderLayer, Qwen3MoeForCausalLM, Qwen3MoeModel, Qwen3MoeSparseMoeBlock, @@ -82,8 +83,18 @@ class Qwen3VLMoeProcessingInfo(Qwen3VLProcessingInfo): } ) class Qwen3MoeLLMModel(Qwen3MoeModel): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__(vllm_config=vllm_config, prefix=prefix) + def __init__( + self, + *, + vllm_config: VllmConfig, + prefix: str = "", + decoder_layer_type: type[torch.nn.Module] = Qwen3MoeDecoderLayer, + ): + super().__init__( + vllm_config=vllm_config, + prefix=prefix, + decoder_layer_type=decoder_layer_type, + ) if not get_pp_group().is_first_rank: assert self.start_layer >= len( vllm_config.model_config.hf_config.vision_config.deepstack_visual_indexes diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index ed2a39d2413..5eeb32ed96d 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -357,6 +357,10 @@ _MULTIMODAL_MODELS = { "interns1", "InternS1ForConditionalGeneration", ), + "InternS1ProForConditionalGeneration": ( + "interns1_pro", + "InternS1ProForConditionalGeneration", + ), "Idefics3ForConditionalGeneration": ( "idefics3", "Idefics3ForConditionalGeneration", From 2a8d84e66d19014c44155ca1ee79b4aa0227734d Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Tue, 3 Feb 2026 13:49:49 +0000 Subject: [PATCH 019/810] Fix Gemma3n audio encoder for Transformers v5 (#33673) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- vllm/model_executor/models/gemma3n_mm.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/models/gemma3n_mm.py b/vllm/model_executor/models/gemma3n_mm.py index 1460a458683..8b5e7b8bb52 100644 --- a/vllm/model_executor/models/gemma3n_mm.py +++ b/vllm/model_executor/models/gemma3n_mm.py @@ -621,10 +621,15 @@ class Gemma3nForConditionalGeneration( # Run on padded features to enable batching input_features = audio_input["input_features_padded"].squeeze(1) input_features_mask = audio_input["input_features_mask"].squeeze(1) - audio_outputs, audio_mask = self.audio_tower( - input_features, ~input_features_mask - ) - audio_features = self.embed_audio(inputs_embeds=audio_outputs) + audio_outputs = self.audio_tower(input_features, ~input_features_mask) + if isinstance(audio_outputs, tuple): + # Transformers v4 + audio_encodings, audio_mask = audio_outputs + else: + # Transformers v5 + audio_encodings = audio_outputs.last_hidden_state + audio_mask = audio_outputs.audio_mel_mask + audio_features = self.embed_audio(inputs_embeds=audio_encodings) # The Gemma3nProcessor expects all audio will be 30s in length and # inserts 188 audio soft tokens into the text to account for this. From 2df2b3499dee2025f3f5aa12fb68ea07013c0aa7 Mon Sep 17 00:00:00 2001 From: Krish Gupta Date: Tue, 3 Feb 2026 19:19:59 +0530 Subject: [PATCH 020/810] Document NixlConnector backend selection via kv_connector_extra_config (#33552) Signed-off-by: KrxGu --- docs/features/nixl_connector_usage.md | 29 +++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/docs/features/nixl_connector_usage.md b/docs/features/nixl_connector_usage.md index af38087e4b3..b8364b237e9 100644 --- a/docs/features/nixl_connector_usage.md +++ b/docs/features/nixl_connector_usage.md @@ -36,6 +36,35 @@ export UCX_NET_DEVICES=all # or specify network devices like "mlx5_0:1,mlx5_1:1 !!! tip When using UCX as the transport backend, NCCL environment variables (like `NCCL_IB_HCA`, `NCCL_SOCKET_IFNAME`) are not applicable to NixlConnector, so configure UCX-specific environment variables instead of NCCL variables. +#### Selecting a NIXL transport backend (plugin) + +NixlConnector can use different NIXL transport backends (plugins). By default, NixlConnector uses UCX as the transport backend. + +To select a different backend, set `kv_connector_extra_config.backends` in `--kv-transfer-config`. + +### Example: using LIBFABRIC backend + +```bash +vllm serve \ + --kv-transfer-config '{ + "kv_connector":"NixlConnector", + "kv_role":"kv_both", + "kv_connector_extra_config":{"backends":["LIBFABRIC"]} + }' +``` + +You can also pass JSON keys individually using dotted arguments, and you can append list elements using `+`: + +```bash +vllm serve \ + --kv-transfer-config.kv_connector NixlConnector \ + --kv-transfer-config.kv_role kv_both \ + --kv-transfer-config.kv_connector_extra_config.backends+ LIBFABRIC +``` + +!!! note + Backend availability depends on how NIXL was built and what plugins are present in your environment. Refer to the [NIXL repository](https://github.com/ai-dynamo/nixl) for available backends and build instructions. + ## Basic Usage (on the same host) ### Producer (Prefiller) Configuration From fbb3cf698123cc1243dae8003b63dfa807ef8b53 Mon Sep 17 00:00:00 2001 From: Kuntai Du Date: Tue, 3 Feb 2026 21:50:15 +0800 Subject: [PATCH 021/810] [Bugfix][Async][Connector] avoid vllm-side double free during async scheduling + request abort + async KV cache transfer (#33377) Signed-off-by: KuntaiDu --- vllm/v1/core/sched/scheduler.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index a4c692b3d60..3f7ac9374e1 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1258,10 +1258,14 @@ class Scheduler(SchedulerInterface): # skip failed or rescheduled requests from KV load failure continue request = self.requests.get(req_id) - if request is None: + if request is None or request.is_finished(): # The request is already finished. This can happen if the # request is aborted while the model is executing it (e.g., - # in pipeline parallelism). + # in pipeline parallelism or in async scheduling). + # NOTE(Kuntai): When delay_free_blocks=True (for async KV + # cache transfer in KV connector), the aborted request will not + # be set to None (in order to finish async KV transfer). + # In this case, we use is_finished() to check. continue req_index = model_runner_output.req_id_to_index[req_id] From 4bc913aeeca39a304e4ace51febf55f142c8c86e Mon Sep 17 00:00:00 2001 From: shaharmor98 <17088876+shaharmor98@users.noreply.github.com> Date: Tue, 3 Feb 2026 15:52:49 +0200 Subject: [PATCH 022/810] Feat/add nemotron nano v3 tests (#33345) --- .../NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.yaml | 15 +++++++++++++++ .../NVIDIA-Nemotron-3-Nano-30B-A3B-FP8.yaml | 19 +++++++++++++++++++ .../configs/models-large-hopper.txt | 1 + .../lm-eval-harness/configs/models-large.txt | 1 + tests/config/base_model_arch_groundtruth.json | 17 +++++++++++++++++ tests/config/test_model_arch_config.py | 1 + 6 files changed, 54 insertions(+) create mode 100644 .buildkite/lm-eval-harness/configs/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.yaml create mode 100644 .buildkite/lm-eval-harness/configs/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8.yaml diff --git a/.buildkite/lm-eval-harness/configs/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.yaml b/.buildkite/lm-eval-harness/configs/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.yaml new file mode 100644 index 00000000000..c1dbaef6294 --- /dev/null +++ b/.buildkite/lm-eval-harness/configs/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.yaml @@ -0,0 +1,15 @@ +model_name: "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16" +tasks: +- name: "gsm8k" + metrics: + - name: "exact_match,strict-match" + value: 0.695 + - name: "exact_match,flexible-extract" + value: 0.447 +limit: 1319 +num_fewshot: 5 +max_model_len: 262144 +enforce_eager: false +apply_chat_template: true +fewshot_as_multiturn: true +trust_remote_code: true \ No newline at end of file diff --git a/.buildkite/lm-eval-harness/configs/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8.yaml b/.buildkite/lm-eval-harness/configs/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8.yaml new file mode 100644 index 00000000000..a87328fcdcc --- /dev/null +++ b/.buildkite/lm-eval-harness/configs/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8.yaml @@ -0,0 +1,19 @@ +model_name: "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8" +tasks: +- name: "gsm8k" + metrics: + - name: "exact_match,strict-match" + value: 0.7142 + - name: "exact_match,flexible-extract" + value: 0.4579 +env_vars: + VLLM_USE_FLASHINFER_MOE_FP8: "1" + VLLM_FLASHINFER_MOE_BACKEND: "throughput" +limit: 1319 +num_fewshot: 5 +max_model_len: 262144 +kv_cache_dtype: fp8 +enforce_eager: false +apply_chat_template: true +fewshot_as_multiturn: true +trust_remote_code: true diff --git a/.buildkite/lm-eval-harness/configs/models-large-hopper.txt b/.buildkite/lm-eval-harness/configs/models-large-hopper.txt index 5552391d9ea..2b6c0b5e64d 100644 --- a/.buildkite/lm-eval-harness/configs/models-large-hopper.txt +++ b/.buildkite/lm-eval-harness/configs/models-large-hopper.txt @@ -1 +1,2 @@ Qwen3-235B-A22B-Instruct-2507-FP8.yaml +NVIDIA-Nemotron-3-Nano-30B-A3B-FP8.yaml diff --git a/.buildkite/lm-eval-harness/configs/models-large.txt b/.buildkite/lm-eval-harness/configs/models-large.txt index 37eeac85c93..385031b74f8 100644 --- a/.buildkite/lm-eval-harness/configs/models-large.txt +++ b/.buildkite/lm-eval-harness/configs/models-large.txt @@ -3,3 +3,4 @@ Meta-Llama-3-70B-Instruct.yaml Mixtral-8x7B-Instruct-v0.1.yaml Qwen2-57B-A14-Instruct.yaml DeepSeek-V2-Lite-Chat.yaml +NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.yaml diff --git a/tests/config/base_model_arch_groundtruth.json b/tests/config/base_model_arch_groundtruth.json index 3401198ad7d..81534886dcb 100644 --- a/tests/config/base_model_arch_groundtruth.json +++ b/tests/config/base_model_arch_groundtruth.json @@ -355,5 +355,22 @@ "is_deepseek_mla": true, "is_multimodal_model": false, "dtype": "torch.float32" + }, + "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16": { + "architectures": [ + "NemotronHForCausalLM" + ], + "model_type": "nemotron_h", + "text_model_type": "nemotron_h", + "hidden_size": 2688, + "total_num_hidden_layers": 52, + "total_num_attention_heads": 32, + "head_size": 128, + "vocab_size": 131072, + "total_num_kv_heads": 2, + "num_experts": 128, + "is_deepseek_mla": false, + "is_multimodal_model": false, + "dtype": "torch.bfloat16" } } diff --git a/tests/config/test_model_arch_config.py b/tests/config/test_model_arch_config.py index f28ed173305..fbae31331be 100644 --- a/tests/config/test_model_arch_config.py +++ b/tests/config/test_model_arch_config.py @@ -14,6 +14,7 @@ from vllm.transformers_utils.model_arch_config_convertor import ( BASE_TRUST_REMOTE_CODE_MODELS = { "nvidia/Llama-3_3-Nemotron-Super-49B-v1", + "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16", "XiaomiMiMo/MiMo-7B-RL", # Excluded: Not available online right now # "FreedomIntelligence/openPangu-Ultra-MoE-718B-V1.1", From f3d8a3467111d861cb814152f9c5c8aeaff335c2 Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Tue, 3 Feb 2026 22:43:47 +0800 Subject: [PATCH 023/810] [Bugfix] Do not add extra \n for image-only cases when constructing multimodal text prompts. (#33647) Signed-off-by: wang.yuqi --- vllm/entrypoints/chat_utils.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/vllm/entrypoints/chat_utils.py b/vllm/entrypoints/chat_utils.py index c77c18a5887..0077a897d53 100644 --- a/vllm/entrypoints/chat_utils.py +++ b/vllm/entrypoints/chat_utils.py @@ -1164,7 +1164,10 @@ def _get_full_multimodal_text_prompt( # NOTE: Default behaviour: we always add missing placeholders # at the front of the prompt, if interleave_strings=False - return "\n".join(missing_placeholders + [text_prompt]) + if text_prompt: + return "\n".join(missing_placeholders + [text_prompt]) + else: + return "\n".join(missing_placeholders) # No need to validate using Pydantic again From 5c4f2dd6ef2009d81f4a765b5c2a7278fc389ef3 Mon Sep 17 00:00:00 2001 From: Shanshan Shen <467638484@qq.com> Date: Tue, 3 Feb 2026 22:47:41 +0800 Subject: [PATCH 024/810] [MM] Pass `prefix` parameter to MMEncoderAttention (#33674) Signed-off-by: shen-shanshan <467638484@qq.com> --- vllm/model_executor/models/aimv2.py | 5 ++++- vllm/model_executor/models/blip.py | 5 ++++- vllm/model_executor/models/glm4_1v.py | 1 + vllm/model_executor/models/glm4v.py | 5 ++++- .../model_executor/models/idefics2_vision_model.py | 5 ++++- vllm/model_executor/models/intern_vit.py | 5 ++++- vllm/model_executor/models/interns1_vit.py | 14 ++++++++++++-- vllm/model_executor/models/mllama4.py | 5 ++++- vllm/model_executor/models/molmo.py | 6 +++++- vllm/model_executor/models/molmo2.py | 1 + vllm/model_executor/models/qwen2_5_vl.py | 1 + vllm/model_executor/models/qwen2_vl.py | 1 + .../models/qwen3_omni_moe_thinker.py | 1 + vllm/model_executor/models/step3_vl.py | 7 ++++++- vllm/model_executor/models/step_vl.py | 7 ++++++- 15 files changed, 58 insertions(+), 11 deletions(-) diff --git a/vllm/model_executor/models/aimv2.py b/vllm/model_executor/models/aimv2.py index 5b8ead4c7b7..d6716c8e580 100644 --- a/vllm/model_executor/models/aimv2.py +++ b/vllm/model_executor/models/aimv2.py @@ -127,7 +127,10 @@ class AIMv2Attention(nn.Module): self.num_heads_per_partition = divide(self.num_heads, self.tp_size) self.attn = MMEncoderAttention( - self.num_heads_per_partition, self.head_dim, self.scale + self.num_heads_per_partition, + self.head_dim, + self.scale, + prefix=prefix, ) def forward(self, x: torch.Tensor) -> torch.Tensor: diff --git a/vllm/model_executor/models/blip.py b/vllm/model_executor/models/blip.py index ac9ae49f03e..ad8f7c1af54 100644 --- a/vllm/model_executor/models/blip.py +++ b/vllm/model_executor/models/blip.py @@ -123,7 +123,10 @@ class BlipAttention(nn.Module): self.num_heads_per_partition = divide(self.num_heads, self.tp_size) self.attn = MMEncoderAttention( - self.num_heads_per_partition, self.head_dim, self.scale + self.num_heads_per_partition, + self.head_dim, + self.scale, + prefix=prefix, ) def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): diff --git a/vllm/model_executor/models/glm4_1v.py b/vllm/model_executor/models/glm4_1v.py index 0321851d1a2..b2886e85a3b 100644 --- a/vllm/model_executor/models/glm4_1v.py +++ b/vllm/model_executor/models/glm4_1v.py @@ -296,6 +296,7 @@ class Glm4vVisionAttention(nn.Module): num_heads=self.num_attention_heads_per_partition, head_size=self.hidden_size_per_attention_head, scale=self.hidden_size_per_attention_head**-0.5, + prefix=prefix, ) self.apply_rotary_emb = ApplyRotaryEmb(enforce_enable=True) diff --git a/vllm/model_executor/models/glm4v.py b/vllm/model_executor/models/glm4v.py index ae64913c5a0..8bad386fa8f 100644 --- a/vllm/model_executor/models/glm4v.py +++ b/vllm/model_executor/models/glm4v.py @@ -136,7 +136,10 @@ class EVA2CLIPAttention(nn.Module): ) self.attn = MMEncoderAttention( - self.num_heads_per_rank, self.head_dim, self.scale + self.num_heads_per_rank, + self.head_dim, + self.scale, + prefix=prefix, ) self.output_dropout = torch.nn.Dropout(config.dropout_prob) diff --git a/vllm/model_executor/models/idefics2_vision_model.py b/vllm/model_executor/models/idefics2_vision_model.py index d5fc6e31513..d6f93a9d49b 100644 --- a/vllm/model_executor/models/idefics2_vision_model.py +++ b/vllm/model_executor/models/idefics2_vision_model.py @@ -163,7 +163,10 @@ class Idefics2VisionAttention(nn.Module): ) # Use unified MMEncoderAttention with Flash Attention support self.attn = MMEncoderAttention( - self.num_heads_per_partition, self.head_dim, self.scale + self.num_heads_per_partition, + self.head_dim, + self.scale, + prefix=prefix, ) def forward( diff --git a/vllm/model_executor/models/intern_vit.py b/vllm/model_executor/models/intern_vit.py index 41ca5c29733..8cacfe06eae 100644 --- a/vllm/model_executor/models/intern_vit.py +++ b/vllm/model_executor/models/intern_vit.py @@ -212,7 +212,10 @@ class InternParallelAttention(nn.Module): ) self.attn = MMEncoderAttention( - self.num_heads_per_partition, self.head_dim, self.scale + self.num_heads_per_partition, + self.head_dim, + self.scale, + prefix=prefix, ) def _apply_qk_norm(self, q: torch.Tensor, k: torch.Tensor): diff --git a/vllm/model_executor/models/interns1_vit.py b/vllm/model_executor/models/interns1_vit.py index 195bb96817f..421e0ffd4dd 100644 --- a/vllm/model_executor/models/interns1_vit.py +++ b/vllm/model_executor/models/interns1_vit.py @@ -170,6 +170,7 @@ class InternSdpaAttention(nn.Module): config: PretrainedConfig, *, num_dummy_heads: int = 0, + prefix: str = "", ) -> None: super().__init__() @@ -215,7 +216,12 @@ class InternSdpaAttention(nn.Module): self.projection_layer = nn.Linear(self.dummy_dim, self.embed_dim) # Use unified MMEncoderAttention with automatic backend selection - self.attn = MMEncoderAttention(self.num_heads, self.head_dim, self.scale) + self.attn = MMEncoderAttention( + self.num_heads, + self.head_dim, + self.scale, + prefix=prefix, + ) def forward(self, x: torch.Tensor) -> torch.Tensor: """x shape: (B, N, C)""" @@ -313,7 +319,11 @@ class InternS1VisionLayer(nn.Module): num_dummy_heads: int, prefix: str = "", ): - return InternSdpaAttention(config, num_dummy_heads=num_dummy_heads) + return InternSdpaAttention( + config, + num_dummy_heads=num_dummy_heads, + prefix=prefix, + ) def forward( self, diff --git a/vllm/model_executor/models/mllama4.py b/vllm/model_executor/models/mllama4.py index 54b58299bc2..52fdeddf4b0 100644 --- a/vllm/model_executor/models/mllama4.py +++ b/vllm/model_executor/models/mllama4.py @@ -254,7 +254,10 @@ class Llama4VisionAttention(nn.Module): self.scaling = self.head_dim**-0.5 self.attn = MMEncoderAttention( - self.num_local_heads, self.head_dim, self.scaling + self.num_local_heads, + self.head_dim, + self.scaling, + prefix=prefix, ) if use_data_parallel: diff --git a/vllm/model_executor/models/molmo.py b/vllm/model_executor/models/molmo.py index 7ea06fd85ae..b1330d92d75 100644 --- a/vllm/model_executor/models/molmo.py +++ b/vllm/model_executor/models/molmo.py @@ -231,7 +231,11 @@ class MultiHeadDotProductAttention(nn.Module): self.scale = self.head_dim**-0.5 self.attn = MMEncoderAttention( - self.num_heads, self.head_dim, self.scale, num_kv_heads=self.num_kv_heads + self.num_heads, + self.head_dim, + self.scale, + num_kv_heads=self.num_kv_heads, + prefix=prefix, ) def forward( diff --git a/vllm/model_executor/models/molmo2.py b/vllm/model_executor/models/molmo2.py index cc718d6d52b..f9664f32e4e 100644 --- a/vllm/model_executor/models/molmo2.py +++ b/vllm/model_executor/models/molmo2.py @@ -611,6 +611,7 @@ class ImagePoolingAttention(nn.Module): self.head_dim, self.scale, num_kv_heads=self.num_kv_heads, + prefix=prefix, ) def forward_sdpa( diff --git a/vllm/model_executor/models/qwen2_5_vl.py b/vllm/model_executor/models/qwen2_5_vl.py index 0310c5415dc..c06beb97fac 100644 --- a/vllm/model_executor/models/qwen2_5_vl.py +++ b/vllm/model_executor/models/qwen2_5_vl.py @@ -345,6 +345,7 @@ class Qwen2_5_VisionAttention(nn.Module): num_heads=self.num_attention_heads_per_partition, head_size=self.hidden_size_per_attention_head, scale=self.hidden_size_per_attention_head**-0.5, + prefix=prefix, ) self.apply_rotary_emb = ApplyRotaryEmb(enforce_enable=True) diff --git a/vllm/model_executor/models/qwen2_vl.py b/vllm/model_executor/models/qwen2_vl.py index c7c26c20672..6169e72df4e 100644 --- a/vllm/model_executor/models/qwen2_vl.py +++ b/vllm/model_executor/models/qwen2_vl.py @@ -319,6 +319,7 @@ class Qwen2VisionAttention(nn.Module): num_heads=self.num_attention_heads_per_partition, head_size=self.hidden_size_per_attention_head, scale=self.hidden_size_per_attention_head**-0.5, + prefix=prefix, ) self.apply_rotary_emb = ApplyRotaryEmb(enforce_enable=True) diff --git a/vllm/model_executor/models/qwen3_omni_moe_thinker.py b/vllm/model_executor/models/qwen3_omni_moe_thinker.py index 4d797528f7c..96294158c75 100755 --- a/vllm/model_executor/models/qwen3_omni_moe_thinker.py +++ b/vllm/model_executor/models/qwen3_omni_moe_thinker.py @@ -194,6 +194,7 @@ class Qwen3OmniMoeAudioAttention(nn.Module): num_heads=self.num_local_heads, head_size=self.head_dim, scale=self.scaling, + prefix=prefix, ) def forward( diff --git a/vllm/model_executor/models/step3_vl.py b/vllm/model_executor/models/step3_vl.py index 1cbb54c8425..fe2bb1ac6a6 100644 --- a/vllm/model_executor/models/step3_vl.py +++ b/vllm/model_executor/models/step3_vl.py @@ -759,7 +759,12 @@ class Step3VisionAttention(nn.Module): ) # Use unified MMEncoderAttention with automatic backend selection - self.attn = MMEncoderAttention(self.num_heads, self.head_dim, self.scale) + self.attn = MMEncoderAttention( + self.num_heads, + self.head_dim, + self.scale, + prefix=prefix, + ) def forward( self, diff --git a/vllm/model_executor/models/step_vl.py b/vllm/model_executor/models/step_vl.py index de7db5daace..31b266a7e7e 100644 --- a/vllm/model_executor/models/step_vl.py +++ b/vllm/model_executor/models/step_vl.py @@ -220,7 +220,12 @@ class PerceptionEncoderVisionAttention(nn.Module): prefix=f"{prefix}.out_proj", disable_tp=use_data_parallel, ) - self.attn = MMEncoderAttention(self.num_heads, self.head_dim, self.scale) + self.attn = MMEncoderAttention( + self.num_heads, + self.head_dim, + self.scale, + prefix=prefix, + ) self.rope = PerceptionEncoderRope2D( dim=self.head_dim, max_grid_height=max_grid_height, From f0d525171557e3fe74e8e6df52257f9d66831d3f Mon Sep 17 00:00:00 2001 From: Patrick von Platen Date: Tue, 3 Feb 2026 16:22:34 +0100 Subject: [PATCH 025/810] [Voxtral models] Skip warm-up to skip confusing error message in warm-up (#33576) Signed-off-by: Patrick von Platen Co-authored-by: Cyrus Leung --- vllm/entrypoints/openai/translations/speech_to_text.py | 7 ++++--- vllm/model_executor/models/voxtral.py | 3 +++ vllm/model_executor/models/voxtral_realtime.py | 3 +++ 3 files changed, 10 insertions(+), 3 deletions(-) diff --git a/vllm/entrypoints/openai/translations/speech_to_text.py b/vllm/entrypoints/openai/translations/speech_to_text.py index e80d6b9a253..58bfb3e970f 100644 --- a/vllm/entrypoints/openai/translations/speech_to_text.py +++ b/vllm/entrypoints/openai/translations/speech_to_text.py @@ -138,6 +138,9 @@ class OpenAISpeechToText(OpenAIServing): if not supports_transcription(self.model_cls): return + if getattr(self.model_cls, "skip_warmup_audio_preprocessing", False): + return + try: warmup_start = time.perf_counter() logger.info("Warming up audio preprocessing libraries...") @@ -150,9 +153,7 @@ class OpenAISpeechToText(OpenAIServing): _ = librosa.get_duration(y=dummy_audio, sr=self.asr_config.sample_rate) # Warm up mel-spectrogram computation with model-specific parameters - from vllm.transformers_utils.processor import ( - cached_processor_from_config, - ) + from vllm.transformers_utils.processor import cached_processor_from_config processor = cached_processor_from_config(self.model_config) feature_extractor = None diff --git a/vllm/model_executor/models/voxtral.py b/vllm/model_executor/models/voxtral.py index 86ee981478c..942d91e4454 100644 --- a/vllm/model_executor/models/voxtral.py +++ b/vllm/model_executor/models/voxtral.py @@ -335,6 +335,9 @@ class VoxtralForConditionalGeneration( nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA, SupportsTranscription ): supported_languages = ISO639_1_SUPPORTED_LANGS + # transformers' currently has limited support for MistralCommon backend + # and cached_get_processor. Let's skip until fixed + skip_warmup_audio_preprocessing = True packed_modules_mapping = { "qkv_proj": ["q_proj", "k_proj", "v_proj"], diff --git a/vllm/model_executor/models/voxtral_realtime.py b/vllm/model_executor/models/voxtral_realtime.py index 82801b6eba4..6c4d20d3537 100644 --- a/vllm/model_executor/models/voxtral_realtime.py +++ b/vllm/model_executor/models/voxtral_realtime.py @@ -218,6 +218,9 @@ class VoxtralRealtimeBuffer: @support_torch_compile class VoxtralRealtimeGeneration(VoxtralForConditionalGeneration, SupportsRealtime): requires_raw_input_tokens = True + # transformers' currently has limited support for MistralCommon backend + # and cached_get_processor. Let's skip until fixed + skip_warmup_audio_preprocessing = True def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__(vllm_config=vllm_config, prefix=prefix) From 18e7cbbb158a86bdc76585e64ada795bf1c0d435 Mon Sep 17 00:00:00 2001 From: Cyrus Leung Date: Tue, 3 Feb 2026 23:57:56 +0800 Subject: [PATCH 026/810] [Bugfix] Fix startup hang for Granite Speech (#33699) Signed-off-by: DarkLight1337 --- vllm/multimodal/budget.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/vllm/multimodal/budget.py b/vllm/multimodal/budget.py index 3fbec3d3996..1380ec1ba7c 100644 --- a/vllm/multimodal/budget.py +++ b/vllm/multimodal/budget.py @@ -54,17 +54,17 @@ class MultiModalBudget: self.max_model_len = model_config.max_model_len self.max_num_reqs = scheduler_config.max_num_seqs - cache = mm_registry.processor_only_cache_from_config(vllm_config) - processor = mm_registry.create_processor(model_config, cache=cache) - - self.cache = cache - self.mm_limits = mm_limits = processor.info.allowed_mm_limits - - active_modalities = { - modality for modality, limit in mm_limits.items() if limit > 0 - } - with set_default_torch_num_threads(): # Avoid hang during startup + cache = mm_registry.processor_only_cache_from_config(vllm_config) + processor = mm_registry.create_processor(model_config, cache=cache) + + self.cache = cache + self.mm_limits = mm_limits = processor.info.allowed_mm_limits + + active_modalities = { + modality for modality, limit in mm_limits.items() if limit > 0 + } + all_mm_max_toks_per_item = get_mm_max_toks_per_item( model_config, mm_registry, From 0d6ccf68fa2c439e17d02f26c4044ed5df7f7099 Mon Sep 17 00:00:00 2001 From: dtc Date: Wed, 4 Feb 2026 00:08:25 +0800 Subject: [PATCH 027/810] [P/D] rework mooncake connector and introduce its bootstrap server (#31034) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Tianchen Ding Co-authored-by: Nicolò Lucchesi --- docs/features/disagg_prefill.md | 3 +- docs/features/mooncake_connector_usage.md | 25 +- .../disaggregated_serving/README.md | 1 + .../mooncake_connector_proxy.py | 376 +++++++++ .../run_mooncake_connector.sh | 222 +++++ .../kv_transfer/kv_connector/factory.py | 2 +- .../kv_connector/v1/mooncake/__init__.py | 0 .../v1/{ => mooncake}/mooncake_connector.py | 771 +++++++++++++----- .../v1/mooncake/mooncake_utils.py | 127 +++ 9 files changed, 1324 insertions(+), 203 deletions(-) create mode 100644 examples/online_serving/disaggregated_serving/mooncake_connector/mooncake_connector_proxy.py create mode 100644 examples/online_serving/disaggregated_serving/mooncake_connector/run_mooncake_connector.sh create mode 100644 vllm/distributed/kv_transfer/kv_connector/v1/mooncake/__init__.py rename vllm/distributed/kv_transfer/kv_connector/v1/{ => mooncake}/mooncake_connector.py (51%) create mode 100644 vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_utils.py diff --git a/docs/features/disagg_prefill.md b/docs/features/disagg_prefill.md index df69849bb92..af5f77747fa 100644 --- a/docs/features/disagg_prefill.md +++ b/docs/features/disagg_prefill.md @@ -19,12 +19,13 @@ Two main reasons: Please refer to [examples/online_serving/disaggregated_prefill.sh](../../examples/online_serving/disaggregated_prefill.sh) for the example usage of disaggregated prefilling. -Now supports 5 types of connectors: +Now supports 6 types of connectors: - **ExampleConnector**: refer to [examples/offline_inference/disaggregated-prefill-v1/run.sh](../../examples/offline_inference/disaggregated-prefill-v1/run.sh) for the example usage of ExampleConnector disaggregated prefilling. - **LMCacheConnectorV1**: refer to [examples/others/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh](../../examples/others/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh) for the example usage of LMCacheConnectorV1 disaggregated prefilling which uses NIXL as the underlying KV transmission. - **NixlConnector**: refer to [tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh](../../tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh) for the example usage of NixlConnector disaggregated prefilling which support fully async send/recv. For detailed usage guide, see [NixlConnector Usage Guide](nixl_connector_usage.md). - **P2pNcclConnector**: refer to [examples/online_serving/disaggregated_serving_p2p_nccl_xpyd/disagg_example_p2p_nccl_xpyd.sh](../../examples/online_serving/disaggregated_serving_p2p_nccl_xpyd/disagg_example_p2p_nccl_xpyd.sh) for the example usage of P2pNcclConnector disaggregated prefilling. +- **MooncakeConnector**: refer to [examples/online_serving/disaggregated_serving/mooncake_connector/run_mooncake_connector.sh](../../examples/online_serving/disaggregated_serving/mooncake_connector/run_mooncake_connector.sh) for the example usage of ExampleConnector disaggregated prefilling. For detailed usage guide, see [MooncakeConnector Usage Guide](mooncake_connector_usage.md). - **MultiConnector**: take advantage of the kv_connector_extra_config: dict[str, Any] already present in KVTransferConfig to stash all the connectors we want in an ordered list of kwargs.such as: ```bash diff --git a/docs/features/mooncake_connector_usage.md b/docs/features/mooncake_connector_usage.md index 653ea29ad94..0e2478924ea 100644 --- a/docs/features/mooncake_connector_usage.md +++ b/docs/features/mooncake_connector_usage.md @@ -31,11 +31,9 @@ vllm serve Qwen/Qwen2.5-7B-Instruct --port 8020 --kv-transfer-config '{"kv_conne ### Proxy ```bash -python tests/v1/kv_connector/nixl_integration/toy_proxy_server.py --prefiller-host 192.168.0.2 --prefiller-port 8010 --decoder-host 192.168.0.3 --decoder-port 8020 +python examples/online_serving/disaggregated_serving/mooncake_connector/mooncake_connector_proxy.py --prefill http://192.168.0.2:8010 --decode http://192.168.0.3:8020 ``` -> NOTE: The Mooncake Connector currently uses the proxy from nixl_integration. This will be replaced with a self-developed proxy in the future. - Now you can send requests to the proxy server through port 8000. ## Environment Variables @@ -43,16 +41,29 @@ Now you can send requests to the proxy server through port 8000. - `VLLM_MOONCAKE_BOOTSTRAP_PORT`: Port for Mooncake bootstrap server - Default: 8998 - Required only for prefiller instances - - Each vLLM worker needs a unique port on its host; using the same port number across different hosts is fine - - For TP/DP deployments, each worker's port on a node is computed as: base_port + dp_rank * tp_size + tp_rank - - Used for the decoder notifying the prefiller + - For headless instances, must be the same as the master instance + - Each instance needs a unique port on its host; using the same port number across different hosts is fine - `VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT`: Timeout (in seconds) for automatically releasing the prefiller’s KV cache for a particular request. (Optional) - Default: 480 - If a request is aborted and the decoder has not yet notified the prefiller, the prefill instance will release its KV-cache blocks after this timeout to avoid holding them indefinitely. -## KV Role Options +## KV Transfer Config + +### KV Role Options - **kv_producer**: For prefiller instances that generate KV caches - **kv_consumer**: For decoder instances that consume KV caches from prefiller - **kv_both**: Enables symmetric functionality where the connector can act as both producer and consumer. This provides flexibility for experimental setups and scenarios where the role distinction is not predetermined. + +### kv_connector_extra_config + +- **num_workers**: Size of thread pool for one prefiller worker to transfer KV caches by mooncake. (default 10) +- **mooncake_protocol**: Mooncake connector protocol. (default "rdma") + +## Example Scripts/Code + +Refer to these example scripts in the vLLM repository: + +- [run_mooncake_connector.sh](../../examples/online_serving/disaggregated_serving/mooncake_connector/run_mooncake_connector.sh) +- [mooncake_connector_proxy.py](../../examples/online_serving/disaggregated_serving/mooncake_connector/mooncake_connector_proxy.py) diff --git a/examples/online_serving/disaggregated_serving/README.md b/examples/online_serving/disaggregated_serving/README.md index 090afd7515e..1e328429934 100644 --- a/examples/online_serving/disaggregated_serving/README.md +++ b/examples/online_serving/disaggregated_serving/README.md @@ -6,3 +6,4 @@ This example contains scripts that demonstrate the disaggregated serving feature - `disagg_proxy_demo.py` - Demonstrates XpYd (X prefill instances, Y decode instances). - `kv_events.sh` - Demonstrates KV cache event publishing. +- `mooncake_connector` - A proxy demo for MooncakeConnector. diff --git a/examples/online_serving/disaggregated_serving/mooncake_connector/mooncake_connector_proxy.py b/examples/online_serving/disaggregated_serving/mooncake_connector/mooncake_connector_proxy.py new file mode 100644 index 00000000000..09880a32aa3 --- /dev/null +++ b/examples/online_serving/disaggregated_serving/mooncake_connector/mooncake_connector_proxy.py @@ -0,0 +1,376 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import argparse +import asyncio +import ipaddress +import itertools +import os +import urllib +import uuid +from contextlib import asynccontextmanager +from typing import Any + +import httpx +from fastapi import FastAPI, HTTPException, Request +from fastapi.responses import StreamingResponse + + +def maybe_wrap_ipv6_address(address: str) -> str: + try: + ipaddress.IPv6Address(address) + return f"[{address}]" + except ValueError: + return address + + +def make_http_path(host: str, port: int) -> str: + return f"http://{host}:{port}" + + +def prefiller_cycle(prefill_clients: list[Any]): + while True: + for prefill_client in prefill_clients: + for i in range(prefill_client["dp_size"]): + yield prefill_client, i + + +async def get_prefiller_info(prefill_clients: list, ready: asyncio.Event): + for prefill_client in prefill_clients: + while True: + try: + # Wait for prefill service to be ready + response = await prefill_client["client"].get("/health") + response.raise_for_status() + except Exception: + await asyncio.sleep(1) + continue + + response = await prefill_client["client"].get( + prefill_client["bootstrap_addr"] + "/query" + ) + response.raise_for_status() + data = response.json() + break + + for dp_rank, dp_entry in data.items(): + prefill_client["dp_engine_id"][int(dp_rank)] = dp_entry["engine_id"] + dp_size = len(data) + prefill_client["dp_size"] = dp_size + print(f"Inited prefiller {prefill_client['url']} with dp_size={dp_size}") + + ready.set() + print("All prefiller instances are ready.") + + +@asynccontextmanager +async def lifespan(app: FastAPI): + """ + Lifespan context manager to handle startup and shutdown events. + """ + # Startup: Initialize client pools for prefiller and decoder services + app.state.prefill_clients = [] + app.state.decode_clients = [] + app.state.ready = asyncio.Event() + + # Create prefill clients + for i, (url, bootstrap_port) in enumerate(global_args.prefill): + parsed_url = urllib.parse.urlparse(url) + hostname = maybe_wrap_ipv6_address(parsed_url.hostname) + app.state.prefill_clients.append( + { + "client": httpx.AsyncClient( + timeout=None, + base_url=url, + limits=httpx.Limits( + max_connections=None, + max_keepalive_connections=None, + ), + ), + "url": url, + "bootstrap_addr": make_http_path(hostname, bootstrap_port or 8998), + "dp_engine_id": {}, + } + ) + + # Create decode clients + for i, url in enumerate(global_args.decode): + parsed_url = urllib.parse.urlparse(url) + hostname = maybe_wrap_ipv6_address(parsed_url.hostname) + app.state.decode_clients.append( + { + "client": httpx.AsyncClient( + timeout=None, + base_url=url, + limits=httpx.Limits( + max_connections=None, + max_keepalive_connections=None, + ), + ), + } + ) + + asyncio.create_task(get_prefiller_info(app.state.prefill_clients, app.state.ready)) + + # Initialize round-robin iterators + app.state.prefill_iterator = prefiller_cycle(app.state.prefill_clients) + app.state.decode_iterator = itertools.cycle(range(len(app.state.decode_clients))) + + print( + f"Got {len(app.state.prefill_clients)} prefill clients " + f"and {len(app.state.decode_clients)} decode clients." + ) + + yield + + # Shutdown: Close all clients + for client_info in app.state.prefill_clients: + await client_info["client"].aclose() + + for client_info in app.state.decode_clients: + await client_info["client"].aclose() + + +# Update FastAPI app initialization to use lifespan +app = FastAPI(lifespan=lifespan) + + +def parse_args(): + parser = argparse.ArgumentParser() + + parser.add_argument("--port", type=int, default=8000) + # Always use 127.0.0.1 as localhost binds to IPv6 which is blocked on CI + parser.add_argument("--host", type=str, default="127.0.0.1") + + # For prefiller instances + parser.add_argument( + "--prefill", + nargs="+", + action="append", + dest="prefill_raw", + metavar=("URL", "bootstrap_port"), + help=( + "Prefill server URL and optional bootstrap port. " + "Can be specified multiple times. " + "Format: --prefill URL [BOOTSTRAP_PORT]. " + "BOOTSTRAP_PORT can be a port number, " + "'none', or omitted (defaults to none)." + ), + ) + + # For decoder instances + parser.add_argument( + "--decode", + nargs=1, + action="append", + dest="decode_raw", + metavar=("URL",), + help="Decode server URL. Can be specified multiple times.", + ) + + args = parser.parse_args() + args.prefill = _parse_prefill_urls(args.prefill_raw) + args.decode = _parse_decode_urls(args.decode_raw) + + return args + + +# From sglang router_args.py +def _parse_prefill_urls(prefill_list): + """Parse prefill URLs from --prefill arguments. + + Format: --prefill URL [BOOTSTRAP_PORT] + Example: + --prefill http://prefill1:8080 9000 # With bootstrap port + --prefill http://prefill2:8080 none # Explicitly no bootstrap port + --prefill http://prefill3:8080 # Defaults to no bootstrap port + """ + if not prefill_list: + return [] + + prefill_urls = [] + for prefill_args in prefill_list: + url = prefill_args[0] + + # Handle optional bootstrap port + if len(prefill_args) >= 2: + bootstrap_port_str = prefill_args[1] + # Handle 'none' as None + if bootstrap_port_str.lower() == "none": + bootstrap_port = None + else: + try: + bootstrap_port = int(bootstrap_port_str) + except ValueError as e: + raise ValueError( + f"Invalid bootstrap port: {bootstrap_port_str}. Must be a number or 'none'" # noqa: E501 + ) from e + else: + # No bootstrap port specified, default to None + bootstrap_port = None + + prefill_urls.append((url, bootstrap_port)) + + return prefill_urls + + +def _parse_decode_urls(decode_list): + """Parse decode URLs from --decode arguments. + + Format: --decode URL + Example: --decode http://decode1:8081 --decode http://decode2:8081 + """ + if not decode_list: + return [] + + # decode_list is a list of single-element lists due to nargs=1 + return [url[0] for url in decode_list] + + +def get_next_client(app, service_type: str): + """ + Get the next client in round-robin fashion. + + Args: + app: The FastAPI app instance + service_type: Either 'prefill' or 'decode' + + Returns: + The next client to use + """ + if service_type == "prefill": + return next(app.state.prefill_iterator) + elif service_type == "decode": + client_idx = next(app.state.decode_iterator) + return app.state.decode_clients[client_idx] + else: + raise ValueError(f"Unknown service type: {service_type}") + + +async def send_request_to_service( + client_info: dict, dp_rank: int, endpoint: str, req_data: dict, request_id: str +): + """ + Send a request to a service using a client from the pool. + """ + req_data = req_data.copy() + req_data["kv_transfer_params"] = { + "do_remote_decode": True, + "do_remote_prefill": False, + "transfer_id": f"xfer-{request_id}", + } + req_data["stream"] = False + req_data["max_tokens"] = 1 + if "max_completion_tokens" in req_data: + req_data["max_completion_tokens"] = 1 + if "stream_options" in req_data: + del req_data["stream_options"] + headers = { + "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}", + "X-Request-Id": request_id, + "X-data-parallel-rank": str(dp_rank), + } + + response = await client_info["client"].post( + endpoint, json=req_data, headers=headers + ) + response.raise_for_status() + + # CRITICAL: Release connection back to pool + await response.aclose() + + +async def stream_service_response( + prefill_client_info: dict, + prefill_dp_rank: int, + decode_client_info: dict, + endpoint: str, + req_data: dict, + request_id: str, +): + """ + Asynchronously stream response from a service using a client from the pool. + """ + headers = { + "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}", + "X-Request-Id": request_id, + } + + req_data["kv_transfer_params"] = { + "do_remote_decode": False, + "do_remote_prefill": True, + "remote_bootstrap_addr": prefill_client_info["bootstrap_addr"], + "remote_engine_id": prefill_client_info["dp_engine_id"][prefill_dp_rank], + "transfer_id": f"xfer-{request_id}", + } + + async with decode_client_info["client"].stream( + "POST", endpoint, json=req_data, headers=headers + ) as response: + response.raise_for_status() + async for chunk in response.aiter_bytes(): + yield chunk + + +async def _handle_completions(api: str, request: Request): + if not app.state.ready.is_set(): + raise HTTPException(status_code=503, detail="Service Unavailable") + + try: + req_data = await request.json() + request_id = str(uuid.uuid4()) + + # Get the next prefill client in round-robin fashion + prefill_client_info, prefill_dp_rank = get_next_client(request.app, "prefill") + + # Send request to prefill service + asyncio.create_task( + send_request_to_service( + prefill_client_info, prefill_dp_rank, api, req_data, request_id + ) + ) + + decode_client_info = get_next_client(request.app, "decode") + + # Stream response from decode service + async def generate_stream(): + async for chunk in stream_service_response( + prefill_client_info, + prefill_dp_rank, + decode_client_info, + api, + req_data, + request_id=request_id, + ): + yield chunk + + return StreamingResponse(generate_stream(), media_type="application/json") + + except Exception as e: + import sys + import traceback + + exc_info = sys.exc_info() + print(f"Error occurred in disagg prefill proxy server - {api} endpoint") + print(e) + print("".join(traceback.format_exception(*exc_info))) + raise + + +@app.post("/v1/completions") +async def handle_completions(request: Request): + return await _handle_completions("/v1/completions", request) + + +@app.post("/v1/chat/completions") +async def handle_chat_completions(request: Request): + return await _handle_completions("/v1/chat/completions", request) + + +if __name__ == "__main__": + global global_args + global_args = parse_args() + + import uvicorn + + uvicorn.run(app, host=global_args.host, port=global_args.port) diff --git a/examples/online_serving/disaggregated_serving/mooncake_connector/run_mooncake_connector.sh b/examples/online_serving/disaggregated_serving/mooncake_connector/run_mooncake_connector.sh new file mode 100644 index 00000000000..e38d377c331 --- /dev/null +++ b/examples/online_serving/disaggregated_serving/mooncake_connector/run_mooncake_connector.sh @@ -0,0 +1,222 @@ +#!/bin/bash + +# ============================================================================= +# vLLM Disaggregated Serving Script for Mooncake Connector +# ============================================================================= +# This script demonstrates disaggregated prefill and decode serving using +# Mooncake Connector. +# +# Configuration can be customized via environment variables: +# MODEL: Model to serve +# PREFILL_GPUS: Comma-separated GPU IDs for prefill servers +# DECODE_GPUS: Comma-separated GPU IDs for decode servers +# PREFILL_PORTS: Comma-separated ports for prefill servers +# BOOTSTRAP_PORTS: Bootstrap server port launched by prefill servers +# DECODE_PORTS: Comma-separated ports for decode servers +# PROXY_PORT: Proxy server port used to setup P/D disaggregated connection. +# TIMEOUT_SECONDS: Server startup timeout +# ============================================================================= + +# Configuration - can be overridden via environment variables +MODEL=${MODEL:-Qwen/Qwen2.5-7B-Instruct} +TIMEOUT_SECONDS=${TIMEOUT_SECONDS:-1200} +PROXY_PORT=${PROXY_PORT:-8000} + +PREFILL_GPUS=${PREFILL_GPUS:-0} +DECODE_GPUS=${DECODE_GPUS:-1} +PREFILL_PORTS=${PREFILL_PORTS:-8010} +BOOTSTRAP_PORTS=${BOOTSTRAP_PORTS:-8998} +DECODE_PORTS=${DECODE_PORTS:-8020} + +echo "Warning: Mooncake Connector support for vLLM v1 is experimental and subject to change." +echo "" +echo "Architecture Configuration:" +echo " Model: $MODEL" +echo " Prefill GPUs: $PREFILL_GPUS, Ports: $PREFILL_PORTS, Bootstrap Port:$BOOTSTRAP_PORTS" +echo " Decode GPUs: $DECODE_GPUS, Ports: $DECODE_PORTS" +echo " Proxy Port: $PROXY_PORT" +echo " Timeout: ${TIMEOUT_SECONDS}s" +echo "" + +PIDS=() + +# Switch to the directory of the current script +cd "$(dirname "${BASH_SOURCE[0]}")" + +check_required_files() { + local files=("mooncake_connector_proxy.py") + for file in "${files[@]}"; do + if [[ ! -f "$file" ]]; then + echo "Required file $file not found in $(pwd)" + exit 1 + fi + done +} + +check_hf_token() { + if [ -z "$HF_TOKEN" ]; then + echo "HF_TOKEN is not set. Please set it to your Hugging Face token." + echo "Example: export HF_TOKEN=your_token_here" + exit 1 + fi + if [[ "$HF_TOKEN" != hf_* ]]; then + echo "HF_TOKEN is not a valid Hugging Face token. Please set it to your Hugging Face token." + exit 1 + fi + echo "HF_TOKEN is set and valid." +} + +check_num_gpus() { + # Check if the number of GPUs are >=2 via nvidia-smi + num_gpus=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l) + if [ "$num_gpus" -lt 2 ]; then + echo "You need at least 2 GPUs to run disaggregated prefill." + exit 1 + else + echo "Found $num_gpus GPUs." + fi +} + +ensure_python_library_installed() { + echo "Checking if $1 is installed..." + if ! python3 -c "import $1" > /dev/null 2>&1; then + echo "$1 is not installed. Please install it via pip install $1." + exit 1 + else + echo "$1 is installed." + fi +} + +cleanup() { + echo "Stopping everything…" + trap - INT TERM # prevent re-entrancy + pkill -9 -f "mooncake_connector_proxy.py" + kill -- -$$ # negative PID == "this whole process-group" + wait # reap children so we don't leave zombies + exit 0 +} + +wait_for_server() { + local port=$1 + local timeout_seconds=$TIMEOUT_SECONDS + local start_time=$(date +%s) + + echo "Waiting for server on port $port..." + + while true; do + if curl -s "localhost:${port}/v1/completions" > /dev/null; then + echo "Server on port $port is ready." + return 0 + fi + + local now=$(date +%s) + if (( now - start_time >= timeout_seconds )); then + echo "Timeout waiting for server on port $port" + return 1 + fi + + sleep 1 + done +} + +main() { + check_required_files + check_hf_token + check_num_gpus + ensure_python_library_installed vllm + ensure_python_library_installed mooncake.engine + + trap cleanup INT + trap cleanup USR1 + trap cleanup TERM + + echo "Launching disaggregated serving components..." + echo "Please check the log files for detailed output:" + echo " - prefill*.log: Prefill server logs" + echo " - decode*.log: Decode server logs" + echo " - proxy.log: Proxy server log" + + # Parse GPU and port arrays + IFS=',' read -ra PREFILL_GPU_ARRAY <<< "$PREFILL_GPUS" + IFS=',' read -ra DECODE_GPU_ARRAY <<< "$DECODE_GPUS" + IFS=',' read -ra PREFILL_PORT_ARRAY <<< "$PREFILL_PORTS" + IFS=',' read -ra BOOTSTRAP_PORT_ARRAY <<< "$BOOTSTRAP_PORTS" + IFS=',' read -ra DECODE_PORT_ARRAY <<< "$DECODE_PORTS" + + proxy_param="" + + # ============================================================================= + # Launch Prefill Servers (X Producers) + # ============================================================================= + echo "" + echo "Starting ${#PREFILL_GPU_ARRAY[@]} prefill server(s)..." + for i in "${!PREFILL_GPU_ARRAY[@]}"; do + local gpu_id=${PREFILL_GPU_ARRAY[$i]} + local port=${PREFILL_PORT_ARRAY[$i]} + local bootstrap_port=${BOOTSTRAP_PORT_ARRAY[$i]} + + echo " Prefill server $((i+1)): GPU $gpu_id, Port $port, Bootstrap Port $bootstrap_port" + VLLM_MOONCAKE_BOOTSTRAP_PORT=$bootstrap_port CUDA_VISIBLE_DEVICES=$gpu_id vllm serve $MODEL \ + --port $port \ + --kv-transfer-config \ + "{\"kv_connector\":\"MooncakeConnector\",\"kv_role\":\"kv_producer\"}" > prefill$((i+1)).log 2>&1 & + PIDS+=($!) + proxy_param="${proxy_param} --prefill http://0.0.0.0:${port} $bootstrap_port" + done + + # ============================================================================= + # Launch Decode Servers (Y Decoders) + # ============================================================================= + echo "" + echo "Starting ${#DECODE_GPU_ARRAY[@]} decode server(s)..." + for i in "${!DECODE_GPU_ARRAY[@]}"; do + local gpu_id=${DECODE_GPU_ARRAY[$i]} + local port=${DECODE_PORT_ARRAY[$i]} + + echo " Decode server $((i+1)): GPU $gpu_id, Port $port" + CUDA_VISIBLE_DEVICES=$gpu_id vllm serve $MODEL \ + --port $port \ + --kv-transfer-config \ + "{\"kv_connector\":\"MooncakeConnector\",\"kv_role\":\"kv_consumer\"}" > decode$((i+1)).log 2>&1 & + PIDS+=($!) + proxy_param="${proxy_param} --decode http://0.0.0.0:${port}" + done + + # ============================================================================= + # Launch Proxy Server + # ============================================================================= + echo "" + echo "Starting proxy server on port $PROXY_PORT..." + python3 mooncake_connector_proxy.py $proxy_param --port $PROXY_PORT > proxy.log 2>&1 & + PIDS+=($!) + + # ============================================================================= + # Wait for All Servers to Start + # ============================================================================= + echo "" + echo "Waiting for all servers to start..." + for port in "${PREFILL_PORT_ARRAY[@]}" "${DECODE_PORT_ARRAY[@]}"; do + if ! wait_for_server $port; then + echo "Failed to start server on port $port" + cleanup + exit 1 + fi + done + + echo "" + echo "All servers are up. Starting benchmark..." + + # ============================================================================= + # Run Benchmark + # ============================================================================= + vllm bench serve --port $PROXY_PORT --seed $(date +%s) \ + --backend vllm --model $MODEL \ + --dataset-name random --random-input-len 7500 --random-output-len 200 \ + --num-prompts 200 --burstiness 100 --request-rate 2 | tee benchmark.log + + echo "Benchmarking done. Cleaning up..." + + cleanup +} + +main diff --git a/vllm/distributed/kv_transfer/kv_connector/factory.py b/vllm/distributed/kv_transfer/kv_connector/factory.py index 3933f6d6569..1ceac39711b 100644 --- a/vllm/distributed/kv_transfer/kv_connector/factory.py +++ b/vllm/distributed/kv_transfer/kv_connector/factory.py @@ -198,6 +198,6 @@ KVConnectorFactory.register_connector( ) KVConnectorFactory.register_connector( "MooncakeConnector", - "vllm.distributed.kv_transfer.kv_connector.v1.mooncake_connector", + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector", "MooncakeConnector", ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/__init__.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py similarity index 51% rename from vllm/distributed/kv_transfer/kv_connector/v1/mooncake_connector.py rename to vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index b2b6411f0e0..f105d34928f 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py @@ -6,8 +6,10 @@ import time from collections import defaultdict from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass +from enum import IntEnum from typing import TYPE_CHECKING, Any +import httpx import msgspec import numpy as np import torch @@ -17,6 +19,7 @@ import zmq.asyncio from vllm import envs from vllm.config import VllmConfig from vllm.distributed.kv_transfer.kv_connector.utils import ( + EngineId, TpKVTopology, get_current_attn_backend, ) @@ -25,10 +28,15 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorMetadata, KVConnectorRole, ) +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_utils import ( + MooncakeBootstrapServer, + RegisterWorkerPayload, +) from vllm.distributed.parallel_state import ( + get_pp_group, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, - get_tp_group, + is_local_first_rank, ) from vllm.forward_context import ForwardContext from vllm.logger import init_logger @@ -43,7 +51,7 @@ try: except ImportError as e: raise ImportError( "Please install mooncake by following the instructions at " - "https://github.com/kvcache-ai/Mooncake/blob/main/doc/en/build.md " # noqa: E501 + "https://github.com/kvcache-ai/Mooncake/blob/main/doc/en/build.md " "to run VLLM with MooncakeTransferEngine." ) from e @@ -52,46 +60,75 @@ if TYPE_CHECKING: from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.request import Request -EngineId = str -ReqId = str - -TRANS_DONE = b"trans_done" -TRANS_ERROR = b"trans_error" +ReqId = str # Internal scheduler request ID +TransferId = str # KV transfer coordination ID (shared by P/D) logger = init_logger(__name__) -class MooncakeAgentMetadata( +class MooncakeXferMetadata( msgspec.Struct, omit_defaults=True, # type: ignore[call-arg] - # required for @cached_property. - dict=True, ): remote_hostname: str remote_port: int - request_ids: list[ReqId] + remote_tp_size: int + remote_tp_rank: int + req_blocks: dict[ReqId, tuple[TransferId, list[int]]] kv_caches_base_addr: list[int] - block_ids: list[list[int]] + + +class MooncakeXferResponseStatus(IntEnum): + # Transfer finished + FINISH = 0 + # Continue to receive + CONTINUE = 1 + # Something wrong, see err_msg + ERROR = 2 + + +class MooncakeXferResponse( + msgspec.Struct, + omit_defaults=True, # type: ignore[call-arg] +): + status: MooncakeXferResponseStatus + ok_reqs: list[ReqId] | None = None + err_reqs: list[ReqId] | None = None + err_msg: str | None = None @dataclass -class RecvReqMeta: +class PullReqMeta: + d_req_id: ReqId + transfer_id: TransferId local_block_ids: list[int] - remote_host: str - remote_port: int + remote_engine_id: EngineId + remote_bootstrap_addr: str + # Set expire time to avoid infinitely sending requests. + expire_time: float = float("inf") + # Designed for one D pairing to multiple P + pull_tasks_count: int = 0 @dataclass class SendBlockMeta: + p_req_id: ReqId + transfer_id: TransferId local_block_ids: list[int] ready: asyncio.Event expire_time: float = float("inf") + need_send: int = 0 + sent: int = 0 + sending: int = 0 class MooncakeConnectorMetadata(KVConnectorMetadata): def __init__(self): - self.reqs_to_recv: dict[ReqId, RecvReqMeta] = {} - self.reqs_to_send: dict[ReqId, list[int]] = {} + # Use (engine_id, dp_rank) to group reqs with same dp. + # See comments in MooncakeBootstrapServer. + self.reqs_to_recv: dict[EngineId, dict[ReqId, PullReqMeta]] = defaultdict(dict) + self.reqs_to_send: dict[ReqId, tuple[TransferId, list[int]]] = {} + self.reqs_not_processed: set[TransferId] = set() def add_new_req( self, @@ -100,14 +137,18 @@ class MooncakeConnectorMetadata(KVConnectorMetadata): kv_transfer_params: dict[str, Any], load_remote_cache: bool = True, ): + transfer_id = kv_transfer_params["transfer_id"] if load_remote_cache: - self.reqs_to_recv[request_id] = RecvReqMeta( + remote_engine_id = kv_transfer_params["remote_engine_id"] + self.reqs_to_recv[remote_engine_id][request_id] = PullReqMeta( + d_req_id=request_id, local_block_ids=local_block_ids, - remote_host=kv_transfer_params["remote_host"], - remote_port=kv_transfer_params["remote_port"], + remote_engine_id=remote_engine_id, + remote_bootstrap_addr=kv_transfer_params["remote_bootstrap_addr"], + transfer_id=transfer_id, ) else: - self.reqs_to_send[request_id] = local_block_ids + self.reqs_to_send[request_id] = (transfer_id, local_block_ids) class MooncakeConnector(KVConnectorBase_V1): @@ -209,19 +250,24 @@ class MooncakeConnectorScheduler: def __init__(self, vllm_config: VllmConfig, engine_id: str): self.vllm_config = vllm_config - self.engine_id: EngineId = engine_id - self.side_channel_host = get_ip() - self.side_channel_port = get_mooncake_side_channel_port(vllm_config) assert vllm_config.kv_transfer_config - self.kv_role = vllm_config.kv_transfer_config.kv_role + self.is_kv_producer: bool = ( + vllm_config.kv_transfer_config.kv_role == "kv_producer" + ) + self.is_kv_consumer: bool = ( + vllm_config.kv_transfer_config.kv_role == "kv_consumer" + ) logger.info("Initializing Mooncake Transfer Engine Scheduler %s", engine_id) # Requests that need to start recv/send. # New requests are added by update_state_after_alloc in # the scheduler. Used to make metadata passed to Worker. self._reqs_need_recv: dict[ReqId, tuple[Request, list[int]]] = {} - self._reqs_need_send: dict[ReqId, list[int]] = {} + self._reqs_need_send: dict[ReqId, tuple[Request, list[int]]] = {} + # Reqs to remove from processed set because they're not to send after + # remote prefill or aborted. + self._reqs_not_processed: set[TransferId] = set() def get_num_new_matched_tokens( self, request: "Request", num_computed_tokens: int @@ -249,8 +295,12 @@ class MooncakeConnectorScheduler: params, ) - if params is not None and params.get("do_remote_prefill"): + if not params: + return 0, False + + if params.get("do_remote_prefill"): # Remote prefill: get all prompt blocks from remote. + assert not self.is_kv_producer token_ids = request.prompt_token_ids or [] count = len(token_ids) - num_computed_tokens if count > 0: @@ -265,7 +315,8 @@ class MooncakeConnectorScheduler: params = request.kv_transfer_params logger.debug( "MooncakeConnector update_state_after_alloc: " - "num_external_tokens=%s, kv_transfer_params=%s", + "req_id=%s num_external_tokens=%s, kv_transfer_params=%s", + request.request_id, num_external_tokens, params, ) @@ -274,8 +325,11 @@ class MooncakeConnectorScheduler: return if params.get("do_remote_prefill"): - assert self.kv_role != "kv_producer" - if all(p in params for p in ("remote_host", "remote_port")): + assert not self.is_kv_producer + if all( + p in params + for p in ("remote_engine_id", "remote_bootstrap_addr", "transfer_id") + ): # If remote_blocks and num_external_tokens = 0, we have # a full prefix cache hit on the D worker. We need to call # send_notif in _read_blocks to free the memory on the P. @@ -294,8 +348,12 @@ class MooncakeConnectorScheduler: params["do_remote_prefill"] = False elif params.get("do_remote_decode"): - # Add an empty list to worker to create event. - self._reqs_need_send[request.request_id] = [] + assert not self.is_kv_consumer + if not params.get("transfer_id"): + logger.warning("Missing transfer_id in kv_transfer_params from router!") + else: + # Add an empty list to worker to create event. + self._reqs_need_send[request.request_id] = (request, []) def build_connector_meta( self, @@ -303,8 +361,8 @@ class MooncakeConnectorScheduler: ) -> KVConnectorMetadata: meta = MooncakeConnectorMetadata() - # Loop through scheduled reqs and convert to RecvReqMeta. - if self.kv_role != "kv_producer": + # Loop through scheduled reqs and convert to PullReqMeta. + if not self.is_kv_producer: for req_id, (req, block_ids) in self._reqs_need_recv.items(): assert req.kv_transfer_params is not None meta.add_new_req( @@ -314,15 +372,18 @@ class MooncakeConnectorScheduler: ) self._reqs_need_recv.clear() - if self.kv_role != "kv_consumer": - for req_id, block_ids in self._reqs_need_send.items(): + if not self.is_kv_consumer: + for req_id, (req, block_ids) in self._reqs_need_send.items(): + assert req.kv_transfer_params is not None meta.add_new_req( request_id=req_id, local_block_ids=block_ids, - kv_transfer_params={}, + kv_transfer_params=req.kv_transfer_params, load_remote_cache=False, ) self._reqs_need_send.clear() + meta.reqs_not_processed = self._reqs_not_processed + self._reqs_not_processed = set() return meta @@ -338,12 +399,13 @@ class MooncakeConnectorScheduler: params = request.kv_transfer_params logger.debug( - "MooncakeConnector request_finished, request_status=%s, " + "MooncakeConnector request_finished, req_id=%s, request_status=%s, " "kv_transfer_params=%s", + request.request_id, request.status, params, ) - if not params: + if not params or not params.get("transfer_id"): return False, None if params.get("do_remote_prefill"): @@ -353,32 +415,30 @@ class MooncakeConnectorScheduler: # To avoid stranding the prefill blocks in the prefill instance, # we must add empty block_ids to _reqs_need_recv so that our # worker side will notify and free blocks in the prefill instance. - assert self.kv_role != "kv_producer" + assert not self.is_kv_producer self._reqs_need_recv[request.request_id] = (request, []) params["do_remote_prefill"] = False return False, None - if ( - not params.get("do_remote_decode") - or request.status != RequestStatus.FINISHED_LENGTH_CAPPED - ): + if not params.get("do_remote_decode"): return False, None - assert self.kv_role != "kv_consumer" + assert not self.is_kv_consumer + + if request.status != RequestStatus.FINISHED_LENGTH_CAPPED: + # Also include the case of a P/D Prefill request with immediate + # block free (eg abort). Stop tracking this request. + self._reqs_not_processed.add(params["transfer_id"]) + return False, None # TODO: check whether block_ids actually ever be 0. If not we could # remove the conditional below delay_free_blocks = len(block_ids) > 0 if delay_free_blocks: - self._reqs_need_send[request.request_id] = block_ids + self._reqs_need_send[request.request_id] = (request, block_ids) - return delay_free_blocks, dict( - do_remote_prefill=True, - do_remote_decode=False, - remote_host=self.side_channel_host, - remote_port=self.side_channel_port, - ) + return delay_free_blocks, None class MooncakeConnectorWorker: @@ -391,7 +451,18 @@ class MooncakeConnectorWorker: self.engine = TransferEngine() self.hostname = get_ip() - protocol = self.vllm_config.kv_transfer_config.kv_connector_extra_config.get( # type: ignore[union-attr] + + assert (kv_transfer_config := vllm_config.kv_transfer_config) + self.is_kv_producer: bool = kv_transfer_config.kv_role == "kv_producer" + self.is_kv_consumer: bool = kv_transfer_config.kv_role == "kv_consumer" + self.num_sender_workers = kv_transfer_config.kv_connector_extra_config.get( + "num_workers", 10 + ) + # Create more tasks than workers to keep the thread pool saturated. + # Tasks can await async events, so a surplus (2x is a robust heuristic) + # prevents workers from idling. + self.num_sender_tasks = self.num_sender_workers * 2 + protocol = kv_transfer_config.kv_connector_extra_config.get( # type: ignore[union-attr] "mooncake_protocol", "rdma" ) logger.info( @@ -409,33 +480,31 @@ class MooncakeConnectorWorker: self.rpc_port, ) - # Mooncake handshake port. - self.side_channel_port: int = get_mooncake_side_channel_port(vllm_config) - + self._remote_agents: dict[EngineId, dict[int, dict[int, str]]] = {} + self._pending_bootstrap_querys: dict[str, asyncio.Event] = {} + self.side_channel_port: int = 0 # we will bind it in register_kv_caches() self.engine_id: EngineId = engine_id self.tp_rank = get_tensor_model_parallel_rank() - self.world_size = get_tensor_model_parallel_world_size() - self.tp_group = get_tp_group() + self.tp_size = get_tensor_model_parallel_world_size() self.num_blocks = 0 - assert vllm_config.kv_transfer_config - self.kv_role = vllm_config.kv_transfer_config.kv_role - self.num_sender_workers = ( - vllm_config.kv_transfer_config.kv_connector_extra_config.get( - "num_workers", 10 + assert (parallel_config := vllm_config.parallel_config) + dp_rank = parallel_config.data_parallel_index + dp_local_rank = parallel_config.data_parallel_rank_local + self.dp_rank = dp_local_rank if parallel_config.local_engines_only else dp_rank + pp_size = vllm_config.parallel_config.pipeline_parallel_size + if pp_size > 1: + raise ValueError( + "Mooncake Transfer Engine does not support pipeline parallelism yet." ) - ) - # Create more tasks than workers to keep the thread pool saturated. - # Tasks can await async events, so a surplus (2x is a robust heuristic) - # prevents workers from idling. - self.num_sender_tasks = self.num_sender_workers * 2 + self.pp_rank = get_pp_group().rank_in_group self.kv_caches_base_addr: list[int] = [] self.device_kv_caches: dict[str, torch.Tensor] = {} - self.reqs_need_send: dict[ReqId, SendBlockMeta] = {} + self.reqs_need_send: dict[TransferId, SendBlockMeta] = {} # For kv_both, we will act both prefiller and decoder. - if self.kv_role != "kv_consumer": + if not self.is_kv_consumer: # Background threads for sending kvcaches to D. self._sender_executor = ThreadPoolExecutor( max_workers=self.num_sender_workers, @@ -454,7 +523,15 @@ class MooncakeConnectorWorker: ) self._sender_listener_t.start() - if self.kv_role != "kv_producer": + # Start bootstrap server on global rank 0. + if should_launch_bootstrap_server(vllm_config): + _, port = get_mooncake_bootstrap_addr(vllm_config) + self.bootstrap_server = MooncakeBootstrapServer( + vllm_config, "0.0.0.0", port + ) + self.bootstrap_server.start() + + if not self.is_kv_producer: self.receiver_loop = asyncio.new_event_loop() self._mooncake_receiver_t = threading.Thread( target=_async_loop, args=(self.receiver_loop,), daemon=True @@ -478,7 +555,7 @@ class MooncakeConnectorWorker: logger.debug("Detected attention backend %s", self.backend_name) logger.debug("Detected kv cache layout %s", self.kv_cache_layout) - self._tp_size: dict[EngineId, int] = {self.engine_id: self.world_size} + self._tp_size: dict[EngineId, int] = {self.engine_id: self.tp_size} self._block_size: dict[EngineId, int] = {self.engine_id: self.block_size} self.kv_topo = TpKVTopology( tp_rank=self.tp_rank, @@ -492,7 +569,8 @@ class MooncakeConnectorWorker: self.async_zmq_ctx = zmq.asyncio.Context() self._encoder = msgspec.msgpack.Encoder() - self._decoder = msgspec.msgpack.Decoder(MooncakeAgentMetadata) + self._xfer_meta_decoder = msgspec.msgpack.Decoder(MooncakeXferMetadata) + self._xfer_resp_decoder = msgspec.msgpack.Decoder(MooncakeXferResponse) def __del__(self): self.shutdown() @@ -500,26 +578,62 @@ class MooncakeConnectorWorker: def shutdown(self): """Cleanup background threads on destruction.""" self.async_zmq_ctx.term() - if self.kv_role != "kv_consumer": + if not self.is_kv_consumer: self._sender_executor.shutdown(wait=False) if self.sender_loop.is_running(): self.sender_loop.call_soon_threadsafe(self.sender_loop.stop) self._sender_listener_t.join() - if self.kv_role != "kv_producer" and self.receiver_loop.is_running(): + if should_launch_bootstrap_server(self.vllm_config): + self.bootstrap_server.shutdown() + if not self.is_kv_producer and self.receiver_loop.is_running(): self.receiver_loop.call_soon_threadsafe(self.receiver_loop.stop) self._mooncake_receiver_t.join() - async def _mooncake_sender_listener( - self, ready_event: threading.Event, base_port: int, tp_rank: int - ): + async def register_worker_with_bootstrap(self): + host, port = get_mooncake_bootstrap_addr(self.vllm_config) + url = make_zmq_path("http", host, port) + "/register" + worker_addr = make_zmq_path("tcp", self.hostname, self.side_channel_port) + payload = RegisterWorkerPayload( + engine_id=self.engine_id, + dp_rank=self.dp_rank, + tp_rank=self.tp_rank, + pp_rank=self.pp_rank, + addr=worker_addr, + ) + while True: + try: + async with httpx.AsyncClient() as client: + response = await client.post(url, json=payload.model_dump()) + response.raise_for_status() + logger.debug("Successfully registered with bootstrap server at %s", url) + break + except httpx.ConnectError: + # Bootstrap server not ready, wait for a while and retry. + await asyncio.sleep(1) + except Exception as e: + err_msg = ( + e.response.text if isinstance(e, httpx.HTTPStatusError) else str(e) + ) + logger.error( + "Error registering %s with bootstrap server: %s", payload, err_msg + ) + raise e + + async def _mooncake_sender_listener(self, ready_event: threading.Event): """ Background thread that listens for Mooncake requests, dispatches them to a thread pool, and sends acknowledgments upon completion. """ - path = make_zmq_path("tcp", self.hostname, base_port + tp_rank) - sock = make_zmq_socket(self.async_zmq_ctx, path, zmq.ROUTER) - logger.debug("Mooncake sender starting listening on path: %s", path) + sock = self.async_zmq_ctx.socket(zmq.ROUTER) + self.side_channel_port = sock.bind_to_random_port(f"tcp://{self.hostname}") + logger.debug( + "Mooncake sender starting listening on path: tcp://%s:%d", + self.hostname, + self.side_channel_port, + ) + + await self.register_worker_with_bootstrap() # Create async worker tasks that process items from the queue sender_tasks = [ @@ -531,7 +645,7 @@ class MooncakeConnectorWorker: try: while True: - identity, _, metadata_bytes = await sock.recv_multipart() + identity, metadata_bytes = await sock.recv_multipart() await self.sender_worker_queue.put((identity, metadata_bytes)) except zmq.ContextTerminated: logger.debug("ZMQ context terminated, exiting Mooncake sender thread.") @@ -549,12 +663,16 @@ class MooncakeConnectorWorker: try: identity, metadata_bytes = await self.sender_worker_queue.get() try: - metadata = self._decoder.decode(metadata_bytes) - await self.send_kv_to_decode(metadata) - await sock.send_multipart((identity, b"", TRANS_DONE)) + metadata = self._xfer_meta_decoder.decode(metadata_bytes) + await self.send_kv_to_decode(identity, sock, metadata) except Exception as e: logger.error("Error processing Mooncake xfer request: %s", e) - await sock.send_multipart((identity, b"", TRANS_ERROR)) + error_response = MooncakeXferResponse( + status=MooncakeXferResponseStatus.ERROR, err_msg=str(e) + ) + await sock.send_multipart( + (identity, self._encoder.encode(error_response)) + ) finally: self.sender_worker_queue.task_done() except asyncio.CancelledError: @@ -562,55 +680,169 @@ class MooncakeConnectorWorker: except Exception as e: logger.error("Error in _sender_worker: %s", e) - async def send_kv_to_decode(self, meta: MooncakeAgentMetadata): - send_reqs: list[tuple[ReqId, SendBlockMeta]] = [] - for req_id in meta.request_ids: - send_meta = self.reqs_need_send.get(req_id) - if send_meta is None: - logger.warning("Request %s not found in reqs_need_send", req_id) - return - # Mark it as not expired. We will send it now. - send_meta.expire_time = float("inf") - send_reqs.append((req_id, send_meta)) + async def send_kv_to_decode( + self, identity: bytes, sock: zmq.asyncio.Socket, meta: MooncakeXferMetadata + ): + pending_reqs: dict[ReqId, SendBlockMeta] = {} + remote_tp_ranks = self.kv_topo.get_target_remote_ranks(meta.remote_tp_size) + if self.tp_rank not in remote_tp_ranks: + # This D worker does not pair with the P worker. + msg = f"This P tp_rank {self.tp_rank} not in remote D target ranks {remote_tp_ranks}" # noqa: E501 + logger.error(msg) + response = MooncakeXferResponse( + status=MooncakeXferResponseStatus.ERROR, + err_msg=msg, + ) + await sock.send_multipart((identity, self._encoder.encode(response))) + return + for d_req_id, (transfer_id, _) in meta.req_blocks.items(): + if transfer_id not in self.reqs_need_send: + # This req is not enqueued in P side yet, create it here. + self.reqs_need_send[transfer_id] = SendBlockMeta( + p_req_id="", + transfer_id=transfer_id, + local_block_ids=[], + ready=asyncio.Event(), + ) + send_meta = self.reqs_need_send[transfer_id] + pending_reqs[d_req_id] = send_meta - src_ptrs, dst_ptrs, lengths = await self._build_transfer_params(send_reqs, meta) - remote_session = f"{meta.remote_hostname}:{meta.remote_port}" - ret_value = await self.sender_loop.run_in_executor( - self._sender_executor, - self._send_blocks, - remote_session, - src_ptrs, - dst_ptrs, - lengths, - ) + async def wait_and_ret( + d_req_id: ReqId, send_meta: SendBlockMeta + ) -> tuple[ReqId, SendBlockMeta]: + await send_meta.ready.wait() + return d_req_id, send_meta - if ret_value != 0: - raise RuntimeError(f"Error in batch_transfer_sync_write: {ret_value}") + wait_tasks = [ + asyncio.create_task(wait_and_ret(d_req_id, send_meta)) + for d_req_id, send_meta in pending_reqs.items() + ] - for req_id in meta.request_ids: - del self.reqs_need_send[req_id] + while wait_tasks: + done, pending = await asyncio.wait( + wait_tasks, + timeout=envs.VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT, + return_when=asyncio.FIRST_COMPLETED, + ) - self.finished_sending_reqs.update(meta.request_ids) + if not done: + # Timeout, abort all pending requests. + for task in wait_tasks: + task.cancel() + logger.warning( + "Timeout waiting for P side ready: %s", list(pending_reqs) + ) + response = MooncakeXferResponse( + status=MooncakeXferResponseStatus.FINISH, + err_reqs=list(pending_reqs), + err_msg="Timeout waiting for P side ready.", + ) + await sock.send_multipart((identity, self._encoder.encode(response))) + break + + wait_tasks = list(pending) + response_status = ( + MooncakeXferResponseStatus.CONTINUE + if wait_tasks + else MooncakeXferResponseStatus.FINISH + ) + ready_reqs: list[tuple[ReqId, SendBlockMeta]] = [] + for task in done: + d_req_id, send_meta = task.result() + del pending_reqs[d_req_id] + # Do we still in reqs_need_send (not expired)? + if send_meta.transfer_id in self.reqs_need_send: + # Mark it sending to avoid expiration. + send_meta.sending += 1 + if not send_meta.need_send: + self.resolve_need_send(send_meta, remote_tp_ranks) + ready_reqs.append((d_req_id, send_meta)) + else: + # Otherwise (expired, very unlikely), just forget it. + logger.warning( + "Request %s expired before sending on P side.", d_req_id + ) + + src_ptrs, dst_ptrs, lengths, err_reqs = await self._build_transfer_params( + ready_reqs, meta + ) + + if err_reqs: + response = MooncakeXferResponse( + status=response_status, + err_reqs=err_reqs, + err_msg="P num blocks less than D", + ) + await sock.send_multipart((identity, self._encoder.encode(response))) + + if src_ptrs: + remote_session = f"{meta.remote_hostname}:{meta.remote_port}" + ret_value = await self.sender_loop.run_in_executor( + self._sender_executor, + self._send_blocks, + remote_session, + src_ptrs, + dst_ptrs, + lengths, + ) + + if ret_value != 0: + err_reqs = [] + for d_req_id, send_meta in ready_reqs: + send_meta.sending -= 1 + err_reqs.append(d_req_id) + # Do best effort to transfer the remaining reqs. + response = MooncakeXferResponse( + status=response_status, + err_reqs=err_reqs, + err_msg=f"Mooncake transfer engine returned {ret_value}", + ) + await sock.send_multipart( + (identity, self._encoder.encode(response)) + ) + continue + + for d_req_id, send_meta in ready_reqs: + # TODO: for heterogeneous TP (one P pairs to multiple D), + # we need to check whether all headers are sent. + # If not, we should set expire_time to normal and skip the below. + send_meta.sending -= 1 + send_meta.sent += 1 + if send_meta.sent == send_meta.need_send: + del self.reqs_need_send[send_meta.transfer_id] + self.finished_sending_reqs.add(send_meta.p_req_id) + + response = MooncakeXferResponse( + status=response_status, + ok_reqs=[d_req_id for d_req_id, _ in ready_reqs], + ) + await sock.send_multipart((identity, self._encoder.encode(response))) + + def resolve_need_send(self, send_meta: SendBlockMeta, remote_tp_ranks: list[int]): + # Prepare for heterogeneous TP (one P pairs to multiple D) + send_meta.need_send = len(remote_tp_ranks) + if send_meta.need_send != 1: + logger.error("Mooncake: Heterogeneous TP is not supported yet.") + raise NotImplementedError( + "Mooncake: Heterogeneous TP is not supported yet." + ) async def _build_transfer_params( self, - send_reqs: list[tuple[ReqId, SendBlockMeta]], - agent_meta: MooncakeAgentMetadata, - ) -> tuple[list[int], list[int], list[int]]: + ready_reqs: list[tuple[ReqId, SendBlockMeta]], + agent_meta: MooncakeXferMetadata, + ) -> tuple[list[int], list[int], list[int], list[ReqId]]: src_ptrs = [] dst_ptrs = [] lengths = [] + err_reqs: list[ReqId] = [] local_base_addr = self.kv_caches_base_addr remote_base_addr = agent_meta.kv_caches_base_addr block_len = self.block_len remote_session = f"{agent_meta.remote_hostname}:{agent_meta.remote_port}" - assert len(send_reqs) == len(agent_meta.block_ids) - for (req_id, send_meta), remote_block_ids in zip( - send_reqs, agent_meta.block_ids - ): - await send_meta.ready.wait() - + for d_req_id, send_meta in ready_reqs: + _, remote_block_ids = agent_meta.req_blocks[d_req_id] num_remote_blocks = len(remote_block_ids) if num_remote_blocks == 0: continue @@ -618,7 +850,15 @@ class MooncakeConnectorWorker: local_block_ids = send_meta.local_block_ids # Partial prefix cache hit: just read uncomputed blocks. num_local_blocks = len(local_block_ids) - assert num_local_blocks >= num_remote_blocks + if num_local_blocks < num_remote_blocks: + logger.error( + "req %s: local blocks(%d) less than remote blocks(%d)!", + d_req_id, + num_local_blocks, + num_remote_blocks, + ) + err_reqs.append(d_req_id) + continue if num_local_blocks > num_remote_blocks: local_block_ids = local_block_ids[-num_remote_blocks:] @@ -643,12 +883,12 @@ class MooncakeConnectorWorker: logger.debug( "Sending kv_caches for request %s (%d blocks) to %s", - req_id, + d_req_id, num_remote_blocks, remote_session, ) - return src_ptrs, dst_ptrs, lengths + return src_ptrs, dst_ptrs, lengths, err_reqs def _send_blocks( self, @@ -722,15 +962,12 @@ class MooncakeConnectorWorker: ) # No need to launch server for D node. - if self.kv_role == "kv_consumer": + if self.is_kv_consumer: return ready_event = threading.Event() asyncio.run_coroutine_threadsafe( - self._mooncake_sender_listener( - ready_event, self.side_channel_port, self.tp_rank - ), - self.sender_loop, + self._mooncake_sender_listener(ready_event), self.sender_loop ) ready_event.wait() # Wait for listener ZMQ socket to be ready. @@ -745,21 +982,25 @@ class MooncakeConnectorWorker: # Handle timeout to avoid stranding blocks on remote. now = time.perf_counter() - expired_reqs = [ - req_id - for req_id, send_meta in self.reqs_need_send.items() - if send_meta.expire_time < now - ] - for req_id in expired_reqs: - logger.warning( - "Request %s timed out after %d seconds without " - "being sent. Freeing its blocks on the producer side.", - req_id, - envs.VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT, - ) - del self.reqs_need_send[req_id] - if expired_reqs: - finished_sending_reqs.update(expired_reqs) + + expired_transfer_id = [] + for transfer_id, send_meta in self.reqs_need_send.items(): + if ( + send_meta.p_req_id + and send_meta.expire_time < now + and send_meta.sending == 0 + ): + logger.warning( + "Request %s timed out after %d seconds without " + "being sent. Freeing its blocks on the producer side.", + send_meta.p_req_id, + envs.VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT, + ) + finished_sending_reqs.add(send_meta.p_req_id) + expired_transfer_id.append(transfer_id) + + for transfer_id in expired_transfer_id: + del self.reqs_need_send[transfer_id] return finished_sending_reqs @@ -771,12 +1012,12 @@ class MooncakeConnectorWorker: """ recv_fut = None send_fut = None - if self.kv_role != "kv_producer": + if not self.is_kv_producer: recv_fut = asyncio.run_coroutine_threadsafe( self.fetch_finished_recving_reqs(), self.receiver_loop ) - if self.kv_role != "kv_consumer": + if not self.is_kv_consumer: send_fut = asyncio.run_coroutine_threadsafe( self.fetch_finished_sending_reqs(), self.sender_loop ) @@ -795,69 +1036,174 @@ class MooncakeConnectorWorker: return finished_sending_reqs or None, finished_recving_reqs or None - async def receive_kv(self, path: str, req_blocks: list[tuple[str, list[int]]]): - req_ids, block_ids = map(list, zip(*req_blocks)) - metadata = MooncakeAgentMetadata( + async def receive_kv_from_single_worker( + self, + worker_addr: str, + pull_metas: dict[ReqId, PullReqMeta], + ): + req_ids = set(pull_metas) + metadata = MooncakeXferMetadata( remote_hostname=self.hostname, remote_port=self.rpc_port, - request_ids=req_ids, + remote_tp_size=self.tp_size, + remote_tp_rank=self.tp_rank, + req_blocks={ + req_id: (pull_meta.transfer_id, pull_meta.local_block_ids) + for req_id, pull_meta in pull_metas.items() + }, kv_caches_base_addr=self.kv_caches_base_addr, - block_ids=block_ids, ) encoded_data = self._encoder.encode(metadata) logger.debug( - "Size of encoded MooncakeAgentMetadata: %d bytes", len(encoded_data) + "Size of encoded MooncakeXferMetadata: %d bytes", len(encoded_data) + ) + logger.debug( + "Sending kv transfer request for %s on path: %s", req_ids, worker_addr ) - logger.debug("Sending kv transfer request for %s on path: %s", req_ids, path) # Send query for the request. - sock: zmq.asyncio.Socket = make_zmq_socket( - self.async_zmq_ctx, path, zmq.REQ, bind=False, linger=0 - ) - sock.setsockopt(zmq.RCVTIMEO, 60000) try: - await sock.send(encoded_data) - ret_msg = await sock.recv() - if ret_msg != TRANS_DONE: - logger.error( - "Error happens during tranfering kvcache for %s, see logs in prefiller.", # noqa: E501 - req_ids, + with make_zmq_socket( + self.async_zmq_ctx, worker_addr, zmq.DEALER, bind=False, linger=0 + ) as sock: + # If something goes wrong, let P wait timeout first (in asyncio.wait()). + sock.setsockopt( + zmq.RCVTIMEO, (envs.VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT + 60) * 1000 ) - return + await sock.send(encoded_data) + while True: + ret_msg = await sock.recv() + response = self._xfer_resp_decoder.decode(ret_msg) + if response.status == MooncakeXferResponseStatus.ERROR: + logger.error( + "Error happens during tranfering kvcache for %s: %s", + req_ids, + response.err_msg, + ) + return + self.process_pulling_result(response, pull_metas) + if response.status == MooncakeXferResponseStatus.FINISH: + break except zmq.ContextTerminated: logger.debug("ZMQ context terminated, exiting Mooncake receiver thread.") except Exception as e: - logger.error("MooncakeAgentMetadata transfer failed for %s: %s", req_ids, e) + logger.error("MooncakeXferMetadata transfer failed for %s: %s", req_ids, e) return - finally: - sock.close() - self.finished_recving_reqs.update(req_ids) + def process_pulling_result( + self, + response: MooncakeXferResponse, + pull_metas: dict[ReqId, PullReqMeta], + ): + ok_reqs: list[ReqId] = response.ok_reqs or [] - logger.debug("pulling kv_caches for %s finished", req_ids) + for req_id in ok_reqs: + pull_meta = pull_metas[req_id] + # No race because we are in async loop. + pull_meta.pull_tasks_count -= 1 + if pull_meta.pull_tasks_count == 0: + self.finished_recving_reqs.add(pull_meta.d_req_id) - def group_kv_pull(self, metadata: MooncakeConnectorMetadata): - kv_pulls = defaultdict(list) - for req_id, meta in metadata.reqs_to_recv.items(): - logger.debug( - "start_load_kv for request %s from remote engine. " - "Num local_block_ids: %s.", - req_id, - len(meta.local_block_ids), + if ok_reqs: + logger.debug("pulling kv_caches for %s finished", ok_reqs) + + if response.err_reqs: + logger.error( + "pulling kv_caches for %s failed: %s", + response.err_reqs, + response.err_msg, ) - path = make_zmq_path( - "tcp", meta.remote_host, meta.remote_port + self.tp_rank - ) - kv_pulls[path].append((req_id, meta.local_block_ids)) - return kv_pulls + async def _connect_to_prefiller_bootstrap(self, remote_bootstrap_addr: str): + url = remote_bootstrap_addr + "/query" + try: + async with httpx.AsyncClient() as client: + response = await client.get(url) + response.raise_for_status() + data: dict = response.json() + for _, dp_entry in data.items(): + remote_engine_id = dp_entry["engine_id"] + self._remote_agents[remote_engine_id] = { + int(tp_rank): { + int(pp_rank): worker_addr + for pp_rank, worker_addr in tp_entry.items() + } + for tp_rank, tp_entry in dp_entry["worker_addr"].items() + } + self._tp_size[remote_engine_id] = len(dp_entry["worker_addr"]) + except Exception as e: + logger.error( + "Failed to connect to bootstrap server %s: %s", + remote_bootstrap_addr, + e, + ) + + # Always notify others regardless of connection success or failure. + self._pending_bootstrap_querys[remote_bootstrap_addr].set() + del self._pending_bootstrap_querys[remote_bootstrap_addr] + + def receive_kv( + self, + remote_engine_id: EngineId, + pull_metas: dict[ReqId, PullReqMeta], + ): + remote_tp_ranks = self.kv_topo.get_target_remote_ranks_from_engine_id( + remote_engine_id + ) + count = len(remote_tp_ranks) + if count != 1: + logger.error("Mooncake: Heterogeneous TP is not supported yet.") + raise NotImplementedError( + "Mooncake: Heterogeneous TP is not supported yet." + ) + for pull_meta in pull_metas.values(): + pull_meta.pull_tasks_count = count + for remote_tp_rank in remote_tp_ranks: + worker_addr = self._remote_agents[remote_engine_id][remote_tp_rank][0] + asyncio.create_task( + self.receive_kv_from_single_worker(worker_addr, pull_metas) + ) + + async def handle_new_engine_id( + self, + remote_engine_id: EngineId, + pull_metas: dict[ReqId, PullReqMeta], + ): + remote_bootstrap_addr = next(iter(pull_metas.values())).remote_bootstrap_addr + if remote_bootstrap_addr not in self._pending_bootstrap_querys: + self._pending_bootstrap_querys[remote_bootstrap_addr] = asyncio.Event() + await self._connect_to_prefiller_bootstrap(remote_bootstrap_addr) + else: + await self._pending_bootstrap_querys[remote_bootstrap_addr].wait() + + if remote_engine_id not in self._remote_agents: + logger.error( + "Failed to find remote engine_id %s from bootstrap server %s", + remote_engine_id, + remote_bootstrap_addr, + ) + return + + self.receive_kv(remote_engine_id, pull_metas) + + async def _start_load_kv( + self, reqs_to_recv: dict[EngineId, dict[ReqId, PullReqMeta]] + ): + for remote_engine_id, pull_metas in reqs_to_recv.items(): + if remote_engine_id not in self._remote_agents: + asyncio.create_task( + self.handle_new_engine_id(remote_engine_id, pull_metas) + ) + else: + self.receive_kv(remote_engine_id, pull_metas) async def record_send_reqs(self, metadata: MooncakeConnectorMetadata): - for req_id, block_ids in metadata.reqs_to_send.items(): + for p_req_id, (transfer_id, block_ids) in metadata.reqs_to_send.items(): if block_ids: # Already gone through request_finished() - send_meta = self.reqs_need_send[req_id] + send_meta = self.reqs_need_send[transfer_id] + send_meta.p_req_id = p_req_id send_meta.local_block_ids = block_ids send_meta.expire_time = ( time.perf_counter() + envs.VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT @@ -866,20 +1212,29 @@ class MooncakeConnectorWorker: else: # From update_state_after_alloc(), # but not reach request_finished() yet - self.reqs_need_send[req_id] = SendBlockMeta( - local_block_ids=[], - ready=asyncio.Event(), - ) + # This may be already created by send_kv_to_decode() + # when D is sending MooncakeXferMetadata. + if transfer_id not in self.reqs_need_send: + self.reqs_need_send[transfer_id] = SendBlockMeta( + p_req_id=p_req_id, + transfer_id=transfer_id, + local_block_ids=[], + ready=asyncio.Event(), + ) + for transfer_id in metadata.reqs_not_processed: + send_meta = self.reqs_need_send.pop(transfer_id) + if send_meta: + assert not send_meta.ready.is_set() def start_load_kv(self, metadata: MooncakeConnectorMetadata): - if self.kv_role != "kv_producer": - kv_pulls = self.group_kv_pull(metadata) - for path, req_blocks in kv_pulls.items(): - asyncio.run_coroutine_threadsafe( - self.receive_kv(path, req_blocks), self.receiver_loop - ) + if not self.is_kv_producer and metadata.reqs_to_recv: + asyncio.run_coroutine_threadsafe( + self._start_load_kv(metadata.reqs_to_recv), self.receiver_loop + ) - if self.kv_role != "kv_consumer": + if not self.is_kv_consumer and ( + metadata.reqs_to_send or metadata.reqs_not_processed + ): asyncio.run_coroutine_threadsafe( self.record_send_reqs(metadata), self.sender_loop ) @@ -914,3 +1269,31 @@ def get_mooncake_side_channel_port(vllm_config: VllmConfig) -> int: def _async_loop(loop: asyncio.AbstractEventLoop): asyncio.set_event_loop(loop) loop.run_forever() + + +def should_launch_bootstrap_server(vllm_config: VllmConfig) -> bool: + assert (parallel_config := vllm_config.parallel_config) + # In hybrid or external LB mode, + # each instance should have its own bootstrap server. + # + # In internal LB mode, + # only the real global first rank need to launch the bootstrap server. + return is_local_first_rank() and ( + parallel_config.local_engines_only or parallel_config.data_parallel_index == 0 + ) + + +def get_mooncake_bootstrap_addr(vllm_config: VllmConfig) -> tuple[str, int]: + """ + Returns the address of the Mooncake bootstrap server. + This is only used by prefillers to register workers. + Decoders should get addr from kv_transfer_params. + """ + assert (parallel_config := vllm_config.parallel_config) + if parallel_config.local_engines_only: + # In hybrid or external LB mode, connect to local server. + host = "127.0.0.1" + else: + host = parallel_config.data_parallel_master_ip + port = envs.VLLM_MOONCAKE_BOOTSTRAP_PORT + return (host, port) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_utils.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_utils.py new file mode 100644 index 00000000000..d1a9946709d --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_utils.py @@ -0,0 +1,127 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import threading +import time +from dataclasses import dataclass + +import uvicorn +from fastapi import FastAPI, HTTPException +from pydantic import BaseModel + +from vllm.config import VllmConfig +from vllm.distributed.kv_transfer.kv_connector.utils import EngineId +from vllm.logger import init_logger + +WorkerAddr = str + +logger = init_logger(__name__) + + +class RegisterWorkerPayload(BaseModel): + engine_id: EngineId + dp_rank: int + tp_rank: int + pp_rank: int + addr: WorkerAddr + + +@dataclass +class EngineEntry: + engine_id: EngineId + # {tp_rank: {pp_rank: worker_addr}} + worker_addr: dict[int, dict[int, WorkerAddr]] + + +class MooncakeBootstrapServer: + """ + A centralized server running on the global rank 0 prefiller worker. + Prefiller workers register their connection info (IP, port, ranks) here. + """ + + def __init__(self, vllm_config: VllmConfig, host: str, port: int): + self.workers: dict[int, EngineEntry] = {} + + self.host = host + self.port = port + self.app = FastAPI() + self._register_routes() + self.server_thread: threading.Thread | None = None + self.server: uvicorn.Server | None = None + + def __del__(self): + self.shutdown() + + def _register_routes(self): + # All methods are async. No need to use lock to protect data. + self.app.post("/register")(self.register_worker) + self.app.get("/query", response_model=dict[int, EngineEntry])(self.query) + + def start(self): + if self.server_thread: + return + + config = uvicorn.Config(app=self.app, host=self.host, port=self.port) + self.server = uvicorn.Server(config=config) + self.server_thread = threading.Thread( + target=self.server.run, name="mooncake_bootstrap_server", daemon=True + ) + self.server_thread.start() + while not self.server.started: + time.sleep(0.1) # Wait for the server to start + logger.info("Mooncake Bootstrap Server started at %s:%d", self.host, self.port) + + def shutdown(self): + if self.server_thread is None or self.server is None or not self.server.started: + return + + self.server.should_exit = True + self.server_thread.join() + logger.info("Mooncake Bootstrap Server stopped.") + + async def register_worker(self, payload: RegisterWorkerPayload): + """Handles registration of a prefiller worker.""" + if payload.dp_rank not in self.workers: + self.workers[payload.dp_rank] = EngineEntry( + engine_id=payload.engine_id, + worker_addr={}, + ) + + dp_entry = self.workers[payload.dp_rank] + if dp_entry.engine_id != payload.engine_id: + raise HTTPException( + status_code=400, + detail=( + f"Engine ID mismatch for dp_rank={payload.dp_rank}: " + f"expected {dp_entry.engine_id}, got {payload.engine_id}" + ), + ) + if payload.tp_rank not in dp_entry.worker_addr: + dp_entry.worker_addr[payload.tp_rank] = {} + + tp_entry = dp_entry.worker_addr[payload.tp_rank] + if payload.pp_rank in tp_entry: + raise HTTPException( + status_code=400, + detail=( + f"Worker with dp_rank={payload.dp_rank}, " + f"tp_rank={payload.tp_rank}, pp_rank={payload.pp_rank} " + f"is already registered at " + f"{tp_entry[payload.pp_rank]}, " + f"but still want to register at {payload.addr}" + ), + ) + + tp_entry[payload.pp_rank] = payload.addr + logger.debug( + "Registered worker: engine_id=%s, dp_rank=%d, tp_rank=%d, pp_rank=%d at %s", + payload.engine_id, + payload.dp_rank, + payload.tp_rank, + payload.pp_rank, + payload.addr, + ) + + return {"status": "ok"} + + async def query(self) -> dict[int, EngineEntry]: + return self.workers From 2267cb1cfd838a192f47ff677d91164fb5cb2862 Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Tue, 3 Feb 2026 09:08:47 -0700 Subject: [PATCH 028/810] [Attention][FA3] Update FA3 to include new swizzle optimization (#23465) Signed-off-by: Lucas Wilkinson --- cmake/external_projects/vllm_flash_attn.cmake | 2 +- vllm/v1/attention/backends/flash_attn.py | 7 ++++++- vllm/v1/attention/backends/mla/flashattn_mla.py | 7 ++++++- 3 files changed, 13 insertions(+), 3 deletions(-) diff --git a/cmake/external_projects/vllm_flash_attn.cmake b/cmake/external_projects/vllm_flash_attn.cmake index b51934a3ab2..dbdfd5e8144 100644 --- a/cmake/external_projects/vllm_flash_attn.cmake +++ b/cmake/external_projects/vllm_flash_attn.cmake @@ -38,7 +38,7 @@ else() FetchContent_Declare( vllm-flash-attn GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git - GIT_TAG 188be16520ceefdc625fdf71365585d2ee348fe2 + GIT_TAG 2adfc8c2177c5b0e8ddeedfd5a8990d80eb496ff 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/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index 9275725314e..232b0b0daff 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -308,10 +308,15 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad self.compilation_config.cudagraph_mode.has_full_cudagraphs() ) self.max_cudagraph_size = self.compilation_config.max_cudagraph_capture_size + max_num_seqs = vllm_config.scheduler_config.max_num_seqs if self.use_full_cuda_graph and self.aot_schedule: + # Times 4 due to: + # https://github.com/vllm-project/flash-attention/blob/3223650ccabe622a0fcae65eec706a50186a89f7/hopper/flash_api.cpp#L650-L653 + # For some tests max_cudagraph_size > max_num_seqs, + # so we need to use the larger one. self.scheduler_metadata = torch.zeros( - vllm_config.scheduler_config.max_num_seqs + 1, + max(self.max_cudagraph_size or 0, max_num_seqs) * 4 + 1, dtype=torch.int32, device=self.device, ) diff --git a/vllm/v1/attention/backends/mla/flashattn_mla.py b/vllm/v1/attention/backends/mla/flashattn_mla.py index e160d325568..f0ba259362f 100644 --- a/vllm/v1/attention/backends/mla/flashattn_mla.py +++ b/vllm/v1/attention/backends/mla/flashattn_mla.py @@ -127,10 +127,15 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata] self.compilation_config.cudagraph_mode.has_full_cudagraphs() ) self.max_cudagraph_size = self.compilation_config.max_cudagraph_capture_size + max_num_seqs = vllm_config.scheduler_config.max_num_seqs if self.use_full_cuda_graph and self.fa_aot_schedule: + # Times 4 due to: + # https://github.com/vllm-project/flash-attention/blob/3223650ccabe622a0fcae65eec706a50186a89f7/hopper/flash_api.cpp#L650-L653 + # For some tests max_cudagraph_size > max_num_seqs, + # so we need to use the larger one. self.scheduler_metadata = torch.zeros( - vllm_config.scheduler_config.max_num_seqs + 1, + max(self.max_cudagraph_size or 0, max_num_seqs) * 4 + 1, dtype=torch.int32, device=self.device, ) From b1bb18de8d12cac63e3bbb0f59b3726fcb68dc80 Mon Sep 17 00:00:00 2001 From: Richard Zou Date: Tue, 3 Feb 2026 09:12:11 -0800 Subject: [PATCH 029/810] [torch.compile] Significantly speed up cold start times (#33641) Signed-off-by: Richard Zou --- tests/compile/test_cold_start.py | 9 ++--- vllm/compilation/backends.py | 50 ++++++++++++++++++-------- vllm/compilation/compiler_interface.py | 3 -- 3 files changed, 41 insertions(+), 21 deletions(-) diff --git a/tests/compile/test_cold_start.py b/tests/compile/test_cold_start.py index 1d24d18397b..dd770c58364 100644 --- a/tests/compile/test_cold_start.py +++ b/tests/compile/test_cold_start.py @@ -37,12 +37,13 @@ def test_moe_compilation_cold_start(monkeypatch, use_fresh_inductor_cache): # The forward pass consists of 32 transformer layers. # Then, we split on the attention operation. This results in # 33 subgraphs (not including the attention operation). - # The 33 subgraphs then get standalone_compile'd. + # We then standalone_compile the unique subgraphs. # # There are actually only 3 unique subgraphs for this model # (all of its transformer layers are the same modulo weights); # this is true for most vLLM models. - # So we test that during cold start, the aot_autograd cache - # misses for 3 subgraphs and hits for the rest. + # So we test that during cold start, only 3 subgraphs are compiled + # These 3 subgraphs should cache miss, and then there should be + # no other compilation (so no cache hits). assert counters["aot_autograd"]["autograd_cache_miss"] == 3 - assert counters["aot_autograd"]["autograd_cache_hit"] == 30 + assert counters["aot_autograd"]["autograd_cache_hit"] == 0 diff --git a/vllm/compilation/backends.py b/vllm/compilation/backends.py index 89981fc2996..38ba97c7fae 100644 --- a/vllm/compilation/backends.py +++ b/vllm/compilation/backends.py @@ -121,7 +121,7 @@ class CompilerManager: and compiling the graph. The cache is a dict mapping - `(runtime_shape, graph_index, backend_name)` + `(runtime_shape, graph_hash, backend_name)` to `any_data` returned from the compiler. When serializing the cache, we save it to a Python file @@ -130,7 +130,7 @@ class CompilerManager: """ def __init__(self, compilation_config: CompilationConfig) -> None: - self.cache: dict[tuple[Range, int, str], Any] = dict() + self.cache: dict[tuple[Range, str, str], Any] = dict() self.is_cache_updated = False self.compilation_config = compilation_config self.compiler = make_compiler(compilation_config) @@ -173,6 +173,7 @@ class CompilerManager: self.disable_cache = disable_cache self.cache_dir = cache_dir self.cache_file_path = os.path.join(cache_dir, "vllm_compile_cache.py") + self.loaded_cache_entries: dict[tuple[Range, str, str], Any] = {} if not disable_cache and os.path.exists(self.cache_file_path): # load the cache from the file @@ -186,9 +187,9 @@ class CompilerManager: if not isinstance(value, ty): raise TypeError(f"Expected {ty} but got {type(value)} for {value}") - def parse_key(key: Any) -> tuple[Range, int, str]: - range_tuple, graph_index, compiler_name = key - check_type(graph_index, int) + def parse_key(key: Any) -> tuple[Range, str, str]: + range_tuple, graph_hash, compiler_name = key + check_type(graph_hash, str) check_type(compiler_name, str) if isinstance(range_tuple, tuple): start, end = range_tuple @@ -196,7 +197,7 @@ class CompilerManager: check_type(end, int) range_tuple = Range(start=start, end=end) check_type(range_tuple, Range) - return range_tuple, graph_index, compiler_name + return range_tuple, graph_hash, compiler_name self.cache = {parse_key(key): value for key, value in cache.items()} @@ -216,18 +217,25 @@ class CompilerManager: self, graph: fx.GraphModule, example_inputs: list[Any], - graph_index: int, + graph_hash: str, compile_range: Range, ) -> Callable[..., Any] | None: - if (compile_range, graph_index, self.compiler.name) not in self.cache: + key = (compile_range, graph_hash, self.compiler.name) + # See if we've already loaded this cache entry + if key in self.loaded_cache_entries: + return self.loaded_cache_entries[key] + # Otherwise, go load it from disk + if key not in self.cache: return None - handle = self.cache[(compile_range, graph_index, self.compiler.name)] + handle = self.cache[key] compiled_graph = self.compiler.load( - handle, graph, example_inputs, graph_index, compile_range + handle, graph, example_inputs, compile_range ) + self.loaded_cache_entries[key] = compiled_graph logger.debug( - "Directly load the %s-th graph for compile range %sfrom %s via handle %s", - graph_index, + "Directly load the graph (hash %s) for compile range " + "%sfrom %s via handle %s", + graph_hash, str(compile_range), self.compiler.name, handle, @@ -249,12 +257,22 @@ class CompilerManager: global compilation_start_time compilation_start_time = time.time() + from torch._functorch._aot_autograd.autograd_cache import ( + AOTAutogradCachePickler, + sanitize_gm_for_cache, + ) + + with sanitize_gm_for_cache(graph): + pickler = AOTAutogradCachePickler(graph) + dumped_graph = pickler.dumps(graph) + graph_hash = hashlib.sha256(dumped_graph).hexdigest() + compilation_counter.num_backend_compilations += 1 compiled_graph = None # try to load from the cache - compiled_graph = self.load(graph, example_inputs, graph_index, compile_range) + compiled_graph = self.load(graph, example_inputs, graph_hash, compile_range) if compiled_graph is not None: if graph_index == num_graphs - 1: # after loading the last graph for this shape, record the time. @@ -290,9 +308,13 @@ class CompilerManager: assert compiled_graph is not None, "Failed to compile the graph" + self.loaded_cache_entries[(compile_range, graph_hash, self.compiler.name)] = ( + compiled_graph + ) + # store the artifact in the cache if is_compile_cache_enabled(additional_inductor_config) and handle is not None: - self.cache[(compile_range, graph_index, self.compiler.name)] = handle + self.cache[(compile_range, graph_hash, self.compiler.name)] = handle compilation_counter.num_cache_entries_updated += 1 self.is_cache_updated = True if graph_index == 0: diff --git a/vllm/compilation/compiler_interface.py b/vllm/compilation/compiler_interface.py index 60650353971..875e628d686 100644 --- a/vllm/compilation/compiler_interface.py +++ b/vllm/compilation/compiler_interface.py @@ -101,7 +101,6 @@ class CompilerInterface: handle: Any, graph: fx.GraphModule, example_inputs: list[Any], - graph_index: int, compile_range: Range, ) -> Callable[..., Any]: """ @@ -302,7 +301,6 @@ class InductorStandaloneAdaptor(CompilerInterface): handle: Any, graph: fx.GraphModule, example_inputs: list[Any], - graph_index: int, compile_range: Range, ) -> Callable[..., Any]: assert isinstance(handle, tuple) @@ -527,7 +525,6 @@ class InductorAdaptor(CompilerInterface): handle: Any, graph: fx.GraphModule, example_inputs: list[Any], - graph_index: int, compile_range: Range, ) -> Callable[..., Any]: assert isinstance(handle, tuple) From 61e632aea15f76fd1c46354b00f9cac62cd28c4e Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Tue, 3 Feb 2026 17:40:59 +0000 Subject: [PATCH 030/810] Turn `@config` into a `dataclass_transform` (#31541) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- tests/engine/test_arg_utils.py | 26 ++-- tests/test_config.py | 5 +- tests/tools/test_config_validator.py | 26 ++-- tests/v1/e2e/test_spec_decode.py | 2 +- .../test_backend_guidance.py | 7 +- tools/pre_commit/validate_config.py | 26 ++-- vllm/config/__init__.py | 2 + vllm/config/attention.py | 2 - vllm/config/cache.py | 2 - vllm/config/compilation.py | 6 +- vllm/config/device.py | 4 +- vllm/config/ec_transfer.py | 3 - vllm/config/kv_events.py | 2 - vllm/config/kv_transfer.py | 3 - vllm/config/load.py | 2 - vllm/config/lora.py | 4 +- vllm/config/model.py | 4 +- vllm/config/multimodal.py | 1 - vllm/config/observability.py | 2 - vllm/config/parallel.py | 7 -- vllm/config/pooler.py | 3 - vllm/config/profiler.py | 2 - vllm/config/scheduler.py | 2 - vllm/config/speculative.py | 2 - vllm/config/speech_to_text.py | 3 - vllm/config/structured_outputs.py | 2 - vllm/config/utils.py | 111 ++++++++++++------ vllm/config/vllm.py | 19 +-- vllm/entrypoints/openai/cli_args.py | 3 - .../configs/speculators/algos.py | 17 +-- .../configs/speculators/base.py | 34 +++--- vllm/v1/spec_decode/draft_model.py | 10 +- 32 files changed, 153 insertions(+), 191 deletions(-) diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 2acb38bc9a1..d1986e0a44f 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -3,11 +3,11 @@ import json from argparse import ArgumentError -from contextlib import nullcontext -from dataclasses import dataclass, field +from contextlib import AbstractContextManager, nullcontext from typing import Annotated, Literal import pytest +from pydantic import Field from vllm.config import AttentionConfig, CompilationConfig, config from vllm.engine.arg_utils import ( @@ -96,7 +96,7 @@ def test_get_type(type_hints, type, expected): ], ) def test_literal_to_kwargs(type_hints, expected): - context = nullcontext() + context: AbstractContextManager[object] = nullcontext() if expected is Exception: context = pytest.raises(expected) with context: @@ -104,14 +104,12 @@ def test_literal_to_kwargs(type_hints, expected): @config -@dataclass class NestedConfig: field: int = 1 """field""" @config -@dataclass class DummyConfig: regular_bool: bool = True """Regular bool with default True""" @@ -119,23 +117,23 @@ class DummyConfig: """Optional bool with default None""" optional_literal: Literal["x", "y"] | None = None """Optional literal with default None""" - tuple_n: tuple[int, ...] = field(default_factory=lambda: (1, 2, 3)) + tuple_n: tuple[int, ...] = Field(default_factory=lambda: (1, 2, 3)) """Tuple with variable length""" - tuple_2: tuple[int, int] = field(default_factory=lambda: (1, 2)) + tuple_2: tuple[int, int] = Field(default_factory=lambda: (1, 2)) """Tuple with fixed length""" - list_n: list[int] = field(default_factory=lambda: [1, 2, 3]) + list_n: list[int] = Field(default_factory=lambda: [1, 2, 3]) """List with variable length""" - list_literal: list[Literal[1, 2]] = field(default_factory=list) + list_literal: list[Literal[1, 2]] = Field(default_factory=list) """List with literal choices""" - list_union: list[str | type[object]] = field(default_factory=list) + list_union: list[str | type[object]] = Field(default_factory=list) """List with union type""" - set_n: set[int] = field(default_factory=lambda: {1, 2, 3}) + set_n: set[int] = Field(default_factory=lambda: {1, 2, 3}) """Set with variable length""" literal_literal: Literal[Literal[1], Literal[2]] = 1 """Literal of literals with default 1""" - json_tip: dict = field(default_factory=dict) + json_tip: dict = Field(default_factory=dict) """Dict which will be JSON in CLI""" - nested_config: NestedConfig = field(default_factory=NestedConfig) + nested_config: NestedConfig = Field(default_factory=NestedConfig) """Nested config""" @@ -195,7 +193,7 @@ def test_get_kwargs(): json_tip = "Should either be a valid JSON string or JSON keys" assert json_tip in kwargs["json_tip"]["help"] # nested config should construct the nested config - assert kwargs["nested_config"]["type"]('{"field": 2}') == NestedConfig(2) + assert kwargs["nested_config"]["type"]('{"field": 2}') == NestedConfig(2) # type: ignore[call-arg] @pytest.mark.parametrize( diff --git a/tests/test_config.py b/tests/test_config.py index 1676598b164..f3c3003a00c 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -66,9 +66,6 @@ class _TestConfigFields: def test_get_field(): - with pytest.raises(ValueError): - get_field(_TestConfigFields, "a") - b = get_field(_TestConfigFields, "b") assert isinstance(b, Field) assert b.default is MISSING @@ -188,7 +185,7 @@ def test_get_pooling_config(): ) def test_get_pooling_config_from_args(): model_id = "sentence-transformers/all-MiniLM-L12-v2" - pooler_config = PoolerConfig(seq_pooling_type="CLS", normalize=True) + pooler_config = PoolerConfig(seq_pooling_type="CLS", use_activation=False) model_config = ModelConfig(model_id, pooler_config=pooler_config) assert asdict(model_config.pooler_config) == asdict(pooler_config) diff --git a/tests/tools/test_config_validator.py b/tests/tools/test_config_validator.py index d6104dc6d2e..e317bf91178 100644 --- a/tests/tools/test_config_validator.py +++ b/tests/tools/test_config_validator.py @@ -7,31 +7,22 @@ import pytest from tools.pre_commit.validate_config import validate_ast -_TestConfig1 = """ +_TestConfig1 = ''' @config class _TestConfig1: - pass -""" - -_TestConfig2 = ''' -@config -@dataclass -class _TestConfig2: a: int """docstring""" ''' -_TestConfig3 = """ +_TestConfig2 = """ @config -@dataclass -class _TestConfig3: +class _TestConfig2: a: int = 1 """ -_TestConfig4 = ''' +_TestConfig3 = ''' @config -@dataclass -class _TestConfig4: +class _TestConfig3: a: Union[Literal[1], Literal[2]] = 1 """docstring""" ''' @@ -40,10 +31,9 @@ class _TestConfig4: @pytest.mark.parametrize( ("test_config", "expected_error"), [ - (_TestConfig1, "must be a dataclass"), - (_TestConfig2, "must have a default"), - (_TestConfig3, "must have a docstring"), - (_TestConfig4, "must use a single Literal"), + (_TestConfig1, "must have a default"), + (_TestConfig2, "must have a docstring"), + (_TestConfig3, "must use a single Literal"), ], ) def test_config(test_config, expected_error): diff --git a/tests/v1/e2e/test_spec_decode.py b/tests/v1/e2e/test_spec_decode.py index 02e1529142b..4905a4120a2 100644 --- a/tests/v1/e2e/test_spec_decode.py +++ b/tests/v1/e2e/test_spec_decode.py @@ -766,8 +766,8 @@ def assert_draft_model_correctness(args: ArgsTest, enforce_eager: bool): "max_model_len": args.max_model_len, "enforce_eager": enforce_eager, "draft_tensor_parallel_size": args.draft_tensor_parallel_size, - "max_num_seqs": 100, # limit cudagraph capture runtime }, + max_num_seqs=100, # limit cudagraph capture runtime max_model_len=args.max_model_len, gpu_memory_utilization=args.gpu_memory_utilization, tensor_parallel_size=args.target_tensor_parallel_size, diff --git a/tests/v1/structured_output/test_backend_guidance.py b/tests/v1/structured_output/test_backend_guidance.py index 4c01560fc88..362f75c49d0 100644 --- a/tests/v1/structured_output/test_backend_guidance.py +++ b/tests/v1/structured_output/test_backend_guidance.py @@ -26,11 +26,8 @@ def test_backend_guidance_rollback_terminated(): # guidance backend. In that case we are in a stopped state, but # it should be reverted in case EOS is not accepted by the target # model. - vllm_config = VllmConfig( - decoding_config=StructuredOutputsConfig( - backend="guidance", - ) - ) + structured_outputs_config = StructuredOutputsConfig(backend="guidance") + vllm_config = VllmConfig(structured_outputs_config=structured_outputs_config) tokenizer = AutoTokenizer.from_pretrained(TOKENIZER) backend = GuidanceBackend( diff --git a/tools/pre_commit/validate_config.py b/tools/pre_commit/validate_config.py index fb6f0e6a928..7da32bc6b48 100644 --- a/tools/pre_commit/validate_config.py +++ b/tools/pre_commit/validate_config.py @@ -54,24 +54,18 @@ class ConfigValidator(ast.NodeVisitor): def __init__(self): ... def visit_ClassDef(self, node): - # Validate class with both @config and @dataclass decorators - decorators = [ - id - for d in node.decorator_list - if ( - isinstance(d, ast.Name) - and ((id := d.id) == "config" or id == "dataclass") - ) - or ( - isinstance(d, ast.Call) - and (isinstance(d.func, ast.Name) and (id := d.func.id) == "dataclass") - ) - ] + # Validate classes with a @config decorator + decorators = set() + for decorator in node.decorator_list: + if isinstance(decorator, ast.Call): + decorator = decorator.func + if isinstance(decorator, ast.Name) and decorator.id == "config": + decorators.add(decorator.id) - if set(decorators) == {"config", "dataclass"}: + if decorators == {"config"}: validate_class(node) - elif set(decorators) == {"config"}: - fail(f"Class {node.name} with config decorator must be a dataclass.", node) + elif "config" in decorators: + fail(f"config decorator for {node.name} should be used alone", node) self.generic_visit(node) diff --git a/vllm/config/__init__.py b/vllm/config/__init__.py index 7f6565053ee..b2044c6e1d0 100644 --- a/vllm/config/__init__.py +++ b/vllm/config/__init__.py @@ -36,6 +36,7 @@ from vllm.config.utils import ( config, get_attr_docs, is_init_field, + replace, update_config, ) from vllm.config.vllm import ( @@ -101,6 +102,7 @@ __all__ = [ "config", "get_attr_docs", "is_init_field", + "replace", "update_config", # From vllm.config.vllm "VllmConfig", diff --git a/vllm/config/attention.py b/vllm/config/attention.py index ee072fb1c86..9379b2878ba 100644 --- a/vllm/config/attention.py +++ b/vllm/config/attention.py @@ -4,14 +4,12 @@ from typing import Any, Literal from pydantic import field_validator -from pydantic.dataclasses import dataclass from vllm.config.utils import config from vllm.v1.attention.backends.registry import AttentionBackendEnum @config -@dataclass class AttentionConfig: """Configuration for attention mechanisms in vLLM.""" diff --git a/vllm/config/cache.py b/vllm/config/cache.py index abf10e21d40..bf121e544c8 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -6,7 +6,6 @@ from dataclasses import field from typing import TYPE_CHECKING, Any, Literal from pydantic import Field, SkipValidation, field_validator -from pydantic.dataclasses import dataclass from vllm.config.utils import config from vllm.logger import init_logger @@ -37,7 +36,6 @@ KVOffloadingBackend = Literal["native", "lmcache"] @config -@dataclass class CacheConfig: """Configuration for the KV cache.""" diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index 7a69629f707..556254a6501 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -8,8 +8,7 @@ from dataclasses import field from pathlib import Path from typing import TYPE_CHECKING, Any, ClassVar, Literal -from pydantic import ConfigDict, Field, TypeAdapter, field_validator -from pydantic.dataclasses import dataclass +from pydantic import Field, TypeAdapter, field_validator import vllm.envs as envs from vllm.compilation.inductor_pass import CallableInductorPass, InductorPass @@ -96,7 +95,6 @@ class CUDAGraphMode(enum.Enum): @config -@dataclass(config=ConfigDict(extra="forbid")) class PassConfig: """Configuration for custom Inductor passes. @@ -267,7 +265,6 @@ class DynamicShapesType(str, enum.Enum): @config -@dataclass(config=ConfigDict(extra="forbid")) class DynamicShapesConfig: """Configuration to control/debug torch compile dynamic shapes.""" @@ -311,7 +308,6 @@ class DynamicShapesConfig: @config -@dataclass(config=ConfigDict(extra="forbid")) class CompilationConfig: """Configuration for compilation. diff --git a/vllm/config/device.py b/vllm/config/device.py index 85662ddff76..c20e4d0f288 100644 --- a/vllm/config/device.py +++ b/vllm/config/device.py @@ -6,7 +6,6 @@ from typing import Any, Literal import torch from pydantic import ConfigDict, SkipValidation -from pydantic.dataclasses import dataclass from vllm.config.utils import config from vllm.utils.hashing import safe_hash @@ -14,8 +13,7 @@ from vllm.utils.hashing import safe_hash Device = Literal["auto", "cuda", "cpu", "tpu", "xpu"] -@config -@dataclass(config=ConfigDict(arbitrary_types_allowed=True)) +@config(config=ConfigDict(arbitrary_types_allowed=True)) class DeviceConfig: """Configuration for the device to use for vLLM execution.""" diff --git a/vllm/config/ec_transfer.py b/vllm/config/ec_transfer.py index d95236f818a..c7f56557f9b 100644 --- a/vllm/config/ec_transfer.py +++ b/vllm/config/ec_transfer.py @@ -5,8 +5,6 @@ import uuid from dataclasses import field from typing import Any, Literal, get_args -from pydantic.dataclasses import dataclass - from vllm.config.utils import config ECProducer = Literal["ec_producer"] @@ -15,7 +13,6 @@ ECRole = Literal[ECProducer, ECConsumer] @config -@dataclass class ECTransferConfig: """Configuration for distributed EC cache transfer.""" diff --git a/vllm/config/kv_events.py b/vllm/config/kv_events.py index ce46cc03c39..94da54c78a6 100644 --- a/vllm/config/kv_events.py +++ b/vllm/config/kv_events.py @@ -5,13 +5,11 @@ from typing import Literal from pydantic import Field -from pydantic.dataclasses import dataclass from vllm.config.utils import config @config -@dataclass class KVEventsConfig: """Configuration for KV event publishing.""" diff --git a/vllm/config/kv_transfer.py b/vllm/config/kv_transfer.py index 98cea821c67..fe3b218fbe9 100644 --- a/vllm/config/kv_transfer.py +++ b/vllm/config/kv_transfer.py @@ -5,8 +5,6 @@ import uuid from dataclasses import field from typing import Any, Literal, get_args -from pydantic.dataclasses import dataclass - from vllm.config.utils import config from vllm.utils.hashing import safe_hash @@ -16,7 +14,6 @@ KVRole = Literal[KVProducer, KVConsumer] @config -@dataclass class KVTransferConfig: """Configuration for distributed KV cache transfer.""" diff --git a/vllm/config/load.py b/vllm/config/load.py index 579a0bc3102..64a269e9885 100644 --- a/vllm/config/load.py +++ b/vllm/config/load.py @@ -4,7 +4,6 @@ from typing import TYPE_CHECKING, Any from pydantic import Field, field_validator -from pydantic.dataclasses import dataclass from vllm.config.utils import config from vllm.logger import init_logger @@ -21,7 +20,6 @@ logger = init_logger(__name__) @config -@dataclass class LoadConfig: """Configuration for loading the model weights.""" diff --git a/vllm/config/lora.py b/vllm/config/lora.py index f15beffe1df..0d310c87e50 100644 --- a/vllm/config/lora.py +++ b/vllm/config/lora.py @@ -5,7 +5,6 @@ from typing import TYPE_CHECKING, Any, Literal import torch from pydantic import ConfigDict, Field, model_validator -from pydantic.dataclasses import dataclass from typing_extensions import Self from vllm.config.utils import config @@ -26,8 +25,7 @@ MaxLoRARanks = Literal[1, 8, 16, 32, 64, 128, 256, 320, 512] LoRAExtraVocabSize = Literal[256, 512] -@config -@dataclass(config=ConfigDict(arbitrary_types_allowed=True)) +@config(config=ConfigDict(arbitrary_types_allowed=True)) class LoRAConfig: """Configuration for LoRA.""" diff --git a/vllm/config/model.py b/vllm/config/model.py index 48ff44ac9fd..3bb8e71770a 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -9,7 +9,6 @@ from typing import TYPE_CHECKING, Any, Literal, cast, get_args import torch from pydantic import ConfigDict, Field, field_validator, model_validator -from pydantic.dataclasses import dataclass import vllm.envs as envs from vllm.config.model_arch import ( @@ -97,8 +96,7 @@ AttnTypeStr = Literal[ ] -@config -@dataclass(config=ConfigDict(arbitrary_types_allowed=True)) +@config(config=ConfigDict(arbitrary_types_allowed=True)) class ModelConfig: """Configuration for the model.""" diff --git a/vllm/config/multimodal.py b/vllm/config/multimodal.py index f4e834f6406..48eea6f4ef5 100644 --- a/vllm/config/multimodal.py +++ b/vllm/config/multimodal.py @@ -51,7 +51,6 @@ DummyOptions: TypeAlias = ( @config -@dataclass class MultiModalConfig: """Controls the behavior of multimodal models.""" diff --git a/vllm/config/observability.py b/vllm/config/observability.py index 9700c911708..3871759125c 100644 --- a/vllm/config/observability.py +++ b/vllm/config/observability.py @@ -6,7 +6,6 @@ from typing import Any, Literal, cast from packaging.version import parse from pydantic import Field, field_validator, model_validator -from pydantic.dataclasses import dataclass from vllm import version from vllm.config.utils import config @@ -16,7 +15,6 @@ DetailedTraceModules = Literal["model", "worker", "all"] @config -@dataclass class ObservabilityConfig: """Configuration for observability - metrics and tracing.""" diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index fa1aa03121b..131db50f191 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -3,12 +3,10 @@ import os from collections.abc import Callable -from dataclasses import replace from typing import TYPE_CHECKING, Any, Literal import torch from pydantic import Field, field_validator, model_validator -from pydantic.dataclasses import dataclass from torch.distributed import ProcessGroup, ReduceOp from typing_extensions import Self @@ -50,7 +48,6 @@ All2AllBackend = Literal[ @config -@dataclass class EPLBConfig: """Configuration for Expert Parallel Load Balancing (EP).""" @@ -94,7 +91,6 @@ class EPLBConfig: @config -@dataclass class ParallelConfig: """Configuration for the distributed execution.""" @@ -715,6 +711,3 @@ class ParallelConfig: ) return self - - def replace(self, **kwargs) -> Self: - return replace(self, **kwargs) diff --git a/vllm/config/pooler.py b/vllm/config/pooler.py index 6d87ec908f7..75cdc90feaa 100644 --- a/vllm/config/pooler.py +++ b/vllm/config/pooler.py @@ -3,8 +3,6 @@ from typing import Any, Literal, get_args -from pydantic.dataclasses import dataclass - from vllm.config.utils import config from vllm.logger import init_logger from vllm.utils.hashing import safe_hash @@ -19,7 +17,6 @@ TOK_POOLING_TYPES: tuple[TokenPoolingType, ...] = get_args(TokenPoolingType) @config -@dataclass class PoolerConfig: """Controls the behavior of output pooling in pooling models.""" diff --git a/vllm/config/profiler.py b/vllm/config/profiler.py index 425f3fb6bcd..b3b8844f77f 100644 --- a/vllm/config/profiler.py +++ b/vllm/config/profiler.py @@ -5,7 +5,6 @@ import os from typing import Any, Literal from pydantic import Field, model_validator -from pydantic.dataclasses import dataclass from typing_extensions import Self from vllm.config.utils import config @@ -32,7 +31,6 @@ def _is_uri_path(path: str) -> bool: @config -@dataclass class ProfilerConfig: """Dataclass which contains profiler config for the engine.""" diff --git a/vllm/config/scheduler.py b/vllm/config/scheduler.py index 5ff9fc930a5..5e44eb84f36 100644 --- a/vllm/config/scheduler.py +++ b/vllm/config/scheduler.py @@ -6,7 +6,6 @@ from dataclasses import InitVar from typing import TYPE_CHECKING, Any, ClassVar, Literal, cast from pydantic import Field, field_validator -from pydantic.dataclasses import dataclass from typing_extensions import Self from vllm.config.utils import config @@ -24,7 +23,6 @@ SchedulerPolicy = Literal["fcfs", "priority"] @config -@dataclass class SchedulerConfig: """Scheduler configuration.""" diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 966d168b47a..ed3dbefb397 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -5,7 +5,6 @@ import ast from typing import TYPE_CHECKING, Any, Literal, get_args from pydantic import Field, SkipValidation, model_validator -from pydantic.dataclasses import dataclass from typing_extensions import Self from vllm.config.model import ModelConfig @@ -55,7 +54,6 @@ SpeculativeMethod = Literal[ @config -@dataclass class SpeculativeConfig: """Configuration for speculative decoding.""" diff --git a/vllm/config/speech_to_text.py b/vllm/config/speech_to_text.py index fe3532c9742..0233d36576c 100644 --- a/vllm/config/speech_to_text.py +++ b/vllm/config/speech_to_text.py @@ -2,13 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from pydantic.dataclasses import dataclass - from vllm.config.utils import config @config -@dataclass class SpeechToTextConfig: """Configuration for speech-to-text models.""" diff --git a/vllm/config/structured_outputs.py b/vllm/config/structured_outputs.py index 8c060c816fd..c4db15989f3 100644 --- a/vllm/config/structured_outputs.py +++ b/vllm/config/structured_outputs.py @@ -4,7 +4,6 @@ from typing import Any, Literal from pydantic import model_validator -from pydantic.dataclasses import dataclass from typing_extensions import Self from vllm.config.utils import config @@ -16,7 +15,6 @@ StructuredOutputsBackend = Literal[ @config -@dataclass class StructuredOutputsConfig: """Dataclass which contains structured outputs config for the engine.""" diff --git a/vllm/config/utils.py b/vllm/config/utils.py index 9288948c5de..e8c866f0223 100644 --- a/vllm/config/utils.py +++ b/vllm/config/utils.py @@ -10,14 +10,17 @@ import json import pathlib import textwrap from collections.abc import Callable, Mapping, Sequence, Set -from dataclasses import MISSING, Field, dataclass, field, fields, is_dataclass, replace +from dataclasses import MISSING, Field, field, fields, is_dataclass from itertools import pairwise -from typing import TYPE_CHECKING, Any, Protocol, TypeVar +from typing import TYPE_CHECKING, Any, Protocol, TypeVar, cast import regex as re import torch +from pydantic import ConfigDict +from pydantic.dataclasses import dataclass +from pydantic.fields import Field as PydanticField from pydantic.fields import FieldInfo -from typing_extensions import runtime_checkable +from typing_extensions import dataclass_transform, runtime_checkable from vllm.logger import init_logger @@ -29,23 +32,39 @@ else: DataclassInstance = Any ConfigType = type[DataclassInstance] -ConfigT = TypeVar("ConfigT", bound=ConfigType) +ConfigT = TypeVar("ConfigT", bound=DataclassInstance) -def config(cls: ConfigT) -> ConfigT: - """ - A decorator that ensures all fields in a dataclass have default values - and that each field has a docstring. +@dataclass_transform(field_specifiers=(PydanticField,)) +def config( + cls: type[ConfigT] | None = None, + *, + config: ConfigDict | None = None, + **kwargs: Any, +) -> type[ConfigT] | Callable[[type[ConfigT]], type[ConfigT]]: + """Decorator to create a pydantic dataclass with default config. The default config + for the dataclass forbids extra fields. - If a `ConfigT` is used as a CLI argument itself, the `type` keyword argument - provided by `get_kwargs` will be - `pydantic.TypeAdapter(ConfigT).validate_json(cli_arg)` which treats the - `cli_arg` as a JSON string which gets validated by `pydantic`. + All config classes in vLLM should use this decorator. - Config validation is performed by the tools/pre_commit/validate_config.py - script, which is invoked during the pre-commit checks. - """ - return cls + Args: + cls: The class to decorate + config: The pydantic ConfigDict to use. If provided, it will be merged with + the default config. + **kwargs: Additional arguments to pass to pydantic.dataclass.""" + # Extra fields are forbidden by default + merged_config = ConfigDict(extra="forbid") + if config is not None: + merged_config.update(config) + + def decorator(cls): + return dataclass(cls, config=merged_config, **kwargs) + + # Called with arguments: @config(config=...) + if cls is None: + return decorator + # Called without arguments: @config + return decorator(cls) def get_field(cls: ConfigType, name: str) -> Field: @@ -53,24 +72,46 @@ def get_field(cls: ConfigType, name: str) -> Field: default factory fields in `EngineArgs`.""" if not is_dataclass(cls): raise TypeError("The given class is not a dataclass.") - cls_fields = {f.name: f for f in fields(cls)} - if name not in cls_fields: - raise ValueError(f"Field '{name}' not found in {cls.__name__}.") - named_field: Field = cls_fields[name] - if (default_factory := named_field.default_factory) is not MISSING: - return field(default_factory=default_factory) - if (default := named_field.default) is not MISSING: - if isinstance(default, FieldInfo): - # Handle pydantic.Field defaults - if default.default_factory is not None: - return field(default_factory=default.default_factory) - else: - default = default.default - return field(default=default) + try: + named_field = next(f for f in fields(cls) if f.name == name) + except StopIteration as e: + raise ValueError(f"Field '{name}' not found in {cls.__name__}.") from e - raise ValueError( - f"{cls.__name__}.{name} must have a default value or default factory." - ) + # The arguments to copy to the new field + default = named_field.default + default_factory = named_field.default_factory + init = named_field.init + + # Handle pydantic.Field + if isinstance(default, FieldInfo): + if default.init is not None: + init = default.init + if default.default_factory is not None: + default_factory = cast(Callable[[], Any], default.default_factory) + default = MISSING + else: + default = default.default + + if default is MISSING and default_factory is MISSING: + logger.warning_once( + "%s.%s has no default or default factory.", cls.__name__, name + ) + return field(default=default, default_factory=default_factory, init=init) + + +def is_init_field(cls: ConfigType, name: str) -> bool: + return get_field(cls, name).init + + +def replace(dataclass_instance: ConfigT, /, **kwargs) -> ConfigT: + """Like [`dataclasses.replace`](https://docs.python.org/3/library/dataclasses.html#dataclasses.replace), + but compatible with Pydantic dataclasses which use `pydantic.fields.Field` instead + of `dataclasses.field`""" + cls = type(dataclass_instance) + dataclass_dict = dataclass_instance.__dict__ + dataclass_dict = {k: v for k, v in dataclass_dict.items() if is_init_field(cls, k)} + dataclass_dict.update(kwargs) + return cls(**dataclass_dict) def getattr_iter( @@ -172,10 +213,6 @@ def get_attr_docs(cls: type[Any]) -> dict[str, str]: return out -def is_init_field(cls: ConfigType, name: str) -> bool: - return next(f for f in fields(cls) if f.name == name).init - - @runtime_checkable class SupportsHash(Protocol): def compute_hash(self) -> str: ... diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index ea133856360..846ed50e0bd 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -9,7 +9,7 @@ import tempfile import threading import time from contextlib import contextmanager -from dataclasses import is_dataclass, replace +from dataclasses import is_dataclass from datetime import datetime from enum import IntEnum from functools import lru_cache @@ -18,10 +18,8 @@ from typing import TYPE_CHECKING, Any, TypeVar, get_args import torch from pydantic import ConfigDict, Field, model_validator -from pydantic.dataclasses import dataclass import vllm.envs as envs -from vllm.config.speculative import EagleModelTypes from vllm.logger import enable_trace_function_call, init_logger from vllm.transformers_utils.runai_utils import is_runai_obj_uri from vllm.utils import random_uuid @@ -41,9 +39,9 @@ from .observability import ObservabilityConfig from .parallel import ParallelConfig from .profiler import ProfilerConfig from .scheduler import SchedulerConfig -from .speculative import SpeculativeConfig +from .speculative import EagleModelTypes, SpeculativeConfig from .structured_outputs import StructuredOutputsConfig -from .utils import SupportsHash, config +from .utils import SupportsHash, config, replace if TYPE_CHECKING: from transformers import PretrainedConfig @@ -187,8 +185,7 @@ OPTIMIZATION_LEVEL_TO_CONFIG = { } -@config -@dataclass(config=ConfigDict(arbitrary_types_allowed=True)) +@config(config=ConfigDict(arbitrary_types_allowed=True)) class VllmConfig: """Dataclass which contains all vllm-related configuration. This simplifies passing around the distinct configurations in the codebase. @@ -1395,14 +1392,6 @@ class VllmConfig: path = self.compilation_config.debug_dump_path / append_path return path - def replace(self, **kwargs): - """ - Replace attributes of the config, and 'recompute' the config. - dataclass.replace() calls __init__() and __post_init__(), source: - https://docs.python.org/3/library/dataclasses.html#dataclasses.replace - """ - return replace(self, **kwargs) - def __str__(self): return ( f"model={self.model_config.model!r}, " diff --git a/vllm/entrypoints/openai/cli_args.py b/vllm/entrypoints/openai/cli_args.py index 808c2a90882..983040a89dc 100644 --- a/vllm/entrypoints/openai/cli_args.py +++ b/vllm/entrypoints/openai/cli_args.py @@ -13,8 +13,6 @@ from collections.abc import Sequence from dataclasses import field from typing import Any, Literal -from pydantic.dataclasses import dataclass - import vllm.envs as envs from vllm.config import config from vllm.engine.arg_utils import AsyncEngineArgs, optional_type @@ -69,7 +67,6 @@ class LoRAParserAction(argparse.Action): @config -@dataclass class FrontendArgs: """Arguments for the OpenAI-compatible frontend server.""" diff --git a/vllm/transformers_utils/configs/speculators/algos.py b/vllm/transformers_utils/configs/speculators/algos.py index 88bce3d4f79..60bb5d588b9 100644 --- a/vllm/transformers_utils/configs/speculators/algos.py +++ b/vllm/transformers_utils/configs/speculators/algos.py @@ -13,9 +13,10 @@ def register_speculator(name): @register_speculator("eagle3") -def update_eagle3(config_dict: dict, vllm_config: dict) -> None: +def update_eagle3(config_dict: dict, pre_trained_config: dict) -> None: """ - Apply Eagle-3 specific configuration transformations. + Apply Eagle-3 specific configuration transformations to the `dict` used to + construct the Transformers PreTrainedConfig. Eagle-3 specific fields: - draft_vocab_size: Size of the draft model's vocabulary @@ -27,12 +28,14 @@ def update_eagle3(config_dict: dict, vllm_config: dict) -> None: predictions. This is the standard field used in Eagle3 checkpoints. """ - vllm_config["draft_vocab_size"] = config_dict.get("draft_vocab_size") + pre_trained_config["draft_vocab_size"] = config_dict.get("draft_vocab_size") if config_dict.get("target_hidden_size") is not None: - vllm_config["target_hidden_size"] = config_dict["target_hidden_size"] - vllm_config["norm_before_residual"] = config_dict.get("norm_before_residual", True) - vllm_config["architectures"] = ["Eagle3LlamaForCausalLM"] + pre_trained_config["target_hidden_size"] = config_dict["target_hidden_size"] + pre_trained_config["norm_before_residual"] = config_dict.get( + "norm_before_residual", True + ) + pre_trained_config["architectures"] = ["Eagle3LlamaForCausalLM"] if config_dict.get("eagle_aux_hidden_state_layer_ids"): - vllm_config["eagle_aux_hidden_state_layer_ids"] = config_dict[ + pre_trained_config["eagle_aux_hidden_state_layer_ids"] = config_dict[ "eagle_aux_hidden_state_layer_ids" ] diff --git a/vllm/transformers_utils/configs/speculators/base.py b/vllm/transformers_utils/configs/speculators/base.py index bf3a5d41319..a57350b0972 100644 --- a/vllm/transformers_utils/configs/speculators/base.py +++ b/vllm/transformers_utils/configs/speculators/base.py @@ -24,13 +24,16 @@ class SpeculatorsConfig(PretrainedConfig): """Load speculators Eagle config and convert to vLLM format.""" config_dict, _ = cls.get_config_dict(pretrained_model_name_or_path, **kwargs) - vllm_config = cls.extract_vllm_speculative_config(config_dict) + vllm_config = cls.extract_transformers_pre_trained_config(config_dict) return cls(**vllm_config) @classmethod - def extract_vllm_speculative_config( + def extract_transformers_pre_trained_config( cls, config_dict: dict[str, Any] ) -> dict[str, Any]: + """ + Extract standard Transformers PreTrainedConfig config from speculators config. + """ speculators_model_type = config_dict.get("speculators_model_type") if speculators_model_type not in SUPPORTED_SPECULATORS_TYPES: raise ValueError( @@ -38,15 +41,23 @@ class SpeculatorsConfig(PretrainedConfig): "Please ensure you're loading a speculators-format model." ) + # Start with transformer layer configuration if present + pre_trained_config = config_dict.get("transformer_layer_config", {}) + # Apply anything specific to the supported algorithm + algo_updater = SUPPORTED_SPECULATORS_TYPES[speculators_model_type] + algo_updater(config_dict=config_dict, pre_trained_config=pre_trained_config) + return pre_trained_config + + @classmethod + def extract_vllm_speculative_config( + cls, config_dict: dict[str, Any] + ) -> dict[str, Any]: + """Extract vLLM speculative config from speculators config.""" # validate fields # TODO: @dsikka - use speculators pydantic model to validate cls.validate_speculators_config(config_dict=config_dict) # Convert from speculators config -> format that can be ingested by vLLM - vllm_config = cls.build_vllm_speculative_config(config_dict=config_dict) - # Apply anything specific to the supported algorithm - algo_updater = SUPPORTED_SPECULATORS_TYPES[speculators_model_type] - algo_updater(config_dict=config_dict, vllm_config=vllm_config) - return vllm_config + return cls.build_vllm_speculative_config(config_dict=config_dict) @classmethod def validate_speculators_config(cls, config_dict: dict[str, Any]) -> None: @@ -101,14 +112,7 @@ class SpeculatorsConfig(PretrainedConfig): ) # Build base vLLM speculative configuration - vllm_config = { + return { "method": config_dict.get("speculators_model_type"), "num_speculative_tokens": num_speculative_tokens, - "target_model": spec_config.get("verifier")["name_or_path"], } - - # Merge transformer layer configuration if present - transformer_config = config_dict.get("transformer_layer_config", {}) - vllm_config.update(transformer_config) - - return vllm_config diff --git a/vllm/v1/spec_decode/draft_model.py b/vllm/v1/spec_decode/draft_model.py index 9c675401358..18e98b26761 100644 --- a/vllm/v1/spec_decode/draft_model.py +++ b/vllm/v1/spec_decode/draft_model.py @@ -4,7 +4,7 @@ from typing import Any import torch -from vllm.config import VllmConfig, get_layers_from_vllm_config +from vllm.config import VllmConfig, get_layers_from_vllm_config, replace from vllm.logger import init_logger from vllm.model_executor.layers.attention import Attention from vllm.model_executor.model_loader import get_model @@ -191,10 +191,12 @@ def create_vllm_config_for_draft_model( old = target_model_vllm_config assert old.speculative_config is not None, "speculative_config is not set" old_spec_config = old.speculative_config - new_parallel_config = old_spec_config.draft_parallel_config.replace( - rank=old.parallel_config.rank + new_parallel_config = replace( + old_spec_config.draft_parallel_config, + rank=old.parallel_config.rank, ) - new: VllmConfig = old.replace( + new: VllmConfig = replace( + old, quant_config=None, # quant_config is recomputed in __init__() model_config=old_spec_config.draft_model_config, parallel_config=new_parallel_config, From a372f3f40afd0aed802242ce59b6a2640d4ef59e Mon Sep 17 00:00:00 2001 From: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com> Date: Wed, 4 Feb 2026 00:10:31 +0400 Subject: [PATCH 031/810] [MISC] Fix Tensor Parallelism for Quantized Mamba Models with n_groups=1 (#33257) Signed-off-by: Vadim Gimpelson --- .../layers/mamba/mamba_mixer2.py | 202 ++++++++---------- 1 file changed, 84 insertions(+), 118 deletions(-) diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index a620495b7a5..f602d9b6219 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -17,7 +17,6 @@ from vllm.forward_context import ForwardContext, get_forward_context from vllm.model_executor.custom_op import CustomOp from vllm.model_executor.layers.linear import ( ColumnParallelLinear, - MergedColumnParallelLinear, RowParallelLinear, ) from vllm.model_executor.layers.mamba.abstract import MambaBase @@ -40,6 +39,7 @@ from vllm.model_executor.model_loader.weight_utils import ( composed_weight_loader, sharded_weight_loader, ) +from vllm.model_executor.parameter import BasevLLMParameter from vllm.model_executor.utils import set_weight_attrs from vllm.platforms import current_platform from vllm.utils.torch_utils import direct_register_custom_op @@ -280,13 +280,6 @@ class MambaMixer2(MambaBase, CustomOp): "then num_groups must equal 1." ) - assert ( - (n_groups % self.tp_size == 0) or self.tp_size == 1 or quant_config is None - ), ( - "Tensor parallel currently supported for quantized models only " - "if tensor parallel world size divides num groups." - ) - self.ssm_state_size = ssm_state_size self.conv_kernel_size = conv_kernel_size self.activation = activation @@ -308,121 +301,94 @@ class MambaMixer2(MambaBase, CustomOp): self.groups_ssm_state_size = self.n_groups * self.ssm_state_size self.conv_dim = intermediate_size + 2 * self.groups_ssm_state_size - if n_groups % self.tp_size == 0: - self.conv1d = MergedColumnParallelLinear( - input_size=conv_kernel_size, - output_sizes=[ - intermediate_size, - self.groups_ssm_state_size, - self.groups_ssm_state_size, - ], - bias=use_conv_bias, - quant_config=None, - prefix=f"{prefix}.conv1d", - ) + # Use ColumnParallelLinear with custom weight loaders for both cases: + # - When n_groups % tp_size == 0: standard sharding without duplication + # - When n_groups == 1: groups are duplicated across TP ranks + # The custom weight loader handles both cases correctly. - self.in_proj = MergedColumnParallelLinear( - input_size=hidden_size, - output_sizes=[ - intermediate_size, - intermediate_size, - self.groups_ssm_state_size, - self.groups_ssm_state_size, - self.num_heads, - ], - bias=use_bias, - quant_config=quant_config, - prefix=f"{prefix}.in_proj", - ) - else: - # This is the n_groups == 1 case, - # where we need to duplicate groups if TP>1. + self.conv1d = ColumnParallelLinear( + input_size=conv_kernel_size, + output_size=self.conv_dim, + bias=use_conv_bias, + quant_config=None, + prefix=f"{prefix}.conv1d", + ) - self.conv1d = ColumnParallelLinear( - input_size=conv_kernel_size, - output_size=self.conv_dim, - bias=use_conv_bias, - quant_config=None, - prefix=f"{prefix}.conv1d", - ) + self.in_proj = ColumnParallelLinear( + input_size=hidden_size, + output_size=intermediate_size + self.conv_dim + self.num_heads, + bias=use_bias, + quant_config=quant_config, + prefix=f"{prefix}.in_proj", + ) - self.in_proj = ColumnParallelLinear( - input_size=hidden_size, - output_size=intermediate_size + self.conv_dim + self.num_heads, - bias=use_bias, - quant_config=quant_config, - prefix=f"{prefix}.in_proj", - ) + # Configure shard settings for the custom weight loader: + # - group_shard_settings handles group duplication when n_groups == 1 + # - When n_groups % tp_size == 0, extra=0 and duplicate_groups=False + group_shard_settings = ( + self.groups_ssm_state_size, # expected model size + (self.n_groups - n_groups) * self.ssm_state_size, # extra dims assigned + n_groups == 1, # duplicate groups when n_groups == 1 + ) + intermediate_settings = (intermediate_size, 0, False) + head_settings = (self.num_heads, 0, False) - # - because in_proj is a concatenation of 3 weights, we - # need to interleave them before sharding - # - use the custom weight loader mamba_v2_sharded_weight_loader - # for conv1d.bias, covn1d.weight and in_proj.weight - # - need to set these settings, to assign the groups - # to the head shards - group_shard_settings = ( - self.groups_ssm_state_size, # expected model size - (self.n_groups - n_groups) * self.ssm_state_size, # extra dims assigned - n_groups == 1, # if there was only one group - ) - intermediate_settings = (intermediate_size, 0, False) - head_settings = (self.num_heads, 0, False) - - # - the weight already has a "weight_loader" attribute - # which set_weight_attrs will raise if we do not - # delete before trying to override it - # - ditto for the other two weights below - delattr(self.conv1d.bias, "weight_loader") - set_weight_attrs( - self.conv1d.bias, - { - "weight_loader": mamba_v2_sharded_weight_loader( - [ - intermediate_settings, - group_shard_settings, - group_shard_settings, - ], - self.tp_size, - tp_rank, - ) - }, - ) - - delattr(self.conv1d.weight, "weight_loader") - set_weight_attrs( - self.conv1d.weight, - { - "weight_loader": mamba_v2_sharded_weight_loader( - [ - intermediate_settings, - group_shard_settings, - group_shard_settings, - ], - self.tp_size, - tp_rank, - ) - }, - ) - - if quant_config is None: - # - quant layers do not have a weight loader - delattr(self.in_proj.weight, "weight_loader") - set_weight_attrs( - self.in_proj.weight, - { - "weight_loader": mamba_v2_sharded_weight_loader( - [ - intermediate_settings, # for gate - intermediate_settings, - group_shard_settings, - group_shard_settings, - head_settings, # for dt - ], - self.tp_size, - tp_rank, - ) - }, + # Apply custom weight loaders for conv1d (bias and weight) + delattr(self.conv1d.bias, "weight_loader") + set_weight_attrs( + self.conv1d.bias, + { + "weight_loader": mamba_v2_sharded_weight_loader( + [ + intermediate_settings, + group_shard_settings, + group_shard_settings, + ], + self.tp_size, + tp_rank, ) + }, + ) + + delattr(self.conv1d.weight, "weight_loader") + set_weight_attrs( + self.conv1d.weight, + { + "weight_loader": mamba_v2_sharded_weight_loader( + [ + intermediate_settings, + group_shard_settings, + group_shard_settings, + ], + self.tp_size, + tp_rank, + ) + }, + ) + + # Create the custom weight loader for in_proj + mamba_loader = mamba_v2_sharded_weight_loader( + [ + intermediate_settings, # for gate + intermediate_settings, + group_shard_settings, + group_shard_settings, + head_settings, # for dt + ], + self.tp_size, + tp_rank, + ) + + # Apply the custom weight loader to in_proj.weight + # Works for both non-quantized (Parameter) and quantized + # (ModelWeightParameter which extends BasevLLMParameter) + if isinstance(self.in_proj.weight, BasevLLMParameter): + # For BasevLLMParameter subclasses (quantized layers like FP8) + self.in_proj.weight.weight_loader = mamba_loader + else: + # For standard Parameter (non-quantized layers) + delattr(self.in_proj.weight, "weight_loader") + set_weight_attrs(self.in_proj.weight, {"weight_loader": mamba_loader}) # unsqueeze to fit conv1d weights shape into the linear weights shape. # Can't do this in `weight_loader` since it already exists in From 3f7662d6505e441026e668ba78a2207d669f4f32 Mon Sep 17 00:00:00 2001 From: Patrick von Platen Date: Tue, 3 Feb 2026 22:03:28 +0100 Subject: [PATCH 032/810] [Voxtral Realtime] Change name (#33716) Signed-off-by: Patrick von Platen --- examples/online_serving/openai_realtime_client.py | 4 ++-- examples/online_serving/openai_realtime_microphone_client.py | 4 ++-- tests/entrypoints/openai/test_realtime_validation.py | 2 +- tests/models/multimodal/generation/test_voxtral_realtime.py | 2 +- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/examples/online_serving/openai_realtime_client.py b/examples/online_serving/openai_realtime_client.py index 5aa31e3e5a3..17335bd238b 100644 --- a/examples/online_serving/openai_realtime_client.py +++ b/examples/online_serving/openai_realtime_client.py @@ -7,7 +7,7 @@ audio transcription by uploading an audio file. Before running this script, you must start the vLLM server with a realtime-capable model, for example: - vllm serve mistralai/Voxtral-Mini-3B-Realtime-2602 --enforce-eager + vllm serve mistralai/Voxtral-Mini-4B-Realtime-2602 --enforce-eager Requirements: - vllm with audio support @@ -126,7 +126,7 @@ if __name__ == "__main__": parser.add_argument( "--model", type=str, - default="mistralai/Voxtral-Mini-3B-Realtime-2602", + default="mistralai/Voxtral-Mini-4B-Realtime-2602", help="Model that is served and should be pinged.", ) parser.add_argument( diff --git a/examples/online_serving/openai_realtime_microphone_client.py b/examples/online_serving/openai_realtime_microphone_client.py index fc80b1c50cb..9a48f1466cc 100644 --- a/examples/online_serving/openai_realtime_microphone_client.py +++ b/examples/online_serving/openai_realtime_microphone_client.py @@ -5,7 +5,7 @@ Minimal Gradio demo for real-time speech transcription using the vLLM Realtime A Start the vLLM server first: - vllm serve mistralai/Voxtral-Mini-3B-Realtime-2602 --enforce-eager + vllm serve mistralai/Voxtral-Mini-4B-Realtime-2602 --enforce-eager Then run this script: @@ -166,7 +166,7 @@ if __name__ == "__main__": parser.add_argument( "--model", type=str, - default="mistralai/Voxtral-Mini-3B-Realtime-2602", + default="mistralai/Voxtral-Mini-4B-Realtime-2602", help="Model that is served and should be pinged.", ) parser.add_argument( diff --git a/tests/entrypoints/openai/test_realtime_validation.py b/tests/entrypoints/openai/test_realtime_validation.py index e0868a87dc9..7f12bcaca3d 100644 --- a/tests/entrypoints/openai/test_realtime_validation.py +++ b/tests/entrypoints/openai/test_realtime_validation.py @@ -24,7 +24,7 @@ MISTRAL_FORMAT_ARGS = [ "mistral", ] -MODEL_NAME = "mistralai/Voxtral-Mini-3B-Realtime-2602" +MODEL_NAME = "mistralai/Voxtral-Mini-4B-Realtime-2602" def _audio_to_base64_pcm16(path: str, target_sr: int = 16000) -> str: diff --git a/tests/models/multimodal/generation/test_voxtral_realtime.py b/tests/models/multimodal/generation/test_voxtral_realtime.py index a8fe162f844..d162f80ffa6 100644 --- a/tests/models/multimodal/generation/test_voxtral_realtime.py +++ b/tests/models/multimodal/generation/test_voxtral_realtime.py @@ -19,7 +19,7 @@ from vllm.engine.arg_utils import AsyncEngineArgs from vllm.inputs.data import TokensPrompt from vllm.v1.engine.async_llm import AsyncLLM, StreamingInput -MODEL_NAME = "mistralai/Voxtral-Mini-3B-Realtime-2602" +MODEL_NAME = "mistralai/Voxtral-Mini-4B-Realtime-2602" ENGINE_CONFIG = dict( model=MODEL_NAME, max_model_len=8192, From 2a99c5a6c86daef8c766ba2dbf05c385b192c64b Mon Sep 17 00:00:00 2001 From: Michael Goin Date: Tue, 3 Feb 2026 16:26:51 -0500 Subject: [PATCH 033/810] [Bugfix] Disable TRTLLM FP8 MoE if router_logits_dtype==float32 and routing_method!=DeepSeekV3 (#33613) Signed-off-by: mgoin --- .../layers/fused_moe/flashinfer_trtllm_moe.py | 35 ++++++++++++++++--- .../layers/fused_moe/oracle/fp8.py | 10 +++--- .../compressed_tensors_moe.py | 15 ++------ .../model_executor/layers/quantization/fp8.py | 15 ++------ vllm/model_executor/models/minimax_m2.py | 1 + 5 files changed, 43 insertions(+), 33 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py index 43e02d51043..0182cfc195f 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py @@ -98,7 +98,23 @@ def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bo return not moe_parallel_config.enable_eplb -def is_supported_config_trtllm( +def _supports_router_logits_dtype( + router_logits_dtype: torch.dtype | None, + routing_method: RoutingMethodType, +) -> bool: + """ + The FlashInfer TRTLLM FP8 kernel expects bfloat16 router_logits by default. + Only DeepSeekV3 routing supports float32 router_logits (which is converted + internally in the kernel). + """ + if router_logits_dtype == torch.float32: + # Only DeepSeekV3 routing handles float32 logits + # https://github.com/flashinfer-ai/flashinfer/issues/2469 + return routing_method == RoutingMethodType.DeepSeekV3 + return True + + +def is_supported_config_trtllm_fp8( moe_config: FusedMoEConfig, weight_key: QuantKey | None, activation_key: QuantKey | None, @@ -127,6 +143,10 @@ def is_supported_config_trtllm( return False, _make_reason("routing method") elif activation_format != mk.FusedMoEActivationFormat.Standard: return False, _make_reason("activation format") + elif not _supports_router_logits_dtype( + moe_config.router_logits_dtype, moe_config.routing_method + ): + return False, _make_reason("float32 router_logits with non-DeepSeekV3 routing") return True, None @@ -161,7 +181,7 @@ def is_supported_config_trtllm_bf16( def flashinfer_fused_moe_blockscale_fp8( routing_logits: torch.Tensor, - routing_bias: torch.Tensor, + routing_bias: torch.Tensor | None, x: torch.Tensor, w13_weight: torch.Tensor, w13_weight_scale_inv: torch.Tensor, @@ -175,7 +195,7 @@ def flashinfer_fused_moe_blockscale_fp8( expert_offset: int, local_num_experts: int, block_shape: list[int], - routing_method_type: int = int(RoutingMethodType.DeepSeekV3), + routing_method_type: int, routed_scaling: float | None = 1.0, ) -> torch.Tensor: from vllm.utils.flashinfer import flashinfer_trtllm_fp8_block_scale_moe @@ -188,6 +208,13 @@ def flashinfer_fused_moe_blockscale_fp8( # Routing kernel expects #experts <= #threads 512 assert global_num_experts <= 512 + # The DeepSeekV3 routing method requires float32 router logits. + if routing_method_type == RoutingMethodType.DeepSeekV3: + routing_logits = routing_logits.to(torch.float32) + + if routing_bias is not None: + routing_bias = routing_bias.to(x.dtype) + a_q, a_sf = per_token_group_quant_fp8(x, block_shape[1]) # NOTE: scales of hidden states have to be transposed! a_sf_t = a_sf.t().contiguous() @@ -215,7 +242,7 @@ def flashinfer_fused_moe_blockscale_fp8( def flashinfer_fused_moe_blockscale_fp8_fake( routing_logits: torch.Tensor, - routing_bias: torch.Tensor, + routing_bias: torch.Tensor | None, x: torch.Tensor, w13_weight: torch.Tensor, w13_weight_scale_inv: torch.Tensor, diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index 15fc6e237bf..70c2516747f 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -18,7 +18,7 @@ from vllm.model_executor.layers.fused_moe.config import ( fp8_w8a16_moe_quant_config, ) from vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe import ( - is_supported_config_trtllm, + is_supported_config_trtllm_fp8, ) from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( FlashinferMoeBackend, @@ -213,7 +213,7 @@ def select_fp8_moe_backend( if fi_backend == FlashinferMoeBackend.TENSORRT_LLM: backend = Fp8MoeBackend.FLASHINFER_TRTLLM - supported, reason = is_supported_config_trtllm( + supported, reason = is_supported_config_trtllm_fp8( config, weight_key, activation_key, activation_format ) if supported: @@ -240,7 +240,7 @@ def select_fp8_moe_backend( ]: if backend == Fp8MoeBackend.FLASHINFER_TRTLLM: k_cls = None - supported, reason = is_supported_config_trtllm( + supported, reason = is_supported_config_trtllm_fp8( config, weight_key, activation_key, @@ -309,7 +309,7 @@ def select_fp8_moe_backend( for backend in AVAILABLE_BACKENDS: if backend == Fp8MoeBackend.FLASHINFER_TRTLLM: k_cls = None - supported, reason = is_supported_config_trtllm( + supported, reason = is_supported_config_trtllm_fp8( config, weight_key, activation_key, @@ -482,7 +482,7 @@ def make_fp8_moe_kernel( ) assert prepare_finalize is not None - logger.info_once("Using %s", prepare_finalize.__class__.__name__) + logger.info_once("Using %s", prepare_finalize.__class__.__name__, scope="local") # Create Experts. if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts: diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py index dbfa8fb9bd7..5152c5cccab 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py @@ -27,7 +27,6 @@ from vllm.model_executor.layers.fused_moe import ( from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, FusedMoEQuantConfig, - RoutingMethodType, int4_w4a16_moe_quant_config, int4_w4afp8_moe_quant_config, int8_w8a8_moe_quant_config, @@ -1027,17 +1026,9 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod): if self.block_quant: import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe # noqa: E501, F401 - e_score_correction_bias = ( - layer.e_score_correction_bias.to(x.dtype) - if layer.e_score_correction_bias is not None - else None - ) - routing_method_type = layer.routing_method_type return torch.ops.vllm.flashinfer_fused_moe_blockscale_fp8( - routing_logits=router_logits.to(torch.float32) - if routing_method_type == RoutingMethodType.DeepSeekV3 - else router_logits, - routing_bias=e_score_correction_bias, + routing_logits=router_logits, + routing_bias=layer.e_score_correction_bias, x=x, w13_weight=layer.w13_weight, w13_weight_scale_inv=layer.w13_weight_scale, @@ -1051,7 +1042,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod): expert_offset=layer.ep_rank * layer.local_num_experts, local_num_experts=layer.local_num_experts, block_shape=self.weight_block_size, - routing_method_type=routing_method_type, + routing_method_type=layer.routing_method_type, routed_scaling=layer.routed_scaling_factor, ) else: diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 9b7d654335d..53bdb972bcd 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -26,7 +26,6 @@ from vllm.model_executor.layers.fused_moe import ( ) from vllm.model_executor.layers.fused_moe.config import ( FusedMoEQuantConfig, - RoutingMethodType, ) from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( @@ -980,17 +979,9 @@ class Fp8MoEMethod(FusedMoEMethodBase): if self.block_quant: import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe # noqa: E501, F401 - e_score_correction_bias = ( - layer.e_score_correction_bias.to(x.dtype) - if layer.e_score_correction_bias is not None - else None - ) - routing_method_type = layer.routing_method_type return torch.ops.vllm.flashinfer_fused_moe_blockscale_fp8( - routing_logits=router_logits.to(torch.float32) - if routing_method_type == RoutingMethodType.DeepSeekV3 - else router_logits, - routing_bias=e_score_correction_bias, + routing_logits=router_logits, + routing_bias=layer.e_score_correction_bias, x=x, w13_weight=layer.w13_weight, w13_weight_scale_inv=layer.w13_weight_scale_inv, @@ -1004,7 +995,7 @@ class Fp8MoEMethod(FusedMoEMethodBase): expert_offset=layer.ep_rank * layer.local_num_experts, local_num_experts=layer.local_num_experts, block_shape=self.weight_block_size, - routing_method_type=routing_method_type, + routing_method_type=layer.routing_method_type, routed_scaling=layer.routed_scaling_factor, ) else: diff --git a/vllm/model_executor/models/minimax_m2.py b/vllm/model_executor/models/minimax_m2.py index 7583be200ff..2dc0f33cc1c 100644 --- a/vllm/model_executor/models/minimax_m2.py +++ b/vllm/model_executor/models/minimax_m2.py @@ -107,6 +107,7 @@ class MiniMaxM2MoE(nn.Module): renormalize=True, quant_config=quant_config, prefix=f"{prefix}.experts", + router_logits_dtype=torch.float32, ) self.gate = ReplicatedLinear( From bd8da29a66ea8c0e0f208cab7ba0b6be640f3faa Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Tue, 3 Feb 2026 18:29:48 -0500 Subject: [PATCH 034/810] [Bugfix] Fix sparse MLA metadata building (#33579) Signed-off-by: Matthew Bonanni --- .../layers/attention/mla_attention.py | 53 ++++++++----------- 1 file changed, 22 insertions(+), 31 deletions(-) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 501b939c11b..1b719330e2a 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -522,22 +522,6 @@ class MLAAttention(nn.Module, AttentionLayerBase): k_c_normed = k_c_normed[:num_actual_toks, ...] k_pe = k_pe[:num_actual_toks, ...] - assert ( - attn_metadata.num_decodes is not None - and attn_metadata.num_prefills is not None - and attn_metadata.num_decode_tokens is not None - ) - - has_decode = attn_metadata.num_decodes > 0 - has_prefill = attn_metadata.num_prefills > 0 - num_decode_tokens = attn_metadata.num_decode_tokens - - decode_q = q[:num_decode_tokens] - - prefill_q = q[num_decode_tokens:] - prefill_k_pe = k_pe[num_decode_tokens:] - prefill_k_c_normed = k_c_normed[num_decode_tokens:] - # write the latent and rope to kv cache if kv_cache.numel() > 0: ops.concat_and_cache_mla( @@ -555,27 +539,32 @@ class MLAAttention(nn.Module, AttentionLayerBase): # Sparse MLA impls only support forward_mqa (decode-style attention) is_sparse_impl = isinstance(self.impl, SparseMLAAttentionImpl) - if has_prefill and not is_sparse_impl: + if is_sparse_impl: + num_mqa_tokens = q.size(0) + num_mha_tokens = 0 + else: + assert ( + attn_metadata.num_decodes is not None + and attn_metadata.num_prefills is not None + and attn_metadata.num_decode_tokens is not None + ) + num_mqa_tokens = attn_metadata.num_decode_tokens + num_mha_tokens = q.size(0) - num_mqa_tokens + + if num_mha_tokens > 0: self.impl.forward_mha( - prefill_q, - prefill_k_c_normed, - prefill_k_pe, + q[num_mqa_tokens:], + k_c_normed[num_mqa_tokens:], + k_pe[num_mqa_tokens:], kv_cache, attn_metadata, self._k_scale, - output=output[num_decode_tokens:], + output=output[num_mqa_tokens:], ) - if has_decode or (has_prefill and is_sparse_impl): - # For sparse impl, we always use forward_mqa for all tokens - # For non-sparse impl, we only use forward_mqa for decode tokens - if is_sparse_impl: - mqa_q = q - mqa_output_slice = output - else: - assert attn_metadata.decode is not None - mqa_q = decode_q - mqa_output_slice = output[:num_decode_tokens] + if num_mqa_tokens > 0: + mqa_q = q[:num_mqa_tokens] + mqa_output_slice = output[:num_mqa_tokens] mqa_q_nope, mqa_q_pe = mqa_q.split( [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 @@ -644,6 +633,8 @@ class MLAAttention(nn.Module, AttentionLayerBase): mqa_q = get_dcp_group().all_gather(mqa_q, dim=1) # call decode attn + if not is_sparse_impl: + assert attn_metadata.decode is not None attn_out, lse = self.impl.forward_mqa(mqa_q, kv_cache, attn_metadata, self) # correct dcp attn_out with lse. From 655efb3e69bb18d150a88c0d726ca2b49f22cbdd Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Tue, 3 Feb 2026 18:30:47 -0500 Subject: [PATCH 035/810] [Dependency] Remove comments of ray in dependency files (#33351) Signed-off-by: yewentao256 --- requirements/cuda.txt | 2 +- requirements/rocm.txt | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/requirements/cuda.txt b/requirements/cuda.txt index 8e9a9063877..378f8460574 100644 --- a/requirements/cuda.txt +++ b/requirements/cuda.txt @@ -4,7 +4,7 @@ numba == 0.61.2 # Required for N-gram speculative decoding # Dependencies for NVIDIA GPUs -ray[cgraph]>=2.48.0 # Ray Compiled Graph, required for pipeline parallelism in V1. +ray[cgraph]>=2.48.0 torch==2.9.1 torchaudio==2.9.1 # These must be updated alongside torch diff --git a/requirements/rocm.txt b/requirements/rocm.txt index 5aeb16599a3..dc407995a21 100644 --- a/requirements/rocm.txt +++ b/requirements/rocm.txt @@ -5,7 +5,7 @@ numba == 0.61.2 # Required for N-gram speculative decoding # Dependencies for AMD GPUs datasets -ray[cgraph]>=2.48.0 # Ray Compiled Graph, required for pipeline parallelism in V1. +ray[cgraph]>=2.48.0 peft pytest-asyncio tensorizer==2.10.1 From 52ee21021a87735d46c4245c60bc0be42dd58c73 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Tue, 3 Feb 2026 15:34:41 -0800 Subject: [PATCH 036/810] [BugFix][Spec Decoding] Fix negative accepted tokens metric crash (#33729) Signed-off-by: Nick Hill --- tests/v1/core/test_scheduler.py | 60 +++++++++++++++++++++++++++++++++ vllm/v1/core/sched/scheduler.py | 2 +- 2 files changed, 61 insertions(+), 1 deletion(-) diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 537a02464d0..580cc70ff34 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -885,6 +885,66 @@ def test_schedule_spec_decoding_stats(spec_tokens, output_tokens, expected): assert stats.num_accepted_tokens_per_pos == expected[3] +def test_spec_decoding_stats_empty_output(): + """Test that spec decoding stats handle empty output tokens gracefully. + + This is a regression test for a bug where empty sampled_token_ids + would cause num_accepted = len([]) - 1 = -1, leading to a + ValueError when incrementing a Prometheus counter with a negative value. + """ + num_spec_tokens = 3 + scheduler = create_scheduler(num_speculative_tokens=num_spec_tokens) + requests = create_requests(num_requests=1, num_tokens=1) + request = requests[0] + req_id = request.request_id + + scheduler.add_request(request) + + # Initial schedule (prefill) + output = scheduler.schedule() + assert len(output.scheduled_new_reqs) == 1 + + # Complete the prefill with a sampled token + model_runner_output = ModelRunnerOutput( + req_ids=[req_id], + req_id_to_index={req_id: 0}, + sampled_token_ids=[[0]], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ) + scheduler.update_from_output(output, model_runner_output) + + # Add draft tokens for speculation + draft_token_ids = DraftTokenIds([req_id], [[1, 2, 3]]) + scheduler.update_draft_token_ids(draft_token_ids) + + # Schedule the speculated tokens for validation + output = scheduler.schedule() + assert req_id in output.scheduled_spec_decode_tokens + assert len(output.scheduled_spec_decode_tokens[req_id]) == 3 + + # Simulate empty output tokens (e.g., due to request abortion or error) + # This would previously cause num_accepted = -1 and crash + model_runner_output = ModelRunnerOutput( + req_ids=[req_id], + req_id_to_index={req_id: 0}, + sampled_token_ids=[[]], # Empty output tokens + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ) + + # This should not raise an error + engine_core_outputs = scheduler.update_from_output(output, model_runner_output) + + # Spec decoding stats should be None since no tokens were generated + scheduler_stats = ( + engine_core_outputs[0].scheduler_stats if engine_core_outputs else None + ) + assert scheduler_stats is None or scheduler_stats.spec_decoding_stats is None + + def _assert_right_scheduler_output( output: SchedulerOutput, num_requests: int, diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 3f7ac9374e1..83c965f233a 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1276,7 +1276,7 @@ class Scheduler(SchedulerInterface): scheduled_spec_token_ids = ( scheduler_output.scheduled_spec_decode_tokens.get(req_id) ) - if scheduled_spec_token_ids: + if scheduled_spec_token_ids and generated_token_ids: num_draft_tokens = len(scheduled_spec_token_ids) num_accepted = len(generated_token_ids) - 1 num_rejected = num_draft_tokens - num_accepted From 1b8fe6f7c4b96ec57172f4fd268341a03c12499b Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Wed, 4 Feb 2026 09:48:40 +0800 Subject: [PATCH 037/810] [Frontend][4/n] Make pooling entrypoints request schema consensus | ScoreRequest (#33060) Signed-off-by: wang.yuqi --- .../pooling/score/vision_rerank_api_online.py | 23 +++ .../pooling/score/vision_score_api_online.py | 38 ++++ .../pooling/score/test_online_score.py | 29 +++ .../pooling/score/test_online_score_vision.py | 131 ++++++++++++- vllm/entrypoints/llm.py | 110 ++++------- vllm/entrypoints/pooling/score/protocol.py | 39 ++-- vllm/entrypoints/pooling/score/serving.py | 178 +++++++++--------- vllm/entrypoints/pooling/score/utils.py | 89 +++++++-- 8 files changed, 432 insertions(+), 205 deletions(-) diff --git a/examples/pooling/score/vision_rerank_api_online.py b/examples/pooling/score/vision_rerank_api_online.py index 875971f1aef..f5a7c1018c5 100644 --- a/examples/pooling/score/vision_rerank_api_online.py +++ b/examples/pooling/score/vision_rerank_api_online.py @@ -89,6 +89,29 @@ def main(args): response = requests.post(rerank_url, json=prompt) pprint.pprint(response.json()) + print("Query: string & Document: text + image url") + prompt = { + "model": model, + "query": query, + "documents": {"content": [documents[0], documents[1]]}, + } + response = requests.post(rerank_url, json=prompt) + pprint.pprint(response.json()) + + print("Query: string & Document: list") + prompt = { + "model": model, + "query": query, + "documents": [ + document, + {"content": [documents[0]]}, + {"content": [documents[1]]}, + {"content": [documents[0], documents[1]]}, + ], + } + response = requests.post(rerank_url, json=prompt) + pprint.pprint(response.json()) + if __name__ == "__main__": args = parse_args() diff --git a/examples/pooling/score/vision_score_api_online.py b/examples/pooling/score/vision_score_api_online.py index df8218a8b2a..7942ddaedb0 100644 --- a/examples/pooling/score/vision_score_api_online.py +++ b/examples/pooling/score/vision_score_api_online.py @@ -92,6 +92,44 @@ def main(args): response = requests.post(score_url, json=prompt) pprint.pprint(response.json()) + print("Query: string & Document: text + image url") + prompt = { + "model": model, + "queries": query, + "documents": {"content": [documents[0], documents[1]]}, + } + response = requests.post(score_url, json=prompt) + pprint.pprint(response.json()) + + print("Query: string & Document: list") + prompt = { + "model": model, + "queries": query, + "documents": [ + document, + {"content": [documents[0]]}, + {"content": [documents[1]]}, + {"content": [documents[0], documents[1]]}, + ], + } + response = requests.post(score_url, json=prompt) + pprint.pprint(response.json()) + + print("Query: list & Document: list") + data = [ + document, + {"content": [documents[0]]}, + {"content": [documents[1]]}, + {"content": [documents[0], documents[1]]}, + ] + prompt = { + "model": model, + "queries": data, + "documents": data, + } + response = requests.post(score_url, json=prompt) + pprint.pprint(response.json()) + if __name__ == "__main__": args = parse_args() diff --git a/tests/entrypoints/pooling/score/test_online_score.py b/tests/entrypoints/pooling/score/test_online_score.py index e1cc074e885..c8b3347780c 100644 --- a/tests/entrypoints/pooling/score/test_online_score.py +++ b/tests/entrypoints/pooling/score/test_online_score.py @@ -90,6 +90,35 @@ class TestModel: for i in range(len(vllm_outputs)): assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01) + def test_queries_str_items_str( + self, server: RemoteOpenAIServer, model: dict[str, Any], runner + ): + queries = "What is the capital of France?" + items = "The capital of France is Paris." + + score_response = requests.post( + server.url_for("score"), + json={ + "model": model["name"], + "queries": queries, + "items": items, + }, + ) + score_response.raise_for_status() + score = ScoreResponse.model_validate(score_response.json()) + + assert score.id is not None + assert score.data is not None + assert len(score.data) == 1 + + vllm_outputs = [d.score for d in score.data] + + text_pairs = [[queries, items]] + hf_outputs = run_transformers(runner, model, text_pairs) + + for i in range(len(vllm_outputs)): + assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01) + def test_text_1_str_text_2_str( self, server: RemoteOpenAIServer, model: dict[str, Any], runner ): diff --git a/tests/entrypoints/pooling/score/test_online_score_vision.py b/tests/entrypoints/pooling/score/test_online_score_vision.py index 0f19498c05b..9e9bc3fec88 100644 --- a/tests/entrypoints/pooling/score/test_online_score_vision.py +++ b/tests/entrypoints/pooling/score/test_online_score_vision.py @@ -5,7 +5,7 @@ import pytest import requests from tests.utils import VLLM_PATH, RemoteOpenAIServer -from vllm.entrypoints.pooling.score.protocol import ScoreResponse +from vllm.entrypoints.pooling.score.protocol import RerankResponse, ScoreResponse from vllm.multimodal.utils import encode_image_url, fetch_image MODEL_NAME = "Qwen/Qwen3-VL-Reranker-2B" @@ -16,11 +16,12 @@ HF_OVERRIDES = { } query = "A cat standing in the snow." +document = "This product was excellent and exceeded my expectations." image_url = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/cat_snow.jpg" documents = [ { "type": "text", - "text": query, + "text": document, }, { "type": "image_url", @@ -32,6 +33,11 @@ documents = [ }, ] +TEXT_VS_TEXT = 0.10040374100208282 +TEXT_VS_IMAGE = 0.7423753142356873 +TEXT_VS_TEXT_PLUS_IMAGE = 0.5298863053321838 +TOL = 0.05 + @pytest.fixture(scope="module") def server(): @@ -50,15 +56,12 @@ def server(): def test_score_api_queries_str_documents_str(server: RemoteOpenAIServer): - queries = "What is the capital of France?" - documents = "The capital of France is Paris." - score_response = requests.post( server.url_for("score"), json={ "model": MODEL_NAME, - "queries": queries, - "documents": documents, + "queries": query, + "documents": document, }, ) score_response.raise_for_status() @@ -67,6 +70,8 @@ def test_score_api_queries_str_documents_str(server: RemoteOpenAIServer): assert score.id is not None assert score.data is not None assert len(score.data) == 1 + assert score.usage.prompt_tokens == 81 + assert score.data[0].score == pytest.approx(TEXT_VS_TEXT, rel=TOL) def test_score_api_queries_str_documents_text_content(server: RemoteOpenAIServer): @@ -84,6 +89,8 @@ def test_score_api_queries_str_documents_text_content(server: RemoteOpenAIServer assert score.id is not None assert score.data is not None assert len(score.data) == 1 + assert score.usage.prompt_tokens == 81 + assert score.data[0].score == pytest.approx(TEXT_VS_TEXT, rel=TOL) def test_score_api_queries_str_documents_image_url_content(server: RemoteOpenAIServer): @@ -101,6 +108,8 @@ def test_score_api_queries_str_documents_image_url_content(server: RemoteOpenAIS assert score.id is not None assert score.data is not None assert len(score.data) == 1 + assert score.usage.prompt_tokens == 98 + assert score.data[0].score == pytest.approx(TEXT_VS_IMAGE, rel=TOL) def test_score_api_queries_str_documents_image_base64_content( @@ -120,3 +129,111 @@ def test_score_api_queries_str_documents_image_base64_content( assert score.id is not None assert score.data is not None assert len(score.data) == 1 + assert score.usage.prompt_tokens == 98 + assert score.data[0].score == pytest.approx(TEXT_VS_IMAGE, rel=TOL) + + +def test_score_api_queries_str_documents_image_url_plus_text_content( + server: RemoteOpenAIServer, +): + score_response = requests.post( + server.url_for("score"), + json={ + "model": MODEL_NAME, + "queries": query, + "documents": {"content": [documents[0], documents[1]]}, + }, + ) + score_response.raise_for_status() + score = ScoreResponse.model_validate(score_response.json()) + + assert score.id is not None + assert score.data is not None + assert len(score.data) == 1 + assert score.usage.prompt_tokens == 108 + assert score.data[0].score == pytest.approx(TEXT_VS_TEXT_PLUS_IMAGE, rel=TOL) + + +def test_score_api_queries_str_documents_list(server: RemoteOpenAIServer): + score_response = requests.post( + server.url_for("score"), + json={ + "model": MODEL_NAME, + "queries": query, + "documents": [ + document, + {"content": [documents[0]]}, + {"content": [documents[1]]}, + {"content": [documents[0], documents[1]]}, + ], + }, + ) + score_response.raise_for_status() + score = ScoreResponse.model_validate(score_response.json()) + + assert score.id is not None + assert score.data is not None + assert len(score.data) == 4 + assert score.usage.prompt_tokens == 368 + assert score.data[0].score == pytest.approx(TEXT_VS_TEXT, rel=TOL) + assert score.data[1].score == pytest.approx(TEXT_VS_TEXT, rel=TOL) + assert score.data[2].score == pytest.approx(TEXT_VS_IMAGE, rel=TOL) + assert score.data[3].score == pytest.approx(TEXT_VS_TEXT_PLUS_IMAGE, rel=TOL) + + +def test_rerank_api_queries_str_documents_list(server: RemoteOpenAIServer): + rerank_response = requests.post( + server.url_for("rerank"), + json={ + "model": MODEL_NAME, + "query": query, + "documents": [ + document, + {"content": [documents[0]]}, + {"content": [documents[1]]}, + {"content": [documents[0], documents[1]]}, + ], + }, + ) + rerank_response.raise_for_status() + rerank = RerankResponse.model_validate(rerank_response.json()) + + assert rerank.id is not None + assert rerank.model is not None + assert rerank.usage is not None + assert len(rerank.results) == 4 + + rerank.results.sort(key=lambda x: x.index) + assert rerank.results[0].relevance_score == pytest.approx(TEXT_VS_TEXT, rel=TOL) + assert rerank.results[1].relevance_score == pytest.approx(TEXT_VS_TEXT, rel=TOL) + assert rerank.results[2].relevance_score == pytest.approx(TEXT_VS_IMAGE, rel=TOL) + assert rerank.results[3].relevance_score == pytest.approx( + TEXT_VS_TEXT_PLUS_IMAGE, rel=TOL + ) + + +def test_score_api_queries_list_documents_list(server: RemoteOpenAIServer): + score_response = requests.post( + server.url_for("score"), + json={ + "model": MODEL_NAME, + "queries": [query] * 4, + "documents": [ + document, + {"content": [documents[0]]}, + {"content": [documents[1]]}, + {"content": [documents[0], documents[1]]}, + ], + }, + ) + score_response.raise_for_status() + score = ScoreResponse.model_validate(score_response.json()) + + assert score.id is not None + assert score.data is not None + assert len(score.data) == 4 + assert score.usage.prompt_tokens == 368 + assert score.data[0].score == pytest.approx(TEXT_VS_TEXT, rel=TOL) + assert score.data[1].score == pytest.approx(TEXT_VS_TEXT, rel=TOL) + assert score.data[2].score == pytest.approx(TEXT_VS_IMAGE, rel=TOL) + assert score.data[3].score == pytest.approx(TEXT_VS_TEXT_PLUS_IMAGE, rel=TOL) diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index 5ee86ee72f6..f3f774bef36 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -40,12 +40,12 @@ from vllm.entrypoints.chat_utils import ( ChatTemplateContentFormatOption, ) from vllm.entrypoints.pooling.score.utils import ( - ScoreContentPartParam, + ScoreData, ScoreMultiModalParam, _cosine_similarity, - _validate_score_input_lens, compress_token_type_ids, get_score_prompt, + validate_score_input, ) from vllm.entrypoints.utils import log_non_default_args from vllm.inputs import ( @@ -1326,8 +1326,8 @@ class LLM: def _embedding_score( self, - text_1: list[SingletonPrompt], - text_2: list[SingletonPrompt], + data_1: list[ScoreData], + data_2: list[ScoreData], *, use_tqdm: bool | Callable[..., tqdm], pooling_params: PoolingParams | None, @@ -1336,8 +1336,16 @@ class LLM: ) -> list[ScoringRequestOutput]: tokenizer = self.get_tokenizer() + input_texts: list[str] = [] + for text in data_1 + data_2: + if not isinstance(text, str): + raise NotImplementedError( + "Embedding scores currently do not support multimodal input." + ) + input_texts.append(text) + encoded_output = self.encode( - text_1 + text_2, + input_texts, use_tqdm=use_tqdm, lora_request=lora_request, pooling_params=pooling_params, @@ -1345,8 +1353,8 @@ class LLM: tokenization_kwargs=tokenization_kwargs, ) - encoded_output_1 = encoded_output[0 : len(text_1)] - encoded_output_2 = encoded_output[len(text_1) :] + encoded_output_1 = encoded_output[0 : len(data_1)] + encoded_output_2 = encoded_output[len(data_1) :] if len(encoded_output_1) == 1: encoded_output_1 = encoded_output_1 * len(encoded_output_2) @@ -1362,8 +1370,8 @@ class LLM: def _cross_encoding_score( self, - data_1: list[str] | list[ScoreContentPartParam], - data_2: list[str] | list[ScoreContentPartParam], + data_1: list[ScoreData], + data_2: list[ScoreData], *, use_tqdm: bool | Callable[..., tqdm], pooling_params: PoolingParams | None, @@ -1424,8 +1432,14 @@ class LLM: def score( self, - data_1: SingletonPrompt | Sequence[SingletonPrompt] | ScoreMultiModalParam, - data_2: SingletonPrompt | Sequence[SingletonPrompt] | ScoreMultiModalParam, + data_1: SingletonPrompt + | Sequence[SingletonPrompt] + | ScoreMultiModalParam + | list[ScoreMultiModalParam], + data_2: SingletonPrompt + | Sequence[SingletonPrompt] + | ScoreMultiModalParam + | list[ScoreMultiModalParam], /, *, use_tqdm: bool | Callable[..., tqdm] = True, @@ -1501,73 +1515,23 @@ class LLM: "chat_template is only supported for cross-encoder models." ) - # the tokenizer for models such as - # "cross-encoder/ms-marco-MiniLM-L-6-v2" doesn't support passing - # lists of tokens to the `text` and `text_pair` kwargs - tokenizer = self.get_tokenizer() + is_multimodal_model = model_config.is_multimodal_model + architecture = model_config.architecture - if not model_config.is_multimodal_model: - - def check_data_type( - data: SingletonPrompt - | Sequence[SingletonPrompt] - | ScoreMultiModalParam, - ): - if isinstance(data, dict) and "content" in data: - raise ValueError( - "ScoreMultiModalParam is not supported " - f"for {model_config.architecture}" - ) - - check_data_type(data_1) - check_data_type(data_2) - - def ensure_str(prompt: SingletonPrompt): - if isinstance(prompt, dict): - if "multi_modal_data" in prompt: - raise ValueError( - "Multi-modal prompt is not supported for scoring" - ) - elif "prompt_token_ids" in prompt: - prompt = tokenizer.decode( - cast(TokensPrompt, prompt)["prompt_token_ids"] - ) - elif "prompt" in prompt: - prompt = cast(TextPrompt, prompt)["prompt"] - assert type(prompt) is str - return prompt - - if isinstance(data_1, (str, dict)): - # Convert a single prompt to a list. - data_1 = [data_1] # type: ignore[list-item] - - data_1 = [ensure_str(t) for t in data_1] - - if isinstance(data_2, (str, dict)): - # Convert a single prompt to a list. - data_2 = [data_2] # type: ignore[list-item] - - data_2 = [ensure_str(t) for t in data_2] - - if isinstance(data_1, dict) and "content" in data_1: - data_1 = data_1.get("content") # type: ignore[assignment] - elif isinstance(data_1, str): - data_1 = [data_1] - - if isinstance(data_2, dict) and "content" in data_2: - data_2 = data_2.get("content") # type: ignore[assignment] - elif isinstance(data_2, str): - data_2 = [data_2] - - _validate_score_input_lens(data_1, data_2) # type: ignore[arg-type] + score_data_1, score_data_2 = validate_score_input( + data_1, # type: ignore[arg-type] + data_2, # type: ignore[arg-type] + is_multimodal_model=is_multimodal_model, + architecture=architecture, + ) tok_params = self._get_cmpl_tok_params(tokenization_kwargs) encode_kwargs = tok_params.get_encode_kwargs() if model_config.is_cross_encoder: return self._cross_encoding_score( - data_1, # type: ignore[arg-type] - data_2, # type: ignore[arg-type] + score_data_1, + score_data_2, use_tqdm=use_tqdm, pooling_params=pooling_params, lora_request=lora_request, @@ -1576,8 +1540,8 @@ class LLM: ) else: return self._embedding_score( - data_1, # type: ignore[arg-type] - data_2, # type: ignore[arg-type] + score_data_1, + score_data_2, use_tqdm=use_tqdm, pooling_params=pooling_params, lora_request=lora_request, diff --git a/vllm/entrypoints/pooling/score/protocol.py b/vllm/entrypoints/pooling/score/protocol.py index 8f30126b3f5..1a7b0520327 100644 --- a/vllm/entrypoints/pooling/score/protocol.py +++ b/vllm/entrypoints/pooling/score/protocol.py @@ -14,7 +14,8 @@ from vllm.entrypoints.pooling.base.protocol import ( ) from vllm.entrypoints.pooling.score.utils import ( ScoreContentPartParam, - ScoreMultiModalParam, + ScoreInput, + ScoreInputs, ) from vllm.renderers import TokenizeParams from vllm.utils import random_uuid @@ -47,13 +48,13 @@ class ScoreRequestMixin(PoolingBasicRequestMixin, ClassifyRequestMixin): class ScoreDataRequest(ScoreRequestMixin): - data_1: list[str] | str | ScoreMultiModalParam - data_2: list[str] | str | ScoreMultiModalParam + data_1: ScoreInputs + data_2: ScoreInputs class ScoreQueriesDocumentsRequest(ScoreRequestMixin): - queries: list[str] | str | ScoreMultiModalParam - documents: list[str] | str | ScoreMultiModalParam + queries: ScoreInputs + documents: ScoreInputs @property def data_1(self): @@ -64,9 +65,22 @@ class ScoreQueriesDocumentsRequest(ScoreRequestMixin): return self.documents +class ScoreQueriesItemsRequest(ScoreRequestMixin): + queries: ScoreInputs + items: ScoreInputs + + @property + def data_1(self): + return self.queries + + @property + def data_2(self): + return self.items + + class ScoreTextRequest(ScoreRequestMixin): - text_1: list[str] | str | ScoreMultiModalParam - text_2: list[str] | str | ScoreMultiModalParam + text_1: ScoreInputs + text_2: ScoreInputs @property def data_1(self): @@ -78,13 +92,16 @@ class ScoreTextRequest(ScoreRequestMixin): ScoreRequest: TypeAlias = ( - ScoreQueriesDocumentsRequest | ScoreDataRequest | ScoreTextRequest + ScoreQueriesDocumentsRequest + | ScoreQueriesItemsRequest + | ScoreDataRequest + | ScoreTextRequest ) class RerankRequest(PoolingBasicRequestMixin, ClassifyRequestMixin): - query: str | ScoreMultiModalParam - documents: list[str] | ScoreMultiModalParam + query: ScoreInput + documents: ScoreInputs top_n: int = Field(default_factory=lambda: 0) # --8<-- [start:rerank-extra-params] @@ -108,7 +125,7 @@ class RerankRequest(PoolingBasicRequestMixin, ClassifyRequestMixin): class RerankDocument(BaseModel): text: str | None = None - multi_modal: ScoreContentPartParam | None = None + multi_modal: list[ScoreContentPartParam] | None = None class RerankResult(BaseModel): diff --git a/vllm/entrypoints/pooling/score/serving.py b/vllm/entrypoints/pooling/score/serving.py index 1bd28f8bdcc..c32f5470d45 100644 --- a/vllm/entrypoints/pooling/score/serving.py +++ b/vllm/entrypoints/pooling/score/serving.py @@ -27,12 +27,12 @@ from vllm.entrypoints.pooling.score.protocol import ( ScoreResponseData, ) from vllm.entrypoints.pooling.score.utils import ( - ScoreContentPartParam, - ScoreMultiModalParam, + ScoreData, + ScoreInputs, _cosine_similarity, - _validate_score_input_lens, compress_token_type_ids, get_score_prompt, + validate_score_input, ) from vllm.inputs.data import TokensPrompt from vllm.logger import init_logger @@ -65,15 +65,32 @@ class ServingScores(OpenAIServing): self._tokenizer_executor = ThreadPoolExecutor(max_workers=1) + self.is_cross_encoder = self.model_config.is_cross_encoder + self.is_multimodal_model = self.model_config.is_multimodal_model + self.architecture = self.model_config.architecture + + if self.is_cross_encoder: + self._score_func = self._cross_encoding_score + else: + self._score_func = self._embedding_score + async def _embedding_score( self, - data_1: list[str], - data_2: list[str], + data_1: list[ScoreData], + data_2: list[ScoreData], request: RerankRequest | ScoreRequest, request_id: str, lora_request: LoRARequest | None | None = None, trace_headers: Mapping[str, str] | None = None, ) -> list[PoolingRequestOutput] | ErrorResponse: + input_texts: list[str] = [] + for text in data_1 + data_2: + if not isinstance(text, str): + raise NotImplementedError( + "Embedding scores currently do not support multimodal input." + ) + input_texts.append(text) + model_config = self.model_config tokenizer = self.renderer.get_tokenizer() @@ -82,8 +99,6 @@ class ServingScores(OpenAIServing): executor=self._tokenizer_executor, ) - input_texts = data_1 + data_2 - tokenization_kwargs = request.build_tok_params(model_config).get_encode_kwargs() tokenized_prompts = await asyncio.gather( *(encode_async(t, **tokenization_kwargs) for t in input_texts) @@ -157,60 +172,30 @@ class ServingScores(OpenAIServing): return final_res_batch - def _preprocess_score( - self, - request: RerankRequest | ScoreRequest, - tokenizer: TokenizerLike, - tokenization_kwargs: dict[str, Any], - data_1: str | ScoreContentPartParam, - data_2: str | ScoreContentPartParam, - ) -> tuple[str, TokensPrompt]: - model_config = self.model_config - - full_prompt, engine_prompt = get_score_prompt( - model_config=model_config, - data_1=data_1, - data_2=data_2, - tokenizer=tokenizer, - tokenization_kwargs=tokenization_kwargs, - score_template=self.score_template, - ) - self._validate_input(request, engine_prompt["prompt_token_ids"], full_prompt) - if request.mm_processor_kwargs is not None: - engine_prompt["mm_processor_kwargs"] = request.mm_processor_kwargs - - return full_prompt, engine_prompt - async def _cross_encoding_score( self, - data_1: list[str] | list[ScoreContentPartParam], - data_2: list[str] | list[ScoreContentPartParam], + data_1: list[ScoreData], + data_2: list[ScoreData], request: RerankRequest | ScoreRequest, request_id: str, lora_request: LoRARequest | None | None = None, trace_headers: Mapping[str, str] | None = None, ) -> list[PoolingRequestOutput] | ErrorResponse: - model_config = self.model_config tokenizer = self.renderer.get_tokenizer() + if isinstance(tokenizer, MistralTokenizer): + raise ValueError("MistralTokenizer not supported for cross-encoding") - request_prompts: list[str] = [] - engine_prompts: list[TokensPrompt] = [] + model_config = self.model_config if len(data_1) == 1: data_1 = data_1 * len(data_2) - if isinstance(tokenizer, MistralTokenizer): - raise ValueError("MistralTokenizer not supported for cross-encoding") - tok_kwargs = request.build_tok_params(model_config).get_encode_kwargs() - input_pairs = [(t1, t2) for t1, t2 in zip(data_1, data_2)] - preprocess_async = make_async( self._preprocess_score, executor=self._tokenizer_executor, ) - preprocessed_prompts = await asyncio.gather( *( preprocess_async( @@ -224,6 +209,8 @@ class ServingScores(OpenAIServing): ) ) + request_prompts: list[str] = [] + engine_prompts: list[TokensPrompt] = [] for full_prompt, engine_prompt in preprocessed_prompts: request_prompts.append(full_prompt) engine_prompts.append(engine_prompt) @@ -278,10 +265,33 @@ class ServingScores(OpenAIServing): return [out for out in final_res_batch if out is not None] + def _preprocess_score( + self, + request: RerankRequest | ScoreRequest, + tokenizer: TokenizerLike, + tokenization_kwargs: dict[str, Any], + data_1: ScoreData, + data_2: ScoreData, + ) -> tuple[str, TokensPrompt]: + model_config = self.model_config + full_prompt, engine_prompt = get_score_prompt( + model_config=model_config, + data_1=data_1, + data_2=data_2, + tokenizer=tokenizer, + tokenization_kwargs=tokenization_kwargs, + score_template=self.score_template, + ) + self._validate_input(request, engine_prompt["prompt_token_ids"], full_prompt) + if request.mm_processor_kwargs is not None: + engine_prompt["mm_processor_kwargs"] = request.mm_processor_kwargs + + return full_prompt, engine_prompt + async def _run_scoring( self, - data_1: list[str] | str | ScoreMultiModalParam, - data_2: list[str] | str | ScoreMultiModalParam, + data_1: ScoreInputs, + data_2: ScoreInputs, request: ScoreRequest | RerankRequest, request_id: str, raw_request: Request | None = None, @@ -294,44 +304,21 @@ class ServingScores(OpenAIServing): else await self._get_trace_headers(raw_request.headers) ) - if not self.model_config.is_multimodal_model and ( - isinstance(data_1, dict) or isinstance(data_2, dict) - ): - raise ValueError( - f"MultiModalParam is not supported for {self.model_config.architecture}" # noqa: E501 - ) + score_data_1, score_data_2 = validate_score_input( + data_1, + data_2, + is_multimodal_model=self.is_multimodal_model, + architecture=self.architecture, + ) - if isinstance(data_1, str): - data_1 = [data_1] - elif isinstance(data_1, dict): - data_1 = data_1.get("content") # type: ignore[assignment] - - if isinstance(data_2, str): - data_2 = [data_2] - elif isinstance(data_2, dict): - data_2 = data_2.get("content") # type: ignore[assignment] - - _validate_score_input_lens(data_1, data_2) # type: ignore[arg-type] - - if self.model_config.is_cross_encoder: - return await self._cross_encoding_score( - data_1=data_1, # type: ignore[arg-type] - data_2=data_2, # type: ignore[arg-type] - request=request, - request_id=request_id, - lora_request=lora_request, - trace_headers=trace_headers, - ) - - else: - return await self._embedding_score( - data_1=data_1, # type: ignore[arg-type] - data_2=data_2, # type: ignore[arg-type] - request=request, - request_id=request_id, - lora_request=lora_request, - trace_headers=trace_headers, - ) + return await self._score_func( + data_1=score_data_1, + data_2=score_data_2, + request=request, + request_id=request_id, + lora_request=lora_request, + trace_headers=trace_headers, + ) async def create_score( self, @@ -391,15 +378,6 @@ class ServingScores(OpenAIServing): request_id = f"rerank-{self._base_request_id(raw_request)}" documents = request.documents - top_n = ( - request.top_n - if request.top_n > 0 - else ( - len(documents) - if isinstance(documents, list) - else len(documents["content"]) - ) - ) try: final_res_batch = await self._run_scoring( @@ -412,6 +390,8 @@ class ServingScores(OpenAIServing): if isinstance(final_res_batch, ErrorResponse): return final_res_batch + top_n = request.top_n if request.top_n > 0 else len(final_res_batch) + return self.request_output_to_rerank_response( final_res_batch, request_id, @@ -465,22 +445,32 @@ class ServingScores(OpenAIServing): final_res_batch: list[PoolingRequestOutput], request_id: str, model_name: str, - documents: list[str] | ScoreMultiModalParam, + documents: ScoreInputs, top_n: int, ) -> RerankResponse: """ Convert the output of do_rank to a RerankResponse """ + + if not isinstance(documents, list): + documents = [documents] + results: list[RerankResult] = [] num_prompt_tokens = 0 for idx, final_res in enumerate(final_res_batch): classify_res = ScoringRequestOutput.from_base(final_res) + document = documents[idx] + if isinstance(document, str): + rerank_document = RerankDocument(text=document) + else: + rerank_document = RerankDocument( + multi_modal=document.get("content", []) + ) + result = RerankResult( index=idx, - document=RerankDocument(text=documents[idx]) - if isinstance(documents, list) - else RerankDocument(multi_modal=documents["content"][idx]), + document=rerank_document, relevance_score=classify_res.outputs.score, ) results.append(result) diff --git a/vllm/entrypoints/pooling/score/utils.py b/vllm/entrypoints/pooling/score/utils.py index 4387ab2216a..bf3bfe8a878 100644 --- a/vllm/entrypoints/pooling/score/utils.py +++ b/vllm/entrypoints/pooling/score/utils.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Iterable from typing import Any, TypeAlias, cast from torch.nn import CosineSimilarity @@ -10,12 +11,13 @@ from vllm.entrypoints.chat_utils import ( BaseMultiModalItemTracker, ChatCompletionContentPartImageEmbedsParam, ChatCompletionContentPartImageParam, + ChatCompletionContentPartParam, ChatCompletionContentPartTextParam, ChatCompletionContentPartVideoParam, ChatTemplateResolutionError, + ConversationMessage, MultiModalItemTracker, - _ContentPart, - _parse_chat_message_content_part, + _parse_chat_message_content_parts, ) from vllm.inputs import TokensPrompt from vllm.model_executor.models.interfaces import supports_score_template @@ -46,6 +48,13 @@ class ScoreMultiModalParam(TypedDict, total=False): """The multimodal contents""" +# Raw input data with content key in ScoreMultiModalParam. +ScoreInput = str | ScoreMultiModalParam +ScoreInputs = ScoreInput | list[ScoreInput] +# Score data without content key. +ScoreData = str | list[ScoreContentPartParam] + + def _cosine_similarity( tokenizer: TokenizerLike, embed_1: list[PoolingRequestOutput], @@ -77,8 +86,8 @@ def _cosine_similarity( def _validate_score_input_lens( - data_1: list[str] | list[ScoreContentPartParam], - data_2: list[str] | list[ScoreContentPartParam], + data_1: list[ScoreData], + data_2: list[ScoreData], ): len_1 = len(data_1) len_2 = len(data_2) @@ -91,19 +100,56 @@ def _validate_score_input_lens( raise ValueError("At least one text_pair element must be given") +def _validate_mm_score_input( + data: list[ScoreInput], + is_multimodal_model: bool, + architecture: str, +) -> list[ScoreData]: + out: list[ScoreData] = [] + for d in data: + if isinstance(d, str): + out.append(d) + else: + if not is_multimodal_model: + raise ValueError(f"MultiModalParam is not supported for {architecture}") + content = cast(list[ScoreContentPartParam], d.get("content", [])) + out.append(content) + return out + + +def validate_score_input( + data_1: ScoreInputs, + data_2: ScoreInputs, + is_multimodal_model: bool, + architecture: str, +) -> tuple[list[ScoreData], list[ScoreData]]: + if not isinstance(data_1, list): + data_1 = [data_1] + + if not isinstance(data_2, list): + data_2 = [data_2] + + score_input_1 = _validate_mm_score_input(data_1, is_multimodal_model, architecture) + score_input_2 = _validate_mm_score_input(data_2, is_multimodal_model, architecture) + _validate_score_input_lens(score_input_1, score_input_2) + return score_input_1, score_input_2 + + def parse_score_data( - data_1: str | ScoreContentPartParam, - data_2: str | ScoreContentPartParam, + data_1: ScoreData, + data_2: ScoreData, model_config: ModelConfig, ) -> tuple[str, str, MultiModalDataDict | None, MultiModalUUIDDict | None]: mm_tracker = MultiModalItemTracker(model_config) - content_1 = _parse_score_content(data_1, mm_tracker) - content_2 = _parse_score_content(data_2, mm_tracker) + content_1 = _parse_score_content("query", data_1, mm_tracker) + content_2 = _parse_score_content("document", data_2, mm_tracker) - def ensure_str(content: _ContentPart | None) -> str: - if content is not None and isinstance(content, str): - return cast(str, content) + def ensure_str(content: list[ConversationMessage]) -> str: + assert len(content) == 1 + prompt = content[0]["content"] + if prompt is not None and isinstance(prompt, str): + return cast(str, prompt) else: raise ValueError(f"Only string content is supported, but got {content}.") @@ -115,19 +161,22 @@ def parse_score_data( def _parse_score_content( - data: str | ScoreContentPartParam, + role: str, + data: ScoreData, mm_tracker: BaseMultiModalItemTracker, -) -> _ContentPart | None: +) -> list[ConversationMessage]: + parts: Iterable[ChatCompletionContentPartParam] if isinstance(data, str): - part = ChatCompletionContentPartTextParam(type="text", text=data) + parts = [ChatCompletionContentPartTextParam(type="text", text=data)] else: - part = data + parts = cast(Iterable[ChatCompletionContentPartParam], data) mm_parser = mm_tracker.create_parser() - parse_res = _parse_chat_message_content_part( - part, - mm_parser, + parse_res = _parse_chat_message_content_parts( + role=role, + parts=parts, + mm_tracker=mm_tracker, wrap_dicts=False, interleave_strings=False, ) @@ -184,8 +233,8 @@ def get_score_prompt( model_config: ModelConfig, tokenizer: TokenizerLike, tokenization_kwargs: dict[str, Any], - data_1: str | ScoreContentPartParam, - data_2: str | ScoreContentPartParam, + data_1: ScoreData, + data_2: ScoreData, score_template: str | None = None, ) -> tuple[str, TokensPrompt]: prompt_1, prompt_2, mm_data, mm_uuids = parse_score_data( From 02080179a3fc92c339b93040838f44b72313e07b Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Wed, 4 Feb 2026 10:17:37 +0800 Subject: [PATCH 038/810] [Bugfix] Fix torchrun PP broadcast deadlock with async scheduling (#33701) Signed-off-by: Isotr0py --- tests/distributed/test_torchrun_example.py | 3 --- tests/distributed/test_torchrun_example_moe.py | 3 --- vllm/v1/worker/gpu_model_runner.py | 5 ++++- 3 files changed, 4 insertions(+), 7 deletions(-) diff --git a/tests/distributed/test_torchrun_example.py b/tests/distributed/test_torchrun_example.py index 35400951966..f415409d7b3 100644 --- a/tests/distributed/test_torchrun_example.py +++ b/tests/distributed/test_torchrun_example.py @@ -32,9 +32,6 @@ llm = LLM( gpu_memory_utilization=random.uniform(0.7, 0.9), swap_space=random.randint(1, 4), seed=0, - # FIXME(Isotr0py): async scheduling causes deadlock - # on torchrun with PP, need to investigate further. - async_scheduling=False, ) outputs = llm.generate(prompts, sampling_params) diff --git a/tests/distributed/test_torchrun_example_moe.py b/tests/distributed/test_torchrun_example_moe.py index 25f55a968c1..1aa7f179357 100644 --- a/tests/distributed/test_torchrun_example_moe.py +++ b/tests/distributed/test_torchrun_example_moe.py @@ -39,9 +39,6 @@ llm = LLM( gpu_memory_utilization=random.uniform(0.7, 0.9), swap_space=random.randint(1, 4), seed=0, - # FIXME(Isotr0py): async scheduling causes deadlock - # on torchrun with PP, need to investigate further. - async_scheduling=False, ) outputs = llm.generate(prompts, sampling_params) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 49211c6805c..39ac6bce820 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -3666,7 +3666,10 @@ class GPUModelRunner( ) if self.use_async_scheduling: pp = get_pp_group() - if pp.world_size > 1 and pp.is_last_rank: + # For torchrun external_launcher PP mode with broadcast_pp_output=True, + # PP outputs have been broadcasted to all ranks at logits computation. + # Therefore, here is no need to send sampled token ids again in this case. + if not self.broadcast_pp_output and pp.world_size > 1 and pp.is_last_rank: self._pp_broadcast_prev_sampled_token_ids( sampler_output.sampled_token_ids ) From e1bf04b6c27a070859264290ffdccbf333f27fa6 Mon Sep 17 00:00:00 2001 From: Andrew Xia Date: Tue, 3 Feb 2026 21:59:03 -0500 Subject: [PATCH 039/810] [1/N] Initial Implementation of Parser for ResponsesAPI (#32712) Signed-off-by: Andrew Xia Co-authored-by: Andrew Xia --- tests/entrypoints/openai/test_chat_error.py | 1 + .../openai/test_completion_error.py | 1 + tests/entrypoints/openai/test_serving_chat.py | 1 + .../openai/chat_completion/serving.py | 9 +- vllm/entrypoints/openai/engine/serving.py | 43 +-- vllm/entrypoints/openai/responses/serving.py | 38 +- vllm/parser/__init__.py | 39 ++ vllm/parser/abstract_parser.py | 341 ++++++++++++++++++ vllm/parser/minimax_m2_parser.py | 52 +++ vllm/parser/parser_manager.py | 308 ++++++++++++++++ 10 files changed, 772 insertions(+), 61 deletions(-) create mode 100644 vllm/parser/__init__.py create mode 100644 vllm/parser/abstract_parser.py create mode 100644 vllm/parser/minimax_m2_parser.py create mode 100644 vllm/parser/parser_manager.py diff --git a/tests/entrypoints/openai/test_chat_error.py b/tests/entrypoints/openai/test_chat_error.py index 7b15421fb55..de5d96d5ba1 100644 --- a/tests/entrypoints/openai/test_chat_error.py +++ b/tests/entrypoints/openai/test_chat_error.py @@ -36,6 +36,7 @@ class MockHFConfig: class MockModelConfig: task = "generate" runner_type = "generate" + model = MODEL_NAME tokenizer = MODEL_NAME trust_remote_code = False tokenizer_mode = "auto" diff --git a/tests/entrypoints/openai/test_completion_error.py b/tests/entrypoints/openai/test_completion_error.py index 01c4e567c9f..b60397cd769 100644 --- a/tests/entrypoints/openai/test_completion_error.py +++ b/tests/entrypoints/openai/test_completion_error.py @@ -36,6 +36,7 @@ class MockHFConfig: class MockModelConfig: task = "generate" runner_type = "generate" + model = MODEL_NAME tokenizer = MODEL_NAME trust_remote_code = False tokenizer_mode = "auto" diff --git a/tests/entrypoints/openai/test_serving_chat.py b/tests/entrypoints/openai/test_serving_chat.py index b966e7dd7e3..4365075f62e 100644 --- a/tests/entrypoints/openai/test_serving_chat.py +++ b/tests/entrypoints/openai/test_serving_chat.py @@ -511,6 +511,7 @@ class MockHFConfig: class MockModelConfig: task = "generate" runner_type = "generate" + model = MODEL_NAME tokenizer = MODEL_NAME trust_remote_code = False tokenizer_mode = "auto" diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index e618b11adf6..21bc0f44245 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -71,6 +71,7 @@ from vllm.inputs.data import EmbedsPrompt, TokensPrompt from vllm.logger import init_logger from vllm.logprobs import Logprob from vllm.outputs import CompletionOutput, RequestOutput +from vllm.parser import ParserManager from vllm.sampling_params import BeamSearchParams, SamplingParams from vllm.tokenizers import TokenizerLike from vllm.tokenizers.mistral import ( @@ -131,13 +132,15 @@ class OpenAIServingChat(OpenAIServing): self.logits_processors = self.model_config.logits_processors # set up reasoning parser - self.reasoning_parser = self._get_reasoning_parser( + self.reasoning_parser = ParserManager.get_reasoning_parser( reasoning_parser_name=reasoning_parser ) # set up tool use self.enable_auto_tools: bool = enable_auto_tools - self.tool_parser = self._get_tool_parser( - tool_parser_name=tool_parser, enable_auto_tools=enable_auto_tools + self.tool_parser = ParserManager.get_tool_parser( + tool_parser_name=tool_parser, + enable_auto_tools=enable_auto_tools, + model_name=self.model_config.model, ) self.exclude_tools_when_tool_choice_none = exclude_tools_when_tool_choice_none diff --git a/vllm/entrypoints/openai/engine/serving.py b/vllm/entrypoints/openai/engine/serving.py index 7f9300a1ac0..801c7dcd52a 100644 --- a/vllm/entrypoints/openai/engine/serving.py +++ b/vllm/entrypoints/openai/engine/serving.py @@ -107,11 +107,10 @@ from vllm.lora.request import LoRARequest from vllm.multimodal import MultiModalDataDict from vllm.outputs import CompletionOutput, PoolingRequestOutput, RequestOutput from vllm.pooling_params import PoolingParams -from vllm.reasoning import ReasoningParser, ReasoningParserManager from vllm.renderers import ChatParams, TokenizeParams, merge_kwargs from vllm.sampling_params import BeamSearchParams, SamplingParams from vllm.tokenizers import TokenizerLike -from vllm.tool_parsers import ToolParser, ToolParserManager +from vllm.tool_parsers import ToolParser from vllm.tracing import ( contains_trace_headers, extract_trace_headers, @@ -246,46 +245,6 @@ class OpenAIServing: self.model_config = self.models.model_config self.max_model_len = self.model_config.max_model_len - def _get_tool_parser( - self, tool_parser_name: str | None = None, enable_auto_tools: bool = False - ) -> Callable[[TokenizerLike], ToolParser] | None: - """Get the tool parser based on the name.""" - parser = None - if not enable_auto_tools or tool_parser_name is None: - return parser - logger.info('"auto" tool choice has been enabled.') - - try: - if tool_parser_name == "pythonic" and self.model_config.model.startswith( - "meta-llama/Llama-3.2" - ): - logger.warning( - "Llama3.2 models may struggle to emit valid pythonic tool calls" - ) - parser = ToolParserManager.get_tool_parser(tool_parser_name) - except Exception as e: - raise TypeError( - "Error: --enable-auto-tool-choice requires " - f"tool_parser:'{tool_parser_name}' which has not " - "been registered" - ) from e - return parser - - def _get_reasoning_parser( - self, - reasoning_parser_name: str, - ) -> Callable[[TokenizerLike], ReasoningParser] | None: - """Get the reasoning parser based on the name.""" - parser = None - if not reasoning_parser_name: - return None - try: - parser = ReasoningParserManager.get_reasoning_parser(reasoning_parser_name) - assert parser is not None - except Exception as e: - raise TypeError(f"{reasoning_parser_name=} has not been registered") from e - return parser - async def beam_search( self, prompt: PromptType, diff --git a/vllm/entrypoints/openai/responses/serving.py b/vllm/entrypoints/openai/responses/serving.py index cd6aa48c30d..32cce3ef4cf 100644 --- a/vllm/entrypoints/openai/responses/serving.py +++ b/vllm/entrypoints/openai/responses/serving.py @@ -123,6 +123,7 @@ from vllm.logger import init_logger from vllm.logprobs import Logprob as SampleLogprob from vllm.logprobs import SampleLogprobs from vllm.outputs import CompletionOutput +from vllm.parser import ParserManager from vllm.sampling_params import SamplingParams, StructuredOutputsParams from vllm.tokenizers import TokenizerLike from vllm.utils import random_uuid @@ -217,8 +218,13 @@ class OpenAIServingResponses(OpenAIServing): self.chat_template_content_format: Final = chat_template_content_format self.enable_log_outputs = enable_log_outputs - self.reasoning_parser = self._get_reasoning_parser( - reasoning_parser_name=reasoning_parser + # Set up the unified parser - either a unified parser or fall back to + # separate parsers accessed through the parser interface + self.parser = ParserManager.get_parser( + tool_parser_name=tool_parser, + reasoning_parser_name=reasoning_parser, + enable_auto_tools=enable_auto_tools, + model_name=self.model_config.model, ) self.enable_prompt_tokens_details = enable_prompt_tokens_details self.enable_force_include_usage = enable_force_include_usage @@ -263,10 +269,6 @@ class OpenAIServingResponses(OpenAIServing): self.tool_call_id_type = "random" self.enable_auto_tools = enable_auto_tools - # set up tool use - self.tool_parser = self._get_tool_parser( - tool_parser_name=tool_parser, enable_auto_tools=enable_auto_tools - ) # HACK(woosuk): This is a hack. We should use a better store. # FIXME: If enable_store=True, this may cause a memory leak since we # never remove responses from the store. @@ -469,9 +471,13 @@ class OpenAIServingResponses(OpenAIServing): context = ParsableContext( response_messages=messages, tokenizer=tokenizer, - reasoning_parser_cls=self.reasoning_parser, + reasoning_parser_cls=self.parser.reasoning_parser_cls + if self.parser + else None, request=request, - tool_parser_cls=self.tool_parser, + tool_parser_cls=self.parser.tool_parser_cls + if self.parser + else None, available_tools=available_tools, chat_template=self.chat_template, chat_template_content_format=self.chat_template_content_format, @@ -479,8 +485,8 @@ class OpenAIServingResponses(OpenAIServing): else: context = SimpleContext() - if self.reasoning_parser is not None: - reasoning_parser = self.reasoning_parser(tokenizer) + if self.parser and self.parser.reasoning_parser_cls is not None: + reasoning_parser = self.parser.reasoning_parser_cls(tokenizer) if ( isinstance( struct_out := sampling_params.structured_outputs, @@ -617,7 +623,7 @@ class OpenAIServingResponses(OpenAIServing): default_template_content_format=self.chat_template_content_format, default_template_kwargs=None, tool_dicts=tool_dicts, - tool_parser=self.tool_parser, + tool_parser=self.parser.tool_parser_cls if self.parser else None, ) return messages, engine_prompts @@ -909,9 +915,9 @@ class OpenAIServingResponses(OpenAIServing): final_output: CompletionOutput, tokenizer: TokenizerLike, ) -> list[ResponseOutputItem]: - if self.reasoning_parser: + if self.parser and self.parser.reasoning_parser_cls: try: - reasoning_parser = self.reasoning_parser(tokenizer) + reasoning_parser = self.parser.reasoning_parser_cls(tokenizer) except RuntimeError as e: logger.exception("Error in reasoning parser creation.") raise e @@ -958,7 +964,7 @@ class OpenAIServingResponses(OpenAIServing): tokenizer=tokenizer, content=content, enable_auto_tools=self.enable_auto_tools, - tool_parser_cls=self.tool_parser, + tool_parser_cls=self.parser.tool_parser_cls if self.parser else None, ) if content or (self.use_harmony and tool_calls): @@ -1339,8 +1345,8 @@ class OpenAIServingResponses(OpenAIServing): current_output_index = 0 current_item_id = "" reasoning_parser = None - if self.reasoning_parser: - reasoning_parser = self.reasoning_parser(tokenizer) + if self.parser and self.parser.reasoning_parser_cls: + reasoning_parser = self.parser.reasoning_parser_cls(tokenizer) previous_text = "" previous_token_ids: list[int] = [] first_delta_sent = False diff --git a/vllm/parser/__init__.py b/vllm/parser/__init__.py new file mode 100644 index 00000000000..8bce3e912cc --- /dev/null +++ b/vllm/parser/__init__.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm.parser.abstract_parser import ( + DelegatingParser, + Parser, + _WrappedParser, +) +from vllm.parser.parser_manager import ParserManager + +__all__ = [ + "Parser", + "DelegatingParser", + "ParserManager", + "_WrappedParser", +] + +_PARSERS_TO_REGISTER = { + "minimax_m2": ( # name + "minimax_m2_parser", # filename + "MiniMaxM2Parser", # class_name + ), +} + +# Register lazy parsers +ParserManager.register_lazy_module( + name="minimax_m2", + module_path="vllm.parser.minimax_m2_parser", + class_name="MiniMaxM2Parser", +) + + +def register_lazy_parsers(): + for name, (file_name, class_name) in _PARSERS_TO_REGISTER.items(): + module_path = f"vllm.parser.{file_name}" + ParserManager.register_lazy_module(name, module_path, class_name) + + +register_lazy_parsers() diff --git a/vllm/parser/abstract_parser.py b/vllm/parser/abstract_parser.py new file mode 100644 index 00000000000..f5cd1430a18 --- /dev/null +++ b/vllm/parser/abstract_parser.py @@ -0,0 +1,341 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from abc import abstractmethod +from collections.abc import Sequence +from functools import cached_property + +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, +) +from vllm.entrypoints.openai.engine.protocol import ( + DeltaMessage, + ExtractedToolCallInformation, +) +from vllm.entrypoints.openai.responses.protocol import ( + ResponsesRequest, +) +from vllm.reasoning.abs_reasoning_parsers import ReasoningParser +from vllm.tokenizers import TokenizerLike +from vllm.tool_parsers.abstract_tool_parser import ToolParser + + +class Parser: + """ + Abstract Parser class that unifies ReasoningParser and ToolParser into + a single interface for parsing model output. + + This class provides a unified way to handle both reasoning extraction + (e.g., chain-of-thought content in tags) and tool call extraction + (e.g., function calls in XML/JSON format) from model outputs. + + Subclasses can either: + 1. Override the abstract methods directly for custom parsing logic + 2. Set `reasoning_parser` and `tool_parser` properties to delegate to + existing parser implementations + + Class Attributes: + reasoning_parser_cls: The ReasoningParser class to use (for compatibility + with code that needs the class, not instance). + tool_parser_cls: The ToolParser class to use (for compatibility with + code that needs the class, not instance). + """ + + # Class-level parser classes for compatibility with existing patterns + # Subclasses should override these if they use specific parser classes + reasoning_parser_cls: type[ReasoningParser] | None = None + tool_parser_cls: type[ToolParser] | None = None + + def __init__(self, tokenizer: TokenizerLike, *args, **kwargs): + """ + Initialize the Parser. + + Args: + tokenizer: The tokenizer used by the model. This is required for + token-based parsing operations. + """ + self.model_tokenizer = tokenizer + self._reasoning_parser: ReasoningParser | None = None + self._tool_parser: ToolParser | None = None + + @cached_property + def vocab(self) -> dict[str, int]: + """Get the vocabulary mapping from tokens to IDs.""" + return self.model_tokenizer.get_vocab() + + @property + def reasoning_parser(self) -> ReasoningParser | None: + """The underlying reasoning parser, if any.""" + return self._reasoning_parser + + @reasoning_parser.setter + def reasoning_parser(self, parser: ReasoningParser | None) -> None: + self._reasoning_parser = parser + + @property + def tool_parser(self) -> ToolParser | None: + """The underlying tool parser, if any.""" + return self._tool_parser + + @tool_parser.setter + def tool_parser(self, parser: ToolParser | None) -> None: + self._tool_parser = parser + + # ========== Reasoning Parser Methods ========== + + @abstractmethod + def is_reasoning_end(self, input_ids: list[int]) -> bool: + """ + Check if the reasoning content ends in the input_ids. + + Used by structured engines like `xgrammar` to check if the + reasoning content ends in the model output. + + Args: + input_ids: The token IDs of the model output. + + Returns: + True if the reasoning content ends in the input_ids. + """ + + def is_reasoning_end_streaming( + self, input_ids: list[int], delta_ids: list[int] + ) -> bool: + """ + Check if the reasoning content ends during a decode step. + + Args: + input_ids: The entire model output token IDs. + delta_ids: The last few computed tokens at the current decode step. + + Returns: + True if the reasoning content ends in the delta_ids. + """ + return self.is_reasoning_end(input_ids) + + @abstractmethod + def extract_content_ids(self, input_ids: list[int]) -> list[int]: + """ + Extract content token IDs from the input_ids. + + This extracts the non-reasoning content (e.g., everything after + the tag). + + Args: + input_ids: The token IDs of the model output. + + Returns: + The extracted content token IDs. + """ + + @abstractmethod + def extract_reasoning( + self, + model_output: str, + request: ChatCompletionRequest | ResponsesRequest, + ) -> tuple[str | None, str | None]: + """ + Extract reasoning content from a complete model-generated string. + + Used for non-streaming responses where we have the entire model + response available before sending to the client. + + Args: + model_output: The complete model-generated string. + request: The request object used to generate the output. + + Returns: + A tuple of (reasoning_content, response_content). + """ + + @abstractmethod + def extract_reasoning_streaming( + self, + previous_text: str, + current_text: str, + delta_text: str, + previous_token_ids: Sequence[int], + current_token_ids: Sequence[int], + delta_token_ids: Sequence[int], + ) -> DeltaMessage | None: + """ + Extract reasoning content from a streaming delta message. + + Args: + previous_text: Text from all previous tokens. + current_text: Text including the current delta. + delta_text: The new text in this delta. + previous_token_ids: Token IDs from previous generation. + current_token_ids: All token IDs including current. + delta_token_ids: The new token IDs in this delta. + + Returns: + A DeltaMessage with reasoning and/or content fields, or None. + """ + + # ========== Tool Parser Methods ========== + + def adjust_request(self, request: ChatCompletionRequest) -> ChatCompletionRequest: + """ + Adjust the request parameters for tool calling. + + Can be overridden by subclasses to modify request parameters + (e.g., setting structured output schemas for tool calling). + + Args: + request: The original request. + + Returns: + The adjusted request. + """ + return request + + @abstractmethod + def extract_tool_calls( + self, + model_output: str, + request: ChatCompletionRequest, + ) -> ExtractedToolCallInformation: + """ + Extract tool calls from a complete model-generated string. + + Used for non-streaming responses. + + Args: + model_output: The complete model-generated string. + request: The request object used to generate the output. + + Returns: + ExtractedToolCallInformation containing the tool calls. + """ + + @abstractmethod + def extract_tool_calls_streaming( + self, + previous_text: str, + current_text: str, + delta_text: str, + previous_token_ids: Sequence[int], + current_token_ids: Sequence[int], + delta_token_ids: Sequence[int], + request: ChatCompletionRequest, + ) -> DeltaMessage | None: + """ + Extract tool calls from a streaming delta message. + + Args: + previous_text: Text from all previous tokens. + current_text: Text including the current delta. + delta_text: The new text in this delta. + previous_token_ids: Token IDs from previous generation. + current_token_ids: All token IDs including current. + delta_token_ids: The new token IDs in this delta. + request: The request object. + + Returns: + A DeltaMessage with tool_calls field, or None. + """ + + +class DelegatingParser(Parser): + """ + A Parser implementation that delegates to separate ReasoningParser and + ToolParser instances. + + This is the recommended base class for creating model-specific parsers + that combine existing reasoning and tool parser implementations. + Subclasses should set `self._reasoning_parser` and `self._tool_parser` + in their `__init__` method. + + If either parser is None, the corresponding methods will return default + values (no reasoning extraction, no tool calls). + """ + + def extract_reasoning( + self, + model_output: str, + request: ChatCompletionRequest | ResponsesRequest, + ) -> tuple[str | None, str | None]: + if self._reasoning_parser is None: + return None, model_output + return self._reasoning_parser.extract_reasoning(model_output, request) + + def extract_reasoning_streaming( + self, + previous_text: str, + current_text: str, + delta_text: str, + previous_token_ids: Sequence[int], + current_token_ids: Sequence[int], + delta_token_ids: Sequence[int], + ) -> DeltaMessage | None: + if self._reasoning_parser is None: + return DeltaMessage(content=delta_text) + return self._reasoning_parser.extract_reasoning_streaming( + previous_text, + current_text, + delta_text, + previous_token_ids, + current_token_ids, + delta_token_ids, + ) + + def extract_tool_calls( + self, + model_output: str, + request: ChatCompletionRequest, + ) -> ExtractedToolCallInformation: + if self._tool_parser is None: + return ExtractedToolCallInformation( + tools_called=False, tool_calls=[], content=model_output + ) + return self._tool_parser.extract_tool_calls(model_output, request) + + def extract_tool_calls_streaming( + self, + previous_text: str, + current_text: str, + delta_text: str, + previous_token_ids: Sequence[int], + current_token_ids: Sequence[int], + delta_token_ids: Sequence[int], + request: ChatCompletionRequest, + ) -> DeltaMessage | None: + if self._tool_parser is None: + return None + return self._tool_parser.extract_tool_calls_streaming( + previous_text, + current_text, + delta_text, + previous_token_ids, + current_token_ids, + delta_token_ids, + request, + ) + + +class _WrappedParser(DelegatingParser): + """ + A DelegatingParser subclass that instantiates parsers from class attributes. + + This class is used to dynamically create a parser that wraps individual + ReasoningParser and ToolParser classes. The class attributes + `reasoning_parser_cls` and `tool_parser_cls` should be set before + instantiation. + + Usage: + _WrappedParser.reasoning_parser_cls = MyReasoningParser + _WrappedParser.tool_parser_cls = MyToolParser + parser = _WrappedParser(tokenizer) + """ + + reasoning_parser_cls: type[ReasoningParser] | None = None + tool_parser_cls: type[ToolParser] | None = None + + def __init__(self, tokenizer: TokenizerLike): + super().__init__(tokenizer) + # Instantiate the underlying parsers from class attributes + if self.__class__.reasoning_parser_cls is not None: + self._reasoning_parser = self.__class__.reasoning_parser_cls(tokenizer) + if self.__class__.tool_parser_cls is not None: + self._tool_parser = self.__class__.tool_parser_cls(tokenizer) diff --git a/vllm/parser/minimax_m2_parser.py b/vllm/parser/minimax_m2_parser.py new file mode 100644 index 00000000000..ee092d4f542 --- /dev/null +++ b/vllm/parser/minimax_m2_parser.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +MiniMax M2 Parser - A unified parser for MiniMax M2 models. + +This parser combines the existing MiniMaxM2ReasoningParser and +MinimaxM2ToolParser into a single unified interface by delegating +to those implementations. +""" + +from vllm.logger import init_logger +from vllm.parser.abstract_parser import DelegatingParser +from vllm.reasoning.minimax_m2_reasoning_parser import MiniMaxM2ReasoningParser +from vllm.tokenizers import TokenizerLike +from vllm.tool_parsers.minimax_m2_tool_parser import MinimaxM2ToolParser + +logger = init_logger(__name__) + + +class MiniMaxM2Parser(DelegatingParser): + """ + Unified parser for MiniMax M2 models that handles both reasoning + extraction and tool call parsing. + + This parser delegates to the existing implementations: + - MiniMaxM2ReasoningParser for reasoning extraction + - MinimaxM2ToolParser for tool call parsing + + MiniMax M2 models have two special behaviors: + 1. Reasoning: They don't generate start token, only end + token. All content before is reasoning, content after is the + actual response. + 2. Tool Calls: They use ... tags + with ... and ... + syntax. + """ + + # Class-level parser classes for compatibility + reasoning_parser_cls = MiniMaxM2ReasoningParser + tool_parser_cls = MinimaxM2ToolParser + + def __init__(self, tokenizer: TokenizerLike): + super().__init__(tokenizer) + + # Initialize the underlying parsers + self._reasoning_parser = MiniMaxM2ReasoningParser(tokenizer) + self._tool_parser = MinimaxM2ToolParser(tokenizer) + + logger.debug( + "vLLM Successfully initialized parser %s!", self.__class__.__name__ + ) diff --git a/vllm/parser/parser_manager.py b/vllm/parser/parser_manager.py new file mode 100644 index 00000000000..4331eba9884 --- /dev/null +++ b/vllm/parser/parser_manager.py @@ -0,0 +1,308 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import importlib +import os +from collections.abc import Callable +from typing import TYPE_CHECKING + +from vllm.logger import init_logger +from vllm.utils.collection_utils import is_list_of +from vllm.utils.import_utils import import_from_path + +if TYPE_CHECKING: + from vllm.parser.abstract_parser import Parser + from vllm.reasoning import ReasoningParser + from vllm.tool_parsers import ToolParser + +logger = init_logger(__name__) + + +class ParserManager: + """ + Central registry for Parser implementations. + + Supports two registration modes: + - Eager registration via `register_module` + - Lazy registration via `register_lazy_module` + """ + + parsers: dict[str, type[Parser]] = {} + lazy_parsers: dict[str, tuple[str, str]] = {} # name -> (module_path, class_name) + + @classmethod + def get_parser_internal(cls, name: str) -> type[Parser]: + """ + Retrieve a registered or lazily registered Parser class. + + Args: + name: The registered name of the parser. + + Returns: + The Parser class. + + Raises: + KeyError: If no parser is found under the given name. + """ + if name in cls.parsers: + return cls.parsers[name] + + if name in cls.lazy_parsers: + return cls._load_lazy_parser(name) + + registered = ", ".join(cls.list_registered()) + raise KeyError(f"Parser '{name}' not found. Available parsers: {registered}") + + @classmethod + def _load_lazy_parser(cls, name: str) -> type[Parser]: + """Import and register a lazily loaded parser.""" + from vllm.parser.abstract_parser import Parser + + module_path, class_name = cls.lazy_parsers[name] + try: + mod = importlib.import_module(module_path) + parser_cls = getattr(mod, class_name) + if not issubclass(parser_cls, Parser): + raise TypeError( + f"{class_name} in {module_path} is not a Parser subclass." + ) + cls.parsers[name] = parser_cls # cache + return parser_cls + except Exception as e: + logger.exception( + "Failed to import lazy parser '%s' from %s: %s", + name, + module_path, + e, + ) + raise + + @classmethod + def _register_module( + cls, + module: type[Parser], + module_name: str | list[str] | None = None, + force: bool = True, + ) -> None: + """Register a Parser class immediately.""" + from vllm.parser.abstract_parser import Parser + + if not issubclass(module, Parser): + raise TypeError( + f"module must be subclass of Parser, but got {type(module)}" + ) + + if module_name is None: + module_names = [module.__name__] + elif isinstance(module_name, str): + module_names = [module_name] + elif is_list_of(module_name, str): + module_names = module_name + else: + raise TypeError("module_name must be str, list[str], or None.") + + for name in module_names: + if not force and name in cls.parsers: + existed = cls.parsers[name] + raise KeyError(f"{name} is already registered at {existed.__module__}") + cls.parsers[name] = module + + @classmethod + def register_lazy_module(cls, name: str, module_path: str, class_name: str) -> None: + """ + Register a lazy module mapping for delayed import. + + Example: + ParserManager.register_lazy_module( + name="minimax_m2", + module_path="vllm.parser.minimax_m2_parser", + class_name="MiniMaxM2Parser", + ) + """ + cls.lazy_parsers[name] = (module_path, class_name) + + @classmethod + def register_module( + cls, + name: str | list[str] | None = None, + force: bool = True, + module: type[Parser] | None = None, + ) -> type[Parser] | Callable[[type[Parser]], type[Parser]]: + """ + Register a Parser class. + + Can be used as a decorator or called directly. + + Usage: + @ParserManager.register_module("my_parser") + class MyParser(Parser): + ... + + Or: + ParserManager.register_module(module=MyParser) + """ + if not isinstance(force, bool): + raise TypeError(f"force must be a boolean, but got {type(force)}") + + # Immediate registration + if module is not None: + cls._register_module(module=module, module_name=name, force=force) + return module + + # Decorator usage + def _decorator(obj: type[Parser]) -> type[Parser]: + module_path = obj.__module__ + class_name = obj.__name__ + + if isinstance(name, str): + names = [name] + elif is_list_of(name, str): + names = name + else: + names = [class_name] + + for n in names: + cls.lazy_parsers[n] = (module_path, class_name) + + return obj + + return _decorator + + @classmethod + def list_registered(cls) -> list[str]: + """Return names of all registered parsers.""" + return sorted(set(cls.parsers.keys()) | set(cls.lazy_parsers.keys())) + + @classmethod + def import_parser(cls, plugin_path: str) -> None: + """Import a user-defined parser from an arbitrary path.""" + module_name = os.path.splitext(os.path.basename(plugin_path))[0] + try: + import_from_path(module_name, plugin_path) + except Exception: + logger.exception( + "Failed to load module '%s' from %s.", module_name, plugin_path + ) + + @classmethod + def get_tool_parser( + cls, + tool_parser_name: str | None = None, + enable_auto_tools: bool = False, + model_name: str | None = None, + ) -> type[ToolParser] | None: + """Get the tool parser based on the name.""" + from vllm.tool_parsers import ToolParserManager + + parser: type[ToolParser] | None = None + if not enable_auto_tools or tool_parser_name is None: + return parser + logger.info('"auto" tool choice has been enabled.') + + try: + if ( + tool_parser_name == "pythonic" + and model_name + and model_name.startswith("meta-llama/Llama-3.2") + ): + logger.warning( + "Llama3.2 models may struggle to emit valid pythonic tool calls" + ) + parser = ToolParserManager.get_tool_parser(tool_parser_name) + except Exception as e: + raise TypeError( + "Error: --enable-auto-tool-choice requires " + f"tool_parser:'{tool_parser_name}' which has not " + "been registered" + ) from e + return parser + + @classmethod + def get_reasoning_parser( + cls, + reasoning_parser_name: str | None, + ) -> type[ReasoningParser] | None: + """Get the reasoning parser based on the name.""" + from vllm.reasoning import ReasoningParserManager + + parser: type[ReasoningParser] | None = None + if not reasoning_parser_name: + return None + try: + parser = ReasoningParserManager.get_reasoning_parser(reasoning_parser_name) + assert parser is not None + except Exception as e: + raise TypeError(f"{reasoning_parser_name=} has not been registered") from e + return parser + + @classmethod + def get_parser( + cls, + tool_parser_name: str | None = None, + reasoning_parser_name: str | None = None, + enable_auto_tools: bool = False, + model_name: str | None = None, + ) -> type[Parser] | None: + """ + Get a unified Parser that handles both reasoning and tool parsing. + + This method checks if a unified Parser exists that can handle both + reasoning extraction and tool call parsing. If no unified parser + exists, it creates a DelegatingParser that wraps the individual + reasoning and tool parsers. + + Args: + tool_parser_name: The name of the tool parser. + reasoning_parser_name: The name of the reasoning parser. + enable_auto_tools: Whether auto tool choice is enabled. + model_name: The model name for parser-specific warnings. + + Returns: + A Parser class, or None if neither parser is specified. + """ + from vllm.parser.abstract_parser import _WrappedParser + + if not tool_parser_name and not reasoning_parser_name: + return None + + # Strategy 1: If both names match, check for a unified parser with that name + if tool_parser_name and tool_parser_name == reasoning_parser_name: + try: + parser = cls.get_parser_internal(tool_parser_name) + logger.info( + "Using unified parser '%s' for both reasoning and tool parsing.", + tool_parser_name, + ) + return parser + except KeyError: + pass # No unified parser with this name + + # Strategy 2: Check for parser with either name + for name in [tool_parser_name, reasoning_parser_name]: + if name: + try: + parser = cls.get_parser_internal(name) + logger.info( + "Using unified parser '%s' for reasoning and tool parsing.", + name, + ) + return parser + except KeyError: + pass + + # Strategy 3: Create a DelegatingParser with the individual parser classes + reasoning_parser_cls = cls.get_reasoning_parser(reasoning_parser_name) + tool_parser_cls = cls.get_tool_parser( + tool_parser_name, enable_auto_tools, model_name + ) + + if reasoning_parser_cls is None and tool_parser_cls is None: + return None + + # Set the class-level attributes on the imported _WrappedParser + _WrappedParser.reasoning_parser_cls = reasoning_parser_cls + _WrappedParser.tool_parser_cls = tool_parser_cls + + return _WrappedParser From 4dffc5e044317326b9e2b2fd2a019c499d63c427 Mon Sep 17 00:00:00 2001 From: R3hankhan Date: Wed, 4 Feb 2026 09:07:15 +0530 Subject: [PATCH 040/810] [CPU] Split attention dispatch by head_dim alignment (#32161) Signed-off-by: Rehan Khan --- cmake/cpu_extension.cmake | 13 ++ csrc/cpu/cpu_attn.cpp | 125 +++----------- csrc/cpu/cpu_attn_amx.hpp | 2 +- csrc/cpu/cpu_attn_neon.hpp | 2 +- csrc/cpu/generate_cpu_attn_dispatch.py | 203 +++++++++++++++++++++++ tests/kernels/attention/test_cpu_attn.py | 3 +- 6 files changed, 241 insertions(+), 107 deletions(-) create mode 100644 csrc/cpu/generate_cpu_attn_dispatch.py diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index 6da4f6c0cdc..c9813a73d91 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -359,6 +359,19 @@ else() add_compile_definitions(-DVLLM_NUMA_DISABLED) endif() +# +# Generate CPU attention dispatch header +# +message(STATUS "Generating CPU attention dispatch header") +execute_process( + COMMAND ${Python_EXECUTABLE} ${CMAKE_SOURCE_DIR}/csrc/cpu/generate_cpu_attn_dispatch.py + WORKING_DIRECTORY ${CMAKE_SOURCE_DIR}/csrc/cpu + RESULT_VARIABLE GEN_RESULT +) +if(NOT GEN_RESULT EQUAL 0) + message(FATAL_ERROR "Failed to generate CPU attention dispatch header") +endif() + # # _C extension # diff --git a/csrc/cpu/cpu_attn.cpp b/csrc/cpu/cpu_attn.cpp index 374fc2ee6dd..641f95a2b1d 100644 --- a/csrc/cpu/cpu_attn.cpp +++ b/csrc/cpu/cpu_attn.cpp @@ -1,79 +1,4 @@ -#include "cpu_attn_vec.hpp" -#include "cpu_attn_vec16.hpp" - -#ifdef CPU_CAPABILITY_AMXBF16 - #include "cpu_attn_amx.hpp" - #define AMX_DISPATCH(...) \ - case cpu_attention::ISA::AMX: { \ - using attn_impl = cpu_attention::AttentionImpl; \ - return __VA_ARGS__(); \ - } -#else - #define AMX_DISPATCH(...) case cpu_attention::ISA::AMX: -#endif - -#ifdef __aarch64__ - #include "cpu_attn_neon.hpp" - // NEON requires head_dim to be a multiple of 32 - #define NEON_DISPATCH(...) \ - case cpu_attention::ISA::NEON: { \ - using attn_impl = cpu_attention::AttentionImpl; \ - return __VA_ARGS__(); \ - } -#else - #define NEON_DISPATCH(...) case cpu_attention::ISA::NEON: -#endif // #ifdef __aarch64__ - -#define CPU_ATTN_DISPATCH_CASE(HEAD_DIM, ...) \ - case HEAD_DIM: { \ - constexpr size_t head_dim = HEAD_DIM; \ - return __VA_ARGS__(); \ - } - -#define CPU_ATTN_DISPATCH_CASE_HEADDIM(HEAD_DIM, ...) \ - [&] { \ - switch (HEAD_DIM) { \ - CPU_ATTN_DISPATCH_CASE(32, __VA_ARGS__) \ - CPU_ATTN_DISPATCH_CASE(64, __VA_ARGS__) \ - CPU_ATTN_DISPATCH_CASE(80, __VA_ARGS__) \ - CPU_ATTN_DISPATCH_CASE(96, __VA_ARGS__) \ - CPU_ATTN_DISPATCH_CASE(112, __VA_ARGS__) \ - CPU_ATTN_DISPATCH_CASE(128, __VA_ARGS__) \ - CPU_ATTN_DISPATCH_CASE(160, __VA_ARGS__) \ - CPU_ATTN_DISPATCH_CASE(192, __VA_ARGS__) \ - CPU_ATTN_DISPATCH_CASE(224, __VA_ARGS__) \ - CPU_ATTN_DISPATCH_CASE(256, __VA_ARGS__) \ - default: { \ - TORCH_CHECK(false, "Invalid CPU attention head_dim: " + \ - std::to_string(HEAD_DIM)); \ - } \ - } \ - }() - -#define CPU_ATTN_DISPATCH_IMPL(ISA_TYPE, ...) \ - [&] { \ - switch (ISA_TYPE) { \ - AMX_DISPATCH(__VA_ARGS__) \ - NEON_DISPATCH(__VA_ARGS__) \ - case cpu_attention::ISA::VEC: { \ - using attn_impl = \ - cpu_attention::AttentionImpl; \ - return __VA_ARGS__(); \ - } \ - case cpu_attention::ISA::VEC16: { \ - using attn_impl = \ - cpu_attention::AttentionImpl; \ - return __VA_ARGS__(); \ - } \ - default: { \ - TORCH_CHECK(false, "Invalid CPU attention ISA type."); \ - } \ - } \ - }() +#include "cpu_attn_dispatch_generated.h" torch::Tensor get_scheduler_metadata( const int64_t num_req, const int64_t num_heads_q, @@ -122,16 +47,14 @@ torch::Tensor get_scheduler_metadata( input.enable_kv_split = enable_kv_split; VLLM_DISPATCH_FLOATING_TYPES(dtype, "get_scheduler_metadata", [&]() { - CPU_ATTN_DISPATCH_CASE_HEADDIM(head_dim, [&] { - CPU_ATTN_DISPATCH_IMPL(isa, [&]() { - input.elem_size = sizeof(scalar_t); - input.q_buffer_elem_size = sizeof(attn_impl::q_buffer_t); - input.logits_buffer_elem_size = sizeof(attn_impl::logits_buffer_t); - input.output_buffer_elem_size = - sizeof(attn_impl::partial_output_buffer_t); - input.max_num_q_per_iter = attn_impl::MaxQHeadNumPerIteration; - input.kv_block_alignment = attn_impl::BlockSizeAlignment; - }); + CPU_ATTN_DISPATCH(head_dim, isa, [&]() { + input.elem_size = sizeof(scalar_t); + input.q_buffer_elem_size = sizeof(attn_impl::q_buffer_t); + input.logits_buffer_elem_size = sizeof(attn_impl::logits_buffer_t); + input.output_buffer_elem_size = + sizeof(attn_impl::partial_output_buffer_t); + input.max_num_q_per_iter = attn_impl::MaxQHeadNumPerIteration; + input.kv_block_alignment = attn_impl::BlockSizeAlignment; }); }); @@ -184,18 +107,14 @@ void cpu_attn_reshape_and_cache( VLLM_DISPATCH_FLOATING_TYPES( key.scalar_type(), "cpu_attn_reshape_and_cache", [&]() { - CPU_ATTN_DISPATCH_CASE_HEADDIM(head_dim, [&] { - CPU_ATTN_DISPATCH_IMPL(isa_tag, [&]() { - attn_impl::reshape_and_cache( - key.data_ptr(), value.data_ptr(), - key_cache.data_ptr(), - value_cache.data_ptr(), - slot_mapping.data_ptr(), token_num, - key_token_num_stride, value_token_num_stride, head_num, - key_head_num_stride, value_head_num_stride, num_blocks, - num_blocks_stride, cache_head_num_stride, block_size, - block_size_stride); - }); + CPU_ATTN_DISPATCH(head_dim, isa_tag, [&]() { + attn_impl::reshape_and_cache( + key.data_ptr(), value.data_ptr(), + key_cache.data_ptr(), value_cache.data_ptr(), + slot_mapping.data_ptr(), token_num, key_token_num_stride, + value_token_num_stride, head_num, key_head_num_stride, + value_head_num_stride, num_blocks, num_blocks_stride, + cache_head_num_stride, block_size, block_size_stride); }); }); } @@ -257,12 +176,10 @@ void cpu_attention_with_kv_cache( VLLM_DISPATCH_FLOATING_TYPES( query.scalar_type(), "cpu_attention_with_kv_cache", [&]() { - CPU_ATTN_DISPATCH_CASE_HEADDIM(query.size(2), [&] { - CPU_ATTN_DISPATCH_IMPL(input.metadata->isa, [&]() { - TORCH_CHECK_EQ(input.block_size % attn_impl::BlockSizeAlignment, 0); - cpu_attention::AttentionMainLoop mainloop; - mainloop(&input); - }); + CPU_ATTN_DISPATCH(query.size(2), input.metadata->isa, [&]() { + TORCH_CHECK_EQ(input.block_size % attn_impl::BlockSizeAlignment, 0); + cpu_attention::AttentionMainLoop mainloop; + mainloop(&input); }); }); } diff --git a/csrc/cpu/cpu_attn_amx.hpp b/csrc/cpu/cpu_attn_amx.hpp index 78be05e8dc8..8da458b9911 100644 --- a/csrc/cpu/cpu_attn_amx.hpp +++ b/csrc/cpu/cpu_attn_amx.hpp @@ -377,7 +377,7 @@ class AttentionImpl { const int32_t q_heads_per_kv, const int64_t q_num_stride, const int64_t q_head_stride, const float scale) { constexpr int64_t bytes_per_head = head_dim * sizeof(scalar_t); - // static_assert(bytes_per_head % AMX_TILE_ROW_BYTES == 0); + static_assert(bytes_per_head % AMX_TILE_ROW_BYTES == 0); constexpr int64_t head_size_block_num = bytes_per_head / AMX_TILE_ROW_BYTES; constexpr int64_t head_elem_num_pre_block = AMX_TILE_ROW_BYTES / sizeof(scalar_t); diff --git a/csrc/cpu/cpu_attn_neon.hpp b/csrc/cpu/cpu_attn_neon.hpp index e9ecd1d3290..827f0cfbc71 100644 --- a/csrc/cpu/cpu_attn_neon.hpp +++ b/csrc/cpu/cpu_attn_neon.hpp @@ -264,7 +264,7 @@ class AttentionImpl { constexpr static ISA ISAType = ISA::NEON; constexpr static bool scale_on_logits = false; // apply scale on q_buffer - // static_assert(HeadDim % HeadDimAlignment == 0); + static_assert(HeadDim % HeadDimAlignment == 0); // the gemm micro kernel is Mx8 static_assert(HeadDimAlignment % 8 == 0); static_assert(BlockSizeAlignment % 8 == 0); diff --git a/csrc/cpu/generate_cpu_attn_dispatch.py b/csrc/cpu/generate_cpu_attn_dispatch.py new file mode 100644 index 00000000000..85f21544df2 --- /dev/null +++ b/csrc/cpu/generate_cpu_attn_dispatch.py @@ -0,0 +1,203 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Generate CPU attention dispatch switch cases and kernel instantiations. +""" + +import os + +# Head dimensions divisible by 32 (support all ISAs) +HEAD_DIMS_32 = [32, 64, 96, 128, 160, 192, 224, 256] + +# Head dimensions divisible by 16 but not 32 (VEC16 only) +HEAD_DIMS_16 = [80, 112] + +# ISA types +ISA_TYPES = { + "AMX": 0, + "VEC": 1, + "VEC16": 2, + "NEON": 3, +} + +# ISAs supported for head_dims divisible by 32 +ISA_FOR_32 = ["AMX", "NEON", "VEC", "VEC16"] + +# ISAs supported for head_dims divisible by 16 only +ISA_FOR_16 = ["VEC16"] + + +def encode_params(head_dim: int, isa_type: str) -> int: + """Encode head_dim and ISA type into a single int64_t.""" + isa_val = ISA_TYPES[isa_type] + # Encoding: (head_dim << 8) | isa_type + # This allows head_dim up to 2^56 - 1 and 256 ISA types + return (head_dim << 8) | isa_val + + +def generate_cases_for_isa_group(isa_list: list[str]) -> str: + """Generate switch cases for a specific ISA group.""" + cases = [] + + # Generate cases for head_dims divisible by 32 + for head_dim in HEAD_DIMS_32: + for isa in isa_list: + if isa not in ISA_FOR_32: + continue + encoded = encode_params(head_dim, isa) + case_str = ( + f""" case {encoded}LL: {{ """ + f"""/* head_dim={head_dim}, isa={isa} */ \\""" + f""" + constexpr size_t head_dim = {head_dim}; \\""" + f""" + using attn_impl = cpu_attention::AttentionImpl<""" + f"""cpu_attention::ISA::{isa}, \\""" + f""" + """ + f"""scalar_t, head_dim>; \\""" + f""" + return __VA_ARGS__(); \\""" + f""" + }} \\""" + ) + cases.append(case_str) + + # Generate cases for head_dims divisible by 16 only + for head_dim in HEAD_DIMS_16: + for isa in isa_list: + encoded = encode_params(head_dim, isa) + case_str = ( + f""" case {encoded}LL: {{ """ + f"""/* head_dim={head_dim}, isa={isa} """ + f"""(using VEC16) */ \\""" + f""" + constexpr size_t head_dim = {head_dim}; \\""" + f""" + using attn_impl = cpu_attention::AttentionImpl<""" + f"""cpu_attention::ISA::VEC16, \\""" + f""" + """ + f"""scalar_t, head_dim>; \\""" + f""" + return __VA_ARGS__(); \\""" + f""" + }} \\""" + ) + cases.append(case_str) + + return "\n".join(cases) + + +def generate_helper_function() -> str: + """Generate helper function to encode parameters.""" + return """ +inline int64_t encode_cpu_attn_params(int64_t head_dim, cpu_attention::ISA isa) { + return (head_dim << 8) | static_cast(isa); +} +""" + + +def generate_header_file() -> str: + """Generate the complete header file content.""" + header = """// auto generated by generate_cpu_attn_dispatch.py +// clang-format off + +#ifndef CPU_ATTN_DISPATCH_GENERATED_H +#define CPU_ATTN_DISPATCH_GENERATED_H + +#include "cpu_attn_vec.hpp" +#include "cpu_attn_vec16.hpp" + +#ifdef CPU_CAPABILITY_AMXBF16 + #include "cpu_attn_amx.hpp" +#endif + +#ifdef __aarch64__ + #include "cpu_attn_neon.hpp" +#endif + +""" + + header += generate_helper_function() + + # Generate dispatch macro with conditional compilation for different ISA sets + header += """ +// Dispatch macro using encoded parameters +""" + + # x86_64 with AMX + header += """#if defined(CPU_CAPABILITY_AMXBF16) +#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, ...) \\ + [&] { \\ + int64_t encoded_params = encode_cpu_attn_params(HEAD_DIM, ISA_TYPE); \\ + switch (encoded_params) { \\ +""" + header += generate_cases_for_isa_group(["AMX", "VEC", "VEC16"]) + header += """ + default: { \\ + TORCH_CHECK(false, "Unsupported CPU attention configuration: head_dim=" + \\ + std::to_string(HEAD_DIM) + " isa=" + \\ + std::to_string(static_cast(ISA_TYPE))); \\ + } \\ + } \\ + }() + +""" + + # ARM64 with NEON + header += """#elif defined(__aarch64__) +#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, ...) \\ + [&] { \\ + int64_t encoded_params = encode_cpu_attn_params(HEAD_DIM, ISA_TYPE); \\ + switch (encoded_params) { \\ +""" + header += generate_cases_for_isa_group(["NEON", "VEC", "VEC16"]) + header += """ + default: { \\ + TORCH_CHECK(false, "Unsupported CPU attention configuration: head_dim=" + \\ + std::to_string(HEAD_DIM) + " isa=" + \\ + std::to_string(static_cast(ISA_TYPE))); \\ + } \\ + } \\ + }() + +""" + + # Fallback: VEC and VEC16 only + header += """#else +#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, ...) \\ + [&] { \\ + int64_t encoded_params = encode_cpu_attn_params(HEAD_DIM, ISA_TYPE); \\ + switch (encoded_params) { \\ +""" + header += generate_cases_for_isa_group(["VEC", "VEC16"]) + header += """ + default: { \\ + TORCH_CHECK(false, "Unsupported CPU attention configuration: head_dim=" + \\ + std::to_string(HEAD_DIM) + " isa=" + \\ + std::to_string(static_cast(ISA_TYPE))); \\ + } \\ + } \\ + }() + +#endif /* CPU_CAPABILITY_AMXBF16 / __aarch64__ */ + +#endif // CPU_ATTN_DISPATCH_GENERATED_H +""" + + return header + + +def main(): + output_path = os.path.join( + os.path.dirname(__file__), "cpu_attn_dispatch_generated.h" + ) + + with open(output_path, "w") as f: + f.write(generate_header_file()) + + +if __name__ == "__main__": + main() diff --git a/tests/kernels/attention/test_cpu_attn.py b/tests/kernels/attention/test_cpu_attn.py index ef0099f635a..9636dfb95ab 100644 --- a/tests/kernels/attention/test_cpu_attn.py +++ b/tests/kernels/attention/test_cpu_attn.py @@ -26,6 +26,7 @@ NUM_HEADS = [ (9, 3), ] HEAD_SIZES = [96, 128] +HEAD_SIZES_VEC16 = [96, 80, 112, 128] QTYPES = [torch.bfloat16, torch.half, torch.float32] SLIDING_WINDOWS = [None, 256] NUM_BLOCKS = [ @@ -432,7 +433,7 @@ def test_varlen_with_paged_kv_normal_amx( @pytest.mark.parametrize("seq_lens", SEQ_LENS) @pytest.mark.parametrize("num_heads", NUM_HEADS) -@pytest.mark.parametrize("head_size", HEAD_SIZES) +@pytest.mark.parametrize("head_size", HEAD_SIZES_VEC16) @pytest.mark.parametrize("block_size", [48]) @pytest.mark.parametrize("sliding_window", SLIDING_WINDOWS) @pytest.mark.parametrize("dtype", [torch.bfloat16]) From 9fb27dd3b3530cbf958a24ee2dcfc907acbf5455 Mon Sep 17 00:00:00 2001 From: Shanshan Shen <467638484@qq.com> Date: Wed, 4 Feb 2026 12:07:30 +0800 Subject: [PATCH 041/810] [MM] Align the prefix of MMEncoderAttention with Attention (#33750) Signed-off-by: shen-shanshan <467638484@qq.com> --- vllm/model_executor/models/aimv2.py | 2 +- vllm/model_executor/models/blip.py | 2 +- vllm/model_executor/models/glm4_1v.py | 2 +- vllm/model_executor/models/glm4v.py | 2 +- vllm/model_executor/models/glm_ocr.py | 1 + vllm/model_executor/models/idefics2_vision_model.py | 2 +- vllm/model_executor/models/intern_vit.py | 2 +- vllm/model_executor/models/interns1_vit.py | 2 +- vllm/model_executor/models/mllama4.py | 2 +- vllm/model_executor/models/molmo.py | 2 +- vllm/model_executor/models/molmo2.py | 2 +- vllm/model_executor/models/openpangu_vl.py | 1 + vllm/model_executor/models/qwen2_5_vl.py | 2 +- vllm/model_executor/models/qwen2_vl.py | 2 +- vllm/model_executor/models/qwen3_omni_moe_thinker.py | 2 +- vllm/model_executor/models/step3_vl.py | 2 +- vllm/model_executor/models/step_vl.py | 2 +- 17 files changed, 17 insertions(+), 15 deletions(-) diff --git a/vllm/model_executor/models/aimv2.py b/vllm/model_executor/models/aimv2.py index d6716c8e580..63cb9c96e2e 100644 --- a/vllm/model_executor/models/aimv2.py +++ b/vllm/model_executor/models/aimv2.py @@ -130,7 +130,7 @@ class AIMv2Attention(nn.Module): self.num_heads_per_partition, self.head_dim, self.scale, - prefix=prefix, + prefix=f"{prefix}.attn", ) def forward(self, x: torch.Tensor) -> torch.Tensor: diff --git a/vllm/model_executor/models/blip.py b/vllm/model_executor/models/blip.py index ad8f7c1af54..73b0b8af930 100644 --- a/vllm/model_executor/models/blip.py +++ b/vllm/model_executor/models/blip.py @@ -126,7 +126,7 @@ class BlipAttention(nn.Module): self.num_heads_per_partition, self.head_dim, self.scale, - prefix=prefix, + prefix=f"{prefix}.attn", ) def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): diff --git a/vllm/model_executor/models/glm4_1v.py b/vllm/model_executor/models/glm4_1v.py index b2886e85a3b..5333042cbf9 100644 --- a/vllm/model_executor/models/glm4_1v.py +++ b/vllm/model_executor/models/glm4_1v.py @@ -296,7 +296,7 @@ class Glm4vVisionAttention(nn.Module): num_heads=self.num_attention_heads_per_partition, head_size=self.hidden_size_per_attention_head, scale=self.hidden_size_per_attention_head**-0.5, - prefix=prefix, + prefix=f"{prefix}.attn", ) self.apply_rotary_emb = ApplyRotaryEmb(enforce_enable=True) diff --git a/vllm/model_executor/models/glm4v.py b/vllm/model_executor/models/glm4v.py index 8bad386fa8f..56504029dc5 100644 --- a/vllm/model_executor/models/glm4v.py +++ b/vllm/model_executor/models/glm4v.py @@ -139,7 +139,7 @@ class EVA2CLIPAttention(nn.Module): self.num_heads_per_rank, self.head_dim, self.scale, - prefix=prefix, + prefix=f"{prefix}.attn", ) self.output_dropout = torch.nn.Dropout(config.dropout_prob) diff --git a/vllm/model_executor/models/glm_ocr.py b/vllm/model_executor/models/glm_ocr.py index 90c6baacbd2..d037431403b 100644 --- a/vllm/model_executor/models/glm_ocr.py +++ b/vllm/model_executor/models/glm_ocr.py @@ -137,6 +137,7 @@ class GlmOcrVisionAttention(nn.Module): num_heads=self.num_attention_heads_per_partition, head_size=self.hidden_size_per_attention_head, scale=self.hidden_size_per_attention_head**-0.5, + prefix=f"{prefix}.attn", ) self.apply_rotary_emb = ApplyRotaryEmb(enforce_enable=True) diff --git a/vllm/model_executor/models/idefics2_vision_model.py b/vllm/model_executor/models/idefics2_vision_model.py index d6f93a9d49b..b90afbe5abb 100644 --- a/vllm/model_executor/models/idefics2_vision_model.py +++ b/vllm/model_executor/models/idefics2_vision_model.py @@ -166,7 +166,7 @@ class Idefics2VisionAttention(nn.Module): self.num_heads_per_partition, self.head_dim, self.scale, - prefix=prefix, + prefix=f"{prefix}.attn", ) def forward( diff --git a/vllm/model_executor/models/intern_vit.py b/vllm/model_executor/models/intern_vit.py index 8cacfe06eae..8161473641c 100644 --- a/vllm/model_executor/models/intern_vit.py +++ b/vllm/model_executor/models/intern_vit.py @@ -215,7 +215,7 @@ class InternParallelAttention(nn.Module): self.num_heads_per_partition, self.head_dim, self.scale, - prefix=prefix, + prefix=f"{prefix}.attn", ) def _apply_qk_norm(self, q: torch.Tensor, k: torch.Tensor): diff --git a/vllm/model_executor/models/interns1_vit.py b/vllm/model_executor/models/interns1_vit.py index 421e0ffd4dd..533f0681c1d 100644 --- a/vllm/model_executor/models/interns1_vit.py +++ b/vllm/model_executor/models/interns1_vit.py @@ -220,7 +220,7 @@ class InternSdpaAttention(nn.Module): self.num_heads, self.head_dim, self.scale, - prefix=prefix, + prefix=f"{prefix}.attn", ) def forward(self, x: torch.Tensor) -> torch.Tensor: diff --git a/vllm/model_executor/models/mllama4.py b/vllm/model_executor/models/mllama4.py index 52fdeddf4b0..58f63597a7d 100644 --- a/vllm/model_executor/models/mllama4.py +++ b/vllm/model_executor/models/mllama4.py @@ -257,7 +257,7 @@ class Llama4VisionAttention(nn.Module): self.num_local_heads, self.head_dim, self.scaling, - prefix=prefix, + prefix=f"{prefix}.attn", ) if use_data_parallel: diff --git a/vllm/model_executor/models/molmo.py b/vllm/model_executor/models/molmo.py index b1330d92d75..1ee1776569d 100644 --- a/vllm/model_executor/models/molmo.py +++ b/vllm/model_executor/models/molmo.py @@ -235,7 +235,7 @@ class MultiHeadDotProductAttention(nn.Module): self.head_dim, self.scale, num_kv_heads=self.num_kv_heads, - prefix=prefix, + prefix=f"{prefix}.attn", ) def forward( diff --git a/vllm/model_executor/models/molmo2.py b/vllm/model_executor/models/molmo2.py index f9664f32e4e..9d996a93b05 100644 --- a/vllm/model_executor/models/molmo2.py +++ b/vllm/model_executor/models/molmo2.py @@ -611,7 +611,7 @@ class ImagePoolingAttention(nn.Module): self.head_dim, self.scale, num_kv_heads=self.num_kv_heads, - prefix=prefix, + prefix=f"{prefix}.attn", ) def forward_sdpa( diff --git a/vllm/model_executor/models/openpangu_vl.py b/vllm/model_executor/models/openpangu_vl.py index 239ef81d3ee..d7df2cbb4cf 100644 --- a/vllm/model_executor/models/openpangu_vl.py +++ b/vllm/model_executor/models/openpangu_vl.py @@ -125,6 +125,7 @@ class OpenPanguVisionAttention(nn.Module): num_heads=self.num_attention_heads_per_partition, head_size=self.hidden_size_per_attention_head, scale=self.hidden_size_per_attention_head**-0.5, + prefix=f"{prefix}.attn", ) self.apply_rotary_emb = ApplyRotaryEmb(enforce_enable=True) diff --git a/vllm/model_executor/models/qwen2_5_vl.py b/vllm/model_executor/models/qwen2_5_vl.py index c06beb97fac..c2c52fa6626 100644 --- a/vllm/model_executor/models/qwen2_5_vl.py +++ b/vllm/model_executor/models/qwen2_5_vl.py @@ -345,7 +345,7 @@ class Qwen2_5_VisionAttention(nn.Module): num_heads=self.num_attention_heads_per_partition, head_size=self.hidden_size_per_attention_head, scale=self.hidden_size_per_attention_head**-0.5, - prefix=prefix, + prefix=f"{prefix}.attn", ) self.apply_rotary_emb = ApplyRotaryEmb(enforce_enable=True) diff --git a/vllm/model_executor/models/qwen2_vl.py b/vllm/model_executor/models/qwen2_vl.py index 6169e72df4e..d911fb1dd94 100644 --- a/vllm/model_executor/models/qwen2_vl.py +++ b/vllm/model_executor/models/qwen2_vl.py @@ -319,7 +319,7 @@ class Qwen2VisionAttention(nn.Module): num_heads=self.num_attention_heads_per_partition, head_size=self.hidden_size_per_attention_head, scale=self.hidden_size_per_attention_head**-0.5, - prefix=prefix, + prefix=f"{prefix}.attn", ) self.apply_rotary_emb = ApplyRotaryEmb(enforce_enable=True) diff --git a/vllm/model_executor/models/qwen3_omni_moe_thinker.py b/vllm/model_executor/models/qwen3_omni_moe_thinker.py index 96294158c75..9500ce2e2bf 100755 --- a/vllm/model_executor/models/qwen3_omni_moe_thinker.py +++ b/vllm/model_executor/models/qwen3_omni_moe_thinker.py @@ -194,7 +194,7 @@ class Qwen3OmniMoeAudioAttention(nn.Module): num_heads=self.num_local_heads, head_size=self.head_dim, scale=self.scaling, - prefix=prefix, + prefix=f"{prefix}.attn", ) def forward( diff --git a/vllm/model_executor/models/step3_vl.py b/vllm/model_executor/models/step3_vl.py index fe2bb1ac6a6..f3993348b30 100644 --- a/vllm/model_executor/models/step3_vl.py +++ b/vllm/model_executor/models/step3_vl.py @@ -763,7 +763,7 @@ class Step3VisionAttention(nn.Module): self.num_heads, self.head_dim, self.scale, - prefix=prefix, + prefix=f"{prefix}.attn", ) def forward( diff --git a/vllm/model_executor/models/step_vl.py b/vllm/model_executor/models/step_vl.py index 31b266a7e7e..4669771f4bc 100644 --- a/vllm/model_executor/models/step_vl.py +++ b/vllm/model_executor/models/step_vl.py @@ -224,7 +224,7 @@ class PerceptionEncoderVisionAttention(nn.Module): self.num_heads, self.head_dim, self.scale, - prefix=prefix, + prefix=f"{prefix}.attn", ) self.rope = PerceptionEncoderRope2D( dim=self.head_dim, From 2647163674720242cee78dae3f7c7b98539ed029 Mon Sep 17 00:00:00 2001 From: Huy Do Date: Tue, 3 Feb 2026 20:37:51 -0800 Subject: [PATCH 042/810] Save startup benchmark results as a list of values (#33629) Signed-off-by: Huy Do --- vllm/benchmarks/lib/utils.py | 6 ++++++ vllm/benchmarks/startup.py | 8 ++++---- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/vllm/benchmarks/lib/utils.py b/vllm/benchmarks/lib/utils.py index 6d8bfd021fe..d3b6be8690c 100644 --- a/vllm/benchmarks/lib/utils.py +++ b/vllm/benchmarks/lib/utils.py @@ -47,6 +47,12 @@ def convert_to_pytorch_benchmark_format( return records for name, benchmark_values in metrics.items(): + if not isinstance(benchmark_values, list): + raise TypeError( + f"benchmark_values for metric '{name}' must be a list, " + f"but got {type(benchmark_values).__name__}" + ) + record = { "benchmark": { "name": "vLLM benchmark", diff --git a/vllm/benchmarks/startup.py b/vllm/benchmarks/startup.py index dabda547733..005625f61b1 100644 --- a/vllm/benchmarks/startup.py +++ b/vllm/benchmarks/startup.py @@ -101,7 +101,7 @@ def save_to_pytorch_benchmark_format( cold_startup_records = convert_to_pytorch_benchmark_format( args=args, metrics={ - "avg_cold_startup_time": results["avg_cold_startup_time"], + "avg_cold_startup_time": [results["avg_cold_startup_time"]], }, extra_info={ "cold_startup_times": results["cold_startup_times"], @@ -114,7 +114,7 @@ def save_to_pytorch_benchmark_format( cold_compilation_records = convert_to_pytorch_benchmark_format( args=args, metrics={ - "avg_cold_compilation_time": results["avg_cold_compilation_time"], + "avg_cold_compilation_time": [results["avg_cold_compilation_time"]], }, extra_info={ "cold_compilation_times": results["cold_compilation_times"], @@ -129,7 +129,7 @@ def save_to_pytorch_benchmark_format( warm_startup_records = convert_to_pytorch_benchmark_format( args=args, metrics={ - "avg_warm_startup_time": results["avg_warm_startup_time"], + "avg_warm_startup_time": [results["avg_warm_startup_time"]], }, extra_info={ "warm_startup_times": results["warm_startup_times"], @@ -142,7 +142,7 @@ def save_to_pytorch_benchmark_format( warm_compilation_records = convert_to_pytorch_benchmark_format( args=args, metrics={ - "avg_warm_compilation_time": results["avg_warm_compilation_time"], + "avg_warm_compilation_time": [results["avg_warm_compilation_time"]], }, extra_info={ "warm_compilation_times": results["warm_compilation_times"], From eb5ed207437d31daf09cd4ae51e993fc1efa6928 Mon Sep 17 00:00:00 2001 From: Michael Goin Date: Wed, 4 Feb 2026 00:24:14 -0500 Subject: [PATCH 043/810] [Bugfix] Define router_logits_dtype for remaining MoE models (#33737) Signed-off-by: mgoin --- vllm/model_executor/models/afmoe.py | 1 + vllm/model_executor/models/bailing_moe.py | 1 + vllm/model_executor/models/flex_olmo.py | 2 +- vllm/model_executor/models/longcat_flash.py | 7 ++++--- vllm/model_executor/models/mimo_v2_flash.py | 1 + vllm/model_executor/models/step3p5.py | 1 + 6 files changed, 9 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/models/afmoe.py b/vllm/model_executor/models/afmoe.py index 6f5d7d766cf..9b3d9fb2290 100644 --- a/vllm/model_executor/models/afmoe.py +++ b/vllm/model_executor/models/afmoe.py @@ -142,6 +142,7 @@ class AfmoeMoE(nn.Module): e_score_correction_bias=self.expert_bias, enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, + router_logits_dtype=torch.float32, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: diff --git a/vllm/model_executor/models/bailing_moe.py b/vllm/model_executor/models/bailing_moe.py index fc10f790e4d..7725dfa2a88 100644 --- a/vllm/model_executor/models/bailing_moe.py +++ b/vllm/model_executor/models/bailing_moe.py @@ -300,6 +300,7 @@ class BailingMoE(nn.Module): num_expert_group=self.n_group, topk_group=self.topk_group, use_grouped_topk=self.use_grouped_topk, + router_logits_dtype=self.router_dtype, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: diff --git a/vllm/model_executor/models/flex_olmo.py b/vllm/model_executor/models/flex_olmo.py index 11d0949a798..a2e2adc2a6b 100644 --- a/vllm/model_executor/models/flex_olmo.py +++ b/vllm/model_executor/models/flex_olmo.py @@ -71,7 +71,6 @@ class FlexOlmoMoE(nn.Module): prefix=f"{prefix}.gate", ) - # Gate always runs at half / full precision for now. self.experts = FusedMoE( num_experts=hf_config.num_experts, top_k=hf_config.num_experts_per_tok, @@ -82,6 +81,7 @@ class FlexOlmoMoE(nn.Module): quant_config=None, tp_size=tp_size, prefix=f"{prefix}.experts", + router_logits_dtype=torch.float32, ) self.top_k = hf_config.num_experts_per_tok diff --git a/vllm/model_executor/models/longcat_flash.py b/vllm/model_executor/models/longcat_flash.py index f8b426df0aa..32408e7c3e3 100644 --- a/vllm/model_executor/models/longcat_flash.py +++ b/vllm/model_executor/models/longcat_flash.py @@ -236,9 +236,9 @@ class FlashMLP(nn.Module): class LongcatRouter(nn.Module): def __init__( self, - config, - zero_expert_num=0, - rounter_params_dtype=torch.bfloat16, + config: FlashConfig, + zero_expert_num: int, + rounter_params_dtype: torch.dtype, prefix: str = "", ): super().__init__() @@ -309,6 +309,7 @@ class LongcatMoe(nn.Module): prefix=f"{prefix}.experts", enable_eplb=enable_eplb, routed_scaling_factor=config.routed_scaling_factor, + router_logits_dtype=self.rounter_params_dtype, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: diff --git a/vllm/model_executor/models/mimo_v2_flash.py b/vllm/model_executor/models/mimo_v2_flash.py index f7640746aab..f74ce59ab68 100644 --- a/vllm/model_executor/models/mimo_v2_flash.py +++ b/vllm/model_executor/models/mimo_v2_flash.py @@ -174,6 +174,7 @@ class MiMoV2MoE(nn.Module): num_expert_group=config.n_group, topk_group=config.topk_group, scoring_func="sigmoid", + router_logits_dtype=self.gate_dtype, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: diff --git a/vllm/model_executor/models/step3p5.py b/vllm/model_executor/models/step3p5.py index f0d7b4a75a9..8019dbdbee1 100644 --- a/vllm/model_executor/models/step3p5.py +++ b/vllm/model_executor/models/step3p5.py @@ -388,6 +388,7 @@ class FusedMoEBlock(nn.Module): routed_scaling_factor=config.moe_router_scaling_factor, enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, + router_logits_dtype=torch.float32, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: From 5e1e0a0fbdf52dd21bfa6bc4dd6e88487a917d23 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Wed, 4 Feb 2026 00:25:11 -0500 Subject: [PATCH 044/810] [Refactor] Remove unused dead code (#33718) Signed-off-by: yewentao256 --- vllm/collect_env.py | 8 ------ vllm/compilation/fx_utils.py | 15 ---------- vllm/config/utils.py | 54 ------------------------------------ 3 files changed, 77 deletions(-) diff --git a/vllm/collect_env.py b/vllm/collect_env.py index c042fc1d707..0cf5681bcf5 100644 --- a/vllm/collect_env.py +++ b/vllm/collect_env.py @@ -146,14 +146,6 @@ def run_and_parse_first_match(run_lambda, command, regex): return match.group(1) -def run_and_return_first_line(run_lambda, command): - """Run command using run_lambda and returns first line if output is not empty.""" - rc, out, _ = run_lambda(command) - if rc != 0: - return None - return out.split("\n")[0] - - def get_conda_packages(run_lambda, patterns=None): if patterns is None: patterns = DEFAULT_CONDA_PATTERNS diff --git a/vllm/compilation/fx_utils.py b/vllm/compilation/fx_utils.py index 5c2e7ac93e6..a87ffea7837 100644 --- a/vllm/compilation/fx_utils.py +++ b/vllm/compilation/fx_utils.py @@ -18,21 +18,6 @@ def is_auto_func(node: fx.Node, op: OpOverload) -> bool: return is_func(node, auto_functionalized) and node.args[0] == op -# Returns the first specified node with the given op (if it exists) -def find_specified_fn_maybe(nodes: Iterable[fx.Node], op: OpOverload) -> fx.Node | None: - for node in nodes: - if node.target == op: - return node - return None - - -# Returns the first specified node with the given op -def find_specified_fn(nodes: Iterable[fx.Node], op: OpOverload) -> fx.Node: - node = find_specified_fn_maybe(nodes, op) - assert node is not None, f"Could not find {op} in nodes {nodes}" - return node - - # Returns the first auto_functionalized node with the given op (if it exists) def find_auto_fn_maybe(nodes: Iterable[fx.Node], op: OpOverload) -> fx.Node | None: for node in nodes: diff --git a/vllm/config/utils.py b/vllm/config/utils.py index e8c866f0223..bd2a741e3d8 100644 --- a/vllm/config/utils.py +++ b/vllm/config/utils.py @@ -14,7 +14,6 @@ from dataclasses import MISSING, Field, field, fields, is_dataclass from itertools import pairwise from typing import TYPE_CHECKING, Any, Protocol, TypeVar, cast -import regex as re import torch from pydantic import ConfigDict from pydantic.dataclasses import dataclass @@ -144,34 +143,6 @@ def getattr_iter( return default_factory() if default_factory is not None else default -def contains_object_print(text: str) -> bool: - """ - Check if the text looks like a printed Python object, e.g. - contains any substring matching the pattern: "at 0xFFFFFFF>" - We match against 0x followed by 2-16 hex chars (there's - a max of 16 on a 64-bit system). - - Args: - text (str): The text to check - - Returns: - result (bool): `True` if a match is found, `False` otherwise. - """ - pattern = r"at 0x[a-fA-F0-9]{2,16}>" - match = re.search(pattern, text) - return match is not None - - -def assert_hashable(text: str) -> bool: - if not contains_object_print(text): - return True - raise AssertionError( - f"vLLM tried to hash some configs that may have Python objects ids " - f"in them. This is a bug, please file an issue. " - f"Text being hashed: {text}" - ) - - def get_attr_docs(cls: type[Any]) -> dict[str, str]: """ Get any docstrings placed after attribute assignments in a class body. @@ -354,31 +325,6 @@ def hash_factors(items: dict[str, object]) -> str: return hashlib.sha256(json.dumps(items, sort_keys=True).encode()).hexdigest() -def handle_deprecated( - config: ConfigT, - old_name: str, - new_name_or_names: str | list[str], - removal_version: str, -) -> None: - old_val = getattr(config, old_name) - if old_val is None: - return - - if isinstance(new_name_or_names, str): - new_names = [new_name_or_names] - else: - new_names = new_name_or_names - - msg = ( - f"{old_name} is deprecated and will be removed in {removal_version}. " - f"Use {', '.join(new_names)} instead." - ) - logger.warning(msg) - - for new_name in new_names: - setattr(config, new_name, old_val) - - @dataclass class Range: """ From 45f8fd6f979de059e2d8a03ea2ddc40c06394931 Mon Sep 17 00:00:00 2001 From: Frank Wang <41319051+frankwang28@users.noreply.github.com> Date: Tue, 3 Feb 2026 21:27:34 -0800 Subject: [PATCH 045/810] [Feature] Enable `TRITON_ATTN` for Batch Invariance (#33688) Signed-off-by: frankwang28 --- docs/features/batch_invariance.md | 1 + tests/v1/determinism/utils.py | 1 + vllm/model_executor/layers/batch_invariant.py | 9 ++++++--- vllm/v1/attention/ops/triton_unified_attention.py | 6 +++++- 4 files changed, 13 insertions(+), 4 deletions(-) diff --git a/docs/features/batch_invariance.md b/docs/features/batch_invariance.md index 0144e2f71f0..72224c96cfd 100644 --- a/docs/features/batch_invariance.md +++ b/docs/features/batch_invariance.md @@ -108,6 +108,7 @@ Batch invariance has been tested and verified on the following models: - **Qwen3 (MoE)**: `Qwen/Qwen3-30B-A3B`, `Qwen/Qwen3-Next-80B-A3B-Instruct` - **Qwen2.5**: `Qwen/Qwen2.5-0.5B-Instruct`, `Qwen/Qwen2.5-1.5B-Instruct`, `Qwen/Qwen2.5-3B-Instruct`, `Qwen/Qwen2.5-7B-Instruct`, `Qwen/Qwen2.5-14B-Instruct`, `Qwen/Qwen2.5-32B-Instruct` - **Llama 3**: `meta-llama/Llama-3.1-8B-Instruct`, `meta-llama/Llama-3.2-1B-Instruct` +- **GPT-OSS**: `openai/gpt-oss-20b`, `openai/gpt-oss-120b` Other models may also work, but these have been explicitly validated. If you encounter issues with a specific model, please report them on the [GitHub issue tracker](https://github.com/vllm-project/vllm/issues/new/choose). diff --git a/tests/v1/determinism/utils.py b/tests/v1/determinism/utils.py index 5066315762e..ca3ccab5eff 100644 --- a/tests/v1/determinism/utils.py +++ b/tests/v1/determinism/utils.py @@ -18,6 +18,7 @@ skip_unsupported = pytest.mark.skipif( BACKENDS: list[str] = [ "FLASH_ATTN", + "TRITON_ATTN", "TRITON_MLA", ] diff --git a/vllm/model_executor/layers/batch_invariant.py b/vllm/model_executor/layers/batch_invariant.py index 3f44608ab9e..fcfadd60f5c 100644 --- a/vllm/model_executor/layers/batch_invariant.py +++ b/vllm/model_executor/layers/batch_invariant.py @@ -1003,8 +1003,11 @@ def vllm_is_batch_invariant() -> bool: def override_envs_for_invariance( attention_backend: AttentionBackendEnum | None, ): - supported_backends = [ + decode_invariant_backends = [ AttentionBackendEnum.FLASH_ATTN, # best supported backend + AttentionBackendEnum.TRITON_ATTN, + ] + supported_backends = decode_invariant_backends + [ # FlashInfer temporarily disabled due to invariant CTA sizes. # See FlashInfer issue #2424 # AttentionBackendEnum.FLASHINFER, @@ -1025,9 +1028,9 @@ def override_envs_for_invariance( "one of the supported backends before enabling batch_invariant." ) raise RuntimeError(error) - if attention_backend != supported_backends[0]: + if attention_backend not in decode_invariant_backends: warning = ( - "You are using a decode-invariant form of batch invariance. " + "You are using a non-decode-invariant form of batch invariance. " "This will not be invariant between prefill and decode." ) logger.warning_once(warning, scope="local") diff --git a/vllm/v1/attention/ops/triton_unified_attention.py b/vllm/v1/attention/ops/triton_unified_attention.py index 6855233ee94..4ddd47c6dd6 100644 --- a/vllm/v1/attention/ops/triton_unified_attention.py +++ b/vllm/v1/attention/ops/triton_unified_attention.py @@ -10,10 +10,12 @@ import torch from vllm.logger import init_logger +from vllm.model_executor.layers.batch_invariant import vllm_is_batch_invariant from vllm.platforms import current_platform from vllm.triton_utils import tl, triton logger = init_logger(__name__) +is_batch_invariant = vllm_is_batch_invariant() float8_info = torch.finfo(current_platform.fp8_dtype()) @@ -972,7 +974,8 @@ def unified_attention( # Launch the 2D kernel if # 1. No intermediate tiled softmax buffers for the 3D kernel have been allocated, or # 2. The batch includes at least one prefill request, or - # 3. The number of sequences exceeds the configured threshold + # 3. The number of sequences exceeds the configured threshold, or + # 4. Batch invariance is enabled if ( seq_threshold_3D is None or num_par_softmax_segments is None @@ -981,6 +984,7 @@ def unified_attention( or softmax_segm_expsum is None or max_seqlen_q > 1 or num_seqs > seq_threshold_3D + or is_batch_invariant ): kernel_unified_attention_2d[ ( From 90d74ebaa47fcecdcd8ef72338dda47b7cb6fbf0 Mon Sep 17 00:00:00 2001 From: Cyrus Leung Date: Wed, 4 Feb 2026 13:51:52 +0800 Subject: [PATCH 046/810] [Deprecation] Remove `_get_data_parser` in MM processor (#33757) Signed-off-by: DarkLight1337 --- vllm/multimodal/processing/processor.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/vllm/multimodal/processing/processor.py b/vllm/multimodal/processing/processor.py index dfce4dab2b6..fe697c5ceab 100644 --- a/vllm/multimodal/processing/processor.py +++ b/vllm/multimodal/processing/processor.py @@ -988,16 +988,15 @@ class BaseMultiModalProcessor(ABC, Generic[_I]): self.dummy_inputs = dummy_inputs self.cache = cache + # TODO: Remove in v0.18 if hasattr(self, "_get_data_parser"): - logger.warning_once( - "BaseMultiModalProcessor._get_data_parser is deprecated " - "and will be removed in v0.16." - "You should override `info.build_data_parser` instead." + raise ValueError( + "BaseMultiModalProcessor._get_data_parser has been " + "moved to `BaseProcessingInfo.build_data_parser` in v0.16. " + "You should override `BaseProcessingInfo.build_data_parser` instead." ) - self.data_parser = self._get_data_parser() # type: ignore - else: - self.data_parser = self.info.get_data_parser() + self.data_parser = self.info.get_data_parser() @property @deprecated("Will be removed in v0.17. Use `info.supported_mm_limits` instead.") From d88a1df699f68e5284fe3a3170f8ae292a3e9c3f Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Wed, 4 Feb 2026 00:58:21 -0500 Subject: [PATCH 047/810] [Deprecation] Deprecate profiling envs (#33722) Signed-off-by: yewentao256 --- docs/usage/security.md | 4 ++-- vllm/envs.py | 47 ------------------------------------------ 2 files changed, 2 insertions(+), 49 deletions(-) diff --git a/docs/usage/security.md b/docs/usage/security.md index 0a54221ec56..bb920ff43b1 100644 --- a/docs/usage/security.md +++ b/docs/usage/security.md @@ -178,7 +178,7 @@ These endpoints are **only available when the environment variable `VLLM_SERVER_ - `/is_sleeping` - Check if engine is sleeping - `/collective_rpc` - Execute arbitrary RPC methods on the engine (extremely dangerous) -**Profiler endpoints (only when `VLLM_TORCH_PROFILER_DIR` or `VLLM_TORCH_CUDA_PROFILE` are set):** +**Profiler endpoints (only when profiling is enabled via `--profiler-config`):** These endpoints are only available when profiling is enabled and should only be used for local development: @@ -207,7 +207,7 @@ An attacker who can reach the vLLM HTTP server can: - Cache manipulation that can disrupt service - Detailed server configuration disclosure -Similarly, never enable profiler endpoints (`VLLM_TORCH_PROFILER_DIR` or `VLLM_TORCH_CUDA_PROFILE`) in production. +Similarly, never enable profiler endpoints in production. **Be cautious with `--enable-tokenizer-info-endpoint`:** Only enable the `/tokenizer_info` endpoint if you need to expose tokenizer configuration information. This endpoint reveals chat templates and tokenizer settings that may contain sensitive implementation details or prompt engineering strategies. diff --git a/vllm/envs.py b/vllm/envs.py index f9aaa4f380c..caddf0b7642 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -855,53 +855,6 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_LORA_RESOLVER_HF_REPO_LIST": lambda: os.getenv( "VLLM_LORA_RESOLVER_HF_REPO_LIST", None ), - # Enables torch CUDA profiling if set to 1. - # Deprecated, see profiler_config. - "VLLM_TORCH_CUDA_PROFILE": lambda: os.getenv("VLLM_TORCH_CUDA_PROFILE"), - # Enables torch profiler if set. - # Deprecated, see profiler_config. - "VLLM_TORCH_PROFILER_DIR": lambda: os.getenv("VLLM_TORCH_PROFILER_DIR"), - # Enable torch profiler to record shapes if set to 1. - # Deprecated, see profiler_config. - "VLLM_TORCH_PROFILER_RECORD_SHAPES": lambda: ( - os.getenv("VLLM_TORCH_PROFILER_RECORD_SHAPES") - ), - # Enable torch profiler to profile memory if set to 1. - # Deprecated, see profiler_config. - "VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY": lambda: ( - os.getenv("VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY") - ), - # Enable torch profiler to profile stack if set to 1. - # Deprecated, see profiler_config. - "VLLM_TORCH_PROFILER_WITH_STACK": lambda: ( - os.getenv("VLLM_TORCH_PROFILER_WITH_STACK") - ), - # Enable torch profiler to profile flops if set to 1. - # Deprecated, see profiler_config. - "VLLM_TORCH_PROFILER_WITH_FLOPS": lambda: ( - os.getenv("VLLM_TORCH_PROFILER_WITH_FLOPS") - ), - # Disable torch profiling of the AsyncLLMEngine process if set to 1. - # Deprecated, see profiler_config. - "VLLM_TORCH_PROFILER_DISABLE_ASYNC_LLM": lambda: ( - os.getenv("VLLM_TORCH_PROFILER_DISABLE_ASYNC_LLM") - ), - # Delay number of iterations before starting profiling when using - # the torch/torch CUDA profiler. If set to 0, will start profiling immediately. - # Deprecated, see profiler_config. - "VLLM_PROFILER_DELAY_ITERS": lambda: (os.getenv("VLLM_PROFILER_DELAY_ITERS")), - # Maximum number of iterations to profile when using the torch/torch CUDA profiler. - # If set to 0, will not limit the number of iterations. - "VLLM_PROFILER_MAX_ITERS": lambda: os.getenv("VLLM_PROFILER_MAX_ITERS"), - # Control whether torch profiler gzip-compresses profiling files. - # Deprecated, see profiler_config. - "VLLM_TORCH_PROFILER_USE_GZIP": lambda: os.getenv("VLLM_TORCH_PROFILER_USE_GZIP"), - # Control whether torch profiler dumps the self_cuda_time_total table. - # Set to 0 to disable dumping the table. - # Deprecated, see profiler_config. - "VLLM_TORCH_PROFILER_DUMP_CUDA_TIME_TOTAL": lambda: ( - os.getenv("VLLM_TORCH_PROFILER_DUMP_CUDA_TIME_TOTAL") - ), # If set, vLLM will use Triton implementations of AWQ. "VLLM_USE_TRITON_AWQ": lambda: bool(int(os.getenv("VLLM_USE_TRITON_AWQ", "0"))), # If set, allow loading or unloading lora adapters in runtime, From 08e094997ecb28b4c5ad4dd28fd6fbb48adba279 Mon Sep 17 00:00:00 2001 From: Matt <156021403+mawong-amd@users.noreply.github.com> Date: Wed, 4 Feb 2026 00:51:33 -0600 Subject: [PATCH 048/810] [Hardware][AMD][CI] Refactor AMD tests to properly use BuildKite parallelism (#32745) Signed-off-by: Matthew Wong --- .../scripts/hardware_ci/run-amd-test.sh | 57 ++----------------- .buildkite/test-amd.yaml | 18 +++--- 2 files changed, 14 insertions(+), 61 deletions(-) diff --git a/.buildkite/scripts/hardware_ci/run-amd-test.sh b/.buildkite/scripts/hardware_ci/run-amd-test.sh index 6f4a0decfc0..f3690939667 100755 --- a/.buildkite/scripts/hardware_ci/run-amd-test.sh +++ b/.buildkite/scripts/hardware_ci/run-amd-test.sh @@ -87,7 +87,7 @@ mkdir -p "${HF_CACHE}" HF_MOUNT="/root/.cache/huggingface" commands=$@ -echo "Commands:$commands" +echo "Raw commands: $commands" commands=${commands//"pytest -v -s basic_correctness/test_basic_correctness.py"/"pytest -v -s basic_correctness/test_basic_correctness.py"} @@ -169,6 +169,9 @@ if [[ $commands == *" entrypoints/llm "* ]]; then --ignore=entrypoints/llm/test_prompt_validation.py "} fi +commands=$(echo "$commands" | sed 's/ \\ / /g') +echo "Final commands: $commands" + # --ignore=entrypoints/openai/test_encoder_decoder.py \ # --ignore=entrypoints/openai/test_embedding.py \ # --ignore=entrypoints/openai/test_oot_registration.py @@ -176,7 +179,6 @@ fi # --ignore=entrypoints/openai/test_models.py <= Fails on MI250 but passes on MI300 as of 2025-03-13 -PARALLEL_JOB_COUNT=8 MYPYTHONPATH=".." # Test that we're launching on the machine that has @@ -187,56 +189,7 @@ if [[ -z "$render_gid" ]]; then exit 1 fi -# check if the command contains shard flag, we will run all shards in parallel because the host have 8 GPUs. -if [[ $commands == *"--shard-id="* ]]; then - # assign job count as the number of shards used - commands=$(echo "$commands" | sed -E "s/--num-shards[[:blank:]]*=[[:blank:]]*[0-9]*/--num-shards=${PARALLEL_JOB_COUNT} /g" | sed 's/ \\ / /g') - for GPU in $(seq 0 $(($PARALLEL_JOB_COUNT-1))); do - # assign shard-id for each shard - commands_gpu=$(echo "$commands" | sed -E "s/--shard-id[[:blank:]]*=[[:blank:]]*[0-9]*/--shard-id=${GPU} /g" | sed 's/ \\ / /g') - echo "Shard ${GPU} commands:$commands_gpu" - echo "Render devices: $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES" - docker run \ - --device /dev/kfd $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES \ - --network=host \ - --shm-size=16gb \ - --group-add "$render_gid" \ - --rm \ - -e HIP_VISIBLE_DEVICES="${GPU}" \ - -e HF_TOKEN \ - -e AWS_ACCESS_KEY_ID \ - -e AWS_SECRET_ACCESS_KEY \ - -v "${HF_CACHE}:${HF_MOUNT}" \ - -e "HF_HOME=${HF_MOUNT}" \ - -e "PYTHONPATH=${MYPYTHONPATH}" \ - --name "${container_name}_${GPU}" \ - "${image_name}" \ - /bin/bash -c "${commands_gpu}" \ - |& while read -r line; do echo ">>Shard $GPU: $line"; done & - PIDS+=($!) - done - #wait for all processes to finish and collect exit codes - for pid in "${PIDS[@]}"; do - wait "${pid}" - STATUS+=($?) - done - at_least_one_shard_with_tests=0 - for st in "${STATUS[@]}"; do - if [[ ${st} -ne 0 ]] && [[ ${st} -ne 5 ]]; then - echo "One of the processes failed with $st" - exit "${st}" - elif [[ ${st} -eq 5 ]]; then - echo "Shard exited with status 5 (no tests collected) - treating as success" - else # This means st is 0 - at_least_one_shard_with_tests=1 - fi - done - if [[ ${#STATUS[@]} -gt 0 && ${at_least_one_shard_with_tests} -eq 0 ]]; then - echo "All shards reported no tests collected. Failing the build." - exit 1 - fi - -elif [[ $commands == *"VLLM_TEST_GROUP_NAME=mi325_4-2-node-tests-4-gpus-in-total"* ]]; then +if [[ $commands == *"VLLM_TEST_GROUP_NAME=mi325_4-2-node-tests-4-gpus-in-total"* ]]; then export DCKR_VER=$(docker --version | sed 's/Docker version \(.*\), build .*/\1/') diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index a14dcd030ec..ee7c6ab0a5d 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -542,7 +542,7 @@ steps: - label: LoRA Test %N # 20min each timeout_in_minutes: 30 mirror_hardwares: [amdexperimental] - agent_pool: mi325_8 + agent_pool: mi325_1 # grade: Blocking source_file_dependencies: - vllm/lora @@ -636,7 +636,7 @@ steps: - label: Kernels Attention Test %N # 23min timeout_in_minutes: 35 mirror_hardwares: [amdexperimental, amdproduction] - agent_pool: mi325_8 + agent_pool: mi325_1 # grade: Blocking source_file_dependencies: - csrc/attention/ @@ -651,7 +651,7 @@ steps: - label: Kernels Quantization Test %N # 64min timeout_in_minutes: 90 mirror_hardwares: [amdexperimental] - agent_pool: mi325_8 + agent_pool: mi325_1 # grade: Blocking source_file_dependencies: - csrc/quantization/ @@ -664,7 +664,7 @@ steps: - label: Kernels MoE Test %N # 40min timeout_in_minutes: 60 mirror_hardwares: [amdexperimental, amdproduction] - agent_pool: mi325_8 + agent_pool: mi325_1 # grade: Blocking source_file_dependencies: - csrc/quantization/cutlass_w8a8/moe/ @@ -742,7 +742,7 @@ steps: - label: Benchmarks # 11min timeout_in_minutes: 20 mirror_hardwares: [amdexperimental, amdproduction] - agent_pool: mi325_8 + agent_pool: mi325_1 # grade: Blocking working_dir: "/vllm-workspace/.buildkite" source_file_dependencies: @@ -753,7 +753,7 @@ steps: - label: Benchmarks CLI Test # 7min timeout_in_minutes: 20 mirror_hardwares: [amdexperimental, amdproduction] - agent_pool: mi325_8 + agent_pool: mi325_1 # grade: Blocking source_file_dependencies: - vllm/ @@ -827,7 +827,7 @@ steps: - label: Basic Models Tests (Extra Initialization) %N timeout_in_minutes: 45 mirror_hardwares: [amdexperimental, amdproduction] - agent_pool: mi325_8 + agent_pool: mi325_1 # grade: Blocking torch_nightly: true source_file_dependencies: @@ -888,7 +888,7 @@ steps: - label: Language Models Tests (Extra Standard) %N timeout_in_minutes: 45 mirror_hardwares: [amdexperimental] - agent_pool: mi325_8 + agent_pool: mi325_1 # grade: Blocking torch_nightly: true source_file_dependencies: @@ -909,7 +909,7 @@ steps: - label: Language Models Tests (Hybrid) %N timeout_in_minutes: 75 mirror_hardwares: [amdexperimental] - agent_pool: mi325_8 + agent_pool: mi325_1 # grade: Blocking torch_nightly: true source_file_dependencies: From 4403e3ed4c880365978d1716e5f3a8dbe6a6af31 Mon Sep 17 00:00:00 2001 From: zhanqiuhu <49648934+ZhanqiuHu@users.noreply.github.com> Date: Wed, 4 Feb 2026 02:46:48 -0500 Subject: [PATCH 049/810] [Metrics] Add labeled prompt token metrics for P/D disaggregation (#33290) Add labeled Prometheus metrics to distinguish where prompt tokens come from in P/D disaggregated deployments. In P/D disaggregation, decode instances receive KV cache from prefill instances. Currently, decode reports inflated prompt throughput because it counts all prompt tokens as "computed", even though most were transferred. This PR adds labeled metrics so users can understand actual compute work vs transferred work: vllm:prompt_tokens_by_source_total{source="local_compute"} # Tokens prefilled locally vllm:prompt_tokens_by_source_total{source="external_kv_transfer"} # Tokens received via KV transfer vllm:prompt_tokens_by_source_total{source="local_cache_hit"} # Tokens from local prefix cache vllm:prompt_tokens_cached_total # Total cached (local + external, -1 when all Signed-off-by: Zhanqiu Hu --- tests/v1/metrics/test_stats.py | 104 +++++++++++++++++++++++++++++++- vllm/v1/core/sched/scheduler.py | 1 + vllm/v1/engine/__init__.py | 4 +- vllm/v1/metrics/loggers.py | 47 ++++++++++++++- vllm/v1/metrics/stats.py | 76 ++++++++++++++++++++++- 5 files changed, 227 insertions(+), 5 deletions(-) diff --git a/tests/v1/metrics/test_stats.py b/tests/v1/metrics/test_stats.py index 7d902bbc6fc..d49874adc99 100644 --- a/tests/v1/metrics/test_stats.py +++ b/tests/v1/metrics/test_stats.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from vllm.v1.engine import FinishReason -from vllm.v1.metrics.stats import IterationStats, RequestStateStats +from vllm.v1.metrics.stats import IterationStats, PromptTokenStats, RequestStateStats def test_iteration_stats_repr(): @@ -107,3 +107,105 @@ def test_prefill_kv_computed_edge_cases(): finished_req2.num_cached_tokens, 0 ) assert prefill_kv_computed2 == 0 # All cached, nothing computed + + +def test_prompt_token_stats_all_computed(): + """Test all tokens computed locally, no caching.""" + stats = PromptTokenStats() + + # Case 1: No caching (All tokens computed locally) + stats.update_from_output( + num_cached_tokens=0, + num_external_computed_tokens=0, + prompt_len=1000, + ) + + assert stats.computed == 1000 + assert stats.local_cache_hit == 0 + assert stats.external_kv_transfer == 0 + assert stats.total == 1000 + + +def test_prompt_token_stats_partial_local_cache(): + """Test partial local prefix cache hit.""" + stats = PromptTokenStats() + + # Case 2: Partial local cache + stats.update_from_output( + num_cached_tokens=300, + num_external_computed_tokens=0, + prompt_len=1000, + ) + + assert stats.computed == 700 + assert stats.local_cache_hit == 300 + assert stats.external_kv_transfer == 0 + + +def test_prompt_token_stats_partial_external_transfer(): + """Test partial external KV transfer.""" + stats = PromptTokenStats() + + # Case 3: Partial external transfer + stats.update_from_output( + num_cached_tokens=500, + num_external_computed_tokens=500, + prompt_len=1000, + ) + + assert stats.computed == 500 + assert stats.local_cache_hit == 0 + assert stats.external_kv_transfer == 500 + + +def test_prompt_token_stats_mixed_sources(): + """Test mix of local cache and external transfer.""" + stats = PromptTokenStats() + + # Case 4: Mixed sources + stats.update_from_output( + num_cached_tokens=600, + num_external_computed_tokens=200, + prompt_len=1000, + ) + + assert stats.computed == 400 + assert stats.local_cache_hit == 400 + assert stats.external_kv_transfer == 200 + + +def test_prompt_token_stats_full_local_cache_recompute(): + """Test full local cache triggers last token recomputation. + + When all tokens are cached, the scheduler reduces num_cached_tokens by 1 + to force the model to recompute the last token. + """ + stats = PromptTokenStats() + + # Case 5: Full local cache (999 cached after reduction, 1 recomputed) + stats.update_from_output( + num_cached_tokens=999, + num_external_computed_tokens=0, + prompt_len=1000, + ) + + assert stats.computed == 1 + assert stats.local_cache_hit == 1000 + assert stats.recomputed_tokens == 1 + + +def test_prompt_token_stats_full_external_transfer_recompute(): + """Test full external transfer triggers last token recomputation.""" + stats = PromptTokenStats() + + # Case 6: Full external transfer (999 cached after reduction, 1 recomputed) + stats.update_from_output( + num_cached_tokens=999, + num_external_computed_tokens=1000, + prompt_len=1000, + ) + + assert stats.computed == 1 + assert stats.local_cache_hit == 0 + assert stats.external_kv_transfer == 1000 + assert stats.recomputed_tokens == 1 diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 83c965f233a..9308f9ed1c1 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1378,6 +1378,7 @@ class Scheduler(SchedulerInterface): kv_transfer_params=kv_transfer_params, trace_headers=request.trace_headers, num_cached_tokens=request.num_cached_tokens, + num_external_computed_tokens=request.num_external_computed_tokens, routed_experts=routed_experts, num_nans_in_logits=request.num_nans_in_logits, ) diff --git a/vllm/v1/engine/__init__.py b/vllm/v1/engine/__init__.py index e8e44746bf4..5328a673554 100644 --- a/vllm/v1/engine/__init__.py +++ b/vllm/v1/engine/__init__.py @@ -139,8 +139,10 @@ class EngineCoreOutput( kv_transfer_params: dict[str, Any] | None = None trace_headers: Mapping[str, str] | None = None - # The number of tokens with prefix cache hits. + # The number of tokens with prefix cache hits (local + external). num_cached_tokens: int = 0 + # The number of tokens computed remotely (original count from connector). + num_external_computed_tokens: int = 0 routed_experts: np.ndarray | None = None # The number of NaNs in logits. # A value greater than 0 indicates that the output is corrupted. diff --git a/vllm/v1/metrics/loggers.py b/vllm/v1/metrics/loggers.py index 3a080f01a4d..49b97e8f37a 100644 --- a/vllm/v1/metrics/loggers.py +++ b/vllm/v1/metrics/loggers.py @@ -25,6 +25,7 @@ from vllm.v1.metrics.stats import ( CachingMetrics, IterationStats, MultiModalCacheStats, + PromptTokenStats, SchedulerStats, ) from vllm.v1.spec_decode.metrics import SpecDecodingLogging, SpecDecodingProm @@ -136,7 +137,8 @@ class LoggingStatLogger(StatLoggerBase): def _track_iteration_stats(self, iteration_stats: IterationStats): # Save tracked stats for token counters. - self.num_prompt_tokens += iteration_stats.num_prompt_tokens + # Use computed tokens for prompt throughput (excludes cached/transferred) + self.num_prompt_tokens += iteration_stats.prompt_token_stats.computed self.num_generation_tokens += iteration_stats.num_generation_tokens self.num_corrupted_reqs += iteration_stats.num_corrupted_reqs self.num_preemptions += iteration_stats.num_preempted_reqs @@ -590,6 +592,41 @@ class PrometheusStatLogger(AggregateStatLoggerBase): counter_prompt_tokens, engine_indexes, model_name ) + # Labeled prompt token counters by source + counter_prompt_tokens_by_source = self._counter_cls( + name="vllm:prompt_tokens_by_source", + documentation="Number of prompt tokens by source.", + labelnames=labelnames + ["source"], + ) + self.counter_prompt_tokens_by_source: dict[str, dict[int, Counter]] = {} + for source in PromptTokenStats.ALL_SOURCES: + self.counter_prompt_tokens_by_source[source] = { + idx: counter_prompt_tokens_by_source.labels( + model_name, str(idx), source + ) + for idx in engine_indexes + } + + # Cached prompt tokens counter + counter_prompt_tokens_cached = self._counter_cls( + name="vllm:prompt_tokens_cached", + documentation="Number of cached prompt tokens (local + external).", + labelnames=labelnames, + ) + self.counter_prompt_tokens_cached = make_per_engine( + counter_prompt_tokens_cached, engine_indexes, model_name + ) + + # Recomputed tokens (last token recomputed when entire prompt is cached) + counter_prompt_tokens_recomputed = self._counter_cls( + name="vllm:prompt_tokens_recomputed", + documentation="Number of cached tokens recomputed for forward pass.", + labelnames=labelnames, + ) + self.counter_prompt_tokens_recomputed = make_per_engine( + counter_prompt_tokens_recomputed, engine_indexes, model_name + ) + counter_generation_tokens = self._counter_cls( name="vllm:generation_tokens", documentation="Number of generation tokens processed.", @@ -1070,6 +1107,14 @@ class PrometheusStatLogger(AggregateStatLoggerBase): iteration_stats.num_preempted_reqs ) self.counter_prompt_tokens[engine_idx].inc(iteration_stats.num_prompt_tokens) + # Labeled prompt token counters by source + pts = iteration_stats.prompt_token_stats + for source in PromptTokenStats.ALL_SOURCES: + self.counter_prompt_tokens_by_source[source][engine_idx].inc( + pts.get_by_source(source) + ) + self.counter_prompt_tokens_cached[engine_idx].inc(pts.cached_tokens) + self.counter_prompt_tokens_recomputed[engine_idx].inc(pts.recomputed_tokens) self.counter_generation_tokens[engine_idx].inc( iteration_stats.num_generation_tokens ) diff --git a/vllm/v1/metrics/stats.py b/vllm/v1/metrics/stats.py index 3404a720e96..1b7ee105ebf 100644 --- a/vllm/v1/metrics/stats.py +++ b/vllm/v1/metrics/stats.py @@ -231,13 +231,76 @@ class FinishedRequestStats: num_cached_tokens: int = 0 +@dataclass +class PromptTokenStats: + """Breakdown of prompt tokens by source. + + Fields: + computed: Tokens prefilled locally (actual compute work). + local_cache_hit: Tokens from local prefix cache. + external_kv_transfer: Tokens from external KV transfer. + cached_tokens: Tokens skipped during prefill (from scheduler). + recomputed_tokens: Cached tokens that were recomputed (see below). + total: Total prompt tokens. + + Invariants: + computed + local_cache_hit + external_kv_transfer - recomputed_tokens = total + local_cache_hit + external_kv_transfer - recomputed_tokens = cached_tokens + """ + + ALL_SOURCES: tuple[str, ...] = ( + "local_compute", + "local_cache_hit", + "external_kv_transfer", + ) + + computed: int = 0 + local_cache_hit: int = 0 + external_kv_transfer: int = 0 + cached_tokens: int = 0 + recomputed_tokens: int = 0 + total: int = 0 + + def update_from_output( + self, + num_cached_tokens: int, + num_external_computed_tokens: int, + prompt_len: int, + ) -> None: + """Update stats from a prefill output.""" + # When all tokens are cached, the scheduler reduces num_cached_tokens + # by 1 to force the model to recompute the last token, since the model + # needs at least one input token to run a forward pass. + recomputed = 1 if (num_cached_tokens + 1 == prompt_len) else 0 + + self.computed += prompt_len - num_cached_tokens + self.external_kv_transfer += num_external_computed_tokens + self.local_cache_hit += ( + num_cached_tokens + recomputed - num_external_computed_tokens + ) + self.cached_tokens += num_cached_tokens + self.recomputed_tokens += recomputed + self.total += prompt_len + + def get_by_source(self, source: str) -> int: + """Get token count by source label.""" + source_map = { + "local_compute": self.computed, + "local_cache_hit": self.local_cache_hit, + "external_kv_transfer": self.external_kv_transfer, + } + if source not in source_map: + raise ValueError(f"Unknown source: {source}") + return source_map[source] + + class IterationStats: """Stats associated with a single set of EngineCoreOutputs.""" def __init__(self): self.iteration_timestamp = time.time() self.num_generation_tokens = 0 - self.num_prompt_tokens = 0 + self.prompt_token_stats = PromptTokenStats() self.num_preempted_reqs = 0 self.finished_requests: list[FinishedRequestStats] = [] self.max_num_generation_tokens_iter: list[int] = [] @@ -250,6 +313,11 @@ class IterationStats: field_to_value_str = ", ".join(f"{k}={v}" for k, v in vars(self).items()) return f"{self.__class__.__name__}({field_to_value_str})" + @property + def num_prompt_tokens(self) -> int: + """Total prompt tokens (for backward compatibility).""" + return self.prompt_token_stats.total + def _time_since(self, start: float) -> float: """Calculate an interval relative to this iteration's timestamp.""" return self.iteration_timestamp - start @@ -268,7 +336,11 @@ class IterationStats: self.num_generation_tokens += num_new_generation_tokens if is_prefilling: - self.num_prompt_tokens += prompt_len + self.prompt_token_stats.update_from_output( + num_cached_tokens=output.num_cached_tokens, + num_external_computed_tokens=output.num_external_computed_tokens, + prompt_len=prompt_len, + ) first_token_latency = self._time_since(req_stats.arrival_time) self.time_to_first_tokens_iter.append(first_token_latency) From 061da6bcf7d369cc2fb55a56b8bdc91fd9bc96d5 Mon Sep 17 00:00:00 2001 From: Kunshang Ji Date: Wed, 4 Feb 2026 16:40:17 +0800 Subject: [PATCH 050/810] [XPU] remove common path warning log (#33769) Signed-off-by: Kunshang Ji --- vllm/platforms/__init__.py | 5 ----- vllm/platforms/xpu.py | 9 ++++++++- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/vllm/platforms/__init__.py b/vllm/platforms/__init__.py index 758409ae1a3..2630df62d33 100644 --- a/vllm/platforms/__init__.py +++ b/vllm/platforms/__init__.py @@ -140,11 +140,6 @@ def xpu_platform_plugin() -> str | None: XPUPlatform.dist_backend = dist_backend logger.debug("Confirmed %s backend is available.", XPUPlatform.dist_backend) - else: - logger.warning( - "xccl is not enabled in this torch build, " - "communication is not available." - ) if hasattr(torch, "xpu") and torch.xpu.is_available(): is_xpu = True diff --git a/vllm/platforms/xpu.py b/vllm/platforms/xpu.py index 6e299f30ee6..3a0ea8b122c 100644 --- a/vllm/platforms/xpu.py +++ b/vllm/platforms/xpu.py @@ -34,7 +34,7 @@ class XPUPlatform(Platform): # Intel XPU's device key is "GPU" for Ray. # see https://github.com/ray-project/ray/blob/6a5eb5865eeb9ccf058a79b44f107e327e360673/python/ray/_private/accelerators/intel_gpu.py#L20 # noqa: E501 ray_device_key: str = "GPU" - dist_backend: str = "ccl" # ccl | xccl + dist_backend: str = "xccl" # xccl only device_control_env_var: str = "ZE_AFFINITY_MASK" @classmethod @@ -223,6 +223,13 @@ class XPUPlatform(Platform): @classmethod def get_device_communicator_cls(cls) -> str: + from vllm.utils.torch_utils import supports_xccl + + if not supports_xccl(): + logger.warning( + "xccl is not enabled in this torch build, communication" + " is not available." + ) return "vllm.distributed.device_communicators.xpu_communicator.XpuCommunicator" # noqa @classmethod From 4c8d1bf361c0ad4066f2c90f636b5fabe80a94eb Mon Sep 17 00:00:00 2001 From: Augusto Yao Date: Wed, 4 Feb 2026 18:04:11 +0800 Subject: [PATCH 051/810] use ORJSONResponse when available to improve the efficiency of request process (#33548) Signed-off-by: augusto.yjh --- vllm/entrypoints/pooling/embed/api_router.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/vllm/entrypoints/pooling/embed/api_router.py b/vllm/entrypoints/pooling/embed/api_router.py index 50a4018857a..c252bb43cd8 100644 --- a/vllm/entrypoints/pooling/embed/api_router.py +++ b/vllm/entrypoints/pooling/embed/api_router.py @@ -1,5 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import importlib.util +from functools import lru_cache from http import HTTPStatus from fastapi import APIRouter, Depends, Request @@ -15,9 +17,24 @@ from vllm.entrypoints.pooling.embed.protocol import ( ) from vllm.entrypoints.pooling.embed.serving import OpenAIServingEmbedding from vllm.entrypoints.utils import load_aware_call, with_cancellation +from vllm.logger import init_logger router = APIRouter() +logger = init_logger(__name__) + + +@lru_cache(maxsize=1) +def _get_json_response_cls(): + if importlib.util.find_spec("orjson") is not None: + from fastapi.responses import ORJSONResponse + + return ORJSONResponse + logger.warning_once( + "To make v1/embeddings API fast, please install orjson by `pip install orjson`" + ) + return JSONResponse + def embedding(request: Request) -> OpenAIServingEmbedding | None: return request.app.state.openai_serving_embedding @@ -54,7 +71,7 @@ async def create_embedding( content=generator.model_dump(), status_code=generator.error.code ) elif isinstance(generator, EmbeddingResponse): - return JSONResponse(content=generator.model_dump()) + return _get_json_response_cls()(content=generator.model_dump()) elif isinstance(generator, EmbeddingBytesResponse): return StreamingResponse( content=generator.content, From f79f777803ae70bc3ee2a4cfd2822e26a378dcdb Mon Sep 17 00:00:00 2001 From: Kunshang Ji Date: Wed, 4 Feb 2026 18:12:25 +0800 Subject: [PATCH 052/810] [XPU][2/N] add support unquantized moe support for xpu (#33659) Signed-off-by: Kunshang Ji --- .../scripts/hardware_ci/run-xpu-test.sh | 2 + requirements/xpu.txt | 2 +- .../layers/fused_moe/__init__.py | 4 + .../layers/fused_moe/oracle/unquantized.py | 11 +- .../fused_moe/unquantized_fused_moe_method.py | 34 +---- .../layers/fused_moe/xpu_fused_moe.py | 120 ++++++++++++++++++ 6 files changed, 139 insertions(+), 34 deletions(-) create mode 100644 vllm/model_executor/layers/fused_moe/xpu_fused_moe.py diff --git a/.buildkite/scripts/hardware_ci/run-xpu-test.sh b/.buildkite/scripts/hardware_ci/run-xpu-test.sh index 36775152f1e..33a7bd5e955 100644 --- a/.buildkite/scripts/hardware_ci/run-xpu-test.sh +++ b/.buildkite/scripts/hardware_ci/run-xpu-test.sh @@ -39,6 +39,8 @@ docker run \ python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend ray python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend mp python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --attention-backend=TRITON_ATTN + python3 examples/offline_inference/basic/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 + python3 examples/offline_inference/basic/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 --enable-expert-parallel cd tests pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py pytest -v -s v1/engine diff --git a/requirements/xpu.txt b/requirements/xpu.txt index 6fde5b8f916..f15f0dcd1d8 100644 --- a/requirements/xpu.txt +++ b/requirements/xpu.txt @@ -15,4 +15,4 @@ torch==2.10.0+xpu torchaudio torchvision -vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.0/vllm_xpu_kernels-0.1.0-cp312-cp312-linux_x86_64.whl \ No newline at end of file +vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.1/vllm_xpu_kernels-0.1.1-cp312-cp312-linux_x86_64.whl \ No newline at end of file diff --git a/vllm/model_executor/layers/fused_moe/__init__.py b/vllm/model_executor/layers/fused_moe/__init__.py index 03be6f8b6a4..edf7544b9e0 100644 --- a/vllm/model_executor/layers/fused_moe/__init__.py +++ b/vllm/model_executor/layers/fused_moe/__init__.py @@ -100,6 +100,9 @@ if HAS_TRITON: from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import ( TritonOrDeepGemmExperts, ) + from vllm.model_executor.layers.fused_moe.xpu_fused_moe import ( + XPUExperts, + ) __all__ += [ "AiterExperts", @@ -117,6 +120,7 @@ if HAS_TRITON: "DeepGemmExperts", "BatchedDeepGemmExperts", "TritonOrDeepGemmExperts", + "XPUExperts", ] else: # Some model classes directly use the custom ops. Add placeholders diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index e79670f9de5..a8754d6d6e4 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -46,7 +46,6 @@ class UnquantizedMoeBackend(Enum): UNSUPPORTED_BACKEND = [ UnquantizedMoeBackend.FLASHINFER_TRTLLM, UnquantizedMoeBackend.CPU, - UnquantizedMoeBackend.XPU, UnquantizedMoeBackend.TPU, UnquantizedMoeBackend.OOT, ] @@ -196,4 +195,14 @@ def make_unquantized_moe_kernel( quant_config=quant_config, ), ) + elif backend == UnquantizedMoeBackend.XPU: + from vllm.model_executor.layers.fused_moe import XPUExperts + + kernel = mk.FusedMoEModularKernel( + MoEPrepareAndFinalizeNoEP(), + XPUExperts( + moe_config=moe_config, + quant_config=quant_config, + ), + ) return kernel, use_inplace diff --git a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py index 2ddaf272b14..6fdd8ecf79b 100644 --- a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py +++ b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py @@ -40,7 +40,7 @@ from vllm.model_executor.utils import replace_parameter, set_weight_attrs from vllm.platforms import current_platform from vllm.platforms.interface import CpuArchEnum -if current_platform.is_cuda_alike(): +if current_platform.is_cuda_alike() or current_platform.is_xpu(): from .fused_batched_moe import BatchedTritonExperts from .fused_moe import TritonExperts else: @@ -71,7 +71,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp): self.kernel: mk.FusedMoEModularKernel | None = None self._is_monolithic = ( current_platform.is_cpu() - or current_platform.is_xpu() or self.unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM ) @@ -82,8 +81,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp): """Select the monolithic implementation based on platform.""" if current_platform.is_cpu(): return self.forward_monolithic_cpu - elif current_platform.is_xpu(): - return self.forward_monolithic_xpu else: return self.forward_monolithic_cuda @@ -256,16 +253,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp): ) layer.w13_weight = Parameter(w13_weights_shuffled, requires_grad=False) layer.w2_weight = Parameter(w2_weights_shuffled, requires_grad=False) - elif self.unquantized_backend == UnquantizedMoeBackend.XPU: - import intel_extension_for_pytorch as ipex - - ep_rank_start = self.moe.ep_rank * self.moe.num_local_experts - self.ipex_fusion = ipex.llm.modules.GatedMLPMOE( - layer.w13_weight, - layer.w2_weight, - use_prepack=True, - experts_start_id=ep_rank_start, - ) elif self.unquantized_backend == UnquantizedMoeBackend.CPU: from vllm.model_executor.layers.fused_moe import cpu_fused_moe @@ -297,7 +284,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp): self.cpu_fused_moe = cpu_fused_moe.CPUFusedMOE(layer) else: self.cpu_fused_moe = cpu_fused_moe.CPUFusedMOE(layer) - elif current_platform.is_cuda_alike(): + elif current_platform.is_cuda_alike() or current_platform.is_xpu(): self._setup_kernel( layer=layer, w13=layer.w13_weight, @@ -399,20 +386,3 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp): layer.apply_router_weight_on_input, layer.activation, ) - - def forward_monolithic_xpu( - self, - layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821 - x: torch.Tensor, - router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - return self.ipex_fusion( - x, - layer.use_grouped_topk, - layer.top_k, - router_logits, - layer.renormalize, - layer.topk_group, - layer.num_expert_group, - custom_routing_function=layer.custom_routing_function, - ) diff --git a/vllm/model_executor/layers/fused_moe/xpu_fused_moe.py b/vllm/model_executor/layers/fused_moe/xpu_fused_moe.py new file mode 100644 index 00000000000..cfb88f6afb1 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/xpu_fused_moe.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEParallelConfig, +) +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceNoOP, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kFp8StaticTensorSym, +) +from vllm.platforms import current_platform + +if current_platform.is_xpu(): + from vllm_xpu_kernels.fused_moe_interface import xpu_fused_moe + + +class XPUExperts(mk.FusedMoEPermuteExpertsUnpermute): + @property + def expects_unquantized_inputs(self) -> bool: + return True + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + @staticmethod + def _supports_current_device() -> bool: + return current_platform.is_xpu() + + @staticmethod + def _supports_no_act_and_mul() -> bool: + return False + + @staticmethod + def _supports_activation(activation: str) -> bool: + return activation in ["silu", "gelu", "swigluoai"] + + @staticmethod + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: + return True + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + # TODO: dispatch based on device. + SUPPORTED_W_A = [ + (None, None), + (kFp8StaticTensorSym, None), + ] + return (weight_key, activation_key) in SUPPORTED_W_A + + def supports_chunking(self) -> bool: + return False + + def supports_expert_map(self) -> bool: + return True + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + return TopKWeightAndReduceNoOP() + + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: str, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + workspace1 = (0,) + workspace2 = (0,) + output = (M, K) + return (workspace1, workspace2, output) + + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: str, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool, + ): + topk = topk_ids.size(-1) + xpu_fused_moe( + hidden_states=hidden_states, + w13=w1, + w13_scales=a1q_scale, + w13_bias=self.w1_bias, + w2=w2, + w2_scales=a2_scale, + w2_bias=self.w2_bias, + topk_weights=topk_weights, + topk_ids=topk_ids, + n_experts_per_token=topk, + activation=activation, + num_experts=self.moe_config.num_local_experts, + ep_rank=self.moe_config.ep_rank, + ep_size=self.moe_config.ep_size, + output=output, + ) + return From bcd2f74c0d1e85a2da4dcb41849ad75a7e3fdaf4 Mon Sep 17 00:00:00 2001 From: Zhengxu Chen Date: Wed, 4 Feb 2026 05:12:53 -0500 Subject: [PATCH 053/810] [compile] Clean up AOT compile bypass on evaluate_guards. (#33578) Signed-off-by: zhxchen17 --- vllm/compilation/wrapper.py | 8 -------- 1 file changed, 8 deletions(-) diff --git a/vllm/compilation/wrapper.py b/vllm/compilation/wrapper.py index 4c1b0466670..850ddae9ab9 100644 --- a/vllm/compilation/wrapper.py +++ b/vllm/compilation/wrapper.py @@ -151,14 +151,6 @@ class TorchCompileWithNoGuardsWrapper: "requires VLLM_USE_BYTECODE_HOOK=0. " ) - if envs.VLLM_USE_AOT_COMPILE: - # disabled until https://github.com/pytorch/pytorch/pull/169239 - # is picked up. - assert ds_type != DynamicShapesType.BACKED, ( - "evaluate_guards for backed shapes requires " - "VLLM_USE_AOT_COMPILE=False. " - ) - options["guard_filter_fn"] = lambda x: [ entry.guard_type == "SHAPE_ENV" for entry in x ] From a208439537a9071668b99dc8089db0dd8995034a Mon Sep 17 00:00:00 2001 From: Zhengxu Chen Date: Wed, 4 Feb 2026 05:56:45 -0500 Subject: [PATCH 054/810] [compile] Remove runner type from ignored caching factor list. (#33712) Signed-off-by: zhxchen17 --- vllm/config/model.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/vllm/config/model.py b/vllm/config/model.py index 3bb8e71770a..2686df4c23e 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -105,8 +105,8 @@ class ModelConfig: content for `model_name` tag in metrics output when `served_model_name` is not specified.""" model_weights: str = "" - """Original model weights path. Used when the model is pulled from object - storage (e.g., RunAI) to preserve the original URI while `model` points to + """Original model weights path. Used when the model is pulled from object + storage (e.g., RunAI) to preserve the original URI while `model` points to the local directory.""" runner: RunnerOption = "auto" """The type of model runner to use. Each vLLM instance only supports one @@ -324,7 +324,6 @@ class ModelConfig: the final hidden states. """ ignored_factors = { - "runner", "convert", "tokenizer", "tokenizer_mode", From 8e3269086916d81b010a2d6209784a106ac2994a Mon Sep 17 00:00:00 2001 From: Or Ozeri Date: Wed, 4 Feb 2026 13:16:34 +0200 Subject: [PATCH 055/810] [KV Connector][BugFix] scheduler: Delay freeing blocks of aborted async loads (#32255) Fixes a not-yet-reported case where it was possible for blocks to be freed by an abort before an async transfer completed, resulting in corrupted KV data. Signed-off-by: Or Ozeri --- tests/v1/core/test_scheduler.py | 49 +++++++++++++++++ .../unit/test_offloading_connector.py | 55 ++++++++++++++++++- vllm/v1/core/sched/scheduler.py | 27 +++++++-- 3 files changed, 124 insertions(+), 7 deletions(-) diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 580cc70ff34..063d0a644fa 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -3473,3 +3473,52 @@ def test_prepend_skipped_requests_order(): # verify waiting order is preserved assert list(scheduler.waiting) == expected_waiting_reqs + + +def test_abort_request_waiting_for_remote_kvs(): + scheduler = create_scheduler(use_kv_connector=True) + + # add a single request + request = create_requests(num_requests=1)[0] + scheduler.add_request(request) + + # set request to waiting for remote KVs, and abort it + request.status = RequestStatus.WAITING_FOR_REMOTE_KVS + scheduler.finish_requests((request.request_id,), RequestStatus.FINISHED_ABORTED) + assert request.status == RequestStatus.FINISHED_ABORTED + + # verify request is not deleted + assert request.request_id in scheduler.requests + + # finish recving request + scheduler_output = scheduler.schedule() + model_runner_output = ModelRunnerOutput( + req_ids=[], + req_id_to_index={}, + kv_connector_output=KVConnectorOutput(finished_recving={request.request_id}), + ) + scheduler.update_from_output(scheduler_output, model_runner_output) + + # assert request is deleted + assert request.request_id not in scheduler.requests + assert not scheduler.finished_recving_kv_req_ids + + +def test_abort_request_finished_recving(): + scheduler = create_scheduler(use_kv_connector=True) + + # add a single request + request = create_requests(num_requests=1)[0] + scheduler.add_request(request) + + # set request to waiting for remote KVs, finished but not yet updated + request.status = RequestStatus.WAITING_FOR_REMOTE_KVS + scheduler.finished_recving_kv_req_ids.add(request.request_id) + + # abort request + scheduler.finish_requests((request.request_id,), RequestStatus.FINISHED_ABORTED) + assert request.status == RequestStatus.FINISHED_ABORTED + + # verify request is deleted + assert request.request_id not in scheduler.requests + assert not scheduler.finished_recving_kv_req_ids diff --git a/tests/v1/kv_connector/unit/test_offloading_connector.py b/tests/v1/kv_connector/unit/test_offloading_connector.py index 1805f009db0..5b84202a581 100644 --- a/tests/v1/kv_connector/unit/test_offloading_connector.py +++ b/tests/v1/kv_connector/unit/test_offloading_connector.py @@ -42,7 +42,7 @@ from vllm.v1.kv_offload.worker.worker import ( TransferSpec, ) from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, KVConnectorOutput -from vllm.v1.request import Request +from vllm.v1.request import Request, RequestStatus from .utils import ( EOS_TOKEN_ID, @@ -355,7 +355,7 @@ class RequestRunner: self.scheduler.update_from_output(scheduler_output, model_runner_output) if ( - prev_token_id is EOS_TOKEN_ID + prev_token_id == EOS_TOKEN_ID and prev_token_id != token_id and self.scheduler.requests ): @@ -730,6 +730,57 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner): assert transfer_jobs == list(runner.offloading_spec.handler.transfer_specs) +def test_abort_loading_requests(request_runner): + offloaded_block_size = 12 + gpu_block_size = 4 + num_gpu_blocks = 100 + + runner = request_runner( + offloaded_block_size=offloaded_block_size, + gpu_block_size=gpu_block_size, + num_gpu_blocks=num_gpu_blocks, + ) + + # store 1 blocks + runner.new_request(token_ids=[0] * offloaded_block_size) + runner.manager.prepare_store.side_effect = ( + lambda block_hashes: generate_store_output(block_hashes) + ) + runner.run( + decoded_tokens=[EOS_TOKEN_ID], + expected_stored_gpu_block_indexes=(0, 1, 2), + ) + + # start a request to load the first block, but don't complete + runner.scheduler.reset_prefix_cache() + runner.new_request(token_ids=[0] * offloaded_block_size) + runner.manager.lookup.return_value = 1 + runner.run( + decoded_tokens=[], + complete_transfers=False, + ) + + # request triggered a load + transfer_jobs = list(runner.offloading_spec.handler.transfer_specs) + assert transfer_jobs + + # abort request + req_id = str(runner.req_id) + runner.scheduler.finish_requests((req_id,), RequestStatus.FINISHED_ABORTED) + + # verify request is not deleted + assert req_id in runner.scheduler.requests + + # complete loading request + runner.run( + decoded_tokens=[], + expected_loaded_gpu_block_indexes=(0, 1, 2), + ) + + # assert request is deleted + assert req_id not in runner.scheduler.requests + + class TestOffloadingConnectorStats: """Tests for OffloadingConnector stats reconstruction and operations.""" diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 9308f9ed1c1..1544d847c8f 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1670,19 +1670,30 @@ class Scheduler(SchedulerInterface): # Second pass: set status and free requests for request in valid_requests: - request.status = finished_status - self._free_request(request) + delay_free_blocks = False + if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS: + delay_free_blocks = ( + request.request_id not in self.finished_recving_kv_req_ids + ) + self.finished_recving_kv_req_ids.discard(request.request_id) + self.failed_recving_kv_req_ids.discard(request.request_id) - def _free_request(self, request: Request) -> dict[str, Any] | None: + request.status = finished_status + self._free_request(request, delay_free_blocks=delay_free_blocks) + + def _free_request( + self, request: Request, delay_free_blocks: bool = False + ) -> dict[str, Any] | None: assert request.is_finished() - delay_free_blocks, kv_xfer_params = self._connector_finished(request) + connector_delay_free_blocks, kv_xfer_params = self._connector_finished(request) self.encoder_cache_manager.free(request) request_id = request.request_id self.finished_req_ids.add(request_id) if self.finished_req_ids_dict is not None: self.finished_req_ids_dict[request.client_index].add(request_id) + delay_free_blocks |= connector_delay_free_blocks if not delay_free_blocks: self._free_blocks(request) @@ -1954,7 +1965,13 @@ class Scheduler(SchedulerInterface): # KV Connector:: update recv and send status from last step. for req_id in kv_connector_output.finished_recving or (): logger.debug("Finished recving KV transfer for request %s", req_id) - self.finished_recving_kv_req_ids.add(req_id) + assert req_id in self.requests + req = self.requests[req_id] + if req.status == RequestStatus.WAITING_FOR_REMOTE_KVS: + self.finished_recving_kv_req_ids.add(req_id) + else: + assert RequestStatus.is_finished(req.status) + self._free_blocks(self.requests[req_id]) for req_id in kv_connector_output.finished_sending or (): logger.debug("Finished sending KV transfer for request %s", req_id) assert req_id in self.requests From 824058076c56164a3772a5f5829bd9662507e5a3 Mon Sep 17 00:00:00 2001 From: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com> Date: Wed, 4 Feb 2026 15:20:52 +0400 Subject: [PATCH 056/810] [PERF] Change GDN Attention State Layout from [N, HV, K, V] to [N, HV, V, K] (#33291) Signed-off-by: Vadim Gimpelson --- vllm/model_executor/layers/fla/ops/chunk.py | 8 +-- .../layers/fla/ops/chunk_delta_h.py | 70 +++++++++---------- vllm/model_executor/layers/fla/ops/chunk_o.py | 8 +-- .../layers/fla/ops/fused_recurrent.py | 32 ++++----- vllm/model_executor/layers/fla/ops/kda.py | 2 +- .../layers/mamba/mamba_utils.py | 2 +- 6 files changed, 61 insertions(+), 61 deletions(-) diff --git a/vllm/model_executor/layers/fla/ops/chunk.py b/vllm/model_executor/layers/fla/ops/chunk.py index 4c8bf9f4399..958464b6941 100644 --- a/vllm/model_executor/layers/fla/ops/chunk.py +++ b/vllm/model_executor/layers/fla/ops/chunk.py @@ -138,11 +138,11 @@ def chunk_gated_delta_rule( Scale factor for the RetNet attention scores. If not provided, it will default to `1 / sqrt(K)`. Default: `None`. initial_state (Optional[torch.Tensor]): - Initial state of shape `[N, H, K, V]` for `N` input sequences. + Initial state of shape `[N, H, V, K]` for `N` input sequences. For equal-length input sequences, `N` equals the batch size `B`. Default: `None`. output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + Whether to output the final state of shape `[N, H, V, K]`. Default: `False`. cu_seqlens (torch.LongTensor): Cumulative sequence lengths of shape `[N+1]` used for variable-length training, consistent with the FlashAttention API. @@ -154,7 +154,7 @@ def chunk_gated_delta_rule( o (torch.Tensor): Outputs of shape `[B, T, H, V]` if `head_first=False` else `[B, H, T, V]`. final_state (torch.Tensor): - Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. + Final state of shape `[N, H, V, K]` if `output_final_state=True` else `None`. Examples:: >>> import torch @@ -168,7 +168,7 @@ def chunk_gated_delta_rule( >>> v = torch.randn(B, T, H, V, dtype=torch.bfloat16, device='cuda') >>> beta = torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda').sigmoid() >>> g = F.logsigmoid(torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda')) - >>> h0 = torch.randn(B, H, K, V, dtype=torch.bfloat16, device='cuda') + >>> h0 = torch.randn(B, H, V, K, dtype=torch.bfloat16, device='cuda') >>> o, ht = chunk_gated_delta_rule( q, k, v, g, beta, initial_state=h0, diff --git a/vllm/model_executor/layers/fla/ops/chunk_delta_h.py b/vllm/model_executor/layers/fla/ops/chunk_delta_h.py index f0b78b65c4a..98a3d61e436 100644 --- a/vllm/model_executor/layers/fla/ops/chunk_delta_h.py +++ b/vllm/model_executor/layers/fla/ops/chunk_delta_h.py @@ -81,70 +81,70 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( NT = tl.cdiv(T, BT) boh = i_n * NT - # [BK, BV] - b_h1 = tl.zeros([64, BV], dtype=tl.float32) + # [BV, BK] + b_h1 = tl.zeros([BV, 64], dtype=tl.float32) if K > 64: - b_h2 = tl.zeros([64, BV], dtype=tl.float32) + b_h2 = tl.zeros([BV, 64], dtype=tl.float32) if K > 128: - b_h3 = tl.zeros([64, BV], dtype=tl.float32) + b_h3 = tl.zeros([BV, 64], dtype=tl.float32) if K > 192: - b_h4 = tl.zeros([64, BV], dtype=tl.float32) + b_h4 = tl.zeros([BV, 64], dtype=tl.float32) # calculate offset - h += ((boh * H + i_h) * K * V).to(tl.int64) + h += ((boh * H + i_h) * V * K).to(tl.int64) v += ((bos * H + i_h) * V).to(tl.int64) k += ((bos * Hg + i_h // (H // Hg)) * K).to(tl.int64) w += ((bos * H + i_h) * K).to(tl.int64) if SAVE_NEW_VALUE: v_new += ((bos * H + i_h) * V).to(tl.int64) stride_v = H * V - stride_h = H * K * V + stride_h = H * V * K stride_k = Hg * K stride_w = H * K if USE_INITIAL_STATE: - h0 = h0 + i_nh * K * V + h0 = h0 + i_nh * V * K if STORE_FINAL_STATE: - ht = ht + i_nh * K * V + ht = ht + i_nh * V * K # load initial state if USE_INITIAL_STATE: - p_h0_1 = tl.make_block_ptr(h0, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) + p_h0_1 = tl.make_block_ptr(h0, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32) if K > 64: p_h0_2 = tl.make_block_ptr( - h0, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0) + h0, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0) ) b_h2 += tl.load(p_h0_2, boundary_check=(0, 1)).to(tl.float32) if K > 128: p_h0_3 = tl.make_block_ptr( - h0, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0) + h0, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0) ) b_h3 += tl.load(p_h0_3, boundary_check=(0, 1)).to(tl.float32) if K > 192: p_h0_4 = tl.make_block_ptr( - h0, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0) + h0, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0) ) b_h4 += tl.load(p_h0_4, boundary_check=(0, 1)).to(tl.float32) # main recurrence for i_t in range(NT): p_h1 = tl.make_block_ptr( - h + i_t * stride_h, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0) + h + i_t * stride_h, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0) ) tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1)) if K > 64: p_h2 = tl.make_block_ptr( - h + i_t * stride_h, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0) + h + i_t * stride_h, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0) ) tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1)) if K > 128: p_h3 = tl.make_block_ptr( - h + i_t * stride_h, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0) + h + i_t * stride_h, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0) ) tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1)) if K > 192: p_h4 = tl.make_block_ptr( - h + i_t * stride_h, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0) + h + i_t * stride_h, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0) ) tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1)) @@ -152,25 +152,25 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( w, (T, K), (stride_w, 1), (i_t * BT, 0), (BT, 64), (1, 0) ) b_w = tl.load(p_w, boundary_check=(0, 1)) - b_v = tl.dot(b_w, b_h1.to(b_w.dtype)) + b_v = tl.dot(b_w, tl.trans(b_h1).to(b_w.dtype)) if K > 64: p_w = tl.make_block_ptr( w, (T, K), (stride_w, 1), (i_t * BT, 64), (BT, 64), (1, 0) ) b_w = tl.load(p_w, boundary_check=(0, 1)) - b_v += tl.dot(b_w, b_h2.to(b_w.dtype)) + b_v += tl.dot(b_w, tl.trans(b_h2).to(b_w.dtype)) if K > 128: p_w = tl.make_block_ptr( w, (T, K), (stride_w, 1), (i_t * BT, 128), (BT, 64), (1, 0) ) b_w = tl.load(p_w, boundary_check=(0, 1)) - b_v += tl.dot(b_w, b_h3.to(b_w.dtype)) + b_v += tl.dot(b_w, tl.trans(b_h3).to(b_w.dtype)) if K > 192: p_w = tl.make_block_ptr( w, (T, K), (stride_w, 1), (i_t * BT, 192), (BT, 64), (1, 0) ) b_w = tl.load(p_w, boundary_check=(0, 1)) - b_v += tl.dot(b_w, b_h4.to(b_w.dtype)) + b_v += tl.dot(b_w, tl.trans(b_h4).to(b_w.dtype)) p_v = tl.make_block_ptr( v, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) ) @@ -207,7 +207,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( mask=(o_k1 < K), other=0.0, ) - b_h1 *= exp(b_gk_last1)[:, None] + b_h1 *= exp(b_gk_last1)[None, :] if K > 64: o_k2 = 64 + o_k1 b_gk_last2 = tl.load( @@ -215,7 +215,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( mask=(o_k2 < K), other=0.0, ) - b_h2 *= exp(b_gk_last2)[:, None] + b_h2 *= exp(b_gk_last2)[None, :] if K > 128: o_k3 = 128 + o_k1 b_gk_last3 = tl.load( @@ -223,7 +223,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( mask=(o_k3 < K), other=0.0, ) - b_h3 *= exp(b_gk_last3)[:, None] + b_h3 *= exp(b_gk_last3)[None, :] if K > 192: o_k4 = 192 + o_k1 b_gk_last4 = tl.load( @@ -231,49 +231,49 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( mask=(o_k4 < K), other=0.0, ) - b_h4 *= exp(b_gk_last4)[:, None] + b_h4 *= exp(b_gk_last4)[None, :] b_v = b_v.to(k.dtype.element_ty) p_k = tl.make_block_ptr( k, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1) ) b_k = tl.load(p_k, boundary_check=(0, 1)) - b_h1 += tl.dot(b_k, b_v) + b_h1 += tl.trans(tl.dot(b_k, b_v)) if K > 64: p_k = tl.make_block_ptr( k, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1) ) b_k = tl.load(p_k, boundary_check=(0, 1)) - b_h2 += tl.dot(b_k, b_v) + b_h2 += tl.trans(tl.dot(b_k, b_v)) if K > 128: p_k = tl.make_block_ptr( k, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1) ) b_k = tl.load(p_k, boundary_check=(0, 1)) - b_h3 += tl.dot(b_k, b_v) + b_h3 += tl.trans(tl.dot(b_k, b_v)) if K > 192: p_k = tl.make_block_ptr( k, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1) ) b_k = tl.load(p_k, boundary_check=(0, 1)) - b_h4 += tl.dot(b_k, b_v) + b_h4 += tl.trans(tl.dot(b_k, b_v)) # epilogue if STORE_FINAL_STATE: - p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) + p_ht = tl.make_block_ptr(ht, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) if K > 64: p_ht = tl.make_block_ptr( - ht, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0) + ht, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0) ) tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) if K > 128: p_ht = tl.make_block_ptr( - ht, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0) + ht, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0) ) tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) if K > 192: p_ht = tl.make_block_ptr( - ht, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0) + ht, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0) ) tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) @@ -312,9 +312,9 @@ def chunk_gated_delta_rule_fwd_h( ) assert K <= 256, "current kernel does not support head dimension larger than 256." - h = k.new_empty(B, NT, H, K, V) + h = k.new_empty(B, NT, H, V, K) final_state = ( - k.new_empty(N, H, K, V, dtype=torch.float32) if output_final_state else None + k.new_empty(N, H, V, K, dtype=torch.float32) if output_final_state else None ) v_new = torch.empty_like(u) if save_new_value else None diff --git a/vllm/model_executor/layers/fla/ops/chunk_o.py b/vllm/model_executor/layers/fla/ops/chunk_o.py index 4e8e04c1d48..2ccf1d4e254 100644 --- a/vllm/model_executor/layers/fla/ops/chunk_o.py +++ b/vllm/model_executor/layers/fla/ops/chunk_o.py @@ -85,7 +85,7 @@ def chunk_fwd_kernel_o( k += (bos * Hg + i_h // (H // Hg)) * K v += (bos * H + i_h) * V o += (bos * H + i_h) * V - h += (i_tg * H + i_h).to(tl.int64) * K * V + h += (i_tg * H + i_h).to(tl.int64) * V * K b_o = tl.zeros([BT, BV], dtype=tl.float32) b_A = tl.zeros([BT, BT], dtype=tl.float32) @@ -98,17 +98,17 @@ def chunk_fwd_kernel_o( k, (K, T), (1, Hg * K), (i_k * BK, i_t * BT), (BK, BT), (0, 1) ) p_h = tl.make_block_ptr( - h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0) + h, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0) ) # [BT, BK] b_q = tl.load(p_q, boundary_check=(0, 1)) # [BK, BT] b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BK, BV] + # [BV, BK] b_h = tl.load(p_h, boundary_check=(0, 1)) # [BT, BK] @ [BK, BV] -> [BT, BV] - b_o += tl.dot(b_q, b_h) + b_o += tl.dot(b_q, tl.trans(b_h)) # [BT, BK] @ [BK, BT] -> [BT, BT] b_A += tl.dot(b_q, b_k) diff --git a/vllm/model_executor/layers/fla/ops/fused_recurrent.py b/vllm/model_executor/layers/fla/ops/fused_recurrent.py index 66540f066d6..67d77e88294 100644 --- a/vllm/model_executor/layers/fla/ops/fused_recurrent.py +++ b/vllm/model_executor/layers/fla/ops/fused_recurrent.py @@ -97,9 +97,9 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( mask_k = o_k < K mask_v = o_v < V - mask_h = mask_k[:, None] & mask_v[None, :] + mask_h = mask_v[:, None] & mask_k[None, :] - b_h = tl.zeros([BK, BV], dtype=tl.float32) + b_h = tl.zeros([BV, BK], dtype=tl.float32) if USE_INITIAL_STATE: if IS_CONTINUOUS_BATCHING: if IS_SPEC_DECODING: @@ -115,8 +115,8 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( return p_h0 = h0 + state_idx * stride_init_state_token else: - p_h0 = h0 + bos * HV * K * V - p_h0 = p_h0 + i_hv * K * V + o_k[:, None] * V + o_v[None, :] + p_h0 = h0 + bos * HV * V * K + p_h0 = p_h0 + i_hv * V * K + o_v[:, None] * K + o_k[None, :] b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) for i_t in range(0, T): @@ -128,24 +128,24 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6) b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6) b_q = b_q * scale - # [BK, BV] + # [BV, BK] if not IS_KDA: b_g = tl.load(p_g).to(tl.float32) b_h *= exp(b_g) else: b_gk = tl.load(p_gk).to(tl.float32) - b_h *= exp(b_gk[:, None]) + b_h *= exp(b_gk[None, :]) # [BV] - b_v -= tl.sum(b_h * b_k[:, None], 0) + b_v -= tl.sum(b_h * b_k[None, :], 1) if IS_BETA_HEADWISE: b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32) else: b_beta = tl.load(p_beta).to(tl.float32) b_v *= b_beta - # [BK, BV] - b_h += b_k[:, None] * b_v[None, :] + # [BV, BK] + b_h += b_v[:, None] * b_k[None, :] # [BV] - b_o = tl.sum(b_h * b_q[:, None], 0) + b_o = tl.sum(b_h * b_q[None, :], 1) tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) # keep the states for multi-query tokens @@ -157,11 +157,11 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( # Only store if state index is valid (not PAD_SLOT_ID) if final_state_idx >= 0: p_ht = ht + final_state_idx * stride_final_state_token - p_ht = p_ht + i_hv * K * V + o_k[:, None] * V + o_v[None, :] + p_ht = p_ht + i_hv * V * K + o_v[:, None] * K + o_k[None, :] tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) else: p_ht = ht + (bos + i_t) * stride_final_state_token - p_ht = p_ht + i_hv * K * V + o_k[:, None] * V + o_v[None, :] + p_ht = p_ht + i_hv * V * K + o_v[:, None] * K + o_k[None, :] tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) p_q += H * K @@ -202,7 +202,7 @@ def fused_recurrent_gated_delta_rule_fwd( if inplace_final_state: final_state = initial_state else: - final_state = q.new_empty(T, HV, K, V, dtype=initial_state.dtype) + final_state = q.new_empty(T, HV, V, K, dtype=initial_state.dtype) stride_init_state_token = initial_state.stride(0) stride_final_state_token = final_state.stride(0) @@ -318,7 +318,7 @@ def fused_recurrent_gated_delta_rule( Scale factor for the RetNet attention scores. If not provided, it will default to `1 / sqrt(K)`. Default: `None`. initial_state (Optional[torch.Tensor]): - Initial state of shape `[N, HV, K, V]` for `N` input sequences. + Initial state of shape `[N, HV, V, K]` for `N` input sequences. For equal-length input sequences, `N` equals the batch size `B`. Default: `None`. inplace_final_state: bool: @@ -336,7 +336,7 @@ def fused_recurrent_gated_delta_rule( o (torch.Tensor): Outputs of shape `[B, T, HV, V]`. final_state (torch.Tensor): - Final state of shape `[N, HV, K, V]`. + Final state of shape `[N, HV, V, K]`. Examples:: >>> import torch @@ -350,7 +350,7 @@ def fused_recurrent_gated_delta_rule( >>> v = torch.randn(B, T, HV, V, device='cuda') >>> g = F.logsigmoid(torch.rand(B, T, HV, device='cuda')) >>> beta = torch.rand(B, T, HV, device='cuda').sigmoid() - >>> h0 = torch.randn(B, HV, K, V, device='cuda') + >>> h0 = torch.randn(B, HV, V, K, device='cuda') >>> o, ht = fused_gated_recurrent_delta_rule( q, k, v, g, beta, initial_state=h0, diff --git a/vllm/model_executor/layers/fla/ops/kda.py b/vllm/model_executor/layers/fla/ops/kda.py index 700f287ca45..7145933e7ed 100644 --- a/vllm/model_executor/layers/fla/ops/kda.py +++ b/vllm/model_executor/layers/fla/ops/kda.py @@ -55,7 +55,7 @@ def fused_recurrent_kda_fwd( if inplace_final_state: final_state = initial_state else: - final_state = q.new_empty(T, HV, K, V, dtype=initial_state.dtype) + final_state = q.new_empty(T, HV, V, K, dtype=initial_state.dtype) stride_init_state_token = initial_state.stride(0) stride_final_state_token = final_state.stride(0) diff --git a/vllm/model_executor/layers/mamba/mamba_utils.py b/vllm/model_executor/layers/mamba/mamba_utils.py index 816f76bfa06..7181ada1c2e 100644 --- a/vllm/model_executor/layers/mamba/mamba_utils.py +++ b/vllm/model_executor/layers/mamba/mamba_utils.py @@ -191,8 +191,8 @@ class MambaStateShapeCalculator: temporal_state_shape = ( divide(num_v_heads, tp_world_size), - head_k_dim, head_v_dim, + head_k_dim, ) return conv_state_shape, temporal_state_shape From f8516a1ab95febcf131a37478914031f50fdd9db Mon Sep 17 00:00:00 2001 From: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Date: Wed, 4 Feb 2026 07:15:29 -0500 Subject: [PATCH 057/810] [Bugfix][Model] Fix audio-in-video support for Qwen2.5-Omni and Qwen3-Omni (#33605) Signed-off-by: linyueqian Signed-off-by: Roger Wang Co-authored-by: Roger Wang --- .../models/qwen2_5_omni_thinker.py | 126 +++++++++++++++++- .../models/qwen3_omni_moe_thinker.py | 58 ++++++-- 2 files changed, 172 insertions(+), 12 deletions(-) diff --git a/vllm/model_executor/models/qwen2_5_omni_thinker.py b/vllm/model_executor/models/qwen2_5_omni_thinker.py index 5152a73de06..3b50ae74d05 100644 --- a/vllm/model_executor/models/qwen2_5_omni_thinker.py +++ b/vllm/model_executor/models/qwen2_5_omni_thinker.py @@ -113,6 +113,95 @@ except (ImportError, ModuleNotFoundError): logger = init_logger(__name__) +def check_interleaved_audio_video( + is_video: torch.Tensor, + is_audio: torch.Tensor, + num_video: int, + num_audio: int, +) -> bool: + """ + Check if video and audio positions are interleaved in the multimodal region. + + Returns: + True if video and audio tokens are interleaved, False otherwise. + """ + if num_video == 0 or num_audio == 0: + return False + + video_pos = is_video.nonzero(as_tuple=True)[0] + audio_pos = is_audio.nonzero(as_tuple=True)[0] + + return ( + video_pos[0].item() < audio_pos[-1].item() + and audio_pos[0].item() < video_pos[-1].item() + ) + + +def merge_interleaved_embeddings( + inputs_embeds: torch.Tensor, + multimodal_embeddings: "MultiModalEmbeddings", + is_video: torch.Tensor, + is_audio: torch.Tensor, + is_multimodal: torch.Tensor, + num_video: int, + num_audio: int, +) -> torch.Tensor: + """ + Merge embeddings for interleaved audio-in-video sequences. + + When use_audio_in_video=True, video and audio tokens are interleaved in + the token sequence, but embeddings are provided as separate contiguous + tensors (video first, then audio). This function reorders video and audio + embeddings to match sequence position order and scatters them efficiently. + + Args: + inputs_embeds: The input embeddings tensor to merge into. + multimodal_embeddings: List of embedding tensors (video, audio, other). + is_video: Boolean mask for video token positions. + is_audio: Boolean mask for audio token positions. + is_multimodal: Boolean mask for all multimodal token positions. + num_video: Total count of video tokens. + num_audio: Total count of audio tokens. + + Returns: + The merged inputs_embeds tensor with multimodal embeddings scattered + to their correct positions. + """ + # Categorize embeddings by modality based on token counts. + # Embeddings come grouped by modality but order varies (e.g., image, video, audio + # or video, audio depending on input kwargs order). + video_embeds: list[torch.Tensor] = [] + audio_embeds: list[torch.Tensor] = [] + other_embeds: list[torch.Tensor] = [] + video_remaining = num_video + audio_remaining = num_audio + + for emb in multimodal_embeddings: + n = emb.shape[0] + if video_remaining > 0 and n <= video_remaining: + video_embeds.append(emb) + video_remaining -= n + elif audio_remaining > 0 and n <= audio_remaining: + audio_embeds.append(emb) + audio_remaining -= n + else: + other_embeds.append(emb) + + # Scatter each modality to its positions + if video_embeds: + video_positions = is_video.nonzero(as_tuple=True)[0] + inputs_embeds[video_positions] = torch.cat(video_embeds, dim=0) + if audio_embeds: + audio_positions = is_audio.nonzero(as_tuple=True)[0] + inputs_embeds[audio_positions] = torch.cat(audio_embeds, dim=0) + if other_embeds: + other_mask = is_multimodal & ~is_video & ~is_audio + other_positions = other_mask.nonzero(as_tuple=True)[0] + inputs_embeds[other_positions] = torch.cat(other_embeds, dim=0) + + return inputs_embeds + + class Qwen2_5OmniAudioFeatureInputs(TensorSchema): """ Dimensions: @@ -1286,17 +1375,48 @@ class Qwen2_5OmniThinkerForConditionalGeneration( is_multimodal: torch.Tensor | None = None, handle_oov_mm_token: bool = False, ) -> torch.Tensor: - # This is to satisfy the type checker for each overload + from .utils import _merge_multimodal_embeddings + if multimodal_embeddings is None or is_multimodal is None: return super().embed_input_ids(input_ids) - return super().embed_input_ids( + inputs_embeds = self._embed_text_input_ids( input_ids, - multimodal_embeddings=multimodal_embeddings, + self.get_language_model().embed_input_ids, is_multimodal=is_multimodal, handle_oov_mm_token=handle_oov_mm_token, ) + if len(multimodal_embeddings) == 0: + return inputs_embeds + + # Check for audio-in-video: interleaved video and audio tokens + # in the multimodal region. + video_token_id = self.config.video_token_index + audio_token_id = self.config.audio_token_index + + is_video = is_multimodal & (input_ids == video_token_id) + is_audio = is_multimodal & (input_ids == audio_token_id) + + num_video = is_video.sum().item() + num_audio = is_audio.sum().item() + + if check_interleaved_audio_video(is_video, is_audio, num_video, num_audio): + return merge_interleaved_embeddings( + inputs_embeds, + multimodal_embeddings, + is_video, + is_audio, + is_multimodal, + num_video, + num_audio, + ) + + # Default: standard merge (no interleaving) + return _merge_multimodal_embeddings( + inputs_embeds, multimodal_embeddings, is_multimodal + ) + def forward( self, input_ids: torch.Tensor | None, diff --git a/vllm/model_executor/models/qwen3_omni_moe_thinker.py b/vllm/model_executor/models/qwen3_omni_moe_thinker.py index 9500ce2e2bf..93a17f0c8c2 100755 --- a/vllm/model_executor/models/qwen3_omni_moe_thinker.py +++ b/vllm/model_executor/models/qwen3_omni_moe_thinker.py @@ -92,6 +92,8 @@ from .qwen2_5_omni_thinker import ( Qwen2_5OmniConditionalGenerationMixin, Qwen2_5OmniThinkerDummyInputsBuilder, Qwen2_5OmniThinkerMultiModalProcessor, + check_interleaved_audio_video, + merge_interleaved_embeddings, ) from .qwen2_5_vl import ( Qwen2_5_VisionAttention, @@ -1780,6 +1782,19 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( if multimodal_embeddings is None or len(multimodal_embeddings) == 0: return inputs_embeds + # Detect interleaved audio-in-video early, since it affects + # both the deepstack path and the final embedding merge. + video_token_id = self.config.video_token_id + audio_token_id = self.config.audio_token_id + is_video = is_multimodal & (input_ids == video_token_id) + is_audio = is_multimodal & (input_ids == audio_token_id) + num_video = is_video.sum().item() + num_audio = is_audio.sum().item() + + is_interleaved = check_interleaved_audio_video( + is_video, is_audio, num_video, num_audio + ) + deepstack_input_embeds = None # split the feat dim to obtain multi-scale visual feature has_vision_embeddings = [ @@ -1791,14 +1806,18 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( ): multiscale_len = len(self.visual.deepstack_visual_indexes) multimodal_embeddings_multiscale = [] - is_vision = torch.zeros_like(is_multimodal) - mm_positions = torch.nonzero(is_multimodal, as_tuple=True)[0] - mm_position_idx = 0 + + if is_interleaved: + # Use input_ids-based mask for correct vision positions + # when audio and video tokens are interleaved. + is_vision = is_video.clone() + else: + is_vision = torch.zeros_like(is_multimodal) + mm_positions = torch.nonzero(is_multimodal, as_tuple=True)[0] + mm_position_idx = 0 + for index, embeddings in enumerate(multimodal_embeddings): num_tokens = embeddings.shape[0] - current_positions = mm_positions[ - mm_position_idx : mm_position_idx + num_tokens - ] # Vision embeddings if embeddings.shape[-1] != self.config.text_config.hidden_size: @@ -1809,13 +1828,22 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( ) multimodal_embeddings[index] = embeddings_main multimodal_embeddings_multiscale.append(embeddings_multiscale) - is_vision[current_positions] = True + if not is_interleaved: + current_positions = mm_positions[ + mm_position_idx : mm_position_idx + num_tokens + ] + is_vision[current_positions] = True # Audio embeddings else: - is_vision[current_positions] = False + if not is_interleaved: + current_positions = mm_positions[ + mm_position_idx : mm_position_idx + num_tokens + ] + is_vision[current_positions] = False - mm_position_idx += num_tokens + if not is_interleaved: + mm_position_idx += num_tokens deepstack_input_embeds = inputs_embeds.new_zeros( inputs_embeds.size(0), multiscale_len * inputs_embeds.size(1) @@ -1834,6 +1862,18 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( ) self._set_deepstack_input_embeds(deepstack_input_embeds) + if is_interleaved: + return merge_interleaved_embeddings( + inputs_embeds, + multimodal_embeddings, + is_video, + is_audio, + is_multimodal, + num_video, + num_audio, + ) + + # Default: standard merge (no interleaving) inputs_embeds = _merge_multimodal_embeddings( inputs_embeds=inputs_embeds, multimodal_embeddings=multimodal_embeddings, From e57ef99b409ba651695a515f6022c9badf25b2a2 Mon Sep 17 00:00:00 2001 From: Cyrus Leung Date: Wed, 4 Feb 2026 20:23:01 +0800 Subject: [PATCH 058/810] [Model] Apply #32631 for recent models (#33785) Signed-off-by: DarkLight1337 --- vllm/model_executor/models/eagle2_5_vl.py | 33 +++++++++++----------- vllm/model_executor/models/funaudiochat.py | 3 -- vllm/model_executor/models/openpangu_vl.py | 33 +++++++++++----------- vllm/model_executor/models/qwen3_asr.py | 27 +++++++++--------- 4 files changed, 46 insertions(+), 50 deletions(-) diff --git a/vllm/model_executor/models/eagle2_5_vl.py b/vllm/model_executor/models/eagle2_5_vl.py index 3ce9b9c4d23..19d21de5b3e 100644 --- a/vllm/model_executor/models/eagle2_5_vl.py +++ b/vllm/model_executor/models/eagle2_5_vl.py @@ -222,22 +222,24 @@ class Eagle2_5_VLForConditionalGeneration( self.select_layer = getattr(config, "select_layer", -1) - # Vision encoder (SigLIP) - self.vision_model = self._init_vision_model( - config, - quant_config=quant_config, - prefix=maybe_prefix(prefix, "vision_model"), - ) + with self._mark_tower_model(vllm_config, "image"): + # Vision encoder (SigLIP) + self.vision_model = self._init_vision_model( + config, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "vision_model"), + ) - # Language model (Qwen2) - self.language_model = init_vllm_registered_model( - vllm_config=vllm_config, - hf_config=config.text_config, - prefix=maybe_prefix(prefix, "language_model"), - ) + # MLP projection + self.mlp1 = self._init_mlp1(config) - # MLP projection - self.mlp1 = self._init_mlp1(config) + with self._mark_language_model(vllm_config): + # Language model (Qwen2) + self.language_model = init_vllm_registered_model( + vllm_config=vllm_config, + hf_config=config.text_config, + prefix=maybe_prefix(prefix, "language_model"), + ) self.img_context_token_id = None @@ -399,9 +401,6 @@ class Eagle2_5_VLForConditionalGeneration( ] return image_embeds.split(image_feature_sizes) - def get_language_model(self) -> torch.nn.Module: - return self.language_model - def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: """Embed multimodal inputs.""" image_input = self._parse_and_validate_image_input(**kwargs) diff --git a/vllm/model_executor/models/funaudiochat.py b/vllm/model_executor/models/funaudiochat.py index 16afec3cf1e..b7b8659a4c9 100644 --- a/vllm/model_executor/models/funaudiochat.py +++ b/vllm/model_executor/models/funaudiochat.py @@ -820,9 +820,6 @@ class FunAudioChatForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor self.language_model.make_empty_intermediate_tensors ) - def get_language_model(self) -> torch.nn.Module: - return self.language_model - def _get_continuous_audio_features( self, input_features: torch.Tensor, diff --git a/vllm/model_executor/models/openpangu_vl.py b/vllm/model_executor/models/openpangu_vl.py index d7df2cbb4cf..e9288e6ddb1 100644 --- a/vllm/model_executor/models/openpangu_vl.py +++ b/vllm/model_executor/models/openpangu_vl.py @@ -843,20 +843,24 @@ class OpenPanguVLForConditionalGeneration( self.config = config self.vllm_config = vllm_config quant_config = vllm_config.quant_config - self.visual = OpenPanguVisionTransformer( - vision_config=config.vision_config, - out_hidden_size=config.vision_config.out_hidden_size, - hidden_size=config.hidden_size, - norm_eps=getattr(config.vision_config, "rms_norm_eps", 1e-6), - quant_config=self._maybe_ignore_quant_config(quant_config), - prefix=maybe_prefix(prefix, "visual"), - ) - self.language_model = init_vllm_registered_model( - vllm_config=vllm_config, - prefix=maybe_prefix("openpangu", "language_model"), - architectures=["PanguEmbeddedForCausalLM"], - ) + with self._mark_tower_model(vllm_config, {"image", "video"}): + self.visual = OpenPanguVisionTransformer( + vision_config=config.vision_config, + out_hidden_size=config.vision_config.out_hidden_size, + hidden_size=config.hidden_size, + norm_eps=getattr(config.vision_config, "rms_norm_eps", 1e-6), + quant_config=self._maybe_ignore_quant_config(quant_config), + prefix=maybe_prefix(prefix, "visual"), + ) + + with self._mark_language_model(vllm_config): + self.language_model = init_vllm_registered_model( + vllm_config=vllm_config, + prefix=maybe_prefix("openpangu", "language_model"), + architectures=["PanguEmbeddedForCausalLM"], + ) + self.make_empty_intermediate_tensors = ( self.language_model.make_empty_intermediate_tensors ) @@ -1008,9 +1012,6 @@ class OpenPanguVLForConditionalGeneration( ) return mm_input_by_modality - def get_language_model(self) -> torch.nn.Module: - return self.language_model - def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings | None: mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs) if not mm_input_by_modality: diff --git a/vllm/model_executor/models/qwen3_asr.py b/vllm/model_executor/models/qwen3_asr.py index e63e03e23e4..9dac8d75b43 100644 --- a/vllm/model_executor/models/qwen3_asr.py +++ b/vllm/model_executor/models/qwen3_asr.py @@ -296,19 +296,21 @@ class Qwen3ASRForConditionalGeneration( multimodal_config = vllm_config.model_config.multimodal_config self.config = thinker_config self.multimodal_config = multimodal_config - - self.audio_tower = Qwen3OmniMoeAudioEncoder( - thinker_config.audio_config, - prefix=maybe_prefix(prefix, "audio_tower"), - ) self.quant_config = quant_config - self.language_model = Qwen3ForCausalLM( - vllm_config=vllm_config.with_hf_config( - thinker_config.text_config, architectures=["Qwen3ForCausalLM"] - ), - prefix=maybe_prefix(prefix, "language_model"), - ) + with self._mark_tower_model(vllm_config, "audio"): + self.audio_tower = Qwen3OmniMoeAudioEncoder( + thinker_config.audio_config, + prefix=maybe_prefix(prefix, "audio_tower"), + ) + + with self._mark_language_model(vllm_config): + self.language_model = Qwen3ForCausalLM( + vllm_config=vllm_config.with_hf_config( + thinker_config.text_config, architectures=["Qwen3ForCausalLM"] + ), + prefix=maybe_prefix(prefix, "language_model"), + ) self.make_empty_intermediate_tensors = ( self.language_model.make_empty_intermediate_tensors @@ -363,9 +365,6 @@ class Qwen3ASRForConditionalGeneration( ) return audio_features.split(audio_output_lengths.tolist()) - def get_language_model(self) -> torch.nn.Module: - return self.language_model - def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings | None: mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs) if not mm_input_by_modality: From f67ee8b859215df4b521c67b9f26e27f30c9739f Mon Sep 17 00:00:00 2001 From: Chauncey Date: Wed, 4 Feb 2026 20:30:36 +0800 Subject: [PATCH 059/810] [Perf] Optimize chat completion streaming performance (#33782) Signed-off-by: chaunceyjiang --- .../openai/chat_completion/serving.py | 31 ++++++++++--------- 1 file changed, 17 insertions(+), 14 deletions(-) diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index 21bc0f44245..48fb666484a 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -679,6 +679,7 @@ class OpenAIServingChat(OpenAIServing): # For reasoning parser and tool call all enabled added_content_delta_arr = [False] * num_choices reasoning_end_arr = [False] * num_choices + prompt_is_reasoning_end_arr: list[bool | None] = [None] * num_choices else: all_previous_token_ids = None @@ -824,6 +825,16 @@ class OpenAIServingChat(OpenAIServing): i = output.index tool_parser = tool_parsers[i] + if ( + self.reasoning_parser + and res.prompt_token_ids + and prompt_is_reasoning_end_arr[i] is None + ): + # only check once per choice, because prompt_token_ids + # are the same for all deltas in that choice + prompt_is_reasoning_end_arr[i] = ( + reasoning_parser.is_reasoning_end(res.prompt_token_ids) + ) if finish_reason_sent[i]: continue @@ -926,13 +937,11 @@ class OpenAIServingChat(OpenAIServing): # i.e {"enable_thinking": False}, # set reasoning status to end. # Only keep 'content', remove 'reasoning'. - if reasoning_parser.is_reasoning_end( - as_list(output.token_ids) - ) or ( - res.prompt_token_ids - and reasoning_parser.is_reasoning_end( - res.prompt_token_ids + if ( + reasoning_parser.is_reasoning_end( + as_list(output.token_ids) ) + or prompt_is_reasoning_end_arr[i] ): reasoning_end_arr[i] = True if delta_message and delta_message.content: @@ -991,8 +1000,7 @@ class OpenAIServingChat(OpenAIServing): if ( self.reasoning_parser is not None and not reasoning_end_arr[i] - and res.prompt_token_ids - and reasoning_parser.is_reasoning_end(res.prompt_token_ids) + and prompt_is_reasoning_end_arr[i] ): reasoning_end_arr[i] = True @@ -1049,12 +1057,7 @@ class OpenAIServingChat(OpenAIServing): # When encountering think end id in prompt_token_ids # i.e {"enable_thinking": False}, # set reasoning status to end. - if ( - res.prompt_token_ids - and reasoning_parser.is_reasoning_end( - res.prompt_token_ids - ) - ): + if prompt_is_reasoning_end_arr[i]: reasoning_end_arr[i] = True current_token_ids = output_token_ids # Don't update current_text, keep it as is from delta From 32a02c7ca29180f70c1c8c73d0f57445231b17b5 Mon Sep 17 00:00:00 2001 From: Cyrus Leung Date: Wed, 4 Feb 2026 21:35:39 +0800 Subject: [PATCH 060/810] Apply #33621 to main (#33758) Signed-off-by: Zachary Aristei Co-authored-by: zaristei2 Co-authored-by: Zachary Aristei --- requirements/common.txt | 2 +- requirements/rocm-test.txt | 2 +- requirements/test.txt | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/requirements/common.txt b/requirements/common.txt index 365e30ed4e4..bc170d90b6d 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -11,7 +11,7 @@ transformers >= 4.56.0, < 5 tokenizers >= 0.21.1 # Required for fast incremental detokenization. protobuf # Required by LlamaTokenizer, gRPC. fastapi[standard] >= 0.115.0 # Required by FastAPI's form models in the OpenAI API server's audio transcriptions endpoint. -aiohttp +aiohttp >= 3.13.3 openai >= 1.99.1 # For Responses API with reasoning content pydantic >= 2.12.0 prometheus_client >= 0.18.0 diff --git a/requirements/rocm-test.txt b/requirements/rocm-test.txt index 835d8b183e4..26c28385270 100644 --- a/requirements/rocm-test.txt +++ b/requirements/rocm-test.txt @@ -14,7 +14,7 @@ pytest-shard==0.1.2 # Async/HTTP dependencies anyio==4.6.2.post1 # via httpx, starlette -aiohttp==3.13.0 +aiohttp==3.13.3 # via gpt-oss httpx==0.27.2 # HTTP testing diff --git a/requirements/test.txt b/requirements/test.txt index 97f580848ba..f72df06a0a7 100644 --- a/requirements/test.txt +++ b/requirements/test.txt @@ -12,7 +12,7 @@ affine==2.4.0 # via rasterio aiohappyeyeballs==2.6.1 # via aiohttp -aiohttp==3.13.0 +aiohttp==3.13.3 # via # aiohttp-cors # datasets From 1d367a738e9098ad4af1f6865747914ccd2c65ca Mon Sep 17 00:00:00 2001 From: Micah Williamson Date: Wed, 4 Feb 2026 07:36:29 -0600 Subject: [PATCH 061/810] [Bugfix][ROCm] Include float8_e4m3fnuz in NCCL Dtype Dispatching (#33713) Signed-off-by: Micah Williamson --- vllm/distributed/device_communicators/pynccl_wrapper.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/distributed/device_communicators/pynccl_wrapper.py b/vllm/distributed/device_communicators/pynccl_wrapper.py index 3b11595b4e4..78b3328f48d 100644 --- a/vllm/distributed/device_communicators/pynccl_wrapper.py +++ b/vllm/distributed/device_communicators/pynccl_wrapper.py @@ -93,7 +93,7 @@ class ncclDataTypeEnum: return cls.ncclFloat64 if dtype == torch.bfloat16: return cls.ncclBfloat16 - if dtype == torch.float8_e4m3fn: + if dtype == current_platform.fp8_dtype(): return cls.ncclFloat8e4m3 raise ValueError( f"Unsupported dtype {dtype}: should be one of " From 711edaf0d089a15df5fa2b99248c516e53929bd2 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Wed, 4 Feb 2026 09:34:32 -0500 Subject: [PATCH 062/810] [Perf] Optimize spec decoding + async scheduling, 1.5% Throughput improvement (#33612) Signed-off-by: yewentao256 Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: Nick Hill --- vllm/v1/core/sched/async_scheduler.py | 9 +++++++-- vllm/v1/core/sched/scheduler.py | 11 +++++++---- 2 files changed, 14 insertions(+), 6 deletions(-) diff --git a/vllm/v1/core/sched/async_scheduler.py b/vllm/v1/core/sched/async_scheduler.py index 3c66a23208e..23c610f3bde 100644 --- a/vllm/v1/core/sched/async_scheduler.py +++ b/vllm/v1/core/sched/async_scheduler.py @@ -10,6 +10,11 @@ logger = init_logger(__name__) class AsyncScheduler(Scheduler): + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + # reusable read-only placeholder list for speculative decoding. + self._spec_token_placeholders: list[int] = [-1] * self.num_spec_tokens + def _update_after_schedule(self, scheduler_output: SchedulerOutput) -> None: super()._update_after_schedule(scheduler_output) has_structured_output_requests = False @@ -31,9 +36,9 @@ class AsyncScheduler(Scheduler): # The request will generate a new token plus num_spec_tokens # in this scheduling step. request.num_output_placeholders += 1 + cur_num_spec_tokens - # Add placeholders for the new tokens in spec_token_ids. + # Add placeholders for the new draft/spec tokens. # We will update the actual spec token ids in the worker process. - request.spec_token_ids = [-1] * self.num_spec_tokens + request.spec_token_ids = self._spec_token_placeholders scheduler_output.has_structured_output_requests = has_structured_output_requests scheduler_output.pending_structured_output_tokens = ( diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 1544d847c8f..869b53601b1 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -487,9 +487,11 @@ class Scheduler(SchedulerInterface): - request.num_output_placeholders ) if num_scheduled_spec_tokens > 0: - # Trim spec_token_ids list to num_scheduled_spec_tokens. - del request.spec_token_ids[num_scheduled_spec_tokens:] - scheduled_spec_decode_tokens[request_id] = request.spec_token_ids + spec_token_ids = request.spec_token_ids + if len(spec_token_ids) > num_scheduled_spec_tokens: + spec_token_ids = spec_token_ids[:num_scheduled_spec_tokens] + scheduled_spec_decode_tokens[request.request_id] = spec_token_ids + # New spec tokens will be set in `update_draft_token_ids` before the # next step when applicable. request.spec_token_ids = [] @@ -887,7 +889,8 @@ class Scheduler(SchedulerInterface): self.encoder_cache_manager.free(request) request.status = RequestStatus.PREEMPTED request.num_computed_tokens = 0 - request.spec_token_ids.clear() + if request.spec_token_ids: + request.spec_token_ids = [] request.num_preemptions += 1 if self.log_stats: request.record_event(EngineCoreEventType.PREEMPTED, timestamp) From 80f921ba4bab2ea251d149305ea0f912c6fc218a Mon Sep 17 00:00:00 2001 From: Cyrus Leung Date: Wed, 4 Feb 2026 23:56:02 +0800 Subject: [PATCH 063/810] [Bugfix] Fix `normalize` still being passed to `PoolerConfig` (#33794) Signed-off-by: DarkLight1337 --- tests/models/language/pooling/test_embedding.py | 2 +- vllm/entrypoints/llm.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/models/language/pooling/test_embedding.py b/tests/models/language/pooling/test_embedding.py index 982dc73f677..e105195afe0 100644 --- a/tests/models/language/pooling/test_embedding.py +++ b/tests/models/language/pooling/test_embedding.py @@ -54,7 +54,7 @@ def test_models( vllm_extra_kwargs = {} if model == "ssmits/Qwen2-7B-Instruct-embed-base": vllm_extra_kwargs["pooler_config"] = PoolerConfig( - seq_pooling_type="MEAN", normalize=False + seq_pooling_type="MEAN", use_activation=False ) max_model_len: int | None = 512 diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index f3f774bef36..24545de19cb 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -174,8 +174,8 @@ class LLM: multi-modal processor obtained from `AutoProcessor.from_pretrained`. The available overrides depend on the model that is being run. For example, for Phi-3-Vision: `{"num_crops": 4}`. - pooler_config: Initialize non-default pooling config for the pooling - model. e.g. `PoolerConfig(seq_pooling_type="MEAN", normalize=False)`. + pooler_config: Initialize non-default pooling config for the pooling model, + e.g., `PoolerConfig(seq_pooling_type="MEAN", use_activation=False)`. compilation_config: Either an integer or a dictionary. If it is an integer, it is used as the mode of compilation optimization. If it is a dictionary, it can specify the full compilation configuration. From 87d9a261664705e0c9635014b4e2d49eddc8a056 Mon Sep 17 00:00:00 2001 From: jiangkuaixue123 Date: Thu, 5 Feb 2026 00:41:45 +0800 Subject: [PATCH 064/810] [Bugfix] Fix ubatch wrapper num_tokens calculate (#33694) Signed-off-by: jiangkuaixue123 --- vllm/v1/worker/gpu_ubatch_wrapper.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/vllm/v1/worker/gpu_ubatch_wrapper.py b/vllm/v1/worker/gpu_ubatch_wrapper.py index 97f8b92cec4..765427683a1 100644 --- a/vllm/v1/worker/gpu_ubatch_wrapper.py +++ b/vllm/v1/worker/gpu_ubatch_wrapper.py @@ -412,9 +412,7 @@ class UBatchWrapper: attn_metadata = forward_context.attn_metadata slot_mapping = forward_context.slot_mapping - num_tokens = ( - ubatch_slices[0].token_slice.stop - ubatch_slices[0].token_slice.start - ) * 2 + num_tokens = sum(ubatch_slice.num_tokens for ubatch_slice in ubatch_slices) input_ids = kwargs["input_ids"] positions = kwargs["positions"] intermediate_tensors = kwargs["intermediate_tensors"] From 0e922986222d9c25ce320fbdd0aff20279c2cb93 Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Wed, 4 Feb 2026 09:41:57 -0700 Subject: [PATCH 065/810] [Misc] Delay deprecation of CommonAttentionMetadata properties (#33801) Signed-off-by: Lucas Wilkinson --- vllm/v1/attention/backend.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 13082608c47..49eb91576ed 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -347,7 +347,7 @@ class CommonAttentionMetadata: """ Prefer using device seq_lens directly to avoid implicit H<>D sync. If a CPU copy is needed, use `seq_lens.cpu()` instead. - Will be removed in a future release (v0.15.0) + Will be removed in a future release, please migrate as soon as possible. """ ) def seq_lens_cpu(self) -> torch.Tensor: @@ -361,7 +361,7 @@ class CommonAttentionMetadata: Prefer using device seq_lens directly to avoid implicit H<>D sync which breaks full async scheduling. If a CPU copy is needed, it can be derived from query_start_loc_cpu and seq_lens. - Will be removed in a future release (v0.15.0) + Will be removed in a future release, please migrate as soon as possible. """ ) def num_computed_tokens_cpu(self) -> torch.Tensor: From 192ad4648b2066ebdf1fa04ad84f24bdf0cd6533 Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Thu, 5 Feb 2026 01:54:45 +0800 Subject: [PATCH 066/810] [Bugfix] Fix interns1-pro initialization and PP (#33793) Signed-off-by: Isotr0py --- .../multimodal/processing/test_common.py | 4 ++ .../processing/test_tensor_schema.py | 3 ++ tests/models/registry.py | 2 - vllm/model_executor/models/interns1_pro.py | 38 +++++++++++++------ vllm/model_executor/models/qwen3_vl.py | 9 +++-- vllm/model_executor/models/qwen3_vl_moe.py | 9 +++-- 6 files changed, 43 insertions(+), 22 deletions(-) diff --git a/tests/models/multimodal/processing/test_common.py b/tests/models/multimodal/processing/test_common.py index b228898ffb4..ae2ec1bc01e 100644 --- a/tests/models/multimodal/processing/test_common.py +++ b/tests/models/multimodal/processing/test_common.py @@ -124,6 +124,7 @@ MM_DATA_PATCHES = { "glm4v_moe": glm4_1v_patch_mm_data, "glm_ocr": glm4_1v_patch_mm_data, "glmasr": glmasr_patch_mm_data, + "interns1_pro": qwen3_vl_patch_mm_data, "molmo2": qwen3_vl_patch_mm_data, "qwen3_vl": qwen3_vl_patch_mm_data, "qwen3_vl_moe": qwen3_vl_patch_mm_data, @@ -439,6 +440,9 @@ def test_processing_correctness( "Qwen-VL tokenizer requires downloading a font file from " "servers that often refuse connections in CI" ) + if model_id == "internlm/Intern-S1-Pro": + # FIXME(Isotr0py): Fix later. + pytest.skip("Tokenization issue. Fix later") _test_processing_correctness( model_id, diff --git a/tests/models/multimodal/processing/test_tensor_schema.py b/tests/models/multimodal/processing/test_tensor_schema.py index 8f79936478d..aabd883a49b 100644 --- a/tests/models/multimodal/processing/test_tensor_schema.py +++ b/tests/models/multimodal/processing/test_tensor_schema.py @@ -160,6 +160,9 @@ def test_model_tensor_schema(model_id: str): pytest.skip( "Kimi-K2.5's offline inference has issues about vision chunks. Fix later." ) + if model_id == "internlm/Intern-S1-Pro": + # FIXME(Isotr0py): Fix later. + pytest.skip("Intern-S1-Pro has issue to pass the test.") model_info = HF_EXAMPLE_MODELS.find_hf_info(model_id) model_info.check_available_online(on_fail="skip") diff --git a/tests/models/registry.py b/tests/models/registry.py index c38637c1c67..cbd07cbc1c7 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -758,8 +758,6 @@ _MULTIMODAL_EXAMPLE_MODELS = { "InternS1ProForConditionalGeneration": _HfExamplesInfo( "internlm/Intern-S1-Pro", trust_remote_code=True, - min_transformers_version="5.0.0", - is_available_online=False, ), "InternVLChatModel": _HfExamplesInfo( "OpenGVLab/InternVL2-1B", diff --git a/vllm/model_executor/models/interns1_pro.py b/vllm/model_executor/models/interns1_pro.py index 60c92cddab3..c5cd1339938 100644 --- a/vllm/model_executor/models/interns1_pro.py +++ b/vllm/model_executor/models/interns1_pro.py @@ -32,7 +32,6 @@ import torch from torch import nn from transformers import AutoProcessor, PretrainedConfig -from vllm.attention.layer import Attention from vllm.config import CacheConfig, VllmConfig from vllm.distributed import ( get_ep_group, @@ -41,8 +40,8 @@ from vllm.distributed import ( ) 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 FusedMoE -from vllm.model_executor.layers.fused_moe.config import RoutingMethodType from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, @@ -188,7 +187,6 @@ class InternS1ProMoeSparseMoeBlock(nn.Module): enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, is_sequence_parallel=self.is_sequence_parallel, - routing_method_type=RoutingMethodType.Renormalize, custom_routing_function=self._custom_routing_function, ) @@ -479,7 +477,7 @@ class InternS1ProMoeLLMModel(Qwen3MoeLLMModel): class InternS1ProMoeLLMForCausalLM(Qwen3MoeForCausalLM): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() + super(Qwen3MoeForCausalLM, self).__init__() self.config = vllm_config.model_config.hf_config.text_config self.quant_config = vllm_config.quant_config self.model = InternS1ProMoeLLMModel( @@ -567,15 +565,10 @@ class InternS1ProForConditionalGeneration( "lm_head.": "language_model.lm_head.", "model.language_model.": "language_model.model.", }, - orig_to_new_suffix={ - # Handle FOPE rotary embeddings - ".rotary_emb.sin_coef": ".layers.0.self_attn.rotary_emb.sin_coef", - ".rotary_emb.cos_coef": ".layers.0.self_attn.rotary_emb.cos_coef", - }, ) def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() + super(Qwen3VLForConditionalGeneration, self).__init__() config: PretrainedConfig = vllm_config.model_config.hf_config multimodal_config = vllm_config.model_config.multimodal_config @@ -595,7 +588,6 @@ class InternS1ProForConditionalGeneration( self.visual = Qwen3_VisionTransformer( config.vision_config, norm_eps=getattr(config, "rms_norm_eps", 1e-6), - multimodal_config=multimodal_config, prefix=maybe_prefix(prefix, "visual"), ) @@ -624,10 +616,32 @@ class InternS1ProForConditionalGeneration( # Set MoE hyperparameters self.set_moe_parameters() + def get_frope_params_map(self) -> str: + mapper = {} + for name, params in self.language_model.model.named_parameters(): + if "rotary_emb.sin_coef" in name: + mapper["language_model.model.rotary_emb.sin_coef"] = ( + f"language_model.model.{name}" + ) + if "rotary_emb.cos_coef" in name: + mapper["language_model.model.rotary_emb.cos_coef"] = ( + f"language_model.model.{name}" + ) + return mapper + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): """load weights""" skip_prefixes = ["model.time_series."] if self.visual is None: skip_prefixes.append("visual.") + # 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_suffix=self.get_frope_params_map(), + ) loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes) - return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + return loader.load_weights(weights, mapper=weights_mapper) diff --git a/vllm/model_executor/models/qwen3_vl.py b/vllm/model_executor/models/qwen3_vl.py index 102d846090c..34ff881aad7 100644 --- a/vllm/model_executor/models/qwen3_vl.py +++ b/vllm/model_executor/models/qwen3_vl.py @@ -1114,10 +1114,11 @@ class Qwen3VLMultiModalProcessor(BaseMultiModalProcessor[Qwen3VLProcessingInfo]) class Qwen3LLMModel(Qwen3Model): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__(vllm_config=vllm_config, prefix=prefix) - if not get_pp_group().is_first_rank: - assert self.start_layer >= len( - vllm_config.model_config.hf_config.vision_config.deepstack_visual_indexes - ), ( + vision_config = vllm_config.model_config.hf_config.vision_config + if not get_pp_group().is_first_rank and hasattr( + vision_config, "deepstack_visual_indexes" + ): + assert self.start_layer >= len(vision_config.deepstack_visual_indexes), ( "start_layer should be greater than or equal to " "len(deepstack_visual_indexes)" ) diff --git a/vllm/model_executor/models/qwen3_vl_moe.py b/vllm/model_executor/models/qwen3_vl_moe.py index af8536e3f19..8ac2dc945c4 100644 --- a/vllm/model_executor/models/qwen3_vl_moe.py +++ b/vllm/model_executor/models/qwen3_vl_moe.py @@ -95,10 +95,11 @@ class Qwen3MoeLLMModel(Qwen3MoeModel): prefix=prefix, decoder_layer_type=decoder_layer_type, ) - if not get_pp_group().is_first_rank: - assert self.start_layer >= len( - vllm_config.model_config.hf_config.vision_config.deepstack_visual_indexes - ), ( + vision_config = vllm_config.model_config.hf_config.vision_config + if not get_pp_group().is_first_rank and hasattr( + vision_config, "deepstack_visual_indexes" + ): + assert self.start_layer >= len(vision_config.deepstack_visual_indexes), ( "start_layer should be greater than or equal to " "len(deepstack_visual_indexes)" ) From 2f6d17cb2f4a49e29aae5c3c1ff64623d66d0257 Mon Sep 17 00:00:00 2001 From: kourosh hakhamaneshi <31483498+kouroshHakha@users.noreply.github.com> Date: Wed, 4 Feb 2026 10:09:14 -0800 Subject: [PATCH 067/810] [rocm][ray] Fix: Unify Ray device visibility handling across CUDA and ROCm (#33308) Signed-off-by: Kourosh Hakhamaneshi --- docker/Dockerfile.rocm | 2 -- tests/config/test_config_generation.py | 5 +++++ vllm/platforms/cuda.py | 3 +++ vllm/platforms/interface.py | 5 +++++ vllm/platforms/rocm.py | 5 +++++ vllm/v1/executor/ray_executor.py | 24 ++++++++++++++++++++++++ vllm/v1/worker/gpu_worker.py | 1 + vllm/v1/worker/worker_base.py | 6 ------ 8 files changed, 43 insertions(+), 8 deletions(-) diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 0178d23b73f..3409f04a1bf 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -15,8 +15,6 @@ FROM ${BASE_IMAGE} AS base ARG ARG_PYTORCH_ROCM_ARCH ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}} -ENV RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES=1 -ENV RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES=1 # Install some basic utilities RUN apt-get update -q -y && apt-get install -q -y \ diff --git a/tests/config/test_config_generation.py b/tests/config/test_config_generation.py index 61c3df0a234..23ceb920cae 100644 --- a/tests/config/test_config_generation.py +++ b/tests/config/test_config_generation.py @@ -60,6 +60,11 @@ def test_ray_runtime_env(monkeypatch: pytest.MonkeyPatch): runtime_env = { "env_vars": { "TEST_ENV_VAR": "test_value", + # In future ray versions, this will be default, so when setting a + # task or actor with num_gpus=None/0, the visible devices env var + # won't be overridden resulting in no GPUs being visible on a gpu + # machine. + "RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO": "0", }, } diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index fbe791f8a8e..0c0bd7db3d9 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -102,6 +102,9 @@ class CudaPlatformBase(Platform): ray_device_key: str = "GPU" dist_backend: str = "nccl" device_control_env_var: str = "CUDA_VISIBLE_DEVICES" + ray_noset_device_env_vars: list[str] = [ + "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES", + ] @property def supported_dtypes(self) -> list[torch.dtype]: diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index f0e7ee0da51..c3b189e013e 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -116,6 +116,11 @@ class Platform: # https://github.com/ray-project/ray/tree/master/python/ray/_private/accelerators # noqa device_control_env_var: str = "VLLM_DEVICE_CONTROL_ENV_VAR_PLACEHOLDER" + # environment variables that need to be set to 1 to prevent ray from + # setting the visible devices e.g. + # RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES + ray_noset_device_env_vars: list[str] = [] + # The torch.compile backend for compiling simple and # standalone functions. The default value is "inductor" to keep # the same behavior as PyTorch. diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 6f4c235bb89..2a9bd53e471 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -194,6 +194,11 @@ class RocmPlatform(Platform): dist_backend: str = "nccl" # rocm shares the same device control env var as CUDA device_control_env_var: str = "CUDA_VISIBLE_DEVICES" + ray_noset_device_env_vars: list[str] = [ + "RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES", + "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES", + "RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES", + ] supported_quantization: list[str] = [ "awq", diff --git a/vllm/v1/executor/ray_executor.py b/vllm/v1/executor/ray_executor.py index c8c6185b6c2..a1f69c47818 100644 --- a/vllm/v1/executor/ray_executor.py +++ b/vllm/v1/executor/ray_executor.py @@ -69,6 +69,8 @@ class RayDistributedExecutor(Executor): "VLLM_HOST_PORT", "LOCAL_RANK", "CUDA_VISIBLE_DEVICES", + "HIP_VISIBLE_DEVICES", + "ROCR_VISIBLE_DEVICES", } # These non-vLLM env vars are copied from the driver to workers @@ -146,6 +148,14 @@ class RayDistributedExecutor(Executor): return ray_remote_kwargs + def _update_noset_device_env_vars(self, ray_remote_kwargs): + runtime_env = ray_remote_kwargs.setdefault("runtime_env", {}) + env_vars = runtime_env.setdefault("env_vars", {}) + env_vars.update( + {env_var: "1" for env_var in current_platform.ray_noset_device_env_vars} + ) + return ray_remote_kwargs + # child class could overwrite this to return actual env vars. def _get_env_vars_to_be_updated(self): return self._env_vars_for_all_workers @@ -169,6 +179,11 @@ class RayDistributedExecutor(Executor): ray_remote_kwargs ) + # The way ray actors are setup in vllm is that the visible devices are + # not set by actors, they are left unset by ray. Internally we index + # the right gpu with local_rank. This is similar to how mp mode works. + self._update_noset_device_env_vars(ray_remote_kwargs) + # Create the workers. bundle_indices: list[int] if envs.VLLM_RAY_BUNDLE_INDICES: @@ -303,6 +318,15 @@ class RayDistributedExecutor(Executor): ) # Set environment variables for the driver and workers. + # We set CUDA_VISIBLE_DEVICES to ALL GPUs on the node for each worker. + # This is needed because: + # 1. Ray's compiled DAG needs to find the allocated GPU in + # CUDA_VISIBLE_DEVICES. + # 2. vLLM's communication layer (NCCL, CustomAllreduce) needs to see + # all GPUs for P2P checks and communication setup. Though if it was + # just this reason, we could have also just kept the visible devices + # unset. + # Each worker will use local_rank to index into the visible devices. all_args_to_update_environment_variables = [ { current_platform.device_control_env_var: ",".join( diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index f6e59526ed4..b451db3826f 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -209,6 +209,7 @@ class Worker(WorkerBase): f"be less than or equal to the number of visible devices " f"({visible_device_count})." ) + self.device = torch.device(f"cuda:{self.local_rank}") current_platform.set_device(self.device) diff --git a/vllm/v1/worker/worker_base.py b/vllm/v1/worker/worker_base.py index d34eb5253ff..eed371e988b 100644 --- a/vllm/v1/worker/worker_base.py +++ b/vllm/v1/worker/worker_base.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import os from collections.abc import Callable from typing import TYPE_CHECKING, Any, TypeVar @@ -221,11 +220,6 @@ class WorkerWrapperBase: envs_list: list[dict[str, str]], ) -> None: envs = envs_list[self.rpc_rank] - key = "CUDA_VISIBLE_DEVICES" - if key in envs and key in os.environ: - # overwriting CUDA_VISIBLE_DEVICES is desired behavior - # suppress the warning in `update_environment_variables` - del os.environ[key] update_environment_variables(envs) def init_worker(self, all_kwargs: list[dict[str, Any]]) -> None: From 6e98f6d8b64984d5e3a8a9b323321f1095bac8b3 Mon Sep 17 00:00:00 2001 From: Taeksang Kim Date: Thu, 5 Feb 2026 05:11:39 +0900 Subject: [PATCH 068/810] Implement zero-copy GQA for multimodal and CPU (#33732) Signed-off-by: Taeksang Kim --- .../layers/attention/mm_encoder_attention.py | 16 ++++------------ vllm/model_executor/models/molmo2.py | 13 +------------ vllm/v1/attention/backends/cpu_attn.py | 5 +---- vllm/v1/attention/ops/vit_attn_wrappers.py | 16 ++++++++++++---- 4 files changed, 18 insertions(+), 32 deletions(-) diff --git a/vllm/model_executor/layers/attention/mm_encoder_attention.py b/vllm/model_executor/layers/attention/mm_encoder_attention.py index 35c10ec0bd9..f26d89f40c3 100644 --- a/vllm/model_executor/layers/attention/mm_encoder_attention.py +++ b/vllm/model_executor/layers/attention/mm_encoder_attention.py @@ -80,7 +80,7 @@ class MMEncoderAttention(CustomOp): def enabled(cls) -> bool: return True - def maybe_reshape_qkv_to_4d( + def view_qkv_to_4d( self, query: torch.Tensor, key: torch.Tensor, @@ -97,11 +97,6 @@ class MMEncoderAttention(CustomOp): key = key.view(bsz, kv_len, self.num_kv_heads, self.head_size) value = value.view(bsz, kv_len, self.num_kv_heads, self.head_size) - if (num_repeat := self.num_queries_per_kv) > 1: - # Handle MQA and GQA - key = torch.repeat_interleave(key, num_repeat, dim=2) - value = torch.repeat_interleave(value, num_repeat, dim=2) - return query, key, value def _forward_sdpa( @@ -119,9 +114,7 @@ class MMEncoderAttention(CustomOp): kv_len = key.size(1) is_reshaped = query.dim() != 4 - query, key, value = self.maybe_reshape_qkv_to_4d( - query, key, value, bsz, q_len, kv_len - ) + query, key, value = self.view_qkv_to_4d(query, key, value, bsz, q_len, kv_len) output = vit_torch_sdpa_wrapper( q=query, @@ -129,6 +122,7 @@ class MMEncoderAttention(CustomOp): v=value, scale=self.scale, cu_seqlens=cu_seqlens, + enable_gqa=self.num_heads > self.num_kv_heads, ) if is_reshaped: output = output.reshape(bsz, q_len, -1) @@ -154,9 +148,7 @@ class MMEncoderAttention(CustomOp): kv_len = key.size(1) is_reshaped = query.dim() != 4 - query, key, value = self.maybe_reshape_qkv_to_4d( - query, key, value, bsz, q_len, kv_len - ) + query, key, value = self.view_qkv_to_4d(query, key, value, bsz, q_len, kv_len) output = vit_flash_attn_wrapper( q=query, diff --git a/vllm/model_executor/models/molmo2.py b/vllm/model_executor/models/molmo2.py index 9d996a93b05..30f639c8bee 100644 --- a/vllm/model_executor/models/molmo2.py +++ b/vllm/model_executor/models/molmo2.py @@ -628,18 +628,6 @@ class ImagePoolingAttention(nn.Module): key = key.view(bsz, kv_len, self.num_kv_heads, self.head_dim) value = value.view(bsz, kv_len, self.num_kv_heads, self.head_dim) - if self.num_heads != self.num_kv_heads: - key = torch.repeat_interleave( - key, - self.num_heads // self.num_kv_heads, - dim=2, - ) - value = torch.repeat_interleave( - value, - self.num_heads // self.num_kv_heads, - dim=2, - ) - query, key, value = (x.transpose(1, 2) for x in (query, key, value)) out = F.scaled_dot_product_attention( @@ -648,6 +636,7 @@ class ImagePoolingAttention(nn.Module): value, attn_mask=attn_mask, is_causal=False, + enable_gqa=self.num_heads > self.num_kv_heads, ).transpose(1, 2) return out.reshape(bsz, q_len, -1) diff --git a/vllm/v1/attention/backends/cpu_attn.py b/vllm/v1/attention/backends/cpu_attn.py index 3eb9b478230..a2f2c6aeb92 100644 --- a/vllm/v1/attention/backends/cpu_attn.py +++ b/vllm/v1/attention/backends/cpu_attn.py @@ -398,10 +398,6 @@ class CPUAttentionBackendImpl(AttentionImpl): key = key.movedim(0, key.dim() - 2) value = value.movedim(0, value.dim() - 2) - if self.num_kv_heads != self.num_heads: - key = key.repeat_interleave(self.num_queries_per_kv, dim=-3) - value = value.repeat_interleave(self.num_queries_per_kv, dim=-3) - causal_attn = attn_type == AttentionType.DECODER sdpa_start_loc = attn_metadata.sdpa_start_loc.numpy() # type: ignore @@ -418,6 +414,7 @@ class CPUAttentionBackendImpl(AttentionImpl): dropout_p=0.0, is_causal=causal_attn and mask is None, scale=self.scale, + enable_gqa=self.num_heads > self.num_kv_heads, ) .squeeze(0) .movedim(query.dim() - 2, 0) diff --git a/vllm/v1/attention/ops/vit_attn_wrappers.py b/vllm/v1/attention/ops/vit_attn_wrappers.py index f077a61c984..32fcb35111d 100644 --- a/vllm/v1/attention/ops/vit_attn_wrappers.py +++ b/vllm/v1/attention/ops/vit_attn_wrappers.py @@ -115,13 +115,16 @@ def apply_sdpa( k: torch.Tensor, v: torch.Tensor, scale: float | None = None, + enable_gqa: bool = False, ) -> torch.Tensor: """ Input shape: (batch_size x seq_len x num_heads x head_size) """ q, k, v = (einops.rearrange(x, "b s h d -> b h s d") for x in [q, k, v]) - output = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0, scale=scale) + output = F.scaled_dot_product_attention( + q, k, v, dropout_p=0.0, scale=scale, enable_gqa=enable_gqa + ) output = einops.rearrange(output, "b h s d -> b s h d ") return output @@ -134,6 +137,7 @@ def torch_sdpa_wrapper( v: torch.Tensor, scale: float | None = None, cu_seqlens: torch.Tensor | None = None, + enable_gqa: bool = False, ) -> torch.Tensor: # Never remove the contiguous logic for ROCm # Without it, hallucinations occur with the backend @@ -143,7 +147,7 @@ def torch_sdpa_wrapper( v = v.contiguous() if cu_seqlens is None: - return apply_sdpa(q, k, v, scale=scale) + return apply_sdpa(q, k, v, scale=scale, enable_gqa=enable_gqa) outputs = [] @@ -152,7 +156,7 @@ def torch_sdpa_wrapper( k_chunks = torch.split(k, lens, dim=1) v_chunks = torch.split(v, lens, dim=1) for q_i, k_i, v_i in zip(q_chunks, k_chunks, v_chunks): - output_i = apply_sdpa(q_i, k_i, v_i, scale=scale) + output_i = apply_sdpa(q_i, k_i, v_i, scale=scale, enable_gqa=enable_gqa) outputs.append(output_i) context_layer = torch.cat(outputs, dim=1) return context_layer @@ -164,6 +168,7 @@ def torch_sdpa_wrapper_fake( v: torch.Tensor, scale: float | None, cu_seqlens: torch.Tensor | None, + enable_gqa: bool = False, ) -> torch.Tensor: return torch.empty_like(q) @@ -181,5 +186,8 @@ def vit_torch_sdpa_wrapper( v: torch.Tensor, scale: float | None = None, cu_seqlens: torch.Tensor | None = None, + enable_gqa: bool = False, ) -> torch.Tensor: - return torch.ops.vllm.torch_sdpa_wrapper(q, k, v, scale, cu_seqlens) + return torch.ops.vllm.torch_sdpa_wrapper( + q, k, v, scale, cu_seqlens, enable_gqa=enable_gqa + ) From 4292c90a2a188121ccbfd132def62031283d9d8a Mon Sep 17 00:00:00 2001 From: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com> Date: Wed, 4 Feb 2026 21:17:41 +0100 Subject: [PATCH 069/810] [Bugfix] Support `RotaryEmbedding` CustomOp for gpt-oss (#33800) Signed-off-by: simondanielsson --- .../compile/test_rotary_embedding_compile.py | 68 +++++++++++++++++++ .../layers/rotary_embedding/base.py | 35 ++++++---- .../rotary_embedding/deepseek_scaling_rope.py | 4 +- .../layers/rotary_embedding/mrope.py | 8 +-- 4 files changed, 97 insertions(+), 18 deletions(-) create mode 100644 tests/compile/test_rotary_embedding_compile.py diff --git a/tests/compile/test_rotary_embedding_compile.py b/tests/compile/test_rotary_embedding_compile.py new file mode 100644 index 00000000000..76f5382534e --- /dev/null +++ b/tests/compile/test_rotary_embedding_compile.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +import vllm.envs as envs +from vllm.compilation.decorators import support_torch_compile +from vllm.config import ( + CompilationConfig, + ModelConfig, + VllmConfig, + set_current_vllm_config, +) +from vllm.config.compilation import CompilationMode, CUDAGraphMode +from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.platforms import current_platform + + +@support_torch_compile +class RotaryEmbeddingCompileModule(torch.nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + self.rotary_emb = get_rope( + head_size=32, + max_position=128, + dtype=torch.float32, + rope_parameters={"rope_type": "default", "rope_theta": 10000}, + is_neox_style=True, + ) + + def forward( + self, positions: torch.Tensor, query: torch.Tensor, key: torch.Tensor + ) -> torch.Tensor: + q_rot, k_rot = self.rotary_emb(positions, query, key) + return q_rot + k_rot + + +@pytest.mark.skipif(current_platform.is_cpu(), reason="Requires GPU for torch.compile") +def test_rotary_embedding_torch_compile_with_custom_op(monkeypatch): + # Ensure env toggles take effect for this test only. + # The bytecode hook is required to detect buffer mutation in compiled code, + # and AOT compile bypasses that hook entirely. + envs.disable_envs_cache() + monkeypatch.setenv("VLLM_USE_BYTECODE_HOOK", "1") + monkeypatch.setenv("VLLM_USE_AOT_COMPILE", "0") + + device = "cuda" + positions = torch.arange(16, device=device) + query = torch.randn(16, 32, device=device, dtype=torch.bfloat16) + key = torch.randn(16, 32, device=device, dtype=torch.bfloat16) + + vllm_config = VllmConfig( + model_config=ModelConfig(dtype=torch.bfloat16), + compilation_config=CompilationConfig( + mode=CompilationMode.VLLM_COMPILE, + backend="inductor", + custom_ops=["+rotary_embedding"], + cudagraph_mode=CUDAGraphMode.NONE, + cudagraph_num_of_warmups=0, + ), + ) + + with set_current_vllm_config(vllm_config): + model = RotaryEmbeddingCompileModule(vllm_config=vllm_config) + model(positions, query, key) + assert model._compiled_bytecode is not None + assert "update" not in model._compiled_bytecode.co_names diff --git a/vllm/model_executor/layers/rotary_embedding/base.py b/vllm/model_executor/layers/rotary_embedding/base.py index 2147e00d2db..1e306339249 100644 --- a/vllm/model_executor/layers/rotary_embedding/base.py +++ b/vllm/model_executor/layers/rotary_embedding/base.py @@ -86,14 +86,23 @@ class RotaryEmbeddingBase(CustomOp): cache = torch.cat((cos, sin), dim=-1) return cache - def _match_cos_sin_cache_dtype(self, query: torch.Tensor) -> None: + def _match_cos_sin_cache_dtype(self, query: torch.Tensor) -> torch.Tensor: # __setattr__ in nn.Module (called by `self.cos_sin_cache = ...`) # is expensive, so avoid calling it if possible + cos_sin_cache = self.cos_sin_cache if ( - self.cos_sin_cache.device != query.device - or self.cos_sin_cache.dtype != query.dtype + cos_sin_cache.device == query.device + and self.cos_sin_cache.dtype == query.dtype ): - self.cos_sin_cache = self.cos_sin_cache.to(query.device, dtype=query.dtype) + return cos_sin_cache + + cos_sin_cache = cos_sin_cache.to(query.device, dtype=query.dtype) + # Avoid mutating buffers during torch.compile (cudagraph) tracing. + if torch.compiler.is_compiling(): + return cos_sin_cache + + self.cos_sin_cache = cos_sin_cache + return cos_sin_cache def get_cos_sin(self, seqlen: int) -> tuple[torch.Tensor, torch.Tensor]: cos_sin = self.cos_sin_cache[:seqlen] @@ -172,13 +181,14 @@ class RotaryEmbedding(RotaryEmbeddingBase): key: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor | None]: """A PyTorch-native implementation of forward().""" + cos_sin_cache = self._match_cos_sin_cache_dtype(query) return self.forward_static( positions, query, key, self.head_size, self.rotary_dim, - self.cos_sin_cache, + cos_sin_cache, self.is_neox_style, ) @@ -201,7 +211,7 @@ class RotaryEmbedding(RotaryEmbeddingBase): from vllm import _custom_ops as ops - self._match_cos_sin_cache_dtype(query) + cos_sin_cache = self._match_cos_sin_cache_dtype(query) # ops.rotary_embedding() is an in-place operation # that updates the query and key tensors. @@ -210,7 +220,7 @@ class RotaryEmbedding(RotaryEmbeddingBase): query, key, self.head_size, - self.cos_sin_cache, + cos_sin_cache, self.is_neox_style, ) return query, key @@ -222,12 +232,12 @@ class RotaryEmbedding(RotaryEmbeddingBase): key: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor | None]: if self.is_rocm_triton_rotary_embed_enabled: - self._match_cos_sin_cache_dtype(query) + cos_sin_cache = self._match_cos_sin_cache_dtype(query) rocm_aiter_ops.triton_rotary_embed( positions, query, key, - self.cos_sin_cache, + cos_sin_cache, self.head_size, self.rotary_dim, self.is_neox_style, @@ -249,12 +259,13 @@ class RotaryEmbedding(RotaryEmbeddingBase): else: from vllm import _custom_ops as ops + cos_sin_cache = self._match_cos_sin_cache_dtype(query) ops.rotary_embedding( positions, query, key, self.head_size, - self.cos_sin_cache, + cos_sin_cache, self.is_neox_style, ) return query, key @@ -267,7 +278,7 @@ class RotaryEmbedding(RotaryEmbeddingBase): ) -> tuple[torch.Tensor, torch.Tensor | None]: from vllm import _custom_ops as ops - self._match_cos_sin_cache_dtype(query) + cos_sin_cache = self._match_cos_sin_cache_dtype(query) # ops.rotary_embedding() is an in-place operation # that updates the query and key tensors. @@ -276,7 +287,7 @@ class RotaryEmbedding(RotaryEmbeddingBase): query, key, self.head_size, - self.cos_sin_cache, + cos_sin_cache, self.is_neox_style, ) return query, key diff --git a/vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py b/vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py index 9be9caacb45..c3abdc1563b 100644 --- a/vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py +++ b/vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py @@ -120,14 +120,14 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbeddingBase): ) -> tuple[torch.Tensor, torch.Tensor | None]: """PyTorch-native implementation equivalent to forward().""" assert key is not None - self._match_cos_sin_cache_dtype(query) + cos_sin_cache = self._match_cos_sin_cache_dtype(query) query_rot = query[..., : self.rotary_dim] key_rot = key[..., : self.rotary_dim] if self.rotary_dim < self.head_size: query_pass = query[..., self.rotary_dim :] key_pass = key[..., self.rotary_dim :] - cos_sin = self.cos_sin_cache[ + cos_sin = cos_sin_cache[ torch.add(positions, offsets) if offsets is not None else positions ] cos, sin = cos_sin.chunk(2, dim=-1) diff --git a/vllm/model_executor/layers/rotary_embedding/mrope.py b/vllm/model_executor/layers/rotary_embedding/mrope.py index a74bf092b18..52f3c333d7f 100644 --- a/vllm/model_executor/layers/rotary_embedding/mrope.py +++ b/vllm/model_executor/layers/rotary_embedding/mrope.py @@ -277,9 +277,9 @@ class MRotaryEmbedding(RotaryEmbeddingBase): assert positions.ndim == 1 or positions.ndim == 2 assert key is not None - self._match_cos_sin_cache_dtype(query) + cos_sin_cache = self._match_cos_sin_cache_dtype(query) num_tokens = positions.shape[-1] - cos_sin = self.cos_sin_cache[positions] + cos_sin = cos_sin_cache[positions] cos, sin = cos_sin.chunk(2, dim=-1) if positions.ndim == 2: assert self.mrope_section @@ -329,9 +329,9 @@ class MRotaryEmbedding(RotaryEmbeddingBase): assert positions.ndim == 1 or positions.ndim == 2 assert key is not None - self._match_cos_sin_cache_dtype(query) + cos_sin_cache = self._match_cos_sin_cache_dtype(query) num_tokens = positions.shape[-1] - cos_sin = self.cos_sin_cache[positions] + cos_sin = cos_sin_cache[positions] cos, sin = cos_sin.chunk(2, dim=-1) query_shape = query.shape key_shape = key.shape From 535de06cb1d90ed1c48246a512e74c87fe1768e4 Mon Sep 17 00:00:00 2001 From: Muhammad Hashmi <105992724+mu-hashmi@users.noreply.github.com> Date: Wed, 4 Feb 2026 13:17:47 -0800 Subject: [PATCH 070/810] [Model] Add transcription support for Qwen3-Omni (#29828) Signed-off-by: Muhammad Hashmi Signed-off-by: NickLucche Co-authored-by: NickLucche --- docs/contributing/model/transcription.md | 1 + docs/models/supported_models.md | 1 + .../models/qwen3_omni_moe_thinker.py | 104 +++++++++++++++++- 3 files changed, 104 insertions(+), 2 deletions(-) diff --git a/docs/contributing/model/transcription.md b/docs/contributing/model/transcription.md index fca941acd50..7fe010e5fd7 100644 --- a/docs/contributing/model/transcription.md +++ b/docs/contributing/model/transcription.md @@ -251,6 +251,7 @@ No extra registration is required beyond having your model class available via t - Whisper encoder–decoder (audio-only): [vllm/model_executor/models/whisper.py](../../../vllm/model_executor/models/whisper.py) - Voxtral decoder-only (audio embeddings + LLM): [vllm/model_executor/models/voxtral.py](../../../vllm/model_executor/models/voxtral.py). Make sure to have installed `mistral-common[audio]`. - Gemma3n decoder-only with fixed instruction prompt: [vllm/model_executor/models/gemma3n_mm.py](../../../vllm/model_executor/models/gemma3n_mm.py) +- Qwen3-Omni multimodal with audio embeddings: [vllm/model_executor/models/qwen3_omni_moe_thinker.py](../../../vllm/model_executor/models/qwen3_omni_moe_thinker.py) ## Test with the API diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index a96abd891fb..e07e17ec50d 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -781,6 +781,7 @@ Speech2Text models trained specifically for Automatic Speech Recognition. | `GlmAsrForConditionalGeneration` | GLM-ASR | `zai-org/GLM-ASR-Nano-2512` | ✅︎ | ✅︎ | | `GraniteSpeechForConditionalGeneration` | Granite Speech | `ibm-granite/granite-speech-3.3-2b`, `ibm-granite/granite-speech-3.3-8b`, etc. | ✅︎ | ✅︎ | | `Qwen3ASRForConditionalGeneration` | Qwen3-ASR | `Qwen/Qwen3-ASR-1.7B`, etc. | | ✅︎ | +| `Qwen3OmniMoeThinkerForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct`, etc. | | ✅︎ | | `VoxtralForConditionalGeneration` | Voxtral (Mistral format) | `mistralai/Voxtral-Mini-3B-2507`, `mistralai/Voxtral-Small-24B-2507`, etc. | ✅︎ | ✅︎ | | `WhisperForConditionalGeneration` | Whisper | `openai/whisper-small`, `openai/whisper-large-v3-turbo`, etc. | | | diff --git a/vllm/model_executor/models/qwen3_omni_moe_thinker.py b/vllm/model_executor/models/qwen3_omni_moe_thinker.py index 93a17f0c8c2..b065030319e 100755 --- a/vllm/model_executor/models/qwen3_omni_moe_thinker.py +++ b/vllm/model_executor/models/qwen3_omni_moe_thinker.py @@ -24,7 +24,7 @@ from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence from functools import partial -from typing import Any +from typing import Any, Literal, cast import numpy as np import torch @@ -48,8 +48,9 @@ from transformers import __version__ as TRANSFORMERS_VERSION # isort: on from vllm.compilation.decorators import support_torch_compile -from vllm.config import VllmConfig +from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size +from vllm.inputs.data import PromptType from vllm.logger import init_logger from vllm.model_executor.layers.activation import _ACTIVATION_REGISTRY from vllm.model_executor.layers.attention.mm_encoder_attention import ( @@ -79,6 +80,7 @@ from vllm.multimodal.processing.processor import ( PromptUpdateDetails, ) from vllm.sequence import IntermediateTensors +from vllm.transformers_utils.processor import cached_processor_from_config from vllm.v1.attention.backends.registry import AttentionBackendEnum from .interfaces import ( @@ -86,6 +88,7 @@ from .interfaces import ( SupportsMRoPE, SupportsMultiModal, SupportsPP, + SupportsTranscription, ) from .qwen2_5_omni_thinker import ( Qwen2_5OmniAudioFeatureInputs, @@ -110,6 +113,29 @@ from .vision import get_vit_attn_backend logger = init_logger(__name__) +# Speech input languages supported by Qwen3-Omni +# From: https://huggingface.co/Qwen/Qwen3-Omni-30B-A3B-Instruct +ISO639_1_SUPPORTED_LANGS = { + "en": "English", + "zh": "Chinese", + "ko": "Korean", + "ja": "Japanese", + "de": "German", + "ru": "Russian", + "it": "Italian", + "fr": "French", + "es": "Spanish", + "pt": "Portuguese", + "ms": "Malay", + "nl": "Dutch", + "id": "Indonesian", + "tr": "Turkish", + "vi": "Vietnamese", + "yue": "Cantonese", + "ar": "Arabic", + "ur": "Urdu", +} + def _get_feat_extract_output_lengths(input_lengths: torch.Tensor): input_lengths_leave = input_lengths % 100 @@ -1572,6 +1598,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( SupportsPP, SupportsMRoPE, Qwen3OmniMoeConditionalGenerationMixin, + SupportsTranscription, ): hf_to_vllm_mapper = WeightsMapper( orig_to_new_prefix={ @@ -1593,6 +1620,8 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( ], } + supported_languages = ISO639_1_SUPPORTED_LANGS + @classmethod def get_placeholder_str(cls, modality: str, i: int) -> str | None: if modality.startswith("image"): @@ -2085,6 +2114,77 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( total_tokens = num_video + audio_len return np.concatenate(pos_ids_list, axis=1), total_tokens + @classmethod + def get_speech_to_text_config( + cls, model_config: ModelConfig, task_type: str + ) -> SpeechToTextConfig: + processor = cached_processor_from_config( + model_config, processor_cls=Qwen3OmniMoeProcessor + ) + return SpeechToTextConfig( + max_audio_clip_s=processor.feature_extractor.chunk_length, + sample_rate=processor.feature_extractor.sampling_rate, + min_energy_split_window_size=None, + ) + + @classmethod + def get_generation_prompt( + cls, + audio: np.ndarray, + stt_config: SpeechToTextConfig, + model_config: ModelConfig, + language: str | None, + task_type: Literal["transcribe", "translate"], + request_prompt: str, + to_language: str | None, + ) -> PromptType: + """ + Construct a transcription/translation prompt for Qwen3-Omni. + """ + # Transcribe this audio [into ] | for transcription + # Translate this audio [from into ] | for translation + instruction = "Transcribe" if task_type == "transcribe" else "Translate" + instruction += " this audio" + + # Default to_language to English for translation + if task_type == "translate" and to_language is None: + to_language = "en" + + # Get full language names from supported_languages mapping + full_lang_name = cls.supported_languages.get(language, "") + full_lang_name_to = cls.supported_languages.get(to_language, "") + + if task_type == "transcribe" and full_lang_name: + instruction += f" into {full_lang_name}" + elif task_type == "translate": + if full_lang_name: + instruction += f" from {full_lang_name}" + if full_lang_name_to: + instruction += f" into {full_lang_name_to}" + + instruction += "." + + if request_prompt: + instruction += f" {request_prompt}" + + processor = cached_processor_from_config( + model_config, processor_cls=Qwen3OmniMoeProcessor + ) + # Audio placeholder format: <|audio_start|><|audio_pad|><|audio_end|> + audio_placeholder = "<|audio_start|><|audio_pad|><|audio_end|>" + user_content = f"{audio_placeholder}{instruction}" + + messages = [{"role": "user", "content": user_content}] + prompt = processor.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + ) + + audio_data = (audio, stt_config.sample_rate) + prompts_dict = {"multi_modal_data": {"audio": audio_data}, "prompt": prompt} + return cast(PromptType, prompts_dict) + def get_mrope_input_positions( self, input_tokens: list[int], From 9f14c9224d3d6664e2f5a2e7fecd012fd048fcb1 Mon Sep 17 00:00:00 2001 From: Richard Zou Date: Wed, 4 Feb 2026 13:59:59 -0800 Subject: [PATCH 071/810] Revert "[torch.compile] Significantly speed up cold start times" (#33820) Signed-off-by: Richard Zou --- tests/compile/test_cold_start.py | 9 +++-- vllm/compilation/backends.py | 50 ++++++++------------------ vllm/compilation/compiler_interface.py | 3 ++ 3 files changed, 21 insertions(+), 41 deletions(-) diff --git a/tests/compile/test_cold_start.py b/tests/compile/test_cold_start.py index dd770c58364..1d24d18397b 100644 --- a/tests/compile/test_cold_start.py +++ b/tests/compile/test_cold_start.py @@ -37,13 +37,12 @@ def test_moe_compilation_cold_start(monkeypatch, use_fresh_inductor_cache): # The forward pass consists of 32 transformer layers. # Then, we split on the attention operation. This results in # 33 subgraphs (not including the attention operation). - # We then standalone_compile the unique subgraphs. + # The 33 subgraphs then get standalone_compile'd. # # There are actually only 3 unique subgraphs for this model # (all of its transformer layers are the same modulo weights); # this is true for most vLLM models. - # So we test that during cold start, only 3 subgraphs are compiled - # These 3 subgraphs should cache miss, and then there should be - # no other compilation (so no cache hits). + # So we test that during cold start, the aot_autograd cache + # misses for 3 subgraphs and hits for the rest. assert counters["aot_autograd"]["autograd_cache_miss"] == 3 - assert counters["aot_autograd"]["autograd_cache_hit"] == 0 + assert counters["aot_autograd"]["autograd_cache_hit"] == 30 diff --git a/vllm/compilation/backends.py b/vllm/compilation/backends.py index 38ba97c7fae..89981fc2996 100644 --- a/vllm/compilation/backends.py +++ b/vllm/compilation/backends.py @@ -121,7 +121,7 @@ class CompilerManager: and compiling the graph. The cache is a dict mapping - `(runtime_shape, graph_hash, backend_name)` + `(runtime_shape, graph_index, backend_name)` to `any_data` returned from the compiler. When serializing the cache, we save it to a Python file @@ -130,7 +130,7 @@ class CompilerManager: """ def __init__(self, compilation_config: CompilationConfig) -> None: - self.cache: dict[tuple[Range, str, str], Any] = dict() + self.cache: dict[tuple[Range, int, str], Any] = dict() self.is_cache_updated = False self.compilation_config = compilation_config self.compiler = make_compiler(compilation_config) @@ -173,7 +173,6 @@ class CompilerManager: self.disable_cache = disable_cache self.cache_dir = cache_dir self.cache_file_path = os.path.join(cache_dir, "vllm_compile_cache.py") - self.loaded_cache_entries: dict[tuple[Range, str, str], Any] = {} if not disable_cache and os.path.exists(self.cache_file_path): # load the cache from the file @@ -187,9 +186,9 @@ class CompilerManager: if not isinstance(value, ty): raise TypeError(f"Expected {ty} but got {type(value)} for {value}") - def parse_key(key: Any) -> tuple[Range, str, str]: - range_tuple, graph_hash, compiler_name = key - check_type(graph_hash, str) + def parse_key(key: Any) -> tuple[Range, int, str]: + range_tuple, graph_index, compiler_name = key + check_type(graph_index, int) check_type(compiler_name, str) if isinstance(range_tuple, tuple): start, end = range_tuple @@ -197,7 +196,7 @@ class CompilerManager: check_type(end, int) range_tuple = Range(start=start, end=end) check_type(range_tuple, Range) - return range_tuple, graph_hash, compiler_name + return range_tuple, graph_index, compiler_name self.cache = {parse_key(key): value for key, value in cache.items()} @@ -217,25 +216,18 @@ class CompilerManager: self, graph: fx.GraphModule, example_inputs: list[Any], - graph_hash: str, + graph_index: int, compile_range: Range, ) -> Callable[..., Any] | None: - key = (compile_range, graph_hash, self.compiler.name) - # See if we've already loaded this cache entry - if key in self.loaded_cache_entries: - return self.loaded_cache_entries[key] - # Otherwise, go load it from disk - if key not in self.cache: + if (compile_range, graph_index, self.compiler.name) not in self.cache: return None - handle = self.cache[key] + handle = self.cache[(compile_range, graph_index, self.compiler.name)] compiled_graph = self.compiler.load( - handle, graph, example_inputs, compile_range + handle, graph, example_inputs, graph_index, compile_range ) - self.loaded_cache_entries[key] = compiled_graph logger.debug( - "Directly load the graph (hash %s) for compile range " - "%sfrom %s via handle %s", - graph_hash, + "Directly load the %s-th graph for compile range %sfrom %s via handle %s", + graph_index, str(compile_range), self.compiler.name, handle, @@ -257,22 +249,12 @@ class CompilerManager: global compilation_start_time compilation_start_time = time.time() - from torch._functorch._aot_autograd.autograd_cache import ( - AOTAutogradCachePickler, - sanitize_gm_for_cache, - ) - - with sanitize_gm_for_cache(graph): - pickler = AOTAutogradCachePickler(graph) - dumped_graph = pickler.dumps(graph) - graph_hash = hashlib.sha256(dumped_graph).hexdigest() - compilation_counter.num_backend_compilations += 1 compiled_graph = None # try to load from the cache - compiled_graph = self.load(graph, example_inputs, graph_hash, compile_range) + compiled_graph = self.load(graph, example_inputs, graph_index, compile_range) if compiled_graph is not None: if graph_index == num_graphs - 1: # after loading the last graph for this shape, record the time. @@ -308,13 +290,9 @@ class CompilerManager: assert compiled_graph is not None, "Failed to compile the graph" - self.loaded_cache_entries[(compile_range, graph_hash, self.compiler.name)] = ( - compiled_graph - ) - # store the artifact in the cache if is_compile_cache_enabled(additional_inductor_config) and handle is not None: - self.cache[(compile_range, graph_hash, self.compiler.name)] = handle + self.cache[(compile_range, graph_index, self.compiler.name)] = handle compilation_counter.num_cache_entries_updated += 1 self.is_cache_updated = True if graph_index == 0: diff --git a/vllm/compilation/compiler_interface.py b/vllm/compilation/compiler_interface.py index 875e628d686..60650353971 100644 --- a/vllm/compilation/compiler_interface.py +++ b/vllm/compilation/compiler_interface.py @@ -101,6 +101,7 @@ class CompilerInterface: handle: Any, graph: fx.GraphModule, example_inputs: list[Any], + graph_index: int, compile_range: Range, ) -> Callable[..., Any]: """ @@ -301,6 +302,7 @@ class InductorStandaloneAdaptor(CompilerInterface): handle: Any, graph: fx.GraphModule, example_inputs: list[Any], + graph_index: int, compile_range: Range, ) -> Callable[..., Any]: assert isinstance(handle, tuple) @@ -525,6 +527,7 @@ class InductorAdaptor(CompilerInterface): handle: Any, graph: fx.GraphModule, example_inputs: list[Any], + graph_index: int, compile_range: Range, ) -> Callable[..., Any]: assert isinstance(handle, tuple) From ce498a6d61a083b1bbbd1cc92754961b43860ff6 Mon Sep 17 00:00:00 2001 From: Sage Moore Date: Wed, 4 Feb 2026 14:02:46 -0800 Subject: [PATCH 072/810] Change the type signature of MixtureOfExperts.expert_weights to MutableSequence[Sequence[Tensor]] (#33573) Signed-off-by: Sage Moore Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> --- vllm/distributed/eplb/rebalance_execute.py | 16 +++++++++------- vllm/model_executor/models/interfaces.py | 11 +++++++++-- 2 files changed, 18 insertions(+), 9 deletions(-) diff --git a/vllm/distributed/eplb/rebalance_execute.py b/vllm/distributed/eplb/rebalance_execute.py index b7b6c11b239..72bbe1c5d9c 100644 --- a/vllm/distributed/eplb/rebalance_execute.py +++ b/vllm/distributed/eplb/rebalance_execute.py @@ -6,7 +6,7 @@ The actual execution of the rearrangement. This involves the exchange of expert weights between GPUs. """ -from collections.abc import Iterable, Sequence +from collections.abc import Sequence from dataclasses import dataclass import numpy as np @@ -153,7 +153,7 @@ def move_to_buffer( num_local_experts: int, old_indices: np.ndarray, new_indices: np.ndarray, - expert_weights: Iterable[torch.Tensor], + expert_weights: Sequence[torch.Tensor], expert_weights_buffers: Sequence[torch.Tensor], cuda_stream: torch.cuda.Stream | None, ep_group: ProcessGroup, @@ -355,7 +355,7 @@ def move_to_buffer( def move_from_buffer( - expert_weights: Iterable[torch.Tensor], + expert_weights: Sequence[torch.Tensor], expert_weights_buffers: list[torch.Tensor], is_unchanged: np.ndarray, is_received_locally: np.ndarray, @@ -436,7 +436,7 @@ def move_from_buffer( async def transfer_layer( old_global_expert_indices: torch.Tensor, new_global_expert_indices: torch.Tensor, - expert_weights: Sequence[Iterable[torch.Tensor]], + expert_weights: Sequence[Sequence[torch.Tensor]], expert_weights_buffer: Sequence[torch.Tensor], ep_group: ProcessGroup, is_profile: bool = False, @@ -488,7 +488,8 @@ async def transfer_layer( assert old_global_expert_indices.shape[1] == new_global_expert_indices.shape[1] num_moe_layers, num_physical_experts = old_global_expert_indices.shape assert len(expert_weights) == num_moe_layers - num_local_physical_experts = next(iter(expert_weights[0])).shape[0] + assert len(expert_weights[0]) >= 1 + num_local_physical_experts = expert_weights[0][0].shape[0] assert new_global_expert_indices.shape == (num_moe_layers, num_physical_experts) assert num_physical_experts == ep_size * num_local_physical_experts @@ -510,7 +511,7 @@ async def transfer_layer( def rearrange_expert_weights_inplace( old_global_expert_indices: torch.Tensor, new_global_expert_indices: torch.Tensor, - expert_weights: Sequence[Iterable[torch.Tensor]], + expert_weights: Sequence[Sequence[torch.Tensor]], ep_group: ProcessGroup, is_profile: bool = False, rank_mapping: dict[int, int] | None = None, @@ -553,8 +554,9 @@ def rearrange_expert_weights_inplace( num_moe_layers, num_physical_experts = old_global_expert_indices.shape assert len(expert_weights) == num_moe_layers + assert len(expert_weights[0]) >= 1 - num_local_physical_experts = next(iter(expert_weights[0])).shape[0] + num_local_physical_experts = expert_weights[0][0].shape[0] assert new_global_expert_indices.shape == (num_moe_layers, num_physical_experts) ep_size = ep_group.size() diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index f05231356d5..c97a9faf6c5 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -2,7 +2,14 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio -from collections.abc import AsyncGenerator, Callable, Iterable, Mapping, MutableSequence +from collections.abc import ( + AsyncGenerator, + Callable, + Iterable, + Mapping, + MutableSequence, + Sequence, +) from contextlib import ExitStack, contextmanager, nullcontext from typing import ( TYPE_CHECKING, @@ -818,7 +825,7 @@ class MixtureOfExperts(Protocol): Check if the model is a mixture of experts (MoE) model. """ - expert_weights: MutableSequence[Iterable[Tensor]] + expert_weights: MutableSequence[Sequence[Tensor]] """ Expert weights saved in this rank. From fa4e0fb028460cf5f4eb9cc90e206d0d6f35b026 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Wed, 4 Feb 2026 15:40:22 -0800 Subject: [PATCH 073/810] [Core] Don't schedule spec tokens with prefill chunks (#33652) Signed-off-by: Nick Hill --- tests/v1/core/test_scheduler.py | 94 +++++++++++++++++++++++++ vllm/v1/core/sched/async_scheduler.py | 31 +++----- vllm/v1/core/sched/scheduler.py | 12 ++++ vllm/v1/request.py | 3 + vllm/v1/worker/gpu/spec_decode/utils.py | 18 ++--- 5 files changed, 129 insertions(+), 29 deletions(-) diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 063d0a644fa..b29df468f0c 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -945,6 +945,100 @@ def test_spec_decoding_stats_empty_output(): assert scheduler_stats is None or scheduler_stats.spec_decoding_stats is None +def test_no_spec_tokens_scheduled_for_prefill_chunks(): + """Test that draft tokens are ignored for prefill chunk requests. + + When a request is being prefilled in chunks (chunked prefill), draft tokens + from `update_draft_token_ids` should be ignored until the prefill is complete. + + The bug manifests when: + - A prefill chunk is scheduled + - Draft tokens are provided via update_draft_token_ids + - The next schedule has enough budget to include spec tokens + + Without the fix, spec tokens would incorrectly be scheduled with the + remaining prefill tokens. With the fix, draft tokens are ignored for + prefill chunks. + """ + num_spec_tokens = 3 + # Use budget of 50, with 80 token prompt: + # - First chunk: 50 tokens + # - Second chunk: 30 remaining + potentially 3 spec tokens = 33 + # Without fix: num_scheduled_spec_tokens = 33 + 50 - 80 = 3 (BUG!) + # With fix: spec_token_ids cleared, so no spec tokens scheduled + scheduler = create_scheduler( + num_speculative_tokens=num_spec_tokens, + max_num_batched_tokens=50, + enable_chunked_prefill=True, + ) + requests = create_requests(num_requests=1, num_tokens=80) + req = requests[0] + scheduler.add_request(req) + + # First schedule - prefill chunk (50 of 80 tokens) + output = scheduler.schedule() + assert len(output.scheduled_new_reqs) == 1 + assert output.num_scheduled_tokens[req.request_id] == 50 + + # Update from output (no sampled token since still prefilling) + req_to_index = {req.request_id: 0} + model_runner_output = ModelRunnerOutput( + req_ids=[req.request_id], + req_id_to_index=req_to_index, + sampled_token_ids=[[]], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ) + scheduler.update_from_output(output, model_runner_output) + + # Provide draft tokens while request is still in prefill. + # The fix ensures these are ignored for prefill chunks. + draft_token_ids = DraftTokenIds([req.request_id], [[1, 2, 3]]) + scheduler.update_draft_token_ids(draft_token_ids) + + # Second schedule - remaining 30 tokens of prefill + output = scheduler.schedule() + # KEY ASSERTION: Should schedule exactly the remaining 30 prefill tokens, + # NOT 33 (30 + 3 spec). Without the fix, this would be 33. + assert output.num_scheduled_tokens[req.request_id] == 30, ( + f"Expected 30 tokens (remaining prefill only), " + f"got {output.num_scheduled_tokens[req.request_id]}. " + "Spec tokens should not be scheduled with prefill chunks." + ) + # No spec tokens should be in the output + assert req.request_id not in output.scheduled_spec_decode_tokens, ( + "Spec tokens should not be scheduled with prefill chunks" + ) + + # Update from output with a sampled token (prefill complete) + model_runner_output = ModelRunnerOutput( + req_ids=[req.request_id], + req_id_to_index=req_to_index, + sampled_token_ids=[[42]], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ) + scheduler.update_from_output(output, model_runner_output) + + # Now provide draft tokens - should be accepted since prefill is complete + draft_token_ids = DraftTokenIds([req.request_id], [[1, 2, 3]]) + scheduler.update_draft_token_ids(draft_token_ids) + + # spec_token_ids SHOULD be set after prefill is complete + assert req.spec_token_ids == [1, 2, 3], ( + f"spec_token_ids should be set after prefill, got {req.spec_token_ids}" + ) + + # Third schedule - decode phase with spec tokens + output = scheduler.schedule() + # 1 new token + 3 spec tokens = 4 + assert output.num_scheduled_tokens[req.request_id] == 4 + assert req.request_id in output.scheduled_spec_decode_tokens + assert len(output.scheduled_spec_decode_tokens[req.request_id]) == num_spec_tokens + + def _assert_right_scheduler_output( output: SchedulerOutput, num_requests: int, diff --git a/vllm/v1/core/sched/async_scheduler.py b/vllm/v1/core/sched/async_scheduler.py index 23c610f3bde..0b3958dbcf5 100644 --- a/vllm/v1/core/sched/async_scheduler.py +++ b/vllm/v1/core/sched/async_scheduler.py @@ -17,33 +17,22 @@ class AsyncScheduler(Scheduler): def _update_after_schedule(self, scheduler_output: SchedulerOutput) -> None: super()._update_after_schedule(scheduler_output) - has_structured_output_requests = False - pending_structured_output_tokens = False spec_decode_tokens = scheduler_output.scheduled_spec_decode_tokens for req_id in scheduler_output.num_scheduled_tokens: request = self.requests[req_id] - has_structured_output_requests |= request.use_structured_output - pending_structured_output_tokens |= ( + if request.is_prefill_chunk: + continue + + scheduler_output.pending_structured_output_tokens |= ( request.use_structured_output and request.num_output_placeholders > 0 ) + # The request will generate a new token plus num_spec_tokens + # in this scheduling step. cur_num_spec_tokens = len(spec_decode_tokens.get(req_id, ())) - if ( - request.num_computed_tokens - == request.num_tokens - + request.num_output_placeholders - + cur_num_spec_tokens - ): - # The request will generate a new token plus num_spec_tokens - # in this scheduling step. - request.num_output_placeholders += 1 + cur_num_spec_tokens - # Add placeholders for the new draft/spec tokens. - # We will update the actual spec token ids in the worker process. - request.spec_token_ids = self._spec_token_placeholders - - scheduler_output.has_structured_output_requests = has_structured_output_requests - scheduler_output.pending_structured_output_tokens = ( - pending_structured_output_tokens - ) + request.num_output_placeholders += 1 + cur_num_spec_tokens + # Add placeholders for the new draft/spec tokens. + # We will update the actual spec token ids in the worker process. + request.spec_token_ids = self._spec_token_placeholders def _update_request_with_output( self, request: Request, new_token_ids: list[int] diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 869b53601b1..88d1a78df00 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -912,6 +912,12 @@ class Scheduler(SchedulerInterface): for req_id, num_scheduled_token in num_scheduled_tokens.items(): request = self.requests[req_id] request.num_computed_tokens += num_scheduled_token + request.is_prefill_chunk = request.num_computed_tokens < ( + request.num_tokens + request.num_output_placeholders + ) + scheduler_output.has_structured_output_requests |= ( + request.use_structured_output + ) # NOTE: _free_encoder_inputs relies on num_computed_tokens, which # may be updated again in _update_from_output for speculative @@ -1562,6 +1568,12 @@ class Scheduler(SchedulerInterface): # The request may have been finished. Skip. continue + if request.is_prefill_chunk: + # Ignore draft tokens for prefill chunks. + if request.spec_token_ids: + request.spec_token_ids = [] + continue + # Add newly generated spec token ids to the request. if self.structured_output_manager.should_advance(request): metadata = request.structured_output_request diff --git a/vllm/v1/request.py b/vllm/v1/request.py index 117478a92c1..e9d3df4421e 100644 --- a/vllm/v1/request.py +++ b/vllm/v1/request.py @@ -147,6 +147,9 @@ class Request: # The number of tokens with prefix cache hits. self.num_cached_tokens = -1 + # True if this request is scheduled as a non-final prefill chunk. + self.is_prefill_chunk = False + # The number of NaNs in logits. A value greater than 0 # indicates that the output is corrupted self.num_nans_in_logits = 0 diff --git a/vllm/v1/worker/gpu/spec_decode/utils.py b/vllm/v1/worker/gpu/spec_decode/utils.py index ddeb99a7180..e1fa21aeb8a 100644 --- a/vllm/v1/worker/gpu/spec_decode/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/utils.py @@ -16,21 +16,21 @@ class DraftTokensHandler: self.req_ids: list[str] = [] self.draft_tokens_np: np.ndarray | None = None + self.num_draft_tokens: int = 0 def set_draft_tokens( self, input_batch: InputBatch, draft_tokens: torch.Tensor ) -> None: + self.req_ids = input_batch.req_ids + self.num_draft_tokens = draft_tokens.shape[1] if not input_batch.has_structured_output_reqs: # No draft token validation needs to be performed by # the scheduler for this batch. - if self.req_ids: - self.req_ids = [] self.draft_tokens_np = None return # For spec decoding + structured outputs, we must transfer the # draft tokens back to the scheduler for grammar validation. - self.req_ids = input_batch.req_ids current_stream = torch.cuda.current_stream(self.device) self.copy_stream.wait_stream(current_stream) with torch.cuda.stream(self.copy_stream): @@ -38,8 +38,10 @@ class DraftTokensHandler: self.copy_event.record() def get_draft_tokens(self) -> DraftTokenIds | None: - if self.draft_tokens_np is None: - return None - - self.copy_event.synchronize() - return DraftTokenIds(self.req_ids, self.draft_tokens_np.tolist()) + if self.draft_tokens_np is not None: + self.copy_event.synchronize() + draft_token_ids = self.draft_tokens_np.tolist() + else: + # This case only happens when async scheduling is disabled. + draft_token_ids = [[-1] * self.num_draft_tokens for _ in self.req_ids] + return DraftTokenIds(self.req_ids, draft_token_ids) From 439afa4eea14db2be232a9ce78eacc2c7bbfac77 Mon Sep 17 00:00:00 2001 From: Ilya Boytsov Date: Thu, 5 Feb 2026 01:05:13 +0100 Subject: [PATCH 074/810] feat: Add ColBERT late interaction model support (#33686) Signed-off-by: Ilya Boytsov Signed-off-by: Ilya Boytsov Co-authored-by: Cyrus Leung Co-authored-by: wang.yuqi --- docs/models/pooling_models.md | 56 ++++ .../pooling/score/colbert_rerank_online.py | 57 ++++ .../pooling/score/test_online_colbert.py | 154 +++++++++++ tests/models/language/pooling/test_colbert.py | 247 ++++++++++++++++++ tests/models/registry.py | 1 + vllm/config/model.py | 5 + vllm/entrypoints/llm.py | 97 ++++++- vllm/entrypoints/pooling/__init__.py | 12 +- vllm/entrypoints/pooling/score/serving.py | 140 ++++++++++ vllm/entrypoints/pooling/score/utils.py | 18 ++ vllm/model_executor/models/colbert.py | 152 +++++++++++ vllm/model_executor/models/interfaces.py | 34 +++ vllm/model_executor/models/registry.py | 4 + 13 files changed, 974 insertions(+), 3 deletions(-) create mode 100644 examples/pooling/score/colbert_rerank_online.py create mode 100644 tests/entrypoints/pooling/score/test_online_colbert.py create mode 100644 tests/models/language/pooling/test_colbert.py create mode 100644 vllm/model_executor/models/colbert.py diff --git a/docs/models/pooling_models.md b/docs/models/pooling_models.md index c1355fe49b5..0555eac41ad 100644 --- a/docs/models/pooling_models.md +++ b/docs/models/pooling_models.md @@ -307,6 +307,62 @@ An OpenAI client example can be found here: [examples/pooling/embed/openai_embed ## Specific models +### ColBERT Late Interaction Models + +[ColBERT](https://arxiv.org/abs/2004.12832) (Contextualized Late Interaction over BERT) is a retrieval model that uses per-token embeddings and MaxSim scoring for document ranking. Unlike single-vector embedding models, ColBERT retains token-level representations and computes relevance scores through late interaction, providing better accuracy while being more efficient than cross-encoders. + +vLLM supports ColBERT models for reranking tasks, automatically applying MaxSim scoring for query-document relevance: + +```shell +vllm serve answerdotai/answerai-colbert-small-v1 +``` + +Currently supports ColBERT models with standard BERT encoders (e.g., `answerdotai/answerai-colbert-small-v1`, `colbert-ir/colbertv2.0`). + +ColBERT models with modified encoder architectures are not yet supported, including BERT variants with rotary embeddings (e.g., `jinaai/jina-colbert-v2`) or other custom encoders (e.g., `LiquidAI/LFM2-ColBERT-350M`). + +If your standard BERT ColBERT model's config doesn't specify the architecture as `HF_ColBERT`, override it with: + +```shell +vllm serve your-colbert-model --hf-overrides '{"architectures": ["HF_ColBERT"]}' +``` + +Then you can use the rerank endpoint: + +```shell +curl -s http://localhost:8000/rerank -H "Content-Type: application/json" -d '{ + "model": "answerdotai/answerai-colbert-small-v1", + "query": "What is machine learning?", + "documents": [ + "Machine learning is a subset of artificial intelligence.", + "Python is a programming language.", + "Deep learning uses neural networks." + ] +}' +``` + +Or the score endpoint: + +```shell +curl -s http://localhost:8000/score -H "Content-Type: application/json" -d '{ + "model": "answerdotai/answerai-colbert-small-v1", + "text_1": "What is machine learning?", + "text_2": ["Machine learning is a subset of AI.", "The weather is sunny."] +}' +``` + +You can also get the raw token embeddings using the pooling endpoint with `token_embed` task: + +```shell +curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{ + "model": "answerdotai/answerai-colbert-small-v1", + "input": "What is machine learning?", + "task": "token_embed" +}' +``` + +An example can be found here: [examples/pooling/score/colbert_rerank_online.py](../../examples/pooling/score/colbert_rerank_online.py) + ### BAAI/bge-m3 The `BAAI/bge-m3` model comes with extra weights for sparse and colbert embeddings but unfortunately in its `config.json` diff --git a/examples/pooling/score/colbert_rerank_online.py b/examples/pooling/score/colbert_rerank_online.py new file mode 100644 index 00000000000..b9223e79157 --- /dev/null +++ b/examples/pooling/score/colbert_rerank_online.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Example of using ColBERT late interaction model for reranking. + +ColBERT (Contextualized Late Interaction over BERT) uses per-token embeddings +and MaxSim scoring for document reranking, providing better accuracy than +single-vector models while being more efficient than cross-encoders. + +Start the server with: + vllm serve answerdotai/answerai-colbert-small-v1 + +Then run this script: + python colbert_rerank_online.py +""" + +import json + +import requests + +url = "http://127.0.0.1:8000/rerank" + +headers = {"accept": "application/json", "Content-Type": "application/json"} + +data = { + "model": "answerdotai/answerai-colbert-small-v1", + "query": "What is machine learning?", + "documents": [ + "Machine learning is a subset of artificial intelligence.", + "Python is a programming language.", + "Deep learning uses neural networks for complex tasks.", + "The weather today is sunny.", + ], +} + + +def main(): + response = requests.post(url, headers=headers, json=data) + + if response.status_code == 200: + print("ColBERT Rerank Request successful!") + result = response.json() + print(json.dumps(result, indent=2)) + + # Show ranked results + print("\nRanked documents (most relevant first):") + for item in result["results"]: + doc_idx = item["index"] + score = item["relevance_score"] + print(f" Score {score:.4f}: {data['documents'][doc_idx]}") + else: + print(f"Request failed with status code: {response.status_code}") + print(response.text) + + +if __name__ == "__main__": + main() diff --git a/tests/entrypoints/pooling/score/test_online_colbert.py b/tests/entrypoints/pooling/score/test_online_colbert.py new file mode 100644 index 00000000000..a7b404d0fde --- /dev/null +++ b/tests/entrypoints/pooling/score/test_online_colbert.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Online API tests for ColBERT late interaction scoring.""" + +import pytest +import requests + +from tests.utils import RemoteOpenAIServer +from vllm.entrypoints.pooling.score.protocol import RerankResponse, ScoreResponse + +# ColBERT model - using answerai-colbert-small-v1 as it's a smaller model +MODEL_NAME = "answerdotai/answerai-colbert-small-v1" +COLBERT_DIM = 96 # This model uses 96-dimensional output +DTYPE = "half" +MAX_MODEL_LEN = 512 + + +@pytest.fixture(scope="module") +def server(): + args = [ + "--max-model-len", + str(MAX_MODEL_LEN), + ] + + with RemoteOpenAIServer(MODEL_NAME, args) as remote_server: + yield remote_server + + +@pytest.mark.parametrize("model_name", [MODEL_NAME]) +def test_colbert_rerank(server: RemoteOpenAIServer, model_name: str): + """Test ColBERT rerank endpoint.""" + query = "What is the capital of France?" + documents = [ + "The capital of Brazil is Brasilia.", + "The capital of France is Paris.", + ] + + rerank_response = requests.post( + server.url_for("rerank"), + json={ + "model": model_name, + "query": query, + "documents": documents, + }, + ) + rerank_response.raise_for_status() + rerank = RerankResponse.model_validate(rerank_response.json()) + + assert rerank.id is not None + assert rerank.results is not None + assert len(rerank.results) == 2 + + # The relevant document (Paris) should have higher score + paris_result = next(r for r in rerank.results if r.index == 1) + brazil_result = next(r for r in rerank.results if r.index == 0) + + assert paris_result.relevance_score > brazil_result.relevance_score + + +@pytest.mark.parametrize("model_name", [MODEL_NAME]) +def test_colbert_rerank_top_n(server: RemoteOpenAIServer, model_name: str): + """Test ColBERT rerank with top_n parameter.""" + query = "What is the capital of France?" + documents = [ + "The capital of Brazil is Brasilia.", + "The capital of France is Paris.", + "Machine learning is a field of AI.", + ] + + rerank_response = requests.post( + server.url_for("rerank"), + json={ + "model": model_name, + "query": query, + "documents": documents, + "top_n": 2, + }, + ) + rerank_response.raise_for_status() + rerank = RerankResponse.model_validate(rerank_response.json()) + + assert len(rerank.results) == 2 + # Top result should be about Paris (index 1) + assert rerank.results[0].index == 1 + + +@pytest.mark.parametrize("model_name", [MODEL_NAME]) +def test_colbert_score(server: RemoteOpenAIServer, model_name: str): + """Test ColBERT score endpoint.""" + text_1 = "What is the capital of France?" + text_2 = ["The capital of France is Paris.", "Python is a language."] + + score_response = requests.post( + server.url_for("score"), + json={ + "model": model_name, + "text_1": text_1, + "text_2": text_2, + }, + ) + score_response.raise_for_status() + score = ScoreResponse.model_validate(score_response.json()) + + assert score.id is not None + assert score.data is not None + assert len(score.data) == 2 + + # The relevant document should have higher score + assert score.data[0].score > score.data[1].score + + +@pytest.mark.parametrize("model_name", [MODEL_NAME]) +def test_colbert_token_embed(server: RemoteOpenAIServer, model_name: str): + """Test ColBERT token_embed task via pooling endpoint.""" + text = "What is the capital of France?" + + pooling_response = requests.post( + server.url_for("pooling"), + json={ + "model": model_name, + "input": text, + "task": "token_embed", + }, + ) + pooling_response.raise_for_status() + pooling = pooling_response.json() + + assert "data" in pooling + assert len(pooling["data"]) == 1 + + # Token embeddings should be 2D + embeddings = pooling["data"][0]["data"] + assert isinstance(embeddings, list) + assert len(embeddings) > 0 # Should have tokens + assert len(embeddings[0]) == COLBERT_DIM + + +@pytest.mark.parametrize("model_name", [MODEL_NAME]) +def test_colbert_embed_not_supported(server: RemoteOpenAIServer, model_name: str): + """Test that ColBERT model does not support 'embed' task.""" + text = "What is the capital of France?" + + pooling_response = requests.post( + server.url_for("pooling"), + json={ + "model": model_name, + "input": text, + "task": "embed", + }, + ) + + # Should return error + assert pooling_response.status_code == 400 + assert "Task embed is not supported" in pooling_response.text diff --git a/tests/models/language/pooling/test_colbert.py b/tests/models/language/pooling/test_colbert.py new file mode 100644 index 00000000000..fa77b8c2680 --- /dev/null +++ b/tests/models/language/pooling/test_colbert.py @@ -0,0 +1,247 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for ColBERT late interaction scoring.""" + +import pytest +import torch + +from vllm.entrypoints.pooling.score.utils import compute_maxsim_score + +# ColBERT model - using answerai-colbert-small-v1 as it's a smaller model +# suitable for testing (based on BERT-base) +COLBERT_MODEL = "answerdotai/answerai-colbert-small-v1" +COLBERT_DIM = 96 # This model uses 96-dimensional output + +TEXTS_1 = [ + "What is the capital of France?", + "What is the capital of Germany?", +] + +TEXTS_2 = [ + "The capital of France is Paris.", + "The capital of Germany is Berlin.", +] + +DTYPE = "half" + + +@pytest.fixture(scope="module") +def colbert_model_name(): + return COLBERT_MODEL + + +def test_colbert_token_embed(vllm_runner, colbert_model_name): + """Test that ColBERT model produces token embeddings.""" + with vllm_runner( + colbert_model_name, + runner="pooling", + dtype=DTYPE, + max_model_len=512, + enforce_eager=True, + ) as vllm_model: + # Get token embeddings for a single text + outputs = vllm_model.token_embed([TEXTS_1[0]]) + + assert len(outputs) == 1 + # Token embeddings should be 2D: [num_tokens, colbert_dim] + emb = torch.tensor(outputs[0]) + assert emb.dim() == 2 + assert emb.shape[1] == COLBERT_DIM + # Should have at least a few tokens + assert emb.shape[0] > 1 + + +def test_colbert_late_interaction_1_to_1(vllm_runner, colbert_model_name): + """Test ColBERT late interaction scoring with 1:1 query-document pair.""" + with vllm_runner( + colbert_model_name, + runner="pooling", + dtype=DTYPE, + max_model_len=512, + enforce_eager=True, + ) as vllm_model: + # Get token embeddings + q_outputs = vllm_model.token_embed([TEXTS_1[0]]) + d_outputs = vllm_model.token_embed([TEXTS_2[0]]) + + q_emb = torch.tensor(q_outputs[0]) + d_emb = torch.tensor(d_outputs[0]) + + # Compute MaxSim manually + manual_score = compute_maxsim_score(q_emb, d_emb).item() + + # Use the score API (which should internally use _late_interaction_score) + vllm_scores = vllm_model.score(TEXTS_1[0], TEXTS_2[0]) + + 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): + """Test ColBERT late interaction scoring with 1:N query-documents.""" + with vllm_runner( + colbert_model_name, + runner="pooling", + dtype=DTYPE, + max_model_len=512, + enforce_eager=True, + ) as vllm_model: + # Get token embeddings + q_outputs = vllm_model.token_embed([TEXTS_1[0]]) + d_outputs = vllm_model.token_embed(TEXTS_2) + + q_emb = torch.tensor(q_outputs[0]) + + # Compute MaxSim manually for each document + manual_scores = [] + for d_out in d_outputs: + d_emb = torch.tensor(d_out) + manual_scores.append(compute_maxsim_score(q_emb, d_emb).item()) + + # Use the score API + vllm_scores = vllm_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) + + +def test_colbert_late_interaction_N_to_N(vllm_runner, colbert_model_name): + """Test ColBERT late interaction scoring with N:N query-documents.""" + with vllm_runner( + colbert_model_name, + runner="pooling", + dtype=DTYPE, + max_model_len=512, + enforce_eager=True, + ) as vllm_model: + # Get token embeddings + q_outputs = vllm_model.token_embed(TEXTS_1) + d_outputs = vllm_model.token_embed(TEXTS_2) + + # Compute MaxSim manually for each pair + manual_scores = [] + for q_out, d_out in zip(q_outputs, d_outputs): + q_emb = torch.tensor(q_out) + d_emb = torch.tensor(d_out) + manual_scores.append(compute_maxsim_score(q_emb, d_emb).item()) + + # Use the score API + vllm_scores = vllm_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) + + +def test_colbert_relevance_ordering(vllm_runner, colbert_model_name): + """Test that ColBERT scores relevant documents higher than irrelevant ones.""" + query = "What is machine learning?" + documents = [ + "Machine learning is a subset of artificial intelligence.", + "Python is a programming language.", + "Deep learning uses neural networks.", + ] + + with vllm_runner( + colbert_model_name, + runner="pooling", + dtype=DTYPE, + max_model_len=512, + enforce_eager=True, + ) as vllm_model: + scores = vllm_model.score(query, documents) + + assert len(scores) == 3 + # ML-related documents should score higher than unrelated Python doc + # Document 0 (ML definition) should be most relevant + # Document 2 (Deep learning) should also be relevant + # Document 1 (Python) should be least relevant + 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): + """Test that ColBERT model does not support 'embed' task.""" + with ( + vllm_runner( + colbert_model_name, + runner="pooling", + dtype=DTYPE, + max_model_len=512, + enforce_eager=True, + ) as vllm_model, + pytest.raises(ValueError, match="Embedding API is not supported"), + ): + vllm_model.embed([TEXTS_1[0]]) + + +def test_colbert_hf_comparison(vllm_runner, colbert_model_name): + """Test that vLLM ColBERT produces same embeddings as HuggingFace.""" + import torch.nn.functional as F + from huggingface_hub import hf_hub_download + from safetensors.torch import load_file + from transformers import AutoTokenizer, BertModel + + test_texts = [TEXTS_1[0], TEXTS_2[0]] + + # Get vLLM embeddings first (to avoid GPU memory contention) + # Use fp32 to match HuggingFace default precision for fair comparison + with vllm_runner( + colbert_model_name, + runner="pooling", + dtype="float32", + max_model_len=512, + enforce_eager=True, + ) as vllm_model: + vllm_outputs = vllm_model.token_embed(test_texts) + + # Get HuggingFace reference embeddings on CPU + # Load the base BERT model and manually apply the ColBERT linear projection + hf_tokenizer = AutoTokenizer.from_pretrained(colbert_model_name) + hf_bert = BertModel.from_pretrained(colbert_model_name) + hf_bert.eval() + + # Load the ColBERT linear weights from safetensors + weights_path = hf_hub_download(colbert_model_name, filename="model.safetensors") + weights = load_file(weights_path) + linear_weight = weights["linear.weight"] # [96, 384] + + hf_embeddings = [] + for text in test_texts: + inputs = hf_tokenizer(text, return_tensors="pt") + with torch.no_grad(): + outputs = hf_bert(**inputs) + # Get last hidden state: [1, seq_len, 384] + hidden_states = outputs.last_hidden_state + # Apply ColBERT linear projection: [1, seq_len, 96] + token_emb = F.linear(hidden_states, linear_weight) + # L2 normalize + token_emb = F.normalize(token_emb, p=2, dim=-1) + hf_embeddings.append(token_emb.squeeze(0).float()) + + # Compare embeddings + for i, (hf_emb, vllm_out) in enumerate(zip(hf_embeddings, vllm_outputs)): + vllm_emb = torch.tensor(vllm_out).float() + + # Print first few components for debugging + print(f"\n=== Text {i}: '{test_texts[i][:30]}...' ===") + print(f"HF shape: {hf_emb.shape}, vLLM shape: {vllm_emb.shape}") + print(f"HF first token, first 10 dims: {hf_emb[0, :10].tolist()}") + print(f"vLLM first token, first 10 dims: {vllm_emb[0, :10].tolist()}") + print(f"HF last token, first 10 dims: {hf_emb[-1, :10].tolist()}") + print(f"vLLM last token, first 10 dims: {vllm_emb[-1, :10].tolist()}") + + # Should have same shape + assert hf_emb.shape == vllm_emb.shape, ( + f"Shape mismatch for text {i}: HF {hf_emb.shape} vs vLLM {vllm_emb.shape}" + ) + + # Should have same values (with tolerance for fp16) + torch.testing.assert_close( + vllm_emb, + hf_emb, + rtol=1e-2, + atol=1e-2, + msg=f"Embedding mismatch for text {i}", + ) diff --git a/tests/models/registry.py b/tests/models/registry.py index cbd07cbc1c7..ffa4f52f138 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -520,6 +520,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { _EMBEDDING_EXAMPLE_MODELS = { # [Text-only] "BertModel": _HfExamplesInfo("BAAI/bge-base-en-v1.5"), + "HF_ColBERT": _HfExamplesInfo("answerdotai/answerai-colbert-small-v1"), "BgeM3EmbeddingModel": _HfExamplesInfo("BAAI/bge-m3"), "Gemma2Model": _HfExamplesInfo("BAAI/bge-multilingual-gemma2"), "Gemma3TextModel": _HfExamplesInfo("google/embeddinggemma-300m"), diff --git a/vllm/config/model.py b/vllm/config/model.py index 2686df4c23e..86b48418180 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -1411,6 +1411,11 @@ class ModelConfig: self._model_info.supports_cross_encoding or self.convert_type == "classify" ) + @property + def is_late_interaction(self) -> bool: + """Check if model uses late interaction (ColBERT-style) scoring.""" + return self._model_info.supports_late_interaction + @property def is_pp_supported(self) -> bool: return self._model_info.supports_pp diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index 24545de19cb..435ccbee6c7 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -44,6 +44,7 @@ from vllm.entrypoints.pooling.score.utils import ( ScoreMultiModalParam, _cosine_similarity, compress_token_type_ids, + compute_maxsim_score, get_score_prompt, validate_score_input, ) @@ -1368,6 +1369,87 @@ class LLM: items = self.engine_class.validate_outputs(scores, PoolingRequestOutput) return [ScoringRequestOutput.from_base(item) for item in items] + def _late_interaction_score( + self, + data_1: list[ScoreData], + data_2: list[ScoreData], + *, + use_tqdm: bool | Callable[..., tqdm], + pooling_params: PoolingParams | None, + lora_request: list[LoRARequest] | LoRARequest | None, + tokenization_kwargs: dict[str, Any], + ) -> list[ScoringRequestOutput]: + """ + Late interaction scoring (ColBERT MaxSim). + + Encodes queries and documents into per-token embeddings, then computes + MaxSim: sum over query tokens of max similarity to any document token. + """ + from vllm.outputs import PoolingOutput + + tokenizer = self.get_tokenizer() + + # Extract text from ScoreData + text_1: list[str] = [] + for text in data_1: + if not isinstance(text, str): + raise NotImplementedError( + "Late interaction scores currently do not support multimodal input." + ) + text_1.append(text) + + text_2: list[str] = [] + for text in data_2: + if not isinstance(text, str): + raise NotImplementedError( + "Late interaction scores currently do not support multimodal input." + ) + text_2.append(text) + + encoded_output: list[PoolingRequestOutput] = self.encode( + text_1 + text_2, + use_tqdm=use_tqdm, + lora_request=lora_request, + pooling_params=pooling_params, + pooling_task="token_embed", + tokenization_kwargs=tokenization_kwargs, + ) + + encoded_output_1: list[PoolingRequestOutput] = encoded_output[0 : len(text_1)] + encoded_output_2: list[PoolingRequestOutput] = encoded_output[len(text_1) :] + + if len(encoded_output_1) == 1: + encoded_output_1 = encoded_output_1 * len(encoded_output_2) + + # Compute MaxSim scores + scores: list[PoolingRequestOutput] = [] + padding: list[int] = [] + if (pad_token_id := tokenizer.pad_token_id) is not None: + padding = [pad_token_id] + + for emb_1, emb_2 in zip(encoded_output_1, encoded_output_2): + # emb_1.outputs.data: [query_len, dim] + # emb_2.outputs.data: [doc_len, dim] + q_emb = emb_1.outputs.data + d_emb = emb_2.outputs.data + + maxsim_score = compute_maxsim_score(q_emb, d_emb) + + tokens = emb_1.prompt_token_ids + padding + emb_2.prompt_token_ids + + scores.append( + PoolingRequestOutput( + request_id=f"{emb_1.request_id}_{emb_2.request_id}", + outputs=PoolingOutput(data=maxsim_score), + prompt_token_ids=tokens, + num_cached_tokens=emb_1.num_cached_tokens + emb_2.num_cached_tokens, + finished=True, + ) + ) + + items = self.engine_class.validate_outputs(scores, PoolingRequestOutput) + return [ScoringRequestOutput.from_base(item) for item in items] + def _cross_encoding_score( self, data_1: list[ScoreData], @@ -1497,7 +1579,11 @@ class LLM: ) supported_tasks = self.supported_tasks - if all(t not in supported_tasks for t in ("embed", "classify")): + # Late interaction models (e.g., ColBERT) use token_embed for scoring + is_late_interaction = model_config.is_late_interaction + if not is_late_interaction and all( + t not in supported_tasks for t in ("embed", "classify") + ): raise ValueError( "Score API is not supported by this model. " "Try converting the model using " @@ -1538,6 +1624,15 @@ class LLM: tokenization_kwargs=encode_kwargs, score_template=chat_template, ) + elif is_late_interaction: + return self._late_interaction_score( + score_data_1, + score_data_2, + use_tqdm=use_tqdm, + pooling_params=pooling_params, + lora_request=lora_request, + tokenization_kwargs=encode_kwargs, + ) else: return self._embedding_score( score_data_1, diff --git a/vllm/entrypoints/pooling/__init__.py b/vllm/entrypoints/pooling/__init__.py index 737f1efe895..4321e19f94c 100644 --- a/vllm/entrypoints/pooling/__init__.py +++ b/vllm/entrypoints/pooling/__init__.py @@ -37,7 +37,11 @@ def register_pooling_api_routers( app.include_router(embed_router) - if "score" in supported_tasks or "embed" in supported_tasks: + # Score/rerank endpoints are available for: + # - "score" task (cross-encoder models) + # - "embed" task (bi-encoder models) + # - "token_embed" task (late interaction models like ColBERT) + if any(t in supported_tasks for t in ("score", "embed", "token_embed")): from vllm.entrypoints.pooling.score.api_router import router as score_router app.include_router(score_router) @@ -101,6 +105,10 @@ def init_pooling_state( if "classify" in supported_tasks else None ) + # ServingScores handles score/rerank for: + # - "score" task (cross-encoder models) + # - "embed" task (bi-encoder models) + # - "token_embed" task (late interaction models like ColBERT) state.openai_serving_scores = ( ServingScores( engine_client, @@ -109,6 +117,6 @@ def init_pooling_state( score_template=resolved_chat_template, log_error_stack=args.log_error_stack, ) - if ("embed" in supported_tasks or "score" in supported_tasks) + if any(t in supported_tasks for t in ("embed", "score", "token_embed")) else None ) diff --git a/vllm/entrypoints/pooling/score/serving.py b/vllm/entrypoints/pooling/score/serving.py index c32f5470d45..9ef3b9afffb 100644 --- a/vllm/entrypoints/pooling/score/serving.py +++ b/vllm/entrypoints/pooling/score/serving.py @@ -31,6 +31,7 @@ from vllm.entrypoints.pooling.score.utils import ( ScoreInputs, _cosine_similarity, compress_token_type_ids, + compute_maxsim_score, get_score_prompt, validate_score_input, ) @@ -68,9 +69,12 @@ class ServingScores(OpenAIServing): self.is_cross_encoder = self.model_config.is_cross_encoder self.is_multimodal_model = self.model_config.is_multimodal_model self.architecture = self.model_config.architecture + self.is_late_interaction = self.model_config.is_late_interaction if self.is_cross_encoder: self._score_func = self._cross_encoding_score + elif self.is_late_interaction: + self._score_func = self._late_interaction_score else: self._score_func = self._embedding_score @@ -172,6 +176,142 @@ class ServingScores(OpenAIServing): return final_res_batch + async def _late_interaction_score( + self, + data_1: list[ScoreData], + data_2: list[ScoreData], + request: RerankRequest | ScoreRequest, + request_id: str, + lora_request: LoRARequest | None = None, + trace_headers: Mapping[str, str] | None = None, + ) -> list[PoolingRequestOutput] | ErrorResponse: + """ + Late interaction scoring (ColBERT MaxSim). + + Encodes queries and documents into per-token embeddings, then computes + MaxSim: sum over query tokens of max similarity to any document token. + """ + input_texts: list[str] = [] + for text in data_1 + data_2: + if not isinstance(text, str): + raise NotImplementedError( + "Late interaction scores currently do not support multimodal input." + ) + input_texts.append(text) + + model_config = self.model_config + tokenizer = self.renderer.get_tokenizer() + + encode_async = make_async( + tokenizer.encode, + executor=self._tokenizer_executor, + ) + + tokenization_kwargs = request.build_tok_params(model_config).get_encode_kwargs() + tokenized_prompts = await asyncio.gather( + *(encode_async(t, **tokenization_kwargs) for t in input_texts) + ) + + engine_prompts: list[TokensPrompt] = [] + for tok_result, input_text in zip(tokenized_prompts, input_texts): + text_token_prompt = self._validate_input(request, tok_result, input_text) + + engine_prompts.append( + TokensPrompt(prompt_token_ids=text_token_prompt["prompt_token_ids"]) + ) + + # Schedule the request and get the result generator. + generators: list[AsyncGenerator[PoolingRequestOutput, None]] = [] + + # Use token_embed task for late interaction models + from vllm import PoolingParams + + pooling_params = PoolingParams( + task="token_embed", + truncate_prompt_tokens=request.truncate_prompt_tokens, + use_activation=request.use_activation, + ) + + try: + pooling_params.verify("token_embed", self.model_config) + except ValueError as e: + return self.create_error_response(str(e)) + + for i, engine_prompt in enumerate(engine_prompts): + request_id_item = f"{request_id}-{i}" + + self._log_inputs( + request_id_item, + input_texts[i], + params=pooling_params, + lora_request=lora_request, + ) + + generators.append( + self.engine_client.encode( + engine_prompt, + pooling_params, + request_id_item, + lora_request=lora_request, + trace_headers=trace_headers, + priority=request.priority, + ) + ) + + result_generator = merge_async_iterators(*generators) + + # Collect token embeddings + embeddings: list[PoolingRequestOutput | None] = [None] * len(engine_prompts) + + async for i, res in result_generator: + embeddings[i] = res + + # Split into query and document embeddings + emb_data_1: list[PoolingRequestOutput] = [] + emb_data_2: list[PoolingRequestOutput] = [] + + for i in range(0, len(data_1)): + assert (emb := embeddings[i]) is not None + emb_data_1.append(emb) + + for i in range(len(data_1), len(embeddings)): + assert (emb := embeddings[i]) is not None + emb_data_2.append(emb) + + # Expand queries if 1:N scoring + if len(emb_data_1) == 1: + emb_data_1 = emb_data_1 * len(emb_data_2) + + # Compute MaxSim scores + from vllm.outputs import PoolingOutput + + scores: list[PoolingRequestOutput] = [] + padding: list[int] = [] + if (pad_token_id := tokenizer.pad_token_id) is not None: + padding = [pad_token_id] + + for emb_1, emb_2 in zip(emb_data_1, emb_data_2): + # emb_1.outputs.data: [query_len, dim] + # emb_2.outputs.data: [doc_len, dim] + q_emb = emb_1.outputs.data + d_emb = emb_2.outputs.data + + maxsim_score = compute_maxsim_score(q_emb, d_emb) + + tokens = emb_1.prompt_token_ids + padding + emb_2.prompt_token_ids + + scores.append( + PoolingRequestOutput( + request_id=f"{emb_1.request_id}_{emb_2.request_id}", + outputs=PoolingOutput(data=maxsim_score), + prompt_token_ids=tokens, + num_cached_tokens=emb_1.num_cached_tokens + emb_2.num_cached_tokens, + finished=True, + ) + ) + + return scores + async def _cross_encoding_score( self, data_1: list[ScoreData], diff --git a/vllm/entrypoints/pooling/score/utils.py b/vllm/entrypoints/pooling/score/utils.py index bf3bfe8a878..7d00f42f5df 100644 --- a/vllm/entrypoints/pooling/score/utils.py +++ b/vllm/entrypoints/pooling/score/utils.py @@ -3,6 +3,7 @@ from collections.abc import Iterable from typing import Any, TypeAlias, cast +import torch from torch.nn import CosineSimilarity from typing_extensions import Required, TypedDict @@ -34,6 +35,23 @@ ScoreContentPartParam: TypeAlias = ( ) +def compute_maxsim_score(q_emb: torch.Tensor, d_emb: torch.Tensor) -> torch.Tensor: + """ + Compute ColBERT MaxSim score. + + Args: + q_emb: Query token embeddings [query_len, dim] + d_emb: Document token embeddings [doc_len, dim] + + Returns: + MaxSim score (sum over query tokens of max similarity to any doc token) + """ + # [query_len, doc_len] + token_scores = torch.matmul(q_emb, d_emb.T) + # Max over document tokens, sum over query tokens + return token_scores.amax(dim=-1).sum() + + class ScoreMultiModalParam(TypedDict, total=False): """ A specialized parameter type for scoring multimodal content diff --git a/vllm/model_executor/models/colbert.py b/vllm/model_executor/models/colbert.py new file mode 100644 index 00000000000..dbb160556d3 --- /dev/null +++ b/vllm/model_executor/models/colbert.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +ColBERT late interaction model for retrieval and reranking. + +ColBERT uses per-token embeddings and late interaction (MaxSim) scoring +instead of single-vector representations or cross-encoder concatenation. + +Reference: https://arxiv.org/abs/2004.12832 +""" + +from collections.abc import Iterable +from typing import ClassVar, Literal + +import torch +from torch import nn + +from vllm.config import PoolerConfig, VllmConfig +from vllm.model_executor.layers.pooler import Pooler +from vllm.model_executor.layers.pooler.tokwise import pooler_for_token_embed + +from .bert import BertEmbeddingModel, BertModel +from .interfaces_base import default_pooling_type + + +@default_pooling_type(seq_pooling_type="CLS", tok_pooling_type="ALL") +class ColBERTModel(BertEmbeddingModel): + """ColBERT late interaction model for retrieval/reranking. + + This model extends BertEmbeddingModel with a ColBERT-style linear + projection layer for per-token embeddings. It supports only: + - "token_embed" task: Per-token embeddings for late interaction + + ColBERT is fundamentally a per-token embedding model - the linear + projection is trained for per-token representations, not for CLS + pooling. Use a dedicated dense embedding model if you need single- + vector representations. + + The ColBERT scoring (MaxSim) is computed externally, either client-side + or via the late interaction scoring path in ServingScores. + + Attributes: + colbert_linear: Linear projection from hidden_size to colbert_dim + supports_late_interaction: Flag indicating this model uses late + interaction scoring + """ + + # Mark this model as supporting late interaction scoring + supports_late_interaction: ClassVar[Literal[True]] = True + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + # Get config before calling super().__init__ + config = vllm_config.model_config.hf_config + self.hidden_size = config.hidden_size + self.head_dtype = vllm_config.model_config.head_dtype + + # ColBERT dimension - check various config field names used by different + # ColBERT implementations. If not found in config, will be inferred + # from loaded weights in load_weights() + self.colbert_dim: int | None = ( + getattr(config, "colbert_dim", None) + or getattr(config, "dim", None) + or getattr(config, "projection_dim", None) + ) + + # Initialize parent (this will call _build_pooler) + super().__init__(vllm_config=vllm_config, prefix=prefix) + + def _build_model(self, vllm_config: VllmConfig, prefix: str = "") -> BertModel: + return BertModel(vllm_config=vllm_config, prefix=prefix) + + def _build_colbert_linear(self) -> nn.Linear: + """Build the ColBERT linear projection layer.""" + if self.colbert_dim is None: + raise ValueError("colbert_dim must be set before building the linear layer") + return nn.Linear( + self.hidden_size, + self.colbert_dim, + bias=False, + dtype=self.head_dtype, + ) + + def _build_pooler(self, pooler_config: PoolerConfig) -> Pooler: + # ColBERT linear projection: hidden_size -> colbert_dim + # Original ColBERT uses bias=False + # If colbert_dim is not set from config, it will be inferred during + # load_weights and the linear layer will be created there + if self.colbert_dim is not None: + self.colbert_linear = self._build_colbert_linear() + else: + # Placeholder - will be created when weights are loaded + self.colbert_linear = None + + # ColBERT only supports token_embed - it's fundamentally a per-token + # embedding model. + return pooler_for_token_embed( + pooler_config, + projector=self.colbert_linear, + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + def _strip(name: str) -> str: + for p in ("model.", "bert."): + if name.startswith(p): + name = name[len(p) :] + return name + + weights_list = list(weights) + model_side: list[tuple[str, torch.Tensor]] = [] + colbert_side: list[tuple[str, torch.Tensor]] = [] + + for name, weight in weights_list: + stripped = _strip(name) + # Handle different checkpoint naming conventions for ColBERT linear + if stripped in ("linear.weight", "colbert_linear.weight"): + colbert_side.append(("colbert_linear.weight", weight)) + elif stripped.startswith("linear.") or stripped.startswith( + "colbert_linear." + ): + new_name = stripped.replace("linear.", "colbert_linear.") + colbert_side.append((new_name, weight)) + else: + model_side.append((stripped, weight)) + + # Load base BERT weights using BertModel.load_weights which handles QKV fusion + loaded: set[str] = set() + loaded_model = self.model.load_weights(model_side) + loaded.update({"model." + n for n in loaded_model}) + + # Load ColBERT linear weights + if colbert_side: + for name, weight in colbert_side: + if name == "colbert_linear.weight": + # Infer colbert_dim from weights if not set in config + if self.colbert_dim is None: + # Weight shape is [colbert_dim, hidden_size] + self.colbert_dim = weight.shape[0] + # Create the linear layer now that we know the dimension + self.colbert_linear = self._build_colbert_linear() + # Move to the same device as the model's existing parameters + device = next(self.model.parameters()).device + self.colbert_linear.to(device) + # Update the pooler's projector to use the new linear layer + self.pooler.head.projector = self.colbert_linear + + # Load weights directly into the pooler's projector + weight = weight.to(self.pooler.head.projector.weight.device) + self.pooler.head.projector.weight.data.copy_(weight) + loaded.add("pooler.head.projector.weight") + break + + return loaded diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index c97a9faf6c5..2c3ca1a5022 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -981,6 +981,40 @@ def supports_cross_encoding( return is_pooling_model(model) and _supports_cross_encoding(model) +@runtime_checkable +class SupportsLateInteraction(Protocol): + """The interface required for all models that support late interaction. + + Late interaction models (like ColBERT) encode queries and documents + separately into per-token embeddings, then compute similarity via + MaxSim (max over document tokens, sum over query tokens). + """ + + supports_late_interaction: ClassVar[Literal[True]] = True + + +@overload +def supports_late_interaction( + model: type[object], +) -> TypeIs[type[SupportsLateInteraction]]: ... + + +@overload +def supports_late_interaction(model: object) -> TypeIs[SupportsLateInteraction]: ... + + +def _supports_late_interaction( + model: type[object] | object, +) -> TypeIs[type[SupportsLateInteraction]] | TypeIs[SupportsLateInteraction]: + return getattr(model, "supports_late_interaction", False) + + +def supports_late_interaction( + model: type[object] | object, +) -> TypeIs[type[SupportsLateInteraction]] | TypeIs[SupportsLateInteraction]: + return is_pooling_model(model) and _supports_late_interaction(model) + + class SupportsQuant: """The interface required for all models that support quantization.""" diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 5eeb32ed96d..830a615ce0e 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -49,6 +49,7 @@ from .interfaces import ( is_hybrid, requires_raw_input_tokens, supports_cross_encoding, + supports_late_interaction, supports_mamba_prefix_caching, supports_multimodal, supports_multimodal_encoder_tp_data, @@ -205,6 +206,7 @@ _EMBEDDING_MODELS = { # [Text-only] "BertModel": ("bert", "BertEmbeddingModel"), "BertSpladeSparseEmbeddingModel": ("bert", "BertSpladeSparseEmbeddingModel"), + "HF_ColBERT": ("colbert", "ColBERTModel"), "DeciLMForCausalLM": ("nemotron_nas", "DeciLMForCausalLM"), "Gemma2Model": ("gemma2", "Gemma2ForCausalLM"), "Gemma3TextModel": ("gemma3", "Gemma3Model"), @@ -593,6 +595,7 @@ class _ModelInfo: default_seq_pooling_type: SequencePoolingType default_tok_pooling_type: TokenPoolingType supports_cross_encoding: bool + supports_late_interaction: bool supports_multimodal: bool supports_multimodal_raw_input_only: bool requires_raw_input_tokens: bool @@ -616,6 +619,7 @@ class _ModelInfo: default_tok_pooling_type=get_default_tok_pooling_type(model), attn_type=get_attn_type(model), supports_cross_encoding=supports_cross_encoding(model), + supports_late_interaction=supports_late_interaction(model), supports_multimodal=supports_multimodal(model), supports_multimodal_raw_input_only=supports_multimodal_raw_input_only( model From 4d9513537d00a9b6678a2b1ed3c3566a81f7dd77 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Luka=20Govedi=C4=8D?= Date: Wed, 4 Feb 2026 19:09:03 -0500 Subject: [PATCH 075/810] [CI][torch.compile] Reduce e2e fusion test time (#33293) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Luka Govedič Signed-off-by: ProExpertProg Signed-off-by: Luka Govedič --- .buildkite/test-amd.yaml | 45 +-- .buildkite/test-pipeline.yaml | 81 +---- .buildkite/test_areas/compile.yaml | 198 +++++++++-- .buildkite/test_areas/distributed.yaml | 59 +--- .buildkite/test_areas/pytorch.yaml | 7 +- tests/compile/distributed/test_fusions_e2e.py | 321 ------------------ tests/compile/fusion_test_utils.py | 208 ------------ tests/compile/fusions_e2e/__init__.py | 0 tests/compile/fusions_e2e/common.py | 102 ++++++ tests/compile/fusions_e2e/conftest.py | 158 +++++++++ tests/compile/fusions_e2e/models.py | 112 ++++++ tests/compile/fusions_e2e/test_tp1_quant.py | 146 ++++++++ tests/compile/fusions_e2e/test_tp2_ar_rms.py | 199 +++++++++++ .../compile/fusions_e2e/test_tp2_async_tp.py | 143 ++++++++ tests/compile/test_fusion_attn.py | 99 ------ tests/test_config.py | 4 +- vllm/config/vllm.py | 7 +- 17 files changed, 1068 insertions(+), 821 deletions(-) delete mode 100644 tests/compile/distributed/test_fusions_e2e.py delete mode 100644 tests/compile/fusion_test_utils.py create mode 100644 tests/compile/fusions_e2e/__init__.py create mode 100644 tests/compile/fusions_e2e/common.py create mode 100644 tests/compile/fusions_e2e/conftest.py create mode 100644 tests/compile/fusions_e2e/models.py create mode 100644 tests/compile/fusions_e2e/test_tp1_quant.py create mode 100644 tests/compile/fusions_e2e/test_tp2_ar_rms.py create mode 100644 tests/compile/fusions_e2e/test_tp2_async_tp.py diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index ee7c6ab0a5d..0050c615a4b 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -604,9 +604,11 @@ steps: - tests/compile commands: - pytest -v -s compile/fullgraph/test_full_graph.py -k 'not test_fp8_kv_scale_compile' - # Limit to no custom ops to reduce running time - # Wrap with quotes to escape yaml and avoid starting -k string with a - - - "pytest -v -s compile/distributed/test_fusions_e2e.py -k 'TRITON and not +quant_fp8 and not Llama-4'" + # # Limit to no custom ops to reduce running time + # # Wrap with quotes to escape yaml and avoid starting -k string with a - + # - "pytest -v -s compile/distributed/test_fusions_e2e.py -k 'TRITON and not +quant_fp8 and not Llama-4'" + # Old E2E tests were removed in https://github.com/vllm-project/vllm/pull/33293 + # in favor of new tests in fusions_e2e. We avoid replicating the new jobs in this file as it's deprecated. - label: Cudagraph test timeout_in_minutes: 20 @@ -1181,7 +1183,6 @@ steps: - tests/compile/test_fusion_attn.py - tests/compile/test_silu_mul_quant_fusion.py - tests/compile/distributed/test_fusion_all_reduce.py - - tests/compile/distributed/test_fusions_e2e.py - tests/compile/fullgraph/test_full_graph.py commands: - nvidia-smi @@ -1189,33 +1190,16 @@ steps: - pytest -v -s tests/compile/test_silu_mul_quant_fusion.py # this runner has 2 GPUs available even though num_gpus=2 is not set - pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py - # Limit to Inductor partition, no custom ops, and allreduce & attn fusion to reduce running time - # Wrap with quotes to escape yaml - - "pytest -v -s tests/compile/distributed/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm -k 'True and not +quant_fp8 and not +rms_norm'" + + # # Limit to Inductor partition, no custom ops, and allreduce & attn fusion to reduce running time + # # Wrap with quotes to escape yaml + # - "pytest -v -s tests/compile/distributed/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm -k 'True and not +quant_fp8 and not +rms_norm'" + # Old E2E tests were removed in https://github.com/vllm-project/vllm/pull/33293 + # in favor of new tests in fusions_e2e. We avoid replicating the new jobs in this file as it's deprecated. + # test_fp8_kv_scale_compile requires FlashAttention (not supported on default L4/L40) - pytest -v -s tests/compile/fullgraph/test_full_graph.py::test_fp8_kv_scale_compile -- label: Blackwell Fusion E2E Tests # 30 min - timeout_in_minutes: 40 - working_dir: "/vllm-workspace/" - gpu: b200 - optional: true - num_gpus: 2 - source_file_dependencies: - - csrc/quantization/fp4/ - - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py - - vllm/v1/attention/backends/flashinfer.py - - vllm/compilation/ - # can affect pattern matching - - vllm/model_executor/layers/layernorm.py - - vllm/model_executor/layers/activation.py - - vllm/model_executor/layers/quantization/input_quant_fp8.py - - tests/compile/distributed/test_fusions_e2e.py - commands: - - nvidia-smi - # Run all e2e fusion tests - - pytest -v -s tests/compile/distributed/test_fusions_e2e.py - - label: Blackwell GPT-OSS Eval timeout_in_minutes: 60 working_dir: "/vllm-workspace/" @@ -1566,7 +1550,10 @@ steps: - pytest -v -s tests/compile/distributed/test_sequence_parallelism.py - pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py #- pytest -v -s tests/compile/distributed/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm - - "VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/distributed/test_fusions_e2e.py -k 'not Llama-4'" + # - "VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/distributed/test_fusions_e2e.py -k 'not Llama-4'" + # Old E2E tests were removed in https://github.com/vllm-project/vllm/pull/33293 + # in favor of new tests in fusions_e2e. We avoid replicating the new jobs in this file as it's deprecated. + - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/distributed/test_sequence_parallel.py - pytest -v -s tests/distributed/test_context_parallel.py - HIP_VISIBLE_DEVICES=0,1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=allgather_reducescatter --disable-nccl-for-dp-synchronization diff --git a/.buildkite/test-pipeline.yaml b/.buildkite/test-pipeline.yaml index bcd9997a48f..554081f5353 100644 --- a/.buildkite/test-pipeline.yaml +++ b/.buildkite/test-pipeline.yaml @@ -537,9 +537,11 @@ steps: commands: # fp8 kv scales not supported on sm89, tested on Blackwell instead - pytest -v -s compile/fullgraph/test_full_graph.py -k 'not test_fp8_kv_scale_compile' - # Limit to no custom ops to reduce running time - # Wrap with quotes to escape yaml and avoid starting -k string with a - - - "pytest -v -s compile/distributed/test_fusions_e2e.py -k 'TRITON and not +quant_fp8 and not Llama-4'" + # # Limit to no custom ops to reduce running time + # # Wrap with quotes to escape yaml and avoid starting -k string with a - + # - "pytest -v -s compile/distributed/test_fusions_e2e.py -k 'TRITON and not +quant_fp8 and not Llama-4'" + # Old E2E tests were removed in https://github.com/vllm-project/vllm/pull/33293 + # in favor of new tests in fusions_e2e. We avoid replicating the new jobs in this file as it's deprecated. - label: Cudagraph test timeout_in_minutes: 20 @@ -1069,7 +1071,6 @@ steps: - tests/compile/test_fusion_attn.py - tests/compile/test_silu_mul_quant_fusion.py - tests/compile/distributed/test_fusion_all_reduce.py - - tests/compile/distributed/test_fusions_e2e.py - tests/compile/fullgraph/test_full_graph.py commands: - nvidia-smi @@ -1077,75 +1078,15 @@ steps: - pytest -v -s tests/compile/test_silu_mul_quant_fusion.py # this runner has 2 GPUs available even though num_gpus=2 is not set - pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py - # Limit to Inductor partition, no custom ops, and allreduce & attn fusion to reduce running time - # Wrap with quotes to escape yaml - - "pytest -v -s tests/compile/distributed/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm -k 'True and not +quant_fp8 and not +rms_norm'" + # # Limit to Inductor partition, no custom ops, and allreduce & attn fusion to reduce running time + # # Wrap with quotes to escape yaml + # - "pytest -v -s tests/compile/distributed/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm -k 'True and not +quant_fp8 and not +rms_norm'" + # Old E2E tests were removed in https://github.com/vllm-project/vllm/pull/33293 + # in favor of new tests in fusions_e2e. We avoid replicating the new jobs in this file as it's deprecated. + # test_fp8_kv_scale_compile requires FlashAttention (not supported on default L4/L40) - pytest -v -s tests/compile/fullgraph/test_full_graph.py::test_fp8_kv_scale_compile -- label: Blackwell Fusion E2E Tests # 30 min - timeout_in_minutes: 40 - working_dir: "/vllm-workspace/" - gpu: b200 - optional: true - num_gpus: 2 - source_file_dependencies: - - csrc/quantization/fp4/ - - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py - - vllm/v1/attention/backends/flashinfer.py - - vllm/compilation/ - # can affect pattern matching - - vllm/model_executor/layers/layernorm.py - - vllm/model_executor/layers/activation.py - - vllm/model_executor/layers/quantization/input_quant_fp8.py - - tests/compile/distributed/test_fusions_e2e.py - commands: - - nvidia-smi - # Run all e2e fusion tests - - pytest -v -s tests/compile/distributed/test_fusions_e2e.py - -- label: Hopper Fusion E2E Tests (H100) # 10min - timeout_in_minutes: 70 - working_dir: "/vllm-workspace/" - gpu: h100 - optional: true - source_file_dependencies: - - csrc/quantization/fp4/ - - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py - - vllm/v1/attention/backends/flashinfer.py - - vllm/compilation/ - # can affect pattern matching - - vllm/model_executor/layers/layernorm.py - - vllm/model_executor/layers/activation.py - - vllm/model_executor/layers/quantization/input_quant_fp8.py - - tests/compile/test_fusion_attn.py - commands: - - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - # skip Llama-4 since it does not fit on this device - - pytest -v -s tests/compile/test_fusion_attn.py -k 'not Llama-4' - -- label: Hopper Fusion Distributed E2E Tests (2xH100) # 70min - timeout_in_minutes: 70 - working_dir: "/vllm-workspace/" - gpu: h100 - optional: true - num_gpus: 2 - source_file_dependencies: - - csrc/quantization/fp4/ - - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py - - vllm/v1/attention/backends/flashinfer.py - - vllm/compilation/ - # can affect pattern matching - - vllm/model_executor/layers/layernorm.py - - vllm/model_executor/layers/activation.py - - vllm/model_executor/layers/quantization/input_quant_fp8.py - - tests/compile/distributed/test_fusions_e2e.py - commands: - - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - # Run all e2e fusion tests - - pytest -v -s tests/compile/distributed/test_fusions_e2e.py -k 'not Llama-4' - - pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py - - label: Blackwell GPT-OSS Eval timeout_in_minutes: 60 working_dir: "/vllm-workspace/" diff --git a/.buildkite/test_areas/compile.yaml b/.buildkite/test_areas/compile.yaml index 3c6f82fdd56..e8cf9e8bd28 100644 --- a/.buildkite/test_areas/compile.yaml +++ b/.buildkite/test_areas/compile.yaml @@ -2,56 +2,196 @@ group: Compile depends_on: - image-build steps: -- label: Fusion and Compile Tests (B200) +- label: Sequence Parallel Tests (2 GPUs) + timeout_in_minutes: 50 + working_dir: "/vllm-workspace/" + num_devices: 2 + source_file_dependencies: + - vllm/model_executor/layers/ + - vllm/compilation/ + - vllm/v1/worker/ + - vllm/v1/cudagraph_dispatcher.py + - tests/distributed/test_sequence_parallel.py + commands: + - export VLLM_TEST_CLEAN_GPU_MEMORY=1 + - pytest -v -s tests/distributed/test_sequence_parallel.py + +- label: Sequence Parallel Tests (2xH100) + timeout_in_minutes: 50 + working_dir: "/vllm-workspace/" + device: h100 + optional: true + num_devices: 2 + commands: + - export VLLM_TEST_CLEAN_GPU_MEMORY=1 + - pytest -v -s tests/distributed/test_sequence_parallel.py + +- label: Distributed Compile Unit Tests (2xH100) timeout_in_minutes: 40 working_dir: "/vllm-workspace/" + device: h100 + num_devices: 2 + source_file_dependencies: + - vllm/compilation/ + - vllm/model_executor/layers + - tests/compile/distributed/test_fusion_all_reduce.py + - tests/compile/distributed/test_sequence_parallelism.py + - tests/compile/distributed/test_async_tp.py + commands: + - export VLLM_TEST_CLEAN_GPU_MEMORY=1 + - pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py + - pytest -v -s tests/compile/distributed/test_sequence_parallelism.py + - pytest -v -s tests/compile/distributed/test_async_tp.py + +- label: Fusion and Compile Unit Tests (B200) + timeout_in_minutes: 20 + working_dir: "/vllm-workspace/" device: b200 source_file_dependencies: - csrc/quantization/fp4/ - - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py - - vllm/v1/attention/backends/flashinfer.py - - vllm/v1/worker/ - - vllm/v1/cudagraph_dispatcher.py - - vllm/compilation/ - # can affect pattern matching + - vllm/model_executor/layers/quantization/ - vllm/model_executor/layers/layernorm.py - vllm/model_executor/layers/activation.py - - vllm/model_executor/layers/quantization/input_quant_fp8.py + - vllm/model_executor/layers/attention/attention.py + - vllm/v1/attention/backends/flashinfer.py + - vllm/compilation/ # TODO(luka) limit to vllm/compilation/passes - tests/compile/test_fusion_attn.py - tests/compile/test_silu_mul_quant_fusion.py - tests/compile/distributed/test_fusion_all_reduce.py - - tests/compile/distributed/test_fusions_e2e.py - tests/compile/fullgraph/test_full_graph.py commands: + # b200 runners are limited, so we limit the tests to the minimum set only supported on Blackwell - nvidia-smi - - pytest -v -s tests/compile/test_fusion_attn.py + - pytest -v -s tests/compile/test_fusion_attn.py -k FLASHINFER - pytest -v -s tests/compile/test_silu_mul_quant_fusion.py # this runner has 2 GPUs available even though num_devices=2 is not set - pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py - # Limit to Inductor partition, no custom ops, and allreduce & attn fusion to reduce running time - # Wrap with quotes to escape yaml - - "pytest -v -s tests/compile/distributed/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm -k 'True and not +quant_fp8 and not +rms_norm'" # test_fp8_kv_scale_compile requires FlashAttention (not supported on default L4/L40) + # TODO(luka) move to H100 once pass tests run on H100 - pytest -v -s tests/compile/fullgraph/test_full_graph.py::test_fp8_kv_scale_compile -- label: Fusion E2E (2 GPUs)(B200) - timeout_in_minutes: 40 +- label: Fusion E2E Quick (H100) + timeout_in_minutes: 15 working_dir: "/vllm-workspace/" - device: b200 - optional: true - num_devices: 2 + device: h100 + num_devices: 1 source_file_dependencies: - - csrc/quantization/fp4/ - - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py - - vllm/v1/attention/backends/flashinfer.py - - vllm/compilation/ - # can affect pattern matching - - vllm/model_executor/layers/layernorm.py - - vllm/model_executor/layers/activation.py - - vllm/model_executor/layers/quantization/input_quant_fp8.py - - tests/compile/distributed/test_fusions_e2e.py + - csrc/quantization/ + - vllm/model_executor/ + - vllm/v1/attention/ + - vllm/compilation/ + - tests/compile/fusions_e2e/ commands: - nvidia-smi - # Run all e2e fusion tests - - pytest -v -s tests/compile/distributed/test_fusions_e2e.py + # Run all models and attn backends but only Inductor partition and native custom ops + - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "inductor_partition and not +rms_norm and not +quant_fp8" + # Qwen 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" +- label: Fusion E2E Config Sweep (H100) + timeout_in_minutes: 30 + working_dir: "/vllm-workspace/" + device: h100 + num_devices: 1 + source_file_dependencies: + - csrc/quantization/ + - vllm/compilation/ + # can affect pattern matching + - vllm/model_executor/layers/layernorm.py + - vllm/model_executor/layers/activation.py + - vllm/model_executor/layers/attention/attention.py + - vllm/model_executor/layers/quantization/input_quant_fp8.py + - tests/compile/fusions_e2e/ + commands: + - nvidia-smi + # 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) + timeout_in_minutes: 30 + working_dir: "/vllm-workspace/" + device: b200 + num_devices: 1 + optional: true + commands: + - nvidia-smi + # Run all models and attn backends but only Inductor partition and native custom ops + # -k "inductor_partition and not +rms_norm and not +quant_fp8" + # Qwen requires +quant_fp8 as -quant_fp8 rms+quant fusion is not supported + # -k "inductor_partition and not +rms_norm and +quant_fp8 and qwen3" + # Run just llama3 (fp8 & fp4) for all config combinations + # -k "llama-3" + - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "inductor_partition and not +rms_norm and not +quant_fp8" -k "inductor_partition and not +rms_norm and +quant_fp8 and qwen3" -k "llama-3" + +- label: Fusion E2E TP2 Quick (H100) + timeout_in_minutes: 20 + working_dir: "/vllm-workspace/" + device: h100 + num_devices: 2 + source_file_dependencies: + - csrc/quantization/ + - vllm/model_executor/ + - vllm/v1/attention/ + - vllm/compilation/ + - tests/compile/fusions_e2e/ + commands: + - nvidia-smi + # Run all models and attn backends but only Inductor partition and native custom ops + - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "inductor_partition and not +rms_norm and not +quant_fp8" + - pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "inductor_partition and not +rms_norm and not +quant_fp8" + +- label: Fusion E2E TP2 AR-RMS Config Sweep (H100) + timeout_in_minutes: 40 + working_dir: "/vllm-workspace/" + device: h100 + num_devices: 2 + source_file_dependencies: + - csrc/quantization/ + - vllm/compilation/ + # can affect pattern matching + - vllm/model_executor/layers/layernorm.py + - vllm/model_executor/layers/activation.py + - vllm/model_executor/layers/attention/attention.py + - vllm/model_executor/layers/quantization/input_quant_fp8.py + - tests/compile/fusions_e2e/ + commands: + - nvidia-smi + # Run just llama3 (fp4 & 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) + timeout_in_minutes: 40 + working_dir: "/vllm-workspace/" + device: h100 + num_devices: 2 + source_file_dependencies: + - csrc/quantization/ + - vllm/compilation/ + # can affect pattern matching + - vllm/model_executor/layers/layernorm.py + - vllm/model_executor/layers/activation.py + - vllm/model_executor/layers/attention/attention.py + - vllm/model_executor/layers/quantization/input_quant_fp8.py + - tests/compile/fusions_e2e/ + commands: + - nvidia-smi + # 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) + timeout_in_minutes: 20 + working_dir: "/vllm-workspace/" + device: b200 + num_devices: 2 + source_file_dependencies: + - csrc/quantization/ + - vllm/model_executor/ + - vllm/v1/attention/ + - vllm/compilation/ + - tests/compile/fusions_e2e/ + commands: + - nvidia-smi + # Run all models and attn backends but only Inductor partition and native custom ops + # for ar-rms-quant-fp4, also sweep llama3 + - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "inductor_partition and not +rms_norm and not +quant_fp8" -k "Llama-3.1-8B-Instruct-FP4" + - pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "inductor_partition and not +rms_norm and not +quant_fp8" diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index 51e1de3f06c..ae4f45fbf4e 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -16,7 +16,7 @@ steps: - pytest -v -s distributed/test_shm_storage.py - label: Distributed (2 GPUs) - timeout_in_minutes: 90 + timeout_in_minutes: 60 working_dir: "/vllm-workspace/tests" num_devices: 2 source_file_dependencies: @@ -47,7 +47,6 @@ steps: - pytest -v -s ./compile/test_wrapper.py - VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed' - VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed' - - pytest -v -s distributed/test_sequence_parallel.py - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown - pytest -v -s v1/worker/test_worker_memory_snapshot.py @@ -133,25 +132,13 @@ steps: - TARGET_TEST_SUITE=A100 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' - pytest -v -s -x lora/test_mixtral.py -- label: Sequence Parallel Tests (H100) - timeout_in_minutes: 60 - working_dir: "/vllm-workspace/" - device: h100 - optional: true - num_devices: 2 - commands: - - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - # Run sequence parallel tests - - pytest -v -s tests/distributed/test_sequence_parallel.py - - pytest -v -s tests/compile/distributed/test_sequence_parallelism.py - - label: Distributed Tests (2 GPUs)(H100) + timeout_in_minutes: 15 device: h100 optional: true working_dir: "/vllm-workspace/" num_devices: 2 commands: - - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/distributed/test_async_tp.py - pytest -v -s tests/distributed/test_context_parallel.py - VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput - pytest -v -s tests/v1/distributed/test_dbo.py @@ -217,45 +204,3 @@ steps: commands: - pytest -v -s distributed/test_pp_cudagraph.py - pytest -v -s distributed/test_pipeline_parallel.py - -- label: Hopper Fusion E2E Tests (H100) - timeout_in_minutes: 70 - working_dir: "/vllm-workspace/" - device: h100 - optional: true - source_file_dependencies: - - csrc/quantization/fp4/ - - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py - - vllm/v1/attention/backends/flashinfer.py - - vllm/compilation/ - # can affect pattern matching - - vllm/model_executor/layers/layernorm.py - - vllm/model_executor/layers/activation.py - - vllm/model_executor/layers/quantization/input_quant_fp8.py - - tests/compile/test_fusion_attn.py - commands: - - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - # skip Llama-4 since it does not fit on this device - - pytest -v -s tests/compile/test_fusion_attn.py -k 'not Llama-4' - -- label: Hopper Fusion Distributed E2E Tests (2xH100) - timeout_in_minutes: 70 - working_dir: "/vllm-workspace/" - device: h100 - optional: true - num_devices: 2 - source_file_dependencies: - - csrc/quantization/fp4/ - - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py - - vllm/v1/attention/backends/flashinfer.py - - vllm/compilation/ - # can affect pattern matching - - vllm/model_executor/layers/layernorm.py - - vllm/model_executor/layers/activation.py - - vllm/model_executor/layers/quantization/input_quant_fp8.py - - tests/compile/distributed/test_fusions_e2e.py - commands: - - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - # Run all e2e fusion tests - - pytest -v -s tests/compile/distributed/test_fusions_e2e.py -k 'not Llama-4' - - pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py diff --git a/.buildkite/test_areas/pytorch.yaml b/.buildkite/test_areas/pytorch.yaml index 332d5202d83..1ac3eec58d9 100644 --- a/.buildkite/test_areas/pytorch.yaml +++ b/.buildkite/test_areas/pytorch.yaml @@ -18,7 +18,7 @@ steps: - "find compile/ -maxdepth 1 -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'" - label: PyTorch Fullgraph Smoke Test - timeout_in_minutes: 30 + timeout_in_minutes: 35 source_file_dependencies: - vllm/ - tests/compile @@ -30,16 +30,13 @@ steps: - "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_graph.py' -exec pytest -s -v {} \\;" - label: PyTorch Fullgraph - timeout_in_minutes: 40 + timeout_in_minutes: 30 source_file_dependencies: - vllm/ - tests/compile commands: # fp8 kv scales not supported on sm89, tested on Blackwell instead - pytest -v -s compile/fullgraph/test_full_graph.py -k 'not test_fp8_kv_scale_compile' - # Limit to no custom ops to reduce running time - # Wrap with quotes to escape yaml and avoid starting -k string with a - - - "pytest -v -s compile/distributed/test_fusions_e2e.py -k 'TRITON and not +quant_fp8 and not Llama-4'" - label: Pytorch Nightly Dependency Override Check # 2min # if this test fails, it means the nightly torch version is not compatible with some diff --git a/tests/compile/distributed/test_fusions_e2e.py b/tests/compile/distributed/test_fusions_e2e.py deleted file mode 100644 index b9913734dfd..00000000000 --- a/tests/compile/distributed/test_fusions_e2e.py +++ /dev/null @@ -1,321 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from __future__ import annotations - -import logging -from typing import Any - -import pytest -import regex as re - -from tests.compile.fusion_test_utils import ( - CUSTOM_OPS_FP8, - CUSTOM_OPS_QUANT_RMS_NORM, - CUSTOM_OPS_RMS_NORM, - MODELS, - MODELS_FP4, - MODELS_FP8, - MODELS_GROUP_FP8, - Matches, - custom_ops_product, - is_blackwell, - run_model, -) -from tests.v1.attention.utils import AttentionBackendEnum -from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode, PassConfig -from vllm.platforms import current_platform -from vllm.utils.flashinfer import has_flashinfer -from vllm.utils.torch_utils import is_torch_equal_or_newer - -from ...utils import flat_product, multi_gpu_test - - -@multi_gpu_test(num_gpus=2) -@pytest.mark.parametrize( - "model_name, model_kwargs, backend, matches, custom_ops", - # Toggle RMSNorm and QuantFP8 for FP8 models - list( - flat_product( - MODELS_FP8, custom_ops_product(CUSTOM_OPS_FP8, CUSTOM_OPS_RMS_NORM) - ) - ) - # Toggle RMSNorm for FP4 models and unquant models - + list(flat_product(MODELS_FP4 + MODELS, CUSTOM_OPS_RMS_NORM)), -) -@pytest.mark.parametrize("inductor_graph_partition", [True, False]) -@pytest.mark.skipif( - not current_platform.is_cuda() - or not has_flashinfer() - or not current_platform.has_device_capability(90), - reason="allreduce+rmsnorm fusion requires flashinfer", -) -def test_tp2_attn_quant_allreduce_rmsnorm( - model_name: str, - model_kwargs: dict, - backend: AttentionBackendEnum, - matches: Matches, - custom_ops: str, - inductor_graph_partition: bool, - caplog_mp_spawn, - monkeypatch, -): - if inductor_graph_partition and not is_torch_equal_or_newer("2.9.0.dev"): - pytest.skip("Inductor graph partition requires torch>=2.9") - - if "fp4" in model_name.lower() and not is_blackwell(): - pytest.skip("NVFP4 quant requires Blackwell") - - if backend == AttentionBackendEnum.FLASHINFER and not is_blackwell(): - # FlashInfer attn fusion requires Blackwell - matches = matches._replace(attention_fusion=0) - - custom_ops_list = custom_ops.split(",") if custom_ops else [] - - if inductor_graph_partition: - mode = CUDAGraphMode.FULL_AND_PIECEWISE - splitting_ops: list[str] | None = None - else: - mode = CUDAGraphMode.FULL_DECODE_ONLY - splitting_ops = [] - - # Disable, compile cache to make sure custom passes run. - # Otherwise, we can't verify fusion happened through the logs. - monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1") - - # To capture subprocess logs, we need to know whether spawn or fork is used. - # Force spawn as it is more general. - monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") - - model_kwargs["attention_config"] = {"backend": backend.name} - - compilation_config = CompilationConfig( - # Testing properties - use_inductor_graph_partition=inductor_graph_partition, - cudagraph_mode=mode, - custom_ops=custom_ops_list, - splitting_ops=splitting_ops, - # Common - mode=CompilationMode.VLLM_COMPILE, - pass_config=PassConfig( - fuse_attn_quant=True, - eliminate_noops=True, - fuse_allreduce_rms=True, - ), - # Inductor caches custom passes by default as well via uuid - inductor_compile_config={"force_disable_caches": True}, - ) - - with caplog_mp_spawn(logging.DEBUG) as log_holder: - run_model( - compilation_config, model_name, tensor_parallel_size=2, **model_kwargs - ) - log_matches = re.findall( - r"fusion_attn.py:\d+] Fused quant onto (\d+) attention nodes", - log_holder.text, - ) - # 2 for each compile range - # (global compile range can be split due to fuse_allreduce_rmsnorm) - num_compile_ranges = len(compilation_config.get_compile_ranges()) - assert num_compile_ranges in [1, 2] - - assert len(log_matches) == 2 * num_compile_ranges, log_holder.text - - assert all(int(log_match) == matches.attention_fusion for log_match in log_matches) - - log_matches = re.findall( - r"collective_fusion.py:\d+] Replaced (\d+) patterns", - log_holder.text, - ) - assert len(log_matches) == 2, log_holder.text - - assert int(log_matches[0]) == matches.allreduce_fusion - assert int(log_matches[1]) == matches.allreduce_fusion - - log_matches = re.findall( - r"pass_manager.py:\d+] Skipping .*AllReduceFusionPass.* with compile range", - log_holder.text, - ) - assert len(log_matches) == 2 * (num_compile_ranges - 1), log_holder.text - - -@multi_gpu_test(num_gpus=2) -@pytest.mark.parametrize( - "model_name, model_kwargs, backend, matches, custom_ops", - # Toggle RMSNorm and QuantFP8 for FP8 models - list( - flat_product( - MODELS_FP8, custom_ops_product(CUSTOM_OPS_FP8, CUSTOM_OPS_RMS_NORM) - ) - ) - # Toggle RMSNorm for FP4 models and unquant models - + list(flat_product(MODELS_FP4 + MODELS, CUSTOM_OPS_RMS_NORM)), -) -@pytest.mark.parametrize("inductor_graph_partition", [True, False]) -@pytest.mark.skipif( - not current_platform.is_cuda(), - reason="sequence parallel only tested on CUDA", -) -def test_tp2_attn_quant_async_tp( - model_name: str, - model_kwargs: dict, - backend: AttentionBackendEnum, - matches: Matches, - custom_ops: str, - inductor_graph_partition: bool, - caplog_mp_spawn, - monkeypatch, -): - if is_blackwell(): - # TODO: https://github.com/vllm-project/vllm/issues/27893 - pytest.skip("Blackwell is not supported for AsyncTP pass") - - if inductor_graph_partition and not is_torch_equal_or_newer("2.9.0.dev"): - pytest.skip("Inductor graph partition requires torch>=2.9") - - if "fp4" in model_name.lower() and not is_blackwell(): - pytest.skip("NVFP4 quant requires Blackwell") - - if backend == AttentionBackendEnum.FLASHINFER: - if not has_flashinfer(): - pytest.skip("FlashInfer backend requires flashinfer installed") - if not is_blackwell(): - # FlashInfer attn fusion requires Blackwell - matches = matches._replace(attention_fusion=0) - - custom_ops_list = custom_ops.split(",") if custom_ops else [] - - if inductor_graph_partition: - mode = CUDAGraphMode.FULL_AND_PIECEWISE - splitting_ops: list[str] | None = None - else: - mode = CUDAGraphMode.FULL_DECODE_ONLY - splitting_ops = [] - - # Disable, compile cache to make sure custom passes run. - # Otherwise, we can't verify fusion happened through the logs. - monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1") - - # To capture subprocess logs, we need to know whether spawn or fork is used. - # Force spawn as it is more general. - monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") - - model_kwargs["attention_config"] = {"backend": backend.name} - - compilation_config = CompilationConfig( - # Testing properties - use_inductor_graph_partition=inductor_graph_partition, - cudagraph_mode=mode, - custom_ops=custom_ops_list, - splitting_ops=splitting_ops, - # Common - mode=CompilationMode.VLLM_COMPILE, - pass_config=PassConfig( - fuse_attn_quant=True, - eliminate_noops=True, - enable_sp=True, - fuse_gemm_comms=True, - ), - # Inductor caches custom passes by default as well via uuid - inductor_compile_config={"force_disable_caches": True}, - ) - - with caplog_mp_spawn(logging.DEBUG) as log_holder: - run_model( - compilation_config, model_name, tensor_parallel_size=2, **model_kwargs - ) - log_matches = re.findall( - r"fusion_attn.py:\d+] Fused quant onto (\d+) attention nodes", - log_holder.text, - ) - assert len(log_matches) == 2, log_holder.text - - assert int(log_matches[0]) == matches.attention_fusion - assert int(log_matches[1]) == matches.attention_fusion - - log_matches = re.findall( - r"sequence_parallelism.py:\d+] Replaced (\d+) patterns", - log_holder.text, - ) - assert len(log_matches) == 2, log_holder.text - - assert int(log_matches[0]) == matches.sequence_parallel - assert int(log_matches[1]) == matches.sequence_parallel - - log_matches = re.findall( - r"collective_fusion.py:\d+] Replaced (\d+) patterns", - log_holder.text, - ) - assert len(log_matches) == 2, log_holder.text - - assert int(log_matches[0]) == matches.async_tp - assert int(log_matches[1]) == matches.async_tp - - -@pytest.mark.parametrize( - "model_name, model_kwargs, backend, matches, custom_ops", - # Test rms norm+group quant_fp8 fusion - list[tuple[Any, ...]](flat_product(MODELS_GROUP_FP8, CUSTOM_OPS_QUANT_RMS_NORM)), -) -@pytest.mark.parametrize("inductor_graph_partition", [True, False]) -# TODO: remove skip after we fix the fusion thoroughly -@pytest.mark.skipif(is_blackwell(), reason="Temporarily disabled on Blackwell") -def test_rms_group_quant( - model_name: str, - model_kwargs: dict[str, Any], - backend: AttentionBackendEnum, - matches: Matches, - custom_ops: str, - inductor_graph_partition: bool, - caplog_mp_spawn, - monkeypatch, -): - if inductor_graph_partition and not is_torch_equal_or_newer("2.9.0.dev"): - pytest.skip("Inductor graph partition requires torch>=2.9") - - custom_ops_list = custom_ops.split(",") if custom_ops else [] - - if inductor_graph_partition: - mode = CUDAGraphMode.FULL_AND_PIECEWISE - splitting_ops: list[str] | None = None - else: - mode = CUDAGraphMode.FULL_DECODE_ONLY - splitting_ops = [] - - # Disable, compile cache to make sure custom passes run. - # Otherwise, we can't verify fusion happened through the logs. - monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1") - - # To capture subprocess logs, we need to know whether spawn or fork is used. - # Force spawn as it is more general. - monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") - - # TODO: remove this after fusion is fixed - monkeypatch.setenv("VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES", "0") - - model_kwargs["attention_config"] = {"backend": backend.name} - - compilation_config = CompilationConfig( - # Testing properties - custom_ops=custom_ops_list, - use_inductor_graph_partition=inductor_graph_partition, - cudagraph_mode=mode, - splitting_ops=splitting_ops, - # Common - mode=CompilationMode.VLLM_COMPILE, - pass_config=PassConfig( - fuse_norm_quant=True, fuse_act_quant=True, eliminate_noops=True - ), - # Inductor caches custom passes by default as well via uuid - inductor_compile_config={"force_disable_caches": True}, - ) - - with caplog_mp_spawn(logging.DEBUG) as log_holder: - run_model(compilation_config, model_name, **model_kwargs) - - log_matches = re.findall( - r"\[fusion.py:\d+] Replaced (\d+) patterns", - log_holder.text, - ) - assert len(log_matches) == 1, log_holder.text - assert int(log_matches[0]) == matches.rms_quant_norm_fusion diff --git a/tests/compile/fusion_test_utils.py b/tests/compile/fusion_test_utils.py deleted file mode 100644 index ec7b987bfa6..00000000000 --- a/tests/compile/fusion_test_utils.py +++ /dev/null @@ -1,208 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Shared utilities for fusion tests (e.g. test_fusion_attn.py).""" - -from __future__ import annotations - -import itertools -from collections.abc import Iterable -from typing import Any, NamedTuple - -from tests.v1.attention.utils import AttentionBackendEnum -from vllm import LLM, SamplingParams -from vllm.config import CompilationConfig, CUDAGraphMode -from vllm.platforms import current_platform - -is_blackwell = lambda: current_platform.is_device_capability_family(100) -"""Are we running on Blackwell, a lot of tests depend on it""" - - -def has_cuda_graph_wrapper_metadata() -> bool: - from importlib import import_module - - try: - module = import_module("torch._inductor.utils") - module.CUDAGraphWrapperMetadata # noqa B018 - except AttributeError: - return False - return True - - -class Matches(NamedTuple): - attention_fusion: int = 0 - allreduce_fusion: int = 0 - sequence_parallel: int = 0 - async_tp: int = 0 - rms_quant_norm_fusion: int = 0 - - -class ModelBackendTestCase(NamedTuple): - model_name: str - model_kwargs: dict[str, Any] - backend: AttentionBackendEnum - matches: Matches - - -# E2E model test cases -MODELS_FP8: list[ModelBackendTestCase] = [] -MODELS_FP4: list[ModelBackendTestCase] = [] -MODELS: list[ModelBackendTestCase] = [] # tp-only (unquantized) -MODELS_GROUP_FP8: list[ModelBackendTestCase] = [] - -if current_platform.is_cuda(): - MODELS_FP8 = [ - ModelBackendTestCase( - # Use smaller model for L40s in CI - model_name="RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8", - model_kwargs=dict(max_model_len=1024, kv_cache_dtype="fp8"), - backend=AttentionBackendEnum.TRITON_ATTN, - matches=Matches( - attention_fusion=32, - allreduce_fusion=65, - sequence_parallel=65, - async_tp=128, - ), - ), - ModelBackendTestCase( - model_name="nvidia/Llama-4-Scout-17B-16E-Instruct-FP8", - model_kwargs=dict(max_model_len=1024, kv_cache_dtype="fp8"), - # TODO FlashInfer attn broken on Hopper with kvcache=fp8: - # https://github.com/vllm-project/vllm/issues/28568 - backend=AttentionBackendEnum.FLASHINFER - if is_blackwell() - else AttentionBackendEnum.TRITON_ATTN, - matches=Matches( - attention_fusion=48, - allreduce_fusion=96, - sequence_parallel=96, - async_tp=95, # mlp is moe, no fusion there - ), - ), - ] - - MODELS_FP4 = [ - ModelBackendTestCase( - model_name="nvidia/Llama-3.1-8B-Instruct-FP4", - model_kwargs=dict(max_model_len=1024, kv_cache_dtype="fp8"), - backend=AttentionBackendEnum.FLASHINFER, - matches=Matches( - attention_fusion=32, - allreduce_fusion=65, - sequence_parallel=65, - async_tp=128, - ), - ), - ] - - # TP only (unquantized models) - MODELS = [ - ModelBackendTestCase( - model_name="meta-llama/Llama-3.1-8B-Instruct", - model_kwargs=dict(max_model_len=1024), - backend=AttentionBackendEnum.TRITON_ATTN, - matches=Matches( - attention_fusion=0, - allreduce_fusion=65, - sequence_parallel=65, - async_tp=128, - ), - ), - ModelBackendTestCase( - model_name="Qwen/Qwen3-30B-A3B", - model_kwargs=dict(max_model_len=1024), - backend=AttentionBackendEnum.TRITON_ATTN, - matches=Matches( - attention_fusion=0, - allreduce_fusion=97, - sequence_parallel=97, - async_tp=96, # MLP is MoE, half the fusions of dense - ), - ), - ] - - MODELS_GROUP_FP8 = [ - ModelBackendTestCase( - model_name="Qwen/Qwen3-30B-A3B-FP8", - model_kwargs=dict(max_model_len=1024, kv_cache_dtype="fp8"), - backend=AttentionBackendEnum.TRITON_ATTN, - matches=Matches( - rms_quant_norm_fusion=48, - ), - ), - ] - -elif current_platform.is_rocm(): - MODELS_FP8 = [ - ModelBackendTestCase( - model_name="amd/Llama-3.1-8B-Instruct-FP8-KV", - model_kwargs=dict(max_model_len=1024), - backend=AttentionBackendEnum.TRITON_ATTN, - matches=Matches(attention_fusion=32), - ), - ModelBackendTestCase( - model_name="amd/Llama-3.1-8B-Instruct-FP8-KV", - model_kwargs=dict(max_model_len=1024), - backend=AttentionBackendEnum.ROCM_ATTN, - matches=Matches(attention_fusion=32), - ), - ModelBackendTestCase( - model_name="amd/Llama-3.1-8B-Instruct-FP8-KV", - model_kwargs=dict(max_model_len=1024), - backend=AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN, - matches=Matches(attention_fusion=32), - ), - ] - - -# Custom ops toggle lists for parametrization -CUSTOM_OPS_FP8 = ["-quant_fp8", "+quant_fp8"] -CUSTOM_OPS_RMS_NORM = ["-rms_norm", "+rms_norm"] -CUSTOM_OPS_QUANT_RMS_NORM = ["+quant_fp8,+rms_norm"] - - -def custom_ops_product(*custom_ops_lists: list[str]) -> Iterable[str]: - """Generate all combinations of custom ops for parametrization.""" - for op_list in itertools.product(*custom_ops_lists): - yield ",".join(op_list) - - -def run_model(compile_config: int | CompilationConfig, model: str, **model_kwargs): - """Run a model with the given compilation config for E2E fusion tests.""" - compilation_config = ( - compile_config - if isinstance(compile_config, CompilationConfig) - else CompilationConfig(mode=compile_config) - ) - - prompts = [ - "Hello, my name is", - "The president of the United States is", - "The capital of France is", - "The future of AI is", - ] - sampling_params = SamplingParams(temperature=0) - # Allow override from model_kwargs - model_kwargs = {"tensor_parallel_size": 1, **model_kwargs} - model_kwargs = {"disable_custom_all_reduce": True, **model_kwargs} - - # No cudagraphs by default - if compilation_config.cudagraph_mode is None: - compilation_config.cudagraph_mode = CUDAGraphMode.NONE - llm = LLM( - model=model, - compilation_config=compilation_config, - **model_kwargs, - ) - outputs = llm.generate(prompts, sampling_params) - - # Print the outputs. - for output in outputs: - prompt = output.prompt - generated_text = output.outputs[0].text - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") - - # Get the compile ranges split points after vllm config post init - # in order to compute compile ranges correctly - compilation_config.compile_ranges_split_points = ( - llm.llm_engine.vllm_config.compilation_config.compile_ranges_split_points - ) diff --git a/tests/compile/fusions_e2e/__init__.py b/tests/compile/fusions_e2e/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/compile/fusions_e2e/common.py b/tests/compile/fusions_e2e/common.py new file mode 100644 index 00000000000..d950bf5b652 --- /dev/null +++ b/tests/compile/fusions_e2e/common.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import itertools +from collections.abc import Callable, Iterable +from typing import Any, NamedTuple + +import pytest +import regex as re + +from vllm.platforms import current_platform +from vllm.v1.attention.backends.registry import AttentionBackendEnum + + +class Matches(NamedTuple): + # simple pointwise + rms_quant_fusion: int = 0 + act_quant_fusion: int = 0 + norm_rope_fusion: int = 0 + attn_quant_fusion: int = 0 + # distributed + ar_rms_fusion: int = 0 + sequence_parallel: int = 0 + async_tp: int = 0 + + +class ModelFusionInfo(NamedTuple): + model_name: str + matches: Callable[[int], Matches] + """Given number of hidden layers, produces the matches object""" + model_kwargs: dict[str, Any] = {} + hf_overrides: Callable[[int], dict] = lambda n: {"num_hidden_layers": n} + + +class AttentionBackendCase(NamedTuple): + backend: AttentionBackendEnum + model_kwargs: dict[str, Any] = {} + """Additional args required for attn+quant fusion""" + + +is_blackwell = lambda: current_platform.is_device_capability_family(100) +"""Are we running on Blackwell, a lot of tests depend on it""" + + +def custom_ops_combos(*custom_ops: str) -> Iterable[str]: + """Generate all combinations of custom ops for parametrization.""" + custom_ops_lists = [[f"-{op}", f"+{op}"] for op in custom_ops] + for op_list in itertools.product(*custom_ops_lists): + yield ",".join(op_list) + + +# Quick inline validation +assert list(custom_ops_combos("silu_and_mul")) == ["-silu_and_mul", "+silu_and_mul"] +assert list(custom_ops_combos("quant_fp8", "rms_norm")) == [ + "-quant_fp8,-rms_norm", + "-quant_fp8,+rms_norm", + "+quant_fp8,-rms_norm", + "+quant_fp8,+rms_norm", +] + + +def has_cuda_graph_wrapper_metadata() -> bool: + from importlib import import_module + + try: + module = import_module("torch._inductor.utils") + module.CUDAGraphWrapperMetadata # noqa B018 + except AttributeError: + return False + return True + + +INDUCTOR_GRAPH_PARTITION = [ + pytest.param( + True, + marks=pytest.mark.skipif( + not has_cuda_graph_wrapper_metadata(), + reason="torch version does not support Inductor partition", + ), + id="inductor_partition", + ), + pytest.param(False, id="dynamo_partition"), +] + +FUSION_LOG_PATTERNS: dict[str, re.Pattern] = { + "rms_quant_fusion": re.compile( + r"\[(?:compilation/)?fusion.py:\d+] Replaced (\d+) patterns" + ), + "act_quant_fusion": re.compile( + r"activation_quant_fusion.py:\d+] Replaced (\d+) patterns" + ), + "norm_rope_fusion": re.compile( + r"qk_norm_rope_fusion.py:\d+] Fused QK Norm\+RoPE on (\d+) sites" + ), + "attn_quant_fusion": re.compile( + r"fusion_attn.py:\d+] Fused quant onto (\d+) attention nodes" + ), + "ar_rms_fusion": re.compile(r"collective_fusion.py:\d+] Replaced (\d+) patterns"), + "sequence_parallel": re.compile( + r"sequence_parallelism.py:\d+] Replaced (\d+) patterns" + ), + "async_tp": re.compile(r"collective_fusion.py:\d+] Replaced (\d+) patterns"), +} diff --git a/tests/compile/fusions_e2e/conftest.py b/tests/compile/fusions_e2e/conftest.py new file mode 100644 index 00000000000..1d9f6cda9fd --- /dev/null +++ b/tests/compile/fusions_e2e/conftest.py @@ -0,0 +1,158 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import logging + +import pytest +import regex as re + +from vllm import LLM, SamplingParams +from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode + +from .common import FUSION_LOG_PATTERNS, AttentionBackendCase, Matches + + +def run_model(compile_config: int | CompilationConfig, model: str, **model_kwargs): + """Run a model with the given compilation config for E2E fusion tests.""" + compilation_config = ( + compile_config + if isinstance(compile_config, CompilationConfig) + else CompilationConfig(mode=compile_config) + ) + + prompts = [ + "Hello, my name is", + "The president of the United States is", + "The capital of France is", + "The future of AI is", + ] + sampling_params = SamplingParams(temperature=0) + # Allow override from model_kwargs + model_kwargs = {"tensor_parallel_size": 1, **model_kwargs} + model_kwargs = {"disable_custom_all_reduce": True, **model_kwargs} + + # No cudagraphs by default + if compilation_config.cudagraph_mode is None: + compilation_config.cudagraph_mode = CUDAGraphMode.NONE + llm = LLM( + model=model, + compilation_config=compilation_config, + **model_kwargs, + ) + outputs = llm.generate(prompts, sampling_params) + + # Print the outputs. + for output in outputs: + prompt = output.prompt + generated_text = output.outputs[0].text + print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") + + # Get the compile ranges split points after vllm config post init + # in order to compute compile ranges correctly + compilation_config.compile_ranges_split_points = ( + llm.llm_engine.vllm_config.compilation_config.compile_ranges_split_points + ) + + +@pytest.fixture +def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn): + def run( + model_name: str, + matches: Matches, + model_kwargs: dict, + attn_backend: AttentionBackendCase, + compilation_config: dict, + matches_check: list[str], + use_deepgemm: bool = False, + tp_size: int = 1, + ): + monkeypatch.setenv("VLLM_USE_DEEP_GEMM", "1" if use_deepgemm else "0") + + # Disable, compile cache to make sure custom passes run. + # Otherwise, we can't verify fusion happened through the logs. + monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1") + + # To capture subprocess logs, we need to know whether spawn or fork is used. + # Force spawn as it is more general. + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + + model_kwargs = {**attn_backend.model_kwargs, **model_kwargs} + model_kwargs["attention_config"] = {"backend": attn_backend.backend.name} + model_kwargs["tensor_parallel_size"] = tp_size + + # Always compile the full graph instead of piecewise + if not compilation_config["use_inductor_graph_partition"]: + compilation_config["splitting_ops"] = [] + + full_compilation_config = CompilationConfig( + cudagraph_mode=CUDAGraphMode.NONE, + mode=CompilationMode.VLLM_COMPILE, + inductor_compile_config={"force_disable_caches": True}, + **compilation_config, + ) + + with caplog_mp_spawn(logging.DEBUG) as log_holder: + run_model(full_compilation_config, model_name, **model_kwargs) + + num_compile_ranges = len(full_compilation_config.get_compile_ranges()) + assert num_compile_ranges in [1, 2] + + print(f"Compile ranges: {full_compilation_config.get_compile_ranges()}") + print("Fusion results:") + + # Iterate through all so printing happens before asserting + log_matches_dict = {} + for match_name, pattern in FUSION_LOG_PATTERNS.items(): + log_matches_dict[match_name] = list(pattern.findall(log_holder.text)) + print(f"- {match_name}={','.join(log_matches_dict[match_name])}") + + # Now check the matches + for match_name in matches_check: + num_ranges_activated = ( + 1 if match_name == "ar_rms_fusion" else num_compile_ranges + ) + n_expected = tp_size * num_ranges_activated + + log_matches = list(int(ms) for ms in log_matches_dict[match_name]) + assert len(log_matches) == n_expected, ( + f"Could not find {n_expected} {match_name} " + f"(found {len(log_matches)}) in:\n {log_holder.text}" + ) + + expected_matches = getattr(matches, match_name) + + if match_name == "rms_quant_fusion" and "ar_rms_fusion" in matches_check: + # AR+rms+quant takes precedence over rms+quant if activated. + # That means we get full matching where ar+rms+quant was not activated, + # and less where it was + assert sum(m == expected_matches for m in log_matches) == tp_size * ( + num_ranges_activated - 1 + ), "Expecting full rms+quant fusion where ar+rms+quant not activated" + + assert all( + expected_matches - matches.ar_rms_fusion <= m <= expected_matches + for m in log_matches + ), ( + f"Expecting at least {expected_matches - matches.ar_rms_fusion} " + f"where ar+rms+quant was activated" + ) + else: + expected_matches_list = [expected_matches] * n_expected + assert sorted(log_matches) == expected_matches_list, ( + f"{match_name} expected: {expected_matches_list}, " + f"found: {sorted(log_matches)}" + ) + + if match_name == "ar_rms_fusion": + log_matches = re.findall( + r"pass_manager.py:\d+] Skipping " + r".*AllReduceFusionPass.* with compile range", + log_holder.text, + ) + + n_expected = tp_size * (num_compile_ranges - num_ranges_activated) + assert len(log_matches) == n_expected, ( + f'Could not find {n_expected} "Skipping AllReduceFusionPass" ' + f"(found {len(log_matches)}) in:\n {log_holder.text}" + ) + + return run diff --git a/tests/compile/fusions_e2e/models.py b/tests/compile/fusions_e2e/models.py new file mode 100644 index 00000000000..ef9b6be2501 --- /dev/null +++ b/tests/compile/fusions_e2e/models.py @@ -0,0 +1,112 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest + +from vllm.utils.flashinfer import has_flashinfer +from vllm.v1.attention.backends.registry import AttentionBackendEnum + +from .common import AttentionBackendCase, Matches, ModelFusionInfo, is_blackwell + +# Attn backends +FLASHINFER_ATTN = pytest.param( + AttentionBackendCase( + backend=AttentionBackendEnum.FLASHINFER, + model_kwargs=dict(kv_cache_dtype="fp8"), + ), + id="FLASHINFER", + marks=pytest.mark.skipif( + not is_blackwell() or not has_flashinfer(), + reason="FI backend requires Blackwell and FlashInfer", + ), +) + +TRITON_ATTN = pytest.param( + AttentionBackendCase(backend=AttentionBackendEnum.TRITON_ATTN), id="TRITON_ATTN" +) + +# Models +llama3_8b = ModelFusionInfo( + model_name="meta-llama/Llama-3.1-8B-Instruct", + matches=lambda n_layers: Matches( + ar_rms_fusion=n_layers * 2 + 1, + sequence_parallel=n_layers * 2 + 1, + async_tp=n_layers * 4, + ), +) + +llama3_8b_fp8 = ModelFusionInfo( + model_name="RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8", + matches=lambda n_layers: Matches( + rms_quant_fusion=n_layers * 2, + act_quant_fusion=n_layers, + attn_quant_fusion=n_layers, + ar_rms_fusion=n_layers * 2 + 1, + sequence_parallel=n_layers * 2 + 1, + async_tp=n_layers * 4, + ), +) + +llama3_8b_fp4 = ModelFusionInfo( + model_name="nvidia/Llama-3.1-8B-Instruct-FP4", + matches=lambda n_layers: Matches( + rms_quant_fusion=0, + act_quant_fusion=n_layers, + attn_quant_fusion=n_layers, + ar_rms_fusion=n_layers * 2 + 1, + sequence_parallel=n_layers * 2 + 1, + async_tp=n_layers * 4, + ), +) + +# MoEs cannot do act+quant fusion because those ops are hidden from torch.compile. +# MoEs also only expose 1 rms+quant fusion because the quant for up_proj is hidden. +# TODO(luka): https://github.com/vllm-project/vllm/issues/31985 +# Also, for MoEs, gemm+collective fusion only happens for dense GEMMs (o_proj/qkv proj) + +llama4_scout_fp8 = ModelFusionInfo( + model_name="nvidia/Llama-4-Scout-17B-16E-Instruct-FP8", + hf_overrides=lambda n_layers: {"text_config": {"num_hidden_layers": n_layers}}, + matches=lambda n_layers: Matches( + rms_quant_fusion=n_layers, + attn_quant_fusion=n_layers, + ar_rms_fusion=n_layers * 2, + sequence_parallel=n_layers * 2, + async_tp=n_layers * 2 - 1, + ), +) + +llama4_scout_fp4 = ModelFusionInfo( + model_name="nvidia/Llama-4-Scout-17B-16E-Instruct-NVFP4", + hf_overrides=lambda n_layers: {"text_config": {"num_hidden_layers": n_layers}}, + matches=lambda n_layers: Matches( + rms_quant_fusion=0, + attn_quant_fusion=n_layers, + ar_rms_fusion=n_layers * 2, + sequence_parallel=n_layers * 2, + async_tp=n_layers * 2 - 1, + ), +) + +qwen3_a3b = ModelFusionInfo( + model_name="Qwen/Qwen3-30B-A3B", + matches=lambda n_layers: Matches( + norm_rope_fusion=n_layers, + ar_rms_fusion=n_layers * 2 + 1, + sequence_parallel=n_layers * 2 + 1, + async_tp=n_layers * 2, + ), +) + +qwen3_a3b_fp8 = ModelFusionInfo( + model_name="Qwen/Qwen3-30B-A3B-FP8", + matches=lambda n_layers: Matches( + rms_quant_fusion=n_layers, + # TODO broken on Blackwell: + # https://github.com/vllm-project/vllm/issues/33295 + norm_rope_fusion=0 if is_blackwell() else n_layers, + attn_quant_fusion=0, # attn + group quant not supported + ar_rms_fusion=n_layers * 2 + 1, + sequence_parallel=n_layers * 2 + 1, + async_tp=n_layers * 2, + ), +) diff --git a/tests/compile/fusions_e2e/test_tp1_quant.py b/tests/compile/fusions_e2e/test_tp1_quant.py new file mode 100644 index 00000000000..03f102794f8 --- /dev/null +++ b/tests/compile/fusions_e2e/test_tp1_quant.py @@ -0,0 +1,146 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Callable + +import pytest + +from vllm.config import PassConfig + +from .common import ( + INDUCTOR_GRAPH_PARTITION, + AttentionBackendCase, + Matches, + custom_ops_combos, + is_blackwell, +) +from .models import ( + FLASHINFER_ATTN, + TRITON_ATTN, + llama3_8b_fp4, + llama3_8b_fp8, + llama4_scout_fp4, + llama4_scout_fp8, + qwen3_a3b_fp8, +) + + +@pytest.mark.parametrize( + "model_name, matches_fn, model_kwargs, hf_overrides, use_deepgemm", + [ + (*llama3_8b_fp8, False), + (*llama4_scout_fp8, False), + (*qwen3_a3b_fp8, False), + (*qwen3_a3b_fp8, True), + ], +) +@pytest.mark.parametrize("attn_backend", [TRITON_ATTN, FLASHINFER_ATTN]) +@pytest.mark.parametrize("n_layers", [6]) +@pytest.mark.parametrize("custom_ops", custom_ops_combos("quant_fp8", "rms_norm")) +@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION) +def test_tp1_fp8_fusions( + model_name: str, + matches_fn: Callable[[int], Matches], + model_kwargs: dict, + hf_overrides: Callable[[int], dict], + attn_backend: AttentionBackendCase, + n_layers: int, + custom_ops: str, + inductor_graph_partition: bool, + use_deepgemm: bool, + run_e2e_fusion_test, + monkeypatch, +): + if use_deepgemm: + # TODO(luka/eliza) DeepGEMM uses different quants, matching not supported + # - on Blackwell, uses a special quant fp8, currently not supported + # - on Hopper, tma-aligned scales inhibit matching (fix WIP) + pytest.skip("DeepGEMM & quant matching not currently supported") + + matches = matches_fn(n_layers) + + if "qwen" in model_name.lower() and "-quant_fp8" in custom_ops: + # This is why config forces +quant_fp8 by default + pytest.skip("native QuantFP8 matching not supported for group quant") + + # Reduce size of model and skip weight loading time + model_kwargs["hf_overrides"] = hf_overrides(n_layers) + model_kwargs["load_format"] = "dummy" + model_kwargs["max_model_len"] = 1024 + + compilation_config = dict( + use_inductor_graph_partition=inductor_graph_partition, + custom_ops=custom_ops.split(","), + pass_config=PassConfig( + fuse_norm_quant=True, + fuse_act_quant=True, + fuse_attn_quant=True, + enable_qk_norm_rope_fusion=True, + ), + ) + + matches_check = [ + "rms_quant_fusion", + "act_quant_fusion", + "norm_rope_fusion", + "attn_quant_fusion", + ] + + run_e2e_fusion_test( + model_name, + matches, + model_kwargs, + attn_backend, + compilation_config, + matches_check, + use_deepgemm=use_deepgemm, + ) + + +@pytest.mark.parametrize( + "model_name, matches_fn, model_kwargs, hf_overrides", + [llama3_8b_fp4, llama4_scout_fp4], +) +@pytest.mark.parametrize("attn_backend", [FLASHINFER_ATTN]) +@pytest.mark.parametrize("n_layers", [6]) +@pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm")) +@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION) +@pytest.mark.skipif(not is_blackwell(), reason="Blackwell required for fp4") +def test_tp1_fp4_fusions( + model_name: str, + matches_fn: Callable[[int], Matches], + model_kwargs: dict, + hf_overrides: Callable[[int], dict], + attn_backend: AttentionBackendCase, + n_layers: int, + custom_ops: str, + inductor_graph_partition: bool, + run_e2e_fusion_test, +): + matches = matches_fn(n_layers) + + # Reduce size of model and skip weight loading time + model_kwargs["hf_overrides"] = hf_overrides(n_layers) + model_kwargs["load_format"] = "dummy" + model_kwargs["max_model_len"] = 1024 + + compilation_config = dict( + use_inductor_graph_partition=inductor_graph_partition, + custom_ops=custom_ops.split(","), + pass_config=PassConfig( + fuse_norm_quant=True, + fuse_act_quant=True, + fuse_attn_quant=True, + enable_qk_norm_rope_fusion=True, + ), + ) + + matches_check = ["act_quant_fusion", "attn_quant_fusion", "norm_rope_fusion"] + + run_e2e_fusion_test( + model_name, + matches, + model_kwargs, + attn_backend, + compilation_config, + matches_check, + ) diff --git a/tests/compile/fusions_e2e/test_tp2_ar_rms.py b/tests/compile/fusions_e2e/test_tp2_ar_rms.py new file mode 100644 index 00000000000..18b19565c1f --- /dev/null +++ b/tests/compile/fusions_e2e/test_tp2_ar_rms.py @@ -0,0 +1,199 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Callable + +import pytest + +from vllm.config import PassConfig + +from ...utils import multi_gpu_test +from .common import ( + INDUCTOR_GRAPH_PARTITION, + AttentionBackendCase, + Matches, + custom_ops_combos, + is_blackwell, +) +from .models import ( + FLASHINFER_ATTN, + TRITON_ATTN, + llama3_8b, + llama3_8b_fp4, + llama3_8b_fp8, + llama4_scout_fp4, + llama4_scout_fp8, + qwen3_a3b, + qwen3_a3b_fp8, +) + + +@multi_gpu_test(num_gpus=2) +@pytest.mark.parametrize( + "model_name, matches_fn, model_kwargs, hf_overrides", + # qwen3-fp8 should still fuse AR+rms even though group quant is not yet supported + [llama3_8b_fp8, llama4_scout_fp8, qwen3_a3b_fp8], +) +@pytest.mark.parametrize("attn_backend", [TRITON_ATTN, FLASHINFER_ATTN]) +@pytest.mark.parametrize("n_layers", [4]) +@pytest.mark.parametrize("custom_ops", custom_ops_combos("quant_fp8", "rms_norm")) +@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION) +def test_tp2_ar_rms_fp8_fusions( + model_name: str, + matches_fn: Callable[[int], Matches], + model_kwargs: dict, + hf_overrides: Callable[[int], dict], + attn_backend: AttentionBackendCase, + n_layers: int, + custom_ops: str, + inductor_graph_partition: bool, + run_e2e_fusion_test, + monkeypatch, +): + matches = matches_fn(n_layers) + + if "qwen" in model_name.lower() and "-quant_fp8" in custom_ops: + # This is why config forces +quant_fp8 by default + pytest.skip("native QuantFP8 matching not supported for group quant") + + # Reduce size of model and skip weight loading time + model_kwargs["hf_overrides"] = hf_overrides(n_layers) + model_kwargs["load_format"] = "dummy" + model_kwargs["max_model_len"] = 1024 + + compilation_config = dict( + use_inductor_graph_partition=inductor_graph_partition, + custom_ops=custom_ops.split(","), + pass_config=PassConfig( + fuse_norm_quant=True, + fuse_act_quant=True, + fuse_attn_quant=True, + enable_qk_norm_rope_fusion=True, + fuse_allreduce_rms=True, + ), + ) + + matches_check = [ + "rms_quant_fusion", + "act_quant_fusion", + "norm_rope_fusion", + "attn_quant_fusion", + "ar_rms_fusion", + ] + + run_e2e_fusion_test( + model_name, + matches, + model_kwargs, + attn_backend, + compilation_config, + matches_check, + tp_size=2, + ) + + +@multi_gpu_test(num_gpus=2) +@pytest.mark.parametrize( + "model_name, matches_fn, model_kwargs, hf_overrides", + [llama3_8b_fp4, llama4_scout_fp4], +) +@pytest.mark.parametrize("attn_backend", [FLASHINFER_ATTN]) +@pytest.mark.parametrize("n_layers", [4]) +@pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm")) +@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION) +@pytest.mark.skipif(not is_blackwell(), reason="Blackwell required for fp4") +def test_tp2_ar_rms_fp4_fusions( + model_name: str, + matches_fn: Callable[[int], Matches], + model_kwargs: dict, + hf_overrides: Callable[[int], dict], + attn_backend: AttentionBackendCase, + n_layers: int, + custom_ops: str, + inductor_graph_partition: bool, + run_e2e_fusion_test, + monkeypatch, +): + matches = matches_fn(n_layers) + + # Reduce size of model and skip weight loading time + model_kwargs["hf_overrides"] = hf_overrides(n_layers) + model_kwargs["load_format"] = "dummy" + model_kwargs["max_model_len"] = 1024 + + compilation_config = dict( + use_inductor_graph_partition=inductor_graph_partition, + custom_ops=custom_ops.split(","), + pass_config=PassConfig( + fuse_act_quant=True, + fuse_attn_quant=True, + fuse_allreduce_rms=True, + ), + ) + + matches_check = [ + "act_quant_fusion", + "attn_quant_fusion", + "ar_rms_fusion", + ] + + run_e2e_fusion_test( + model_name, + matches, + model_kwargs, + attn_backend, + compilation_config, + matches_check, + tp_size=2, + ) + + +@multi_gpu_test(num_gpus=2) +@pytest.mark.parametrize( + "model_name, matches_fn, model_kwargs, hf_overrides", + [llama3_8b, qwen3_a3b], +) +@pytest.mark.parametrize("attn_backend", [TRITON_ATTN]) +@pytest.mark.parametrize("n_layers", [4]) +@pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm")) +@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION) +def test_tp2_ar_rms_fusions( + model_name: str, + matches_fn: Callable[[int], Matches], + model_kwargs: dict, + hf_overrides: Callable[[int], dict], + attn_backend: AttentionBackendCase, + n_layers: int, + custom_ops: str, + inductor_graph_partition: bool, + run_e2e_fusion_test, +): + matches = matches_fn(n_layers) + + # Reduce size of model and skip weight loading time + model_kwargs["hf_overrides"] = hf_overrides(n_layers) + model_kwargs["load_format"] = "dummy" + model_kwargs["max_model_len"] = 1024 + + compilation_config = dict( + use_inductor_graph_partition=inductor_graph_partition, + custom_ops=custom_ops.split(","), + pass_config=PassConfig( + enable_qk_norm_rope_fusion=True, + fuse_allreduce_rms=True, + ), + ) + + matches_check = [ + "norm_rope_fusion", + "ar_rms_fusion", + ] + + run_e2e_fusion_test( + model_name, + matches, + model_kwargs, + attn_backend, + compilation_config, + matches_check, + tp_size=2, + ) diff --git a/tests/compile/fusions_e2e/test_tp2_async_tp.py b/tests/compile/fusions_e2e/test_tp2_async_tp.py new file mode 100644 index 00000000000..4769ca1e0b6 --- /dev/null +++ b/tests/compile/fusions_e2e/test_tp2_async_tp.py @@ -0,0 +1,143 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Callable + +import pytest + +from vllm.config import PassConfig + +from ...utils import multi_gpu_test +from .common import ( + INDUCTOR_GRAPH_PARTITION, + AttentionBackendCase, + Matches, + custom_ops_combos, + is_blackwell, +) +from .models import ( + FLASHINFER_ATTN, + TRITON_ATTN, + llama3_8b, + llama3_8b_fp8, + llama4_scout_fp8, + qwen3_a3b, +) + + +@multi_gpu_test(num_gpus=2) +@pytest.mark.parametrize( + "model_name, matches_fn, model_kwargs, hf_overrides", + [llama3_8b_fp8, llama4_scout_fp8], +) +@pytest.mark.parametrize("attn_backend", [TRITON_ATTN, FLASHINFER_ATTN]) +@pytest.mark.parametrize("n_layers", [4]) +@pytest.mark.parametrize("custom_ops", custom_ops_combos("quant_fp8", "rms_norm")) +@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION) +def test_tp2_async_tp_fp8_fusions( + model_name: str, + matches_fn: Callable[[int], Matches], + model_kwargs: dict, + hf_overrides: Callable[[int], dict], + attn_backend: AttentionBackendCase, + n_layers: int, + custom_ops: str, + inductor_graph_partition: bool, + run_e2e_fusion_test, + monkeypatch, +): + matches = matches_fn(n_layers) + + if is_blackwell(): + # Disable FlashInfer scaled_mm FP8 as it's not supported in async tp patterns + monkeypatch.setenv("VLLM_DISABLED_KERNELS", "FlashInferFP8ScaledMMLinearKernel") + + # Reduce size of model and skip weight loading time + model_kwargs["hf_overrides"] = hf_overrides(n_layers) + model_kwargs["load_format"] = "dummy" + model_kwargs["max_model_len"] = 1024 + + compilation_config = dict( + use_inductor_graph_partition=inductor_graph_partition, + custom_ops=custom_ops.split(","), + pass_config=PassConfig( + fuse_norm_quant=True, + fuse_act_quant=True, + fuse_attn_quant=True, + enable_qk_norm_rope_fusion=True, + enable_sp=True, + fuse_gemm_comms=True, + ), + ) + + matches_check = [ + "rms_quant_fusion", + "act_quant_fusion", + "norm_rope_fusion", + "attn_quant_fusion", + "sequence_parallel", + "async_tp", + ] + + run_e2e_fusion_test( + model_name, + matches, + model_kwargs, + attn_backend, + compilation_config, + matches_check, + tp_size=2, + ) + + +@multi_gpu_test(num_gpus=2) +@pytest.mark.parametrize( + "model_name, matches_fn, model_kwargs, hf_overrides", + [llama3_8b, qwen3_a3b], +) +@pytest.mark.parametrize("attn_backend", [TRITON_ATTN]) +@pytest.mark.parametrize("n_layers", [4]) +@pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm")) +@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION) +def test_tp2_async_tp_fusions( + model_name: str, + matches_fn: Callable[[int], Matches], + model_kwargs: dict, + hf_overrides: Callable[[int], dict], + attn_backend: AttentionBackendCase, + n_layers: int, + custom_ops: str, + inductor_graph_partition: bool, + run_e2e_fusion_test, +): + matches = matches_fn(n_layers) + + # Reduce size of model and skip weight loading time + model_kwargs["hf_overrides"] = hf_overrides(n_layers) + model_kwargs["load_format"] = "dummy" + model_kwargs["max_model_len"] = 1024 + + compilation_config = dict( + use_inductor_graph_partition=inductor_graph_partition, + custom_ops=custom_ops.split(","), + pass_config=PassConfig( + enable_qk_norm_rope_fusion=True, + enable_sp=True, + fuse_gemm_comms=True, + ), + ) + + matches_check = [ + "norm_rope_fusion", + "sequence_parallel", + "async_tp", + ] + + run_e2e_fusion_test( + model_name, + matches, + model_kwargs, + attn_backend, + compilation_config, + matches_check, + tp_size=2, + ) diff --git a/tests/compile/test_fusion_attn.py b/tests/compile/test_fusion_attn.py index 50492a5693d..6515c52228e 100644 --- a/tests/compile/test_fusion_attn.py +++ b/tests/compile/test_fusion_attn.py @@ -1,23 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import copy -import logging -from typing import Any import pytest -import regex as re import torch._dynamo from tests.compile.backend import LazyInitPass, TestBackend -from tests.compile.fusion_test_utils import ( - CUSTOM_OPS_FP8, - MODELS_FP4, - MODELS_FP8, - Matches, - has_cuda_graph_wrapper_metadata, - is_blackwell, - run_model, -) from tests.utils import flat_product from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant @@ -31,7 +19,6 @@ from vllm.config import ( CacheConfig, CompilationConfig, CompilationMode, - CUDAGraphMode, ModelConfig, PassConfig, SchedulerConfig, @@ -47,7 +34,6 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( ) from vllm.platforms import current_platform from vllm.utils.flashinfer import has_flashinfer -from vllm.utils.torch_utils import is_torch_equal_or_newer from vllm.v1.attention.backend import AttentionMetadata from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.kv_cache_interface import AttentionSpec @@ -501,88 +487,3 @@ def test_attention_quant_pattern( # Check that results are close torch.testing.assert_close(result_unfused, result_fused, atol=1e-2, rtol=1e-2) - - -@pytest.mark.parametrize( - "model_name, model_kwargs, backend, matches, custom_ops", - # Test attention+quant_fp8 fusion with custom and torch impls of QuantFP8 - list(flat_product(MODELS_FP8, CUSTOM_OPS_FP8)) - # quant_fp4 only has the custom impl - + list(flat_product(MODELS_FP4, [""])), -) -@pytest.mark.parametrize( - "inductor_graph_partition", - [ - pytest.param( - True, - marks=pytest.mark.skipif( - not has_cuda_graph_wrapper_metadata(), - reason="This test requires" - "torch._inductor.utils.CUDAGraphWrapperMetadata to run", - ), - ), - False, - ], -) -def test_attn_quant( - model_name: str, - model_kwargs: dict[str, Any], - backend: AttentionBackendEnum, - matches: Matches, - custom_ops: str, - inductor_graph_partition: bool, - caplog_mp_spawn, - monkeypatch, -): - if not current_platform.has_device_capability(90): - pytest.skip("test_attn_quant requires H100 (SM90) or B200 (SM100) GPU") - if backend == AttentionBackendEnum.FLASHINFER and ( - not is_blackwell() or not has_flashinfer() - ): - pytest.skip("FlashInfer attn fusion requires Blackwell and flashinfer") - if inductor_graph_partition and not is_torch_equal_or_newer("2.9.0.dev"): - pytest.skip("Inductor graph partition requires torch>=2.9") - - custom_ops_list = custom_ops.split(",") if custom_ops else [] - - if inductor_graph_partition: - mode = CUDAGraphMode.FULL_AND_PIECEWISE - splitting_ops: list[str] | None = None - else: - # FIXME: Llama-4-Scout-17B-16E-Instruct-FP8 + FlashInfer + Blackwell end at - # CUDAGraphMode.NONE here because it derives an attention backend that - # does not support full cudagraphs - mode = CUDAGraphMode.FULL_DECODE_ONLY - splitting_ops = [] - - # Disable, compile cache to make sure custom passes run. - # Otherwise, we can't verify fusion happened through the logs. - monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1") - - # To capture subprocess logs, we need to know whether spawn or fork is used. - # Force spawn as it is more general. - monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") - model_kwargs["attention_config"] = {"backend": backend.name} - - compilation_config = CompilationConfig( - # Testing properties - custom_ops=custom_ops_list, - use_inductor_graph_partition=inductor_graph_partition, - cudagraph_mode=mode, - splitting_ops=splitting_ops, - # Common - mode=CompilationMode.VLLM_COMPILE, - pass_config=PassConfig(fuse_attn_quant=True, eliminate_noops=True), - # Inductor caches custom passes by default as well via uuid - inductor_compile_config={"force_disable_caches": True}, - ) - - with caplog_mp_spawn(logging.DEBUG) as log_holder: - run_model(compilation_config, model_name, **model_kwargs) - - log_matches = re.findall( - r"fusion_attn.py:\d+] Fused quant onto (\d+) attention nodes", - log_holder.text, - ) - assert len(log_matches) == 1, log_holder.text - assert int(log_matches[0]) == matches.attention_fusion diff --git a/tests/test_config.py b/tests/test_config.py index f3c3003a00c..6e2a5966116 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1002,7 +1002,7 @@ def test_vllm_config_explicit_overrides(): assert config.compilation_config.pass_config.fuse_attn_quant is True # Explicit cudagraph mode override on quantized model at O2 - pass_config = PassConfig(fuse_gemm_comms=True) + pass_config = PassConfig(enable_qk_norm_rope_fusion=True) compilation_config = CompilationConfig( cudagraph_mode=CUDAGraphMode.NONE, pass_config=pass_config ) @@ -1012,7 +1012,7 @@ def test_vllm_config_explicit_overrides(): compilation_config=compilation_config, ) assert config.compilation_config.cudagraph_mode == CUDAGraphMode.NONE - assert config.compilation_config.pass_config.fuse_gemm_comms is True + assert config.compilation_config.pass_config.enable_qk_norm_rope_fusion is True # Mode should still use default for O2 assert config.compilation_config.mode == CompilationMode.VLLM_COMPILE diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 846ed50e0bd..93d88730ea1 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -766,7 +766,12 @@ class VllmConfig: if self.compilation_config.pass_config.fuse_gemm_comms: self.compilation_config.pass_config.enable_sp = True if self.compilation_config.pass_config.enable_sp: - if "-rms_norm" in self.compilation_config.custom_ops: + if self.parallel_config.tensor_parallel_size == 1: + logger.warning("Sequence Parallelism requires TP>1, disabling") + self.compilation_config.pass_config.enable_sp = False + self.compilation_config.pass_config.fuse_gemm_comms = False + + elif "-rms_norm" in self.compilation_config.custom_ops: logger.warning( "RMS norm force disabled, sequence parallelism might break" ) From bbe0574d8e51c1c5935aeff9e92040c61d1d59c5 Mon Sep 17 00:00:00 2001 From: zhanqiuhu <49648934+ZhanqiuHu@users.noreply.github.com> Date: Wed, 4 Feb 2026 19:49:18 -0500 Subject: [PATCH 076/810] [Bugfix] Disable TRTLLM attention when KV transfer is enabled (#33192) Signed-off-by: Zhanqiu Hu --- vllm/v1/attention/backends/flashinfer.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index afefc164f5f..7e02aa36ffc 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -573,6 +573,20 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): # try to use fp8 q if kv cache is fp8, and will fall back to model dtype # if TRTLLM attention kernel is not used when building attn metadata can_use_trtllm = can_use_trtllm_attention(self.num_qo_heads, self.num_kv_heads) + + # TRTLLM attention requires strictly contiguous KV cache tensors. + # When KV transfer (P/D disaggregation) is enabled, the KV cache may be + # permuted into non-contiguous views, which causes assertion failures. + self._kv_transfer_enabled = vllm_config.kv_transfer_config is not None + if can_use_trtllm and self._kv_transfer_enabled: + logger.info_once( + "TRTLLM attention is disabled because KV transfer " + "(P/D disaggregation) is enabled. TRTLLM attention requires " + "strictly contiguous KV cache tensors which may not be " + "guaranteed with KV transfer." + ) + can_use_trtllm = False + if ( can_use_trtllm and not vllm_config.attention_config.disable_flashinfer_q_quantization @@ -822,6 +836,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): has_sinks=self.has_sinks, has_spec=uses_spec_reorder, ) + # KV transfer requires non-contiguous KV cache views, incompatible with TRTLLM + if self._kv_transfer_enabled: + prefill_use_trtllm = False decode_use_trtllm = ( self.use_trtllm_decode_attention and self.dcp_world_size <= 1 ) From a7be77beef5f59d9d349818b4f2860483551b255 Mon Sep 17 00:00:00 2001 From: Chauncey Date: Thu, 5 Feb 2026 09:28:36 +0800 Subject: [PATCH 077/810] [Bugfix] fix DeepSeek R1 with CUTLASS MLA Broken on B200 (#33637) Signed-off-by: chaunceyjiang --- vllm/model_executor/layers/attention/mla_attention.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 1b719330e2a..febad382162 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -293,7 +293,6 @@ class MLAAttention(nn.Module, AttentionLayerBase): prefix: str = "", use_sparse: bool = False, indexer: object | None = None, - q_pad_num_heads: int | None = None, **extra_impl_args, ): super().__init__() @@ -308,7 +307,6 @@ class MLAAttention(nn.Module, AttentionLayerBase): self.head_size = kv_lora_rank + qk_rope_head_dim self.layer_name = prefix self.indexer = indexer - self.q_pad_num_heads = q_pad_num_heads self.num_kv_heads = 1 self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim @@ -375,10 +373,9 @@ class MLAAttention(nn.Module, AttentionLayerBase): v_head_dim=self.v_head_dim, kv_b_proj=kv_b_proj, indexer=indexer, - q_pad_num_heads=q_pad_num_heads, **extra_impl_args, ) - + self.q_pad_num_heads = getattr(self.impl, "q_pad_num_heads", None) self.use_direct_call = not current_platform.opaque_attention_op() compilation_config = get_current_vllm_config().compilation_config From 72bb24e2db2acd98a49adcb9e3f1dc6f1bbef4c0 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Wed, 4 Feb 2026 18:07:35 -0800 Subject: [PATCH 078/810] [release] Minor fixes to release annotation (#33849) Signed-off-by: Kevin H. Luu --- .buildkite/scripts/annotate-release.sh | 49 ++++++++++++++------------ 1 file changed, 27 insertions(+), 22 deletions(-) diff --git a/.buildkite/scripts/annotate-release.sh b/.buildkite/scripts/annotate-release.sh index 19e5a9036ca..fe73ea6428e 100755 --- a/.buildkite/scripts/annotate-release.sh +++ b/.buildkite/scripts/annotate-release.sh @@ -27,7 +27,7 @@ aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cpu-cp38- To download and upload the image: \`\`\` -Download images: +# Download images: docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64 docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64 @@ -35,8 +35,12 @@ docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64 docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu130 docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm-base docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm +docker pull public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION} +docker pull public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:v${RELEASE_VERSION} -Tag and push images: +# Tag and push images: + +## CUDA docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64 vllm/vllm-openai:x86_64 docker tag vllm/vllm-openai:x86_64 vllm/vllm-openai:latest-x86_64 @@ -62,34 +66,21 @@ docker tag vllm/vllm-openai:aarch64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-a docker push vllm/vllm-openai:latest-aarch64-cu130 docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130 -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-rocm -docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-rocm vllm/vllm-openai-rocm:latest -docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-rocm vllm/vllm-openai-rocm:v${RELEASE_VERSION}-rocm +## ROCm + +docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm vllm/vllm-openai-rocm:${BUILDKITE_COMMIT} +docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT} vllm/vllm-openai-rocm:latest +docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT} vllm/vllm-openai-rocm:v${RELEASE_VERSION} docker push vllm/vllm-openai-rocm:latest -docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}-rocm +docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION} -Create multi-arch manifest: docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm-base vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:latest-base docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-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 manifest rm vllm/vllm-openai:latest -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 push vllm/vllm-openai:latest -docker manifest push vllm/vllm-openai:v${RELEASE_VERSION} - -docker manifest rm vllm/vllm-openai:latest-cu130 -docker manifest create vllm/vllm-openai:latest-cu130 vllm/vllm-openai:latest-x86_64-cu130 vllm/vllm-openai:latest-aarch64-cu130 -docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130 -docker manifest push vllm/vllm-openai:latest-cu130 -docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu130 - -# CPU images (vllm/vllm-openai-cpu) -docker pull public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION} -docker pull public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:v${RELEASE_VERSION} +## CPU docker tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION} vllm/vllm-openai-cpu:x86_64 docker tag vllm/vllm-openai-cpu:x86_64 vllm/vllm-openai-cpu:latest-x86_64 @@ -103,6 +94,20 @@ docker tag vllm/vllm-openai-cpu:arm64 vllm/vllm-openai-cpu:v${RELEASE_VERSION}-a docker push vllm/vllm-openai-cpu:latest-arm64 docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64 +# Create multi-arch manifest: + +docker manifest rm vllm/vllm-openai:latest +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 push vllm/vllm-openai:latest +docker manifest push vllm/vllm-openai:v${RELEASE_VERSION} + +docker manifest rm vllm/vllm-openai:latest-cu130 +docker manifest create vllm/vllm-openai:latest-cu130 vllm/vllm-openai:latest-x86_64-cu130 vllm/vllm-openai:latest-aarch64-cu130 +docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130 +docker manifest push vllm/vllm-openai:latest-cu130 +docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu130 + docker manifest rm vllm/vllm-openai-cpu:latest || 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 From fb1270f1f821603402a8868e3067b3c3342455e7 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Wed, 4 Feb 2026 21:14:06 -0600 Subject: [PATCH 079/810] [CI][Bugfix]: return McpCall for built-in MCP tools in non-streaming mode (#32762) Signed-off-by: Andreas Karatzas --- .../openai/responses/test_harmony.py | 22 ++++-- .../openai/responses/test_mcp_tools.py | 1 + .../openai/responses/test_parsable_context.py | 8 +- .../openai/responses/test_simple.py | 2 +- tests/utils.py | 78 ++++++++++++++++++- .../openai/parser/harmony_utils.py | 48 ++++++++---- 6 files changed, 131 insertions(+), 28 deletions(-) diff --git a/tests/entrypoints/openai/responses/test_harmony.py b/tests/entrypoints/openai/responses/test_harmony.py index e99a299c12e..b6842f3db1f 100644 --- a/tests/entrypoints/openai/responses/test_harmony.py +++ b/tests/entrypoints/openai/responses/test_harmony.py @@ -62,7 +62,7 @@ async def client(server): async def test_basic(client: OpenAI, model_name: str): response = await client.responses.create( model=model_name, - input="What is 13 * 24?", + input="What is 123 * 456?", ) assert response is not None print("response: ", response) @@ -74,7 +74,7 @@ async def test_basic(client: OpenAI, model_name: str): async def test_basic_with_instructions(client: OpenAI, model_name: str): response = await client.responses.create( model=model_name, - input="What is 13 * 24?", + input="What is 123 * 456?", instructions="Respond in Korean.", ) assert response is not None @@ -116,7 +116,7 @@ async def test_chat(client: OpenAI, model_name: str): {"role": "system", "content": "Respond in Korean."}, {"role": "user", "content": "Hello!"}, {"role": "assistant", "content": "Hello! How can I help you today?"}, - {"role": "user", "content": "What is 13 * 24? Explain your answer."}, + {"role": "user", "content": "What is 123 * 456? Explain your answer."}, ], ) assert response is not None @@ -131,7 +131,7 @@ async def test_chat_with_input_type(client: OpenAI, model_name: str): input=[ { "role": "user", - "content": [{"type": "input_text", "text": "What is 13*24?"}], + "content": [{"type": "input_text", "text": "What is 123 * 456?"}], }, ], ) @@ -200,7 +200,7 @@ async def test_store(client: OpenAI, model_name: str): for store in [True, False]: response = await client.responses.create( model=model_name, - input="What is 13 * 24?", + input="What is 123 * 456?", store=store, ) assert response is not None @@ -219,7 +219,7 @@ async def test_store(client: OpenAI, model_name: str): async def test_background(client: OpenAI, model_name: str): response = await client.responses.create( model=model_name, - input="What is 13 * 24?", + input="What is 123 * 456?", background=True, ) assert response is not None @@ -256,7 +256,7 @@ async def test_background_cancel(client: OpenAI, model_name: str): async def test_stateful_multi_turn(client: OpenAI, model_name: str): response1 = await client.responses.create( model=model_name, - input="What is 13 * 24?", + input="What is 123 * 456?", ) assert response1 is not None assert response1.status == "completed" @@ -361,7 +361,7 @@ async def test_streaming(client: OpenAI, model_name: str, background: bool): # TODO: Add back when web search and code interpreter are available in CI prompts = [ "tell me a story about a cat in 20 words", - "What is 13 * 24? Use python to calculate the result.", + "What is 123 * 456? Use python to calculate the result.", # "When did Jensen found NVIDIA? Search it and answer the year only.", ] @@ -976,6 +976,9 @@ async def test_mcp_code_interpreter_streaming(client: OpenAI, model_name: str, s @pytest.mark.asyncio @pytest.mark.parametrize("model_name", [MODEL_NAME]) +@pytest.mark.dependency( + depends=["test_mcp_code_interpreter_streaming[openai/gpt-oss-20b]"] +) async def test_mcp_tool_multi_turn(client: OpenAI, model_name: str, server): """Test MCP tool calling across multiple turns. @@ -1117,8 +1120,10 @@ async def test_function_call_with_previous_input_messages( model=model_name, input="What is the horoscope for Aquarius today?", tools=tools, + temperature=0.0, extra_body={"enable_response_messages": True}, stream=True, + max_output_tokens=1000, ) response = None @@ -1170,6 +1175,7 @@ async def test_function_call_with_previous_input_messages( stream_response_2 = await client.responses.create( model=model_name, tools=tools, + temperature=0.0, input="", extra_body={ "previous_input_messages": previous_messages, diff --git a/tests/entrypoints/openai/responses/test_mcp_tools.py b/tests/entrypoints/openai/responses/test_mcp_tools.py index 0cc1ac9bac5..9658f5d90ea 100644 --- a/tests/entrypoints/openai/responses/test_mcp_tools.py +++ b/tests/entrypoints/openai/responses/test_mcp_tools.py @@ -160,6 +160,7 @@ class TestMCPEnabled: "No developer messages should be present with valid mcp tool" ) + @pytest.mark.flaky(reruns=3) @pytest.mark.asyncio @pytest.mark.parametrize("model_name", [MODEL_NAME]) async def test_mcp_tool_with_allowed_tools_star( diff --git a/tests/entrypoints/openai/responses/test_parsable_context.py b/tests/entrypoints/openai/responses/test_parsable_context.py index c1f0f435b83..0d50f1251a6 100644 --- a/tests/entrypoints/openai/responses/test_parsable_context.py +++ b/tests/entrypoints/openai/responses/test_parsable_context.py @@ -53,7 +53,7 @@ async def client(server): async def test_basic(client: OpenAI, model_name: str): response = await client.responses.create( model=model_name, - input="What is 13 * 24?", + input="What is 123 * 456?", ) assert response is not None print("response: ", response) @@ -164,7 +164,7 @@ async def test_function_call_first_turn(client: OpenAI, model_name: str): async def test_mcp_tool_call(client: OpenAI, model_name: str): response = await client.responses.create( model=model_name, - input="What is 13 * 24? Use python to calculate the result.", + input="What is 123 * 456? Use python to calculate the result.", tools=[{"type": "code_interpreter", "container": {"type": "auto"}}], extra_body={"enable_response_messages": True}, temperature=0.0, @@ -179,12 +179,12 @@ async def test_mcp_tool_call(client: OpenAI, model_name: str): assert response.output[2].type == "reasoning" # make sure the correct math is in the final output assert response.output[3].type == "message" - assert "312" in response.output[3].content[0].text + assert "56088" in response.output[3].content[0].text # test raw input_messages / output_messages assert len(response.input_messages) == 1 assert len(response.output_messages) == 3 - assert "312" in response.output_messages[2]["message"] + assert "56088" in response.output_messages[2]["message"] @pytest.mark.asyncio diff --git a/tests/entrypoints/openai/responses/test_simple.py b/tests/entrypoints/openai/responses/test_simple.py index 8f07b02a308..a5bec6dfd89 100644 --- a/tests/entrypoints/openai/responses/test_simple.py +++ b/tests/entrypoints/openai/responses/test_simple.py @@ -34,7 +34,7 @@ async def client(server): async def test_basic(client: OpenAI, model_name: str): response = await client.responses.create( model=model_name, - input="What is 13 * 24?", + input="What is 123 * 456?", ) assert response is not None print("response: ", response) diff --git a/tests/utils.py b/tests/utils.py index efeceba63bc..5252115f291 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -197,12 +197,86 @@ class RemoteOpenAIServer: return self def __exit__(self, exc_type, exc_value, traceback): + pid = self.proc.pid + # Graceful shutdown self.proc.terminate() try: - self.proc.wait(8) + self.proc.wait(timeout=15) + print(f"[RemoteOpenAIServer] Server {pid} terminated gracefully") except subprocess.TimeoutExpired: - # force kill if needed + print( + f"[RemoteOpenAIServer] Server {pid} did not respond " + "to SIGTERM, sending SIGKILL" + ) self.proc.kill() + try: + self.proc.wait(timeout=5) + print(f"[RemoteOpenAIServer] Server {pid} killed") + except subprocess.TimeoutExpired as err: + raise RuntimeError( + f"[RemoteOpenAIServer] Failed to kill server process {pid}" + ) from err + # Wait for GPU memory to be released + self._wait_for_gpu_memory_release() + + def _get_gpu_memory_used(self) -> float | None: + """Get total GPU memory used across all visible devices in bytes.""" + try: + if current_platform.is_rocm(): + with _nvml(): + handles = amdsmi_get_processor_handles() + total_used = 0 + for handle in handles: + vram_info = amdsmi_get_gpu_vram_usage(handle) + total_used += vram_info["vram_used"] + return total_used + elif current_platform.is_cuda(): + with _nvml(): + total_used = 0 + device_count = cuda_device_count_stateless() + for i in range(device_count): + handle = nvmlDeviceGetHandleByIndex(i) + mem_info = nvmlDeviceGetMemoryInfo(handle) + total_used += mem_info.used + return total_used + except Exception as e: + print(f"[RemoteOpenAIServer] Could not query GPU memory: {e}") + return None + return None + + def _wait_for_gpu_memory_release(self, timeout: float = 30.0): + """Poll GPU memory until it stabilizes, indicating cleanup is complete.""" + start = time.time() + prev_used: float | None = None + stable_count = 0 + + while time.time() - start < timeout: + used = self._get_gpu_memory_used() + + if used is None: + return # Can't query, assume ok + + if prev_used is not None and abs(used - prev_used) < 100 * 1024 * 1024: + stable_count += 1 + if stable_count >= 3: + used_gb = used / 1e9 + print( + f"[RemoteOpenAIServer] GPU memory stabilized " + f"at {used_gb:.2f} GB" + ) + return + else: + stable_count = 0 + + prev_used = used + time.sleep(0.1) + + last_reading = prev_used / 1e9 if prev_used is not None else 0.0 + raise RuntimeError( + f"[RemoteOpenAIServer] GPU memory did not stabilize within {timeout}s. " + f"Last reading: {last_reading:.2f} GB. " + "Child processes may still be holding GPU memory." + ) def _poll(self) -> int | None: """Subclasses override this method to customize process polling""" diff --git a/vllm/entrypoints/openai/parser/harmony_utils.py b/vllm/entrypoints/openai/parser/harmony_utils.py index 58ba9fee442..3bb81273878 100644 --- a/vllm/entrypoints/openai/parser/harmony_utils.py +++ b/vllm/entrypoints/openai/parser/harmony_utils.py @@ -68,6 +68,14 @@ MCP_BUILTIN_TOOLS: set[str] = { "container", } +# Mapping from built-in tool recipient names to their MCP server labels. +# This ensures consistency between streaming and non-streaming responses. +_BUILTIN_TOOL_TO_MCP_SERVER_LABEL: dict[str, str] = { + "python": "code_interpreter", + "browser": "web_search_preview", + "container": "container", +} + def has_custom_tools(tool_types: set[str]) -> bool: """ @@ -601,7 +609,13 @@ def _parse_mcp_recipient(recipient: str) -> tuple[str, str]: def _parse_mcp_call(message: Message, recipient: str) -> list[ResponseOutputItem]: """Parse MCP calls into MCP call items.""" - server_label, tool_name = _parse_mcp_recipient(recipient) + # Handle built-in tools that need server_label mapping + if recipient in _BUILTIN_TOOL_TO_MCP_SERVER_LABEL: + server_label = _BUILTIN_TOOL_TO_MCP_SERVER_LABEL[recipient] + tool_name = recipient + else: + server_label, tool_name = _parse_mcp_recipient(recipient) + output_items = [] for content in message.content: response_item = McpCall( @@ -630,7 +644,7 @@ def parse_output_message(message: Message) -> list[ResponseOutputItem]: recipient = message.recipient if recipient is not None: - # Browser tool calls + # Browser tool calls (browser.search, browser.open, browser.find) if recipient.startswith("browser."): output_items.append(_parse_browser_tool_call(message, recipient)) @@ -638,10 +652,8 @@ def parse_output_message(message: Message) -> list[ResponseOutputItem]: elif message.channel == "commentary" and recipient.startswith("functions."): output_items.extend(_parse_function_call(message, recipient)) - # Built-in tools are treated as reasoning - elif recipient.startswith(("python", "browser", "container")): - # Built-in tool recipients (python/browser/container) - # generate reasoning output + # Built-in MCP tools (python, browser, container) + elif recipient in _BUILTIN_TOOL_TO_MCP_SERVER_LABEL: output_items.extend(_parse_reasoning(message)) # All other recipients are MCP calls @@ -688,13 +700,23 @@ def parse_remaining_state(parser: StreamableParser) -> list[ResponseOutputItem]: status="in_progress", ) ] - # Built-in tools (python, browser, container) should be treated as reasoning - elif not ( - current_recipient.startswith("python") - or current_recipient.startswith("browser") - or current_recipient.startswith("container") - ): - # All other recipients are MCP calls + # Built-in MCP tools (python, browser, container) + elif current_recipient in _BUILTIN_TOOL_TO_MCP_SERVER_LABEL: + return [ + ResponseReasoningItem( + id=f"rs_{random_uuid()}", + summary=[], + type="reasoning", + content=[ + ResponseReasoningTextContent( + text=parser.current_content, type="reasoning_text" + ) + ], + status=None, + ) + ] + # All other recipients are MCP calls + else: rid = random_uuid() server_label, tool_name = _parse_mcp_recipient(current_recipient) return [ From e3bf79ffa080a5052aa61fce71b70b11fb7f9d1e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Luka=20Govedi=C4=8D?= Date: Wed, 4 Feb 2026 22:54:27 -0500 Subject: [PATCH 080/810] Revert "[Attention][FA3] Update FA3 to include new swizzle optimization" (#33841) --- cmake/external_projects/vllm_flash_attn.cmake | 2 +- vllm/v1/attention/backends/flash_attn.py | 7 +------ vllm/v1/attention/backends/mla/flashattn_mla.py | 7 +------ 3 files changed, 3 insertions(+), 13 deletions(-) diff --git a/cmake/external_projects/vllm_flash_attn.cmake b/cmake/external_projects/vllm_flash_attn.cmake index dbdfd5e8144..b51934a3ab2 100644 --- a/cmake/external_projects/vllm_flash_attn.cmake +++ b/cmake/external_projects/vllm_flash_attn.cmake @@ -38,7 +38,7 @@ else() FetchContent_Declare( vllm-flash-attn GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git - GIT_TAG 2adfc8c2177c5b0e8ddeedfd5a8990d80eb496ff + GIT_TAG 188be16520ceefdc625fdf71365585d2ee348fe2 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/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index 232b0b0daff..9275725314e 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -308,15 +308,10 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad self.compilation_config.cudagraph_mode.has_full_cudagraphs() ) self.max_cudagraph_size = self.compilation_config.max_cudagraph_capture_size - max_num_seqs = vllm_config.scheduler_config.max_num_seqs if self.use_full_cuda_graph and self.aot_schedule: - # Times 4 due to: - # https://github.com/vllm-project/flash-attention/blob/3223650ccabe622a0fcae65eec706a50186a89f7/hopper/flash_api.cpp#L650-L653 - # For some tests max_cudagraph_size > max_num_seqs, - # so we need to use the larger one. self.scheduler_metadata = torch.zeros( - max(self.max_cudagraph_size or 0, max_num_seqs) * 4 + 1, + vllm_config.scheduler_config.max_num_seqs + 1, dtype=torch.int32, device=self.device, ) diff --git a/vllm/v1/attention/backends/mla/flashattn_mla.py b/vllm/v1/attention/backends/mla/flashattn_mla.py index f0ba259362f..e160d325568 100644 --- a/vllm/v1/attention/backends/mla/flashattn_mla.py +++ b/vllm/v1/attention/backends/mla/flashattn_mla.py @@ -127,15 +127,10 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata] self.compilation_config.cudagraph_mode.has_full_cudagraphs() ) self.max_cudagraph_size = self.compilation_config.max_cudagraph_capture_size - max_num_seqs = vllm_config.scheduler_config.max_num_seqs if self.use_full_cuda_graph and self.fa_aot_schedule: - # Times 4 due to: - # https://github.com/vllm-project/flash-attention/blob/3223650ccabe622a0fcae65eec706a50186a89f7/hopper/flash_api.cpp#L650-L653 - # For some tests max_cudagraph_size > max_num_seqs, - # so we need to use the larger one. self.scheduler_metadata = torch.zeros( - max(self.max_cudagraph_size or 0, max_num_seqs) * 4 + 1, + vllm_config.scheduler_config.max_num_seqs + 1, dtype=torch.int32, device=self.device, ) From add9f1fbd920611c2b909fe10d9b44b59650f8b7 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Wed, 4 Feb 2026 20:38:20 -0800 Subject: [PATCH 081/810] [Minor] Include `StreamingInput` in inputs package (#33856) Signed-off-by: Nick Hill --- tests/v1/e2e/test_streaming_input.py | 2 +- tests/v1/streaming_input/test_async_llm_streaming.py | 2 +- vllm/inputs/__init__.py | 2 ++ vllm/v1/engine/async_llm.py | 4 ++-- 4 files changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/v1/e2e/test_streaming_input.py b/tests/v1/e2e/test_streaming_input.py index a1eaa065a63..4c9b43099e4 100644 --- a/tests/v1/e2e/test_streaming_input.py +++ b/tests/v1/e2e/test_streaming_input.py @@ -19,7 +19,7 @@ import pytest import pytest_asyncio from vllm import SamplingParams -from vllm.inputs.data import StreamingInput +from vllm.inputs import StreamingInput from vllm.outputs import RequestOutput from vllm.platforms import current_platform from vllm.sampling_params import RequestOutputKind diff --git a/tests/v1/streaming_input/test_async_llm_streaming.py b/tests/v1/streaming_input/test_async_llm_streaming.py index 99263438717..b5ba757d0a9 100644 --- a/tests/v1/streaming_input/test_async_llm_streaming.py +++ b/tests/v1/streaming_input/test_async_llm_streaming.py @@ -7,7 +7,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest -from vllm.inputs.data import StreamingInput +from vllm.inputs import StreamingInput from vllm.outputs import RequestOutput from vllm.sampling_params import RequestOutputKind, SamplingParams from vllm.v1.engine.async_llm import AsyncLLM diff --git a/vllm/inputs/__init__.py b/vllm/inputs/__init__.py index d9aed70c9b9..0fdb3ab5ea4 100644 --- a/vllm/inputs/__init__.py +++ b/vllm/inputs/__init__.py @@ -12,6 +12,7 @@ from .data import ( PromptType, SingletonInputs, SingletonPrompt, + StreamingInput, TextPrompt, TokenInputs, TokensPrompt, @@ -41,4 +42,5 @@ __all__ = [ "build_explicit_enc_dec_prompt", "to_enc_dec_tuple_list", "zip_enc_dec_prompts", + "StreamingInput", ] diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index f1a3e341fd9..2beb9c4f8c7 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -16,8 +16,7 @@ from vllm import TokensPrompt from vllm.config import VllmConfig from vllm.engine.arg_utils import AsyncEngineArgs from vllm.engine.protocol import EngineClient -from vllm.inputs import PromptType -from vllm.inputs.data import StreamingInput +from vllm.inputs import PromptType, StreamingInput from vllm.logger import init_logger from vllm.lora.request import LoRARequest from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry @@ -461,6 +460,7 @@ class AsyncLLM(EngineClient): self._validate_streaming_input_sampling_params(sp) else: sp = sampling_params + # TODO(nick): Avoid re-validating reused sampling parameters req = self.input_processor.process_inputs( request_id=internal_req_id, prompt=input_chunk.prompt, From 007b183d745f5b37aeb6cdf936c3b590b0c29fde Mon Sep 17 00:00:00 2001 From: rinbaro Date: Wed, 4 Feb 2026 20:50:59 -0800 Subject: [PATCH 082/810] [docs] fix unintentional misspellings (#33863) Signed-off-by: rinbaro --- docs/contributing/model/basic.md | 2 +- docs/contributing/model/multimodal.md | 2 +- docs/getting_started/quickstart.md | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/contributing/model/basic.md b/docs/contributing/model/basic.md index 624f13bf793..ba1f5e43d61 100644 --- a/docs/contributing/model/basic.md +++ b/docs/contributing/model/basic.md @@ -138,7 +138,7 @@ These models should follow the same instructions as case (1), but they should in For case (3), we recommend looking at the implementation of [`MiniMaxText01ForCausalLM`](../../../vllm/model_executor/models/minimax_text_01.py) or [`Lfm2ForCausalLM`](../../../vllm/model_executor/models/lfm2.py) as a reference, which use custom "mamba-like" layers `MiniMaxText01LinearAttention` and `ShortConv` respectively. Please follow the same guidelines as case (2) for implementing these models. -We use "mamba-like" to refer to layers that posses a state that is updated in-place, rather than being appended-to (like KV cache for attention). +We use "mamba-like" to refer to layers that possess a state that is updated in-place, rather than being appended-to (like KV cache for attention). For implementing new custom mamba-like layers, one should inherit from `MambaBase` and implement the methods `get_state_dtype`, `get_state_shape` to calculate the data types and state shapes at runtime, as well as `mamba_type` and `get_attn_backend`. It is also necessary to implement the "attention meta-data" class which handles the meta-data that is common across all layers. Please see [`LinearAttentionMetadata`](../../../vllm/v1/attention/backends/linear_attn.py) or [`ShortConvAttentionMetadata`](../../../vllm/v1/attention/backends/short_conv_attn.py) for examples of this. diff --git a/docs/contributing/model/multimodal.md b/docs/contributing/model/multimodal.md index c876cc47c11..e123e0dcd15 100644 --- a/docs/contributing/model/multimodal.md +++ b/docs/contributing/model/multimodal.md @@ -739,7 +739,7 @@ Each [PromptUpdate][vllm.multimodal.processing.PromptUpdate] instance specifies ``` However, this is not entirely correct. After `FuyuImageProcessor.preprocess_with_tokenizer_info` is called, - a BOS token (``) is also added to the promopt: + a BOS token (``) is also added to the prompt: ??? code diff --git a/docs/getting_started/quickstart.md b/docs/getting_started/quickstart.md index d5c68172ddb..40b6dab067d 100644 --- a/docs/getting_started/quickstart.md +++ b/docs/getting_started/quickstart.md @@ -57,7 +57,7 @@ This guide will help you quickly get started with vLLM to perform: It currently supports Python 3.12, ROCm 7.0 and `glibc >= 2.35`. !!! note - Note that, previously, docker images were published using AMD's docker release pipeline and were located `rocm/vlm-dev`. This is being deprecated by using vLLM's docker release pipeline. + Note that, previously, docker images were published using AMD's docker release pipeline and were located `rocm/vllm-dev`. This is being deprecated by using vLLM's docker release pipeline. === "Google TPU" From c1395f72cd22d97eb39ecd67d9d22f2af3d20bda Mon Sep 17 00:00:00 2001 From: rasmith Date: Wed, 4 Feb 2026 23:05:48 -0600 Subject: [PATCH 083/810] [CI][AMD][BugFix] Ensure VLLM_ROCM_USE_AITER is set so test_rocm_aiter_topk.py can run correctly (#33840) Signed-off-by: Randall Smith --- tests/kernels/moe/test_rocm_aiter_topk.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/tests/kernels/moe/test_rocm_aiter_topk.py b/tests/kernels/moe/test_rocm_aiter_topk.py index d4724d749fc..070d00f6112 100644 --- a/tests/kernels/moe/test_rocm_aiter_topk.py +++ b/tests/kernels/moe/test_rocm_aiter_topk.py @@ -10,22 +10,28 @@ # and the platform is not ROCm. import importlib.util +import os import pytest import torch +from vllm.platforms import current_platform + +if not current_platform.is_rocm(): + pytest.skip("This test can only run on ROCm.", allow_module_level=True) + +# This environment variable must be set so ops will be registered. +os.environ["VLLM_ROCM_USE_AITER"] = "1" + # this import statement is needed to ensure the ops are registered import vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe # noqa: F401 -from vllm.platforms import current_platform # need to import once to ensure the ops are registered # Check if aiter package is installed aiter_available = importlib.util.find_spec("aiter") is not None -pytestmark = pytest.mark.skipif( - not (current_platform.is_rocm() and aiter_available), - reason="AITER ops are only available on ROCm with aiter package installed", -) +if not aiter_available: + pytest.skip("These tests require AITER to run.", allow_module_level=True) def test_rocm_aiter_biased_grouped_topk_custom_op_registration(): From 9595afda183bdd79b0ee2d38d2b0049fe86e6628 Mon Sep 17 00:00:00 2001 From: Andrew Xia Date: Thu, 5 Feb 2026 00:46:15 -0500 Subject: [PATCH 084/810] [2/N] move responses/serving _make_response_output_items logic to parser (#33281) Signed-off-by: Andrew Xia Signed-off-by: Andrew Xia Co-authored-by: Andrew Xia --- vllm/entrypoints/openai/responses/serving.py | 142 ++++--------- vllm/parser/abstract_parser.py | 200 +++++++++++++++++++ 2 files changed, 242 insertions(+), 100 deletions(-) diff --git a/vllm/entrypoints/openai/responses/serving.py b/vllm/entrypoints/openai/responses/serving.py index 32cce3ef4cf..1ed7a79cc87 100644 --- a/vllm/entrypoints/openai/responses/serving.py +++ b/vllm/entrypoints/openai/responses/serving.py @@ -63,7 +63,6 @@ from vllm.engine.protocol import EngineClient from vllm.entrypoints.chat_utils import ( ChatCompletionMessageParam, ChatTemplateContentFormatOption, - make_tool_call_id, ) from vllm.entrypoints.logger import RequestLogger from vllm.entrypoints.mcp.tool_server import ToolServer @@ -915,114 +914,57 @@ class OpenAIServingResponses(OpenAIServing): final_output: CompletionOutput, tokenizer: TokenizerLike, ) -> list[ResponseOutputItem]: - if self.parser and self.parser.reasoning_parser_cls: - try: - reasoning_parser = self.parser.reasoning_parser_cls(tokenizer) - except RuntimeError as e: - logger.exception("Error in reasoning parser creation.") - raise e - - reasoning, content = reasoning_parser.extract_reasoning( - final_output.text, request=request - ) - else: - reasoning = None - content = final_output.text - # Log complete response if output logging is enabled if self.enable_log_outputs and self.request_logger: - output_text = "" - if content: - output_text = content - elif reasoning: - output_text = f"[reasoning: {reasoning}]" - - if output_text: - self.request_logger.log_outputs( - request_id=request.request_id, - outputs=output_text, - output_token_ids=final_output.token_ids, - finish_reason=final_output.finish_reason, - is_streaming=False, - delta=False, - ) - - reasoning_item = None - message_item = None - if reasoning: - reasoning_item = ResponseReasoningItem( - id=f"rs_{random_uuid()}", - summary=[], - type="reasoning", - content=[ - ResponseReasoningTextContent(text=reasoning, type="reasoning_text") - ], - status=None, # NOTE: Only the last output item has status. + self.request_logger.log_outputs( + request_id=request.request_id, + outputs=final_output.text, + output_token_ids=final_output.token_ids, + finish_reason=final_output.finish_reason, + is_streaming=False, + delta=False, ) - tool_calls, content = self._parse_tool_calls_from_content( - request=request, - tokenizer=tokenizer, - content=content, - enable_auto_tools=self.enable_auto_tools, - tool_parser_cls=self.parser.tool_parser_cls if self.parser else None, - ) - if content or (self.use_harmony and tool_calls): - res_text_part = None - if content: - res_text_part = ResponseOutputText( - text=content, - annotations=[], # TODO - type="output_text", - logprobs=( - self._create_response_logprobs( - token_ids=final_output.token_ids, - logprobs=final_output.logprobs, - tokenizer=tokenizer, - top_logprobs=request.top_logprobs, - ) - if request.is_include_output_logprobs() - else None - ), - ) - message_item = ResponseOutputMessage( + # Compute logprobs if requested + logprobs = None + if request.is_include_output_logprobs() and final_output.logprobs: + logprobs = self._create_response_logprobs( + token_ids=final_output.token_ids, + logprobs=final_output.logprobs, + tokenizer=tokenizer, + top_logprobs=request.top_logprobs, + ) + + # Use parser to extract and create response output items + if self.parser: + parser = self.parser(tokenizer) + return parser.extract_response_outputs( + model_output=final_output.text, + request=request, + enable_auto_tools=self.enable_auto_tools, + tool_call_id_type=self.tool_call_id_type, + logprobs=logprobs, + ) + + # Fallback when no parser is configured + return [ + ResponseOutputMessage( id=f"msg_{random_uuid()}", - content=[res_text_part] if res_text_part else [], + content=[ + ResponseOutputText( + text=final_output.text, + annotations=[], + type="output_text", + logprobs=logprobs, + ) + ] + if final_output.text + else [], role="assistant", status="completed", type="message", ) - outputs = [] - - if reasoning_item: - outputs.append(reasoning_item) - if message_item: - outputs.append(message_item) - if tool_calls: - # We use a simple counter for history_tool_call_count because - # we don't track the history of tool calls in the Responses API yet. - # This means that the tool call index will start from 0 for each - # request. - tool_call_items = [] - for history_tool_call_cnt, tool_call in enumerate(tool_calls): - tool_call_items.append( - ResponseFunctionToolCall( - id=f"fc_{random_uuid()}", - call_id=tool_call.id - if tool_call.id - else make_tool_call_id( - id_type=self.tool_call_id_type, - func_name=tool_call.name, - idx=history_tool_call_cnt, - ), - type="function_call", - status="completed", - name=tool_call.name, - arguments=tool_call.arguments, - ) - ) - outputs.extend(tool_call_items) - return outputs + ] def _make_response_output_items_with_harmony( self, diff --git a/vllm/parser/abstract_parser.py b/vllm/parser/abstract_parser.py index f5cd1430a18..aa145bab212 100644 --- a/vllm/parser/abstract_parser.py +++ b/vllm/parser/abstract_parser.py @@ -1,23 +1,46 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json from abc import abstractmethod from collections.abc import Sequence from functools import cached_property +from openai.types.responses import ( + ResponseFunctionToolCall, + ResponseOutputItem, + ResponseOutputMessage, + ResponseOutputText, + ResponseReasoningItem, + ToolChoiceFunction, +) +from openai.types.responses.response_output_text import Logprob +from openai.types.responses.response_reasoning_item import ( + Content as ResponseReasoningTextContent, +) +from pydantic import TypeAdapter + +from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionNamedToolChoiceParam, ChatCompletionRequest, ) from vllm.entrypoints.openai.engine.protocol import ( DeltaMessage, ExtractedToolCallInformation, + FunctionCall, + FunctionDefinition, ) from vllm.entrypoints.openai.responses.protocol import ( ResponsesRequest, ) +from vllm.logger import init_logger from vllm.reasoning.abs_reasoning_parsers import ReasoningParser from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import ToolParser +from vllm.utils import random_uuid + +logger = init_logger(__name__) class Parser: @@ -128,6 +151,33 @@ class Parser: The extracted content token IDs. """ + @abstractmethod + def extract_response_outputs( + self, + model_output: str, + request: ResponsesRequest, + enable_auto_tools: bool = False, + tool_call_id_type: str = "random", + logprobs: list[Logprob] | None = None, + ) -> list[ResponseOutputItem]: + """ + Extract reasoning, content, and tool calls from a complete + model-generated string and return as ResponseOutputItem objects. + + Used for non-streaming responses where we have the entire model + response available before sending to the client. + + Args: + model_output: The complete model-generated string. + request: The request object used to generate the output. + enable_auto_tools: Whether to enable automatic tool call parsing. + tool_call_id_type: Type of tool call ID generation ("random", etc). + logprobs: Pre-computed logprobs for the output text, if any. + + Returns: + A list of ResponseOutputItem objects. + """ + @abstractmethod def extract_reasoning( self, @@ -260,6 +310,156 @@ class DelegatingParser(Parser): return None, model_output return self._reasoning_parser.extract_reasoning(model_output, request) + def extract_response_outputs( + self, + model_output: str, + request: ResponsesRequest, + enable_auto_tools: bool = False, + tool_call_id_type: str = "random", + logprobs: list[Logprob] | None = None, + ) -> list[ResponseOutputItem]: + # First extract reasoning + reasoning, content = self.extract_reasoning(model_output, request) + + # Then parse tool calls from the content + tool_calls, content = self._parse_tool_calls( + request=request, + content=content, + enable_auto_tools=enable_auto_tools, + ) + + # Build output items + outputs: list[ResponseOutputItem] = [] + + # Add reasoning item if present + if reasoning: + reasoning_item = ResponseReasoningItem( + id=f"rs_{random_uuid()}", + summary=[], + type="reasoning", + content=[ + ResponseReasoningTextContent(text=reasoning, type="reasoning_text") + ], + status=None, # NOTE: Only the last output item has status. + ) + outputs.append(reasoning_item) + + # Add message item if there's content + if content: + res_text_part = ResponseOutputText( + text=content, + annotations=[], + type="output_text", + logprobs=logprobs, + ) + message_item = ResponseOutputMessage( + id=f"msg_{random_uuid()}", + content=[res_text_part], + role="assistant", + status="completed", + type="message", + ) + outputs.append(message_item) + + if tool_calls: + # We use a simple counter for history_tool_call_count because + # we don't track the history of tool calls in the Responses API yet. + # This means that the tool call index will start from 0 for each + # request. + for history_tool_call_cnt, tool_call in enumerate(tool_calls): + tool_call_item = ResponseFunctionToolCall( + id=f"fc_{random_uuid()}", + call_id=tool_call.id + if tool_call.id + else make_tool_call_id( + id_type=tool_call_id_type, + func_name=tool_call.name, + idx=history_tool_call_cnt, + ), + type="function_call", + status="completed", + name=tool_call.name, + arguments=tool_call.arguments, + ) + outputs.append(tool_call_item) + + return outputs + + def _parse_tool_calls( + self, + request: ResponsesRequest, + content: str | None, + enable_auto_tools: bool, + ) -> tuple[list[FunctionCall], str | None]: + """ + TODO(qandrew): merge _parse_tool_calls_from_content + for ChatCompletions into this function + Parse tool calls from content based on request tool_choice settings. + + Returns: + A tuple of (function_calls, remaining_content) if tool calls + were parsed + """ + function_calls: list[FunctionCall] = [] + + if request.tool_choice and isinstance(request.tool_choice, ToolChoiceFunction): + # Forced Function Call (Responses API style) + assert content is not None + function_calls.append( + FunctionCall(name=request.tool_choice.name, arguments=content) + ) + return function_calls, None # Clear content since tool is called. + + if request.tool_choice and isinstance( + request.tool_choice, ChatCompletionNamedToolChoiceParam + ): + # Forced Function Call (Chat Completion API style) + assert content is not None + function_calls.append( + FunctionCall(name=request.tool_choice.function.name, arguments=content) + ) + return function_calls, None # Clear content since tool is called. + + if request.tool_choice == "required": + # Required tool calls - parse JSON + assert content is not None + tool_calls = TypeAdapter(list[FunctionDefinition]).validate_json(content) + function_calls.extend( + FunctionCall( + name=tool_call.name, + arguments=json.dumps(tool_call.parameters, ensure_ascii=False), + ) + for tool_call in tool_calls + ) + return function_calls, None # Clear content since tool is called. + + if ( + self._tool_parser is not None + and enable_auto_tools + and (request.tool_choice == "auto" or request.tool_choice is None) + ): + # Automatic Tool Call Parsing + tool_call_info = self._tool_parser.extract_tool_calls( + content if content is not None else "", + request=request, # type: ignore + ) + if tool_call_info is not None and tool_call_info.tools_called: + function_calls.extend( + FunctionCall( + id=tool_call.id, + name=tool_call.function.name, + arguments=tool_call.function.arguments, + ) + for tool_call in tool_call_info.tool_calls + ) + remaining_content = tool_call_info.content + if remaining_content and remaining_content.strip() == "": + remaining_content = None + return function_calls, remaining_content + + # No tool calls + return [], content + def extract_reasoning_streaming( self, previous_text: str, From 07daee132b30140bb7c5b28d7f8c856036d2baad Mon Sep 17 00:00:00 2001 From: "Li, Jiang" Date: Thu, 5 Feb 2026 13:53:48 +0800 Subject: [PATCH 085/810] [CI/Build] Parallelize CPU CI tests (#33778) Signed-off-by: jiang1.li --- .buildkite/hardware_tests/arm.yaml | 8 -- .buildkite/hardware_tests/cpu.yaml | 100 +++++++++++++++ .buildkite/hardware_tests/intel.yaml | 7 -- .../run-cpu-distributed-smoke-test.sh | 26 ++++ .../scripts/hardware_ci/run-cpu-test.sh | 118 ++---------------- vllm/v1/worker/cpu_worker.py | 28 ++++- 6 files changed, 157 insertions(+), 130 deletions(-) delete mode 100644 .buildkite/hardware_tests/arm.yaml create mode 100644 .buildkite/hardware_tests/cpu.yaml create mode 100644 .buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh diff --git a/.buildkite/hardware_tests/arm.yaml b/.buildkite/hardware_tests/arm.yaml deleted file mode 100644 index d39ab4a7e44..00000000000 --- a/.buildkite/hardware_tests/arm.yaml +++ /dev/null @@ -1,8 +0,0 @@ -group: Hardware -steps: - - label: "Arm CPU Test" - soft_fail: true - device: arm_cpu - no_plugin: true - commands: - - bash .buildkite/scripts/hardware_ci/run-cpu-test-arm.sh diff --git a/.buildkite/hardware_tests/cpu.yaml b/.buildkite/hardware_tests/cpu.yaml new file mode 100644 index 00000000000..39a5516967f --- /dev/null +++ b/.buildkite/hardware_tests/cpu.yaml @@ -0,0 +1,100 @@ +group: CPU +depends_on: [] +steps: +- label: CPU-Kernel Tests + depends_on: [] + soft_fail: true + device: intel_cpu + no_plugin: true + source_file_dependencies: + - csrc/cpu/ + - cmake/cpu_extension.cmake + - CMakeLists.txt + - vllm/_custom_ops.py + - tests/kernels/attention/test_cpu_attn.py + - tests/kernels/moe/test_cpu_fused_moe.py + - tests/kernels/test_onednn.py + commands: + - | + bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 20m " + pytest -x -v -s tests/kernels/attention/test_cpu_attn.py + pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py + pytest -x -v -s tests/kernels/test_onednn.py" + +- label: CPU-Language Generation and Pooling Model Tests + depends_on: [] + soft_fail: true + device: intel_cpu + no_plugin: true + source_file_dependencies: + - csrc/cpu/ + - vllm/ + - tests/models/language/generation/ + - tests/models/language/pooling/ + commands: + - | + bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m " + pytest -x -v -s tests/models/language/generation -m cpu_model + pytest -x -v -s tests/models/language/pooling -m cpu_model" + +- label: CPU-Quantization Model Tests + depends_on: [] + soft_fail: true + device: intel_cpu + no_plugin: true + source_file_dependencies: + - csrc/cpu/ + - vllm/model_executor/layers/quantization/cpu_wna16.py + - vllm/model_executor/layers/quantization/gptq_marlin.py + - vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8.py + - vllm/model_executor/layers/quantization/kernels/scaled_mm/cpu.py + - vllm/model_executor/layers/quantization/kernels/mixed_precision/cpu.py + - tests/quantization/test_compressed_tensors.py + - tests/quantization/test_cpu_wna16.py + commands: + - | + bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 20m " + pytest -x -v -s tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_logprobs + pytest -x -v -s tests/quantization/test_cpu_wna16.py" + +- label: CPU-TP/DP/PP Tests + depends_on: [] + soft_fail: true + device: intel_cpu + no_plugin: true + source_file_dependencies: + - csrc/cpu/shm.cpp + - vllm/v1/worker/cpu_worker.py + - vllm/v1/worker/gpu_worker.py + - vllm/v1/worker/cpu_model_runner.py + - vllm/v1/worker/gpu_model_runner.py + - vllm/platforms/cpu.py + - vllm/distributed/parallel_state.py + - vllm/distributed/device_communicators/cpu_communicator.py + commands: + - | + bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 10m " + bash .buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh" + +- label: CPU-Multi-Modal Model Tests %N + depends_on: [] + soft_fail: true + device: intel_cpu + no_plugin: true + source_file_dependencies: + # - vllm/ + - vllm/model_executor/layers/rotary_embedding + - tests/models/multimodal/generation/ + commands: + - | + bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m " + pytest -x -v -s tests/models/multimodal/generation --ignore=tests/models/multimodal/generation/test_pixtral.py -m cpu_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB" + parallelism: 2 + +- label: "Arm CPU Test" + depends_on: [] + soft_fail: true + device: arm_cpu + no_plugin: true + commands: + - bash .buildkite/scripts/hardware_ci/run-cpu-test-arm.sh diff --git a/.buildkite/hardware_tests/intel.yaml b/.buildkite/hardware_tests/intel.yaml index 76bf2e0be28..ba0088b3af6 100644 --- a/.buildkite/hardware_tests/intel.yaml +++ b/.buildkite/hardware_tests/intel.yaml @@ -1,13 +1,6 @@ group: Hardware depends_on: ~ steps: - - label: "Intel CPU Test" - soft_fail: true - device: intel_cpu - no_plugin: true - commands: - - bash .buildkite/scripts/hardware_ci/run-cpu-test.sh - - label: "Intel HPU Test" soft_fail: true device: intel_hpu diff --git a/.buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh b/.buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh new file mode 100644 index 00000000000..3caa49832c3 --- /dev/null +++ b/.buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh @@ -0,0 +1,26 @@ +#!/bin/bash +set -euox pipefail + +echo "--- PP+TP" +vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -pp=2 & +server_pid=$! +timeout 600 bash -c "until curl localhost:8000/v1/models; do sleep 1; done" || exit 1 +vllm bench serve \ + --backend vllm \ + --dataset-name random \ + --model meta-llama/Llama-3.2-3B-Instruct \ + --num-prompts 20 \ + --endpoint /v1/completions +kill -s SIGTERM $server_pid & + +echo "--- DP+TP" +vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 & +server_pid=$! +timeout 600 bash -c "until curl localhost:8000/v1/models; do sleep 1; done" || exit 1 +vllm bench serve \ + --backend vllm \ + --dataset-name random \ + --model meta-llama/Llama-3.2-3B-Instruct \ + --num-prompts 20 \ + --endpoint /v1/completions +kill -s SIGTERM $server_pid & diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test.sh b/.buildkite/scripts/hardware_ci/run-cpu-test.sh index ee6510bf88e..c32b051cabc 100644 --- a/.buildkite/scripts/hardware_ci/run-cpu-test.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-test.sh @@ -2,119 +2,19 @@ # This script build the CPU docker image and run the offline inference inside the container. # It serves a sanity check for compilation and basic model usage. -set -ex +set -euox pipefail # allow to bind to different cores CORE_RANGE=${CORE_RANGE:-48-95} -# used for TP/PP E2E test -OMP_CORE_RANGE=${OMP_CORE_RANGE:-48-95} NUMA_NODE=${NUMA_NODE:-1} +IMAGE_NAME="cpu-test-$NUMA_NODE" +TIMEOUT_VAL=$1 +TEST_COMMAND=$2 -export CMAKE_BUILD_PARALLEL_LEVEL=32 - -# Setup cleanup -remove_docker_container() { - set -e; - docker rm -f cpu-test-"$NUMA_NODE" cpu-test-"$NUMA_NODE"-avx2 || true; -} -trap remove_docker_container EXIT -remove_docker_container - -# Try building the docker image -numactl -C "$CORE_RANGE" -N "$NUMA_NODE" docker build --progress plain --tag cpu-test-"$NUMA_NODE" --target vllm-test -f docker/Dockerfile.cpu . -numactl -C "$CORE_RANGE" -N "$NUMA_NODE" docker build --progress plain --build-arg VLLM_CPU_DISABLE_AVX512="true" --tag cpu-test-"$NUMA_NODE"-avx2 --target vllm-test -f docker/Dockerfile.cpu . +# building the docker image +echo "--- :docker: Building Docker image" +docker build --progress plain --tag "$IMAGE_NAME" --target vllm-test -f docker/Dockerfile.cpu . # Run the image, setting --shm-size=4g for tensor parallel. -docker run -itd --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" --entrypoint /bin/bash -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN --env VLLM_CPU_KVCACHE_SPACE=16 --env VLLM_CPU_CI_ENV=1 -e E2E_OMP_THREADS="$OMP_CORE_RANGE" --shm-size=4g --name cpu-test-"$NUMA_NODE" cpu-test-"$NUMA_NODE" -docker run -itd --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" --entrypoint /bin/bash -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN --env VLLM_CPU_KVCACHE_SPACE=16 --env VLLM_CPU_CI_ENV=1 -e E2E_OMP_THREADS="$OMP_CORE_RANGE" --shm-size=4g --name cpu-test-"$NUMA_NODE"-avx2 cpu-test-"$NUMA_NODE"-avx2 - -function cpu_tests() { - set -e - export NUMA_NODE=$2 - - # list packages - docker exec cpu-test-"$NUMA_NODE"-avx2 bash -c " - set -e - pip list" - - docker exec cpu-test-"$NUMA_NODE" bash -c " - set -e - pip list" - - # offline inference - docker exec cpu-test-"$NUMA_NODE"-avx2 bash -c " - set -e - python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m" - - # Run kernel tests - docker exec cpu-test-"$NUMA_NODE" bash -c " - set -e - pytest -x -v -s tests/kernels/attention/test_cpu_attn.py - pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py - pytest -x -v -s tests/kernels/test_onednn.py" - - # Run basic model test - docker exec cpu-test-"$NUMA_NODE" bash -c " - set -e - # Note: disable until supports V1 - # pytest -x -v -s tests/kernels/attention/test_cache.py -m cpu_model - # pytest -x -v -s tests/kernels/attention/test_mla_decode_cpu.py -m cpu_model - - pytest -x -v -s tests/models/language/generation -m cpu_model - VLLM_CPU_SGL_KERNEL=1 pytest -x -v -s tests/models/language/generation -m cpu_model - - pytest -x -v -s tests/models/language/pooling -m cpu_model - pytest -x -v -s tests/models/multimodal/generation \ - --ignore=tests/models/multimodal/generation/test_pixtral.py \ - -m cpu_model" - - # Run compressed-tensor test - docker exec cpu-test-"$NUMA_NODE" bash -c " - set -e - pytest -x -s -v \ - tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_logprobs" - - # Run AWQ/GPTQ test - docker exec cpu-test-"$NUMA_NODE" bash -c " - set -e - pytest -x -s -v \ - tests/quantization/test_cpu_wna16.py" - - # Run multi-lora tests - docker exec cpu-test-"$NUMA_NODE" bash -c " - set -e - pytest -x -s -v \ - tests/lora/test_qwenvl.py" - - # online serving: tp+pp - docker exec cpu-test-"$NUMA_NODE" bash -c ' - set -e - VLLM_CPU_OMP_THREADS_BIND=$E2E_OMP_THREADS VLLM_CPU_SGL_KERNEL=1 vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -pp=2 & - server_pid=$! - timeout 600 bash -c "until curl localhost:8000/v1/models; do sleep 1; done" || exit 1 - vllm bench serve \ - --backend vllm \ - --dataset-name random \ - --model meta-llama/Llama-3.2-3B-Instruct \ - --num-prompts 20 \ - --endpoint /v1/completions - kill -s SIGTERM $server_pid &' - - # online serving: tp+dp - docker exec cpu-test-"$NUMA_NODE" bash -c ' - set -e - VLLM_CPU_OMP_THREADS_BIND=$E2E_OMP_THREADS VLLM_CPU_SGL_KERNEL=1 vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 & - server_pid=$! - timeout 600 bash -c "until curl localhost:8000/v1/models; do sleep 1; done" || exit 1 - vllm bench serve \ - --backend vllm \ - --dataset-name random \ - --model meta-llama/Llama-3.2-3B-Instruct \ - --num-prompts 20 \ - --endpoint /v1/completions - kill -s SIGTERM $server_pid &' -} - -# All of CPU tests are expected to be finished less than 40 mins. -export -f cpu_tests -timeout 2.5h bash -c "cpu_tests $CORE_RANGE $NUMA_NODE" +docker run --rm --cpuset-cpus=$CORE_RANGE --cpuset-mems=$NUMA_NODE -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 --shm-size=4g $IMAGE_NAME \ + timeout $TIMEOUT_VAL bash -c "set -euox pipefail; echo \"--- Print packages\"; pip list; echo \"--- Running tests\"; ${TEST_COMMAND}" diff --git a/vllm/v1/worker/cpu_worker.py b/vllm/v1/worker/cpu_worker.py index ca82fc8af45..169696ca19a 100644 --- a/vllm/v1/worker/cpu_worker.py +++ b/vllm/v1/worker/cpu_worker.py @@ -136,22 +136,38 @@ class CPUWorker(Worker): the LogicalCPUInfo.id. A selected LogicalCPUInfo list should be returned. """ + # simulate multiple numa nodes, for testing + sim_multi_numa_nodes = os.environ.get("VLLM_CPU_SIM_MULTI_NUMA", "0") != "0" allowed_numa_nodes, logical_cpu_list = ( CpuPlatform.get_allowed_cpu_core_node_list() ) - assert len(allowed_numa_nodes) >= self.parallel_config.world_size, ( + assert ( + len(allowed_numa_nodes) >= self.parallel_config.world_size + or sim_multi_numa_nodes + ), ( f"Not enough allowed NUMA nodes to bind threads of " f"{self.parallel_config.world_size} CPUWorkers. " f"Allowed NUMA nodes are {allowed_numa_nodes}. " "Please try to bind threads manually." ) - # Get CPUs on NUMA node `allowed_numa_nodes[local_rank]` - selected_numa_node = allowed_numa_nodes[self.local_rank] # type: ignore - logical_cpu_list = [ - x for x in logical_cpu_list if x.numa_node == selected_numa_node - ] + if not sim_multi_numa_nodes: + # Get CPUs on NUMA node `allowed_numa_nodes[local_rank]` + selected_numa_node = allowed_numa_nodes[self.local_rank] # type: ignore + logical_cpu_list = [ + x for x in logical_cpu_list if x.numa_node == selected_numa_node + ] + else: + assert len(logical_cpu_list) >= self.parallel_config.world_size + logical_cpu_list = sorted(logical_cpu_list, key=lambda x: x.numa_node) + sim_cpu_num_per_node = ( + len(logical_cpu_list) // self.parallel_config.world_size + ) + start_idx = self.local_rank * sim_cpu_num_per_node + logical_cpu_list = logical_cpu_list[ + start_idx : (start_idx + sim_cpu_num_per_node) + ] # Select CPUs from each physical core via cpu_selector core_to_cpus: dict[int, list[LogicalCPUInfo]] = {} From 1f70313e59bfae08588bb503b69c249a5ebd1e01 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Thu, 5 Feb 2026 00:17:00 -0600 Subject: [PATCH 086/810] [Bugfix] Fix ScoreMultiModalParam multi-document scoring returning single result (#33837) Signed-off-by: Andreas Karatzas Signed-off-by: wang.yuqi Co-authored-by: wang.yuqi --- .../pooling/test_jinavl_reranker.py | 65 ++++++------------- 1 file changed, 21 insertions(+), 44 deletions(-) diff --git a/tests/models/multimodal/pooling/test_jinavl_reranker.py b/tests/models/multimodal/pooling/test_jinavl_reranker.py index ad3ccfa3b9d..fef5b420de6 100644 --- a/tests/models/multimodal/pooling/test_jinavl_reranker.py +++ b/tests/models/multimodal/pooling/test_jinavl_reranker.py @@ -1,6 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import cast import pytest import transformers @@ -117,7 +116,7 @@ def _normalize_image(image_val: str) -> str: def create_score_multimodal_param( content_parts: list[dict], -) -> ScoreMultiModalParam: +) -> list[ScoreMultiModalParam]: """ Create a ScoreMultiModalParam from a list of content dictionaries. @@ -152,7 +151,7 @@ def create_score_multimodal_param( ) ) - return ScoreMultiModalParam(content=formatted_content) + return [ScoreMultiModalParam(content=[content]) for content in formatted_content] def _run_vllm( @@ -198,23 +197,7 @@ def _run_hf( else: raise ValueError("Unsupported query format") - # Separate documents by type - text_docs: list[str] = [] - image_docs: list[str] = [] - text_indices: list[int] = [] - image_indices: list[int] = [] - - for idx, doc in enumerate(document_strs): - if "text" in doc: - text_docs.append(doc["text"]) - text_indices.append(idx) - elif "image" in doc: - image_docs.append(_normalize_image(doc["image"])) - image_indices.append(idx) - else: - raise ValueError(f"Unsupported document format at index {idx}") - - scores: list[None | float] = [None] * len(document_strs) + scores: list[float] = [] with hf_runner( model, @@ -223,30 +206,24 @@ def _run_hf( auto_cls=AutoModel, model_kwargs={"key_mapping": CHECKPOINT_TO_HF_MAPPER}, ) as hf_model: - # Score text documents - if text_docs: - text_scores = hf_model.model.compute_score( - [[query_data, d] for d in text_docs], - max_length=2048, - query_type=query_type, - doc_type="text", - ) - for i, s in zip(text_indices, text_scores): - scores[i] = s - - # Score image documents - if image_docs: - image_scores = hf_model.model.compute_score( - [[query_data, d] for d in image_docs], - max_length=2048, - query_type=query_type, - doc_type="image", - ) - for i, s in zip(image_indices, image_scores): - scores[i] = s - - assert all(s is not None for s in scores) - return cast(list[float], scores) + for doc in document_strs: + if "text" in doc: + score = hf_model.model.compute_score( + [[query_data, doc["text"]]], + max_length=2048, + query_type=query_type, + doc_type="text", + ) + scores.append(score) + elif "image" in doc: + score = hf_model.model.compute_score( + [[query_data, doc["image"]]], + max_length=2048, + query_type=query_type, + doc_type="image", + ) + scores.append(score) + return scores def _run_test( From fd03538bf97cd7f4fedd6da4584c89635878174f Mon Sep 17 00:00:00 2001 From: Fadi Arafeh <115173828+fadara01@users.noreply.github.com> Date: Thu, 5 Feb 2026 06:26:09 +0000 Subject: [PATCH 087/810] [CPU][BugFix] Allow w8a8 oneDNN quantized matmul to support 3D inputs (#33727) Signed-off-by: Fadi Arafeh --- .../layers/quantization/kernels/scaled_mm/cpu.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/quantization/kernels/scaled_mm/cpu.py b/vllm/model_executor/layers/quantization/kernels/scaled_mm/cpu.py index b82f5781c28..3d67a73af43 100644 --- a/vllm/model_executor/layers/quantization/kernels/scaled_mm/cpu.py +++ b/vllm/model_executor/layers/quantization/kernels/scaled_mm/cpu.py @@ -182,6 +182,8 @@ class CPUInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel): x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: + x_shape = x.shape + x = x.reshape(-1, x_shape[-1]) if len(x_shape) > 2 else x w_q, w_s, i_s, i_zp, azp_adj = self._get_layer_params(layer) # ops.scaled_int8_quant supports both dynamic and static quant: @@ -195,7 +197,7 @@ class CPUInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel): n = self.dnnl_handler.n out = torch.empty((m, n), dtype=x.dtype) ops.onednn_scaled_mm(self.dnnl_handler, x_q, out, x_s, x_zp, azp_adj, bias) - + out = out.reshape(x_shape[:-1] + (n,)) if len(x_shape) > 2 else out return out def _apply_weights_sgl( From db6f71d4c9efc4679b05311c9a8fcc594b187c06 Mon Sep 17 00:00:00 2001 From: "Li, Jiang" Date: Thu, 5 Feb 2026 15:07:14 +0800 Subject: [PATCH 088/810] [CI/Build] Fix CPU CI test case title (#33870) Signed-off-by: jiang1.li --- .buildkite/hardware_tests/cpu.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.buildkite/hardware_tests/cpu.yaml b/.buildkite/hardware_tests/cpu.yaml index 39a5516967f..b387cf93502 100644 --- a/.buildkite/hardware_tests/cpu.yaml +++ b/.buildkite/hardware_tests/cpu.yaml @@ -57,7 +57,7 @@ steps: pytest -x -v -s tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_logprobs pytest -x -v -s tests/quantization/test_cpu_wna16.py" -- label: CPU-TP/DP/PP Tests +- label: CPU-Distributed Tests depends_on: [] soft_fail: true device: intel_cpu From 6abb0454adb531de0b081bbf65ccf907e4bd560d Mon Sep 17 00:00:00 2001 From: Chauncey Date: Thu, 5 Feb 2026 15:45:29 +0800 Subject: [PATCH 089/810] [Perf] Optimize the performance of structured output + reasoning (#33557) Signed-off-by: chaunceyjiang --- .../openai/chat_completion/serving.py | 101 +++++++----------- vllm/v1/engine/__init__.py | 2 + vllm/v1/request.py | 4 + vllm/v1/structured_output/__init__.py | 5 +- 4 files changed, 51 insertions(+), 61 deletions(-) diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index 48fb666484a..8ff6865165f 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -72,6 +72,7 @@ from vllm.logger import init_logger from vllm.logprobs import Logprob from vllm.outputs import CompletionOutput, RequestOutput from vllm.parser import ParserManager +from vllm.reasoning import ReasoningParser from vllm.sampling_params import BeamSearchParams, SamplingParams from vllm.tokenizers import TokenizerLike from vllm.tokenizers.mistral import ( @@ -132,7 +133,7 @@ class OpenAIServingChat(OpenAIServing): self.logits_processors = self.model_config.logits_processors # set up reasoning parser - self.reasoning_parser = ParserManager.get_reasoning_parser( + self.reasoning_parser_cls = ParserManager.get_reasoning_parser( reasoning_parser_name=reasoning_parser ) # set up tool use @@ -330,6 +331,24 @@ class OpenAIServingChat(OpenAIServing): for the API specification. This API mimics the OpenAI Chat Completion API. """ + # Streaming response + tokenizer = self.renderer.tokenizer + assert tokenizer is not None + reasoning_parser: ReasoningParser | None = None + try: + if self.reasoning_parser_cls: + # Pass the same chat template kwargs as used in tokenization + chat_template_kwargs = self._prepare_extra_chat_template_kwargs( + request.chat_template_kwargs, + self.default_chat_template_kwargs, + ) + reasoning_parser = self.reasoning_parser_cls( + tokenizer, + chat_template_kwargs=chat_template_kwargs, # type: ignore[call-arg] + ) + except RuntimeError as e: + logger.exception("Error in reasoning parser creation.") + return self.create_error_response(str(e)) result = await self.render_chat_request(request) if isinstance(result, ErrorResponse): return result @@ -427,7 +446,12 @@ class OpenAIServingChat(OpenAIServing): priority=request.priority, data_parallel_rank=data_parallel_rank, ) - + reasoning_ended = None + if reasoning_parser: + reasoning_ended = reasoning_parser.is_reasoning_end( + engine_request.prompt_token_ids or [] # type: ignore[attr-defined] + ) + engine_request.reasoning_ended = reasoning_ended generator = self.engine_client.generate( engine_request, sampling_params, @@ -447,10 +471,6 @@ class OpenAIServingChat(OpenAIServing): assert len(generators) == 1 (result_generator,) = generators - # Streaming response - tokenizer = self.renderer.tokenizer - assert tokenizer is not None - if request.stream: return self.chat_completion_stream_generator( request, @@ -460,6 +480,7 @@ class OpenAIServingChat(OpenAIServing): conversation, tokenizer, request_metadata, + reasoning_parser, ) try: @@ -471,6 +492,7 @@ class OpenAIServingChat(OpenAIServing): conversation, tokenizer, request_metadata, + reasoning_parser, ) except GenerationError as e: return self._convert_generation_error_to_response(e) @@ -630,6 +652,7 @@ class OpenAIServingChat(OpenAIServing): conversation: list[ConversationMessage], tokenizer: TokenizerLike, request_metadata: RequestResponseMetadata, + reasoning_parser: ReasoningParser | None = None, ) -> AsyncGenerator[str, None]: from vllm.tokenizers.mistral import MistralTokenizer @@ -673,7 +696,7 @@ class OpenAIServingChat(OpenAIServing): # Only one of these will be used, thus previous_texts and # all_previous_token_ids will not be used twice in the same iteration. - if tool_choice_auto or self.reasoning_parser: + if tool_choice_auto or reasoning_parser: # These are only required in "auto" tool choice case all_previous_token_ids = [[]] * num_choices # For reasoning parser and tool call all enabled @@ -683,28 +706,6 @@ class OpenAIServingChat(OpenAIServing): else: all_previous_token_ids = None - try: - if self.reasoning_parser: - if tokenizer is None: - raise ValueError( - "Tokenizer not available when `skip_tokenizer_init=True`" - ) - - # Pass the same chat template kwargs as used in tokenization - chat_template_kwargs = self._prepare_extra_chat_template_kwargs( - request.chat_template_kwargs, - self.default_chat_template_kwargs, - ) - reasoning_parser = self.reasoning_parser( - tokenizer, - chat_template_kwargs=chat_template_kwargs or {}, # type: ignore[call-arg] - ) - except RuntimeError as e: - logger.exception("Error in reasoning parser creation.") - data = self.create_streaming_error_response(str(e)) - yield f"data: {data}\n\n" - yield "data: [DONE]\n\n" - return # Prepare the tool parser if it's needed try: if tool_choice_auto and self.tool_parser: @@ -826,7 +827,7 @@ class OpenAIServingChat(OpenAIServing): tool_parser = tool_parsers[i] if ( - self.reasoning_parser + reasoning_parser and res.prompt_token_ids and prompt_is_reasoning_end_arr[i] is None ): @@ -888,7 +889,7 @@ class OpenAIServingChat(OpenAIServing): delta_message: DeltaMessage | None # just update previous_texts and previous_token_ids - if tool_choice_auto or self.reasoning_parser: + if tool_choice_auto or reasoning_parser: assert previous_texts is not None assert all_previous_token_ids is not None previous_text = previous_texts[i] @@ -915,7 +916,7 @@ class OpenAIServingChat(OpenAIServing): # handle streaming deltas for tools with named tool_choice elif tool_choice_function_name: if ( - self.reasoning_parser + reasoning_parser and not reasoning_end_arr[i] and not reasoning_parser.is_reasoning_end( previous_token_ids @@ -952,7 +953,7 @@ class OpenAIServingChat(OpenAIServing): current_text = "" else: # Just to add remaining `content` - if self.reasoning_parser: + if reasoning_parser: delta_text = previous_text + delta_text current_text = "" @@ -998,13 +999,13 @@ class OpenAIServingChat(OpenAIServing): output_token_ids = as_list(output.token_ids) if ( - self.reasoning_parser is not None + reasoning_parser is not None and not reasoning_end_arr[i] and prompt_is_reasoning_end_arr[i] ): reasoning_end_arr[i] = True - if self.reasoning_parser and not reasoning_end_arr[i]: + if reasoning_parser and not reasoning_end_arr[i]: delta_message = ( reasoning_parser.extract_reasoning_streaming( previous_text, @@ -1047,9 +1048,8 @@ class OpenAIServingChat(OpenAIServing): # handle streaming deltas for tools with "auto" tool choice # and reasoning parser - elif tool_choice_auto and self.reasoning_parser: + elif tool_choice_auto and reasoning_parser: assert tool_parser is not None - assert reasoning_parser is not None assert added_content_delta_arr is not None assert reasoning_end_arr is not None output_token_ids = as_list(output.token_ids) @@ -1130,7 +1130,7 @@ class OpenAIServingChat(OpenAIServing): tools_streamed[i] = True # when only reasoning - elif self.reasoning_parser: + elif reasoning_parser: delta_message = reasoning_parser.extract_reasoning_streaming( previous_text, current_text, @@ -1144,9 +1144,7 @@ class OpenAIServingChat(OpenAIServing): delta_message = DeltaMessage(content=delta_text) # update the previous values for the next iteration - if ( - tool_choice_auto or self.reasoning_parser - ) and not self.use_harmony: + if (tool_choice_auto or reasoning_parser) and not self.use_harmony: assert previous_texts is not None assert all_previous_token_ids is not None previous_texts[i] = current_text @@ -1400,6 +1398,7 @@ class OpenAIServingChat(OpenAIServing): conversation: list[ConversationMessage], tokenizer: TokenizerLike, request_metadata: RequestResponseMetadata, + reasoning_parser: ReasoningParser | None = None, ) -> ErrorResponse | ChatCompletionResponse: from vllm.tokenizers.mistral import MistralTokenizer @@ -1494,25 +1493,7 @@ class OpenAIServingChat(OpenAIServing): choices.append(choice_data) continue - if self.reasoning_parser: - try: - if tokenizer is None: - raise ValueError( - "Tokenizer not available when `skip_tokenizer_init=True`" - ) - - # Pass the same chat template kwargs as used in tokenization - chat_template_kwargs = self._prepare_extra_chat_template_kwargs( - request.chat_template_kwargs, - self.default_chat_template_kwargs, - ) - reasoning_parser = self.reasoning_parser( - tokenizer, - chat_template_kwargs=chat_template_kwargs, # type: ignore[call-arg] - ) - except RuntimeError as e: - logger.exception("Error in reasoning parser creation.") - return self.create_error_response(str(e)) + if reasoning_parser: # If the reasoning parser is enabled, # tool calls are extracted exclusively from the content. reasoning, content = reasoning_parser.extract_reasoning( diff --git a/vllm/v1/engine/__init__.py b/vllm/v1/engine/__init__.py index 5328a673554..f5d8ce1ff89 100644 --- a/vllm/v1/engine/__init__.py +++ b/vllm/v1/engine/__init__.py @@ -83,6 +83,8 @@ class EngineCoreRequest( # Used in outputs and to support abort(req_id, internal=False). external_req_id: str | None = None + reasoning_ended: bool | None = None + @property def params(self) -> SamplingParams | PoolingParams: """Return the processed params (sampling or pooling).""" diff --git a/vllm/v1/request.py b/vllm/v1/request.py index e9d3df4421e..8e3684d3c9a 100644 --- a/vllm/v1/request.py +++ b/vllm/v1/request.py @@ -74,6 +74,7 @@ class Request: trace_headers: Mapping[str, str] | None = None, block_hasher: Callable[["Request"], list["BlockHash"]] | None = None, resumable: bool = False, + reasoning_ended: bool | None = None, ) -> None: self.request_id = request_id self.client_index = client_index @@ -86,6 +87,8 @@ class Request: self.structured_output_request = StructuredOutputRequest.from_sampling_params( sampling_params ) + if self.structured_output_request is not None: + self.structured_output_request.reasoning_ended = reasoning_ended self.arrival_time = arrival_time if arrival_time is not None else time.time() self.status = RequestStatus.WAITING @@ -195,6 +198,7 @@ class Request: trace_headers=request.trace_headers, block_hasher=block_hasher, resumable=request.resumable, + reasoning_ended=request.reasoning_ended, ) def append_output_token_ids( diff --git a/vllm/v1/structured_output/__init__.py b/vllm/v1/structured_output/__init__.py index 9b86d69a751..921bee6a647 100644 --- a/vllm/v1/structured_output/__init__.py +++ b/vllm/v1/structured_output/__init__.py @@ -284,12 +284,15 @@ class StructuredOutputManager: # 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. - if self.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: + # 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 = ( self.reasoner.is_reasoning_end(request.prompt_token_ids or []) ) From 2abd97592f947c041ba70329532f0cf62dd8971f Mon Sep 17 00:00:00 2001 From: Mark McLoughlin Date: Thu, 5 Feb 2026 07:57:27 +0000 Subject: [PATCH 090/810] [KV Connector][Metrics] Do not count local prefix cache hits in connector queries (#30522) Signed-off-by: Mark McLoughlin --- tests/v1/core/test_scheduler.py | 79 ++++++++++++++++--- .../unit/test_invalid_blocks_correctness.py | 28 ++++++- vllm/v1/core/sched/scheduler.py | 39 +++++---- 3 files changed, 115 insertions(+), 31 deletions(-) diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index b29df468f0c..a1e3d09d24f 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -1136,7 +1136,7 @@ def _step_until_kv_transfer_finished(scheduler: Scheduler, req_ids: list[str]): prompt_logprobs_dict={}, pooler_output=[], ) - scheduler.update_from_output(output, EMPTY_OUTPUT) + initial_ecos = scheduler.update_from_output(output, EMPTY_OUTPUT) # Simulate KV transfer completion using KVConnectorOutput.finished_recving output = scheduler.schedule() @@ -1156,6 +1156,8 @@ def _step_until_kv_transfer_finished(scheduler: Scheduler, req_ids: list[str]): for req_id in req_ids: assert req_id in scheduler.finished_recving_kv_req_ids + return initial_ecos + @pytest.mark.parametrize("is_async", [False, True]) def test_kv_connector_basic(is_async: bool): @@ -1286,29 +1288,72 @@ def test_kv_connector_basic(is_async: bool): @pytest.mark.parametrize("is_async", [False, True]) -def test_external_prefix_cache_metrics(is_async: bool): +@pytest.mark.parametrize("local_cache_hits", [False, True]) +def test_external_prefix_cache_metrics(is_async: bool, local_cache_hits: bool): """ Verify connector prefix cache metrics are updated correctly when the scheduler processes requests with KV connector hits. """ + BLOCK_SIZE = 16 + if local_cache_hits: + NUM_MATCHED_NEW_TOKENS = BLOCK_SIZE * 2 # 32 tokens + NUM_LOCAL_HITS = NUM_MATCHED_NEW_TOKENS * 2 # 64 tokens + NUM_REQUESTS = 1 + NUM_TOKENS = NUM_LOCAL_HITS * 2 # 128 tokens + else: + NUM_MATCHED_NEW_TOKENS = 4 + NUM_LOCAL_HITS = 0 + NUM_REQUESTS = 2 + NUM_TOKENS = 8 # 8 tokens + # Setup Scheduler. - NUM_MATCHED_NEW_TOKENS = 4 scheduler = create_scheduler( - enable_prefix_caching=False, + enable_prefix_caching=local_cache_hits, use_kv_connector=mock_kv( matched_tokens=NUM_MATCHED_NEW_TOKENS, is_async=is_async ), + block_size=BLOCK_SIZE, ) - # --- Prepare simple requests --- - NUM_REQUESTS = 2 - NUM_TOKENS = 8 + if local_cache_hits: + # First, establish local cache by running a request to completion + requests = create_requests( + num_requests=1, + num_tokens=NUM_LOCAL_HITS, + max_tokens=2, + block_size=BLOCK_SIZE, + ) + req_ids = [] + req_to_index = {} + for i, request in enumerate(requests): + scheduler.add_request(request) + req_ids.append(request.request_id) + req_to_index[request.request_id] = i + + if is_async: + _step_until_kv_transfer_finished(scheduler, req_ids) + + # Run first request to completion to establish local cache + output = scheduler.schedule() + MODEL_RUNNER_OUTPUT = ModelRunnerOutput( + req_ids=req_ids, + req_id_to_index=req_to_index, + sampled_token_ids=[[1000]] * len(req_ids), + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ) + _step_until_done(scheduler, output, MODEL_RUNNER_OUTPUT) + _ = scheduler.schedule() + + # --- Prepare test requests --- MAX_TOKENS = 2 requests = create_requests( num_requests=NUM_REQUESTS, num_tokens=NUM_TOKENS, max_tokens=MAX_TOKENS, + block_size=BLOCK_SIZE, ) req_ids = [] req_to_index = {} @@ -1317,8 +1362,9 @@ def test_external_prefix_cache_metrics(is_async: bool): req_ids.append(request.request_id) req_to_index[request.request_id] = i + initial_ecos = None if is_async: - _step_until_kv_transfer_finished(scheduler, req_ids) + initial_ecos = _step_until_kv_transfer_finished(scheduler, req_ids) # --- Trigger scheduling and simulate model output --- output = scheduler.schedule() @@ -1338,10 +1384,23 @@ def test_external_prefix_cache_metrics(is_async: bool): assert ecos is not None and len(ecos) > 0 assert ecos[0].scheduler_stats is not None - external_stats = ecos[0].scheduler_stats.connector_prefix_cache_stats + if local_cache_hits: + # For async, local cache stats come from the first step + if initial_ecos: + local_stats = initial_ecos[0].scheduler_stats.prefix_cache_stats + else: + local_stats = ecos[0].scheduler_stats.prefix_cache_stats + assert local_stats is not None + assert local_stats.queries == NUM_TOKENS * NUM_REQUESTS + assert local_stats.hits == NUM_LOCAL_HITS * NUM_REQUESTS + + if initial_ecos: + external_stats = initial_ecos[0].scheduler_stats.connector_prefix_cache_stats + else: + external_stats = ecos[0].scheduler_stats.connector_prefix_cache_stats assert external_stats is not None - assert external_stats.queries == NUM_TOKENS * NUM_REQUESTS + assert external_stats.queries == (NUM_TOKENS - NUM_LOCAL_HITS) * NUM_REQUESTS assert external_stats.hits == NUM_MATCHED_NEW_TOKENS * NUM_REQUESTS assert external_stats.requests == NUM_REQUESTS assert external_stats.preempted_requests == 0 diff --git a/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py b/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py index 940f3a98308..6cb2d3ea4d9 100644 --- a/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py +++ b/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py @@ -281,6 +281,17 @@ def test_sync_fail_invalid_blocks_evicted(fail_scheduler: Scheduler): f"(hash should be None), but hash is still {block.block_hash}" ) + # Verify connector prefix cache stats: + # - queries = num_prompt_tokens (total tokens not in local cache) + # - hits = num_external_computed_tokens (tokens loaded externally) + assert engine_outputs.scheduler_stats is not None + stats = engine_outputs.scheduler_stats + assert stats.connector_prefix_cache_stats is not None + conn_stats = stats.connector_prefix_cache_stats + assert conn_stats.requests == 1 + assert conn_stats.queries == num_prompt_tokens + assert conn_stats.hits == num_external_computed_tokens + def test_async_recompute_blocks_not_cached_when_invalid( recompute_scheduler: Scheduler, @@ -364,7 +375,9 @@ def test_async_recompute_blocks_not_cached_when_invalid( with patch.object( recompute_scheduler.kv_cache_manager, "evict_blocks", evict_blocks_spy ): - recompute_scheduler.update_from_output(scheduler_output, model_runner_output) + outputs = recompute_scheduler.update_from_output( + scheduler_output, model_runner_output + ) # verify evict_blocks was NOT called (async blocks excluded from eviction) assert len(evict_blocks_calls) == 0, ( @@ -386,6 +399,19 @@ def test_async_recompute_blocks_not_cached_when_invalid( f"Block {invalid_block_id} hash should be None but is {block.block_hash}" ) + # Verify connector prefix cache stats: + # - queries = num_prompt_tokens (total tokens not in local cache) + # - hits = num_external_computed_tokens (tokens loaded externally) + assert len(outputs) == 1 + engine_outputs = next(iter(outputs.values())) + assert engine_outputs.scheduler_stats is not None + stats = engine_outputs.scheduler_stats + assert stats.connector_prefix_cache_stats is not None + conn_stats = stats.connector_prefix_cache_stats + assert conn_stats.requests == 1 + assert conn_stats.queries == num_prompt_tokens + assert conn_stats.hits == num_external_computed_tokens + # now simulate async transfer completing model_runner_output_2 = create_model_runner_output( reqs=[], diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 88d1a78df00..745d9ffec77 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -586,6 +586,7 @@ class Scheduler(SchedulerInterface): num_external_computed_tokens = 0 load_kv_async = False + connector_prefix_cache_queries, connector_prefix_cache_hits = 0, 0 # Get already-cached tokens. if request.num_computed_tokens == 0: @@ -613,6 +614,11 @@ class Scheduler(SchedulerInterface): request.num_external_computed_tokens = ext_tokens num_external_computed_tokens = ext_tokens + connector_prefix_cache_queries = ( + request.num_tokens - num_new_local_computed_tokens + ) + connector_prefix_cache_hits = num_external_computed_tokens + # Total computed tokens (local + external). num_computed_tokens = ( num_new_local_computed_tokens + num_external_computed_tokens @@ -728,6 +734,15 @@ class Scheduler(SchedulerInterface): self.kv_cache_manager.get_blocks(request_id), num_external_computed_tokens, ) + if ( + self.connector_prefix_cache_stats is not None + and connector_prefix_cache_queries != 0 + ): + self.connector_prefix_cache_stats.record( + num_tokens=connector_prefix_cache_queries, + num_hits=connector_prefix_cache_hits, + preempted=request.num_preemptions > 0, + ) # Request was already popped from self.waiting # unless it was re-added above due to new_blocks being None. @@ -739,8 +754,6 @@ class Scheduler(SchedulerInterface): request.status = RequestStatus.WAITING_FOR_REMOTE_KVS continue - self._update_connector_prefix_cache_stats(request) - self.running.append(request) if self.log_stats: request.record_event( @@ -1805,7 +1818,10 @@ class Scheduler(SchedulerInterface): return None prefix_cache_stats = self.kv_cache_manager.make_prefix_cache_stats() assert prefix_cache_stats is not None - connector_prefix_cache_stats = self._make_connector_prefix_cache_stats() + connector_prefix_cache_stats: PrefixCacheStats | None = None + if self.connector_prefix_cache_stats is not None: + connector_prefix_cache_stats = self.connector_prefix_cache_stats + self.connector_prefix_cache_stats = PrefixCacheStats() eviction_events = ( self.kv_metrics_collector.drain_events() if self.kv_metrics_collector is not None @@ -1866,23 +1882,6 @@ class Scheduler(SchedulerInterface): # KV Connector Related Methods ######################################################################## - def _update_connector_prefix_cache_stats(self, request: Request) -> None: - if self.connector_prefix_cache_stats is None: - return - - self.connector_prefix_cache_stats.record( - num_tokens=request.num_tokens, - num_hits=request.num_external_computed_tokens, - preempted=request.num_preemptions > 0, - ) - - def _make_connector_prefix_cache_stats(self) -> PrefixCacheStats | None: - if self.connector_prefix_cache_stats is None: - return None - stats = self.connector_prefix_cache_stats - self.connector_prefix_cache_stats = PrefixCacheStats() - return stats - def get_kv_connector(self) -> KVConnectorBase_V1 | None: return self.connector From d2f4a71cd54418369f617a174e6c839a71a47ed8 Mon Sep 17 00:00:00 2001 From: Pavani Majety Date: Thu, 5 Feb 2026 01:32:10 -0800 Subject: [PATCH 091/810] [Bugfix] Kimi-K2 grouped_topk usage for Flashinfer monolithic kernels. (#33858) Signed-off-by: Pavani Majety --- vllm/model_executor/models/deepseek_v2.py | 14 +++----------- 1 file changed, 3 insertions(+), 11 deletions(-) diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 4e465f0fe16..f8907ed86ef 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -295,14 +295,6 @@ class DeepseekV2MoE(nn.Module): prefix=f"{prefix}.shared_experts", ) - n_group = getattr(config, "n_group", 1) - topk_group = getattr(config, "topk_group", 1) - use_grouped_topk = True - if (n_group, topk_group) == (1, 1): - n_group = None - topk_group = None - use_grouped_topk = False - self.experts = SharedFusedMoE( shared_experts=self.shared_experts, gate=self.gate, @@ -313,9 +305,9 @@ class DeepseekV2MoE(nn.Module): reduce_results=False, renormalize=config.norm_topk_prob, quant_config=quant_config, - use_grouped_topk=use_grouped_topk, - num_expert_group=n_group, - topk_group=topk_group, + use_grouped_topk=True, + num_expert_group=getattr(config, "n_group", 1), + topk_group=getattr(config, "topk_group", 1), prefix=f"{prefix}.experts", scoring_func=getattr(config, "scoring_func", "softmax"), # we do scaling outside, set factor to 1.0 to avoid double mul From 038914b7c891c0b5b2853ec0574062dc3bea8073 Mon Sep 17 00:00:00 2001 From: Cyrus Leung Date: Thu, 5 Feb 2026 17:33:11 +0800 Subject: [PATCH 092/810] [Refactor] Move `task` outside of `PoolingParams.verify` (#33796) Signed-off-by: DarkLight1337 Signed-off-by: wang.yuqi Co-authored-by: wang.yuqi --- .buildkite/test-amd.yaml | 2 + .buildkite/test-pipeline.yaml | 2 + .buildkite/test_areas/misc.yaml | 2 + .../pooling/classify/test_online.py | 4 +- .../entrypoints/pooling/embed/test_online.py | 4 +- .../pooling/score/test_online_colbert.py | 10 +-- .../pooling/score/test_online_rerank.py | 4 +- tests/test_pooling_params.py | 90 +++++++++---------- vllm/entrypoints/llm.py | 13 +-- vllm/entrypoints/pooling/__init__.py | 3 +- vllm/entrypoints/pooling/base/protocol.py | 24 ----- vllm/entrypoints/pooling/classify/protocol.py | 18 ++++ vllm/entrypoints/pooling/classify/serving.py | 16 ---- vllm/entrypoints/pooling/embed/protocol.py | 34 +++++++ vllm/entrypoints/pooling/embed/serving.py | 26 +----- vllm/entrypoints/pooling/pooling/protocol.py | 4 +- vllm/entrypoints/pooling/pooling/serving.py | 35 +------- vllm/entrypoints/pooling/score/protocol.py | 11 ++- vllm/entrypoints/pooling/score/serving.py | 34 ++----- vllm/pooling_params.py | 16 +--- vllm/v1/engine/async_llm.py | 7 +- vllm/v1/engine/input_processor.py | 34 ++++++- vllm/v1/engine/llm_engine.py | 7 +- vllm/v1/worker/gpu_model_runner.py | 2 +- 24 files changed, 186 insertions(+), 216 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 0050c615a4b..64aaf1eb6ff 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -70,6 +70,7 @@ steps: - vllm/ - tests/test_inputs.py - tests/test_outputs.py + - tests/test_pooling_params.py - tests/multimodal - tests/renderers - tests/standalone_tests/lazy_imports.py @@ -82,6 +83,7 @@ steps: - 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 -m 'cpu_test' multimodal - pytest -v -s renderers - pytest -v -s tokenizers_ diff --git a/.buildkite/test-pipeline.yaml b/.buildkite/test-pipeline.yaml index 554081f5353..a3e25c0f755 100644 --- a/.buildkite/test-pipeline.yaml +++ b/.buildkite/test-pipeline.yaml @@ -63,6 +63,7 @@ steps: - vllm/ - tests/test_inputs.py - tests/test_outputs.py + - tests/test_pooling_params.py - tests/multimodal - tests/renderers - tests/standalone_tests/lazy_imports.py @@ -75,6 +76,7 @@ steps: - 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 -m 'cpu_test' multimodal - pytest -v -s renderers - pytest -v -s tokenizers_ diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index a8cb5cd8690..a01c2296f28 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -122,6 +122,7 @@ steps: - vllm/ - tests/test_inputs.py - tests/test_outputs.py + - tests/test_pooling_params.py - tests/multimodal - tests/renderers - tests/standalone_tests/lazy_imports.py @@ -134,6 +135,7 @@ steps: - 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 -m 'cpu_test' multimodal - pytest -v -s renderers - pytest -v -s tokenizers_ diff --git a/tests/entrypoints/pooling/classify/test_online.py b/tests/entrypoints/pooling/classify/test_online.py index 45712c8425b..e23918fb8db 100644 --- a/tests/entrypoints/pooling/classify/test_online.py +++ b/tests/entrypoints/pooling/classify/test_online.py @@ -469,6 +469,4 @@ async def test_pooling_not_supported( }, ) assert response.json()["error"]["type"] == "BadRequestError" - assert response.json()["error"]["message"].startswith( - f"Task {task} is not supported" - ) + assert response.json()["error"]["message"].startswith(f"Unsupported task: {task!r}") diff --git a/tests/entrypoints/pooling/embed/test_online.py b/tests/entrypoints/pooling/embed/test_online.py index 092a5c008ed..d2a5974b757 100644 --- a/tests/entrypoints/pooling/embed/test_online.py +++ b/tests/entrypoints/pooling/embed/test_online.py @@ -757,6 +757,4 @@ async def test_pooling_not_supported( }, ) assert response.json()["error"]["type"] == "BadRequestError" - assert response.json()["error"]["message"].startswith( - f"Task {task} is not supported" - ) + assert response.json()["error"]["message"].startswith(f"Unsupported task: {task!r}") diff --git a/tests/entrypoints/pooling/score/test_online_colbert.py b/tests/entrypoints/pooling/score/test_online_colbert.py index a7b404d0fde..dcc7dff239e 100644 --- a/tests/entrypoints/pooling/score/test_online_colbert.py +++ b/tests/entrypoints/pooling/score/test_online_colbert.py @@ -138,17 +138,17 @@ def test_colbert_token_embed(server: RemoteOpenAIServer, model_name: str): @pytest.mark.parametrize("model_name", [MODEL_NAME]) def test_colbert_embed_not_supported(server: RemoteOpenAIServer, model_name: str): """Test that ColBERT model does not support 'embed' task.""" + task = "embed" text = "What is the capital of France?" - pooling_response = requests.post( + response = requests.post( server.url_for("pooling"), json={ "model": model_name, "input": text, - "task": "embed", + "task": task, }, ) - # Should return error - assert pooling_response.status_code == 400 - assert "Task embed is not supported" in pooling_response.text + assert response.json()["error"]["type"] == "BadRequestError" + assert response.json()["error"]["message"].startswith(f"Unsupported task: {task!r}") diff --git a/tests/entrypoints/pooling/score/test_online_rerank.py b/tests/entrypoints/pooling/score/test_online_rerank.py index 35687eff0d5..b0e8152aed7 100644 --- a/tests/entrypoints/pooling/score/test_online_rerank.py +++ b/tests/entrypoints/pooling/score/test_online_rerank.py @@ -232,6 +232,4 @@ async def test_pooling_not_supported( }, ) assert response.json()["error"]["type"] == "BadRequestError" - assert response.json()["error"]["message"].startswith( - f"Task {task} is not supported" - ) + assert response.json()["error"]["message"].startswith(f"Unsupported task: {task!r}") diff --git a/tests/test_pooling_params.py b/tests/test_pooling_params.py index 28dedc10e1a..54a577d2bf8 100644 --- a/tests/test_pooling_params.py +++ b/tests/test_pooling_params.py @@ -27,35 +27,24 @@ class MockModelConfig: pooler_config: PoolerConfig -def test_task(): - pooling_params = PoolingParams() - pooling_params.verify(task="score") - - pooling_params = PoolingParams(task="score") - pooling_params.verify(task="score") - - with pytest.raises(ValueError): - pooling_params.verify(task="classify") - - def test_embed(): task = "embed" model_config = MockModelConfig(pooler_config=PoolerConfig(seq_pooling_type="CLS")) - pooling_params = PoolingParams(use_activation=None) - pooling_params.verify(task=task, model_config=model_config) + pooling_params = PoolingParams(task=task, use_activation=None) + pooling_params.verify(model_config) - pooling_params = PoolingParams(use_activation=True) - pooling_params.verify(task=task, model_config=model_config) + pooling_params = PoolingParams(task=task, use_activation=True) + pooling_params.verify(model_config) - pooling_params = PoolingParams(use_activation=False) - pooling_params.verify(task=task, model_config=model_config) + pooling_params = PoolingParams(task=task, use_activation=False) + pooling_params.verify(model_config) invalid_parameters = classify_parameters + step_pooling_parameters for p in set(invalid_parameters) - set(embed_parameters): with pytest.raises(ValueError): - pooling_params = PoolingParams(**{p: True}) - pooling_params.verify(task=task, model_config=model_config) + pooling_params = PoolingParams(task=task, **{p: True}) + pooling_params.verify(model_config) @pytest.mark.parametrize("model_info", EMBEDDING_MODELS) @@ -63,7 +52,6 @@ def test_embed_dimensions(model_info: EmbedModelInfo): task = "embed" model_config = ModelConfig( model_info.name, - task="auto", tokenizer=model_info.name, tokenizer_mode="auto", trust_remote_code=False, @@ -71,37 +59,39 @@ def test_embed_dimensions(model_info: EmbedModelInfo): dtype="float16", ) - pooling_params = PoolingParams(dimensions=None) - pooling_params.verify(task=task, model_config=model_config) + pooling_params = PoolingParams(task=task, dimensions=None) + pooling_params.verify(model_config) with pytest.raises(ValueError): - pooling_params = PoolingParams(dimensions=1) - pooling_params.verify(task=task, model_config=model_config) + pooling_params = PoolingParams(task=task, dimensions=1) + pooling_params.verify(model_config) if model_info.is_matryoshka: assert model_info.matryoshka_dimensions is not None - pooling_params = PoolingParams(dimensions=model_info.matryoshka_dimensions[0]) - pooling_params.verify(task=task, model_config=model_config) + pooling_params = PoolingParams( + task=task, dimensions=model_info.matryoshka_dimensions[0] + ) + pooling_params.verify(model_config) @pytest.mark.parametrize("task", ["score", "classify"]) def test_classify(task): model_config = MockModelConfig(pooler_config=PoolerConfig(seq_pooling_type="CLS")) - pooling_params = PoolingParams(use_activation=None) - pooling_params.verify(task=task, model_config=model_config) + pooling_params = PoolingParams(task=task, use_activation=None) + pooling_params.verify(model_config) - pooling_params = PoolingParams(use_activation=True) - pooling_params.verify(task=task, model_config=model_config) + pooling_params = PoolingParams(task=task, use_activation=True) + pooling_params.verify(model_config) - pooling_params = PoolingParams(use_activation=False) - pooling_params.verify(task=task, model_config=model_config) + pooling_params = PoolingParams(task=task, use_activation=False) + pooling_params.verify(model_config) invalid_parameters = embed_parameters + step_pooling_parameters for p in set(invalid_parameters) - set(classify_parameters): with pytest.raises(ValueError): - pooling_params = PoolingParams(**{p: True}) - pooling_params.verify(task=task, model_config=model_config) + pooling_params = PoolingParams(task=task, **{p: True}) + pooling_params.verify(model_config) @pytest.mark.parametrize("pooling_type", ["ALL", "STEP"]) @@ -111,14 +101,14 @@ def test_token_embed(pooling_type: str): pooler_config=PoolerConfig(tok_pooling_type=pooling_type) ) - pooling_params = PoolingParams(use_activation=None) - pooling_params.verify(task=task, model_config=model_config) + pooling_params = PoolingParams(task=task, use_activation=None) + pooling_params.verify(model_config) - pooling_params = PoolingParams(use_activation=True) - pooling_params.verify(task=task, model_config=model_config) + pooling_params = PoolingParams(task=task, use_activation=True) + pooling_params.verify(model_config) - pooling_params = PoolingParams(use_activation=False) - pooling_params.verify(task=task, model_config=model_config) + pooling_params = PoolingParams(task=task, use_activation=False) + pooling_params.verify(model_config) invalid_parameters = classify_parameters if pooling_type != "STEP": @@ -126,8 +116,8 @@ def test_token_embed(pooling_type: str): for p in set(invalid_parameters) - set(embed_parameters): with pytest.raises(ValueError): - pooling_params = PoolingParams(**{p: True}) - pooling_params.verify(task=task, model_config=model_config) + pooling_params = PoolingParams(task=task, **{p: True}) + pooling_params.verify(model_config) @pytest.mark.parametrize("pooling_type", ["ALL", "STEP"]) @@ -137,14 +127,14 @@ def test_token_classify(pooling_type: str): pooler_config=PoolerConfig(tok_pooling_type=pooling_type) ) - pooling_params = PoolingParams(use_activation=None) - pooling_params.verify(task=task, model_config=model_config) + pooling_params = PoolingParams(task=task, use_activation=None) + pooling_params.verify(model_config) - pooling_params = PoolingParams(use_activation=True) - pooling_params.verify(task=task, model_config=model_config) + pooling_params = PoolingParams(task=task, use_activation=True) + pooling_params.verify(model_config) - pooling_params = PoolingParams(use_activation=False) - pooling_params.verify(task=task, model_config=model_config) + pooling_params = PoolingParams(task=task, use_activation=False) + pooling_params.verify(model_config) invalid_parameters = embed_parameters if pooling_type != "STEP": @@ -152,5 +142,5 @@ def test_token_classify(pooling_type: str): for p in set(invalid_parameters) - set(classify_parameters): with pytest.raises(ValueError): - pooling_params = PoolingParams(**{p: True}) - pooling_params.verify(task=task, model_config=model_config) + pooling_params = PoolingParams(task=task, **{p: True}) + pooling_params.verify(model_config) diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index 435ccbee6c7..fbcf3a43773 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -1135,11 +1135,12 @@ class LLM: # Use default pooling params. pooling_params = PoolingParams() - if pooling_task not in self.supported_tasks: - raise ValueError(f"pooling_task must be one of {self.supported_tasks}.") - for param in as_iter(pooling_params): - param.verify(pooling_task, model_config) + if param.task is None: + param.task = pooling_task + elif param.task != pooling_task: + msg = f"You cannot overwrite {param.task=!r} with {pooling_task=!r}!" + raise ValueError(msg) self._validate_and_add_requests( prompts=prompts, @@ -1472,8 +1473,9 @@ class LLM: if pooling_params is None: pooling_params = PoolingParams(task="score") + elif pooling_params.task is None: + pooling_params.task = "score" - pooling_params.verify("score", model_config) pooling_params_list = list[PoolingParams]() prompts = list[PromptType]() @@ -1836,6 +1838,7 @@ class LLM: lora_request=lora_request, tokenization_kwargs=tokenization_kwargs, priority=priority, + supported_tasks=self.supported_tasks, ) self.llm_engine.add_request( diff --git a/vllm/entrypoints/pooling/__init__.py b/vllm/entrypoints/pooling/__init__.py index 4321e19f94c..1108be175bc 100644 --- a/vllm/entrypoints/pooling/__init__.py +++ b/vllm/entrypoints/pooling/__init__.py @@ -68,7 +68,6 @@ def init_pooling_state( OpenAIServingPooling( engine_client, state.openai_serving_models, - supported_tasks=supported_tasks, request_logger=request_logger, chat_template=resolved_chat_template, chat_template_content_format=args.chat_template_content_format, @@ -76,7 +75,7 @@ def init_pooling_state( log_error_stack=args.log_error_stack, ) ) - if any(task in POOLING_TASKS for task in supported_tasks) + if any(t in supported_tasks for t in POOLING_TASKS) else None ) state.openai_serving_embedding = ( diff --git a/vllm/entrypoints/pooling/base/protocol.py b/vllm/entrypoints/pooling/base/protocol.py index f27025970f6..19654db504b 100644 --- a/vllm/entrypoints/pooling/base/protocol.py +++ b/vllm/entrypoints/pooling/base/protocol.py @@ -6,19 +6,15 @@ from typing import Annotated, Any from pydantic import Field, model_validator -from vllm import PoolingParams from vllm.entrypoints.chat_utils import ( ChatCompletionMessageParam, ChatTemplateContentFormatOption, ) from vllm.entrypoints.openai.engine.protocol import OpenAIBaseModel -from vllm.logger import init_logger from vllm.renderers import ChatParams, merge_kwargs from vllm.utils import random_uuid from vllm.utils.serial_utils import EmbedDType, EncodingFormat, Endianness -logger = init_logger(__name__) - class PoolingBasicRequestMixin(OpenAIBaseModel): # --8<-- [start:pooling-common-params] @@ -185,20 +181,6 @@ class EmbedRequestMixin(EncodingRequestMixin): ) # --8<-- [end:embed-extra-params] - def to_pooling_params(self): - if self.normalize is not None: - logger.warning_once( - "`normalize` is deprecated and will be removed in v0.17. " - "Please pass `use_activation` instead." - ) - self.use_activation = self.normalize - - return PoolingParams( - dimensions=self.dimensions, - use_activation=self.use_activation, - truncate_prompt_tokens=getattr(self, "truncate_prompt_tokens", None), - ) - class ClassifyRequestMixin(OpenAIBaseModel): # --8<-- [start:classify-extra-params] @@ -208,9 +190,3 @@ class ClassifyRequestMixin(OpenAIBaseModel): "`None` uses the pooler's default, which is `True` in most cases.", ) # --8<-- [end:classify-extra-params] - - def to_pooling_params(self): - return PoolingParams( - use_activation=self.use_activation, - truncate_prompt_tokens=getattr(self, "truncate_prompt_tokens", None), - ) diff --git a/vllm/entrypoints/pooling/classify/protocol.py b/vllm/entrypoints/pooling/classify/protocol.py index 33a25335d1f..55641561d5a 100644 --- a/vllm/entrypoints/pooling/classify/protocol.py +++ b/vllm/entrypoints/pooling/classify/protocol.py @@ -6,6 +6,7 @@ from typing import Any, TypeAlias from pydantic import Field +from vllm import PoolingParams from vllm.config import ModelConfig from vllm.entrypoints.openai.engine.protocol import OpenAIBaseModel, UsageInfo from vllm.entrypoints.pooling.base.protocol import ( @@ -14,9 +15,12 @@ from vllm.entrypoints.pooling.base.protocol import ( CompletionRequestMixin, PoolingBasicRequestMixin, ) +from vllm.logger import init_logger from vllm.renderers import TokenizeParams from vllm.utils import random_uuid +logger = init_logger(__name__) + class ClassificationCompletionRequest( PoolingBasicRequestMixin, CompletionRequestMixin, ClassifyRequestMixin @@ -33,6 +37,13 @@ class ClassificationCompletionRequest( max_total_tokens_param="max_model_len", ) + def to_pooling_params(self): + return PoolingParams( + task="classify", + truncate_prompt_tokens=self.truncate_prompt_tokens, + use_activation=self.use_activation, + ) + class ClassificationChatRequest( PoolingBasicRequestMixin, ChatRequestMixin, ClassifyRequestMixin @@ -55,6 +66,13 @@ class ClassificationChatRequest( max_total_tokens_param="max_model_len", ) + def to_pooling_params(self): + return PoolingParams( + task="classify", + truncate_prompt_tokens=self.truncate_prompt_tokens, + use_activation=self.use_activation, + ) + ClassificationRequest: TypeAlias = ( ClassificationCompletionRequest | ClassificationChatRequest diff --git a/vllm/entrypoints/pooling/classify/serving.py b/vllm/entrypoints/pooling/classify/serving.py index d9f7db95381..8cdbbde6d6f 100644 --- a/vllm/entrypoints/pooling/classify/serving.py +++ b/vllm/entrypoints/pooling/classify/serving.py @@ -22,7 +22,6 @@ from vllm.entrypoints.pooling.classify.protocol import ( ) from vllm.logger import init_logger from vllm.outputs import ClassificationOutput -from vllm.pooling_params import PoolingParams logger = init_logger(__name__) @@ -159,18 +158,3 @@ class ServingClassification(OpenAIServing): ) return await self.handle(ctx) # type: ignore[return-value] - - def _create_pooling_params( - self, - ctx: ClassificationServeContext, - ) -> PoolingParams | ErrorResponse: - pooling_params = super()._create_pooling_params(ctx) - if isinstance(pooling_params, ErrorResponse): - return pooling_params - - try: - pooling_params.verify("classify", self.model_config) - except ValueError as e: - return self.create_error_response(str(e)) - - return pooling_params diff --git a/vllm/entrypoints/pooling/embed/protocol.py b/vllm/entrypoints/pooling/embed/protocol.py index 1ab6097e792..61bec5ae0ec 100644 --- a/vllm/entrypoints/pooling/embed/protocol.py +++ b/vllm/entrypoints/pooling/embed/protocol.py @@ -5,6 +5,7 @@ from typing import Any, TypeAlias from pydantic import Field +from vllm import PoolingParams from vllm.config import ModelConfig from vllm.entrypoints.openai.engine.protocol import OpenAIBaseModel, UsageInfo from vllm.entrypoints.pooling.base.protocol import ( @@ -13,9 +14,12 @@ from vllm.entrypoints.pooling.base.protocol import ( EmbedRequestMixin, PoolingBasicRequestMixin, ) +from vllm.logger import init_logger from vllm.renderers import TokenizeParams from vllm.utils import random_uuid +logger = init_logger(__name__) + def _get_max_total_output_tokens( model_config: ModelConfig, @@ -55,6 +59,21 @@ class EmbeddingCompletionRequest( max_output_tokens_param="max_model_len - max_embed_len", ) + def to_pooling_params(self): + if self.normalize is not None: + logger.warning_once( + "`normalize` is deprecated and will be removed in v0.17. " + "Please pass `use_activation` instead." + ) + self.use_activation = self.normalize + + return PoolingParams( + task="embed", + dimensions=self.dimensions, + use_activation=self.use_activation, + truncate_prompt_tokens=self.truncate_prompt_tokens, + ) + class EmbeddingChatRequest( PoolingBasicRequestMixin, ChatRequestMixin, EmbedRequestMixin @@ -82,6 +101,21 @@ class EmbeddingChatRequest( max_output_tokens_param="max_model_len - max_embed_len", ) + def to_pooling_params(self): + if self.normalize is not None: + logger.warning_once( + "`normalize` is deprecated and will be removed in v0.17. " + "Please pass `use_activation` instead." + ) + self.use_activation = self.normalize + + return PoolingParams( + task="embed", + dimensions=self.dimensions, + use_activation=self.use_activation, + truncate_prompt_tokens=self.truncate_prompt_tokens, + ) + EmbeddingRequest: TypeAlias = EmbeddingCompletionRequest | EmbeddingChatRequest diff --git a/vllm/entrypoints/pooling/embed/serving.py b/vllm/entrypoints/pooling/embed/serving.py index a535801351c..e1b77637727 100644 --- a/vllm/entrypoints/pooling/embed/serving.py +++ b/vllm/entrypoints/pooling/embed/serving.py @@ -424,12 +424,6 @@ class OpenAIServingEmbedding(OpenAIServing): if isinstance(pooling_params, ErrorResponse): return pooling_params - # Verify and set the task for pooling params - try: - pooling_params.verify("embed", self.model_config) - except ValueError as e: - return self.create_error_response(str(e)) - if ctx.engine_prompts is None: return self.create_error_response("Engine prompts not available") @@ -463,8 +457,7 @@ class OpenAIServingEmbedding(OpenAIServing): return None except Exception as e: - # TODO: Use a vllm-specific Validation Error - return self.create_error_response(str(e)) + return self.create_error_response(e) async def _collect_batch( self, @@ -634,7 +627,7 @@ class OpenAIServingEmbedding(OpenAIServing): return None except Exception as e: - return self.create_error_response(str(e)) + return self.create_error_response(e) async def create_embedding( self, @@ -661,18 +654,3 @@ class OpenAIServingEmbedding(OpenAIServing): ) return await self.handle(ctx) # type: ignore[return-value] - - def _create_pooling_params( - self, - ctx: EmbeddingServeContext, - ) -> PoolingParams | ErrorResponse: - pooling_params = super()._create_pooling_params(ctx) - if isinstance(pooling_params, ErrorResponse): - return pooling_params - - try: - pooling_params.verify("embed", self.model_config) - except ValueError as e: - return self.create_error_response(str(e)) - - return pooling_params diff --git a/vllm/entrypoints/pooling/pooling/protocol.py b/vllm/entrypoints/pooling/pooling/protocol.py index 4818f851c97..50d769e1f7d 100644 --- a/vllm/entrypoints/pooling/pooling/protocol.py +++ b/vllm/entrypoints/pooling/pooling/protocol.py @@ -53,6 +53,7 @@ class PoolingCompletionRequest( self.use_activation = self.normalize return PoolingParams( + task=self.task, truncate_prompt_tokens=self.truncate_prompt_tokens, use_activation=self.use_activation, dimensions=self.dimensions, @@ -90,6 +91,7 @@ class PoolingChatRequest( self.use_activation = self.normalize return PoolingParams( + task=self.task, truncate_prompt_tokens=self.truncate_prompt_tokens, use_activation=self.use_activation, dimensions=self.dimensions, @@ -104,7 +106,7 @@ class IOProcessorRequest(PoolingBasicRequestMixin, EncodingRequestMixin, Generic task: PoolingTask = "plugin" def to_pooling_params(self): - return PoolingParams() + return PoolingParams(task=self.task) class IOProcessorResponse(OpenAIBaseModel, Generic[T]): diff --git a/vllm/entrypoints/pooling/pooling/serving.py b/vllm/entrypoints/pooling/pooling/serving.py index 423474ca958..faf5a09d4ab 100644 --- a/vllm/entrypoints/pooling/pooling/serving.py +++ b/vllm/entrypoints/pooling/pooling/serving.py @@ -35,7 +35,6 @@ from vllm.entrypoints.pooling.utils import ( ) from vllm.logger import init_logger from vllm.outputs import PoolingRequestOutput -from vllm.tasks import PoolingTask, SupportedTask from vllm.utils.async_utils import merge_async_iterators from vllm.utils.serial_utils import EmbedDType, EncodingFormat, Endianness @@ -48,7 +47,6 @@ class OpenAIServingPooling(OpenAIServing): engine_client: EngineClient, models: OpenAIServingModels, *, - supported_tasks: tuple[SupportedTask, ...], request_logger: RequestLogger | None, chat_template: str | None, chat_template_content_format: ChatTemplateContentFormatOption, @@ -62,7 +60,6 @@ class OpenAIServingPooling(OpenAIServing): log_error_stack=log_error_stack, ) - self.supported_tasks = supported_tasks self.chat_template = chat_template self.chat_template_content_format: Final = chat_template_content_format self.trust_request_chat_template = trust_request_chat_template @@ -152,32 +149,6 @@ class OpenAIServingPooling(OpenAIServing): else: pooling_params = request.to_pooling_params() - pooling_task: PoolingTask - if request.task is None: - if "token_embed" in self.supported_tasks: - pooling_task = "token_embed" - elif "token_classify" in self.supported_tasks: - pooling_task = "token_classify" - elif "plugin" in self.supported_tasks: - pooling_task = "plugin" - else: - return self.create_error_response( - f"pooling_task must be one of {self.supported_tasks}." - ) - else: - pooling_task = request.task - - if pooling_task not in self.supported_tasks: - return self.create_error_response( - f"Task {pooling_task} is not supported, it" - f" must be one of {self.supported_tasks}." - ) - - try: - pooling_params.verify(pooling_task, self.model_config) - except ValueError as e: - return self.create_error_response(str(e)) - for i, engine_prompt in enumerate(engine_prompts): request_id_item = f"{request_id}-{i}" @@ -212,8 +183,7 @@ class OpenAIServingPooling(OpenAIServing): generators.append(generator) except ValueError as e: - # TODO: Use a vllm-specific Validation Error - return self.create_error_response(str(e)) + return self.create_error_response(e) result_generator = merge_async_iterators(*generators) @@ -251,8 +221,7 @@ class OpenAIServingPooling(OpenAIServing): except asyncio.CancelledError: return self.create_error_response("Client disconnected") except ValueError as e: - # TODO: Use a vllm-specific Validation Error - return self.create_error_response(str(e)) + return self.create_error_response(e) return response diff --git a/vllm/entrypoints/pooling/score/protocol.py b/vllm/entrypoints/pooling/score/protocol.py index 1a7b0520327..9fe9f25445b 100644 --- a/vllm/entrypoints/pooling/score/protocol.py +++ b/vllm/entrypoints/pooling/score/protocol.py @@ -18,6 +18,7 @@ from vllm.entrypoints.pooling.score.utils import ( ScoreInputs, ) from vllm.renderers import TokenizeParams +from vllm.tasks import PoolingTask from vllm.utils import random_uuid @@ -40,8 +41,9 @@ class ScoreRequestMixin(PoolingBasicRequestMixin, ClassifyRequestMixin): max_total_tokens_param="max_model_len", ) - def to_pooling_params(self): + def to_pooling_params(self, task: PoolingTask = "score"): return PoolingParams( + task=task, truncate_prompt_tokens=self.truncate_prompt_tokens, use_activation=self.use_activation, ) @@ -122,6 +124,13 @@ class RerankRequest(PoolingBasicRequestMixin, ClassifyRequestMixin): max_total_tokens_param="max_model_len", ) + def to_pooling_params(self, task: PoolingTask = "score"): + return PoolingParams( + task=task, + truncate_prompt_tokens=self.truncate_prompt_tokens, + use_activation=self.use_activation, + ) + class RerankDocument(BaseModel): text: str | None = None diff --git a/vllm/entrypoints/pooling/score/serving.py b/vllm/entrypoints/pooling/score/serving.py index 9ef3b9afffb..12f9bb7efc5 100644 --- a/vllm/entrypoints/pooling/score/serving.py +++ b/vllm/entrypoints/pooling/score/serving.py @@ -118,12 +118,7 @@ class ServingScores(OpenAIServing): # Schedule the request and get the result generator. generators: list[AsyncGenerator[PoolingRequestOutput, None]] = [] - pooling_params = request.to_pooling_params() - - try: - pooling_params.verify("embed", self.model_config) - except ValueError as e: - return self.create_error_response(str(e)) + pooling_params = request.to_pooling_params("embed") for i, engine_prompt in enumerate(engine_prompts): request_id_item = f"{request_id}-{i}" @@ -223,19 +218,7 @@ class ServingScores(OpenAIServing): # Schedule the request and get the result generator. generators: list[AsyncGenerator[PoolingRequestOutput, None]] = [] - # Use token_embed task for late interaction models - from vllm import PoolingParams - - pooling_params = PoolingParams( - task="token_embed", - truncate_prompt_tokens=request.truncate_prompt_tokens, - use_activation=request.use_activation, - ) - - try: - pooling_params.verify("token_embed", self.model_config) - except ValueError as e: - return self.create_error_response(str(e)) + pooling_params = request.to_pooling_params("token_embed") for i, engine_prompt in enumerate(engine_prompts): request_id_item = f"{request_id}-{i}" @@ -358,12 +341,7 @@ class ServingScores(OpenAIServing): # Schedule the request and get the result generator. generators: list[AsyncGenerator[PoolingRequestOutput, None]] = [] - default_pooling_params = request.to_pooling_params() - - try: - default_pooling_params.verify("score", self.model_config) - except ValueError as e: - return self.create_error_response(str(e)) + default_pooling_params = request.to_pooling_params("score") for i, engine_prompt in enumerate(engine_prompts): request_id_item = f"{request_id}-{i}" @@ -497,8 +475,7 @@ class ServingScores(OpenAIServing): except asyncio.CancelledError: return self.create_error_response("Client disconnected") except ValueError as e: - # TODO: Use a vllm-specific Validation Error - return self.create_error_response(str(e)) + return self.create_error_response(e) async def do_rerank( self, request: RerankRequest, raw_request: Request | None = None @@ -542,8 +519,7 @@ class ServingScores(OpenAIServing): except asyncio.CancelledError: return self.create_error_response("Client disconnected") except ValueError as e: - # TODO: Use a vllm-specific Validation Error - return self.create_error_response(str(e)) + return self.create_error_response(e) def request_output_to_score_response( self, diff --git a/vllm/pooling_params.py b/vllm/pooling_params.py index 1beb6906b8c..2251cceefd8 100644 --- a/vllm/pooling_params.py +++ b/vllm/pooling_params.py @@ -72,15 +72,7 @@ class PoolingParams( """Returns a deep copy of the PoolingParams instance.""" return deepcopy(self) - def verify( - self, task: PoolingTask, model_config: "ModelConfig | None" = None - ) -> None: - if self.task is None: - self.task = task - elif self.task != task: - msg = f"You cannot overwrite {self.task=!r} with {task=!r}!" - raise ValueError(msg) - + def verify(self, model_config: "ModelConfig") -> None: # plugin task uses io_processor.parse_request to verify inputs, # skipping PoolingParams verify if self.task == "plugin": @@ -167,7 +159,7 @@ class PoolingParams( if mds is not None: if self.dimensions not in mds: raise ValueError( - f'Model "{model_config.served_model_name}" ' + f"Model {model_config.served_model_name!r} " f"only supports {str(mds)} matryoshka dimensions, " f"use other output dimensions will " f"lead to poor results." @@ -179,7 +171,7 @@ class PoolingParams( if self.use_activation is None: self.use_activation = True else: - raise ValueError(f"Unknown pooling task: {self.task}") + raise ValueError(f"Unknown pooling task: {self.task!r}") def _verify_valid_parameters(self): assert self.task is not None, "task must be set" @@ -194,7 +186,7 @@ class PoolingParams( if invalid_parameters: raise ValueError( - f"Task {self.task} only supports {valid_parameters} " + f"Task {self.task!r} only supports {valid_parameters} " f"parameters, does not support " f"{invalid_parameters} parameters" ) diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 2beb9c4f8c7..c0613137091 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -269,7 +269,11 @@ class AsyncLLM(EngineClient): cancel_task_threadsafe(handler) async def get_supported_tasks(self) -> tuple[SupportedTask, ...]: - return await self.engine_core.get_supported_tasks_async() + if not hasattr(self, "_supported_tasks"): + # Cache the result + self._supported_tasks = await self.engine_core.get_supported_tasks_async() + + return self._supported_tasks async def add_request( self, @@ -355,6 +359,7 @@ class AsyncLLM(EngineClient): trace_headers=trace_headers, priority=priority, data_parallel_rank=data_parallel_rank, + supported_tasks=await self.get_supported_tasks(), ) prompt_text = get_prompt_text(prompt) diff --git a/vllm/v1/engine/input_processor.py b/vllm/v1/engine/input_processor.py index b4de0e50c95..01cf999c7d2 100644 --- a/vllm/v1/engine/input_processor.py +++ b/vllm/v1/engine/input_processor.py @@ -31,6 +31,7 @@ from vllm.multimodal.utils import argsort_mm_positions from vllm.pooling_params import PoolingParams from vllm.renderers import BaseRenderer from vllm.sampling_params import _SAMPLING_EPS, SamplingParams +from vllm.tasks import POOLING_TASKS, SupportedTask from vllm.tokenizers import TokenizerLike from vllm.tokenizers.mistral import MistralTokenizer from vllm.utils import length_from_prompt_token_ids_or_embeds, random_uuid @@ -196,13 +197,41 @@ class InputProcessor: def _validate_params( self, params: SamplingParams | PoolingParams, + # TODO: Validate generation tasks as well once `supported_tasks` + # is passed to all `process_inputs` calls + supported_tasks: tuple[SupportedTask, ...] | None, ): """ Validate supported SamplingParam. Should raise ValueError if unsupported for API Server. """ - if isinstance(params, PoolingParams): + if supported_tasks is None: + raise RuntimeError("`supported_tasks` must be passed for pooling") + + supported_pooling_tasks = [ + task for task in supported_tasks if task in POOLING_TASKS + ] + + if params.task is None: + if not supported_pooling_tasks: + raise ValueError("Pooling tasks are not supported") + + if "token_embed" in supported_pooling_tasks: + params.task = "token_embed" + elif "token_classify" in supported_pooling_tasks: + params.task = "token_classify" + elif "plugin" in supported_pooling_tasks: + params.task = "plugin" + + if params.task not in supported_pooling_tasks: + raise ValueError( + f"Unsupported task: {params.task!r} " + f"Supported tasks: {supported_pooling_tasks}" + ) + + params.verify(self.model_config) + return self._validate_logprobs(params) @@ -498,10 +527,11 @@ class InputProcessor: trace_headers: Mapping[str, str] | None = None, priority: int = 0, data_parallel_rank: int | None = None, + supported_tasks: tuple[SupportedTask, ...] | None = None, resumable: bool = False, ) -> EngineCoreRequest: self._validate_lora(lora_request) - self._validate_params(params) + self._validate_params(params, supported_tasks) parallel_config = self.vllm_config.parallel_config dp_size = parallel_config.data_parallel_size diff --git a/vllm/v1/engine/llm_engine.py b/vllm/v1/engine/llm_engine.py index 0f1b7f46d2f..9cae71a4348 100644 --- a/vllm/v1/engine/llm_engine.py +++ b/vllm/v1/engine/llm_engine.py @@ -201,7 +201,11 @@ class LLMEngine: return outputs def get_supported_tasks(self) -> tuple[SupportedTask, ...]: - return self.engine_core.get_supported_tasks() + if not hasattr(self, "_supported_tasks"): + # Cache the result + self._supported_tasks = self.engine_core.get_supported_tasks() + + return self._supported_tasks def abort_request(self, request_ids: list[str], internal: bool = False) -> None: """Remove request_ids from EngineCore and Detokenizer.""" @@ -245,6 +249,7 @@ class LLMEngine: tokenization_kwargs, trace_headers, priority, + supported_tasks=self.get_supported_tasks(), ) prompt_text = get_prompt_text(prompt) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 39ac6bce820..10d4dfd3309 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -5037,7 +5037,7 @@ class GPUModelRunner( model = cast(VllmModelForPooling, self.get_model()) dummy_pooling_params = PoolingParams(task=task) - dummy_pooling_params.verify(task=task, model_config=self.model_config) + dummy_pooling_params.verify(self.model_config) to_update = model.pooler.get_pooling_updates(task) to_update.apply(dummy_pooling_params) From 3e472e81f99b5bcf494369ee2d26ee9d6ceeffe3 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Thu, 5 Feb 2026 04:01:23 -0600 Subject: [PATCH 093/810] [ROCm][Bugfix][CI] Fix hybrid models and their tests (Mamba/Jamba/Bamba) (#32710) Signed-off-by: Andreas Karatzas Signed-off-by: Matthew Wong Co-authored-by: Matthew Wong --- tests/models/language/generation/test_hybrid.py | 5 +++++ vllm/model_executor/layers/mamba/mamba_mixer.py | 6 ++++++ 2 files changed, 11 insertions(+) diff --git a/tests/models/language/generation/test_hybrid.py b/tests/models/language/generation/test_hybrid.py index c3e6d7899e2..2724f612cee 100644 --- a/tests/models/language/generation/test_hybrid.py +++ b/tests/models/language/generation/test_hybrid.py @@ -8,6 +8,7 @@ import pytest from tests.models.registry import HF_EXAMPLE_MODELS from tests.utils import multi_gpu_test from vllm.engine.arg_utils import EngineArgs +from vllm.platforms import current_platform from vllm.sampling_params import SamplingParams from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher @@ -577,6 +578,10 @@ def test_apc_multiple_prompts_all_cached_outputs( model, max_model_len, tensor_parallel_size=tensor_parallel_size ) vllm_runner_kwargs["mamba_ssm_cache_dtype"] = "float32" + # Reduce the effects of batch variance on ROCm since batch invariance is not + # yet supported. See: https://github.com/vllm-project/vllm/issues/27433 + if current_platform.is_rocm(): + vllm_runner_kwargs["max_num_seqs"] = 4 vllm_outputs_no_cache, _ = _get_vLLM_output( vllm_runner, vllm_runner_kwargs, generated_prompts, max_tokens, num_logprobs diff --git a/vllm/model_executor/layers/mamba/mamba_mixer.py b/vllm/model_executor/layers/mamba/mamba_mixer.py index 134e1dfd628..adc643c385e 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer.py @@ -214,6 +214,12 @@ class MambaMixer(MambaBase, CustomOp): time_step = self.dt_layernorm(time_step.contiguous()) B = self.b_layernorm(B.contiguous()) C = self.c_layernorm(C.contiguous()) + + # ROCm: tensor from split is non-contiguous, causing incorrect + # GEMM results in dt_proj. + if current_platform.is_rocm(): + time_step = time_step.contiguous() + discrete_time_step = self.dt_proj(time_step)[0].transpose(-2, -1) return discrete_time_step, B, C From 8322d4e47f89f7985b9b3b808fc4ba8549d6afcd Mon Sep 17 00:00:00 2001 From: liranschour Date: Thu, 5 Feb 2026 12:17:02 +0200 Subject: [PATCH 094/810] Enable Cross layers KV cache layout at NIXL Connector V2 (#33339) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Liran Schour Signed-off-by: liranschour Co-authored-by: Or Ozeri Co-authored-by: Nicolò Lucchesi Co-authored-by: Nicolò Lucchesi --- docs/features/nixl_connector_usage.md | 9 + .../config_sweep_accuracy_test.sh | 10 +- .../nixl_integration/run_accuracy_test.sh | 15 +- .../kv_connector/unit/test_nixl_connector.py | 235 ++++++++++++++---- .../kv_transfer/kv_connector/utils.py | 49 +++- .../kv_connector/v1/nixl_connector.py | 110 +++++--- 6 files changed, 339 insertions(+), 89 deletions(-) diff --git a/docs/features/nixl_connector_usage.md b/docs/features/nixl_connector_usage.md index b8364b237e9..3fc735efa68 100644 --- a/docs/features/nixl_connector_usage.md +++ b/docs/features/nixl_connector_usage.md @@ -213,6 +213,15 @@ Support use case: Prefill with 'HND' and decode with 'NHD' with experimental con --kv-transfer-config '{..., "enable_permute_local_kv":"True"}' ``` +### Cross layers blocks + +By default, this feature is disabled. On attention backends that support this feature, each logical block is contiguous in physical memory. This reduces the number of buffers that need to be transferred. +To enable this feature: + +```bash +--kv-transfer-config '{..., "kv_connector_extra_config": {"enable_cross_layers_blocks": "True"}}' +``` + ## Example Scripts/Code Refer to these example scripts in the vLLM repository: diff --git a/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh b/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh index 2e25e2f1ac3..cdbcdca546e 100755 --- a/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh +++ b/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh @@ -14,8 +14,8 @@ tp_configs=( "GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=1 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" ) dp_ep_configs=( -"DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA+P-TP1, D-DPEP=2 (TP=1) -"DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA+P-TP2, D-DPEP=2 (TP=1) +"DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA+P-TP1, D-DPEP=2 (TP=1) +"DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA+P-TP2, D-DPEP=2 (TP=1) ) # Select config array based on DP_EP env var @@ -57,3 +57,9 @@ if [[ -n "${FLASHINFER:-}" ]]; then else echo "FLASHINFER not set, skipping FLASHINFER runs." fi + +# Check if cross-layers is enabled (non-empty) +if [[ -n "${CROSS_LAYERS_BLOCKS:-}" ]]; then + echo "CROSS_LAYERS_BLOCKS is set, rerunning with --enable-cross-layers" + run_tests "default backend" "--enable-cross-layers" +fi 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 c2c38f51c50..560ce440703 100755 --- a/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh +++ b/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh @@ -4,6 +4,7 @@ set -xe # Parse command line arguments KV_BUFFER_DEVICE="cuda" # Default to cuda ATTENTION_BACKEND="" # Default to empty (use vllm default) +CROSS_LAYERS_BLOCKS="False" while [[ $# -gt 0 ]]; do case $1 in --kv_buffer_device) @@ -14,6 +15,10 @@ while [[ $# -gt 0 ]]; do ATTENTION_BACKEND="$2" shift 2 ;; + --enable-cross-layers) + CROSS_LAYERS_BLOCKS="True" + shift 1 + ;; *) echo "Unknown option $1" echo "Usage: $0 [--kv_buffer_device ] [--attention-backend ]" @@ -34,11 +39,17 @@ else KV_CONFIG_HETERO_LAYOUT='' fi +if [[ "$CROSS_LAYERS_BLOCKS" == "True" ]]; then + KV_EXTRA_CONFIG=',"kv_connector_extra_config":{"enable_cross_layers_blocks": "True"}' +else + KV_EXTRA_CONFIG='' +fi + # Build the kv-transfer-config once if [[ "$KV_BUFFER_DEVICE" == "cuda" ]]; then - KV_CONFIG='{"kv_connector":"NixlConnector","kv_role":"kv_both"'${KV_CONFIG_HETERO_LAYOUT}'}' + KV_CONFIG='{"kv_connector":"NixlConnector","kv_role":"kv_both"'${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}'}' else - KV_CONFIG="{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\""${KV_CONFIG_HETERO_LAYOUT}"}" + KV_CONFIG="{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\""${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}"}" fi # Models to run diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index a3fe2e21b56..1975d222607 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -18,8 +18,12 @@ import ray import torch from vllm import LLM -from vllm.config import KVTransferConfig -from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator +from vllm.config import KVTransferConfig, set_current_vllm_config +from vllm.distributed.kv_transfer.kv_connector.utils import ( + KVOutputAggregator, + TpKVTopology, + get_current_attn_backend, +) from vllm.distributed.kv_transfer.kv_connector.v1 import nixl_connector from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import ( @@ -46,10 +50,14 @@ from vllm.platforms import current_platform from vllm.platforms.interface import Platform from vllm.sampling_params import SamplingParams from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend +from vllm.v1.attention.backends.utils import set_kv_cache_layout from vllm.v1.engine import EngineCoreRequest from vllm.v1.engine.output_processor import OutputProcessor +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig, KVCacheTensor from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput from vllm.v1.request import RequestStatus +from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin +from vllm.v1.worker.utils import AttentionGroup from .utils import create_request, create_scheduler, create_vllm_config @@ -366,6 +374,7 @@ def test_kv_transfer_handshake(dist_init): # Decode connector will be able to create handshake with the prefill connector. decode_connector = NixlConnector(vllm_config, KVConnectorRole.WORKER) + decode_connector.register_kv_caches(kv_caches) # Here we are testing the retrieval of NIXLAgentMetadata. # Knowing the implementation detail, we override the add_remote_agent @@ -402,6 +411,23 @@ class FakeNixlConnectorWorker(NixlConnectorWorker): self.kv_cache_layout = kv_cache_layout # Mock register_kv_caches attribute needed for tests that do not call it. self.src_xfer_handles_by_block_size = {self.block_size: 1} + test_shape = self.attn_backend.get_kv_cache_shape( + num_blocks=1, block_size=16, num_kv_heads=1, head_size=1 + ) + self.kv_topo = TpKVTopology( + tp_rank=self.tp_rank, + engine_id=self.engine_id, + remote_tp_size=self._tp_size, # shared state + remote_block_size=self._block_size, # shared state + is_mla=self.use_mla, + total_num_kv_heads=self.model_config.get_total_num_kv_heads(), + attn_backend=self.attn_backend, + tensor_shape=test_shape, + ) + + self.compat_hash = compute_nixl_compatibility_hash( + self.vllm_config, self.backend_name, self.kv_topo.cross_layers_blocks + ) def _nixl_handshake( self, host: str, port: int, remote_tp_size: int, expected_engine_id: str @@ -1352,6 +1378,7 @@ def _run_abort_timeout_test(llm: LLM, timeout: int): llm.llm_engine.engine_core.shutdown() +@pytest.mark.parametrize("enable_cross_layers", ["False", "True"]) @pytest.mark.parametrize( "attn_backend", [ @@ -1372,7 +1399,9 @@ def _run_abort_timeout_test(llm: LLM, timeout: int): "TRITON_ATTN", ], ) -def test_register_kv_caches(default_vllm_config, dist_init, attn_backend): +def test_register_kv_caches( + default_vllm_config, dist_init, attn_backend, enable_cross_layers +): """ Test that register_kv_caches() properly calls nixl_wrapper methods with correct data. @@ -1386,6 +1415,12 @@ def test_register_kv_caches(default_vllm_config, dist_init, attn_backend): vllm_config = create_vllm_config(attention_backend=attn_backend) + # Enable cross layers blocks + vllm_config.kv_transfer_config.kv_connector_extra_config[ + "enable_cross_layers_blocks" + ] = enable_cross_layers + set_kv_cache_layout("HND") + # Import the appropriate backend based on the parameter if attn_backend == "FLASH_ATTN": from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend @@ -1400,44 +1435,6 @@ def test_register_kv_caches(default_vllm_config, dist_init, attn_backend): backend_cls = TritonAttentionBackend - # Create test kv cache tensors using proper backend shape - kv_cache_shape = backend_cls.get_kv_cache_shape( - num_blocks=2, block_size=16, num_kv_heads=4, head_size=64 - ) - shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16) - unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16) - kv_caches = { - "layer0": shared_tensor, - "layer1": unique_tensor, - "layer2": shared_tensor, - } - - # Store tensor info for validation - - test_shape = backend_cls.get_kv_cache_shape( - num_blocks=1, block_size=16, num_kv_heads=1, head_size=1 - ) - is_blocks_first = len(test_shape) == 5 and test_shape[0] == 1 - - if is_blocks_first: - expected_tensor_size = shared_tensor.element_size() * shared_tensor.numel() - expected_base_addrs = [ - shared_tensor.data_ptr(), - unique_tensor.data_ptr(), - ] - expected_num_entries = 2 - else: - expected_tensor_size = ( - shared_tensor[0].element_size() * shared_tensor[0].numel() - ) - expected_base_addrs = [ - shared_tensor[0].data_ptr(), - shared_tensor[1].data_ptr(), - unique_tensor[0].data_ptr(), - unique_tensor[1].data_ptr(), - ] - expected_num_entries = 4 - nixl_module = "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector" with ( patch(f"{nixl_module}.NixlWrapper") as mock_nixl_wrapper, @@ -1466,6 +1463,111 @@ def test_register_kv_caches(default_vllm_config, dist_init, attn_backend): # Reassure the shutdown() check that the thread is terminated mock_thread.return_value.is_alive.return_value = False + expected_tensor_size: int + expected_base_addrs: list[int] + expected_num_entries: int + kv_caches: dict[str, torch.Tensor] + assert str(enable_cross_layers).lower() != "true" or ( + (attn_backend not in ("FLASH_ATTN", "FLASHINFER")) + or connector.prefer_cross_layer_blocks + ) + if connector.prefer_cross_layer_blocks: + num_layers = 32 + block_size = 16 + num_blocks = 8 + kv_cache_spec = AttentionSpec( + block_size=block_size, + num_kv_heads=4, + head_size=64, + dtype=torch.bfloat16, + ) + kv_cache_config = KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=kv_cache_spec.page_size_bytes * num_blocks, + shared_by=["dummy-layer"], + ) + for i in range(num_layers) + ], + # allocate_uniform_kv_caches does not use this + kv_cache_groups=[], + ) + + with set_current_vllm_config(vllm_config): + _, cross_layers_kv_cache, _ = ( + KVConnectorModelRunnerMixin.allocate_uniform_kv_caches( + kv_cache_config=kv_cache_config, + attn_groups=[ + [ + AttentionGroup( + backend=backend_cls, + layer_names=[], + kv_cache_spec=kv_cache_spec, + kv_cache_group_id=0, + ) + ] + ], + cache_dtype=torch.bfloat16, + device=torch.cuda.current_device(), + kernel_block_sizes=[block_size], + ) + ) + # Store tensor info for validation + expected_tensor_size = ( + cross_layers_kv_cache.element_size() * cross_layers_kv_cache.numel() + ) + expected_base_addrs = [ + cross_layers_kv_cache.data_ptr(), + ] + expected_num_entries = 1 + + expected_blocks_count = 8 + + kv_caches = {"all-layers": cross_layers_kv_cache} + + else: + # Create test kv cache tensors using proper backend shape + kv_cache_shape = backend_cls.get_kv_cache_shape( + num_blocks=2, block_size=16, num_kv_heads=4, head_size=64 + ) + shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16) + unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16) + kv_caches = { + "layer0": shared_tensor, + "layer1": unique_tensor, + "layer2": shared_tensor, + } + + # Store tensor info for validation + + test_shape = backend_cls.get_kv_cache_shape( + num_blocks=1, block_size=16, num_kv_heads=1, head_size=1 + ) + is_blocks_first = len(test_shape) == 5 and test_shape[0] == 1 + + if is_blocks_first: + expected_tensor_size = ( + shared_tensor.element_size() * shared_tensor.numel() + ) + expected_base_addrs = [ + shared_tensor.data_ptr(), + unique_tensor.data_ptr(), + ] + expected_num_entries = 2 + else: + expected_tensor_size = ( + shared_tensor[0].element_size() * shared_tensor[0].numel() + ) + expected_base_addrs = [ + shared_tensor[0].data_ptr(), + shared_tensor[1].data_ptr(), + unique_tensor[0].data_ptr(), + unique_tensor[1].data_ptr(), + ] + expected_num_entries = 4 + expected_blocks_count = 8 + # Execute register_kv_caches connector.register_kv_caches(kv_caches) @@ -1489,16 +1591,19 @@ def test_register_kv_caches(default_vllm_config, dist_init, attn_backend): blocks_data, _ = mock_wrapper_instance.get_xfer_descs.call_args[0] # Validate blocks_data structure and size - expected_blocks_count = 8 assert len(blocks_data) == expected_blocks_count, ( f"Expected {expected_blocks_count} blocks, got {len(blocks_data)}" ) - num_blocks = 2 - if is_blocks_first: - expected_block_len = expected_tensor_size // num_blocks // 2 - else: + if connector.prefer_cross_layer_blocks: + num_blocks = 8 expected_block_len = expected_tensor_size // num_blocks + else: + num_blocks = 2 + if is_blocks_first: + expected_block_len = expected_tensor_size // num_blocks // 2 + else: + expected_block_len = expected_tensor_size // num_blocks for i, block_entry in enumerate(blocks_data): block_start_addr, block_len, tp_rank = block_entry @@ -1507,6 +1612,8 @@ def test_register_kv_caches(default_vllm_config, dist_init, attn_backend): f"got {block_len}" ) + assert connector.connector_worker.block_size == 16 + class FakePlatform(Platform): device_type: str = "oot" @@ -2049,6 +2156,17 @@ def test_compatibility_hash_validation( ) decode_connector = NixlConnector(local_vllm_config, KVConnectorRole.WORKER) decode_worker = decode_connector.connector_worker + kv_cache_shape = decode_worker.attn_backend.get_kv_cache_shape( + num_blocks=2, block_size=16, num_kv_heads=4, head_size=64 + ) + shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16) + unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16) + kv_caches = { + "layer0": shared_tensor, + "layer1": unique_tensor, + "layer2": shared_tensor, + } + decode_connector.register_kv_caches(kv_caches) remote_config_params: dict[str, Any] = { "model": "facebook/opt-125m", @@ -2071,7 +2189,9 @@ def test_compatibility_hash_validation( ) ) remote_hash = compute_nixl_compatibility_hash( - remote_vllm_config, decode_worker.backend_name + remote_vllm_config, + decode_worker.backend_name, + decode_worker.kv_topo.cross_layers_blocks, ) prefill_block_size = config_overrides.get("block_size", 16) @@ -2150,6 +2270,27 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario) decode_connector = NixlConnector(local_vllm_config, KVConnectorRole.WORKER) decode_worker = decode_connector.connector_worker + backend = get_current_attn_backend(local_vllm_config) + test_shape = backend.get_kv_cache_shape( + num_blocks=1, block_size=16, num_kv_heads=1, head_size=1 + ) + decode_worker.kv_topo = TpKVTopology( + tp_rank=decode_worker.tp_rank, + engine_id=decode_worker.engine_id, + remote_tp_size=decode_worker._tp_size, # shared state + remote_block_size=decode_worker._block_size, # shared state + is_mla=decode_worker.use_mla, + total_num_kv_heads=decode_worker.model_config.get_total_num_kv_heads(), + attn_backend=backend, + tensor_shape=test_shape, + ) + + decode_worker.compat_hash = compute_nixl_compatibility_hash( + decode_worker.vllm_config, + decode_worker.backend_name, + decode_worker.kv_topo.cross_layers_blocks, + ) + if error_scenario == "handshake_decode_error": msg_bytes = b"this is not valid msgpack data" elif error_scenario == "handshake_validation_error": diff --git a/vllm/distributed/kv_transfer/kv_connector/utils.py b/vllm/distributed/kv_transfer/kv_connector/utils.py index fd833e29393..019201ede73 100644 --- a/vllm/distributed/kv_transfer/kv_connector/utils.py +++ b/vllm/distributed/kv_transfer/kv_connector/utils.py @@ -14,6 +14,7 @@ from vllm.config import VllmConfig, get_current_vllm_config, get_layers_from_vll from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory from vllm.logger import init_logger from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.platforms import current_platform from vllm.v1.attention.backend import AttentionBackend from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput @@ -192,8 +193,6 @@ def copy_kv_blocks( dst_device=dst_device, ) - from vllm.platforms import current_platform - if direction == "h2d": copy_fn = current_platform.insert_blocks_to_device else: @@ -316,12 +315,14 @@ class TpKVTopology: attn_backend: type[AttentionBackend] engine_id: EngineId remote_block_size: dict[EngineId, int] + tensor_shape: torch.Size | None = None def __post_init__(self): # Figure out whether the first dimension of the cache is K/V # or num_blocks. This is used to register the memory regions correctly. + _MOCK_BLOCK_SIZE = 16 kv_cache_shape = self.attn_backend.get_kv_cache_shape( - num_blocks=1, block_size=16, num_kv_heads=1, head_size=1 + num_blocks=1, block_size=_MOCK_BLOCK_SIZE, num_kv_heads=1, head_size=1 ) # Non-MLA backends caches have 5 dims [2, num_blocks, H,N,D], # we just mock num_blocks to 1 for the dimension check below. @@ -329,6 +330,36 @@ class TpKVTopology: len(kv_cache_shape) == 5 and kv_cache_shape[0] == 1 ) + self._cross_layers_blocks = False + if self.tensor_shape is not None: + self._cross_layers_blocks = ( + len(self.tensor_shape) == len(kv_cache_shape) + 1 + ) + + if self._cross_layers_blocks: + # prepend layers dimension + _MOCK_NUM_LAYERS = 80 + kv_cache_shape = (_MOCK_NUM_LAYERS,) + kv_cache_shape + try: + kv_cache_stride_order = self.attn_backend.get_kv_cache_stride_order( + include_num_layers_dimension=self._cross_layers_blocks + ) + except (AttributeError, NotImplementedError): + kv_cache_stride_order = tuple(range(len(self.tensor_shape))) + + # In case of cross layers permute kv_cache_shape according to + # stride_order to retrieve physical position of block_size + kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order) + + # In the default non-cross layers layout the block_size position + # is logical while in the cross layers case it is the physical + # position. This matches the shape of the actual kv cache tensors + # passed at register_kv_caches()/register_cross_layers_kv_cache() + block_size_position = kv_cache_shape.index(_MOCK_BLOCK_SIZE) + + assert block_size_position is not None + self._block_size_position = -(len(kv_cache_shape) - block_size_position) + @property def is_kv_layout_blocks_first(self) -> bool: return self._is_kv_layout_blocks_first @@ -336,7 +367,9 @@ class TpKVTopology: @property def split_k_and_v(self) -> bool: # Whether to register regions for K and V separately (when present). - return not (self.is_mla or self.is_kv_layout_blocks_first) + return not ( + self._cross_layers_blocks or self.is_mla or self.is_kv_layout_blocks_first + ) @property def tp_size(self) -> int: @@ -346,6 +379,14 @@ class TpKVTopology: def block_size(self) -> int: return self.remote_block_size[self.engine_id] + @property + def cross_layers_blocks(self) -> bool: + return self._cross_layers_blocks + + @property + def block_size_position(self) -> int: + return self._block_size_position + def tp_ratio( self, remote_tp_size: int, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py index d03d7086039..8ce939ee405 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py @@ -54,7 +54,7 @@ from vllm.forward_context import ForwardContext from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.utils.network_utils import make_zmq_path, make_zmq_socket -from vllm.v1.attention.backend import AttentionMetadata +from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata from vllm.v1.attention.backends.utils import get_kv_cache_layout from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.worker.block_table import BlockTable @@ -173,7 +173,7 @@ class NixlHandshakePayload(KVConnectorHandshakeMetadata): def compute_nixl_compatibility_hash( - vllm_config: VllmConfig, attn_backend_name: str + vllm_config: VllmConfig, attn_backend_name: str, cross_layers_blocks: bool ) -> str: """ Compute compatibility hash for NIXL KV transfer. @@ -216,6 +216,7 @@ def compute_nixl_compatibility_hash( # Attention backend and KV cache dtype affect memory layout "attn_backend_name": attn_backend_name, "cache_dtype": str(cache_config.cache_dtype), + "cross_layers_blocks": cross_layers_blocks, } compat_hash = hash_factors(factors) @@ -298,6 +299,26 @@ class NixlConnectorMetadata(KVConnectorMetadata): class NixlConnector(KVConnectorBase_V1): + @property + def prefer_cross_layer_blocks(self) -> bool: + backend = get_current_attn_backend(self._vllm_config) + if backend.get_name() not in ( + "FLASH_ATTN", + "FLASHINFER", + ): + return False + + # For now there is no benefit to run cross layers when backend + # does not support on HND + if get_kv_cache_layout() != "HND": + return False + + extra_config = self.kv_transfer_config.kv_connector_extra_config + return ( + str(extra_config.get("enable_cross_layers_blocks", "False")).lower() + == "true" + ) + def __init__( self, vllm_config: VllmConfig, @@ -309,7 +330,7 @@ class NixlConnector(KVConnectorBase_V1): assert vllm_config.kv_transfer_config is not None assert vllm_config.kv_transfer_config.engine_id is not None self.engine_id: EngineId = vllm_config.kv_transfer_config.engine_id - + self.kv_transfer_config = vllm_config.kv_transfer_config if role == KVConnectorRole.SCHEDULER: self.connector_scheduler: NixlConnectorScheduler | None = ( NixlConnectorScheduler(vllm_config, self.engine_id) @@ -395,6 +416,16 @@ class NixlConnector(KVConnectorBase_V1): assert self.connector_worker is not None self.connector_worker.register_kv_caches(kv_caches) + def register_cross_layers_kv_cache( + self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend] + ): + assert self.connector_worker is not None + + cross_layer_name = "ALL_LAYERS" + kv_caches = {cross_layer_name: kv_cache} + + self.connector_worker.register_kv_caches(kv_caches) + def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp): assert self.connector_worker is not None self.connector_worker.set_host_xfer_buffer_ops(copy_operation) @@ -976,20 +1007,17 @@ class NixlConnectorWorker: # Get the attention backend from the first layer # NOTE (NickLucche) models with multiple backends are not supported yet - backend = get_current_attn_backend(vllm_config) + self.attn_backend = get_current_attn_backend(vllm_config) - self.backend_name = backend.get_name() + self.backend_name = self.attn_backend.get_name() self.kv_cache_layout = get_kv_cache_layout() self.host_buffer_kv_cache_layout = self.kv_cache_layout logger.debug("Detected attention backend %s", self.backend_name) logger.debug("Detected kv cache layout %s", self.kv_cache_layout) - self.compat_hash = compute_nixl_compatibility_hash( - self.vllm_config, self.backend_name - ) - self.enforce_compat_hash = self.kv_transfer_config.get_from_extra_config( - "enforce_handshake_compat", True - ) + # lazy initialized in register_kv_caches + self.compat_hash: str | None = None + self.kv_topo: TpKVTopology | None = None self._tp_size: dict[EngineId, int] = {self.engine_id: self.world_size} self._block_size: dict[EngineId, int] = {self.engine_id: self.block_size} @@ -998,17 +1026,12 @@ class NixlConnectorWorker: self.consumer_notification_counts_by_req = defaultdict[ReqId, int](int) self.xfer_stats = NixlKVConnectorStats() - self.kv_topo = TpKVTopology( - tp_rank=self.tp_rank, - engine_id=self.engine_id, - remote_tp_size=self._tp_size, # shared state - remote_block_size=self._block_size, # shared state - is_mla=self.use_mla, - total_num_kv_heads=self.model_config.get_total_num_kv_heads(), - attn_backend=backend, - ) self._physical_blocks_per_logical_kv_block = 1 + self.enforce_compat_hash = self.kv_transfer_config.get_from_extra_config( + "enforce_handshake_compat", True + ) + def _nixl_handshake( self, host: str, @@ -1022,6 +1045,7 @@ class NixlConnectorWorker: # Regardless, only handshake with the remote TP rank(s) that current # local rank will read from. Note that With homogeneous TP, # this happens to be the same single rank_i. + assert self.kv_topo is not None p_remote_ranks = self.kv_topo.get_target_remote_ranks(remote_tp_size) remote_rank_to_agent_name = {} path = make_zmq_path("tcp", host, port) @@ -1059,6 +1083,7 @@ class NixlConnectorWorker: ) # Check compatibility hash BEFORE decoding agent metadata + assert self.compat_hash is not None if ( self.enforce_compat_hash and handshake_payload.compatibility_hash != self.compat_hash @@ -1267,6 +1292,20 @@ class NixlConnectorWorker: def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): """Register the KV Cache data in nixl.""" + self.kv_topo = TpKVTopology( + tp_rank=self.tp_rank, + engine_id=self.engine_id, + remote_tp_size=self._tp_size, # shared state + remote_block_size=self._block_size, # shared state + is_mla=self.use_mla, + total_num_kv_heads=self.model_config.get_total_num_kv_heads(), + attn_backend=self.attn_backend, + tensor_shape=next(iter(kv_caches.values())).shape, + ) + self.compat_hash = compute_nixl_compatibility_hash( + self.vllm_config, self.backend_name, self.kv_topo.cross_layers_blocks + ) + if self.use_host_buffer: self.initialize_host_xfer_buffer(kv_caches=kv_caches) assert len(self.host_xfer_buffers) == len(kv_caches), ( @@ -1301,29 +1340,21 @@ class NixlConnectorWorker: # (roughly 8KB vs 5KB). # Conversely for FlashInfer, K and V are registered in the same region # to better exploit the memory layout (ie num_blocks is the first dim). - split_k_and_v = self.kv_topo.split_k_and_v tensor_size_bytes = None - # TODO (NickLucche): Get kernel_block_size in a cleaner way - # NHD default "view" for non-MLA cache - if self.device_type == "cpu": - block_size_position = -2 - else: - block_size_position = -2 if self.use_mla else -3 - # Enable different block lengths for different layers when MLA is used. self.block_len_per_layer = list[int]() self.slot_size_per_layer = list[int]() # HD bytes in kv terms for layer_name, cache_or_caches in xfer_buffers.items(): - cache_list = cache_or_caches if split_k_and_v else [cache_or_caches] - + cache_list = ( + cache_or_caches if self.kv_topo.split_k_and_v else [cache_or_caches] + ) for cache in cache_list: base_addr = cache.data_ptr() if base_addr in seen_base_addresses: continue - kernel_block_size = cache.shape[block_size_position] - + kernel_block_size = cache.shape[self.kv_topo.block_size_position] if self.block_size != kernel_block_size: logger.info_once( "User-specified logical block size (%s) does not match" @@ -1385,6 +1416,7 @@ class NixlConnectorWorker: self.device_kv_caches = kv_caches self.dst_num_blocks[self.engine_id] = self.num_blocks + if self.kv_topo.is_kv_layout_blocks_first: for i in range(len(self.slot_size_per_layer)): assert self.slot_size_per_layer[i] % 2 == 0 @@ -1440,6 +1472,7 @@ class NixlConnectorWorker: block_size=self.block_size, ) # Wrap metadata in payload with hash for defensive decoding + assert self.compat_hash is not None encoder = msgspec.msgpack.Encoder() self.xfer_handshake_metadata = NixlHandshakePayload( compatibility_hash=self.compat_hash, @@ -1461,6 +1494,8 @@ class NixlConnectorWorker: register another local_xfer_handler using remote block len to ensure data copy correctness. """ + assert self.kv_topo is not None + block_size_ratio = self.block_size // block_size blocks_data = [] for i, base_addr in enumerate(self.seen_base_addresses): @@ -1573,6 +1608,7 @@ class NixlConnectorWorker: # remote: | 0| 1| 2| 3| 4| 5| 6| 7| 8| 9|10|11|12| # local origin:| 0| 1| 8| 12| # local mapped:| 0| 1| 2| 3| 4| 5| 6| 7| 8| 9|10|11|12|13|14|15| + assert self.kv_topo is not None block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(engine_id) if engine_id not in self.dst_num_blocks: @@ -1701,6 +1737,7 @@ class NixlConnectorWorker: remote_engine_id = nixl_agent_meta.engine_id assert self._tp_size[remote_engine_id] == remote_tp_size + assert self.kv_topo is not None tp_ratio = self.kv_topo.tp_ratio_from_engine_id(remote_engine_id) block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id( @@ -1837,6 +1874,7 @@ class NixlConnectorWorker: if len(self.device_kv_caches) == 0: return assert block_size_ratio >= 1, "Only nP < nD supported currently." + assert self.kv_topo is not None if self.enable_permute_local_kv and block_size_ratio > 1: logger.debug( "Post-processing device kv cache on receive by converting " @@ -1856,7 +1894,7 @@ class NixlConnectorWorker: block_size_ratio, ) - split_k_and_v = not (self.use_mla or self.kv_topo.is_kv_layout_blocks_first) + split_k_and_v = self.kv_topo.split_k_and_v for block_ids in block_ids_list: indices = torch.tensor(block_ids, device=self.device_type, dtype=torch.long) @@ -1881,6 +1919,7 @@ class NixlConnectorWorker: The scheduler process (via the MultiprocExecutor) will use this output to track which workers are done. """ + assert self.kv_topo is not None done_sending = self._get_new_notifs() done_recving = self._pop_done_transfers(self._recving_transfers) @@ -1950,6 +1989,7 @@ class NixlConnectorWorker: are reading from the same producer (heterogeneous TP scenario), wait for all consumers to be done pulling. """ + assert self.kv_topo is not None notified_req_ids: set[str] = set() for notifs in self.nixl_wrapper.get_new_notifs().values(): for notif in notifs: @@ -2109,7 +2149,7 @@ class NixlConnectorWorker: self._reqs_to_send[req_id] = expiration_time def _read_blocks_for_req(self, req_id: str, meta: ReqMeta): - assert meta.remote is not None + assert meta.remote is not None and self.kv_topo is not None remote_ranks = self.kv_topo.get_target_remote_ranks_from_engine_id( meta.remote.engine_id ) @@ -2182,6 +2222,7 @@ class NixlConnectorWorker: Post a READ point-to-point xfer request from a single local worker to a single remote worker. """ + assert self.kv_topo is not None block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(dst_engine_id) if block_size_ratio > 1: local_block_ids = self.get_mapped_blocks( @@ -2414,6 +2455,7 @@ class NixlConnectorWorker: For FlashInfer, this is half the length of the whole block, as K and V share the same region. """ + assert self.kv_topo is not None if self.kv_topo.is_kv_layout_blocks_first: # For indexing only half (either just the K or V part). block_len = self.block_len_per_layer[layer_idx] // 2 From 59a5cb387ae4c11c73855d505adb0b2c7cd3861d Mon Sep 17 00:00:00 2001 From: jiahanc <173873397+jiahanc@users.noreply.github.com> Date: Thu, 5 Feb 2026 18:23:11 +0800 Subject: [PATCH 095/810] [perf] Integrate flashinfer concat_mla_k (#31171) --- .../layers/attention/mla_attention.py | 20 ++++++-- vllm/utils/flashinfer.py | 47 +++++++++++++++++++ 2 files changed, 64 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index febad382162..862f8493985 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -1885,6 +1885,16 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]): self.indexer = indexer self.q_pad_num_heads = q_pad_num_heads + # Use flashinfer's optimized concat_mla_k kernel when available. + # The kernel is optimized for DeepSeek V3 dimensions: + # num_heads=128, nope_dim=128, rope_dim=64 + self._use_flashinfer_concat_mla_k = ( + has_flashinfer() + and (self.num_heads == 128) + and (self.qk_nope_head_dim == 128) + and (self.qk_rope_head_dim == 64) + ) + if use_trtllm_ragged_deepseek_prefill(): logger.info_once( "Using TRT-LLM ragged DeepSeek prefill for MLA", scope="local" @@ -2192,9 +2202,13 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]): dtype=k_nope.dtype, device=k_nope.device, ) - # Direct copies with efficient broadcasting - k[..., : k_nope.shape[-1]] = k_nope - k[..., k_nope.shape[-1] :] = k_pe + + if self._use_flashinfer_concat_mla_k: + torch.ops.vllm.flashinfer_concat_mla_k(k, k_nope, k_pe) + else: + # Fallback: Direct copies with efficient broadcasting + k[..., : k_nope.shape[-1]] = k_nope + k[..., k_nope.shape[-1] :] = k_pe return k def _compute_prefill_context( diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index f8cb1e14e1e..88e31718adf 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -396,6 +396,53 @@ def use_trtllm_attention( if has_flashinfer(): + from vllm.utils.torch_utils import direct_register_custom_op + + def _flashinfer_concat_mla_k( + k: torch.Tensor, + k_nope: torch.Tensor, + k_pe: torch.Tensor, + ) -> None: + """Custom op wrapper for flashinfer's concat_mla_k. + + This is an in-place operation that concatenates k_nope and k_pe into k. + + The kernel is optimized for DeepSeek V3 dimensions: + - num_heads=128 + - nope_dim=128 + - rope_dim=64 + + Key optimizations: + - Warp-based processing with software pipelining + - Vectorized memory access (int2 for nope, int for rope) + - L2 prefetching for next row while processing current + - Register reuse for rope values across all heads + + Args: + k: Output tensor, shape [num_tokens, num_heads, nope_dim + rope_dim]. + Modified in-place. + k_nope: The nope part of k, shape [num_tokens, num_heads, nope_dim]. + k_pe: The rope part of k (shared), shape [num_tokens, 1, rope_dim]. + This is broadcast to all heads. + """ + from flashinfer.concat_ops import concat_mla_k + + concat_mla_k(k, k_nope, k_pe) + + def _flashinfer_concat_mla_k_fake( + k: torch.Tensor, + k_nope: torch.Tensor, + k_pe: torch.Tensor, + ) -> None: + return + + # Register flashinfer concat_mla_k custom op + direct_register_custom_op( + op_name="flashinfer_concat_mla_k", + op_func=_flashinfer_concat_mla_k, + mutates_args=["k"], # k tensor is modified in-place + fake_impl=_flashinfer_concat_mla_k_fake, + ) @torch.library.custom_op( "vllm::flashinfer_mm_fp4", From a2522839d87d2b81b57458dfdbbcb27afb8191ae Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Thu, 5 Feb 2026 18:29:54 +0800 Subject: [PATCH 096/810] [Bugfix] Fix Kimi-K2.5 NVFP4 checkpoints weight loading (#33876) Signed-off-by: Isotr0py --- vllm/model_executor/models/deepseek_v2.py | 2 +- vllm/model_executor/models/kimi_k25.py | 18 ++++++++++++++---- 2 files changed, 15 insertions(+), 5 deletions(-) diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index f8907ed86ef..464518a3db9 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -1485,7 +1485,7 @@ class DeepseekV2ForCausalLM( param, "weight_loader", default_weight_loader ) weight_loader(param, loaded_weight) - if not is_fusion_moe_shared_experts_layer: + if name is not None and not is_fusion_moe_shared_experts_layer: loaded_params.add(name) return loaded_params diff --git a/vllm/model_executor/models/kimi_k25.py b/vllm/model_executor/models/kimi_k25.py index 191aed8e554..cb07cfe98ba 100644 --- a/vllm/model_executor/models/kimi_k25.py +++ b/vllm/model_executor/models/kimi_k25.py @@ -24,7 +24,11 @@ from transformers.processing_utils import ProcessorMixin from vllm.config import VllmConfig from vllm.config.multimodal import BaseDummyOptions from vllm.logger import init_logger -from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP +from vllm.model_executor.models.interfaces import ( + SupportsMultiModal, + SupportsPP, + SupportsQuant, +) from vllm.model_executor.models.kimi_k25_vit import ( KimiK25MultiModalProjector, MoonViT3dPretrainedModel, @@ -302,7 +306,9 @@ class KimiK25MultiModalProcessor(BaseMultiModalProcessor[KimiK25ProcessingInfo]) info=KimiK25ProcessingInfo, dummy_inputs=KimiK25DummyInputsBuilder, ) -class KimiK25ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP): +class KimiK25ForConditionalGeneration( + nn.Module, SupportsMultiModal, SupportsPP, SupportsQuant +): """Kimi-K2.5 model for conditional generation. Supports both image and video-chunk modalities. @@ -312,8 +318,12 @@ class KimiK25ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP) supports_encoder_tp_data = True - weights_mapper = WeightsMapper( + hf_to_vllm_mapper = WeightsMapper( orig_to_new_prefix={ + # For legacy NVFP4 checkpoint compatibility: + # see https://github.com/vllm-project/vllm/pull/33346#issuecomment-3851475033 + "language_model.layers.": "language_model.model.layers.", + # mm projector "mm_projector.proj.0": "mm_projector.linear_1", "mm_projector.proj.2": "mm_projector.linear_2", } @@ -465,4 +475,4 @@ class KimiK25ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): loader = AutoWeightsLoader(self) - return loader.load_weights(weights, mapper=self.weights_mapper) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) From 7bd42e609d24501f59a8b405229ed91f4ca8037c Mon Sep 17 00:00:00 2001 From: Cyrus Leung Date: Thu, 5 Feb 2026 18:43:42 +0800 Subject: [PATCH 097/810] [Refactor] Clean up input preprocessing (#33687) Signed-off-by: DarkLight1337 --- vllm/inputs/data.py | 4 +- vllm/inputs/parse.py | 16 +++ vllm/inputs/preprocess.py | 291 ++++++++++---------------------------- vllm/multimodal/inputs.py | 4 + 4 files changed, 101 insertions(+), 214 deletions(-) diff --git a/vllm/inputs/data.py b/vllm/inputs/data.py index 315ffddde59..d9f9814eef1 100644 --- a/vllm/inputs/data.py +++ b/vllm/inputs/data.py @@ -12,11 +12,13 @@ from vllm.sampling_params import SamplingParams if TYPE_CHECKING: from vllm.multimodal.inputs import ( MultiModalDataDict, + MultiModalEncDecInputs, MultiModalInputs, MultiModalUUIDDict, ) else: MultiModalDataDict = object + MultiModalEncDecInputs = object MultiModalInputs = object MultiModalUUIDDict = object @@ -241,7 +243,7 @@ class EncoderDecoderInputs(TypedDict): This specifies the required data for encoder-decoder models. """ - encoder: TokenInputs | MultiModalInputs + encoder: TokenInputs | MultiModalEncDecInputs """The inputs for the encoder portion.""" decoder: TokenInputs | MultiModalInputs diff --git a/vllm/inputs/parse.py b/vllm/inputs/parse.py index 5f832afdbf4..7cb1eb4b40f 100644 --- a/vllm/inputs/parse.py +++ b/vllm/inputs/parse.py @@ -69,6 +69,22 @@ def is_explicit_encoder_decoder_prompt( return isinstance(prompt, dict) and "encoder_prompt" in prompt +def split_enc_dec_prompt( + prompt: PromptType, +) -> tuple[SingletonPrompt, SingletonPrompt | None]: + if isinstance(prompt, str): + return prompt, None + + if "encoder_prompt" in prompt and "decoder_prompt" in prompt: + # NOTE: This passes pyright but not mypy + return ( + prompt["encoder_prompt"], # type: ignore[typeddict-item] + prompt["decoder_prompt"], # type: ignore[typeddict-item] + ) + + return prompt, None + + def split_enc_dec_inputs( inputs: ProcessorInputs, ) -> tuple[SingletonInputs | None, SingletonInputs]: diff --git a/vllm/inputs/preprocess.py b/vllm/inputs/preprocess.py index 6edb26a4a39..0a3b0c9468c 100644 --- a/vllm/inputs/preprocess.py +++ b/vllm/inputs/preprocess.py @@ -2,11 +2,12 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Mapping -from typing import Any, cast +from typing import Any from typing_extensions import assert_never from vllm.config import ModelConfig, ObservabilityConfig +from vllm.inputs.parse import split_enc_dec_prompt from vllm.logger import init_logger from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry from vllm.multimodal.cache import BaseMultiModalProcessorCache @@ -27,7 +28,6 @@ from .data import ( EmbedsInputs, EmbedsPrompt, EncoderDecoderInputs, - ExplicitEncoderDecoderPrompt, ProcessorInputs, PromptType, SingletonInputs, @@ -86,30 +86,15 @@ class InputPreprocessor: return self.tokenizer.eos_token_id - def get_decoder_start_token_id(self) -> int | None: + def get_decoder_start_token_id(self) -> int: """ Obtain the decoder start token id employed by an encoder/decoder - model. Returns None for non-encoder/decoder models or if the - model config is unavailable. + model. Raises an error if it is not available. """ - - if not self.model_config.is_encoder_decoder: - logger.warning_once( - "Using None for decoder start token id because " - "this is not an encoder/decoder model." - ) - return None - - if self.model_config is None or self.model_config.hf_config is None: - logger.warning_once( - "Using None for decoder start token id because " - "model config is not available." - ) - return None - dec_start_token_id = getattr( self.model_config.hf_config, "decoder_start_token_id", None ) + if dec_start_token_id is None: logger.warning_once( "Falling back on for decoder start token " @@ -118,48 +103,12 @@ class InputPreprocessor: ) dec_start_token_id = self.get_bos_token_id() + if dec_start_token_id is None: + raise RuntimeError("Cannot find decoder start token id or ") + return dec_start_token_id - def _get_default_enc_dec_decoder_prompt(self) -> list[int]: - """ - Specifically for encoder/decoder models: - generate a default decoder prompt for when - the user specifies only the encoder prompt. - - Encoder/decoder models utilize the decoder - prompt in different ways; as new models are - added, it is intended that this function - will be extended to produce differing - default decoder prompts, depending on the - model variety. - - Absent a special case, the default behavior - of this method is to mirror the behavior of - the HuggingFace (HF) GenerationMixin for a None - decoder prompt, which is to employ a logit processor - setting to force the first decoded token to be . - Here, this behavior is approximated by having the - "default" decoder prompt be . - - However, it is possible that in the future - other models may have different or more - complex logic for the default decoder prompt. - This motivates having a special helper method - for default decoder prompts. - - Returns: - - * prompt_token_ids - """ - - bos_token_id = self.get_bos_token_id() - assert bos_token_id is not None - return [bos_token_id] - - def _prepare_decoder_input_ids_for_generation( - self, - decoder_input_ids: list[int] | None, - ) -> list[int]: + def _prepare_decoder_input_ids(self, decoder_input_ids: list[int]) -> list[int]: """ Prepares `decoder_input_ids` for generation with encoder-decoder models. @@ -176,14 +125,7 @@ class InputPreprocessor: * Processed token list """ - decoder_start_token_id = self.get_decoder_start_token_id() - assert decoder_start_token_id is not None - - if decoder_input_ids is None: - # no decoder prompt input -> - # use decoder_start_token_id as decoder_input_ids - decoder_input_ids = self._get_default_enc_dec_decoder_prompt() if ( len(decoder_input_ids) == 0 @@ -428,111 +370,70 @@ class InputPreprocessor: assert_never(parsed) - def _build_enc_dec_llm_inputs( + def _validate_enc_inputs( + self, + inputs: SingletonInputs, + ) -> TokenInputs | MultiModalEncDecInputs: + if inputs["type"] == "embeds": + raise ValueError( + "Embedding inputs are not supported for encoder-decoder models" + ) + + if inputs["type"] == "multimodal" and "encoder_prompt_token_ids" not in inputs: + raise RuntimeError( + "You should register an encoder-decoder " + "multi-modal processor for encoder-decoder models." + ) + + return inputs # type: ignore[return-value] + + def _validate_dec_inputs( + self, + inputs: SingletonInputs, + ) -> TokenInputs | MultiModalInputs: + if inputs["type"] == "embeds": + raise ValueError( + "Embedding inputs are not supported for encoder-decoder models" + ) + + return inputs + + def _build_enc_dec_inputs( self, encoder_inputs: SingletonInputs, - decoder_inputs: SingletonInputs | None, + decoder_inputs: SingletonInputs | None = None, ) -> EncoderDecoderInputs: - if ( - encoder_inputs["type"] == "embeds" - or decoder_inputs - and decoder_inputs["type"] == "embeds" - ): - raise ValueError( - "Embedding inputs are not supported for encoder-decoder models" - ) - - # Needed for mypy - encoder_inputs = cast(TokenInputs | MultiModalInputs, encoder_inputs) - decoder_inputs = cast(TokenInputs | MultiModalInputs | None, decoder_inputs) - if decoder_inputs is None: - if self.model_config.hf_config.model_type == "whisper": - # For Whisper models, the text prompt should go to the decoder. - # If no explicit encoder/decoder inputs, then copy the prompt - # from the encoder to the decoder. The encoder tokens are later - # overridden by the audio features. - dec_token_ids = encoder_inputs["prompt_token_ids"].copy() - else: - dec_token_ids = self._prepare_decoder_input_ids_for_generation(None) - decoder_inputs = token_inputs(dec_token_ids) - else: - if "multi_modal_data" in decoder_inputs: - raise ValueError( - "Multi-modal decoder inputs of encoder-" - "decoder models are not supported yet" - ) + decoder_inputs = encoder_inputs - dec_token_ids = self._prepare_decoder_input_ids_for_generation( - decoder_inputs["prompt_token_ids"] - ) - decoder_inputs["prompt_token_ids"] = dec_token_ids + enc_inputs = self._validate_enc_inputs(encoder_inputs) + dec_inputs = self._validate_dec_inputs(decoder_inputs) - return EncoderDecoderInputs( - encoder=encoder_inputs, - decoder=decoder_inputs, - ) + enc_inputs_new: TokenInputs | MultiModalEncDecInputs + dec_inputs_new: TokenInputs | MultiModalInputs - def _split_enc_dec_mm_inputs( - self, - inputs: SingletonInputs | MultiModalEncDecInputs, - decoder_inputs_to_override: SingletonInputs | None = None, - ) -> tuple[SingletonInputs, SingletonInputs]: - """ - For encoder/decoder models only: - Separate Encoder/Decoder inputs from a MultiModalEncDecInputs - """ - if ( - inputs["type"] == "embeds" - or decoder_inputs_to_override - and decoder_inputs_to_override["type"] == "embeds" - ): - raise ValueError( - "Embedding inputs are not supported for encoder-decoder models" - ) - - # Needed for mypy - inputs = cast( - TokenInputs | MultiModalInputs | MultiModalEncDecInputs, - inputs, - ) - decoder_inputs_to_override = cast( - TokenInputs | MultiModalInputs | None, - decoder_inputs_to_override, - ) - - encoder_inputs: SingletonInputs - decoder_inputs: SingletonInputs - - if inputs["type"] == "multimodal": # Multimodal data inputs - if "encoder_prompt_token_ids" not in inputs: - raise RuntimeError( - "You should register an encoder-decoder " - "multi-modal processor for encoder-decoder " - "models." - ) - inputs = cast(MultiModalEncDecInputs, inputs) - - encoder_inputs = token_inputs(inputs["encoder_prompt_token_ids"]) - - decoder_prompt_inputs = decoder_inputs_to_override or inputs - decoder_inputs = MultiModalInputs( + if enc_inputs["type"] == "multimodal": + enc_inputs_new = token_inputs(enc_inputs["encoder_prompt_token_ids"]) + dec_inputs_new = MultiModalInputs( type="multimodal", - prompt_token_ids=decoder_prompt_inputs["prompt_token_ids"], - mm_kwargs=inputs["mm_kwargs"], - mm_hashes=inputs["mm_hashes"], - mm_placeholders=inputs["mm_placeholders"], + prompt_token_ids=dec_inputs["prompt_token_ids"], + mm_kwargs=enc_inputs["mm_kwargs"], + mm_hashes=enc_inputs["mm_hashes"], + mm_placeholders=enc_inputs["mm_placeholders"], ) - if cache_salt := inputs.get("cache_salt"): - decoder_inputs["cache_salt"] = cache_salt - - elif inputs["type"] == "token": # Text-only inputs - encoder_inputs = token_inputs(prompt_token_ids=[]) - decoder_inputs = decoder_inputs_to_override or inputs + elif enc_inputs["type"] == "token": + enc_inputs_new = token_inputs(prompt_token_ids=[]) + dec_inputs_new = dec_inputs else: - assert_never(inputs) # type: ignore[arg-type] + assert_never(enc_inputs) - return encoder_inputs, decoder_inputs + dec_inputs_new["prompt_token_ids"] = self._prepare_decoder_input_ids( + dec_inputs_new["prompt_token_ids"] + ) + if cache_salt := enc_inputs.get("cache_salt"): + dec_inputs_new["cache_salt"] = cache_salt + + return EncoderDecoderInputs(encoder=enc_inputs_new, decoder=dec_inputs_new) def _process_encoder_decoder_prompt( self, @@ -574,54 +475,23 @@ class InputPreprocessor: * [`EncoderDecoderInputs`][vllm.inputs.data.EncoderDecoderInputs] instance """ - encoder_inputs: SingletonInputs - decoder_inputs: SingletonInputs | None - if is_explicit_encoder_decoder_prompt(prompt): - # `cast` is needed for mypy, but not pyright - prompt_ = cast(ExplicitEncoderDecoderPrompt, prompt) - encoder_inputs = self._prompt_to_llm_inputs( - prompt_["encoder_prompt"], + encoder_prompt, decoder_prompt = split_enc_dec_prompt(prompt) + + return self._build_enc_dec_inputs( + encoder_inputs=self._prompt_to_llm_inputs( + encoder_prompt, tokenization_kwargs=tokenization_kwargs, mm_uuids=mm_uuids, - ) - if (decoder_input := prompt_["decoder_prompt"]) is None: - decoder_inputs = None - else: - decoder_inputs = self._prompt_to_llm_inputs( - decoder_input, tokenization_kwargs=tokenization_kwargs + ), + decoder_inputs=( + None + if decoder_prompt is None + else self._prompt_to_llm_inputs( + decoder_prompt, + tokenization_kwargs=tokenization_kwargs, ) - # For multimodal model, override decoder prompt from processor - # with explicit decoder prompt. - if self.model_config.is_multimodal_model: - encoder_inputs, decoder_inputs = self._split_enc_dec_mm_inputs( - encoder_inputs, decoder_inputs - ) - else: - # `cast` is needed for mypy, but not pyright - inputs = self._prompt_to_llm_inputs( - cast(SingletonPrompt, prompt), - tokenization_kwargs=tokenization_kwargs, - mm_uuids=mm_uuids, - ) - if self.model_config.is_multimodal_model: - # Encoder-Decoder Multimodal model - encoder_inputs, decoder_inputs = self._split_enc_dec_mm_inputs(inputs) - else: - encoder_inputs = inputs - decoder_inputs = None - - return self._build_enc_dec_llm_inputs(encoder_inputs, decoder_inputs) - - def _build_decoder_only_llm_inputs( - self, - prompt_inputs: DecoderOnlyInputs, - ) -> DecoderOnlyInputs: - if "prompt_token_ids" in prompt_inputs: - prompt_inputs = cast( - TokenInputs | MultiModalInputs, prompt_inputs - ) # Needed for mypy - - return prompt_inputs + ), + ) def _process_decoder_only_prompt( self, @@ -643,15 +513,12 @@ class InputPreprocessor: * [`DecoderOnlyInputs`][vllm.inputs.data.DecoderOnlyInputs] instance """ - - prompt_comps = self._prompt_to_llm_inputs( + return self._prompt_to_llm_inputs( prompt, tokenization_kwargs=tokenization_kwargs, mm_uuids=mm_uuids, ) - return self._build_decoder_only_llm_inputs(prompt_comps) - def _preprocess( self, prompt: PromptType, @@ -673,10 +540,8 @@ class InputPreprocessor: "Cannot pass encoder-decoder prompt to decoder-only models" ) - # Decoder-only operation - # `cast` is needed for mypy, but not pyright return self._process_decoder_only_prompt( - cast(SingletonPrompt, prompt), + prompt, tokenization_kwargs=tokenization_kwargs, mm_uuids=mm_uuids, ) diff --git a/vllm/multimodal/inputs.py b/vllm/multimodal/inputs.py index e50771e9909..262def71220 100644 --- a/vllm/multimodal/inputs.py +++ b/vllm/multimodal/inputs.py @@ -1083,6 +1083,10 @@ class MultiModalEncDecInputs(MultiModalInputs): Represents the outputs of [`EncDecMultiModalProcessor`][vllm.multimodal.processing.EncDecMultiModalProcessor] ready to be passed to vLLM internals. + + Note: Even text-only encoder-decoder models are currently implemented + as multi-modal models for convenience. + (Example: https://github.com/neuralmagic/bart-plugin) """ encoder_prompt_token_ids: list[int] From 1c3a221d3b0f7a82cd9a6d56e10ea360e2435a1c Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Thu, 5 Feb 2026 18:51:22 +0800 Subject: [PATCH 098/810] [Bugfix] Fix corner case of sparse embedding (#33886) Signed-off-by: wang.yuqi --- tests/models/language/pooling/test_bge_m3.py | 10 ++++++++++ vllm/model_executor/layers/pooler/special.py | 2 +- 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/tests/models/language/pooling/test_bge_m3.py b/tests/models/language/pooling/test_bge_m3.py index 5ad1fee037d..2c0c0de346f 100644 --- a/tests/models/language/pooling/test_bge_m3.py +++ b/tests/models/language/pooling/test_bge_m3.py @@ -136,6 +136,16 @@ async def test_bge_m3_api_server_sparse_embedding(client: openai.AsyncOpenAI): ) +@pytest.mark.asyncio +async def test_bge_m3_api_server_sparse_embedding_corner_case( + client: openai.AsyncOpenAI, +): + embeddings = await sparse_embeddings(client, ["Hi"]) + assert len(embeddings) == 1 + assert 2673 in embeddings[0] + assert embeddings[0][2673] == pytest.approx(0.26710861921310425, rel=0.01) + + # https://github.com/FlagOpen/FlagEmbedding/blob/6fd176266f2382878bcc69cd656cff425d52f49b/FlagEmbedding/inference/embedder/encoder_only/m3.py#L163 def colbert_score(q_reps: torch.Tensor, p_reps: torch.Tensor) -> torch.Tensor: token_scores = torch.einsum("in,jn->ij", q_reps, p_reps) diff --git a/vllm/model_executor/layers/pooler/special.py b/vllm/model_executor/layers/pooler/special.py index 707e7c90760..bafa191dbac 100644 --- a/vllm/model_executor/layers/pooler/special.py +++ b/vllm/model_executor/layers/pooler/special.py @@ -165,7 +165,7 @@ class BOSEOSFilter(Pooler): pooled_data = pooled_data[1:] if token_ids[-1] == self.eos_token_id: pooled_data = pooled_data[:-1] - pooled_outputs[i] = pooled_data.squeeze() + pooled_outputs[i] = pooled_data.squeeze(-1) return pooled_outputs From 81a90e52776503c6cbdccd30fbe53f61c9179bdf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Thu, 5 Feb 2026 13:20:25 +0100 Subject: [PATCH 099/810] [Docs] Add bart-plugin to docs (#33905) Signed-off-by: NickLucche --- docs/design/plugin_system.md | 2 +- docs/models/supported_models.md | 10 ++++++++++ docs/usage/v1_guide.md | 9 ++++++--- 3 files changed, 17 insertions(+), 4 deletions(-) diff --git a/docs/design/plugin_system.md b/docs/design/plugin_system.md index 1bf8aea3c66..f8fb099245c 100644 --- a/docs/design/plugin_system.md +++ b/docs/design/plugin_system.md @@ -45,7 +45,7 @@ Every plugin has three parts: ## Types of supported plugins -- **General plugins** (with group name `vllm.general_plugins`): The primary use case for these plugins is to register custom, out-of-the-tree models into vLLM. This is done by calling `ModelRegistry.register_model` to register the model inside the plugin function. +- **General plugins** (with group name `vllm.general_plugins`): The primary use case for these plugins is to register custom, out-of-the-tree models into vLLM. This is done by calling `ModelRegistry.register_model` to register the model inside the plugin function. For an example of an official model plugin, see the [bart-plugin](https://github.com/vllm-project/bart-plugin) which adds support for `BartForConditionalGeneration`. - **Platform plugins** (with group name `vllm.platform_plugins`): The primary use case for these plugins is to register custom, out-of-the-tree platforms into vLLM. The plugin function should return `None` when the platform is not supported in the current environment, or the platform class's fully qualified name when the platform is supported. diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index e07e17ec50d..e69f68feedc 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -174,6 +174,16 @@ class MyConfig(PretrainedConfig): - The `list` in the first element of the `tuple` contains the names of the input arguments - The `list` in the last element of the `tuple` contains the names of the variables the layer outputs to in your modeling code +### Plugins + +Some model architectures are supported via vLLM plugins. These plugins extend vLLM's capabilities through the [plugin system](../design/plugin_system.md). + +| Architecture | Models | Plugin Repository | +|--------------|--------|-------------------| +| `BartForConditionalGeneration` | BART | [bart-plugin](https://github.com/vllm-project/bart-plugin) | + +For other model architectures not natively supported, in particular for Encoder-Decoder models, we recommend following a similar pattern by implementing support through the plugin system. + ## Loading a Model ### Hugging Face Hub diff --git a/docs/usage/v1_guide.md b/docs/usage/v1_guide.md index 8506e01b96d..96850871d0b 100644 --- a/docs/usage/v1_guide.md +++ b/docs/usage/v1_guide.md @@ -134,9 +134,12 @@ Please note that prefix caching is not yet supported for any of the above models #### Encoder-Decoder Models -Whisper is supported. Other models requiring cross-attention between separate -encoder and decoder (e.g., `BartForConditionalGeneration`, -`MllamaForConditionalGeneration`) are no longer supported. +Whisper is supported natively. Other encoder-decoder models are supported via the plugin system: + +- **BART**: `BartForConditionalGeneration` is supported via the official [bart-plugin](https://github.com/vllm-project/bart-plugin). + +For other encoder-decoder models (e.g., `MllamaForConditionalGeneration`), we recommend +following a similar pattern by implementing support through the [plugin system](../design/plugin_system.md). ### Features From 82914d2ae8d0362be06700222f4cd4c5f6b0dc36 Mon Sep 17 00:00:00 2001 From: Mario Hong <86880754+mariohong128@users.noreply.github.com> Date: Fri, 6 Feb 2026 00:04:04 +0800 Subject: [PATCH 100/810] [Bugfix] Fix step3p5 parser when using mtp (#33690) Signed-off-by: mariohong --- .../tool_parsers/test_step3p5_tool_parser.py | 1435 +++++++++++++++++ vllm/tool_parsers/step3p5_tool_parser.py | 25 +- 2 files changed, 1455 insertions(+), 5 deletions(-) create mode 100644 tests/tool_parsers/test_step3p5_tool_parser.py diff --git a/tests/tool_parsers/test_step3p5_tool_parser.py b/tests/tool_parsers/test_step3p5_tool_parser.py new file mode 100644 index 00000000000..6da1e08550a --- /dev/null +++ b/tests/tool_parsers/test_step3p5_tool_parser.py @@ -0,0 +1,1435 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json +from collections.abc import Generator + +import pytest + +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, + ChatCompletionToolsParam, +) +from vllm.entrypoints.openai.engine.protocol import ( + DeltaMessage, + FunctionCall, + ToolCall, +) +from vllm.tokenizers import TokenizerLike, get_tokenizer +from vllm.tokenizers.detokenizer_utils import detokenize_incrementally +from vllm.tool_parsers.step3p5_tool_parser import Step3p5ToolParser + +MODEL = "stepfun-ai/Step-3.5-Flash" + + +@pytest.fixture(scope="module") +def step3p5_tokenizer(): + return get_tokenizer(tokenizer_name=MODEL) + + +@pytest.fixture +def step3p5_tool_parser(step3p5_tokenizer): + return Step3p5ToolParser(step3p5_tokenizer) + + +@pytest.fixture +def sample_tools(): + return [ + ChatCompletionToolsParam( + type="function", + function={ + "name": "get_current_weather", + "description": "Get the current weather", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "The city name"}, + "state": {"type": "string", "description": "The state code"}, + "unit": {"type": "string", "enum": ["fahrenheit", "celsius"]}, + }, + "required": ["city", "state"], + }, + }, + ), + ChatCompletionToolsParam( + type="function", + function={ + "name": "calculate_area", + "description": "Calculate area of a shape", + "parameters": { + "type": "object", + "properties": { + "shape": {"type": "string"}, + "dimensions": {"type": "object"}, + "precision": {"type": "integer"}, + }, + }, + }, + ), + ] + + +def assert_tool_calls( + actual_tool_calls: list[ToolCall], expected_tool_calls: list[ToolCall] +): + assert len(actual_tool_calls) == len(expected_tool_calls) + + for actual_tool_call, expected_tool_call in zip( + actual_tool_calls, expected_tool_calls + ): + assert actual_tool_call.type == "function" + assert actual_tool_call.function.name == expected_tool_call.function.name + assert json.loads(actual_tool_call.function.arguments) == json.loads( + expected_tool_call.function.arguments + ) + + +def stream_delta_message_generator( + step3p5_tool_parser, + step3p5_tokenizer: TokenizerLike, + model_output: str, + request: ChatCompletionRequest | None = None, +) -> Generator[DeltaMessage, None, None]: + all_token_ids = step3p5_tokenizer.encode(model_output, add_special_tokens=False) + + previous_text = "" + previous_tokens = None + prefix_offset = 0 + read_offset = 0 + for i, delta_token in enumerate(all_token_ids): + delta_token_ids = [delta_token] + previous_token_ids = all_token_ids[:i] + current_token_ids = all_token_ids[: i + 1] + + (new_tokens, delta_text, new_prefix_offset, new_read_offset) = ( + detokenize_incrementally( + tokenizer=step3p5_tokenizer, + all_input_ids=current_token_ids, + prev_tokens=previous_tokens, + prefix_offset=prefix_offset, + read_offset=read_offset, + skip_special_tokens=False, + spaces_between_special_tokens=True, + ) + ) + + current_text = previous_text + delta_text + + delta_message = step3p5_tool_parser.extract_tool_calls_streaming( + previous_text, + current_text, + delta_text, + previous_token_ids, + current_token_ids, + delta_token_ids, + request=request, + ) + if delta_message: + yield delta_message + + previous_text = current_text + previous_tokens = ( + previous_tokens + new_tokens if previous_tokens else new_tokens + ) + prefix_offset = new_prefix_offset + read_offset = new_read_offset + + +def stream_delta_message_generator_from_chunks( + step3p5_tool_parser, + step3p5_tokenizer: TokenizerLike, + delta_text_chunks: list[str], + request: ChatCompletionRequest | None = None, +) -> Generator[DeltaMessage, None, None]: + previous_text = "" + previous_token_ids: list[int] = [] + + for delta_text in delta_text_chunks: + delta_token_ids = step3p5_tokenizer.encode(delta_text, add_special_tokens=False) + current_text = previous_text + delta_text + current_token_ids = previous_token_ids + delta_token_ids + + delta_message = step3p5_tool_parser.extract_tool_calls_streaming( + previous_text, + current_text, + delta_text, + previous_token_ids, + current_token_ids, + delta_token_ids, + request=request, + ) + if delta_message: + yield delta_message + + previous_text = current_text + previous_token_ids = current_token_ids + + +def test_extract_tool_calls_no_tools(step3p5_tool_parser): + model_output = "This is a test response without any tool calls" + extracted_tool_calls = step3p5_tool_parser.extract_tool_calls( + model_output, request=None + ) # type: ignore[arg-type] + assert not extracted_tool_calls.tools_called + assert extracted_tool_calls.tool_calls == [] + assert extracted_tool_calls.content == model_output + + +@pytest.mark.parametrize( + ids=[ + "single_tool", + "single_tool_with_content", + "single_tool_multiline_param", + "parallel_tools", + "tool_with_typed_params", + ], + argnames=["model_output", "expected_tool_calls", "expected_content"], + argvalues=[ + ( + """ + + +Dallas + + +TX + + +fahrenheit + + +""", + [ + ToolCall( + function=FunctionCall( + name="get_current_weather", + arguments=json.dumps( + {"city": "Dallas", "state": "TX", "unit": "fahrenheit"} + ), + ) + ) + ], + None, + ), + ( + """Sure! Let me check the weather for you. + + +Dallas + + +TX + + +fahrenheit + + +""", + [ + ToolCall( + function=FunctionCall( + name="get_current_weather", + arguments=json.dumps( + {"city": "Dallas", "state": "TX", "unit": "fahrenheit"} + ), + ) + ) + ], + "Sure! Let me check the weather for you.", + ), + ( + """ + + +rectangle + + +{"width": 10, + "height": 20} + + +2 + + +""", + [ + ToolCall( + function=FunctionCall( + name="calculate_area", + arguments=json.dumps( + { + "shape": "rectangle", + "dimensions": {"width": 10, "height": 20}, + "precision": 2, + } + ), + ) + ) + ], + None, + ), + ( + """ + + +Dallas + + +TX + + +fahrenheit + + + + + + +Orlando + + +FL + + +fahrenheit + + +""", + [ + ToolCall( + function=FunctionCall( + name="get_current_weather", + arguments=json.dumps( + {"city": "Dallas", "state": "TX", "unit": "fahrenheit"} + ), + ) + ), + ToolCall( + function=FunctionCall( + name="get_current_weather", + arguments=json.dumps( + {"city": "Orlando", "state": "FL", "unit": "fahrenheit"} + ), + ) + ), + ], + None, + ), + ( + """Let me calculate that area for you. + + +circle + + +{"radius": 15.5} + + +3 + + +""", + [ + ToolCall( + function=FunctionCall( + name="calculate_area", + arguments=json.dumps( + { + "shape": "circle", + "dimensions": {"radius": 15.5}, + "precision": 3, + } + ), + ) + ) + ], + "Let me calculate that area for you.", + ), + ], +) +def test_extract_tool_calls( + step3p5_tool_parser, + sample_tools, + model_output, + expected_tool_calls, + expected_content, +): + request = ChatCompletionRequest(model=MODEL, messages=[], tools=sample_tools) + extracted_tool_calls = step3p5_tool_parser.extract_tool_calls( + model_output, request=request + ) + assert extracted_tool_calls.tools_called + + assert_tool_calls(extracted_tool_calls.tool_calls, expected_tool_calls) + + assert extracted_tool_calls.content == expected_content + + +def test_extract_tool_calls_fallback_no_tags(step3p5_tool_parser, sample_tools): + """Test fallback parsing when XML tags are missing""" + model_output = """ + +Dallas + + +TX + +""" + + request = ChatCompletionRequest(model=MODEL, messages=[], tools=sample_tools) + extracted_tool_calls = step3p5_tool_parser.extract_tool_calls( + model_output, request=request + ) + + assert extracted_tool_calls.tools_called + assert len(extracted_tool_calls.tool_calls) == 1 + assert extracted_tool_calls.tool_calls[0].function.name == "get_current_weather" + + +def test_extract_tool_calls_type_conversion(step3p5_tool_parser): + """Test parameter type conversion based on tool schema""" + tools = [ + ChatCompletionToolsParam( + type="function", + function={ + "name": "test_types", + "parameters": { + "type": "object", + "properties": { + "int_param": {"type": "integer"}, + "float_param": {"type": "float"}, + "bool_param": {"type": "boolean"}, + "str_param": {"type": "string"}, + "obj_param": {"type": "object"}, + }, + }, + }, + ) + ] + + model_output = """ + + +42 + + +3.14 + + +true + + +hello world + + +{"key": "value"} + + +""" + + request = ChatCompletionRequest(model=MODEL, messages=[], tools=tools) + extracted_tool_calls = step3p5_tool_parser.extract_tool_calls( + model_output, request=request + ) + + args = json.loads(extracted_tool_calls.tool_calls[0].function.arguments) + assert args["int_param"] == 42 + assert args["float_param"] == 3.14 + assert args["bool_param"] is True + assert args["str_param"] == "hello world" + assert args["obj_param"] == {"key": "value"} + + +@pytest.mark.parametrize( + ids=[ + "no_tools", + "single_tool", + "single_tool_with_content", + "single_tool_multiline_param", + "parallel_tools", + "tool_with_typed_params", # Added this test case + ], + argnames=["model_output", "expected_tool_calls", "expected_content"], + argvalues=[ + ("This is a test without tools", [], "This is a test without tools"), + ( + """ + + +Dallas + + +TX + + +fahrenheit + + +""", + [ + ToolCall( + function=FunctionCall( + name="get_current_weather", + arguments=json.dumps( + {"city": "Dallas", "state": "TX", "unit": "fahrenheit"} + ), + ) + ) + ], + None, + ), + ( + """Sure! Let me check the weather for you. + + +Dallas + + +TX + + +fahrenheit + + +""", + [ + ToolCall( + function=FunctionCall( + name="get_current_weather", + arguments=json.dumps( + {"city": "Dallas", "state": "TX", "unit": "fahrenheit"} + ), + ) + ) + ], + "Sure! Let me check the weather for you.", + ), + ( + """ + + +rectangle + + +{"width": 10, + "height": 20} + + +2 + + +""", + [ + ToolCall( + function=FunctionCall( + name="calculate_area", + arguments=json.dumps( + { + "shape": "rectangle", + "dimensions": {"width": 10, "height": 20}, + "precision": 2, + } + ), + ) + ) + ], + None, + ), + ( + """ + + +Dallas + + +TX + + +fahrenheit + + + + + + +Orlando + + +FL + + +celsius + + +""", + [ + ToolCall( + function=FunctionCall( + name="get_current_weather", + arguments=json.dumps( + {"city": "Dallas", "state": "TX", "unit": "fahrenheit"} + ), + ) + ), + ToolCall( + function=FunctionCall( + name="get_current_weather", + arguments=json.dumps( + {"city": "Orlando", "state": "FL", "unit": "celsius"} + ), + ) + ), + ], + None, + ), + # Added tool_with_typed_params test case + ( + """Let me calculate that area for you. + + +circle + + +{"radius": 15.5} + + +3 + + +""", + [ + ToolCall( + function=FunctionCall( + name="calculate_area", + arguments=json.dumps( + { + "shape": "circle", + "dimensions": {"radius": 15.5}, + "precision": 3, + } + ), + ) + ) + ], + "Let me calculate that area for you.", + ), + ], +) +def test_extract_tool_calls_streaming( + step3p5_tool_parser, + step3p5_tokenizer, + sample_tools, + model_output, + expected_tool_calls, + expected_content, +): + """Test incremental streaming behavior including typed parameters""" + request = ChatCompletionRequest(model=MODEL, messages=[], tools=sample_tools) + + other_content = "" + tool_states = {} # Track state per tool index + + for delta_message in stream_delta_message_generator( + step3p5_tool_parser, step3p5_tokenizer, model_output, request + ): + # role should never be streamed from tool parser + assert not delta_message.role + + if delta_message.content: + other_content += delta_message.content + + if delta_message.tool_calls: + for tool_call in delta_message.tool_calls: + idx = tool_call.index + + # Initialize state for new tool + if idx not in tool_states: + tool_states[idx] = { + "id": None, + "name": None, + "arguments": "", + "type": None, + } + + # First chunk should have id, name, and type + if tool_call.id: + tool_states[idx]["id"] = tool_call.id + + if tool_call.type: + assert tool_call.type == "function" + tool_states[idx]["type"] = tool_call.type + + if tool_call.function: + if tool_call.function.name: + # Should only be set once + assert tool_states[idx]["name"] is None + tool_states[idx]["name"] = tool_call.function.name + + if tool_call.function.arguments is not None: + # Accumulate arguments incrementally + tool_states[idx]["arguments"] += tool_call.function.arguments + + # Verify final content + assert other_content == (expected_content or "") # Handle None case + + # Verify we got all expected tool calls + assert len(tool_states) == len(expected_tool_calls) + + # Verify each tool call + for idx, expected_tool in enumerate(expected_tool_calls): + state = tool_states[idx] + assert state["id"] is not None + assert state["type"] == "function" + assert state["name"] == expected_tool.function.name + + # Parse accumulated arguments + arguments_str = state["arguments"] + assert arguments_str is not None + actual_args = json.loads(arguments_str) + expected_args = json.loads(expected_tool.function.arguments) + assert actual_args == expected_args + + +def test_extract_tool_calls_missing_closing_parameter_tag( + step3p5_tool_parser, sample_tools +): + """Test handling of missing closing tag""" + # Using get_current_weather from sample_tools but with malformed XML + model_output = """Let me check the weather for you: + + + +Dallas + +TX + + +fahrenheit + + +""" + + request = ChatCompletionRequest(model=MODEL, messages=[], tools=sample_tools) + extracted_tool_calls = step3p5_tool_parser.extract_tool_calls( + model_output, request=request + ) + + # The parser should handle the malformed XML gracefully + assert extracted_tool_calls.tools_called + assert len(extracted_tool_calls.tool_calls) == 1 + + # Verify the function name is correct + assert extracted_tool_calls.tool_calls[0].function.name == "get_current_weather" + + # Verify the arguments are parsed despite the missing closing tag + args = json.loads(extracted_tool_calls.tool_calls[0].function.arguments) + assert "city" in args + assert args["city"] == "Dallas" + assert args["state"] == "TX" + assert args["unit"] == "fahrenheit" + + # Check that content before the tool call is preserved + assert "Let me check the weather for you:" in extracted_tool_calls.content + + +def test_extract_tool_calls_streaming_missing_closing_tag( + step3p5_tool_parser, step3p5_tokenizer, sample_tools +): + """Test streaming with missing closing tag""" + # Using get_current_weather from sample_tools but with malformed XML + model_output = """Let me check the weather for you: + + + +Dallas + +TX + + +fahrenheit + + +""" + + request = ChatCompletionRequest(model=MODEL, messages=[], tools=sample_tools) + + other_content = "" + tool_states = {} + + for delta_message in stream_delta_message_generator( + step3p5_tool_parser, step3p5_tokenizer, model_output, request + ): + if delta_message.content: + other_content += delta_message.content + + if delta_message.tool_calls: + for tool_call in delta_message.tool_calls: + idx = tool_call.index + + if idx not in tool_states: + tool_states[idx] = { + "id": None, + "name": None, + "arguments": "", + "type": None, + } + + if tool_call.id: + tool_states[idx]["id"] = tool_call.id + + if tool_call.type: + assert tool_call.type == "function" + tool_states[idx]["type"] = tool_call.type + + if tool_call.function: + if tool_call.function.name: + tool_states[idx]["name"] = tool_call.function.name + + if tool_call.function.arguments is not None: + tool_states[idx]["arguments"] += tool_call.function.arguments + + # Verify content was streamed + assert "Let me check the weather for you:" in other_content + + # Verify we got the tool call + assert len(tool_states) == 1 + state = tool_states[0] + assert state["id"] is not None + assert state["type"] == "function" + assert state["name"] == "get_current_weather" + + # Verify arguments were parsed correctly despite missing closing tag + assert state["arguments"] is not None + args = json.loads(state["arguments"]) + assert args["city"] == "Dallas" + assert args["state"] == "TX" + assert args["unit"] == "fahrenheit" + + +def test_extract_tool_calls_streaming_incremental( + step3p5_tool_parser, step3p5_tokenizer, sample_tools +): + """Test that streaming is truly incremental""" + model_output = """I'll check the weather. + + +Dallas + + +TX + + +""" + + request = ChatCompletionRequest(model=MODEL, messages=[], tools=sample_tools) + + chunks = [] + for delta_message in stream_delta_message_generator( + step3p5_tool_parser, step3p5_tokenizer, model_output, request + ): + chunks.append(delta_message) + + # Should have multiple chunks + assert len(chunks) > 3 + + # First chunk(s) should be content + assert chunks[0].content is not None + assert chunks[0].tool_calls is None or chunks[0].tool_calls == [] + + # Should have a chunk with tool header (id, name, type) + header_found = False + for chunk in chunks: + if chunk.tool_calls and chunk.tool_calls[0].id: + header_found = True + assert chunk.tool_calls[0].function.name == "get_current_weather" + assert chunk.tool_calls[0].type == "function" + # Empty initially + assert chunk.tool_calls[0].function.arguments == "" + break + assert header_found + + # Should have chunks with incremental arguments + arg_chunks = [] + for chunk in chunks: + if chunk.tool_calls and chunk.tool_calls[0].function.arguments: + arg_chunks.append(chunk.tool_calls[0].function.arguments) + + # Arguments should be streamed incrementally + assert len(arg_chunks) > 1 + + # Concatenated arguments should form valid JSON + full_args = "".join(arg_chunks) + parsed_args = json.loads(full_args) + assert parsed_args["city"] == "Dallas" + assert parsed_args["state"] == "TX" + + +def test_extract_tool_calls_complex_type_with_single_quote(step3p5_tool_parser): + """Test parameter type conversion based on tool schema""" + tools = [ + ChatCompletionToolsParam( + type="function", + function={ + "name": "test_types", + "parameters": { + "type": "object", + "properties": { + "int_param": {"type": "integer"}, + "float_param": {"type": "float"}, + "bool_param": {"type": "boolean"}, + "str_param": {"type": "string"}, + "obj_param": {"type": "object"}, + }, + }, + }, + ) + ] + + model_output = """ + + +{'key': 'value'} + + +""" + + request = ChatCompletionRequest(model=MODEL, messages=[], tools=tools) + extracted_tool_calls = step3p5_tool_parser.extract_tool_calls( + model_output, request=request + ) + + args = json.loads(extracted_tool_calls.tool_calls[0].function.arguments) + assert args["obj_param"] == {"key": "value"} + + +def test_extract_tool_calls_streaming_mixed_content_and_multiple_tool_calls( + step3p5_tool_parser, step3p5_tokenizer, sample_tools +): + """Test mixed content with multiple complete tool calls. + + Scenario: Model outputs "hello" + complete tool call + "hi" + complete tool call. + Expected: "hello" as content, first tool call parsed (index=0), "hi" as content, + second tool call parsed (index=1). + """ + # Model output: hello + complete tool call + hi + complete tool call + model_output = """hello + + +Dallas + + +TX + + +hi + + +rectangle + + +{"width": 10, "height": 5} + + +""" + + request = ChatCompletionRequest(model=MODEL, messages=[], tools=sample_tools) + + other_content = "" + tool_states = {} + + for delta_message in stream_delta_message_generator( + step3p5_tool_parser, step3p5_tokenizer, model_output, request + ): + if delta_message.content: + other_content += delta_message.content + + if delta_message.tool_calls: + for tool_call in delta_message.tool_calls: + idx = tool_call.index + + if idx not in tool_states: + tool_states[idx] = { + "id": None, + "name": None, + "arguments": "", + "type": None, + } + + if tool_call.id: + tool_states[idx]["id"] = tool_call.id + + if tool_call.type: + assert tool_call.type == "function" + tool_states[idx]["type"] = tool_call.type + + if tool_call.function: + if tool_call.function.name: + tool_states[idx]["name"] = tool_call.function.name + + if tool_call.function.arguments is not None: + tool_states[idx]["arguments"] += tool_call.function.arguments + + # Should have exactly two complete tool calls + assert len(tool_states) == 2, "Should have exactly two complete tool calls" + + # Verify the first tool call (index=0) + assert tool_states[0]["name"] == "get_current_weather" + assert tool_states[0]["arguments"] + args_dict_0 = json.loads(tool_states[0]["arguments"]) + assert args_dict_0["city"] == "Dallas" + assert args_dict_0["state"] == "TX" + + # Verify the second tool call (index=1) + assert tool_states[1]["name"] == "calculate_area" + assert tool_states[1]["arguments"] + args_dict_1 = json.loads(tool_states[1]["arguments"]) + assert args_dict_1["shape"] == "rectangle" + assert isinstance(args_dict_1["dimensions"], dict), "dimensions should be a dict" + assert args_dict_1["dimensions"]["width"] == 10 + assert args_dict_1["dimensions"]["height"] == 5 + # Verify content: should contain "hello", "hi" + assert "hello" in other_content, "Should contain 'hello' as content" + assert "hi" in other_content, "Should contain 'hi' as content" + + # Verify the order: hello should come first, then hi + hello_index = other_content.find("hello") + hi_index = other_content.find("hi") + + assert hello_index >= 0, "'hello' should be in content" + assert hi_index > hello_index, "'hi' should come after 'hello'" + + # Verify that tool call tags are NOT in the content + # We should not see complete tool call structures in content + assert "" not in other_content, ( + "First tool call should not be in content" + ) + assert "" not in other_content, ( + "Second tool call should not be in content" + ) + + +def test_extract_tool_calls_non_streaming_mixed_content_and_multiple_tool_calls( + step3p5_tool_parser, sample_tools +): + """Test non-streaming extraction with mixed content and multiple tool calls. + + Scenario: Model outputs "hello" + complete tool call + "hi" + complete tool call. + Expected: "hello" as content, first tool call parsed (index=0), "hi" as content, + second tool call parsed (index=1) + """ + # Model output: hello + complete tool call + hi + complete tool call + model_output = """hello + + +Dallas + + +TX + + +hi + + +rectangle + + +{"width": 10, "height": 5} + + +""" + + request = ChatCompletionRequest(model=MODEL, messages=[], tools=sample_tools) + + extracted_tool_calls = step3p5_tool_parser.extract_tool_calls( + model_output, request=request + ) + + # Should have exactly two complete tool calls + assert extracted_tool_calls.tools_called + assert len(extracted_tool_calls.tool_calls) == 2, ( + "Should have exactly two complete tool calls" + ) + + # Verify the first tool call (index=0) + assert extracted_tool_calls.tool_calls[0].function.name == "get_current_weather" + args_dict_0 = json.loads(extracted_tool_calls.tool_calls[0].function.arguments) + assert args_dict_0["city"] == "Dallas" + assert args_dict_0["state"] == "TX" + + # Verify the second tool call (index=1) + assert extracted_tool_calls.tool_calls[1].function.name == "calculate_area" + args_dict_1 = json.loads(extracted_tool_calls.tool_calls[1].function.arguments) + assert args_dict_1["shape"] == "rectangle" + assert isinstance(args_dict_1["dimensions"], dict), "dimensions should be a dict" + assert args_dict_1["dimensions"]["width"] == 10 + assert args_dict_1["dimensions"]["height"] == 5 + + # Verify content: should contain "hello", "hi" + assert extracted_tool_calls.content is not None + assert "hello" in extracted_tool_calls.content, "Should contain 'hello' as content" + assert "hi" in extracted_tool_calls.content, "Should contain 'hi' as content" + + # Verify the order: hello should come first, then hi + hello_index = extracted_tool_calls.content.find("hello") + hi_index = extracted_tool_calls.content.find("hi") + + assert hello_index >= 0, "'hello' should be in content" + assert hi_index > hello_index, "'hi' should come after 'hello'" + + # Verify that tool call tags are NOT in the content + assert "" not in extracted_tool_calls.content, ( + "First tool call should not be in content" + ) + assert "" not in extracted_tool_calls.content, ( + "Second tool call should not be in content" + ) + + +def test_extract_tool_calls_streaming_full_input_mixed_content_and_multiple_tool_calls( + step3p5_tool_parser, step3p5_tokenizer, sample_tools +): + """Test streaming with entire input as single delta_text. + + Scenario: Model outputs "hello" + complete tool call + "hi" + complete tool call. + This test simulates the case where the entire input is sent as a single delta_text. + Expected: "hello" as content, first tool call parsed (index=0), "hi" as content, + second tool call parsed (index=1). + """ + # Model output: hello + complete tool call + hi + complete tool call + model_output = """hello + + +Dallas + + +TX + + +hi + + +rectangle + + +{"width": 10, "height": 5} + + +""" + + request = ChatCompletionRequest(model=MODEL, messages=[], tools=sample_tools) + + other_content = "" + tool_states = {} + + # Encode all content tokens at once + all_token_ids = step3p5_tokenizer.encode(model_output, add_special_tokens=False) + eos_token_id = getattr(step3p5_tokenizer, "eos_token_id", None) + + # Include EOS token in delta_token_ids if available + if eos_token_id is not None: + delta_token_ids = all_token_ids + [eos_token_id] + else: + delta_token_ids = all_token_ids + + # current_token_ids includes all content tokens (EOS is not part of the text) + current_token_ids = all_token_ids + previous_token_ids: list[int] = [] + + # Decode all tokens to get the full text + current_text = step3p5_tokenizer.decode( + current_token_ids, skip_special_tokens=False + ) + previous_text = "" + delta_text = current_text + + # Call parser once with all tokens including EOS + delta_result = step3p5_tool_parser.extract_tool_calls_streaming( + previous_text, + current_text, + delta_text, + previous_token_ids, + current_token_ids, + delta_token_ids, + request=request, + ) + + # Process delta result + if delta_result: + if delta_result.content: + other_content += delta_result.content + if delta_result.tool_calls: + for tool_call in delta_result.tool_calls: + idx = tool_call.index + if idx not in tool_states: + tool_states[idx] = { + "id": None, + "name": None, + "arguments": "", + "type": None, + } + if tool_call.id: + tool_states[idx]["id"] = tool_call.id + if tool_call.type: + tool_states[idx]["type"] = tool_call.type + if tool_call.function: + if tool_call.function.name: + tool_states[idx]["name"] = tool_call.function.name + if tool_call.function.arguments is not None: + tool_states[idx]["arguments"] += tool_call.function.arguments + + # Should have exactly two complete tool calls + assert len(tool_states) == 2, "Should have exactly two complete tool calls" + + # Verify the first tool call (index=0) + assert tool_states[0]["name"] == "get_current_weather" + assert tool_states[0]["arguments"] + args_dict_0 = json.loads(tool_states[0]["arguments"]) + assert args_dict_0["city"] == "Dallas" + assert args_dict_0["state"] == "TX" + + # Verify the second tool call (index=1) + assert tool_states[1]["name"] == "calculate_area" + assert tool_states[1]["arguments"] + args_dict_1 = json.loads(tool_states[1]["arguments"]) + assert args_dict_1["shape"] == "rectangle" + assert isinstance(args_dict_1["dimensions"], dict), "dimensions should be a dict" + assert args_dict_1["dimensions"]["width"] == 10 + assert args_dict_1["dimensions"]["height"] == 5 + + # Verify content: should contain "hello", "hi" + assert "hello" in other_content, "Should contain 'hello' as content" + assert "hi" in other_content, "Should contain 'hi' as content" + + # Verify the order: hello should come first, then hi + hello_index = other_content.find("hello") + hi_index = other_content.find("hi") + + assert hello_index >= 0, "'hello' should be in content" + assert hi_index > hello_index, "'hi' should come after 'hello'" + + # Verify that tool call tags are NOT in the content + assert "" not in other_content, ( + "First tool call should not be in content" + ) + assert "" not in other_content, ( + "Second tool call should not be in content" + ) + + +def test_extract_tool_calls_streaming_multiple_tool_calls_no_content_between( + step3p5_tool_parser, step3p5_tokenizer, sample_tools +): + """Test multiple tool calls with no content between them. + + Scenario: Model outputs "hello" + tool call + tool call + Expected: "hello" as content, first tool call parsed (index=0), + second tool call parsed (index=1). + No content should appear between the two tool calls. + """ + # Model output: hello + tool call + tool call (no content between tool calls) + model_output = """hello + + +Dallas + + +TX + + + + + +rectangle + + +{"width": 10, "height": 5} + + +""" + + request = ChatCompletionRequest(model=MODEL, messages=[], tools=sample_tools) + + other_content = "" + tool_states = {} + + for delta_message in stream_delta_message_generator( + step3p5_tool_parser, step3p5_tokenizer, model_output, request + ): + if delta_message.content: + other_content += delta_message.content + + if delta_message.tool_calls: + for tool_call in delta_message.tool_calls: + idx = tool_call.index + + if idx not in tool_states: + tool_states[idx] = { + "id": None, + "name": None, + "arguments": "", + "type": None, + } + + if tool_call.id: + tool_states[idx]["id"] = tool_call.id + + if tool_call.type: + assert tool_call.type == "function" + tool_states[idx]["type"] = tool_call.type + + if tool_call.function: + if tool_call.function.name: + tool_states[idx]["name"] = tool_call.function.name + + if tool_call.function.arguments is not None: + tool_states[idx]["arguments"] += tool_call.function.arguments + + # Should have exactly two complete tool calls + assert len(tool_states) == 2, "Should have exactly two complete tool calls" + + # Verify the first tool call (index=0) + assert tool_states[0]["name"] == "get_current_weather" + assert tool_states[0]["arguments"] + args_dict_0 = json.loads(tool_states[0]["arguments"]) + assert args_dict_0["city"] == "Dallas" + assert args_dict_0["state"] == "TX" + + # Verify the second tool call (index=1) + assert tool_states[1]["name"] == "calculate_area" + assert tool_states[1]["arguments"] + args_dict_1 = json.loads(tool_states[1]["arguments"]) + assert args_dict_1["shape"] == "rectangle" + assert isinstance(args_dict_1["dimensions"], dict), "dimensions should be a dict" + assert args_dict_1["dimensions"]["width"] == 10 + assert args_dict_1["dimensions"]["height"] == 5 + + assert "hello" in other_content, "Should contain 'hello' as content" + + # Verify that tool call tags are NOT in the content + assert "" not in other_content, ( + "First tool call should not be in content" + ) + assert "" not in other_content, ( + "Second tool call should not be in content" + ) + + +def test_extract_tool_calls_streaming_multi_token_chunk_boundary( + step3p5_tool_parser, step3p5_tokenizer, sample_tools +): + """Ensure fallback doesn't close a new tool_call when boundary is in one chunk.""" + request = ChatCompletionRequest(model=MODEL, messages=[], tools=sample_tools) + delta_text_chunks = [ + """ + + +Sys""", + """ + + +""", + """ +<""", + """function=calculate_area> + +rectangle""", + """ + +""", + ] + boundary_chunk = delta_text_chunks[1] + assert len(step3p5_tokenizer.encode(boundary_chunk, add_special_tokens=False)) > 1 + + tool_states = {} + for delta_message in stream_delta_message_generator_from_chunks( + step3p5_tool_parser, step3p5_tokenizer, delta_text_chunks, request + ): + print(delta_message) + if delta_message.tool_calls: + for tool_call in delta_message.tool_calls: + idx = tool_call.index + if idx not in tool_states: + tool_states[idx] = { + "name": None, + "arguments": "", + } + if tool_call.function: + if tool_call.function.name: + tool_states[idx]["name"] = tool_call.function.name + if tool_call.function.arguments is not None: + tool_states[idx]["arguments"] += tool_call.function.arguments + + assert len(tool_states) == 2 + assert all(state["name"] for state in tool_states.values()) + assert tool_states[0]["name"] == "get_current_weather" + assert tool_states[1]["name"] == "calculate_area" + + +def test_extract_tool_calls_non_streaming_multiple_tool_calls_no_content_between( + step3p5_tool_parser, sample_tools +): + """Test non-streaming extraction with tool calls and no content between them. + + Scenario: Model outputs "hello" + tool call + tool call. + Expected: "hello" as content, first tool call parsed (index=0), + second tool call parsed (index=1). + No content should appear between the two tool calls. + """ + # Model output: hello + tool call + tool call (no content between tool calls) + model_output = """hello + + +Dallas + + +TX + + + + + +rectangle + + +{"width": 10, "height": 5} + + +""" + + request = ChatCompletionRequest(model=MODEL, messages=[], tools=sample_tools) + + extracted_tool_calls = step3p5_tool_parser.extract_tool_calls( + model_output, request=request + ) + + # Should have exactly two complete tool calls + assert extracted_tool_calls.tools_called + assert len(extracted_tool_calls.tool_calls) == 2, ( + "Should have exactly two complete tool calls" + ) + + # Verify the first tool call (index=0) + assert extracted_tool_calls.tool_calls[0].function.name == "get_current_weather" + args_dict_0 = json.loads(extracted_tool_calls.tool_calls[0].function.arguments) + assert args_dict_0["city"] == "Dallas" + assert args_dict_0["state"] == "TX" + + # Verify the second tool call (index=1) + assert extracted_tool_calls.tool_calls[1].function.name == "calculate_area" + args_dict_1 = json.loads(extracted_tool_calls.tool_calls[1].function.arguments) + assert args_dict_1["shape"] == "rectangle" + assert isinstance(args_dict_1["dimensions"], dict), "dimensions should be a dict" + assert args_dict_1["dimensions"]["width"] == 10 + assert args_dict_1["dimensions"]["height"] == 5 + + # Verify content: should contain "hello" + assert extracted_tool_calls.content is not None + assert "hello" in extracted_tool_calls.content, "Should contain 'hello' as content" + + # Verify that tool call tags are NOT in the content + assert "" not in extracted_tool_calls.content, ( + "First tool call should not be in content" + ) + assert "" not in extracted_tool_calls.content, ( + "Second tool call should not be in content" + ) diff --git a/vllm/tool_parsers/step3p5_tool_parser.py b/vllm/tool_parsers/step3p5_tool_parser.py index b7c8699a03d..e52c0a706da 100644 --- a/vllm/tool_parsers/step3p5_tool_parser.py +++ b/vllm/tool_parsers/step3p5_tool_parser.py @@ -97,11 +97,26 @@ class StreamingXMLToolCallParser: """ # Record delta count before processing initial_delta_count = len(self.deltas) + entry_call_id = self.current_call_id + entry_tool_call_index = self.tool_call_index self.streaming_buffer += xml_chunk found_elements = self._process_complete_xml_elements() + fallback_call_id = None + if entry_call_id is not None: + if ( + self.current_call_id == entry_call_id + and self.tool_call_index == entry_tool_call_index + ): + fallback_call_id = entry_call_id + elif ( + self.current_call_id is not None + and self.tool_call_index == entry_tool_call_index + 1 + ): + fallback_call_id = self.current_call_id + if found_elements: # If complete elements found, check if end events were missed # some tags may not have been triggered @@ -110,7 +125,7 @@ class StreamingXMLToolCallParser: # If this chunk contains # but didn't generate '}', then complete it if ( - self.current_call_id is not None + fallback_call_id is not None and self.function_end_token in xml_chunk ): # - Added '}' (non-empty parameter ending) @@ -121,7 +136,7 @@ class StreamingXMLToolCallParser: and any( ( tc.function - and tc.id == self.current_call_id + and tc.id == fallback_call_id and isinstance(tc.function.arguments, str) and (tc.function.arguments in ("}", "{}")) ) @@ -139,7 +154,7 @@ class StreamingXMLToolCallParser: # If this chunk contains # but didn't generate final empty delta, then complete it if ( - self.current_call_id is not None + fallback_call_id is not None and self.tool_call_end_token in xml_chunk ): has_toolcall_close = any( @@ -150,7 +165,7 @@ class StreamingXMLToolCallParser: tc.type == "function" and tc.function and tc.function.arguments == "" - and tc.id == self.current_call_id + and tc.id == fallback_call_id ) for tc in td.tool_calls ) @@ -186,7 +201,7 @@ class StreamingXMLToolCallParser: # Only execute when still on the same call as when entered, # to prevent accidentally closing new calls # in multi scenarios - if self.current_call_id is not None and ( + if fallback_call_id is not None and ( self.function_end_token in xml_chunk or self.tool_call_end_token in xml_chunk ): From c1858b7ec8aa571dc0c0e00aded01019cca6a7e6 Mon Sep 17 00:00:00 2001 From: Aaron Hao Date: Thu, 5 Feb 2026 09:13:23 -0800 Subject: [PATCH 101/810] [Feat][RL][1/2] Native Weight Syncing API: NCCL (#31943) Signed-off-by: ahao-anyscale Signed-off-by: Aaron Hao Co-authored-by: SumanthRH --- .buildkite/test-amd.yaml | 7 + .buildkite/test-pipeline.yaml | 9 + .buildkite/test_areas/distributed.yaml | 6 + .../new_weight_syncing/rlhf.py | 208 ++++++++ .../new_weight_syncing/rlhf_async_new_apis.py | 283 +++++++++++ examples/online_serving/rlhf_http.py | 241 ++++++++++ tests/distributed/test_packed_tensor.py | 443 ++++++++++++++++++ tests/distributed/test_weight_transfer.py | 346 ++++++++++++++ .../entrypoints/openai/test_openai_schema.py | 8 + tests/entrypoints/weight_transfer/__init__.py | 3 + .../test_weight_transfer_llm.py | 300 ++++++++++++ vllm/config/__init__.py | 2 + vllm/config/vllm.py | 4 + vllm/config/weight_transfer.py | 15 + vllm/distributed/weight_transfer/__init__.py | 12 + vllm/distributed/weight_transfer/base.py | 158 +++++++ vllm/distributed/weight_transfer/factory.py | 116 +++++ .../weight_transfer/nccl_engine.py | 315 +++++++++++++ .../weight_transfer/packed_tensor.py | 216 +++++++++ vllm/engine/arg_utils.py | 14 + vllm/engine/protocol.py | 14 + vllm/entrypoints/llm.py | 53 +++ vllm/entrypoints/serve/rlhf/api_router.py | 67 ++- .../model_executor/layers/quantization/fp8.py | 9 + .../model_loader/reload/layerwise.py | 5 + vllm/v1/engine/async_llm.py | 45 ++ vllm/v1/worker/gpu_worker.py | 77 +++ 27 files changed, 2974 insertions(+), 2 deletions(-) create mode 100644 examples/offline_inference/new_weight_syncing/rlhf.py create mode 100644 examples/offline_inference/new_weight_syncing/rlhf_async_new_apis.py create mode 100644 examples/online_serving/rlhf_http.py create mode 100644 tests/distributed/test_packed_tensor.py create mode 100644 tests/distributed/test_weight_transfer.py create mode 100644 tests/entrypoints/weight_transfer/__init__.py create mode 100644 tests/entrypoints/weight_transfer/test_weight_transfer_llm.py create mode 100644 vllm/config/weight_transfer.py create mode 100644 vllm/distributed/weight_transfer/__init__.py create mode 100644 vllm/distributed/weight_transfer/base.py create mode 100644 vllm/distributed/weight_transfer/factory.py create mode 100644 vllm/distributed/weight_transfer/nccl_engine.py create mode 100644 vllm/distributed/weight_transfer/packed_tensor.py diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 64aaf1eb6ff..ca3bebcb0c3 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -233,6 +233,7 @@ steps: - tests/compile/fullgraph/test_basic_correctness.py - examples/offline_inference/rlhf.py - examples/offline_inference/rlhf_colocate.py + - examples/offline_inference/new_weight_syncing/ - tests/examples/offline_inference/data_parallel.py - tests/v1/distributed - tests/v1/engine/test_engine_core_client.py @@ -268,10 +269,16 @@ steps: - pytest -v -s distributed/test_symm_mem_allreduce.py # TODO: create a dedicated test section for multi-GPU example tests # when we have multiple distributed example tests + # OLD rlhf examples - pushd ../examples/offline_inference - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf.py - VLLM_ALLOW_INSECURE_SERIALIZATION=1 RAY_DEDUP_LOGS=0 python3 rlhf_colocate.py - popd + # NEW rlhf examples + - pushd ../examples/offline_inference/new_weight_syncing + - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf.py + - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_async_new_apis.py + - popd - label: Distributed Tests (8 GPUs) # 4min timeout_in_minutes: 10 diff --git a/.buildkite/test-pipeline.yaml b/.buildkite/test-pipeline.yaml index a3e25c0f755..b03e4b6d87b 100644 --- a/.buildkite/test-pipeline.yaml +++ b/.buildkite/test-pipeline.yaml @@ -206,6 +206,7 @@ steps: - tests/compile/fullgraph/test_basic_correctness.py - examples/offline_inference/rlhf.py - examples/offline_inference/rlhf_colocate.py + - examples/offline_inference/new_weight_syncing/ - tests/examples/offline_inference/data_parallel.py - tests/v1/distributed - tests/v1/engine/test_engine_core_client.py @@ -240,10 +241,16 @@ steps: - pytest -v -s distributed/test_symm_mem_allreduce.py # TODO: create a dedicated test section for multi-GPU example tests # when we have multiple distributed example tests + # OLD rlhf examples - pushd ../examples/offline_inference - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf.py - VLLM_ALLOW_INSECURE_SERIALIZATION=1 RAY_DEDUP_LOGS=0 python3 rlhf_colocate.py - popd + # NEW rlhf examples + - pushd ../examples/offline_inference/new_weight_syncing + - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf.py + - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_async_new_apis.py + - popd - label: Distributed Tests (8 GPUs) # 4min timeout_in_minutes: 10 @@ -1146,6 +1153,8 @@ steps: - pytest -v -s distributed/test_shm_broadcast.py - pytest -v -s distributed/test_shm_buffer.py - pytest -v -s distributed/test_shm_storage.py + - pytest -v -s distributed/test_packed_tensor.py + - pytest -v -s distributed/test_weight_transfer.py - label: 2 Node Tests (4 GPUs in total) # 16min timeout_in_minutes: 30 diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index ae4f45fbf4e..4fac613c351 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -62,6 +62,7 @@ steps: - tests/compile/fullgraph/test_basic_correctness.py - examples/offline_inference/rlhf.py - examples/offline_inference/rlhf_colocate.py + - examples/offline_inference/new_weight_syncing/ - tests/examples/offline_inference/data_parallel.py - tests/v1/distributed - tests/v1/engine/test_engine_core_client.py @@ -96,9 +97,14 @@ steps: - pytest -v -s distributed/test_symm_mem_allreduce.py # TODO: create a dedicated test section for multi-GPU example tests # when we have multiple distributed example tests + # OLD rlhf examples - cd ../examples/offline_inference - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf.py - VLLM_ALLOW_INSECURE_SERIALIZATION=1 RAY_DEDUP_LOGS=0 python3 rlhf_colocate.py + # NEW rlhf examples + - cd new_weight_syncing + - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf.py + - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_async_new_apis.py - label: Distributed Tests (8 GPUs)(H100) timeout_in_minutes: 10 diff --git a/examples/offline_inference/new_weight_syncing/rlhf.py b/examples/offline_inference/new_weight_syncing/rlhf.py new file mode 100644 index 00000000000..b3a3ca62f5a --- /dev/null +++ b/examples/offline_inference/new_weight_syncing/rlhf.py @@ -0,0 +1,208 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Demonstrates reinforcement learning using vLLM and Ray, +with native weight syncing APIs at engine instance. + +The script separates training and inference workloads onto distinct GPUs +so that Ray can manage process placement and inter-process communication. +A Hugging Face Transformer model occupies one GPU for training, whereas a +2x tensor-parallel vLLM inference engine occupies two GPUs. + +The example performs the following steps: +* Load the training model on one gpu (scheduled via ray) +* Initialize the inference model with dummy weights across + two gpus using vLLM's tensor parallelism and Ray placement groups. +* Generate gibberish from a list of prompts using the randomly initialized + inference engine. +* Update the weights of the training model and broadcast the updated weights + to the inference engine by using a Ray collective RPC group. +* Generating from the list of prompts after weight sync should result + in sensible outputs. + +This example assumes a single-node cluster with three GPUs, but Ray +supports multi-node clusters. vLLM expects the GPUs are only used for vLLM +workloads. Residual GPU activity interferes with vLLM memory profiling and +causes unexpected behavior. +""" + +import os + +import ray +from ray.util.placement_group import placement_group +from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy +from transformers import AutoModelForCausalLM + +from vllm import LLM, SamplingParams +from vllm.config import WeightTransferConfig +from vllm.distributed.weight_transfer.nccl_engine import ( + NCCLWeightTransferEngine, +) +from vllm.utils.network_utils import get_ip, get_open_port + +MODEL_NAME = "facebook/opt-125m" +# MODEL_NAME = "inference-optimization/Qwen3-0.6B-W4A16-G128" + + +class MyLLM(LLM): + """Configure the vLLM worker for Ray placement group execution.""" + + def __init__(self, *args, **kwargs): + os.environ["VLLM_RAY_BUNDLE_INDICES"] = "0,1" + super().__init__(*args, **kwargs) + + +@ray.remote(num_gpus=1) +class TrainModel: + """Ray actor that wraps the training model on a dedicated GPU.""" + + def __init__(self, model_name: str): + self.model = AutoModelForCausalLM.from_pretrained( + model_name, + ).to("cuda:0") + + self.port = get_open_port() + self.master_address = get_ip() + + def get_master_address_and_port(self): + return self.master_address, self.port + + def get_weight_metadata(self): + """Return weight names, dtypes, and shapes for weight transfer.""" + names = [] + dtype_names = [] + shapes = [] + for name, p in self.model.named_parameters(): + names.append(name) + dtype_names.append(str(p.dtype).split(".")[-1]) + shapes.append(list(p.shape)) + return names, dtype_names, shapes + + def init_weight_transfer_group(self, world_size): + """Initialize the NCCL process group for weight transfer.""" + self.model_update_group = NCCLWeightTransferEngine.trainer_init( + dict( + master_address=self.master_address, + master_port=self.port, + world_size=world_size, + ), + ) + + def broadcast_weights(self, packed: bool = True): + """Broadcast weights to the inference engine.""" + NCCLWeightTransferEngine.trainer_send_weights( + iterator=self.model.named_parameters(), + group=self.model_update_group, + packed=packed, + ) + + +# Initialize Ray and set the visible devices. The vLLM engine will +# be placed on GPUs 1 and 2. +ray.init() + +# Create a placement group that reserves GPU 1–2 for the vLLM inference engine. +# Learn more about Ray placement groups: +# https://docs.ray.io/en/latest/placement-groups.html +# Launch the training model actor. Ray's resource scheduler will allocate +# 1 GPU (via num_gpus=1 in the decorator), ensuring pg_inference gets different GPUs. +train_model = TrainModel.remote(MODEL_NAME) + +pg_inference = placement_group([{"GPU": 1, "CPU": 0}] * 2) +ray.get(pg_inference.ready()) +scheduling_inference = PlacementGroupSchedulingStrategy( + placement_group=pg_inference, + placement_group_capture_child_tasks=True, + placement_group_bundle_index=0, +) + +# Launch the vLLM inference engine. The `enforce_eager` flag reduces +# start-up latency. +# Note: Weight transfer APIs (init_weight_transfer_engine, update_weights) +# are now native to vLLM workers. +llm = ray.remote( + num_cpus=0, + num_gpus=0, + scheduling_strategy=scheduling_inference, +)(MyLLM).remote( + model=MODEL_NAME, + enforce_eager=True, + tensor_parallel_size=2, + data_parallel_size=1, + distributed_executor_backend="ray", + weight_transfer_config=WeightTransferConfig(backend="nccl"), + load_format="dummy", + quantization="fp8", +) + +# Generate text from the prompts. +prompts = [ + "Hello, my name is", + "The president of the United States is", + "The capital of France is", + "The future of AI is", +] + +sampling_params = SamplingParams(temperature=0) + +outputs = ray.get(llm.generate.remote(prompts, sampling_params)) + +# Generate text with the initial model. The output is expected to be nonsense +# because the weights are randomly initialized. +print("-" * 50) +for output in outputs: + prompt = output.prompt + generated_text = output.outputs[0].text + print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}") + print("-" * 50) + +# Set up the communication channel between the training process and the +# inference engine. +master_address, master_port = ray.get(train_model.get_master_address_and_port.remote()) + +world_size = ray.get(llm.get_world_size.remote()) + 1 # +1 for the trainer +inference_handle = llm.init_weight_transfer_engine.remote( + dict( + init_info=dict( + master_address=master_address, + master_port=master_port, + rank_offset=1, + world_size=world_size, + ) + ) +) + +# Initialize weight transfer group on both the training actor and inference engine +train_handle = train_model.init_weight_transfer_group.remote(world_size) +ray.get([train_handle, inference_handle]) + +# Synchronize the updated weights to the inference engine using batched API. +# Collect all weight metadata from the training actor +names, dtype_names, shapes = ray.get(train_model.get_weight_metadata.remote()) + +# Issue update_weights call with NCCL-specific update info +# packed=True enables efficient batched tensor broadcasting +inference_handle = llm.update_weights.remote( + dict( + update_info=dict( + names=names, + dtype_names=dtype_names, + shapes=shapes, + packed=True, + ) + ) +) + +# Broadcast all weights from trainer using the weight transfer API +train_handle = train_model.broadcast_weights.remote(packed=True) +ray.get([train_handle, inference_handle]) + +# Generate text with the updated model. The output is expected to be normal +# because the weights are updated. +outputs_updated = ray.get(llm.generate.remote(prompts, sampling_params)) +print("-" * 50) +for output in outputs_updated: + prompt = output.prompt + generated_text = output.outputs[0].text + print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}") + print("-" * 50) diff --git a/examples/offline_inference/new_weight_syncing/rlhf_async_new_apis.py b/examples/offline_inference/new_weight_syncing/rlhf_async_new_apis.py new file mode 100644 index 00000000000..835c16a7f55 --- /dev/null +++ b/examples/offline_inference/new_weight_syncing/rlhf_async_new_apis.py @@ -0,0 +1,283 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Demonstrates async reinforcement learning using vLLM and Ray, +with native weight syncing APIs at engine instance. + +The script separates training and inference workloads onto distinct GPUs +so that Ray can manage process placement and inter-process communication. +A Hugging Face Transformer model occupies one GPU for training, whereas a +2x tensor-parallel vLLM inference engine occupies two GPUs. + +The example performs the following steps: +* Load the training model on one gpu (scheduled via ray) +* Initialize the inference model with dummy weights across + two gpus using vLLM's tensor parallelism and Ray placement groups. +* Generate gibberish from a list of prompts using the randomly initialized + inference engine. +* Pause generation once generation completes for one sequence +* Update the weights of the training model and broadcast the updated weights + to the inference engine by using a Ray collective RPC group. +* Resume generation and print out the results + +This example assumes a single-node cluster with three GPUs, but Ray +supports multi-node clusters. vLLM expects the GPUs are only used for vLLM +workloads. Residual GPU activity interferes with vLLM memory profiling and +causes unexpected behavior. +""" + +import os +import uuid +from dataclasses import asdict + +import ray +import torch +from ray.util.placement_group import placement_group +from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy +from transformers import AutoModelForCausalLM, AutoTokenizer + +import vllm +from vllm import SamplingParams +from vllm.config import WeightTransferConfig +from vllm.distributed.weight_transfer.base import ( + WeightTransferInitRequest, + WeightTransferUpdateRequest, +) +from vllm.distributed.weight_transfer.nccl_engine import ( + NCCLWeightTransferEngine, + NCCLWeightTransferInitInfo, + NCCLWeightTransferUpdateInfo, +) +from vllm.utils.network_utils import get_ip, get_open_port +from vllm.v1.executor import Executor + +MODEL_NAME = "facebook/opt-125m" + + +class MyLLM(vllm.AsyncLLMEngine): + """Configure the vLLM worker for Ray placement group execution.""" + + def __init__(self, **kwargs): + os.environ["VLLM_RAY_BUNDLE_INDICES"] = "0,1" + engine_args = vllm.AsyncEngineArgs(**kwargs) + vllm_config = engine_args.create_engine_config() + executor_class = Executor.get_class(vllm_config) + super().__init__( + vllm_config=vllm_config, + executor_class=executor_class, + log_requests=engine_args.enable_log_requests, + log_stats=not engine_args.disable_log_stats, + ) + + async def generate_with_retry( + self, prompt_token_ids: list[int], sampling_params: vllm.SamplingParams + ) -> vllm.RequestOutput: + finish_reason = "abort" + while finish_reason == "abort": + async for request_output in self.generate( + {"prompt_token_ids": prompt_token_ids}, + sampling_params, + request_id=str(uuid.uuid4()), + ): + output = request_output + finish_reason = output.outputs[0].finish_reason + if finish_reason == "abort": + print( + f"ABORT, prompt_token_ids: {prompt_token_ids}, " + f"generated token_ids: {list(output.outputs[0].token_ids)}" + ) + prompt_token_ids = prompt_token_ids + list(output.outputs[0].token_ids) + return output + + +@ray.remote(num_gpus=1) +class TrainModel: + """Ray actor that wraps the training model on a dedicated GPU.""" + + def __init__(self, model_name: str): + self.model = AutoModelForCausalLM.from_pretrained( + model_name, dtype=torch.bfloat16 + ).to("cuda:0") + self.port = get_open_port() + self.master_address = get_ip() + + def get_master_address_and_port(self): + return self.master_address, self.port + + def get_weight_metadata(self): + """Return weight names, dtypes, and shapes for weight transfer.""" + names = [] + dtype_names = [] + shapes = [] + for name, p in self.model.named_parameters(): + names.append(name) + dtype_names.append(str(p.dtype).split(".")[-1]) + shapes.append(list(p.shape)) + return names, dtype_names, shapes + + def init_weight_transfer_group(self, world_size): + """Initialize the NCCL process group for weight transfer.""" + self.model_update_group = NCCLWeightTransferEngine.trainer_init( + dict( + master_address=self.master_address, + master_port=self.port, + world_size=world_size, + ), + ) + + def broadcast_weights(self, packed: bool = True): + """Broadcast weights to the inference engine.""" + NCCLWeightTransferEngine.trainer_send_weights( + iterator=self.model.named_parameters(), + group=self.model_update_group, + packed=packed, + ) + + +# Initialize Ray and set the visible devices. The vLLM engine will +# be placed on GPUs 1 and 2. +ray.init() + +# Launch the training model actor. Ray's resource scheduler will allocate +# 1 GPU (via num_gpus=1 in the decorator), ensuring pg_inference gets different GPUs. +train_model = TrainModel.remote(MODEL_NAME) + +# Create a placement group that reserves GPU 1–2 for the vLLM inference engine. +# Learn more about Ray placement groups: +# https://docs.ray.io/en/latest/placement-groups.html + +pg_inference = placement_group([{"GPU": 1, "CPU": 0}] * 2) +ray.get(pg_inference.ready()) +scheduling_inference = PlacementGroupSchedulingStrategy( + placement_group=pg_inference, + placement_group_capture_child_tasks=True, + placement_group_bundle_index=0, +) + +# Launch the vLLM inference engine. The `enforce_eager` flag reduces +# start-up latency. +# Note: Weight transfer APIs (init_weight_transfer_engine, update_weights) +# are now native to vLLM workers. +llm = ray.remote( + num_cpus=0, + num_gpus=0, + scheduling_strategy=scheduling_inference, +)(MyLLM).remote( + model=MODEL_NAME, + enforce_eager=True, + tensor_parallel_size=2, + distributed_executor_backend="ray", + load_format="dummy", + weight_transfer_config=WeightTransferConfig(backend="nccl"), +) + +# Generate text from the prompts. +prompts = [ + "My name is", + "The president of the United States is", + "The capital of France is", + "The future of AI is", +] + +# Tokenize prompts to token IDs +tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) +prompt_token_ids_list = [ + tokenizer.encode(prompt, add_special_tokens=False) for prompt in prompts +] + +sampling_params = [ + SamplingParams(temperature=0, max_tokens=2), + SamplingParams(temperature=0, max_tokens=32), + SamplingParams(temperature=0, max_tokens=32), + SamplingParams(temperature=0, max_tokens=32), +] + +# Set up the communication channel between the training process and the +# inference engine. +master_address, master_port = ray.get(train_model.get_master_address_and_port.remote()) + +world_size = 3 # 1 trainer + 2 inference workers (tensor_parallel_size=2) +inference_handle = llm.init_weight_transfer_engine.remote( + WeightTransferInitRequest( + init_info=asdict( + NCCLWeightTransferInitInfo( + master_address=master_address, + master_port=master_port, + rank_offset=1, + world_size=world_size, + ) + ) + ) +) + +# Initialize weight transfer group on both the training actor and inference engine +train_handle = train_model.init_weight_transfer_group.remote(world_size) +ray.get([train_handle, inference_handle]) + + +generation_futures = [ + llm.generate_with_retry.remote(prompt_token_ids, params) + for prompt_token_ids, params in zip(prompt_token_ids_list, sampling_params) +] + +finished, pending = ray.wait(generation_futures, num_returns=1) + +# Pause generation in preparation for weight sync +ray.get(llm.pause_generation.remote(wait_for_inflight_requests=False)) + +# Synchronize the updated weights to the inference engine using batched API. +# Collect all weight metadata from the training actor +names, dtype_names, shapes = ray.get(train_model.get_weight_metadata.remote()) + +# Issue update_weights call with NCCL-specific update info +# packed=True enables efficient batched tensor broadcasting +inference_handle = llm.update_weights.remote( + WeightTransferUpdateRequest( + update_info=asdict( + NCCLWeightTransferUpdateInfo( + names=names, + dtype_names=dtype_names, + shapes=shapes, + packed=True, + ) + ) + ) +) + +# Broadcast all weights from trainer using the weight transfer API +train_handle = train_model.broadcast_weights.remote(packed=True) +ray.get([train_handle, inference_handle]) + +# Resume generation since weight sync is complete +ray.get(llm.resume_generation.remote()) + +# Get outputs separately - finished completed before pause, pending were paused/resumed +finished_outputs = ray.get(finished) +pending_outputs = ray.get(pending) + +# Requests that finished before the pause: all generation used original weights +print("-" * 50) +print("Requests that completed BEFORE weight change:") +print("-" * 50) +for output in finished_outputs: + prompt_text = tokenizer.decode(output.prompt_token_ids) + print(f"Prompt: {prompt_text!r}") + print(f"Generated (with original weights): {output.outputs[0].text!r}") + print("-" * 50) + +# Requests that were paused mid-generation: some text before, some after weight change +print("Requests that were PAUSED and RESUMED after weight change:") +print("-" * 50) +for output in pending_outputs: + # Decode the full prompt token IDs (original + generated before pause) + full_prompt_text = tokenizer.decode(output.prompt_token_ids) + # Find the original prompt by checking which one this output started with + original_prompt = next(p for p in prompts if full_prompt_text.startswith(p)) + # output.prompt_token_ids contains original prompt + tokens generated before pause + # output.outputs[0].text is what was generated after resuming with new weights + text_before_pause = full_prompt_text[len(original_prompt) :] + text_after_pause = output.outputs[0].text + print(f"Original prompt: {original_prompt!r}") + print(f"Generated before weight change: {text_before_pause!r}") + print(f"Generated after weight change: {text_after_pause!r}") + print("-" * 50) diff --git a/examples/online_serving/rlhf_http.py b/examples/online_serving/rlhf_http.py new file mode 100644 index 00000000000..721a038a660 --- /dev/null +++ b/examples/online_serving/rlhf_http.py @@ -0,0 +1,241 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Demonstrates reinforcement learning from human feedback (RLHF) using vLLM +via HTTP API, with native weight syncing APIs. + +Unlike rlhf.py which creates a vLLM instance programmatically, this script +assumes you have already started a vLLM server using `vllm serve`. It uses: +- OpenAI-compatible API for inference requests +- HTTP endpoints for weight transfer control plane +- NCCL for actual weight data transfer + +Prerequisites: + Start a vLLM server with weight transfer enabled: + + $ VLLM_SERVER_DEV_MODE=1 vllm serve facebook/opt-125m \ + --enforce-eager \ + --weight-transfer-config '{"backend": "nccl"}' \ + --load-format dummy + + Then run this script: + + $ python rlhf_http.py + +The example performs the following steps: + +* Load the training model on GPU 0. +* Generate text using the vLLM server via OpenAI-compatible API. The output + is expected to be nonsense because the server is initialized with dummy weights. +* Initialize weight transfer via HTTP endpoint. +* Broadcast the real weights from the training model to the vLLM server + using NCCL. +* Generate text again to show normal output after the weight update. +""" + +import requests +import torch +from openai import OpenAI +from transformers import AutoModelForCausalLM + +from vllm.distributed.weight_transfer.nccl_engine import ( + NCCLWeightTransferEngine, +) +from vllm.utils.network_utils import get_ip, get_open_port + +BASE_URL = "http://localhost:8000" +MODEL_NAME = "facebook/opt-125m" + + +def generate_completions(client: OpenAI, model: str, prompts: list[str]) -> list[str]: + """Generate completions using the OpenAI-compatible API.""" + results = [] + for prompt in prompts: + response = client.completions.create( + model=model, + prompt=prompt, + max_tokens=32, + temperature=0, + ) + results.append(response.choices[0].text) + return results + + +def init_weight_transfer_engine( + base_url: str, + master_address: str, + master_port: int, + rank_offset: int, + world_size: int, +) -> None: + """Initialize weight transfer via HTTP endpoint.""" + url = f"{base_url}/init_weight_transfer_engine" + payload = { + "init_info": dict( + master_address=master_address, + master_port=master_port, + rank_offset=rank_offset, + world_size=world_size, + ) + } + response = requests.post(url, json=payload, timeout=60) + response.raise_for_status() + + +def update_weights( + base_url: str, + names: list[str], + dtype_names: list[str], + shapes: list[list[int]], + packed: bool = False, +) -> None: + """Update weights via HTTP endpoint.""" + url = f"{base_url}/update_weights" + payload = { + "update_info": dict( + names=names, + dtype_names=dtype_names, + shapes=shapes, + packed=packed, + ) + } + response = requests.post(url, json=payload, timeout=300) + response.raise_for_status() + + +def pause_generation(base_url: str) -> None: + """Pause generation via HTTP endpoint.""" + url = f"{base_url}/pause" + response = requests.post(url, timeout=60) + response.raise_for_status() + + +def resume_generation(base_url: str) -> None: + """Resume generation via HTTP endpoint.""" + url = f"{base_url}/resume" + response = requests.post(url, timeout=60) + response.raise_for_status() + + +def get_world_size(base_url: str) -> int: + """Get world size from the vLLM server.""" + url = f"{base_url}/get_world_size" + response = requests.get(url, timeout=10) + response.raise_for_status() + return response.json()["world_size"] + + +def main(): + # Get the inference world size from the vLLM server + inference_world_size = get_world_size(BASE_URL) + world_size = inference_world_size + 1 # +1 for the trainer + device = f"cuda:{inference_world_size}" + torch.cuda.set_device(device) + + # Load the training model + print(f"Loading training model: {MODEL_NAME}") + train_model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, dtype=torch.bfloat16) + train_model.to(device) + + # Create OpenAI client pointing to the vLLM server + client = OpenAI( + base_url=f"{BASE_URL}/v1", + api_key="EMPTY", # vLLM doesn't require an API key by default + ) + + # Test prompts + prompts = [ + "Hello, my name is", + "The president of the United States is", + "The capital of France is", + "The future of AI is", + ] + + # Generate text before weight update. The output is expected to be nonsense + # because the server is initialized with dummy weights. + print("-" * 50) + print("Generating text BEFORE weight update (expect nonsense):") + print("-" * 50) + outputs = generate_completions(client, MODEL_NAME, prompts) + for prompt, generated_text in zip(prompts, outputs): + print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}") + print("-" * 50) + + # Set up the communication channel between the training process and the + # vLLM server. The trainer is rank 0, vLLM worker(s) start at rank_offset. + master_address = get_ip() + master_port = get_open_port() + rank_offset = 1 + + print(f"Initializing weight transfer: master={master_address}:{master_port}") + + # Initialize weight transfer on vLLM server (this is async, server will + # wait for NCCL connection) + import threading + + init_thread = threading.Thread( + target=init_weight_transfer_engine, + args=(BASE_URL, master_address, master_port, rank_offset, world_size), + ) + init_thread.start() + + # Initialize NCCL process group on trainer side + model_update_group = NCCLWeightTransferEngine.trainer_init( + dict( + master_address=master_address, + master_port=master_port, + world_size=world_size, + ), + ) + + # Wait for init_weight_transfer_engine to complete + init_thread.join() + + # Pause generation before weight sync + pause_generation(BASE_URL) + + # Collect weight metadata for the update request + names = [] + dtype_names = [] + shapes = [] + for name, p in train_model.named_parameters(): + names.append(name) + dtype_names.append(str(p.dtype).split(".")[-1]) + shapes.append(list(p.shape)) + + # Start the update_weights call in a separate thread since it will block + # waiting for NCCL broadcasts + # packed=True enables efficient batched tensor broadcasting + update_thread = threading.Thread( + target=update_weights, + args=(BASE_URL, names, dtype_names, shapes, True), # packed=True + ) + update_thread.start() + + # Broadcast all weights from trainer to vLLM workers + print("Broadcasting weights via NCCL...") + NCCLWeightTransferEngine.trainer_send_weights( + iterator=train_model.named_parameters(), + group=model_update_group, + packed=True, + ) + + # Wait for update_weights to complete + update_thread.join() + + # Resume generation after weight sync + resume_generation(BASE_URL) + + # Generate text after weight update. The output is expected to be normal + # because the real weights are now loaded. + print("-" * 50) + print("Generating text AFTER weight update:") + print("-" * 50) + outputs_updated = generate_completions(client, MODEL_NAME, prompts) + for prompt, generated_text in zip(prompts, outputs_updated): + print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}") + print("-" * 50) + + +if __name__ == "__main__": + main() diff --git a/tests/distributed/test_packed_tensor.py b/tests/distributed/test_packed_tensor.py new file mode 100644 index 00000000000..134629e2b79 --- /dev/null +++ b/tests/distributed/test_packed_tensor.py @@ -0,0 +1,443 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for packed tensor broadcasting functionality. + +Unit tests for packed_broadcast_producer and packed_broadcast_consumer. +These utilities enable efficient batched tensor transfer over NCCL. +""" + +import pytest +import torch + +from vllm.distributed.weight_transfer.nccl_engine import NCCLWeightTransferUpdateInfo +from vllm.distributed.weight_transfer.packed_tensor import ( + packed_broadcast_consumer, + packed_broadcast_producer, +) + + +class MockCommunicationGroup: + """Mock communication group for testing producer broadcast operations.""" + + def __init__(self): + self.broadcasted_tensors: list[torch.Tensor] = [] + self.broadcast_count = 0 + self.device = torch.device("cuda:0") + + def broadcast(self, tensor, src): + """Mock broadcast that stores the tensor for later verification.""" + self.broadcasted_tensors.append(tensor.clone()) + self.broadcast_count += 1 + + +class MockConsumerCommunicationGroup: + """Mock communication group for consumer that returns pre-stored tensors.""" + + def __init__(self, tensors_to_return: list[torch.Tensor]): + self.tensors_to_return = tensors_to_return + self.current_index = 0 + self.device = torch.device("cuda:0") + + def broadcast(self, tensor, src): + """Mock broadcast that fills the tensor with pre-stored data.""" + if self.current_index < len(self.tensors_to_return): + tensor.copy_(self.tensors_to_return[self.current_index]) + self.current_index += 1 + + +def create_mock_model_params( + num_layers: int = 3, + dtype: torch.dtype = torch.float32, +) -> list[tuple[str, torch.Tensor]]: + """Create mock model parameters for testing.""" + params = [] + for i in range(num_layers): + params.append((f"layer{i}.weight", torch.randn(10, 20, dtype=dtype))) + params.append((f"layer{i}.bias", torch.randn(10, dtype=dtype))) + return params + + +def create_state_dict_info( + params: list[tuple[str, torch.Tensor]], +) -> dict[str, tuple[tuple[int, ...], torch.dtype]]: + """Create state dict info (name -> (shape, dtype)) from params.""" + return {name: (tuple(tensor.shape), tensor.dtype) for name, tensor in params} + + +# --- Unit Tests: NCCLWeightTransferUpdateInfo packed field --- + + +class TestNCCLWeightTransferUpdateInfoPacked: + """Test NCCLWeightTransferUpdateInfo dataclass packed field.""" + + def test_packed_default_false(self): + """Test that packed defaults to False.""" + info = NCCLWeightTransferUpdateInfo( + names=["layer.weight"], + dtype_names=["float32"], + shapes=[[10, 10]], + ) + assert info.packed is False + + def test_packed_can_be_set_true(self): + """Test that packed can be set to True.""" + info = NCCLWeightTransferUpdateInfo( + names=["layer.weight"], + dtype_names=["float32"], + shapes=[[10, 10]], + packed=True, + ) + assert info.packed is True + + +# --- Unit Tests: packed_broadcast_producer --- + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +class TestPackedBroadcastProducer: + """Test packed_broadcast_producer function.""" + + def test_producer_broadcasts_tensors(self): + """Test that producer broadcasts all tensors.""" + params = create_mock_model_params() + params_cuda = [(name, tensor.cuda()) for name, tensor in params] + + mock_group = MockCommunicationGroup() + + # Use a small target size to force multiple batches + packed_broadcast_producer( + iterator=iter(params_cuda), + group=mock_group, + src=0, + post_iter_func=lambda x: x[1], + buffer_size_bytes=500, + ) + + # Should have broadcasted some tensors + assert mock_group.broadcast_count > 0 + assert len(mock_group.broadcasted_tensors) > 0 + + def test_producer_single_large_tensor(self): + """Test with a single tensor larger than target size.""" + # Create a large tensor + large_tensor = torch.randn(1000, 1000, dtype=torch.float32).cuda() + params = [("large_weight", large_tensor)] + + mock_group = MockCommunicationGroup() + + # Small target size to force the tensor to exceed it + packed_broadcast_producer( + iterator=iter(params), + group=mock_group, + src=0, + post_iter_func=lambda x: x[1], + buffer_size_bytes=100, + ) + + # Should still broadcast the tensor (at least 1 broadcast) + assert mock_group.broadcast_count >= 1 + assert len(mock_group.broadcasted_tensors) >= 1 + + # Verify the total broadcasted size matches the tensor + expected_size = large_tensor.numel() * large_tensor.element_size() + actual_size = sum(t.numel() for t in mock_group.broadcasted_tensors) + assert actual_size == expected_size + + def test_producer_multiple_batches(self): + """Test that tensors are properly batched when exceeding target size.""" + # Create many small tensors + params = [ + (f"weight_{i}", torch.randn(10, 10, dtype=torch.float32).cuda()) + for i in range(20) + ] + + mock_group = MockCommunicationGroup() + + # Small target size to force multiple batches + packed_broadcast_producer( + iterator=iter(params), + group=mock_group, + src=0, + post_iter_func=lambda x: x[1], + buffer_size_bytes=2000, + ) + + # Should have multiple broadcasts + assert mock_group.broadcast_count > 1 + + # Total size should match sum of all tensors + expected_total = sum(t.numel() * t.element_size() for _, t in params) + actual_total = sum(t.numel() for t in mock_group.broadcasted_tensors) + assert actual_total == expected_total + + def test_producer_empty_iterator(self): + """Test producer handles empty iterator gracefully.""" + mock_group = MockCommunicationGroup() + + packed_broadcast_producer( + iterator=iter([]), + group=mock_group, + src=0, + post_iter_func=lambda x: x[1], + buffer_size_bytes=1000, + ) + + # No broadcasts for empty iterator + assert mock_group.broadcast_count == 0 + + +# --- Unit Tests: packed_broadcast_consumer --- + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +class TestPackedBroadcastConsumer: + """Test packed_broadcast_consumer function.""" + + def test_consumer_receives_tensors(self): + """Test that consumer receives and unpacks tensors.""" + params = create_mock_model_params() + params_cuda = [(name, tensor.cuda()) for name, tensor in params] + + buffer_size = 2000 + + # First, run producer to get the broadcasted tensors + producer_group = MockCommunicationGroup() + + packed_broadcast_producer( + iterator=iter(params_cuda), + group=producer_group, + src=0, + post_iter_func=lambda x: x[1], + buffer_size_bytes=buffer_size, + ) + + # Now run consumer with the broadcasted tensors + consumer_group = MockConsumerCommunicationGroup( + producer_group.broadcasted_tensors + ) + + state_dict_info = create_state_dict_info(params_cuda) + + unpacked_tensors = {} + + def post_unpack_func(tensor_list): + for name, tensor in tensor_list: + unpacked_tensors[name] = tensor.clone() + + packed_broadcast_consumer( + iterator=iter(state_dict_info.items()), + group=consumer_group, + src=0, + post_unpack_func=post_unpack_func, + buffer_size_bytes=buffer_size, + ) + + # Verify all parameters were unpacked + assert len(unpacked_tensors) == len(params) + + # Verify each tensor matches the original + for name, original_tensor in params_cuda: + assert name in unpacked_tensors + unpacked = unpacked_tensors[name] + assert unpacked.shape == original_tensor.shape + assert unpacked.dtype == original_tensor.dtype + assert torch.allclose(unpacked, original_tensor, rtol=1e-5, atol=1e-7) + + +# --- Integration Tests: Producer-Consumer Roundtrip --- + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +class TestPackedBroadcastRoundtrip: + """Test producer-consumer roundtrip behavior.""" + + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) + def test_roundtrip_different_dtypes(self, dtype): + """Test roundtrip with different data types.""" + params = create_mock_model_params(num_layers=2, dtype=dtype) + params_cuda = [(name, tensor.cuda()) for name, tensor in params] + + buffer_size = 1000 + producer_group = MockCommunicationGroup() + + packed_broadcast_producer( + iterator=iter(params_cuda), + group=producer_group, + src=0, + post_iter_func=lambda x: x[1], + buffer_size_bytes=buffer_size, + ) + + consumer_group = MockConsumerCommunicationGroup( + producer_group.broadcasted_tensors + ) + + state_dict_info = create_state_dict_info(params_cuda) + unpacked_tensors = {} + + def post_unpack_func(tensor_list): + for name, tensor in tensor_list: + unpacked_tensors[name] = tensor.clone() + + packed_broadcast_consumer( + iterator=iter(state_dict_info.items()), + group=consumer_group, + src=0, + post_unpack_func=post_unpack_func, + buffer_size_bytes=buffer_size, + ) + + # Verify roundtrip preserves data + for name, original_tensor in params_cuda: + assert name in unpacked_tensors + unpacked = unpacked_tensors[name] + assert unpacked.dtype == dtype + assert torch.allclose(unpacked, original_tensor, rtol=1e-4, atol=1e-6) + + def test_roundtrip_mixed_dtypes(self): + """Test roundtrip with mixed data types.""" + # Create params with mixed dtypes + params = [ + ("layer1.weight", torch.randn(10, 20, dtype=torch.float32).cuda()), + ("layer1.bias", torch.randn(10, dtype=torch.float16).cuda()), + ("layer2.weight", torch.randn(20, 30, dtype=torch.bfloat16).cuda()), + ] + + buffer_size = 500 + producer_group = MockCommunicationGroup() + + packed_broadcast_producer( + iterator=iter(params), + group=producer_group, + src=0, + post_iter_func=lambda x: x[1], + buffer_size_bytes=buffer_size, + ) + + consumer_group = MockConsumerCommunicationGroup( + producer_group.broadcasted_tensors + ) + + state_dict_info = create_state_dict_info(params) + unpacked_tensors = {} + + def post_unpack_func(tensor_list): + for name, tensor in tensor_list: + unpacked_tensors[name] = tensor.clone() + + packed_broadcast_consumer( + iterator=iter(state_dict_info.items()), + group=consumer_group, + src=0, + post_unpack_func=post_unpack_func, + buffer_size_bytes=buffer_size, + ) + + # Verify all params roundtrip correctly with correct dtypes + for name, original_tensor in params: + assert name in unpacked_tensors + unpacked = unpacked_tensors[name] + assert unpacked.shape == original_tensor.shape + assert unpacked.dtype == original_tensor.dtype + assert torch.allclose(unpacked, original_tensor, rtol=1e-4, atol=1e-6) + + @pytest.mark.parametrize("target_size", [100, 1000, 10000, 100000]) + def test_roundtrip_different_batch_sizes(self, target_size): + """Test roundtrip with different target batch sizes.""" + params = create_mock_model_params(num_layers=5) + params_cuda = [(name, tensor.cuda()) for name, tensor in params] + + producer_group = MockCommunicationGroup() + + packed_broadcast_producer( + iterator=iter(params_cuda), + group=producer_group, + src=0, + post_iter_func=lambda x: x[1], + buffer_size_bytes=target_size, + ) + + consumer_group = MockConsumerCommunicationGroup( + producer_group.broadcasted_tensors + ) + + state_dict_info = create_state_dict_info(params_cuda) + unpacked_tensors = {} + + def post_unpack_func(tensor_list): + for name, tensor in tensor_list: + unpacked_tensors[name] = tensor.clone() + + packed_broadcast_consumer( + iterator=iter(state_dict_info.items()), + group=consumer_group, + src=0, + post_unpack_func=post_unpack_func, + buffer_size_bytes=target_size, + ) + + # Verify all params roundtrip correctly + assert len(unpacked_tensors) == len(params) + for name, original_tensor in params_cuda: + assert name in unpacked_tensors + assert torch.allclose( + unpacked_tensors[name], original_tensor, rtol=1e-5, atol=1e-7 + ) + + def test_roundtrip_non_contiguous_tensors(self): + """Test roundtrip with non-contiguous tensors from the trainer.""" + # Create non-contiguous tensors (simulating trainer outputs) + # Transposed tensors are non-contiguous + weight1 = torch.randn(20, 10, dtype=torch.float32).cuda().T + # Sliced tensors with step are non-contiguous + weight2 = torch.randn(40, 30, dtype=torch.float16).cuda()[::2, ::2] + # Permuted tensors are non-contiguous + weight3 = torch.randn(5, 10, 15, dtype=torch.bfloat16).cuda().permute(2, 0, 1) + + params = [ + ("layer1.weight", weight1), + ("layer2.weight", weight2), + ("layer3.weight", weight3), + ] + + # Verify tensors are indeed non-contiguous + for name, tensor in params: + assert not tensor.is_contiguous(), f"{name} should be non-contiguous" + + buffer_size = 500 + producer_group = MockCommunicationGroup() + + packed_broadcast_producer( + iterator=iter(params), + group=producer_group, + src=0, + post_iter_func=lambda x: x[1], + buffer_size_bytes=buffer_size, + ) + + consumer_group = MockConsumerCommunicationGroup( + producer_group.broadcasted_tensors + ) + + state_dict_info = create_state_dict_info(params) + unpacked_tensors = {} + + def post_unpack_func(tensor_list): + for name, tensor in tensor_list: + unpacked_tensors[name] = tensor.clone() + + packed_broadcast_consumer( + iterator=iter(state_dict_info.items()), + group=consumer_group, + src=0, + post_unpack_func=post_unpack_func, + buffer_size_bytes=buffer_size, + ) + + # Verify all non-contiguous params roundtrip correctly + for name, original_tensor in params: + assert name in unpacked_tensors + unpacked = unpacked_tensors[name] + assert unpacked.shape == original_tensor.shape + assert unpacked.dtype == original_tensor.dtype + assert torch.allclose(unpacked, original_tensor, rtol=1e-4, atol=1e-6) diff --git a/tests/distributed/test_weight_transfer.py b/tests/distributed/test_weight_transfer.py new file mode 100644 index 00000000000..4c348dd799b --- /dev/null +++ b/tests/distributed/test_weight_transfer.py @@ -0,0 +1,346 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for weight transfer engine backends. + +Unit tests for engine classes (parsing, validation, registry). +Integration test for NCCL weight transfer between processes using Ray. +""" + +from unittest.mock import MagicMock + +import pytest +import ray +import torch + +from vllm.config.parallel import ParallelConfig +from vllm.config.weight_transfer import WeightTransferConfig +from vllm.distributed.weight_transfer import WeightTransferEngineFactory +from vllm.distributed.weight_transfer.nccl_engine import ( + NCCLWeightTransferEngine, + NCCLWeightTransferInitInfo, + NCCLWeightTransferUpdateInfo, +) +from vllm.utils.network_utils import get_open_port + + +def create_mock_parallel_config( + rank: int = 0, + world_size: int = 1, + dp_rank: int = 0, +) -> ParallelConfig: + """Create a mock ParallelConfig for testing.""" + config = MagicMock(spec=ParallelConfig) + config.rank = rank + config.world_size = world_size + config.data_parallel_rank = dp_rank + return config + + +# --- Unit Tests: NCCLWeightTransferUpdateInfo Validation --- + + +class TestNCCLWeightTransferUpdateInfoValidation: + """Test NCCLWeightTransferUpdateInfo dataclass validation.""" + + def test_valid_update_info(self): + """Test creating valid NCCLWeightTransferUpdateInfo.""" + info = NCCLWeightTransferUpdateInfo( + names=["layer.weight", "layer.bias"], + dtype_names=["float32", "float32"], + shapes=[[10, 10], [10]], + ) + assert info.names == ["layer.weight", "layer.bias"] + assert info.dtype_names == ["float32", "float32"] + assert info.shapes == [[10, 10], [10]] + + def test_mismatched_dtype_names_raises(self): + """Test that mismatched dtype_names length raises ValueError.""" + with pytest.raises(ValueError, match="dtype_names"): + NCCLWeightTransferUpdateInfo( + names=["layer.weight", "layer.bias"], + dtype_names=["float32"], # Only one dtype + shapes=[[10, 10], [10]], + ) + + def test_mismatched_shapes_raises(self): + """Test that mismatched shapes length raises ValueError.""" + with pytest.raises(ValueError, match="shapes"): + NCCLWeightTransferUpdateInfo( + names=["layer.weight", "layer.bias"], + dtype_names=["float32", "float32"], + shapes=[[10, 10]], # Only one shape + ) + + def test_empty_lists_valid(self): + """Test that empty lists are valid.""" + info = NCCLWeightTransferUpdateInfo( + names=[], + dtype_names=[], + shapes=[], + ) + assert len(info.names) == 0 + + +# --- Unit Tests: Engine Parsing --- + + +class TestNCCLEngineParsing: + """Test NCCLWeightTransferEngine parsing methods.""" + + def test_parse_init_info_valid(self): + """Test parsing valid init info dict.""" + config = WeightTransferConfig(backend="nccl") + parallel_config = create_mock_parallel_config() + engine = NCCLWeightTransferEngine(config, parallel_config) + + init_info = engine.parse_init_info( + { + "master_address": "127.0.0.1", + "master_port": 12345, + "rank_offset": 1, + "world_size": 3, + } + ) + + assert isinstance(init_info, NCCLWeightTransferInitInfo) + assert init_info.master_address == "127.0.0.1" + assert init_info.master_port == 12345 + assert init_info.rank_offset == 1 + assert init_info.world_size == 3 + + def test_parse_init_info_missing_field_raises(self): + """Test parsing init info with missing required field.""" + config = WeightTransferConfig(backend="nccl") + parallel_config = create_mock_parallel_config() + engine = NCCLWeightTransferEngine(config, parallel_config) + + with pytest.raises(ValueError, match="Invalid init_info"): + engine.parse_init_info( + { + "master_address": "127.0.0.1", + # Missing master_port, rank_offset, world_size + } + ) + + def test_parse_update_info_valid(self): + """Test parsing valid update info dict.""" + config = WeightTransferConfig(backend="nccl") + parallel_config = create_mock_parallel_config() + engine = NCCLWeightTransferEngine(config, parallel_config) + + update_info = engine.parse_update_info( + { + "names": ["w1", "w2"], + "dtype_names": ["float32", "bfloat16"], + "shapes": [[100, 100], [50]], + } + ) + + assert isinstance(update_info, NCCLWeightTransferUpdateInfo) + assert update_info.names == ["w1", "w2"] + assert update_info.dtype_names == ["float32", "bfloat16"] + assert update_info.shapes == [[100, 100], [50]] + + +# --- Unit Tests: Engine Registry --- + + +class TestEngineRegistry: + """Test weight transfer engine registry.""" + + def test_create_engine_nccl(self): + """Test factory creates NCCL engine.""" + config = WeightTransferConfig(backend="nccl") + parallel_config = create_mock_parallel_config() + engine = WeightTransferEngineFactory.create_engine(config, parallel_config) + assert isinstance(engine, NCCLWeightTransferEngine) + + def test_create_engine_invalid_backend(self): + """Test factory raises for invalid backend.""" + config = WeightTransferConfig(backend="invalid") + parallel_config = create_mock_parallel_config() + with pytest.raises(ValueError, match="Invalid weight transfer backend"): + WeightTransferEngineFactory.create_engine(config, parallel_config) + + def test_register_duplicate_raises(self): + """Test registering duplicate engine name raises.""" + with pytest.raises(ValueError, match="already registered"): + WeightTransferEngineFactory.register_engine( + "nccl", NCCLWeightTransferEngine + ) + + +# --- Test receive_weights without init raises --- + + +def test_nccl_receive_weights_without_init_raises(): + """Test that receive_weights raises if init_transfer_engine wasn't called.""" + if torch.cuda.device_count() < 1: + pytest.skip("Need at least 1 GPU for this test") + + config = WeightTransferConfig(backend="nccl") + parallel_config = create_mock_parallel_config() + engine = NCCLWeightTransferEngine(config, parallel_config) + + update_info = NCCLWeightTransferUpdateInfo( + names=["w"], + dtype_names=["float32"], + shapes=[[10]], + ) + + with pytest.raises(RuntimeError, match="not initialized"): + engine.receive_weights(update_info, lambda x: None) + + +# --- Integration Test: NCCL Weight Transfer Between Ray Tasks --- + + +@ray.remote(num_gpus=1) +def trainer_broadcast_tensor( + master_address: str, + master_port: int, + world_size: int, + tensor_shape: list[int], + tensor_dtype: str, +) -> bool: + """Trainer task that broadcasts a tensor via NCCL.""" + import torch + + from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator + from vllm.distributed.utils import StatelessProcessGroup + + # Create process group as rank 0 (trainer) + pg = StatelessProcessGroup.create( + host=master_address, + port=master_port, + rank=0, + world_size=world_size, + ) + # Ray sets CUDA_VISIBLE_DEVICES, so device 0 is the assigned GPU + comm = PyNcclCommunicator(pg, device=0) + + # Create and broadcast the tensor + dtype = getattr(torch, tensor_dtype) + tensor_to_send = torch.ones(tensor_shape, dtype=dtype, device="cuda:0") + comm.broadcast(tensor_to_send, src=0, stream=torch.cuda.current_stream()) + torch.cuda.synchronize() + + return True + + +@ray.remote(num_gpus=1) +def inference_receive_tensor( + master_address: str, + master_port: int, + world_size: int, + tensor_shape: list[int], + tensor_dtype: str, +) -> dict: + """Inference task that receives tensor via NCCLWeightTransferEngine.""" + from unittest.mock import MagicMock + + import torch + + from vllm.config.parallel import ParallelConfig + from vllm.config.weight_transfer import WeightTransferConfig + from vllm.distributed.weight_transfer.nccl_engine import ( + NCCLWeightTransferEngine, + NCCLWeightTransferInitInfo, + NCCLWeightTransferUpdateInfo, + ) + + # Create engine with mock parallel config + config = WeightTransferConfig(backend="nccl") + parallel_config = MagicMock(spec=ParallelConfig) + parallel_config.rank = 0 + parallel_config.world_size = 1 + parallel_config.data_parallel_rank = 0 + + engine = NCCLWeightTransferEngine(config, parallel_config) + + # Initialize the engine (joins as rank 1) + init_info = NCCLWeightTransferInitInfo( + master_address=master_address, + master_port=master_port, + rank_offset=1, # Trainer is rank 0, we become rank 1 + world_size=world_size, + ) + engine.init_transfer_engine(init_info) + + # Receive weights with a no-op load_weights that captures the tensor + received_tensors = [] + + def noop_load_weights(weights: list[tuple[str, torch.Tensor]]): + for name, tensor in weights: + # Clone tensor to keep it after engine cleans up + received_tensors.append((name, tensor.clone())) + + update_info = NCCLWeightTransferUpdateInfo( + names=["test.weight"], + dtype_names=[tensor_dtype], + shapes=[tensor_shape], + ) + engine.receive_weights(update_info, noop_load_weights) + torch.cuda.synchronize() + + # Verify we received the tensor + success = False + received_shape = None + received_sum = None + + if len(received_tensors) == 1: + name, tensor = received_tensors[0] + received_shape = list(tensor.shape) + received_sum = tensor.sum().item() + # Check shape matches and values are all 1s (trainer sends ones) + if received_shape == tensor_shape: + expected_sum = 1.0 * torch.tensor(tensor_shape).prod().item() + if abs(received_sum - expected_sum) < 0.01: + success = True + + engine.shutdown() + + return { + "success": success, + "received_shape": received_shape, + "received_sum": received_sum, + } + + +@pytest.mark.skipif( + torch.cuda.device_count() < 2, + reason="Need at least 2 GPUs to run NCCL weight transfer test.", +) +def test_nccl_weight_transfer_between_processes(): + """Test NCCL weight transfer from trainer to inference process using Ray. + + This test verifies that the NCCLWeightTransferEngine can receive + tensors broadcast by a trainer process via NCCL. + """ + ray.init(ignore_reinit_error=True) + + master_address = "127.0.0.1" + master_port = get_open_port() + world_size = 2 # 1 trainer + 1 inference worker + + # Tensor to transfer: 100x100 ones + tensor_shape = [100, 100] + tensor_dtype = "float32" + + # Start both tasks concurrently - Ray assigns GPUs automatically + inference_future = inference_receive_tensor.remote( + master_address, master_port, world_size, tensor_shape, tensor_dtype + ) + trainer_future = trainer_broadcast_tensor.remote( + master_address, master_port, world_size, tensor_shape, tensor_dtype + ) + + # Wait for both to complete + trainer_result, result = ray.get([trainer_future, inference_future]) + + assert trainer_result, "Trainer should complete successfully" + assert result["success"], ( + f"Weight transfer failed. " + f"Received shape: {result['received_shape']}, " + f"Received sum: {result['received_sum']}" + ) diff --git a/tests/entrypoints/openai/test_openai_schema.py b/tests/entrypoints/openai/test_openai_schema.py index 50d24a40054..1baab9934fd 100644 --- a/tests/entrypoints/openai/test_openai_schema.py +++ b/tests/entrypoints/openai/test_openai_schema.py @@ -139,6 +139,14 @@ def test_openapi_stateless(case: schemathesis.Case): # Skip responses API as it is meant to be stateful. return + # Skip weight transfer endpoints as they require special setup + # (weight_transfer_config) and are meant to be stateful. + if case.operation.path in ( + "/init_weight_transfer_engine", + "/update_weights", + ): + return + timeout = { # requires a longer timeout ("POST", "/v1/chat/completions"): LONG_TIMEOUT_SECONDS, diff --git a/tests/entrypoints/weight_transfer/__init__.py b/tests/entrypoints/weight_transfer/__init__.py new file mode 100644 index 00000000000..6655f891362 --- /dev/null +++ b/tests/entrypoints/weight_transfer/__init__.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + diff --git a/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py b/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py new file mode 100644 index 00000000000..9f2309c765b --- /dev/null +++ b/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py @@ -0,0 +1,300 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for weight transfer APIs via LLM class. + +These tests use a mock weight transfer engine to verify that the API +calls the correct methods with the right arguments, without requiring +actual NCCL communication. +""" + +import os +from collections.abc import Callable +from dataclasses import dataclass +from unittest.mock import patch + +import pytest +import torch + +from vllm import LLM +from vllm.config import WeightTransferConfig +from vllm.distributed.weight_transfer.base import ( + WeightTransferEngine, + WeightTransferInitInfo, + WeightTransferInitRequest, + WeightTransferUpdateInfo, + WeightTransferUpdateRequest, +) + +from ...utils import create_new_process_for_each_test + +# Use a tiny model for fast testing +MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM" + + +# --- Mock Weight Transfer Engine --- + + +@dataclass +class MockInitInfo(WeightTransferInitInfo): + """Mock initialization info.""" + + test_param: str = "test" + + +@dataclass +class MockUpdateInfo(WeightTransferUpdateInfo): + """Mock update info.""" + + names: list[str] | None = None + dtype_names: list[str] | None = None + shapes: list[list[int]] | None = None + + +class MockWeightTransferEngine(WeightTransferEngine[MockInitInfo, MockUpdateInfo]): + """Mock weight transfer engine that tracks method calls.""" + + init_info_cls = MockInitInfo + update_info_cls = MockUpdateInfo + + # Class-level tracking for verification across processes + init_transfer_engine_called: bool = False + receive_weights_called: bool = False + shutdown_called: bool = False + last_init_info: MockInitInfo | None = None + last_update_info: MockUpdateInfo | None = None + + def __init__(self, config, parallel_config): + super().__init__(config, parallel_config) + # Reset tracking on init + MockWeightTransferEngine.init_transfer_engine_called = False + MockWeightTransferEngine.receive_weights_called = False + MockWeightTransferEngine.shutdown_called = False + MockWeightTransferEngine.last_init_info = None + MockWeightTransferEngine.last_update_info = None + + def init_transfer_engine(self, init_info: MockInitInfo) -> None: + MockWeightTransferEngine.init_transfer_engine_called = True + MockWeightTransferEngine.last_init_info = init_info + + def receive_weights( + self, + update_info: MockUpdateInfo, + load_weights: Callable[[list[tuple[str, torch.Tensor]]], None], + ) -> None: + MockWeightTransferEngine.receive_weights_called = True + MockWeightTransferEngine.last_update_info = update_info + # Simulate loading weights by calling load_weights with empty list + # (In real implementation, this would receive and load actual weights) + load_weights([]) + + def shutdown(self) -> None: + MockWeightTransferEngine.shutdown_called = True + + +def mock_create_engine(config, parallel_config): + """Mock factory function that returns our mock engine.""" + return MockWeightTransferEngine(config, parallel_config) + + +# --- Tests --- + + +@create_new_process_for_each_test() +def test_get_world_size_tp1(): + """Test world_size is correctly configured for TP=1.""" + if torch.cuda.device_count() < 1: + pytest.skip("Need at least 1 GPU for this test") + + llm = LLM( + model=MODEL_NAME, + enforce_eager=True, + load_format="dummy", + tensor_parallel_size=1, + weight_transfer_config=WeightTransferConfig(backend="nccl"), + ) + + world_size = llm.llm_engine.vllm_config.parallel_config.world_size + assert world_size == 1 + + +@create_new_process_for_each_test() +def test_init_weight_transfer_engine_calls_engine(): + """Test that init_weight_transfer_engine calls the engine's + init_transfer_engine method.""" + if torch.cuda.device_count() < 1: + pytest.skip("Need at least 1 GPU for this test") + + # Enable insecure serialization to allow pickling functions for collective_rpc + os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1" + + with patch( + "vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine", + mock_create_engine, + ): + llm = LLM( + model=MODEL_NAME, + enforce_eager=True, + load_format="dummy", + tensor_parallel_size=1, + weight_transfer_config=WeightTransferConfig(backend="nccl"), + ) + + # Verify engine was created + def check_engine_exists(self): + return self.weight_transfer_engine is not None + + results = llm.collective_rpc(check_engine_exists) + assert all(results), "Weight transfer engine should be initialized" + + # Call init_weight_transfer_engine + llm.init_weight_transfer_engine( + WeightTransferInitRequest(init_info={"test_param": "hello"}) + ) + + # Verify init_transfer_engine was called on the engine + def check_init_called(self): + engine = self.weight_transfer_engine + return ( + engine.init_transfer_engine_called, + engine.last_init_info.test_param if engine.last_init_info else None, + ) + + results = llm.collective_rpc(check_init_called) + for called, param in results: + assert called, "init_transfer_engine should have been called" + assert param == "hello", f"Expected 'hello', got {param}" + + +@create_new_process_for_each_test() +def test_update_weights_calls_engine(): + """Test that update_weights calls the engine's receive_weights method.""" + if torch.cuda.device_count() < 1: + pytest.skip("Need at least 1 GPU for this test") + + # Enable insecure serialization to allow pickling functions for collective_rpc + os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1" + + with patch( + "vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine", + mock_create_engine, + ): + llm = LLM( + model=MODEL_NAME, + enforce_eager=True, + load_format="dummy", + tensor_parallel_size=1, + weight_transfer_config=WeightTransferConfig(backend="nccl"), + ) + + # First init the weight transfer + llm.init_weight_transfer_engine( + WeightTransferInitRequest(init_info={"test_param": "init"}) + ) + + # Call update_weights + test_names = ["layer.weight", "layer.bias"] + test_dtypes = ["float32", "float32"] + test_shapes = [[10, 10], [10]] + + llm.update_weights( + WeightTransferUpdateRequest( + update_info={ + "names": test_names, + "dtype_names": test_dtypes, + "shapes": test_shapes, + } + ) + ) + + # Verify receive_weights was called with correct info + def check_update_called(self): + engine = self.weight_transfer_engine + if not engine.receive_weights_called: + return False, None, None, None + info = engine.last_update_info + return (True, info.names, info.dtype_names, info.shapes) + + results = llm.collective_rpc(check_update_called) + for called, names, dtypes, shapes in results: + assert called, "receive_weights should have been called" + assert names == test_names + assert dtypes == test_dtypes + assert shapes == test_shapes + + +@create_new_process_for_each_test() +def test_full_weight_transfer_flow(): + """Test the complete weight transfer flow: init -> update.""" + if torch.cuda.device_count() < 1: + pytest.skip("Need at least 1 GPU for this test") + + # Enable insecure serialization to allow pickling functions for collective_rpc + os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1" + + with patch( + "vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine", + mock_create_engine, + ): + llm = LLM( + model=MODEL_NAME, + enforce_eager=True, + load_format="dummy", + tensor_parallel_size=1, + weight_transfer_config=WeightTransferConfig(backend="nccl"), + ) + + # Step 1: Initialize + llm.init_weight_transfer_engine( + WeightTransferInitRequest(init_info={"test_param": "flow_test"}) + ) + + # Step 2: Update weights + llm.update_weights( + WeightTransferUpdateRequest( + update_info={ + "names": ["test.weight"], + "dtype_names": ["bfloat16"], + "shapes": [[100, 100]], + } + ) + ) + + # Verify the full flow completed + def check_flow(self): + engine = self.weight_transfer_engine + return { + "init_called": engine.init_transfer_engine_called, + "update_called": engine.receive_weights_called, + "init_param": ( + engine.last_init_info.test_param if engine.last_init_info else None + ), + "update_names": ( + engine.last_update_info.names if engine.last_update_info else None + ), + } + + results = llm.collective_rpc(check_flow) + for result in results: + assert result["init_called"], "init_transfer_engine should be called" + assert result["update_called"], "receive_weights should be called" + assert result["init_param"] == "flow_test" + assert result["update_names"] == ["test.weight"] + + +@create_new_process_for_each_test() +def test_weight_transfer_config_backend(): + """Test that WeightTransferConfig backend is properly configured.""" + if torch.cuda.device_count() < 1: + pytest.skip("Need at least 1 GPU for this test") + + # Test with nccl backend + llm = LLM( + model=MODEL_NAME, + enforce_eager=True, + load_format="dummy", + tensor_parallel_size=1, + weight_transfer_config=WeightTransferConfig(backend="nccl"), + ) + + config = llm.llm_engine.vllm_config.weight_transfer_config + assert config.backend == "nccl" diff --git a/vllm/config/__init__.py b/vllm/config/__init__.py index b2044c6e1d0..6014f642c57 100644 --- a/vllm/config/__init__.py +++ b/vllm/config/__init__.py @@ -47,6 +47,7 @@ from vllm.config.vllm import ( get_layers_from_vllm_config, set_current_vllm_config, ) +from vllm.config.weight_transfer import WeightTransferConfig # __all__ should only contain classes and functions. # Types and globals should be imported from their respective modules. @@ -111,4 +112,5 @@ __all__ = [ "get_current_vllm_config_or_none", "set_current_vllm_config", "get_layers_from_vllm_config", + "WeightTransferConfig", ] diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 93d88730ea1..4d34c5584ba 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -42,6 +42,7 @@ from .scheduler import SchedulerConfig from .speculative import EagleModelTypes, SpeculativeConfig from .structured_outputs import StructuredOutputsConfig from .utils import SupportsHash, config, replace +from .weight_transfer import WeightTransferConfig if TYPE_CHECKING: from transformers import PretrainedConfig @@ -255,6 +256,9 @@ class VllmConfig: performance. -02 is used by defult. See OptimizationLevel for full description.""" + weight_transfer_config: WeightTransferConfig | None = None + """The configurations for weight transfer during RL training.""" + def compute_hash(self) -> str: """ WARNING: Whenever a new field is added to this config, diff --git a/vllm/config/weight_transfer.py b/vllm/config/weight_transfer.py new file mode 100644 index 00000000000..7ccac13fbfa --- /dev/null +++ b/vllm/config/weight_transfer.py @@ -0,0 +1,15 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from dataclasses import dataclass +from typing import Literal + +from vllm.config.utils import config + + +@config +@dataclass +class WeightTransferConfig: + """Configuration for weight transfer during RL training.""" + + backend: Literal["nccl"] = "nccl" + """The backend to use for weight transfer.""" diff --git a/vllm/distributed/weight_transfer/__init__.py b/vllm/distributed/weight_transfer/__init__.py new file mode 100644 index 00000000000..c96ad0e3bb4 --- /dev/null +++ b/vllm/distributed/weight_transfer/__init__.py @@ -0,0 +1,12 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Weight transfer engines for syncing model weights from trainers +to inference workers. +""" + +from vllm.distributed.weight_transfer.factory import WeightTransferEngineFactory + +__all__ = [ + "WeightTransferEngineFactory", +] diff --git a/vllm/distributed/weight_transfer/base.py b/vllm/distributed/weight_transfer/base.py new file mode 100644 index 00000000000..b87f190fcf7 --- /dev/null +++ b/vllm/distributed/weight_transfer/base.py @@ -0,0 +1,158 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Base class for weight transfer engines.""" + +from abc import ABC, abstractmethod +from collections.abc import Callable +from dataclasses import KW_ONLY, dataclass, field +from typing import Any, Generic, TypeVar + +import torch + +from vllm.config.parallel import ParallelConfig +from vllm.config.weight_transfer import WeightTransferConfig + +TInitInfo = TypeVar("TInitInfo", bound="WeightTransferInitInfo") +TUpdateInfo = TypeVar("TUpdateInfo", bound="WeightTransferUpdateInfo") + + +# Base protocols for backend-specific dataclasses +@dataclass +class WeightTransferInitInfo(ABC): # noqa: B024 + """Base class for backend-specific initialization info.""" + + pass + + +@dataclass +class WeightTransferUpdateInfo(ABC): # noqa: B024 + """Base class for backend-specific weight update info.""" + + _: KW_ONLY + is_checkpoint_format: bool = True + """Set to True if weights are in checkpoint/original model format and need + layerwise processing. Set to False if weights have already been processed + into kernel format (repacking, renaming, etc.).""" + + +# API-level request classes (accept dicts for backend-agnostic serialization) +@dataclass +class WeightTransferInitRequest: + """API-level weight transfer initialization request.""" + + init_info: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class WeightTransferUpdateRequest: + """API-level weight update request.""" + + update_info: dict[str, Any] = field(default_factory=dict) + + +class WeightTransferEngine(ABC, Generic[TInitInfo, TUpdateInfo]): + """ + Base class for weight transfer engines that handle transport of model weights + from a trainer to inference workers. + + This abstraction separates weight transfer transport logic from the worker + implementation, allowing different backends (NCCL, CUDA IPC[TODO], RDMA[TODO]) to be + plugged in. + + Subclasses should define: + init_info_cls: Type of backend-specific initialization info + update_info_cls: Type of backend-specific update info + """ + + # Subclasses should override these class attributes + init_info_cls: type[TInitInfo] + update_info_cls: type[TUpdateInfo] + + def __init__( + self, config: WeightTransferConfig, parallel_config: ParallelConfig + ) -> None: + """ + Initialize the weight transfer engine. + + Args: + config: The configuration for the weight transfer engine + parallel_config: The configuration for the parallel setup + """ + self.config = config + self.parallel_config = parallel_config + + def parse_init_info(self, init_dict: dict[str, Any]) -> TInitInfo: + """ + Construct typed init info from dict with validation. + + Args: + init_dict: Dictionary containing backend-specific initialization parameters + + Returns: + Typed backend-specific init info dataclass + + Raises: + ValueError: If init_dict is invalid for this backend + """ + try: + return self.init_info_cls(**init_dict) + except TypeError as e: + raise ValueError( + f"Invalid init_info for {self.__class__.__name__}: {e}" + ) from e + + def parse_update_info(self, update_dict: dict[str, Any]) -> TUpdateInfo: + """ + Construct typed update info from dict with validation. + + Args: + update_dict: Dictionary containing backend-specific update parameters + + Returns: + Typed backend-specific update info dataclass + + Raises: + ValueError: If update_dict is invalid for this backend + """ + try: + return self.update_info_cls(**update_dict) + except TypeError as e: + raise ValueError( + f"Invalid update_info for {self.__class__.__name__}: {e}" + ) from e + + @abstractmethod + def init_transfer_engine(self, init_info: TInitInfo) -> None: + """ + Initialize the weight transfer mechanism. + This is called once at the beginning of training. + + Args: + init_info: Backend-specific initialization info + """ + raise NotImplementedError + + @abstractmethod + def receive_weights( + self, + update_info: TUpdateInfo, + load_weights: Callable[[list[tuple[str, torch.Tensor]]], None], + ) -> None: + """ + Receive weights from the trainer and load them incrementally. + + Args: + update_info: Backend-specific update info containing parameter metadata + and any backend-specific data + load_weights: Callable that loads weights into the model. Called + incrementally for each weight to avoid OOM. + """ + raise NotImplementedError + + @abstractmethod + def shutdown(self) -> None: + """ + Shutdown the weight transfer engine. + This should be called when the worker is shutting down. + """ + raise NotImplementedError diff --git a/vllm/distributed/weight_transfer/factory.py b/vllm/distributed/weight_transfer/factory.py new file mode 100644 index 00000000000..7235e30d1af --- /dev/null +++ b/vllm/distributed/weight_transfer/factory.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Factory for weight transfer engines with lazy loading.""" + +import importlib +from collections.abc import Callable +from typing import TYPE_CHECKING + +from vllm.distributed.weight_transfer.base import WeightTransferEngine +from vllm.logger import init_logger + +if TYPE_CHECKING: + from vllm.config.parallel import ParallelConfig + from vllm.config.weight_transfer import WeightTransferConfig + +logger = init_logger(__name__) + + +class WeightTransferEngineFactory: + """Factory for creating weight transfer engines with lazy loading. + + This factory implements a registry pattern that supports: + - Lazy loading: Engine modules are only imported when actually needed + - Extensibility: Custom engines can be registered at runtime + - Centralized registration: All built-in engines registered in one place + """ + + _registry: dict[str, Callable[[], type[WeightTransferEngine]]] = {} + + @classmethod + def register_engine( + cls, + name: str, + module_path_or_cls: str | type[WeightTransferEngine], + class_name: str | None = None, + ) -> None: + """Register an engine with lazy-loading or direct class reference. + + Supports two calling conventions: + 1. Lazy loading: register_engine(name, module_path, class_name) + 2. Direct class: register_engine(name, engine_cls) + + Args: + name: The name to register the engine under (e.g., "nccl") + module_path_or_cls: Either a module path string for lazy loading, + or the engine class directly + class_name: Name of the engine class (required if module_path is string) + + Raises: + ValueError: If an engine with the same name is already registered + """ + if name in cls._registry: + raise ValueError(f"Weight transfer engine '{name}' is already registered.") + + if isinstance(module_path_or_cls, str): + # Lazy loading path + module_path = module_path_or_cls + if class_name is None: + raise ValueError( + "class_name is required when registering with module path" + ) + + def loader() -> type[WeightTransferEngine]: + module = importlib.import_module(module_path) + return getattr(module, class_name) + + cls._registry[name] = loader + else: + # Direct class registration + engine_cls = module_path_or_cls + cls._registry[name] = lambda: engine_cls + + @classmethod + def create_engine( + cls, + config: "WeightTransferConfig", + parallel_config: "ParallelConfig", + ) -> WeightTransferEngine: + """Create a weight transfer engine instance. + + Args: + config: Weight transfer configuration containing the backend name + parallel_config: Parallel configuration for the engine + + Returns: + An initialized weight transfer engine instance + + Raises: + ValueError: If the backend is not registered + """ + backend = config.backend + if backend not in cls._registry: + available = list(cls._registry.keys()) + raise ValueError( + f"Invalid weight transfer backend: {backend}. " + f"Available engines: {available}" + ) + engine_cls = cls._registry[backend]() + + logger.info( + "Creating weight transfer engine: %s", + engine_cls.__name__, + ) + + return engine_cls(config, parallel_config) + + +# Register built-in weight transfer engines here. +# Registration should be centralized to ensure lazy loading - +# engine modules are only imported when actually used. + +WeightTransferEngineFactory.register_engine( + "nccl", + "vllm.distributed.weight_transfer.nccl_engine", + "NCCLWeightTransferEngine", +) diff --git a/vllm/distributed/weight_transfer/nccl_engine.py b/vllm/distributed/weight_transfer/nccl_engine.py new file mode 100644 index 00000000000..5c90198bf61 --- /dev/null +++ b/vllm/distributed/weight_transfer/nccl_engine.py @@ -0,0 +1,315 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NCCL-based weight transfer engine.""" + +from collections.abc import Callable, Iterator +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +import torch + +if TYPE_CHECKING: + from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator + +from vllm.config.parallel import ParallelConfig +from vllm.config.weight_transfer import WeightTransferConfig +from vllm.distributed.weight_transfer.base import ( + WeightTransferEngine, + WeightTransferInitInfo, + WeightTransferUpdateInfo, +) +from vllm.distributed.weight_transfer.packed_tensor import ( + DEFAULT_PACKED_BUFFER_SIZE_BYTES, + DEFAULT_PACKED_NUM_BUFFERS, + packed_broadcast_consumer, +) + + +@dataclass +class NCCLWeightTransferInitInfo(WeightTransferInitInfo): + """Initialization info for NCCL weight transfer backend.""" + + master_address: str + master_port: int + rank_offset: int + world_size: int + + +@dataclass +class NCCLWeightTransferUpdateInfo(WeightTransferUpdateInfo): + """Update info for NCCL weight transfer backend.""" + + names: list[str] + dtype_names: list[str] + shapes: list[list[int]] + packed: bool = False + """Whether to use packed tensor broadcasting for efficiency. + When True, multiple tensors are batched together before broadcasting + to reduce NCCL communication overhead.""" + packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES + """Size in bytes for each packed tensor buffer. Default is 1GB. + Both producer and consumer must use the same value.""" + packed_num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS + """Number of buffers for double/triple buffering during packed transfer. + Both producer and consumer must use the same value.""" + + def __post_init__(self): + """Validate that all lists have the same length.""" + num_params = len(self.names) + if len(self.dtype_names) != num_params: + raise ValueError( + f"`dtype_names` should be of the same size as `names`: " + f"got {len(self.dtype_names)} and {len(self.names)}" + ) + if len(self.shapes) != num_params: + raise ValueError( + f"`shapes` should be of the same size as `names`: " + f"got {len(self.shapes)} and {len(self.names)}" + ) + + +class NCCLWeightTransferEngine( + WeightTransferEngine[NCCLWeightTransferInitInfo, NCCLWeightTransferUpdateInfo] +): + """ + Weight transfer engine using NCCL for communication between trainer and workers. + + This implementation uses NCCL broadcast operations to transfer weights from + the trainer (rank 0) to all inference workers in a process group. + """ + + # Define backend-specific dataclass types + init_info_cls = NCCLWeightTransferInitInfo + update_info_cls = NCCLWeightTransferUpdateInfo + + def __init__( + self, config: WeightTransferConfig, parallel_config: ParallelConfig + ) -> None: + """ + Initialize the NCCL weight transfer engine. + + Args: + config: The configuration for the weight transfer engine + parallel_config: The configuration for the parallel setup + """ + super().__init__(config, parallel_config) + self.model_update_group: PyNcclCommunicator | None = None + + def init_transfer_engine(self, init_info: NCCLWeightTransferInitInfo) -> None: + """ + Initialize NCCL process group with the trainer. + + Args: + init_info: NCCL initialization info containing master address, port, + rank offset, and world size + """ + + # Calculate the global rank in the trainer-worker process group + # Must account for data parallel to get unique ranks across all workers + dp_rank = self.parallel_config.data_parallel_rank + world_size_per_dp = self.parallel_config.world_size # TP * PP + rank_within_dp = self.parallel_config.rank + + # Unique rank across all DP groups + worker_rank = dp_rank * world_size_per_dp + rank_within_dp + rank = worker_rank + init_info.rank_offset + # Create stateless process group + self.model_update_group = ( + NCCLWeightTransferEngine._stateless_init_process_group( + init_info.master_address, + init_info.master_port, + rank, + init_info.world_size, + torch.cuda.current_device(), + ) + ) + + def receive_weights( + self, + update_info: NCCLWeightTransferUpdateInfo, + load_weights: Callable[[list[tuple[str, torch.Tensor]]], None], + ) -> None: + """ + Receive weights from trainer via NCCL broadcast and load them incrementally. + + If update_info.packed is True, uses packed tensor broadcasting for + efficient transfer of multiple weights in batches. Otherwise, uses simple + one-by-one broadcasting. + + Args: + update_info: NCCL update info containing parameter names, dtypes, shapes, + and packed flag + load_weights: Callable that loads weights into the model. Called + incrementally for each batch of weights to avoid OOM. + """ + if self.model_update_group is None: + raise RuntimeError( + "NCCL weight transfer not initialized. " + "Call init_transfer_engine() first." + ) + + if update_info.packed: + # Build iterator of (name, (shape, dtype)) from update_info + def state_dict_info_iterator(): + for name, dtype_name, shape in zip( + update_info.names, update_info.dtype_names, update_info.shapes + ): + dtype = getattr(torch, dtype_name) + yield (name, (shape, dtype)) + + packed_broadcast_consumer( + iterator=state_dict_info_iterator(), + group=self.model_update_group, + src=0, + post_unpack_func=load_weights, + buffer_size_bytes=update_info.packed_buffer_size_bytes, + num_buffers=update_info.packed_num_buffers, + ) + else: + # Use simple one-by-one broadcasting + for name, dtype_name, shape in zip( + update_info.names, update_info.dtype_names, update_info.shapes + ): + dtype = getattr(torch, dtype_name) + weight = torch.empty(shape, dtype=dtype, device="cuda") + self.model_update_group.broadcast( + weight, src=0, stream=torch.cuda.current_stream() + ) + load_weights([(name, weight)]) + del weight + + def shutdown(self) -> None: + if self.model_update_group is not None: + # Clean up the communicator by removing the reference + self.model_update_group = None + + @staticmethod + def trainer_send_weights( + iterator: Iterator[tuple[str, torch.Tensor]], + group: Any, + src: int = 0, + post_iter_func: Callable[[tuple[str, torch.Tensor]], torch.Tensor] + | None = None, + packed: bool = False, + stream: torch.cuda.Stream | None = None, + packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES, + packed_num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS, + ) -> None: + """Broadcast weights from trainer to vLLM workers. + + Args: + iterator: Iterator of model parameters. Returns (name, tensor) tuples + group: Process group (PyNcclCommunicator) + src: Source rank (default 0, trainer is typically rank 0) + post_iter_func: Optional function to apply to each (name, tensor) pair + before broadcasting. If None, extracts just the tensor. + packed: Whether to use packed tensor broadcasting for efficiency. + When True, multiple tensors are batched together before + broadcasting to reduce NCCL communication overhead. + stream: CUDA stream to use for broadcasting if packed is False. + If packed is True, new streams will be created for each buffer. + packed_buffer_size_bytes: Size in bytes for each packed tensor buffer. + Must match the value used in NCCLWeightTransferUpdateInfo. + packed_num_buffers: Number of buffers for double/triple buffering. + Must match the value used in NCCLWeightTransferUpdateInfo. + + Example: + >>> from vllm.distributed.weight_transfer.nccl_engine import ( + ... NCCLWeightTransferEngine, + ... ) + >>> param_iter = ((n, p) for n, p in model.named_parameters()) + >>> NCCLWeightTransferEngine.trainer_send_weights( + ... param_iter, group, packed=True + ... ) + """ + if post_iter_func is None: + # Default: extract just the tensor from (name, tensor) tuple + post_iter_func = lambda x: x[1] + + if packed: + # Use packed tensor broadcasting for efficiency + from vllm.distributed.weight_transfer.packed_tensor import ( + packed_broadcast_producer, + ) + + packed_broadcast_producer( + iterator=iterator, + group=group, + src=src, + post_iter_func=post_iter_func, + buffer_size_bytes=packed_buffer_size_bytes, + num_buffers=packed_num_buffers, + ) + else: + # Use simple one-by-one broadcasting + for item in iterator: + tensor = post_iter_func(item) + group.broadcast( + tensor, src=src, stream=stream or torch.cuda.current_stream() + ) + + @staticmethod + def trainer_init( + init_info: NCCLWeightTransferInitInfo | dict, + ) -> "PyNcclCommunicator": + """ + Initialize NCCL process group for trainer-side weight transfer. + + The trainer is always rank 0 in the process group. Uses the current + CUDA device (torch.cuda.current_device()). + + Args: + init_info: Either an NCCLWeightTransferInitInfo object or a dict with keys: + - master_address: str + - master_port: int + - world_size: int + + Returns: + PyNcclCommunicator for weight transfer. + + Example: + >>> from vllm.distributed.weight_transfer.nccl_engine import ( + ... NCCLWeightTransferEngine, + ... ) + >>> group = NCCLWeightTransferEngine.trainer_init( + ... dict( + ... master_address=master_address, + ... master_port=master_port, + ... world_size=world_size, + ... ), + ... ) + """ + if isinstance(init_info, dict): + master_address = init_info["master_address"] + master_port = init_info["master_port"] + world_size = init_info["world_size"] + else: + # NCCLWeightTransferInitInfo object + master_address = init_info.master_address + master_port = init_info.master_port + world_size = init_info.world_size + + # Trainer is always rank 0 + return NCCLWeightTransferEngine._stateless_init_process_group( + master_address, master_port, 0, world_size, torch.cuda.current_device() + ) + + @staticmethod + def _stateless_init_process_group( + master_address, master_port, rank, world_size, device + ): + """ + vLLM provides `StatelessProcessGroup` to create a process group + without considering the global process group in torch.distributed. + It is recommended to create `StatelessProcessGroup`, and then initialize + the data-plane communication (NCCL) between external (train processes) + and vLLM workers. + """ + from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator + from vllm.distributed.utils import StatelessProcessGroup + + pg = StatelessProcessGroup.create( + host=master_address, port=master_port, rank=rank, world_size=world_size + ) + pynccl = PyNcclCommunicator(pg, device=device) + return pynccl diff --git a/vllm/distributed/weight_transfer/packed_tensor.py b/vllm/distributed/weight_transfer/packed_tensor.py new file mode 100644 index 00000000000..1c96d72edac --- /dev/null +++ b/vllm/distributed/weight_transfer/packed_tensor.py @@ -0,0 +1,216 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Packed tensor utilities for efficient weight transfer.""" + +import math +from collections.abc import Callable, Iterator +from typing import Any + +import torch + +# Default values for packed tensor configuration. +# These are imported by NCCLWeightTransferUpdateInfo and trainer_send_weights. +DEFAULT_PACKED_BUFFER_SIZE_BYTES = 1024 * 1024 * 1024 # 1GB +DEFAULT_PACKED_NUM_BUFFERS = 2 + + +def packed_broadcast_producer( + iterator: Iterator[tuple[str, torch.Tensor]], + group: Any, + src: int, + post_iter_func: Callable[[tuple[str, torch.Tensor]], torch.Tensor], + buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES, + num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS, +) -> None: + """Broadcast tensors in a packed manner from trainer to workers. + + Args: + iterator: Iterator of model parameters. Returns a tuple of (name, tensor) + group: Process group (PyNcclCommunicator) + src: Source rank (0 in current implementation) + post_iter_func: Function to apply to each (name, tensor) pair before + packing, should return a tensor + buffer_size_bytes: Size in bytes for each packed tensor buffer. + Both producer and consumer must use the same value. + num_buffers: Number of buffers for double/triple buffering. + Both producer and consumer must use the same value. + + """ + target_packed_tensor_size = buffer_size_bytes + + streams = [torch.cuda.Stream() for _ in range(num_buffers)] + buffer_idx = 0 + + packing_tensor_list: list[list[torch.Tensor]] = [[] for _ in range(num_buffers)] + packing_tensor_sizes: list[int] = [0 for _ in range(num_buffers)] + packed_tensors: list[torch.Tensor] = [ + torch.empty(0, dtype=torch.uint8, device="cuda") for _ in range(num_buffers) + ] + + while True: + # Synchronize the current stream + streams[buffer_idx].synchronize() + # Start tasks for the new buffer in a new stream + with torch.cuda.stream(streams[buffer_idx]): + try: + # Initialize the packing tensor list and sizes + packing_tensor_list[buffer_idx] = [] + packing_tensor_sizes[buffer_idx] = 0 + # Pack the tensors + while True: + # Apply post processing and convert to linearized uint8 tensor + tensor = ( + post_iter_func(next(iterator)) + .contiguous() + .view(torch.uint8) + .view(-1) + ) + packing_tensor_list[buffer_idx].append(tensor) + packing_tensor_sizes[buffer_idx] += tensor.numel() + if packing_tensor_sizes[buffer_idx] > target_packed_tensor_size: + break + # Pack the tensors and call broadcast collective + packed_tensors[buffer_idx] = torch.cat( + packing_tensor_list[buffer_idx], dim=0 + ) + group.broadcast(packed_tensors[buffer_idx], src=src) + # Move to the next buffer + buffer_idx = (buffer_idx + 1) % num_buffers + except StopIteration: + # Do the last broadcast if there are remaining tensors + if len(packing_tensor_list[buffer_idx]) > 0: + packed_tensors[buffer_idx] = torch.cat( + packing_tensor_list[buffer_idx], dim=0 + ) + group.broadcast(packed_tensors[buffer_idx], src=src) + break + + +def packed_broadcast_consumer( + iterator: Iterator[tuple[str, tuple[list[int], torch.dtype]]], + group: Any, + src: int, + post_unpack_func: Callable[[list[tuple[str, torch.Tensor]]], None], + buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES, + num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS, +) -> None: + """Consume packed tensors and unpack them into a list of tensors. + + Args: + iterator: Iterator of parameter metadata. Returns (name, (shape, dtype)) + group: Process group (PyNcclCommunicator) + src: Source rank (0 in current implementation) + post_unpack_func: Function to apply to each list of (name, tensor) after + unpacking + buffer_size_bytes: Size in bytes for each packed tensor buffer. + Both producer and consumer must use the same value. + num_buffers: Number of buffers for double/triple buffering. + Both producer and consumer must use the same value. + + """ + + def unpack_tensor( + packed_tensor: torch.Tensor, + names: list[str], + shapes: list[list[int]], + dtypes: list[torch.dtype], + tensor_sizes: list[int], + ) -> list[tuple[str, torch.Tensor]]: + """Unpack a single tensor into a list of tensors. + + Args: + packed_tensor: The packed torch.uint8 tensor to unpack + names: List of tensor names + shapes: List of tensor shapes + dtypes: List of tensor dtypes + tensor_sizes: List of tensor sizes in bytes + + Returns: + unpacked List[(name, tensor)] + """ + unpacked_tensors = packed_tensor.split(tensor_sizes) + + unpacked_list = [ + (name, tensor.contiguous().view(dtype).view(*shape)) + for name, shape, dtype, tensor in zip( + names, shapes, dtypes, unpacked_tensors + ) + ] + + return unpacked_list + + target_packed_tensor_size = buffer_size_bytes + + streams = [torch.cuda.Stream() for _ in range(num_buffers)] + buffer_idx = 0 + + packing_tensor_meta_data: list[list[tuple[str, list[int], torch.dtype, int]]] = [ + [] for _ in range(num_buffers) + ] + packing_tensor_sizes: list[int] = [0 for _ in range(num_buffers)] + packed_tensors: list[torch.Tensor] = [ + torch.empty(0, dtype=torch.uint8, device="cuda") for _ in range(num_buffers) + ] + + while True: + # Synchronize the current stream + streams[buffer_idx].synchronize() + with torch.cuda.stream(streams[buffer_idx]): + # Initialize the packing tensor meta data + packing_tensor_meta_data[buffer_idx] = [] + packing_tensor_sizes[buffer_idx] = 0 + try: + # Form a packed tensor + while True: + name, (shape, dtype) = next(iterator) + tensor_size = math.prod(shape) * dtype.itemsize + packing_tensor_meta_data[buffer_idx].append( + (name, shape, dtype, tensor_size) + ) + packing_tensor_sizes[buffer_idx] += tensor_size + if packing_tensor_sizes[buffer_idx] > target_packed_tensor_size: + break + # Create a packed tensor and broadcast it + packed_tensors[buffer_idx] = torch.empty( + packing_tensor_sizes[buffer_idx], dtype=torch.uint8, device="cuda" + ) + group.broadcast(packed_tensors[buffer_idx], src=src) + # Load the packed tensor into the model + names, shapes, dtypes, tensor_sizes = zip( + *packing_tensor_meta_data[buffer_idx] + ) + post_unpack_func( + unpack_tensor( + packed_tensors[buffer_idx], + list(names), + list(shapes), + list(dtypes), + list(tensor_sizes), + ) + ) + # Move to the next buffer + buffer_idx = (buffer_idx + 1) % num_buffers + except StopIteration: + # Do the last broadcast if there are remaining tensors + if len(packing_tensor_meta_data[buffer_idx]) > 0: + # Create a packed tensor and broadcast it + packed_tensors[buffer_idx] = torch.empty( + packing_tensor_sizes[buffer_idx], + dtype=torch.uint8, + device="cuda", + ) + group.broadcast(packed_tensors[buffer_idx], src=src) + # Load the packed tensor into the model + names, shapes, dtypes, tensor_sizes = zip( + *packing_tensor_meta_data[buffer_idx] + ) + post_unpack_func( + unpack_tensor( + packed_tensors[buffer_idx], + list(names), + list(shapes), + list(dtypes), + list(tensor_sizes), + ) + ) + break diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index f3e7729f64e..471516e3237 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -54,6 +54,7 @@ from vllm.config import ( SpeculativeConfig, StructuredOutputsConfig, VllmConfig, + WeightTransferConfig, get_attr_docs, ) from vllm.config.cache import ( @@ -581,6 +582,11 @@ class EngineArgs: kv_offloading_backend: KVOffloadingBackend = CacheConfig.kv_offloading_backend tokens_only: bool = False + weight_transfer_config: WeightTransferConfig | None = None + """Configuration for weight transfer during RL training. + Accepts a JSON string or dict with backend-specific options. + Example: '{"backend": "nccl"}'""" + def __post_init__(self): # support `EngineArgs(compilation_config={...})` # without having to manually construct a @@ -591,6 +597,10 @@ class EngineArgs: self.attention_config = AttentionConfig(**self.attention_config) if isinstance(self.eplb_config, dict): self.eplb_config = EPLBConfig(**self.eplb_config) + if isinstance(self.weight_transfer_config, dict): + self.weight_transfer_config = WeightTransferConfig( + **self.weight_transfer_config + ) # Setup plugins from vllm.plugins import load_general_plugins @@ -1189,6 +1199,9 @@ class EngineArgs: vllm_group.add_argument( "--optimization-level", **vllm_kwargs["optimization_level"] ) + vllm_group.add_argument( + "--weight-transfer-config", **vllm_kwargs["weight_transfer_config"] + ) # Other arguments parser.add_argument( @@ -1765,6 +1778,7 @@ class EngineArgs: profiler_config=self.profiler_config, additional_config=self.additional_config, optimization_level=self.optimization_level, + weight_transfer_config=self.weight_transfer_config, ) return config diff --git a/vllm/engine/protocol.py b/vllm/engine/protocol.py index 1502bbff4bf..253cfc42d67 100644 --- a/vllm/engine/protocol.py +++ b/vllm/engine/protocol.py @@ -6,6 +6,10 @@ from collections.abc import AsyncGenerator, Iterable, Mapping from typing import Any from vllm.config import ModelConfig, VllmConfig +from vllm.distributed.weight_transfer.base import ( + WeightTransferInitRequest, + WeightTransferUpdateRequest, +) from vllm.inputs.data import PromptType, StreamingInput from vllm.lora.request import LoRARequest from vllm.outputs import PoolingRequestOutput, RequestOutput @@ -191,3 +195,13 @@ class EngineClient(ABC): async def get_supported_tasks(self) -> tuple[SupportedTask, ...]: """Get supported tasks""" raise NotImplementedError + + async def init_weight_transfer_engine( + self, init_request: WeightTransferInitRequest + ) -> None: + """Initialize weight transfer for RL training.""" + raise NotImplementedError + + async def update_weights(self, request: WeightTransferUpdateRequest) -> None: + """Batched weight update for RL training.""" + raise NotImplementedError diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index fbcf3a43773..a7180d92842 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -34,6 +34,10 @@ from vllm.config.model import ( RunnerOption, TokenizerMode, ) +from vllm.distributed.weight_transfer.base import ( + WeightTransferInitRequest, + WeightTransferUpdateRequest, +) from vllm.engine.arg_utils import EngineArgs from vllm.entrypoints.chat_utils import ( ChatCompletionMessageParam, @@ -360,6 +364,23 @@ class LLM: def get_tokenizer(self) -> TokenizerLike: return self.llm_engine.get_tokenizer() + def get_world_size(self, include_dp: bool = True) -> int: + """Get the world size from the parallel config. + + Args: + include_dp: If True (default), returns the world size including + data parallelism (TP * PP * DP). If False, returns the world + size without data parallelism (TP * PP). + + Returns: + The world size (tensor_parallel_size * pipeline_parallel_size), + optionally multiplied by data_parallel_size if include_dp is True. + """ + parallel_config = self.llm_engine.vllm_config.parallel_config + if include_dp: + return parallel_config.world_size_across_dp + return parallel_config.world_size + def reset_mm_cache(self) -> None: self.input_processor.clear_mm_cache() self.llm_engine.reset_mm_cache() @@ -1903,6 +1924,38 @@ class LLM: # its previous requests. return sorted(outputs, key=lambda x: int(x.request_id)) + def init_weight_transfer_engine( + self, request: WeightTransferInitRequest | dict + ) -> None: + """ + Initialize weight transfer for RL training. + + Args: + request: Weight transfer initialization request with backend-specific info + """ + init_info_dict = ( + request["init_info"] if isinstance(request, dict) else request.init_info + ) + + self.llm_engine.collective_rpc( + "init_weight_transfer_engine", kwargs={"init_info": init_info_dict} + ) + + def update_weights(self, request: WeightTransferUpdateRequest | dict) -> None: + """ + Update the weights of the model. + + Args: + request: Weight update request with backend-specific update info + """ + update_info_dict = ( + request["update_info"] if isinstance(request, dict) else request.update_info + ) + + self.llm_engine.collective_rpc( + "update_weights", kwargs={"update_info": update_info_dict} + ) + def __repr__(self) -> str: """Return a transformers-style hierarchical view of the model.""" # Cache the result to avoid repeated collective_rpc calls diff --git a/vllm/entrypoints/serve/rlhf/api_router.py b/vllm/entrypoints/serve/rlhf/api_router.py index 3b37840ae08..38461b14778 100644 --- a/vllm/entrypoints/serve/rlhf/api_router.py +++ b/vllm/entrypoints/serve/rlhf/api_router.py @@ -1,12 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project - +import json from http import HTTPStatus -from fastapi import APIRouter, FastAPI, Query, Request +from fastapi import APIRouter, FastAPI, HTTPException, Query, Request from fastapi.responses import JSONResponse +import vllm.envs as envs +from vllm.distributed.weight_transfer.base import ( + WeightTransferInitRequest, + WeightTransferUpdateRequest, +) from vllm.engine.protocol import EngineClient from vllm.logger import init_logger @@ -98,5 +103,63 @@ async def is_paused(raw_request: Request) -> JSONResponse: return JSONResponse(content={"is_paused": paused}) +@router.post("/init_weight_transfer_engine") +async def init_weight_transfer_engine(raw_request: Request): + try: + body = await raw_request.json() + except json.JSONDecodeError as e: + raise HTTPException(status_code=400, detail="Invalid JSON format") from e # noqa: B904 + init_info = body.get("init_info") + if init_info is None: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail="Missing 'init_info' in request body", + ) + await engine_client(raw_request).init_weight_transfer_engine( + WeightTransferInitRequest(init_info=init_info) + ) + return JSONResponse(content={"message": "Weight transfer initialized"}) + + +@router.post("/update_weights") +async def update_weights(raw_request: Request): + try: + body = await raw_request.json() + except json.JSONDecodeError as e: + raise HTTPException(status_code=400, detail="Invalid JSON format") from e # noqa: B904 + update_info = body.get("update_info") + if update_info is None: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail="Missing 'update_info' in request body", + ) + await engine_client(raw_request).update_weights( + request=WeightTransferUpdateRequest(update_info=update_info) + ) + return JSONResponse(content={"message": "Weights updated"}) + + +@router.get("/get_world_size") +async def get_world_size( + raw_request: Request, + include_dp: bool = Query(True), +): + """Get the world size from the parallel config. + + Args: + include_dp: If True (default), returns the world size including + data parallelism (TP * PP * DP). If False, returns the world + size without data parallelism (TP * PP). + """ + parallel_config = engine_client(raw_request).vllm_config.parallel_config + if include_dp: + world_size = parallel_config.world_size_across_dp + else: + world_size = parallel_config.world_size + return JSONResponse(content={"world_size": world_size}) + + def attach_router(app: FastAPI): + if not envs.VLLM_SERVER_DEV_MODE: + return app.include_router(router) diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 53bdb972bcd..8b9fe0f3e93 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -649,6 +649,9 @@ class Fp8OnlineLinearMethod(Fp8LinearMethod): ) # Activations not quantized for marlin. + # Prevent duplicate processing (e.g., during weight reload) + layer._already_called_process_weights_after_loading = True + class Fp8MoEMethod(FusedMoEMethodBase): """MoE method for FP8. @@ -908,6 +911,9 @@ class Fp8MoEMethod(FusedMoEMethodBase): layer, w13, w2, w13_scale, w2_scale, w13_input_scale, w2_input_scale ) + # Prevent duplicate processing (e.g., during weight reload) + layer._already_called_process_weights_after_loading = True + def maybe_make_prepare_finalize( self, routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, @@ -1241,6 +1247,9 @@ class Fp8OnlineMoEMethod(Fp8MoEMethod): layer.w2_input_scale, ) + # Prevent duplicate processing (e.g., during weight reload) + layer._already_called_process_weights_after_loading = True + class Fp8KVCacheMethod(BaseKVCacheMethod): """ diff --git a/vllm/model_executor/model_loader/reload/layerwise.py b/vllm/model_executor/model_loader/reload/layerwise.py index f7aaf8a677c..21795e63995 100644 --- a/vllm/model_executor/model_loader/reload/layerwise.py +++ b/vllm/model_executor/model_loader/reload/layerwise.py @@ -216,6 +216,11 @@ def _layerwise_process(layer: torch.nn.Module, info: LayerReloadingInfo): # Materialize layer tensors onto device materialize_layer(layer) + # Reset FP8 online quantization flag so process_weights_after_loading + # will run again during reload + if hasattr(layer, "_already_called_process_weights_after_loading"): + delattr(layer, "_already_called_process_weights_after_loading") + # Unwrap layerwise loading wrappers for param in get_layer_tensors(layer).values(): param.weight_loader = _get_original_loader(param) diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index c0613137091..43d63bcff25 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -14,6 +14,10 @@ import torch import vllm.envs as envs from vllm import TokensPrompt from vllm.config import VllmConfig +from vllm.distributed.weight_transfer.base import ( + WeightTransferInitRequest, + WeightTransferUpdateRequest, +) from vllm.engine.arg_utils import AsyncEngineArgs from vllm.engine.protocol import EngineClient from vllm.inputs import PromptType, StreamingInput @@ -1011,3 +1015,44 @@ class AsyncLLM(EngineClient): @property def dead_error(self) -> BaseException: return EngineDeadError() + + async def init_weight_transfer_engine( + self, request: WeightTransferInitRequest + ) -> None: + """ + Initialize weight transfer for RL training. + + Args: + request: Weight transfer initialization request with backend-specific info + """ + from vllm.distributed.weight_transfer.base import ( + WeightTransferInitRequest, + ) + + if isinstance(request, WeightTransferInitRequest): + init_info_dict = request.init_info + else: + raise TypeError(f"Expected WeightTransferInitRequest, got {type(request)}") + + await self.collective_rpc( + "init_weight_transfer_engine", kwargs={"init_info": init_info_dict} + ) + + async def update_weights(self, request: WeightTransferUpdateRequest) -> None: + """ + Batched weight update for RL training. + + Args: + request: Weight update request with backend-specific update info + """ + + if isinstance(request, WeightTransferUpdateRequest): + update_info_dict = request.update_info + else: + raise TypeError( + f"Expected WeightTransferUpdateRequest, got {type(request)}" + ) + + await self.collective_rpc( + "update_weights", kwargs={"update_info": update_info_dict} + ) diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index b451db3826f..09880f79bf1 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -33,6 +33,7 @@ from vllm.distributed.parallel_state import ( get_pp_group, get_tp_group, ) +from vllm.distributed.weight_transfer import WeightTransferEngineFactory from vllm.logger import init_logger from vllm.lora.request import LoRARequest from vllm.model_executor.models.interfaces import is_mixture_of_experts @@ -89,6 +90,16 @@ class Worker(WorkerBase): # Buffers saved before sleep self._sleep_saved_buffers: dict[str, torch.Tensor] = {} + # Weight transfer engine (initialized on-demand) + self.weight_transfer_engine = ( + WeightTransferEngineFactory.create_engine( + self.vllm_config.weight_transfer_config, + self.vllm_config.parallel_config, + ) + if self.vllm_config.weight_transfer_config is not None + else None + ) + # Torch/CUDA profiler. Enabled and configured through profiler_config. self.profiler: Any | None = None profiler_config = vllm_config.profiler_config @@ -932,6 +943,69 @@ class Worker(WorkerBase): tensorizer_config=tensorizer_config, ) + def init_weight_transfer_engine(self, init_info: dict) -> None: + """ + Initialize weight transfer mechanism. + For NCCL backend, this creates a process group with the trainer. + + Args: + init_info: Dictionary containing backend-specific initialization info + """ + if self.weight_transfer_engine is None: + raise RuntimeError( + "Weight transfer not configured. " + "Please set weight_transfer_config to enable weight transfer." + ) + # Parse dict into backend-specific typed dataclass + typed_init_info = self.weight_transfer_engine.parse_init_info(init_info) + self.weight_transfer_engine.init_transfer_engine(typed_init_info) + + def update_weights(self, update_info: dict) -> None: + """ + Batched weight update from the trainer. + + Args: + update_info: Dictionary containing backend-specific update info + """ + if self.weight_transfer_engine is None: + raise RuntimeError( + "Weight transfer not configured. " + "Please set weight_transfer_config to enable weight transfer." + ) + + # Parse dict into backend-specific typed dataclass + typed_update_info = self.weight_transfer_engine.parse_update_info(update_info) + + model = self.model_runner.model + + if typed_update_info.is_checkpoint_format: + from vllm.model_executor.model_loader.reload import ( + finalize_layerwise_reload, + initialize_layerwise_reload, + ) + + # Use layerwise reload pattern for checkpoint format weights + with torch.device(self.device): + initialize_layerwise_reload(model) + self.weight_transfer_engine.receive_weights( + typed_update_info, + load_weights=model.load_weights, + ) + finalize_layerwise_reload(model, self.model_config) + else: + # Weights are already in kernel format, copy directly + def load_weights_direct( + weights: list[tuple[str, torch.Tensor]], + ) -> None: + for name, weight in weights: + param = model.get_parameter(name) + param.copy_(weight) + + self.weight_transfer_engine.receive_weights( + typed_update_info, + load_weights=load_weights_direct, + ) + def shutdown(self) -> None: # has_kv_transfer_group can be None during interpreter shutdown. if ensure_kv_transfer_shutdown is not None: @@ -939,6 +1013,9 @@ class Worker(WorkerBase): if self.profiler is not None: self.profiler.shutdown() + if weight_transfer_engine := getattr(self, "weight_transfer_engine", None): + weight_transfer_engine.shutdown() + def init_worker_distributed_environment( vllm_config: VllmConfig, From 5b2a9422f0f4cbdd69a9fed1dc1605838314ff81 Mon Sep 17 00:00:00 2001 From: danisereb Date: Thu, 5 Feb 2026 19:25:55 +0200 Subject: [PATCH 102/810] [BugFix] Fix LoRA Fp8 (#33879) Signed-off-by: Daniel Serebrenik --- vllm/lora/layers/fused_moe.py | 22 ++++++++++++++-------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py index c2b35fbb105..4d4e053cffd 100644 --- a/vllm/lora/layers/fused_moe.py +++ b/vllm/lora/layers/fused_moe.py @@ -130,14 +130,20 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA): self.base_layer.ensure_moe_quant_config_init() quant_config = self.base_layer.quant_method.moe_quant_config - prepare_finalize = MoEPrepareAndFinalizeNoEP() - m_fused_moe_fn = FusedMoEModularKernel( - prepare_finalize, - self.base_layer.quant_method.select_gemm_impl( - prepare_finalize, self.base_layer - ), - self.base_layer.shared_experts, - ) + if getattr(self.base_layer.quant_method, "supports_internal_mk", False): + # Use the existing modular kernel from the quant method + m_fused_moe_fn = self.base_layer.quant_method.moe_mk + else: + # Create a new modular kernel via select_gemm_impl + prepare_finalize = MoEPrepareAndFinalizeNoEP() + m_fused_moe_fn = FusedMoEModularKernel( + prepare_finalize, + self.base_layer.quant_method.select_gemm_impl( + prepare_finalize, self.base_layer + ), + self.base_layer.shared_experts, + ) + if quant_config.use_mxfp4_w4a16: assert isinstance( m_fused_moe_fn.fused_experts, (MarlinExperts, UnfusedOAITritonExperts) From af3162d3aaa559a738396baf5b5134c1ab0742f5 Mon Sep 17 00:00:00 2001 From: Benjamin Chislett Date: Thu, 5 Feb 2026 12:37:18 -0500 Subject: [PATCH 103/810] [Spec Decode] Unified Parallel Drafting (#32887) Signed-off-by: Benjamin Chislett --- examples/offline_inference/spec_decode.py | 3 + tests/v1/e2e/test_spec_decode.py | 95 ++--- tests/v1/spec_decode/test_eagle.py | 414 ++++++++++++++++++++- tests/v1/spec_decode/test_mtp.py | 2 +- vllm/config/speculative.py | 7 + vllm/config/vllm.py | 28 +- vllm/model_executor/models/llama_eagle3.py | 42 ++- vllm/v1/attention/backend.py | 7 +- vllm/v1/attention/backends/flashinfer.py | 36 +- vllm/v1/attention/backends/utils.py | 32 -- vllm/v1/spec_decode/draft_model.py | 245 ++---------- vllm/v1/spec_decode/eagle.py | 316 +++++++++++++--- vllm/v1/spec_decode/utils.py | 248 ++++++++++++ vllm/v1/worker/gpu_model_runner.py | 2 +- 14 files changed, 1085 insertions(+), 392 deletions(-) diff --git a/examples/offline_inference/spec_decode.py b/examples/offline_inference/spec_decode.py index 45593b53061..d8c5ece4fa6 100644 --- a/examples/offline_inference/spec_decode.py +++ b/examples/offline_inference/spec_decode.py @@ -75,6 +75,7 @@ def parse_args(): parser.add_argument("--gpu-memory-utilization", type=float, default=0.9) parser.add_argument("--disable-padded-drafter-batch", action="store_true") parser.add_argument("--max-num-seqs", type=int, default=None) + parser.add_argument("--parallel-drafting", action="store_true") parser.add_argument("--allowed-local-media-path", type=str, default="") return parser.parse_args() @@ -121,6 +122,7 @@ def main(args): "model": eagle_dir, "num_speculative_tokens": args.num_spec_tokens, "disable_padded_drafter_batch": args.disable_padded_drafter_batch, + "parallel_drafting": args.parallel_drafting, } elif args.method == "ngram": speculative_config = { @@ -137,6 +139,7 @@ def main(args): "num_speculative_tokens": args.num_spec_tokens, "enforce_eager": args.enforce_eager, "max_model_len": args.max_model_len, + "parallel_drafting": args.parallel_drafting, } elif args.method == "mtp": speculative_config = { diff --git a/tests/v1/e2e/test_spec_decode.py b/tests/v1/e2e/test_spec_decode.py index 4905a4120a2..a141e9da08a 100644 --- a/tests/v1/e2e/test_spec_decode.py +++ b/tests/v1/e2e/test_spec_decode.py @@ -13,15 +13,12 @@ from vllm import LLM, SamplingParams from vllm.assets.base import VLLM_S3_BUCKET_URL from vllm.assets.image import VLM_IMAGES_DIR from vllm.benchmarks.datasets import InstructCoderDataset -from vllm.config.vllm import VllmConfig +from vllm.config import VllmConfig from vllm.distributed import cleanup_dist_env_and_memory from vllm.engine.arg_utils import EngineArgs from vllm.platforms import current_platform from vllm.v1.metrics.reader import Metric -from vllm.v1.spec_decode.draft_model import ( - create_vllm_config_for_draft_model, - merge_toks_kernel, -) +from vllm.v1.spec_decode.utils import create_vllm_config_for_draft_model MTP_SIMILARITY_RATE = 0.8 @@ -625,6 +622,8 @@ class ArgsTest: expected_acceptance_rate: float expected_acceptance_len: float # Defaults + enforce_eager: bool = True + parallel_drafting: bool = False target_tensor_parallel_size: int = 1 draft_tensor_parallel_size: int = 1 max_model_len: int = 1024 @@ -658,7 +657,8 @@ cases = [ @pytest.mark.parametrize("args", cases) @pytest.mark.parametrize("enforce_eager", [True, False]) def test_draft_model_correctness(args: ArgsTest, enforce_eager: bool): - assert_draft_model_correctness(args, enforce_eager) + args.enforce_eager = enforce_eager + assert_draft_model_correctness(args) def test_draft_model_realistic_example(): @@ -668,11 +668,28 @@ def test_draft_model_realistic_example(): dataset="likaixin/InstructCoder", num_speculative_tokens=3, sampling_config=greedy_sampling(), + enforce_eager=False, # values below are not derived, but just prevent a regression expected_acceptance_len=2.8, expected_acceptance_rate=0.55, ) - assert_draft_model_correctness(args, enforce_eager=False) + assert_draft_model_correctness(args) + + +def test_draft_model_parallel_drafting(): + args = ArgsTest( + target_model="Qwen/Qwen3-1.7B", + draft_model="amd/PARD-Qwen3-0.6B", + dataset="likaixin/InstructCoder", + num_speculative_tokens=3, + sampling_config=greedy_sampling(), + parallel_drafting=True, + enforce_eager=False, + # values below are collected from a stable run, with ~5% tolerance + expected_acceptance_len=2.375, + expected_acceptance_rate=0.45, + ) + assert_draft_model_correctness(args) @pytest.mark.parametrize( @@ -691,8 +708,9 @@ def test_draft_model_quantization(models: tuple[str, str], enforce_eager: bool): target_model=tgt_model, draft_model=draft_model, **some_high_acceptance_metrics(), + enforce_eager=enforce_eager, ) - assert_draft_model_correctness(sd_case, enforce_eager) + assert_draft_model_correctness(sd_case) def test_draft_model_tensor_parallelism(): @@ -704,8 +722,9 @@ def test_draft_model_tensor_parallelism(): draft_model="Qwen/Qwen3-0.6B", draft_tensor_parallel_size=2, **some_high_acceptance_metrics(), + enforce_eager=False, ) - assert_draft_model_correctness(sd_case, enforce_eager=False) + assert_draft_model_correctness(sd_case) def test_draft_model_engine_args_tensor_parallelism(): @@ -750,7 +769,7 @@ def test_draft_model_engine_args_rejects_invalid_tp_argname(): engine_args.create_engine_config() -def assert_draft_model_correctness(args: ArgsTest, enforce_eager: bool): +def assert_draft_model_correctness(args: ArgsTest): """Compare the outputs using and not using speculative decoding. In the greedy decoding case, the outputs must match EXACTLY.""" test_prompts: list[Messages] = get_messages( @@ -764,14 +783,15 @@ def assert_draft_model_correctness(args: ArgsTest, enforce_eager: bool): "method": "draft_model", "num_speculative_tokens": args.num_speculative_tokens, "max_model_len": args.max_model_len, - "enforce_eager": enforce_eager, + "enforce_eager": args.enforce_eager, "draft_tensor_parallel_size": args.draft_tensor_parallel_size, + "parallel_drafting": args.parallel_drafting, }, max_num_seqs=100, # limit cudagraph capture runtime max_model_len=args.max_model_len, gpu_memory_utilization=args.gpu_memory_utilization, tensor_parallel_size=args.target_tensor_parallel_size, - enforce_eager=enforce_eager, + enforce_eager=args.enforce_eager, disable_log_stats=False, # enables get_metrics() ) # we don't check the outputs, only check the metrics @@ -813,57 +833,6 @@ def some_high_acceptance_metrics() -> dict: } -def test_merge_toks_kernel(): - device = "cuda" - merged_len = 5 + 2 # len(target_toks) = 5, batch_size = 2 - merged = torch.full((merged_len,), -100, device=device) # -100 is arbitrary - is_rejected_tok = torch.full((merged_len,), True, device=device) - grid = (2,) - merge_toks_kernel[grid]( - target_toks_ptr=torch.tensor([0, 1, 2, 0, 1], device=device), - next_toks_ptr=torch.tensor([3, 2], device=device), - query_start_locs_ptr=torch.tensor([0, 3], device=device), - query_end_locs_ptr=torch.tensor([2, 4], device=device), - out_ptr_merged_toks=merged, - out_ptr_is_rejected_tok=is_rejected_tok, - target_toks_size=5, - rejected_tok_fill=-1, - ) - expected_merged = torch.tensor([0, 1, 2, 3, 0, 1, 2], device=device) - assert torch.allclose(merged, expected_merged) - - expected_rejected_toks = torch.tensor([False] * merged_len, device=device) - assert torch.allclose(is_rejected_tok, expected_rejected_toks) - - -def test_merge_toks_kernel_with_rejected_tokens(): - device = "cuda" - merged_size = 9 + 2 # len(target_toks) = 9, batch_size = 2 - merged = torch.full((merged_size,), -100, device=device) - is_rejected_tok = torch.full((merged_size,), True, device=device) - grid = (2,) - merge_toks_kernel[grid]( - # rejected tokens - # ↓ ↓ ↓ ↓ - target_toks_ptr=torch.tensor([0, 1, 2, 13, 14, 15, 0, 1, 22], device=device), - next_toks_ptr=torch.tensor([3, 2], device=device), - query_start_locs_ptr=torch.tensor([0, 6], device=device), - query_end_locs_ptr=torch.tensor([2, 7], device=device), - out_ptr_merged_toks=merged, - out_ptr_is_rejected_tok=is_rejected_tok, - target_toks_size=9, - rejected_tok_fill=-1, - ) - expected_merged = torch.tensor([0, 1, 2, 3, -1, -1, -1, 0, 1, 2, -1], device=device) - assert torch.allclose(merged, expected_merged) - - expected_rejected_toks = torch.tensor( - [False, False, False, False, True, True, True, False, False, False, True], - device=device, - ) - assert torch.allclose(is_rejected_tok, expected_rejected_toks) - - def compute_acceptance_rate(metrics: list[Metric]) -> float: name2metric = {metric.name: metric for metric in metrics} n_draft_toks = name2metric["vllm:spec_decode_num_draft_tokens"].value # type: ignore diff --git a/tests/v1/spec_decode/test_eagle.py b/tests/v1/spec_decode/test_eagle.py index 3158ff0bda9..8b180168dff 100644 --- a/tests/v1/spec_decode/test_eagle.py +++ b/tests/v1/spec_decode/test_eagle.py @@ -27,6 +27,7 @@ from vllm.config.load import LoadConfig from vllm.model_executor.models.llama import LlamaForCausalLM from vllm.platforms import current_platform from vllm.v1.attention.backends.registry import AttentionBackendEnum +from vllm.v1.spec_decode.draft_model import DraftModelProposer from vllm.v1.spec_decode.eagle import EagleProposer from vllm.v1.spec_decode.metadata import SpecDecodeMetadata from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch @@ -34,6 +35,7 @@ from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch model_dir = "meta-llama/Llama-3.1-8B-Instruct" eagle_dir = "yuhuili/EAGLE-LLaMA3.1-Instruct-8B" eagle3_dir = "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B" +ar_draft_model_dir = "amd/PARD-Llama-3.2-1B" # Compatible with parallel and AR drafting def _create_proposer( @@ -41,11 +43,19 @@ def _create_proposer( num_speculative_tokens: int, attention_backend: str | None = None, speculative_token_tree: list[tuple[int, ...]] | None = None, + parallel_drafting: bool = False, ) -> EagleProposer: model_config = ModelConfig(model=model_dir, runner="generate", max_model_len=100) - # Choose model directory based on method - draft_model_dir = eagle_dir if method == "eagle" else eagle3_dir + # Method-dependent setup + if method == "eagle": + draft_model_dir = eagle_dir + elif method == "eagle3": + draft_model_dir = eagle3_dir + elif method == "draft_model": + draft_model_dir = ar_draft_model_dir + else: + raise ValueError(f"Unknown method: {method}") spec_token_tree_str = None if speculative_token_tree is not None: @@ -59,13 +69,18 @@ def _create_proposer( method=method, num_speculative_tokens=num_speculative_tokens, speculative_token_tree=spec_token_tree_str, + parallel_drafting=parallel_drafting, ) + if parallel_drafting: + # Overwrite pard_token to avoid crash during init + speculative_config.draft_model_config.hf_config.pard_token = 0 + device = current_platform.device_type vllm_config = VllmConfig( model_config=model_config, cache_config=CacheConfig(), speculative_config=speculative_config, - device_config=DeviceConfig(device=current_platform.device_type), + device_config=DeviceConfig(device=device), parallel_config=ParallelConfig(), load_config=LoadConfig(), scheduler_config=SchedulerConfig( @@ -75,7 +90,10 @@ def _create_proposer( attention_config=AttentionConfig(backend=attention_backend), ) - return EagleProposer(vllm_config=vllm_config, device=current_platform.device_type) + if "eagle" in method: + return EagleProposer(vllm_config=vllm_config, device=device) + else: + return DraftModelProposer(vllm_config=vllm_config, device=device) def test_prepare_next_token_ids(): @@ -321,6 +339,390 @@ def test_prepare_inputs_padded(): assert torch.equal(token_indices_to_sample, expected_token_indices_to_sample) +def test_set_inputs_first_pass_default_eagle(): + """ + Test for set_inputs_first_pass without extra input slots (default EAGLE). + + This tests the path where needs_extra_input_slots=False, which is the + default EAGLE pathway. In this case: + - Input IDs are rotated (shifted by one) + - The next_token_ids are inserted at the last position of each request + - Positions are copied as-is + - Hidden states are copied as-is + - The CommonAttentionMetadata is returned unchanged + + Setup: + - 3 requests with query_lens [3, 2, 4] + - Tokens: [a1, a2, a3, b1, b2, c1, c2, c3, c4] + - After rotation: [a2, a3, -, b2, -, c2, c3, c4, -] + - After inserting next_tokens [100, 200, 300]: + [a2, a3, 100, b2, 200, c2, c3, c4, 300] + """ + device = torch.device(current_platform.device_type) + + num_speculative_tokens = 3 + proposer = _create_proposer("eagle", num_speculative_tokens) + + # Setup batch with 3 requests + batch_spec = BatchSpec( + seq_lens=[10, 8, 12], # Arbitrary context lengths + query_lens=[3, 2, 4], + ) + + common_attn_metadata = create_common_attn_metadata( + batch_spec, + block_size=16, + device=device, + ) + + # Input tensors + # Request 0: tokens [10, 11, 12] at positions [7, 8, 9] + # Request 1: tokens [20, 21] at positions [6, 7] + # Request 2: tokens [30, 31, 32, 33] at positions [8, 9, 10, 11] + target_token_ids = torch.tensor( + [10, 11, 12, 20, 21, 30, 31, 32, 33], dtype=torch.int32, device=device + ) + target_positions = torch.tensor( + [7, 8, 9, 6, 7, 8, 9, 10, 11], dtype=torch.int64, device=device + ) + target_hidden_states = torch.randn( + 9, proposer.hidden_size, dtype=proposer.dtype, device=device + ) + next_token_ids = torch.tensor([100, 200, 300], dtype=torch.int32, device=device) + + num_tokens, token_indices_to_sample, output_cad = proposer.set_inputs_first_pass( + target_token_ids=target_token_ids, + next_token_ids=next_token_ids, + target_positions=target_positions, + target_hidden_states=target_hidden_states, + token_indices_to_sample=None, + cad=common_attn_metadata, + num_rejected_tokens_gpu=None, + ) + + assert num_tokens == 9 # Total tokens unchanged + + expected_token_indices_to_sample = torch.tensor( + [2, 4, 8], dtype=torch.int32, device=device + ) + assert torch.equal(token_indices_to_sample, expected_token_indices_to_sample) + + assert output_cad is common_attn_metadata + + # Verify input_ids are rotated and next_tokens inserted + # Original: [10, 11, 12, 20, 21, 30, 31, 32, 33] + # After shift by 1: [11, 12, 12, 21, 21, 31, 32, 33, 33] + # After inserting at last indices [2, 4, 8]: [11, 12, 100, 21, 200, 31, 32, 33, 300] + expected_input_ids = torch.tensor( + [11, 12, 100, 21, 200, 31, 32, 33, 300], dtype=torch.int32, device=device + ) + assert torch.equal(proposer.input_ids[:num_tokens], expected_input_ids) + + # Verify positions are copied as-is + assert torch.equal(proposer.positions[:num_tokens], target_positions) + + # Verify hidden states are copied as-is + assert torch.equal(proposer.hidden_states[:num_tokens], target_hidden_states) + + +def test_set_inputs_first_pass_draft_model(): + """ + Test for set_inputs_first_pass with a draft model (extra input slots, + no shift). + + This tests the path where needs_extra_input_slots=True and + shift_input_ids=False (draft model case). In this case: + - Input IDs are NOT shifted + - Each request gets extra_slots_per_request (1) new slots + - The kernel handles copying tokens and inserting bonus/padding tokens + - A new CommonAttentionMetadata is returned with updated query_start_loc + + Setup: + - 2 requests + - Request 0: tokens [10, 11, 12] at positions [0, 1, 2] + - Only tokens [10, 11] are "valid" (query_end_loc=1), + token 12 is a rejected token from previous speculation + - Request 1: tokens [20, 21] at positions [0, 1], both valid. + - Note: this is less than num_speculative_tokens (2) to ensure + we handle variable lengths correctly. + - next_token_ids: [100, 200] (bonus tokens) + + With extra_slots_per_request=1 and shift=False: + Expected output layout: + Request 0 (indices 0-3): + - idx 0: token 10, pos 0 + - idx 1: token 11, pos 1 + - idx 2: token 100, pos 2 (bonus token) + - idx 3: padding_token_id, is_rejected=True + Request 1 (indices 4-6): + - idx 4: token 20, pos 0 + - idx 5: token 21, pos 1 + - idx 6: token 200, pos 2 (bonus token) + """ + device = torch.device(current_platform.device_type) + + num_speculative_tokens = 2 + block_size = 16 + + # Create a proposer configured as a draft model (pass_hidden_states=False) + # We need to mock this since _create_proposer defaults to EAGLE + proposer = _create_proposer("draft_model", num_speculative_tokens) + + proposer.parallel_drafting_token_id = 0 + proposer.is_rejected_token_mask = torch.zeros( + proposer.max_num_tokens, dtype=torch.bool, device=device + ) + proposer.is_masked_token_mask = torch.zeros( + proposer.max_num_tokens, dtype=torch.bool, device=device + ) + + # Mock the attn_metadata_builder to avoid needing the full model setup + mock_kv_cache_spec = mock.MagicMock() + mock_kv_cache_spec.block_size = block_size + mock_builder = mock.MagicMock() + mock_builder.kv_cache_spec = mock_kv_cache_spec + proposer.attn_metadata_builder = mock_builder + + # Request 0: query_len=3 (but 1 rejected), Request 1: query_len=2 + batch_spec = BatchSpec( + seq_lens=[3, 2], + query_lens=[3, 2], + ) + + common_attn_metadata = create_common_attn_metadata( + batch_spec, + block_size=block_size, + device=device, + arange_block_indices=True, # Use predictable block indices + ) + + # Input tensors + target_token_ids = torch.tensor( + [10, 11, 12, 20, 21], dtype=torch.int32, device=device + ) + target_positions = torch.tensor([0, 1, 2, 0, 1], dtype=torch.int64, device=device) + target_hidden_states = torch.randn( + 5, proposer.hidden_size, dtype=proposer.dtype, device=device + ) + next_token_ids = torch.tensor([100, 200], dtype=torch.int32, device=device) + + num_rejected_tokens_gpu = torch.tensor([1, 0], dtype=torch.int32, device=device) + + num_tokens, token_indices_to_sample, output_cad = proposer.set_inputs_first_pass( + target_token_ids=target_token_ids, + next_token_ids=next_token_ids, + target_positions=target_positions, + target_hidden_states=target_hidden_states, + token_indices_to_sample=None, + cad=common_attn_metadata, + num_rejected_tokens_gpu=num_rejected_tokens_gpu, + ) + + assert proposer.net_num_new_slots_per_request == 1 + assert proposer.needs_extra_input_slots + + # total_output_tokens = total_input_tokens + net_num_new_slots * batch_size + assert num_tokens == 7 + + # Request 0: [10, 11, 100, padding_token (0)] + # Request 1: [20, 21, 200] + # Combined: [10, 11, 100, 0, 20, 21, 200] + expected_input_ids = torch.tensor( + [10, 11, 100, 0, 20, 21, 200], dtype=torch.int32, device=device + ) + assert torch.equal(proposer.input_ids[:num_tokens], expected_input_ids) + + # Verify positions + # Request 0: [0, 1, 2, 0 (don't care)] + # Request 1: [0, 1, 2] + # Combined: [0, 1, 2, 0, 0, 1, 2] + expected_positions = torch.tensor( + [0, 1, 2, 0, 0, 1, 2], dtype=torch.int64, device=device + ) + assert torch.equal( + proposer.positions[:num_tokens], + expected_positions, + ) + + # Verify rejection mask + expected_is_rejected = torch.zeros(7, dtype=torch.bool, device=device) + expected_is_rejected[3] = True # padding token at index 3 + assert torch.equal( + proposer.is_rejected_token_mask[:num_tokens], expected_is_rejected + ) + + # Verify masked token mask (should all be False for non-parallel drafting) + expected_is_masked = torch.zeros(7, dtype=torch.bool, device=device) + assert torch.equal(proposer.is_masked_token_mask[:num_tokens], expected_is_masked) + + # Verify token_indices_to_sample (bonus tokens at indices 2 and 6) + expected_token_indices_to_sample = torch.tensor( + [2, 6], dtype=torch.int32, device=device + ) + assert torch.equal(token_indices_to_sample, expected_token_indices_to_sample) + + # Verify the new CAD has updated query_start_loc + # Original: [0, 3, 5] -> New: [0, 4, 7] (each request gains 1 slot) + expected_query_start_loc = torch.tensor([0, 4, 7], dtype=torch.int32, device=device) + assert torch.equal(output_cad.query_start_loc, expected_query_start_loc) + + +def test_set_inputs_first_pass_parallel_drafting(): + """ + Test for set_inputs_first_pass with parallel drafting (extra input slots, + with shift). + + This tests the path where needs_extra_input_slots=True and + shift_input_ids=True (parallel drafting case). In this case: + - Input IDs ARE shifted (like default EAGLE) + - Each request gets extra_slots_per_request (3) new slots + - Parallel drafting tokens are inserted and marked as masked + - Hidden states are mapped correctly + + Setup: + - 2 requests with query_lens [4, 4] (1 bonus + 3 spec tokens each) + - Request 0: tokens [10, 11, 12, 13] at positions [5, 6, 7, 8] + - Only tokens [10, 11, 12] are "valid", token 13 is rejected + - Request 1: tokens [20, 21, 22, 23] at positions [10, 11, 12, 13], all valid. + - next_token_ids: [100, 200] (bonus tokens) + + With shift_input_ids=True, extra_slots_per_request=3: + Expected output layout: + Request 0 (6 output slots = 4 - 1 + 3): + - idx 0-2: shifted tokens [11, 12, 100] + - idx 3-4: parallel_drafting_tokens, is_masked=True + - idx 5: padding_token, is_rejected=True + Request 1 (6 output slots = 4 - 1 + 3): + - idx 6-8: shifted tokens [21, 22, 23] + - idx 9: bonus token 200 + - idx 10-11: parallel_drafting_tokens, is_masked=True + """ + device = torch.device(current_platform.device_type) + + num_speculative_tokens = 3 + block_size = 16 + + proposer = _create_proposer("eagle", num_speculative_tokens, parallel_drafting=True) + + # Override to simulate parallel drafting behavior + proposer.parallel_drafting_token_id = -2 + proposer.parallel_drafting_hidden_state_tensor = torch.zeros( + proposer.hidden_size, dtype=proposer.dtype, device=device + ) + proposer.is_rejected_token_mask = torch.zeros( + proposer.max_num_tokens, dtype=torch.bool, device=device + ) + proposer.is_masked_token_mask = torch.zeros( + proposer.max_num_tokens, dtype=torch.bool, device=device + ) + + # Mock the attn_metadata_builder + mock_kv_cache_spec = mock.MagicMock() + mock_kv_cache_spec.block_size = block_size + mock_builder = mock.MagicMock() + mock_builder.kv_cache_spec = mock_kv_cache_spec + proposer.attn_metadata_builder = mock_builder + + # Request 0: query_len=4 (1 rejected), Request 1: query_len=4 (all valid) + batch_spec = BatchSpec( + seq_lens=[9, 14], + query_lens=[4, 4], + ) + + common_attn_metadata = create_common_attn_metadata( + batch_spec, + block_size=block_size, + device=device, + arange_block_indices=True, + ) + + # Input tensors + target_token_ids = torch.tensor( + [10, 11, 12, 13, 20, 21, 22, 23], dtype=torch.int32, device=device + ) + target_positions = torch.tensor( + [5, 6, 7, 8, 10, 11, 12, 13], dtype=torch.int64, device=device + ) + target_hidden_states = torch.arange( + 8 * proposer.hidden_size, dtype=proposer.dtype, device=device + ).view(8, proposer.hidden_size) + next_token_ids = torch.tensor([100, 200], dtype=torch.int32, device=device) + + num_rejected_tokens_gpu = torch.tensor([1, 0], dtype=torch.int32, device=device) + + num_tokens, token_indices_to_sample, output_cad = proposer.set_inputs_first_pass( + target_token_ids=target_token_ids, + next_token_ids=next_token_ids, + target_positions=target_positions, + target_hidden_states=target_hidden_states, + token_indices_to_sample=None, + cad=common_attn_metadata, + num_rejected_tokens_gpu=num_rejected_tokens_gpu, + ) + + # total_output_tokens = total_input_tokens + net_num_new_slots * batch_size + # = 8 + 2 * 2 = 12 + assert num_tokens == 12 + + # Request 0: [11, 12, 100, -2, -2, 0(padding)] + # Request 1: [21, 22, 23, 200, -2, -2] + expected_input_ids = torch.tensor( + [11, 12, 100, -2, -2, 0, 21, 22, 23, 200, -2, -2], + dtype=torch.int32, + device=device, + ) + assert torch.equal(proposer.input_ids[:num_tokens], expected_input_ids) + + # Verify positions + # Request 0: [5, 6, 7, 8, 9, 0 (don't care)] + # Request 1: [10, 11, 12, 13, 14, 15] + expected_positions = torch.tensor( + [5, 6, 7, 8, 9, 0, 10, 11, 12, 13, 14, 15], dtype=torch.int64, device=device + ) + assert torch.equal( + proposer.positions[:num_tokens], + expected_positions, + ) + + # Verify rejection mask + expected_is_rejected = torch.zeros(12, dtype=torch.bool, device=device) + expected_is_rejected[5] = True + assert torch.equal( + proposer.is_rejected_token_mask[:num_tokens], expected_is_rejected + ) + + # Verify masked token mask (parallel drafting slots should be masked) + expected_is_masked = torch.zeros(12, dtype=torch.bool, device=device) + expected_is_masked[3] = True + expected_is_masked[4] = True + expected_is_masked[10] = True + expected_is_masked[11] = True + assert torch.equal(proposer.is_masked_token_mask[:num_tokens], expected_is_masked) + + # Verify token_indices_to_sample (bonus + parallel drafting tokens) + # Request 0: bonus at 2, parallel at 3, 4 + # Request 1: bonus at 9, parallel at 10, 11 + expected_token_indices_to_sample = torch.tensor( + [2, 3, 4, 9, 10, 11], dtype=torch.int32, device=device + ) + assert torch.equal(token_indices_to_sample, expected_token_indices_to_sample) + + # Verify the new CAD has updated query_start_loc + # Original query_lens: [4, 4] -> Output: [6, 6] + expected_query_start_loc = torch.tensor( + [0, 6, 12], dtype=torch.int32, device=device + ) + assert torch.equal(output_cad.query_start_loc, expected_query_start_loc) + + # Verify masked positions have the parallel drafting hidden state (zeros) + parallel_drafting_hs = proposer.parallel_drafting_hidden_state_tensor + for i in range(num_tokens): + if expected_is_masked[i]: + assert torch.equal(proposer.hidden_states[i], parallel_drafting_hs), ( + f"Masked position {i} should have parallel drafting hidden state" + ) + + @pytest.mark.parametrize("method", ["eagle", "eagle3"]) @pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform()) @pytest.mark.parametrize("pp_size", [1, 2]) @@ -579,7 +981,7 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch): target_positions=target_positions, target_hidden_states=target_hidden_states, next_token_ids=next_token_ids, - last_token_indices=None, + token_indices_to_sample=None, common_attn_metadata=common_attn_metadata, sampling_metadata=sampling_metadata, ) @@ -737,7 +1139,7 @@ def test_propose_tree(spec_token_tree): target_positions=target_positions, target_hidden_states=target_hidden_states, next_token_ids=next_token_ids, - last_token_indices=None, + token_indices_to_sample=None, common_attn_metadata=common_attn_metadata, sampling_metadata=sampling_metadata, ) diff --git a/tests/v1/spec_decode/test_mtp.py b/tests/v1/spec_decode/test_mtp.py index b33dc58ffe3..16f4fb0befe 100644 --- a/tests/v1/spec_decode/test_mtp.py +++ b/tests/v1/spec_decode/test_mtp.py @@ -204,7 +204,7 @@ def test_mtp_propose(num_speculative_tokens, monkeypatch): target_positions=target_positions, target_hidden_states=target_hidden_states, next_token_ids=next_token_ids, - last_token_indices=None, + token_indices_to_sample=None, common_attn_metadata=common_attn_metadata, sampling_metadata=sampling_metadata, ) diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index ed3dbefb397..5a2fe8eeb43 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -116,9 +116,16 @@ class SpeculativeConfig: """Minimum size of ngram token window when using Ngram proposer, if provided. Defaults to 1.""" + # Alternative drafting strategies speculative_token_tree: str | None = None """Specifies the tree structure for speculative token generation. """ + parallel_drafting: bool = False + """Enable parallel drafting, where all speculative tokens are generated + in parallel rather than sequentially. This can improve performance but + requires the speculative model be trained to support parallel drafting. + Only compatible with EAGLE and draft model methods.""" + # required configuration params passed from engine target_model_config: SkipValidation[ModelConfig] = None # type: ignore """The configuration of the target model.""" diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 4d34c5584ba..0f499c39ead 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -604,10 +604,13 @@ class VllmConfig: # Currently, async scheduling only support eagle speculative # decoding. if self.speculative_config is not None: - if self.speculative_config.method not in get_args(EagleModelTypes): + if ( + self.speculative_config.method not in get_args(EagleModelTypes) + and self.speculative_config.method != "draft_model" + ): raise ValueError( "Currently, async scheduling is only supported " - "with EAGLE/MTP kind of speculative decoding." + "with EAGLE/MTP/Draft Model kind of speculative decoding." ) if self.speculative_config.disable_padded_drafter_batch: raise ValueError( @@ -1298,16 +1301,21 @@ class VllmConfig: computed_compile_ranges_split_points = [] # The upper bound of the compile ranges is the max_num_batched_tokens. - # For speculative decoding with draft model, the compile range must be extended - # by 1 for each sequence. + # For speculative decoding, the compile range must be extended + # - Sequential: + 1 * max_num_seqs (one draft token per iteration) + # - Parallel draft: + num_speculative_tokens * max_num_seqs compile_range_end = self.scheduler_config.max_num_batched_tokens if compile_range_end is not None: - do_extend: bool = ( - self.speculative_config is not None - and self.speculative_config.uses_draft_model() - ) - if do_extend: - compile_range_end += self.scheduler_config.max_num_seqs + if self.speculative_config is not None and ( + self.speculative_config.uses_draft_model() + or self.speculative_config.use_eagle() + ): + multiplier = ( + self.speculative_config.num_speculative_tokens + if self.speculative_config.parallel_drafting + else 1 + ) + compile_range_end += multiplier * self.scheduler_config.max_num_seqs computed_compile_ranges_split_points.append(compile_range_end) diff --git a/vllm/model_executor/models/llama_eagle3.py b/vllm/model_executor/models/llama_eagle3.py index e47a3ee74c6..5f66716d545 100644 --- a/vllm/model_executor/models/llama_eagle3.py +++ b/vllm/model_executor/models/llama_eagle3.py @@ -52,13 +52,16 @@ class LlamaDecoderLayer(LlamaDecoderLayer): # Subsequent layers use hidden_size (only hidden_states, no embeds) qkv_input_size = 2 * self.hidden_size if layer_idx == 0 else self.hidden_size - # override qkv + # Parallel drafting checkpoints may have attention bias enabled + qkv_bias = getattr(config, "attention_bias", False) + + # Override qkv_proj with correct input size and bias setting self.self_attn.qkv_proj = QKVParallelLinear( qkv_input_size, self.self_attn.head_dim, self.self_attn.total_num_heads, self.self_attn.total_num_kv_heads, - bias=False, + bias=qkv_bias, quant_config=quant_config, prefix=maybe_prefix(prefix, "qkv_proj"), ) @@ -293,6 +296,19 @@ class Eagle3LlamaForCausalLM(LlamaForCausalLM): requires_grad=False, ) + self.use_parallel_drafting = vllm_config.speculative_config.parallel_drafting + + if self.use_parallel_drafting: + self.register_buffer( + "mask_hidden", + torch.zeros( + 1, + (3 if self.model.use_aux_hidden_state else 1) + * self.config.hidden_size, + ), + persistent=False, + ) + def embed_input_ids( self, input_ids: torch.Tensor, @@ -347,12 +363,25 @@ class Eagle3LlamaForCausalLM(LlamaForCausalLM): model_weights = {} includes_draft_id_mapping = False includes_embed_tokens = False + includes_mask_hidden = False for name, loaded_weight in weights: if "t2d" in name: continue if "d2t" in name: name = name.replace("d2t", "draft_id_to_target_id") includes_draft_id_mapping = True + elif "mask_hidden" in name: + # Load mask_hidden directly into buffer + if not self.use_parallel_drafting: + logger.warning( + "mask_hidden found in weights but " + "model is not configured for parallel drafting. " + "Skipping loading mask_hidden." + ) + continue + self.mask_hidden.copy_(loaded_weight.view(1, -1)) + includes_mask_hidden = True + continue elif "lm_head" not in name: name = "model." + name if "embed_tokens" in name: @@ -360,7 +389,14 @@ class Eagle3LlamaForCausalLM(LlamaForCausalLM): model_weights[name] = loaded_weight process_eagle_weight(self, name) - skip_substrs = [] + if not includes_mask_hidden and self.use_parallel_drafting: + raise ValueError( + "mask_hidden not found in weights but " + "model is configured for parallel drafting. " + "Please provide mask_hidden in the weights." + ) + + skip_substrs = ["mask_hidden"] if not includes_draft_id_mapping: skip_substrs.append("draft_id_to_target_id") if not includes_embed_tokens: diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 49eb91576ed..9c004d7724d 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -480,9 +480,14 @@ class AttentionMetadataBuilder(ABC, Generic[M]): speculative_config is not None and speculative_config.num_speculative_tokens is not None ): + max_num_queries_for_spec = ( + 1 + + (2 if speculative_config.parallel_drafting else 1) + * speculative_config.num_speculative_tokens + ) self.reorder_batch_threshold = max( self.reorder_batch_threshold, - 1 + speculative_config.num_speculative_tokens, + max_num_queries_for_spec, ) if ( diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 7e02aa36ffc..8e81c6fe965 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -60,7 +60,7 @@ from vllm.v1.attention.backends.utils import ( ) from vllm.v1.attention.ops.common import cp_lse_ag_out_rs from vllm.v1.attention.ops.merge_attn_states import merge_attn_states -from vllm.v1.kv_cache_interface import AttentionSpec +from vllm.v1.kv_cache_interface import AttentionSpec, UniformTypeKVCacheSpecs from vllm.v1.utils import CpuGpuBuffer FLASHINFER_WORKSPACE_BUFFER_SIZE_BATCH_INVARIANT = 2048 * 1024 * 1024 @@ -658,12 +658,36 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): vllm_config: VllmConfig, kv_cache_spec: AttentionSpec, ) -> AttentionCGSupport: - has_trtllm_support = can_use_trtllm_attention( - num_qo_heads=vllm_config.model_config.get_num_attention_heads( - vllm_config.parallel_config - ), - num_kv_heads=kv_cache_spec.num_kv_heads, + """Get the cudagraph support level for FlashInfer attention. + + This depends on whether we can use TRTLLM attention for decodes, since we can + only do UNIFORM_SINGLE_TOKEN_DECODE if it is unavailable. + To check this, we must call can_use_trtllm_attention with the number of KV + heads from the kv_cache_spec. We check all available KV cache specs and + only return UNIFORM_BATCH if all of them support TRTLLM attention. + """ + # For UniformTypeKVCacheSpecs, check all contained specs + kv_specs = ( + kv_cache_spec.kv_cache_specs.values() + if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs) + else [kv_cache_spec] ) + num_qo_heads = vllm_config.model_config.get_num_attention_heads( + vllm_config.parallel_config + ) + has_trtllm_support: bool = len(kv_specs) > 0 + for spec in kv_specs: + if not isinstance(spec, AttentionSpec): + # FlashInfer only applies to attention, so we don't consider other types + # of KV spec (e.g. Mamba) here. This is mostly for type checking. + continue + if not can_use_trtllm_attention( + num_qo_heads=num_qo_heads, + num_kv_heads=spec.num_kv_heads, + ): + has_trtllm_support = False + break + if has_trtllm_support: return AttentionCGSupport.UNIFORM_BATCH else: diff --git a/vllm/v1/attention/backends/utils.py b/vllm/v1/attention/backends/utils.py index dab298f1481..e0aa2c988a2 100644 --- a/vllm/v1/attention/backends/utils.py +++ b/vllm/v1/attention/backends/utils.py @@ -825,38 +825,6 @@ def get_dcp_local_seq_lens( return dcp_local_seq_lens.squeeze(1) -def extend_all_queries_by_1( - common_attn_metadata: CommonAttentionMetadata, - arange: torch.Tensor, - new_slot_mapping: torch.Tensor, -) -> CommonAttentionMetadata: - """ - Creates a new CommonAttentionMetadata with all query lengths increased by 1. - Also all seq lens are increased by 1. - This is useful e.g. in speculative decoding with draft models, where we - extend each sequence by 1 token. - The slot mapping is computed externally, as it requires more information. - """ - cad = common_attn_metadata - # query start loc must be increased by [+0, +1, +2, ..., +batch_size] - new_query_start_loc = cad.query_start_loc + arange[: len(cad.query_start_loc)] - new_query_start_loc_cpu = cad.query_start_loc_cpu + torch.arange( - len(cad.query_start_loc_cpu), dtype=torch.int32 - ) - new_cad = cad.replace( - query_start_loc=new_query_start_loc, - query_start_loc_cpu=new_query_start_loc_cpu, - seq_lens=cad.seq_lens + 1, - # each request is extended by 1 token -> batch_size tokens are added - num_actual_tokens=cad.num_actual_tokens + cad.batch_size(), - # All query lens increase by 1, so max query len increases by 1 - max_query_len=cad.max_query_len + 1, - max_seq_len=cad.max_seq_len + 1, - slot_mapping=new_slot_mapping, - ) - return new_cad - - def mamba_get_block_table_tensor( block_table: torch.Tensor, seq_lens: torch.Tensor, diff --git a/vllm/v1/spec_decode/draft_model.py b/vllm/v1/spec_decode/draft_model.py index 18e98b26761..4361d6f0bc7 100644 --- a/vllm/v1/spec_decode/draft_model.py +++ b/vllm/v1/spec_decode/draft_model.py @@ -1,19 +1,15 @@ # 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 typing_extensions import override -from vllm.config import VllmConfig, get_layers_from_vllm_config, replace +from vllm.config import VllmConfig from vllm.logger import init_logger -from vllm.model_executor.layers.attention import Attention from vllm.model_executor.model_loader import get_model -from vllm.triton_utils import tl, triton -from vllm.v1.attention.backends.utils import ( - CommonAttentionMetadata, - extend_all_queries_by_1, -) -from vllm.v1.spec_decode.eagle import PADDING_SLOT_ID, SpecDecodeBaseProposer +from vllm.v1.spec_decode.eagle import SpecDecodeBaseProposer +from vllm.v1.spec_decode.utils import create_vllm_config_for_draft_model logger = init_logger(__name__) @@ -31,37 +27,9 @@ class DraftModelProposer(SpecDecodeBaseProposer): pass_hidden_states_to_model=False, runner=runner, ) - self._raise_if_multimodal() - self._raise_if_mrope() - self._raise_if_padded_drafter_batch_disabled() self._raise_if_vocab_size_mismatch() self._raise_if_draft_tp_mismatch() - def _block_size(self) -> int: - builder = self._get_attention_metadata_builder() - return builder.kv_cache_spec.block_size - - def _raise_if_multimodal(self): - if self.supports_mm_inputs: - raise NotImplementedError( - "Speculative Decoding with draft models " - "does not support multimodal models yet" - ) - - def _raise_if_mrope(self): - if self.draft_model_config.uses_mrope: - raise NotImplementedError( - "Speculative Decoding with draft models does not support M-RoPE yet" - ) - - def _raise_if_padded_drafter_batch_disabled(self): - if self.speculative_config.disable_padded_drafter_batch: - raise NotImplementedError( - "Speculative Decoding with draft models only supports " - "padded drafter batch. Please don't pass --disable-padded-drafter-batch" - " in the speculative_config." - ) - def _raise_if_vocab_size_mismatch(self): self.speculative_config.verify_equal_vocab_size_if_draft_model() @@ -82,193 +50,26 @@ class DraftModelProposer(SpecDecodeBaseProposer): "Please pass 'draft_tensor_parallel_size' in the speculative_config." ) - def set_inputs_first_pass( - self, - target_token_ids: torch.Tensor, - next_token_ids: torch.Tensor, - target_positions: torch.Tensor, - last_token_indices: torch.Tensor | None, - cad: CommonAttentionMetadata, - num_rejected_tokens_gpu: torch.Tensor | None, - ) -> tuple[int, torch.Tensor, CommonAttentionMetadata]: - batch_size = cad.batch_size() - grid = (batch_size,) - start_locs = cad.query_start_loc[:-1] - end_locs = cad.query_start_loc[1:] - 1 - if num_rejected_tokens_gpu is not None: - end_locs -= num_rejected_tokens_gpu - - num_tokens = target_token_ids.shape[0] + batch_size - is_rejected_tok = torch.empty( - (num_tokens,), device=self.input_ids.device, dtype=torch.bool - ) - merge_toks_kernel[grid]( - target_toks_ptr=target_token_ids, - next_toks_ptr=next_token_ids, - query_start_locs_ptr=start_locs, - query_end_locs_ptr=end_locs, - out_ptr_merged_toks=self.input_ids, - out_ptr_is_rejected_tok=is_rejected_tok, - target_toks_size=target_token_ids.shape[0], - # passing a negative rejected_tok_fill value will raise an error - # when the value is used to index into embeddings. - # Therefore, we pass a valid integer, e.g. 0. - rejected_tok_fill=0, - ) - merge_toks_kernel[grid]( - target_toks_ptr=target_positions, - next_toks_ptr=target_positions[end_locs] + 1, - query_start_locs_ptr=start_locs, - query_end_locs_ptr=end_locs, - out_ptr_merged_toks=self.positions, - out_ptr_is_rejected_tok=is_rejected_tok, - target_toks_size=target_positions.shape[0], - rejected_tok_fill=0, - ) - - # recompute slot mapping - new_slot_mapping = compute_new_slot_mapping( - cad=cad, - new_positions=self.positions[:num_tokens], - is_rejected_token_mask=is_rejected_tok, - block_size=self._block_size(), - max_model_len=self.max_model_len, - ) - # update common_attn_metadata - new_cad: CommonAttentionMetadata = extend_all_queries_by_1( - cad, - arange=self.arange, - new_slot_mapping=new_slot_mapping, - ) - - new_last_token_indices = new_cad.query_start_loc[1:] - 1 - if num_rejected_tokens_gpu is not None: - new_last_token_indices -= num_rejected_tokens_gpu - - return num_tokens, new_last_token_indices, new_cad - - def load_model(self, target_model: Any) -> None: - """Takes target_model to satisfy the type checker.""" - - # This must be computed before loading the draft model - # because that mutates the forward_context of the vllm_config - target_attn_layer_names = set( - get_layers_from_vllm_config(self.vllm_config, Attention).keys() - ) - + @override + def _get_model(self) -> nn.Module: + # Draft models may be quantized or on different parallelism, + # so we load them with a modified vllm config from vllm.compilation.backends import set_model_tag - draft_vllm_config: VllmConfig = create_vllm_config_for_draft_model( - target_model_vllm_config=self.vllm_config - ) - logger.info( - "Starting to load draft model %s. TP=%d, rank=%d", - draft_vllm_config.model_config.model, - draft_vllm_config.parallel_config.tensor_parallel_size, - draft_vllm_config.parallel_config.rank, - ) + temp_vllm_config = create_vllm_config_for_draft_model(self.vllm_config) with set_model_tag("draft_model"): - self.model = get_model(vllm_config=draft_vllm_config, prefix="draft_model") + model = get_model( + vllm_config=temp_vllm_config, + prefix="draft_model", + ) + return model - # This must be computed after loading the draft model - # because that mutates the forward_context of the vllm_config - draft_attn_layer_names = ( - get_layers_from_vllm_config(self.vllm_config, Attention).keys() - - target_attn_layer_names - ) - self.attn_layer_names = list(draft_attn_layer_names) + @override + def _maybe_share_embeddings(self, target_language_model: nn.Module) -> None: + # Draft models don't share embeddings with the target model + pass - -def create_vllm_config_for_draft_model( - target_model_vllm_config: VllmConfig, -) -> VllmConfig: - """The vllm_config is configured for the target model, e.g. - its quant_config and parallel_config. But the draft model is potentially - quantized differently, and has potentially different tensor_parallel_size. - This function creates a new vllm_config configured for the draft model. - The vllm_config is useful when loading the draft model with get_model(). - """ - old = target_model_vllm_config - assert old.speculative_config is not None, "speculative_config is not set" - old_spec_config = old.speculative_config - new_parallel_config = replace( - old_spec_config.draft_parallel_config, - rank=old.parallel_config.rank, - ) - new: VllmConfig = replace( - old, - quant_config=None, # quant_config is recomputed in __init__() - model_config=old_spec_config.draft_model_config, - parallel_config=new_parallel_config, - ) - return new - - -def compute_new_slot_mapping( - cad: CommonAttentionMetadata, - new_positions: torch.Tensor, - is_rejected_token_mask: torch.Tensor, - block_size: int, - max_model_len: int, -): - batch_size, n_blocks_per_req = cad.block_table_tensor.shape - req_indices = torch.arange(batch_size, device=cad.query_start_loc.device) - req_indices = torch.repeat_interleave( - req_indices, cad.naive_query_lens() + 1, output_size=len(new_positions) - ) - # Clamp the positions to prevent an out-of-bounds error when indexing - # into block_table_tensor. - clamped_positions = torch.clamp(new_positions, max=max_model_len - 1) - block_table_indices = ( - req_indices * n_blocks_per_req + clamped_positions // block_size - ) - block_nums = cad.block_table_tensor.view(-1)[block_table_indices] - block_offsets = clamped_positions % block_size - new_slot_mapping = block_nums * block_size + block_offsets - # Mask out the position ids that exceed the max model length. - exceeds_max_model_len = new_positions >= max_model_len - new_slot_mapping.masked_fill_(exceeds_max_model_len, PADDING_SLOT_ID) - # Mask out rejected tokens to prevent saves to the KV cache. - new_slot_mapping.masked_fill_(is_rejected_token_mask, PADDING_SLOT_ID) - return new_slot_mapping - - -@triton.jit -def merge_toks_kernel( - target_toks_ptr, - next_toks_ptr, - query_start_locs_ptr, - query_end_locs_ptr, - out_ptr_merged_toks, - out_ptr_is_rejected_tok, - target_toks_size, - rejected_tok_fill, -): - """ - Merges the `target_toks_ptr` and the `next_toks_ptr` into a new tensor - called `out_ptr_merged_toks`. Rejected tokens are those after the - `query_end_locs_ptr` and before the next `query_start_locs_ptr`. Fills the - rejected tokens positions with the value `rejected_tok_fill`. Also fills a mask - of the rejected tokens in `out_ptr_is_rejected_tok`. - """ - pid = tl.program_id(0) - start_loc = tl.load(query_start_locs_ptr + pid) - is_last_program = pid == tl.num_programs(0) - 1 - if is_last_program: - next_start_loc = target_toks_size.to(tl.int32) - else: - next_start_loc = tl.load(query_start_locs_ptr + pid + 1).to(tl.int32) - - end_loc = tl.load(query_end_locs_ptr + pid) - new_val = tl.load(next_toks_ptr + pid) - for i in range(start_loc, next_start_loc + 1): - if i <= end_loc: # copy existing tokens - old_val = tl.load(target_toks_ptr + i) - tl.store(out_ptr_merged_toks + pid + i, old_val) - tl.store(out_ptr_is_rejected_tok + pid + i, False) - elif i == end_loc + 1: # copy bonus token - tl.store(out_ptr_merged_toks + pid + i, new_val) - tl.store(out_ptr_is_rejected_tok + pid + i, False) - else: # fill rejected tokens - tl.store(out_ptr_merged_toks + pid + i, rejected_tok_fill) - tl.store(out_ptr_is_rejected_tok + pid + i, True) + @override + def _maybe_share_lm_head(self, target_language_model: nn.Module) -> None: + # Draft models don't share lm_head with the target model + pass diff --git a/vllm/v1/spec_decode/eagle.py b/vllm/v1/spec_decode/eagle.py index 45680a7965b..82505645cfc 100644 --- a/vllm/v1/spec_decode/eagle.py +++ b/vllm/v1/spec_decode/eagle.py @@ -43,8 +43,12 @@ from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.sample.sampler import _SAMPLING_EPS from vllm.v1.spec_decode.metadata import SpecDecodeMetadata from vllm.v1.spec_decode.utils import ( + PADDING_SLOT_ID, + compute_new_slot_mapping, + copy_and_expand_eagle_inputs_kernel, eagle_prepare_inputs_padded_kernel, eagle_prepare_next_token_padded_kernel, + extend_all_queries_by_N, ) from vllm.v1.utils import CpuGpuBuffer from vllm.v1.worker.dp_utils import coordinate_batch_across_dp @@ -52,8 +56,6 @@ from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch logger = init_logger(__name__) -PADDING_SLOT_ID = -1 - class SpecDecodeBaseProposer: def __init__( @@ -76,18 +78,35 @@ class SpecDecodeBaseProposer: self.max_model_len = vllm_config.model_config.max_model_len self.dp_rank = vllm_config.parallel_config.data_parallel_rank self.num_speculative_tokens = self.speculative_config.num_speculative_tokens - # The drafter can get longer sequences than the target model. - max_batch_size = vllm_config.scheduler_config.max_num_seqs - self.max_num_tokens = ( - vllm_config.scheduler_config.max_num_batched_tokens + max_batch_size - ) - self.token_arange_np = np.arange(self.max_num_tokens) + # We need to get the hidden size from the draft model config because # the draft model's hidden size can be different from the target model's # hidden size (e.g., Llama 3.3 70B). self.hidden_size = self.draft_model_config.get_hidden_size() self.inputs_embeds_size = self.draft_model_config.get_inputs_embeds_size() + # Unifying eagle, draft model, and parallel drafting support + self.parallel_drafting: bool = self.speculative_config.parallel_drafting + self.extra_slots_per_request = ( + 1 if not self.parallel_drafting else self.num_speculative_tokens + ) + self.net_num_new_slots_per_request = self.extra_slots_per_request - ( + 1 if self.pass_hidden_states_to_model else 0 + ) + self.needs_extra_input_slots = self.net_num_new_slots_per_request > 0 + + self.parallel_drafting_token_id: int = 0 + self.parallel_drafting_hidden_state_tensor: torch.Tensor | None = None + if self.parallel_drafting: + self._init_parallel_drafting_params() + + # The drafter can get longer sequences than the target model. + max_batch_size = vllm_config.scheduler_config.max_num_seqs + self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens + ( + self.net_num_new_slots_per_request * max_batch_size + ) + self.token_arange_np = np.arange(self.max_num_tokens) + # Multi-modal data support self.mm_registry = MULTIMODAL_REGISTRY self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs( @@ -155,6 +174,26 @@ class SpecDecodeBaseProposer: max_num_slots_for_arange, device=device, dtype=torch.int32 ) + if self.needs_extra_input_slots: + self._raise_if_padded_drafter_batch_disabled() + self._raise_if_multimodal() + self._raise_if_mrope() + + self.is_rejected_token_mask: torch.Tensor | None = None + self.is_masked_token_mask: torch.Tensor | None = None + if self.needs_extra_input_slots: + # For draft models and parallel drafting, we need to keep track of + # which tokens are rejected to update the slot mapping with padding slots. + self.is_rejected_token_mask = torch.zeros( + (self.max_num_tokens,), dtype=torch.bool, device=device + ) + # For parallel drafting, we also need to keep track of which tokens + # are parallel-padding tokens used to sample at later positions. + # We populate this tensor even when using draft models for simplicity. + self.is_masked_token_mask = torch.zeros( + (self.max_num_tokens,), dtype=torch.bool, device=device + ) + self.inputs_embeds = torch.zeros( (self.max_num_tokens, self.inputs_embeds_size), dtype=self.dtype, @@ -231,6 +270,49 @@ class SpecDecodeBaseProposer: 1, len(self.tree_choices) + 1, device=device, dtype=torch.int32 ).repeat(max_batch_size, 1) + def _raise_if_padded_drafter_batch_disabled(self): + if self.speculative_config.disable_padded_drafter_batch: + raise NotImplementedError( + "Speculative Decoding with draft models or parallel drafting only " + "supports padded drafter batch. Please unset " + "disable_padded_drafter_batch in the speculative_config." + ) + + def _raise_if_multimodal(self): + if self.supports_mm_inputs: + raise NotImplementedError( + "Speculative Decoding with draft models or parallel drafting " + "does not support multimodal models yet" + ) + + def _raise_if_mrope(self): + if self.draft_model_config.uses_mrope: + raise NotImplementedError( + "Speculative Decoding with draft models or parallel drafting " + "does not support M-RoPE yet" + ) + + def _init_parallel_drafting_params(self): + # For parallel drafting, we need the token ID to use for masked slots + # And for EAGLE + parallel drafting, we need the hidden state tensor to use + # for those masked slots. + + model_hf_config = self.draft_model_config.hf_config + if hasattr(model_hf_config, "pard_token"): + self.parallel_drafting_token_id = model_hf_config.pard_token + elif hasattr(model_hf_config, "ptd_token_id"): + self.parallel_drafting_token_id = model_hf_config.ptd_token_id + else: + raise ValueError( + "For parallel drafting, the draft model config must have " + "`pard_token` or `ptd_token_id` specified in its config.json." + ) + + if self.pass_hidden_states_to_model: + self.parallel_drafting_hidden_state_tensor = torch.empty( + self.hidden_size, dtype=self.dtype, device=self.device + ) + def _get_positions(self, num_tokens: int): if self.uses_mrope: return self.mrope_positions[:, :num_tokens] @@ -296,7 +378,7 @@ class SpecDecodeBaseProposer: target_hidden_states: torch.Tensor, # [batch_size] next_token_ids: torch.Tensor, - last_token_indices: torch.Tensor | None, + token_indices_to_sample: torch.Tensor | None, common_attn_metadata: CommonAttentionMetadata, sampling_metadata: SamplingMetadata, mm_embed_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, @@ -314,12 +396,13 @@ class SpecDecodeBaseProposer: ) assert target_hidden_states.shape[-1] == self.hidden_size - num_tokens, last_token_indices, common_attn_metadata = ( + num_tokens, token_indices_to_sample, common_attn_metadata = ( self.set_inputs_first_pass( target_token_ids=target_token_ids, next_token_ids=next_token_ids, target_positions=target_positions, - last_token_indices=last_token_indices, + target_hidden_states=target_hidden_states, + token_indices_to_sample=token_indices_to_sample, cad=common_attn_metadata, num_rejected_tokens_gpu=num_rejected_tokens_gpu, ) @@ -366,11 +449,6 @@ class SpecDecodeBaseProposer: if num_tokens_across_dp is not None: num_tokens_across_dp[self.dp_rank] = num_input_tokens - if self.pass_hidden_states_to_model: - # target_hidden_states and self.hidden_states can have different - # hidden dims. E.g. large target model and small draft model. - self.hidden_states[:num_tokens] = target_hidden_states - if self.supports_mm_inputs: mm_embeds, is_mm_embed = mm_embed_inputs or (None, None) @@ -411,27 +489,27 @@ class SpecDecodeBaseProposer: else: last_hidden_states, hidden_states = ret_hidden_states - sample_hidden_states = last_hidden_states[last_token_indices] + sample_hidden_states = last_hidden_states[token_indices_to_sample] logits = self.model.compute_logits(sample_hidden_states) # Early exit if there is only one draft token to be generated. - if self.num_speculative_tokens == 1: + if self.num_speculative_tokens == 1 or self.parallel_drafting: draft_token_ids = logits.argmax(dim=-1) - return draft_token_ids.view(-1, 1) + return draft_token_ids.view(-1, self.num_speculative_tokens) if self.uses_mrope: - positions = self.mrope_positions[:, last_token_indices] + positions = self.mrope_positions[:, token_indices_to_sample] else: - positions = self.positions[last_token_indices] + positions = self.positions[token_indices_to_sample] if self.method in ( "deepseek_mtp", "ernie_mtp", "longcat_flash_mtp", "pangu_ultra_moe_mtp", ): - hidden_states = self.hidden_states[last_token_indices] + hidden_states = self.hidden_states[token_indices_to_sample] else: - hidden_states = hidden_states[last_token_indices] + hidden_states = hidden_states[token_indices_to_sample] if isinstance(attn_metadata, TreeAttentionMetadata): # Draft using tree attention. @@ -624,27 +702,139 @@ class SpecDecodeBaseProposer: target_token_ids: torch.Tensor, next_token_ids: torch.Tensor, target_positions: torch.Tensor, - last_token_indices: torch.Tensor | None, + target_hidden_states: torch.Tensor, + token_indices_to_sample: torch.Tensor | None, cad: CommonAttentionMetadata, num_rejected_tokens_gpu: torch.Tensor | None, ) -> tuple[int, torch.Tensor, CommonAttentionMetadata]: - if last_token_indices is None: - last_token_indices = cad.query_start_loc[1:] - 1 + if not self.needs_extra_input_slots: + # Default EAGLE pathway: no reshaping of input tensors needed. + # Simply rotate the input ids and leave the positions unchanged, + # Inserting the next token ids at the last slot in each request. + if token_indices_to_sample is None: + token_indices_to_sample = cad.query_start_loc[1:] - 1 - num_tokens = target_token_ids.shape[0] - # Shift the input ids by one token. - # E.g., [a1, b1, b2, c1, c2, c3] -> [b1, b2, c1, c2, c3, c3] - self.input_ids[: num_tokens - 1] = target_token_ids[1:] - # Replace the last token with the next token. - # E.g., [b1, b2, c1, c2, c3, c3] -> [a2, b2, b3, c2, c3, c4] - self.input_ids[last_token_indices] = next_token_ids + num_tokens = target_token_ids.shape[0] + # Shift the input ids by one token. + # E.g., [a1, b1, b2, c1, c2, c3] -> [b1, b2, c1, c2, c3, c3] + self.input_ids[: num_tokens - 1] = target_token_ids[1:] + # Replace the last token with the next token. + # E.g., [b1, b2, c1, c2, c3, c3] -> [a2, b2, b3, c2, c3, c4] + self.input_ids[token_indices_to_sample] = next_token_ids - # copy inputs to buffer for cudagraph - if self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim == 0: - target_positions = target_positions[0] - self._set_positions(num_tokens, target_positions) + # copy inputs to buffer for cudagraph + if self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim == 0: + target_positions = target_positions[0] + self._set_positions(num_tokens, target_positions) - return num_tokens, last_token_indices, cad + self.hidden_states[:num_tokens] = target_hidden_states + + return num_tokens, token_indices_to_sample, cad + else: + assert self.is_rejected_token_mask is not None + assert self.is_masked_token_mask is not None + # 1. + # Call a custom triton kernel to copy input_ids and positions + # into the correct slots in the preallocated buffers self.input_ids, + # self.positions. + batch_size = cad.batch_size() + # Since we might have to copy a lot of data for prefills, we select the + # block size based on the max query length and limit to max 256 slots/block. + max_num_tokens_per_request = ( + cad.max_query_len + self.net_num_new_slots_per_request + ) + BLOCK_SIZE_TOKENS = min( + 256, triton.next_power_of_2(max_num_tokens_per_request) + ) + num_blocks = ( + max_num_tokens_per_request + BLOCK_SIZE_TOKENS - 1 + ) // BLOCK_SIZE_TOKENS + total_num_input_tokens = target_token_ids.shape[0] + total_num_output_tokens = total_num_input_tokens + ( + self.net_num_new_slots_per_request * batch_size + ) + + token_indices_to_sample = torch.empty( + batch_size * self.extra_slots_per_request, + dtype=torch.int32, + device=self.device, + ) + + # Destination indices to write target_hidden_states into drafting buffer. + out_hidden_state_mapping = torch.empty( + total_num_input_tokens, dtype=torch.int32, device=self.device + ) + + # Kernel grid: one program per request (row) + grid = (batch_size, num_blocks) + query_start_loc = cad.query_start_loc + query_end_loc = cad.query_start_loc[1:] - 1 + if num_rejected_tokens_gpu is not None: + query_end_loc = query_end_loc - num_rejected_tokens_gpu + copy_and_expand_eagle_inputs_kernel[grid]( + # (Padded) Inputs from the target model + target_token_ids_ptr=target_token_ids, + target_positions_ptr=target_positions, + next_token_ids_ptr=next_token_ids, # sampled tokens, one per request + # Outputs to the drafting buffers + out_input_ids_ptr=self.input_ids, + out_positions_ptr=self.positions, # Doesn't support mrope for now + out_is_rejected_token_mask_ptr=self.is_rejected_token_mask, + out_is_masked_token_mask_ptr=self.is_masked_token_mask, + out_new_token_indices_ptr=token_indices_to_sample, + out_hidden_state_mapping_ptr=out_hidden_state_mapping, + # Input metadata + query_start_loc_ptr=query_start_loc, + query_end_loc_ptr=query_end_loc, + padding_token_id=0, + parallel_drafting_token_id=self.parallel_drafting_token_id, + # Sizing info + # Note that we can deduce batch_size for free from the grid size + total_input_tokens=total_num_input_tokens, + num_padding_slots_per_request=self.extra_slots_per_request, + shift_input_ids=self.pass_hidden_states_to_model, + BLOCK_SIZE_TOKENS=BLOCK_SIZE_TOKENS, + ) + if self.pass_hidden_states_to_model: + assert self.parallel_drafting_hidden_state_tensor is not None + self.hidden_states[out_hidden_state_mapping] = target_hidden_states + # Use torch.where to avoid DtoH sync from boolean indexing + mask = self.is_masked_token_mask[:total_num_output_tokens] + torch.where( + mask.unsqueeze(1), + self.parallel_drafting_hidden_state_tensor, + self.hidden_states[:total_num_output_tokens], + out=self.hidden_states[:total_num_output_tokens], + ) + + # 2. + # Recompute the slot mapping based on the new positions and + # rejection mask. + builder = ( + self._get_attention_metadata_builder() + if self.attn_metadata_builder is None + else self.attn_metadata_builder + ) + new_slot_mapping = compute_new_slot_mapping( + cad=cad, + new_positions=self.positions[:total_num_output_tokens], + is_rejected_token_mask=self.is_rejected_token_mask[ + :total_num_output_tokens + ], + block_size=builder.kv_cache_spec.block_size, + num_new_tokens=self.net_num_new_slots_per_request, + max_model_len=self.max_model_len, + ) + + # 3. Update the common attention metadata with the new (meta)data + new_cad = extend_all_queries_by_N( + cad, + N=self.net_num_new_slots_per_request, + arange=self.arange, + new_slot_mapping=new_slot_mapping, + ) + + return total_num_output_tokens, token_indices_to_sample, new_cad def model_returns_tuple(self) -> bool: return self.method not in ("mtp", "draft_model") @@ -1081,8 +1271,21 @@ class SpecDecodeBaseProposer: model = model.module return model.__class__.__name__ + def _get_model(self) -> nn.Module: + """ + Default method to call get_model(). Can be overridden by subclasses which + need to customize model loading. + """ + from vllm.compilation.backends import set_model_tag + + with set_model_tag("eagle_head"): + model = get_model( + vllm_config=self.vllm_config, + model_config=self.speculative_config.draft_model_config, + ) + return model + def load_model(self, target_model: nn.Module) -> None: - draft_model_config = self.speculative_config.draft_model_config target_attn_layer_names = set( get_layers_from_vllm_config( self.vllm_config, @@ -1096,12 +1299,7 @@ class SpecDecodeBaseProposer: ).keys() ) - from vllm.compilation.backends import set_model_tag - - with set_model_tag("eagle_head"): - self.model = get_model( - vllm_config=self.vllm_config, model_config=draft_model_config - ) + self.model = self._get_model() draft_attn_layer_names = ( get_layers_from_vllm_config( @@ -1170,7 +1368,26 @@ class SpecDecodeBaseProposer: else: target_language_model = target_model - # share embed_tokens with the target model if needed + self._maybe_share_embeddings(target_language_model) + self._maybe_share_lm_head(target_language_model) + + if self.parallel_drafting and self.pass_hidden_states_to_model: + assert self.parallel_drafting_hidden_state_tensor is not None + self.parallel_drafting_hidden_state_tensor.copy_( + self.model.combine_hidden_states( + self.model.mask_hidden.view(3 * self.hidden_size) + ) + if self.eagle3_use_aux_hidden_state + else self.model.mask_hidden.view(self.hidden_size) + ) + + def _maybe_share_embeddings(self, target_language_model: nn.Module) -> None: + """ + Some draft models may not have their own embedding layers, and some may + have a duplicate copy of the target model's embedding layers. In these cases, + we share the target model's embedding layers with the draft model to save + memory. + """ if get_pp_group().world_size == 1: inner_model = getattr(target_language_model, "model", None) if inner_model is None: @@ -1233,7 +1450,12 @@ class SpecDecodeBaseProposer: " from the target model." ) - # share lm_head with the target model if needed + def _maybe_share_lm_head(self, target_language_model: nn.Module) -> None: + """ + Some draft models may not have their own LM head, and some may have a + duplicate copy of the target model's LM head. In these cases, we share + the target model's LM head with the draft model to save memory. + """ share_lm_head = False if hasattr(self.model, "has_own_lm_head"): # EAGLE model diff --git a/vllm/v1/spec_decode/utils.py b/vllm/v1/spec_decode/utils.py index 524714db37a..387c6df9bc4 100644 --- a/vllm/v1/spec_decode/utils.py +++ b/vllm/v1/spec_decode/utils.py @@ -1,6 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch + +from vllm.config import VllmConfig, replace from vllm.triton_utils import tl, triton +from vllm.v1.attention.backends.utils import ( + CommonAttentionMetadata, +) + +PADDING_SLOT_ID = -1 @triton.jit @@ -107,3 +115,243 @@ def eagle_prepare_next_token_padded_kernel( tl.store(next_token_ids_ptr + req_idx, backup_token) tl.store(valid_sampled_tokens_count_ptr + req_idx, valid_count) + + +def compute_new_slot_mapping( + cad: CommonAttentionMetadata, + new_positions: torch.Tensor, + is_rejected_token_mask: torch.Tensor, + block_size: int, + num_new_tokens: int, + max_model_len: int, +): + batch_size, n_blocks_per_req = cad.block_table_tensor.shape + req_indices = torch.arange(batch_size, device=cad.query_start_loc.device) + req_indices = torch.repeat_interleave( + req_indices, + cad.naive_query_lens() + num_new_tokens, + output_size=len(new_positions), + ) + # Clamp the positions to prevent an out-of-bounds error when indexing + # into block_table_tensor. + clamped_positions = torch.clamp(new_positions, max=max_model_len - 1) + block_table_indices = ( + req_indices * n_blocks_per_req + clamped_positions // block_size + ) + block_nums = cad.block_table_tensor.view(-1)[block_table_indices] + block_offsets = clamped_positions % block_size + new_slot_mapping = block_nums * block_size + block_offsets + # Mask out the position ids that exceed the max model length. + exceeds_max_model_len = new_positions >= max_model_len + new_slot_mapping.masked_fill_(exceeds_max_model_len, PADDING_SLOT_ID) + # Mask out rejected tokens to prevent saves to the KV cache. + new_slot_mapping.masked_fill_(is_rejected_token_mask, PADDING_SLOT_ID) + return new_slot_mapping + + +def create_vllm_config_for_draft_model( + target_model_vllm_config: VllmConfig, +) -> VllmConfig: + """The vllm_config is configured for the target model, e.g. + its quant_config and parallel_config. But the draft model is potentially + quantized differently, and has potentially different tensor_parallel_size. + This function creates a new vllm_config configured for the drafter. + The vllm_config is useful when loading the draft model with get_model(). + """ + old = target_model_vllm_config + assert old.speculative_config is not None, "speculative_config is not set" + old_spec_config = old.speculative_config + new_parallel_config = replace( + old_spec_config.draft_parallel_config, rank=old.parallel_config.rank + ) + new: VllmConfig = replace( + old, + quant_config=None, + parallel_config=new_parallel_config, + model_config=old_spec_config.draft_model_config, + ) + return new + + +def extend_all_queries_by_N( + common_attn_metadata: CommonAttentionMetadata, + N: int, + arange: torch.Tensor, + new_slot_mapping: torch.Tensor, +) -> CommonAttentionMetadata: + """ + Creates a new CommonAttentionMetadata with all query lengths increased by N. + Also all seq lens are increased by N. + This is useful e.g. in speculative decoding with parallel drafting, where we + extend each sequence by N tokens and predict all tokens in one pass. + The slot mapping is computed externally, as it requires more information. + """ + cad = common_attn_metadata + # query start loc must be increased by [+0, +N, +2N, ..., +batch_size * N] + new_query_start_loc = cad.query_start_loc + N * arange[: len(cad.query_start_loc)] + new_query_start_loc_cpu = cad.query_start_loc_cpu + N * torch.arange( + len(cad.query_start_loc_cpu), dtype=torch.int32 + ) + new_cad = cad.replace( + query_start_loc=new_query_start_loc, + query_start_loc_cpu=new_query_start_loc_cpu, + seq_lens=cad.seq_lens + N, + # each request is extended by N tokens -> batch_size * N tokens are added + num_actual_tokens=cad.num_actual_tokens + cad.batch_size() * N, + # All query lens increase by N, so max query len increases by N + max_query_len=cad.max_query_len + N, + max_seq_len=cad.max_seq_len + N, + slot_mapping=new_slot_mapping, + ) + return new_cad + + +# Unified copy/expand kernel +@triton.jit +def copy_and_expand_eagle_inputs_kernel( + # (Padded) Inputs from the target model + target_token_ids_ptr, # [total_tokens_in_batch] + target_positions_ptr, # [total_tokens_in_batch] + next_token_ids_ptr, # [num_reqs] + # Outputs to the drafting buffers + out_input_ids_ptr, # [total_draft_tokens_in_batch] (output) + out_positions_ptr, # [total_draft_tokens_in_batch] (output) + out_is_rejected_token_mask_ptr, # [total_draft_tokens_in_batch] (output) + out_is_masked_token_mask_ptr, # [total_draft_tokens_in_batch] (output) + out_new_token_indices_ptr, # [num_padding_slots_per_request * num_reqs] (output) + out_hidden_state_mapping_ptr, # [total_tokens_in_batch] + # Input metadata + query_start_loc_ptr, # [num_reqs + 1], last value is the total num input tokens + query_end_loc_ptr, # [num_reqs] + padding_token_id, # tl.int32 + parallel_drafting_token_id, # tl.int32 + # Sizing info + total_input_tokens, # tl.int32 + num_padding_slots_per_request, # tl.int32 + shift_input_ids, # tl.bool + BLOCK_SIZE_TOKENS: tl.constexpr, # Blocks along token dim to handle prefills +): + """ + Copy and expand inputs from the target model to the drafting buffers for Eagle + speculative decoding. This kernel handles padding slots and parallel drafting + tokens, if enabled. + """ + request_idx = tl.program_id(axis=0) + token_batch_idx = tl.program_id(axis=1) + + # Load query locations + query_start_loc = tl.load(query_start_loc_ptr + request_idx) + next_query_start_loc = tl.load(query_start_loc_ptr + request_idx + 1) + query_end_loc = tl.load(query_end_loc_ptr + request_idx) + + # Calculate number of valid tokens to copy and input offset + # With shift_input_ids=True, we skip the first token + # Output layout: each request gets (input_len + num_padding_slots_per_request) slots + # But with shift, we lose one token per request + if shift_input_ids: + num_valid_tokens = query_end_loc - query_start_loc + input_offset = 1 + output_start = query_start_loc + request_idx * ( + num_padding_slots_per_request - 1 + ) + else: + num_valid_tokens = query_end_loc - query_start_loc + 1 + input_offset = 0 + output_start = query_start_loc + request_idx * num_padding_slots_per_request + + # Number of rejected tokens from previous speculation + num_rejected = next_query_start_loc - query_end_loc - 1 + + # Total output tokens for this request + total_output_tokens = ( + num_valid_tokens + num_padding_slots_per_request + num_rejected + ) + + # Process tokens in this block + j = token_batch_idx * BLOCK_SIZE_TOKENS + tl.arange(0, BLOCK_SIZE_TOKENS) + + # Compute masks for different output regions: + # [0, num_valid_tokens): valid tokens copied from input + # [num_valid_tokens]: bonus token from next_token_ids + # (num_valid_tokens, num_valid_tokens + num_padding_slots_per_request): + # parallel drafting slots + # [num_valid_tokens + num_padding_slots_per_request, total_output_tokens): + # rejected slots + in_bounds = j < total_output_tokens + is_valid_region = j < num_valid_tokens + is_bonus_region = j == num_valid_tokens + is_parallel_draft_region = (j > num_valid_tokens) & ( + j < num_valid_tokens + num_padding_slots_per_request + ) + is_rejected_region = j >= num_valid_tokens + num_padding_slots_per_request + + # Compute output indices + out_idx = output_start + j + + # For valid tokens, compute input index + in_idx = query_start_loc + input_offset + j + # Clamp to avoid out-of-bounds access (masked loads still need valid addresses) + in_idx_clamped = tl.minimum(in_idx, total_input_tokens - 1) + + # Load input tokens (masked to valid region) + token_ids = tl.load( + target_token_ids_ptr + in_idx_clamped, mask=is_valid_region & in_bounds, other=0 + ) + + # Load the starting position for this request (first position in the sequence) + start_pos = tl.load(target_positions_ptr + query_start_loc) + + # Load bonus token for this request + bonus_token = tl.load(next_token_ids_ptr + request_idx) + + # Build final token_ids based on region + token_ids = tl.where(is_bonus_region, bonus_token, token_ids) + token_ids = tl.where( + is_parallel_draft_region, parallel_drafting_token_id, token_ids + ) + token_ids = tl.where(is_rejected_region, padding_token_id, token_ids) + + # Build final positions: + # Positions are NOT shifted - they start from the first input position and increment + # Output position j gets start_pos + j + # (e.g., input positions [5,6,7] -> output [5,6,7,8,9,...]) + positions = start_pos + j + # Rejected positions are don't-care, set to 0 + positions = tl.where(is_rejected_region, 0, positions) + + # Compute output masks + is_rejected_out = is_rejected_region & in_bounds + is_masked_out = is_parallel_draft_region & in_bounds + + # Compute indices of new tokens (bonus + parallel drafting) for sampling + # New tokens are at positions + # [num_valid_tokens, num_valid_tokens + num_padding_slots_per_request) + is_new_token_region = (j >= num_valid_tokens) & ( + j < num_valid_tokens + num_padding_slots_per_request + ) + new_token_local_idx = ( + j - num_valid_tokens + ) # 0 for bonus, 1, 2, ... for parallel drafting + new_token_out_idx = ( + request_idx * num_padding_slots_per_request + new_token_local_idx + ) + + # Compute hidden state mapping (source index -> destination index) + # This maps each input position to its corresponding output position + # Hidden states don't get shifted, so we map all input tokens (including rejected) + if shift_input_ids: + num_input_tokens_this_request = next_query_start_loc - query_start_loc + is_input_region = j < num_input_tokens_this_request + src_idx = query_start_loc + j + tl.store(out_hidden_state_mapping_ptr + src_idx, out_idx, mask=is_input_region) + + # Store outputs + tl.store(out_input_ids_ptr + out_idx, token_ids, mask=in_bounds) + tl.store(out_positions_ptr + out_idx, positions, mask=in_bounds) + tl.store(out_is_rejected_token_mask_ptr + out_idx, is_rejected_out, mask=in_bounds) + tl.store(out_is_masked_token_mask_ptr + out_idx, is_masked_out, mask=in_bounds) + tl.store( + out_new_token_indices_ptr + new_token_out_idx, + out_idx, + mask=is_new_token_region & in_bounds, + ) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 10d4dfd3309..6b04774a8b6 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -4090,7 +4090,7 @@ class GPUModelRunner( target_positions=target_positions, target_hidden_states=target_hidden_states, next_token_ids=next_token_ids, - last_token_indices=token_indices_to_sample, + token_indices_to_sample=token_indices_to_sample, sampling_metadata=sampling_metadata, common_attn_metadata=common_attn_metadata, mm_embed_inputs=mm_embed_inputs, From 7d8c6804e2654873cb25d0b23fef178fa5f37237 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Thu, 5 Feb 2026 18:42:40 +0100 Subject: [PATCH 104/810] [Misc] Add debug logs (#33931) Signed-off-by: NickLucche --- vllm/distributed/kv_transfer/kv_connector/utils.py | 2 ++ vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py | 3 +++ 2 files changed, 5 insertions(+) diff --git a/vllm/distributed/kv_transfer/kv_connector/utils.py b/vllm/distributed/kv_transfer/kv_connector/utils.py index 019201ede73..f9367da7371 100644 --- a/vllm/distributed/kv_transfer/kv_connector/utils.py +++ b/vllm/distributed/kv_transfer/kv_connector/utils.py @@ -324,6 +324,7 @@ class TpKVTopology: kv_cache_shape = self.attn_backend.get_kv_cache_shape( num_blocks=1, block_size=_MOCK_BLOCK_SIZE, num_kv_heads=1, head_size=1 ) + logger.debug("Test kv_cache_shape: %s", kv_cache_shape) # Non-MLA backends caches have 5 dims [2, num_blocks, H,N,D], # we just mock num_blocks to 1 for the dimension check below. self._is_kv_layout_blocks_first = ( @@ -337,6 +338,7 @@ class TpKVTopology: ) if self._cross_layers_blocks: + logger.debug("Using cross-layer KV cache") # prepend layers dimension _MOCK_NUM_LAYERS = 80 kv_cache_shape = (_MOCK_NUM_LAYERS,) + kv_cache_shape diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py index 8ce939ee405..3a8400447e8 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py @@ -1354,6 +1354,9 @@ class NixlConnectorWorker: if base_addr in seen_base_addresses: continue + logger.debug( + "Registering layer %s with cache shape: %s", layer_name, cache.shape + ) kernel_block_size = cache.shape[self.kv_topo.block_size_position] if self.block_size != kernel_block_size: logger.info_once( From 1ee95841bd251f9081c3a317984c4dcaa003b3c0 Mon Sep 17 00:00:00 2001 From: zackyoray Date: Thu, 5 Feb 2026 19:51:58 +0200 Subject: [PATCH 105/810] [Bugfix] Fix swapped engine_ids in NIXL Llama 4 local attention path (#33795) Signed-off-by: Yoray Zack --- .../kv_transfer/kv_connector/v1/nixl_connector.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py index 3a8400447e8..c2777b3936b 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py @@ -2323,15 +2323,15 @@ class NixlConnectorWorker: # Get descs ids for the layer. layer_local_desc_ids = self._get_block_descs_ids( - dst_engine_id, + self.engine_id, layer_local_block_ids, layer_idx, + block_size_ratio=block_size_ratio, ) layer_remote_desc_ids = self._get_block_descs_ids( - self.engine_id, + dst_engine_id, layer_remote_block_ids, layer_idx, - block_size_ratio=block_size_ratio, ) local_descs_list.append(layer_local_desc_ids) From a57c8228ffb3ff82b983b32b71ff62a837255129 Mon Sep 17 00:00:00 2001 From: bnellnm <49004751+bnellnm@users.noreply.github.com> Date: Thu, 5 Feb 2026 13:07:18 -0500 Subject: [PATCH 106/810] [Moe Refactor] Make Inplace Flag for FusedMoEModularKernel part of the constructor (#33375) Signed-off-by: Bill Nell Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> --- .../moe/modular_kernel_tools/common.py | 1 + tests/kernels/moe/test_batched_deepgemm.py | 14 +++++++--- tests/kernels/moe/test_block_fp8.py | 5 ++-- tests/kernels/moe/test_cutlass_moe.py | 2 ++ tests/kernels/moe/test_deepep_deepgemm_moe.py | 15 +++++++---- tests/kernels/moe/test_deepep_moe.py | 7 +++-- tests/kernels/moe/test_deepgemm.py | 2 +- tests/kernels/moe/test_flashinfer.py | 2 +- tests/kernels/moe/test_flashinfer_moe.py | 1 + .../moe/test_modular_oai_triton_moe.py | 7 +++-- tests/kernels/moe/test_moe.py | 9 ++++--- tests/kernels/moe/test_nvfp4_moe.py | 1 + tests/kernels/moe/test_pplx_cutlass_moe.py | 1 + tests/kernels/moe/test_pplx_moe.py | 1 + tests/kernels/moe/utils.py | 24 ++++++++++++++++- .../model_executor/layers/fused_moe/config.py | 9 ++++--- .../layers/fused_moe/cutlass_moe.py | 1 + .../layers/fused_moe/fused_marlin_moe.py | 6 ++--- .../layers/fused_moe/fused_moe.py | 26 ++++--------------- .../layers/fused_moe/fused_moe_method_base.py | 4 --- .../fused_moe/fused_moe_modular_method.py | 7 ++--- vllm/model_executor/layers/fused_moe/layer.py | 11 +++++++- .../layers/fused_moe/modular_kernel.py | 9 ++++--- .../layers/fused_moe/oracle/fp8.py | 10 ++++--- .../layers/fused_moe/oracle/nvfp4.py | 1 + .../layers/fused_moe/oracle/unquantized.py | 14 +++++----- .../fused_moe/unquantized_fused_moe_method.py | 7 +---- .../layers/quantization/awq_marlin.py | 1 + .../layers/quantization/bitsandbytes.py | 2 +- .../compressed_tensors_moe.py | 11 ++++---- .../layers/quantization/experts_int8.py | 2 +- .../model_executor/layers/quantization/fp8.py | 7 +---- .../layers/quantization/gptq_marlin.py | 1 + .../layers/quantization/modelopt.py | 4 +-- .../layers/quantization/moe_wna16.py | 2 +- .../layers/quantization/mxfp4.py | 5 +--- .../layers/quantization/quark/quark_moe.py | 9 +++---- 37 files changed, 132 insertions(+), 109 deletions(-) diff --git a/tests/kernels/moe/modular_kernel_tools/common.py b/tests/kernels/moe/modular_kernel_tools/common.py index 327cd44f612..893968b5cd9 100644 --- a/tests/kernels/moe/modular_kernel_tools/common.py +++ b/tests/kernels/moe/modular_kernel_tools/common.py @@ -620,6 +620,7 @@ def make_modular_kernel( modular_kernel = mk.FusedMoEModularKernel( prepare_finalize=prepare_finalize, fused_experts=fused_experts, + inplace=False, ) return modular_kernel diff --git a/tests/kernels/moe/test_batched_deepgemm.py b/tests/kernels/moe/test_batched_deepgemm.py index 081a5fd0b93..2c6c45a5f23 100644 --- a/tests/kernels/moe/test_batched_deepgemm.py +++ b/tests/kernels/moe/test_batched_deepgemm.py @@ -74,7 +74,11 @@ def test_batched_deepgemm_vs_triton( quant_config=quant_config, moe_config=make_dummy_moe_config(), ) - mk_triton = FusedMoEModularKernel(prep_finalize, triton_experts) + mk_triton = FusedMoEModularKernel( + prep_finalize, + triton_experts, + inplace=False, + ) out_triton = mk_triton( hidden_states=a, @@ -82,7 +86,6 @@ def test_batched_deepgemm_vs_triton( w2=w2, topk_weights=topk_weights, topk_ids=topk_ids, - inplace=False, global_num_experts=E, ) @@ -93,7 +96,11 @@ def test_batched_deepgemm_vs_triton( quant_config=quant_config, moe_config=make_dummy_moe_config(), ) - mk_deepgemm = FusedMoEModularKernel(prep_finalize, deepgemm_experts) + mk_deepgemm = FusedMoEModularKernel( + prep_finalize, + deepgemm_experts, + inplace=False, + ) out_deepgemm = mk_deepgemm( hidden_states=a, @@ -101,7 +108,6 @@ def test_batched_deepgemm_vs_triton( w2=w2, topk_weights=topk_weights, topk_ids=topk_ids, - inplace=False, global_num_experts=E, ) diff --git a/tests/kernels/moe/test_block_fp8.py b/tests/kernels/moe/test_block_fp8.py index 508df9e328a..66508568ed2 100644 --- a/tests/kernels/moe/test_block_fp8.py +++ b/tests/kernels/moe/test_block_fp8.py @@ -9,6 +9,7 @@ from tests.kernels.moe.utils import ( make_dummy_moe_config, make_test_quant_config, make_test_weights, + modular_triton_fused_moe, ) from tests.kernels.quant_utils import ( native_per_token_group_quant_fp8, @@ -26,9 +27,6 @@ from vllm.model_executor.layers.fused_moe.config import ( from vllm.model_executor.layers.fused_moe.deep_gemm_moe import ( _valid_deep_gemm_shape, ) -from vllm.model_executor.layers.fused_moe.fused_moe import ( - modular_triton_fused_moe, -) from vllm.model_executor.layers.fused_moe.prepare_finalize import ( MoEPrepareAndFinalizeNoEP, ) @@ -261,6 +259,7 @@ def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed, monkeypatch) moe_config=make_dummy_moe_config(), quant_config=quant_config, ), + inplace=False, ) def deep_gemm_moe_fp8(a, w1, w2, w1_s, w2_s, topk_weights, topk_ids): diff --git a/tests/kernels/moe/test_cutlass_moe.py b/tests/kernels/moe/test_cutlass_moe.py index 3a5a66a383d..d232d00fcbb 100644 --- a/tests/kernels/moe/test_cutlass_moe.py +++ b/tests/kernels/moe/test_cutlass_moe.py @@ -207,6 +207,7 @@ def run_with_expert_maps( ), quant_config=new_quant_config, ), + inplace=False, ) out_tensor = out_tensor + kernel(**kwargs) @@ -266,6 +267,7 @@ def run_8_bit( ), quant_config=quant_config, ), + inplace=False, ) return kernel(**kwargs) diff --git a/tests/kernels/moe/test_deepep_deepgemm_moe.py b/tests/kernels/moe/test_deepep_deepgemm_moe.py index 1bf5ced2e84..11f5357157d 100644 --- a/tests/kernels/moe/test_deepep_deepgemm_moe.py +++ b/tests/kernels/moe/test_deepep_deepgemm_moe.py @@ -194,8 +194,11 @@ def make_ll_modular_kernel( quant_config=quant_config, moe_config=make_dummy_moe_config(), ) - mk = FusedMoEModularKernel(prepare_finalize=a2a, fused_experts=fused_experts) - return mk + return FusedMoEModularKernel( + prepare_finalize=a2a, + fused_experts=fused_experts, + inplace=False, + ) def make_ht_modular_kernel( @@ -224,8 +227,11 @@ def make_ht_modular_kernel( moe_config=make_dummy_moe_config(), quant_config=quant_config, ) - mk = FusedMoEModularKernel(prepare_finalize=a2a, fused_experts=fused_experts) - return mk + return FusedMoEModularKernel( + prepare_finalize=a2a, + fused_experts=fused_experts, + inplace=False, + ) def make_modular_kernel( @@ -318,7 +324,6 @@ def deepep_deepgemm_moe_impl( w2=w2, topk_weights=test_tensors.topk_weights, topk_ids=test_tensors.topk, - inplace=False, activation="silu", global_num_experts=num_experts, expert_map=build_expert_map(), diff --git a/tests/kernels/moe/test_deepep_moe.py b/tests/kernels/moe/test_deepep_moe.py index f740f5bf958..8d3ca165076 100644 --- a/tests/kernels/moe/test_deepep_moe.py +++ b/tests/kernels/moe/test_deepep_moe.py @@ -179,7 +179,11 @@ def make_modular_kernel( quant_config=quant_config, ) - mk = FusedMoEModularKernel(prepare_finalize=a2a, fused_experts=fused_experts) + mk = FusedMoEModularKernel( + prepare_finalize=a2a, + fused_experts=fused_experts, + inplace=False, + ) return mk @@ -256,7 +260,6 @@ def deep_ep_moe_impl( w2=w2, topk_weights=topk_weights_chunk, topk_ids=topk_chunk, - inplace=False, activation="silu", global_num_experts=num_experts, expert_map=build_expert_map(), diff --git a/tests/kernels/moe/test_deepgemm.py b/tests/kernels/moe/test_deepgemm.py index 729b54753b0..7f9bccb739e 100644 --- a/tests/kernels/moe/test_deepgemm.py +++ b/tests/kernels/moe/test_deepgemm.py @@ -115,6 +115,7 @@ def run_single_case(m, n, k, topk, num_experts, block_size): moe_config=make_dummy_moe_config(), quant_config=quant_config, ), + inplace=False, ) # triton reference @@ -135,7 +136,6 @@ def run_single_case(m, n, k, topk, num_experts, block_size): w2=w2, topk_weights=topk_weights, topk_ids=topk_ids, - inplace=False, ) diff = calc_diff(out_deepgemm, out_triton) assert diff < 0.001, f"Diff exceeded 1%: {diff}" diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index 1c512b5b168..e62cf79418c 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -301,6 +301,7 @@ def test_flashinfer_cutlass_moe_fp8_no_graph( moe_config=moe_config, quant_config=quant_config, ), + inplace=False, ) flashinfer_cutlass_output = kernel( @@ -309,7 +310,6 @@ def test_flashinfer_cutlass_moe_fp8_no_graph( td.layer.w2_weight, topk_weights, topk_ids, - inplace=False, activation=activation, global_num_experts=e, expert_map=None, diff --git a/tests/kernels/moe/test_flashinfer_moe.py b/tests/kernels/moe/test_flashinfer_moe.py index 9bb61ddfa0f..113649afe2f 100644 --- a/tests/kernels/moe/test_flashinfer_moe.py +++ b/tests/kernels/moe/test_flashinfer_moe.py @@ -108,6 +108,7 @@ def test_flashinfer_fp4_moe_no_graph( flashinfer_experts = FusedMoEModularKernel( MoEPrepareAndFinalizeNoEP(), FlashInferExperts(moe_config=moe_config, quant_config=quant_config), + inplace=False, ) fi_activation = {"silu_and_mul": "silu", "relu2": "relu2_no_mul"}[activation] diff --git a/tests/kernels/moe/test_modular_oai_triton_moe.py b/tests/kernels/moe/test_modular_oai_triton_moe.py index 38022e0e61b..bebf18ef0aa 100644 --- a/tests/kernels/moe/test_modular_oai_triton_moe.py +++ b/tests/kernels/moe/test_modular_oai_triton_moe.py @@ -180,7 +180,11 @@ def oai_triton_moe_impl( else: fused_experts = OAITritonExperts(make_dummy_moe_config(), quant_config) - mk = FusedMoEModularKernel(MoEPrepareAndFinalizeNoEP(), fused_experts) + mk = FusedMoEModularKernel( + MoEPrepareAndFinalizeNoEP(), + fused_experts, + inplace=False, + ) return mk.forward( hidden_states=x, @@ -188,7 +192,6 @@ def oai_triton_moe_impl( w2=w2, topk_weights=topk_weights, topk_ids=topk_ids, - inplace=True, activation="swigluoai", global_num_experts=num_experts, expert_map=None, diff --git a/tests/kernels/moe/test_moe.py b/tests/kernels/moe/test_moe.py index a304e70fc44..53fb43e3c12 100644 --- a/tests/kernels/moe/test_moe.py +++ b/tests/kernels/moe/test_moe.py @@ -18,7 +18,11 @@ from transformers import MixtralConfig from transformers.models.mixtral.modeling_mixtral import MixtralSparseMoeBlock import vllm.model_executor.layers.fused_moe # noqa -from tests.kernels.moe.utils import fused_moe, make_dummy_moe_config +from tests.kernels.moe.utils import ( + fused_moe, + make_dummy_moe_config, + modular_triton_fused_moe, +) from tests.kernels.utils import opcheck, stack_and_dev, torch_experts, torch_moe from vllm._aiter_ops import rocm_aiter_ops from vllm.config import VllmConfig, set_current_vllm_config @@ -36,9 +40,6 @@ from vllm.model_executor.layers.fused_moe.fused_marlin_moe import ( batched_fused_marlin_moe, fused_marlin_moe, ) -from vllm.model_executor.layers.fused_moe.fused_moe import ( - modular_triton_fused_moe, -) from vllm.model_executor.layers.quantization.utils.marlin_utils import ( marlin_permute_bias, ) diff --git a/tests/kernels/moe/test_nvfp4_moe.py b/tests/kernels/moe/test_nvfp4_moe.py index a22b2088bb0..10678e37624 100644 --- a/tests/kernels/moe/test_nvfp4_moe.py +++ b/tests/kernels/moe/test_nvfp4_moe.py @@ -95,6 +95,7 @@ def test_cutlass_fp4_moe_no_graph( moe_config=make_dummy_moe_config(), quant_config=quant_config, ), + inplace=False, ) cutlass_output = kernel( diff --git a/tests/kernels/moe/test_pplx_cutlass_moe.py b/tests/kernels/moe/test_pplx_cutlass_moe.py index ef37c1c7443..213d28cda77 100644 --- a/tests/kernels/moe/test_pplx_cutlass_moe.py +++ b/tests/kernels/moe/test_pplx_cutlass_moe.py @@ -172,6 +172,7 @@ def pplx_cutlass_moe( fused_cutlass_experts = FusedMoEModularKernel( prepare_finalize, experts, + inplace=False, ) a_chunk = chunk_by_rank(a, rank, world_size).to(device) diff --git a/tests/kernels/moe/test_pplx_moe.py b/tests/kernels/moe/test_pplx_moe.py index 08519087e1c..deb3b9eb4d7 100644 --- a/tests/kernels/moe/test_pplx_moe.py +++ b/tests/kernels/moe/test_pplx_moe.py @@ -592,6 +592,7 @@ def pplx_moe( prepare_finalize, experts, shared_experts, + inplace=False, ) # Note: for now use_compile will error out if the problem size is diff --git a/tests/kernels/moe/utils.py b/tests/kernels/moe/utils.py index 4883085cb83..897bfddce5e 100644 --- a/tests/kernels/moe/utils.py +++ b/tests/kernels/moe/utils.py @@ -7,7 +7,11 @@ import vllm._custom_ops as ops from tests.kernels.quant_utils import per_block_cast_to_int8 from tests.kernels.quantization.nvfp4_utils import FLOAT4_E2M1_MAX, FLOAT8_E4M3_MAX from vllm.model_executor.layers.activation import SiluAndMul -from vllm.model_executor.layers.fused_moe import fused_experts, fused_topk +from vllm.model_executor.layers.fused_moe import ( + TritonExperts, + fused_experts, + fused_topk, +) from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, FusedMoEParallelConfig, @@ -20,6 +24,9 @@ from vllm.model_executor.layers.fused_moe.fused_batched_moe import ( NaiveBatchedExperts, ) from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel +from vllm.model_executor.layers.fused_moe.prepare_finalize import ( + MoEPrepareAndFinalizeNoEP, +) from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input from vllm.utils.deep_gemm import per_block_cast_to_fp8 from vllm.utils.math_utils import round_up @@ -116,6 +123,7 @@ def batched_moe( quant_config=quant_config, moe_config=make_dummy_moe_config(), ), + inplace=False, ) return fused_experts(a, w1, w2, topk_weight, topk_ids) @@ -157,6 +165,7 @@ def naive_batched_moe( quant_config=quant_config, moe_config=make_dummy_moe_config(), ), + inplace=False, ) return fused_experts(a, w1, w2, topk_weight, topk_ids) @@ -554,3 +563,16 @@ def make_shared_experts( return RealMLP(K, N, w1, w2, "silu", quant_config, w1_s=w1_s, w2_s=w2_s) finally: torch.set_default_dtype(old_dtype) + + +def modular_triton_fused_moe( + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + shared_experts: torch.nn.Module | None = None, +) -> FusedMoEModularKernel: + return FusedMoEModularKernel( + MoEPrepareAndFinalizeNoEP(), + TritonExperts(moe_config, quant_config), + shared_experts, + inplace=False, + ) diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 6650367da03..3a8c13b3a19 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1083,13 +1083,16 @@ class FusedMoEConfig: router_logits_dtype: torch.dtype | None = None max_num_tokens: int = envs.VLLM_MOE_DP_CHUNK_SIZE - has_bias: bool = False - is_act_and_mul: bool = True - is_lora_enabled: bool = False + # This flag is used to disable the inplace optimization + # in MoE kernels. If this flag is True then the kernel + # should not be using inplace. If the flag is false, the + # kernel is free to use inplace or not. + disable_inplace: bool = True + def __post_init__(self): if self.dp_size > 1: logger.debug_once( diff --git a/vllm/model_executor/layers/fused_moe/cutlass_moe.py b/vllm/model_executor/layers/fused_moe/cutlass_moe.py index 74f05a2c0f8..ac5a860679f 100644 --- a/vllm/model_executor/layers/fused_moe/cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/cutlass_moe.py @@ -1165,6 +1165,7 @@ def cutlass_moe_w4a8_fp8( quant_config=quant_config, group_size=group_size, ), + inplace=False, ) return fn( diff --git a/vllm/model_executor/layers/fused_moe/fused_marlin_moe.py b/vllm/model_executor/layers/fused_moe/fused_marlin_moe.py index 2e5167bdfd5..8d95665f763 100644 --- a/vllm/model_executor/layers/fused_moe/fused_marlin_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_marlin_moe.py @@ -267,6 +267,7 @@ def fused_marlin_moe( if inplace: assert output is None, "Conflicting request" + assert not disable_inplace() quant_type = ScalarType.from_id(quant_type_id) assert quant_type in [ @@ -356,10 +357,7 @@ def fused_marlin_moe( ).view(-1, topk, K) if output is None: - if inplace and not disable_inplace(): - output = hidden_states - else: - output = torch.empty_like(hidden_states) + output = hidden_states if inplace else torch.empty_like(hidden_states) if moe_sum is None: return torch.sum(moe_output.view(-1, topk, K), dim=1, out=output) diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py index 120b3c2d1e2..e0907368bbc 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe.py @@ -27,9 +27,6 @@ from vllm.model_executor.layers.fused_moe.config import ( from vllm.model_executor.layers.fused_moe.moe_align_block_size import ( moe_align_block_size, ) -from vllm.model_executor.layers.fused_moe.prepare_finalize import ( - MoEPrepareAndFinalizeNoEP, -) from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( TopKWeightAndReduceNoOP, ) @@ -1511,7 +1508,7 @@ def torch_vllm_outplace_fused_experts(**kwargs) -> torch.Tensor: def dispatch_fused_experts_func(inplace: bool) -> Callable[..., torch.Tensor]: - if inplace and not disable_inplace(): + if inplace: return torch_vllm_inplace_fused_experts return torch_vllm_outplace_fused_experts @@ -1534,6 +1531,8 @@ def fused_experts( if quant_config is None: quant_config = FUSED_MOE_UNQUANTIZED_CONFIG + assert not inplace or not disable_inplace() + return dispatch_fused_experts_func(inplace)( hidden_states=hidden_states, w1=w1, @@ -1593,7 +1592,7 @@ def fused_experts_impl( w2: torch.Tensor, topk_weights: torch.Tensor, topk_ids: torch.Tensor, - inplace: bool = False, + inplace: bool, activation: str = "silu", apply_router_weight_on_input: bool = False, use_fp8_w8a8: bool = False, @@ -1712,10 +1711,7 @@ def fused_experts_impl( else: raise ValueError(f"Unsupported compute_type: {hidden_states.dtype}") - if inplace and not disable_inplace(): - out_hidden_states = hidden_states - else: - out_hidden_states = torch.empty_like(hidden_states) + out_hidden_states = hidden_states if inplace else torch.empty_like(hidden_states) if ocp_mx_scheme is not None: # TODO: On platforms for which `current_platform.supports_mx()` is True @@ -2291,15 +2287,3 @@ class TritonWNA16Experts(TritonExperts): # separate function is required for MoE + LoRA self.moe_sum(intermediate_cache3, output) - - -def modular_triton_fused_moe( - moe_config: FusedMoEConfig, - quant_config: FusedMoEQuantConfig, - shared_experts: torch.nn.Module | None = None, -) -> mk.FusedMoEModularKernel: - return mk.FusedMoEModularKernel( - MoEPrepareAndFinalizeNoEP(), - TritonExperts(moe_config, quant_config), - shared_experts, - ) diff --git a/vllm/model_executor/layers/fused_moe/fused_moe_method_base.py b/vllm/model_executor/layers/fused_moe/fused_moe_method_base.py index 3ad56cc4c2d..93db1c54571 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe_method_base.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe_method_base.py @@ -113,10 +113,6 @@ class FusedMoEMethodBase(QuantizeMethodBase): def supports_eplb(self) -> bool: return False - @property - def allow_inplace(self) -> bool: - return False - @property def method_name(self) -> str: return self.__class__.__name__ diff --git a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py index 7a2244a9bc1..c30eeb6dc2d 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py @@ -46,6 +46,7 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp): old_quant_method: FusedMoEMethodBase, prepare_finalize: FusedMoEPrepareAndFinalize, shared_experts: torch.nn.Module | None, + inplace: bool = False, ) -> "FusedMoEModularMethod": return FusedMoEModularMethod( old_quant_method, @@ -54,6 +55,7 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp): old_quant_method.select_gemm_impl(prepare_finalize, moe_layer), shared_experts, moe_parallel_config=moe_layer.moe_parallel_config, + inplace=inplace, ), ) @@ -61,10 +63,6 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp): def supports_eplb(self) -> bool: return self.old_quant_method.supports_eplb - @property - def allow_inplace(self) -> bool: - return self.old_quant_method.allow_inplace - @property def method_name(self) -> str: return self.old_quant_method.method_name @@ -99,7 +97,6 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp): w2=layer.w2_weight, topk_weights=topk_weights, topk_ids=topk_ids, - inplace=self.allow_inplace, activation=layer.activation, global_num_experts=layer.global_num_experts, apply_router_weight_on_input=layer.apply_router_weight_on_input, diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index b092cf6cf0e..3935fe374bc 100755 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -50,6 +50,9 @@ from vllm.model_executor.layers.fused_moe.router.router_factory import ( from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( UnquantizedFusedMoEMethod, ) +from vllm.model_executor.layers.fused_moe.utils import ( + disable_inplace, +) from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, ) @@ -560,6 +563,8 @@ class FusedMoE(CustomOp): activation=activation, device=vllm_config.device_config.device, routing_method=self.routing_method_type, + # TODO: in_dtype == out_dtype? + disable_inplace=disable_inplace() or self.shared_experts is not None, ) if self.use_mori_kernels: assert self.rocm_aiter_fmoe_enabled, ( @@ -650,7 +655,11 @@ class FusedMoE(CustomOp): "%s for %s(%s)", prepare_finalize.__class__.__name__, self, id(self) ) self.quant_method = FusedMoEModularMethod.make( - self, self.quant_method, prepare_finalize, self.shared_experts + self, + self.quant_method, + prepare_finalize, + self.shared_experts, + inplace=not self.moe_config.disable_inplace, ) @property diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 940a2c55f73..598374af295 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -811,11 +811,13 @@ class FusedMoEModularKernel(torch.nn.Module): fused_experts: FusedMoEPermuteExpertsUnpermute, shared_experts: torch.nn.Module | None = None, moe_parallel_config: FusedMoEParallelConfig | None = None, + inplace: bool = False, ): super().__init__() self.prepare_finalize = prepare_finalize self.fused_experts = fused_experts self.shared_experts = shared_experts + self.inplace = inplace # prefer an explicit FusedMoEParallelConfig when available (from # FusedMoE layers / tests). @@ -1292,7 +1294,6 @@ class FusedMoEModularKernel(torch.nn.Module): w2: torch.Tensor, topk_weights: torch.Tensor, topk_ids: torch.Tensor, - inplace: bool = False, activation: str = "silu", global_num_experts: int = -1, expert_map: torch.Tensor | None = None, @@ -1309,8 +1310,6 @@ class FusedMoEModularKernel(torch.nn.Module): - topk_weights (torch.Tensor): The topk weights applied at the end of the layer. - topk_ids (torch.Tensor): A map of row to expert id. - - inplace (bool): If True, perform the operation in-place. - Defaults to False. - activation (str): The activation function to apply after the first MoE layer. - global_num_experts (int): The total number of experts in the global @@ -1326,7 +1325,9 @@ class FusedMoEModularKernel(torch.nn.Module): - torch.Tensor: The output tensor after applying the MoE layer. """ - if inplace and self.shared_experts is None and not disable_inplace(): + if self.inplace: + assert self.shared_experts is None + assert not disable_inplace() output = hidden_states else: output = torch.zeros_like(hidden_states) diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index 70c2516747f..bc0fc9a887e 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -472,7 +472,7 @@ def make_fp8_moe_kernel( fp8_backend: Fp8MoeBackend, routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, shared_experts: torch.nn.Module | None = None, -) -> tuple[mk.FusedMoEModularKernel, bool]: +) -> mk.FusedMoEModularKernel: # Create Prepare/Finalize. prepare_finalize = maybe_make_prepare_finalize( moe=moe_config, @@ -512,8 +512,10 @@ def make_fp8_moe_kernel( else None ), moe_parallel_config=moe_config.moe_parallel_config, + inplace=( + not moe_config.disable_inplace + and fp8_backend != Fp8MoeBackend.FLASHINFER_CUTLASS + ), ) - # TODO(rob): update inplace logic to be part of the kernel. - inplace = fp8_backend != Fp8MoeBackend.FLASHINFER_CUTLASS - return kernel, inplace + return kernel diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 276d231eb2e..dc3ac61ad14 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -437,6 +437,7 @@ def make_nvfp4_moe_kernel( else None ), moe_parallel_config=moe_config.moe_parallel_config, + inplace=False, ) # TODO(rob): update inplace logic to be part of the kernel. diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index a8754d6d6e4..c4a19ecb61a 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -154,11 +154,9 @@ def make_unquantized_moe_kernel( backend: UnquantizedMoeBackend, quant_config: FusedMoEQuantConfig, moe_config: FusedMoEConfig, -) -> tuple[mk.FusedMoEModularKernel | None, bool]: - use_inplace = True - +) -> mk.FusedMoEModularKernel | None: if backend in UNSUPPORTED_BACKEND: - return None, use_inplace + return None if backend == UnquantizedMoeBackend.FLASHINFER_CUTLASS: from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( @@ -171,8 +169,9 @@ def make_unquantized_moe_kernel( moe_config=moe_config, quant_config=quant_config, ), + inplace=False, ) - use_inplace = False + elif backend == UnquantizedMoeBackend.AITER: from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( AiterExperts, @@ -184,6 +183,7 @@ def make_unquantized_moe_kernel( moe_config=moe_config, quant_config=quant_config, ), + inplace=not moe_config.disable_inplace, ) elif backend == UnquantizedMoeBackend.TRITON: from vllm.model_executor.layers.fused_moe import TritonExperts @@ -194,6 +194,7 @@ def make_unquantized_moe_kernel( moe_config=moe_config, quant_config=quant_config, ), + inplace=not moe_config.disable_inplace, ) elif backend == UnquantizedMoeBackend.XPU: from vllm.model_executor.layers.fused_moe import XPUExperts @@ -204,5 +205,6 @@ def make_unquantized_moe_kernel( moe_config=moe_config, quant_config=quant_config, ), + inplace=not moe_config.disable_inplace, ) - return kernel, use_inplace + return kernel diff --git a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py index 6fdd8ecf79b..8a35be78bc4 100644 --- a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py +++ b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py @@ -101,10 +101,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp): def supports_eplb(self) -> bool: return True - @property - def allow_inplace(self) -> bool: - return True - def maybe_make_prepare_finalize( self, routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, @@ -225,7 +221,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp): self.moe_quant_config = self.get_fused_moe_quant_config(layer) assert self.moe_quant_config is not None - self.kernel, self.use_inplace = make_unquantized_moe_kernel( + self.kernel = make_unquantized_moe_kernel( backend=self.unquantized_backend, quant_config=self.moe_quant_config, moe_config=self.moe, @@ -329,7 +325,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp): w2=layer.w2_weight, topk_weights=topk_weights, topk_ids=topk_ids, - inplace=self.use_inplace, activation=layer.activation, apply_router_weight_on_input=layer.apply_router_weight_on_input, global_num_experts=layer.global_num_experts, diff --git a/vllm/model_executor/layers/quantization/awq_marlin.py b/vllm/model_executor/layers/quantization/awq_marlin.py index 163ee78a718..642088a4536 100644 --- a/vllm/model_executor/layers/quantization/awq_marlin.py +++ b/vllm/model_executor/layers/quantization/awq_marlin.py @@ -785,4 +785,5 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase): w2_zeros=layer.w2_qzeros, workspace=layer.workspace, input_dtype=self.input_dtype, + inplace=not self.moe.disable_inplace, ) diff --git a/vllm/model_executor/layers/quantization/bitsandbytes.py b/vllm/model_executor/layers/quantization/bitsandbytes.py index 8b6b1e445f3..2fd567d7fae 100644 --- a/vllm/model_executor/layers/quantization/bitsandbytes.py +++ b/vllm/model_executor/layers/quantization/bitsandbytes.py @@ -515,7 +515,7 @@ class BitsAndBytesMoEMethod(FusedMoEMethodBase): w2=w2, topk_weights=topk_weights, topk_ids=topk_ids, - inplace=True, + inplace=not self.moe.disable_inplace, activation=layer.activation, apply_router_weight_on_input=layer.apply_router_weight_on_input, global_num_experts=layer.global_num_experts, diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py index 5152c5cccab..e25a415a593 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py @@ -357,7 +357,6 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod): layer.w2_weight, topk_weights, topk_ids, - inplace=False, activation=layer.activation, global_num_experts=layer.global_num_experts, expert_map=layer.expert_map, @@ -669,7 +668,6 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod): layer.w2_weight, topk_weights, topk_ids, - inplace=False, activation=layer.activation, global_num_experts=layer.global_num_experts, expert_map=layer.expert_map, @@ -960,7 +958,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod): self.moe_quant_config = self.get_fused_moe_quant_config(layer) if self.moe_quant_config: assert self.experts_cls is not None - self.moe_mk, self.use_inplace = make_fp8_moe_kernel( + self.moe_mk = make_fp8_moe_kernel( moe_quant_config=self.moe_quant_config, moe_config=self.moe, fp8_backend=self.fp8_backend, @@ -1073,7 +1071,6 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod): layer.w2_weight, topk_weights, topk_ids, - inplace=self.use_inplace, activation=layer.activation, global_num_experts=layer.global_num_experts, # TODO(rob): investigate the disable_expert_map introduced by: @@ -1212,7 +1209,7 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod): w2=layer.w2_weight, topk_weights=topk_weights, topk_ids=topk_ids, - inplace=True, + inplace=not self.moe.disable_inplace, activation=layer.activation, apply_router_weight_on_input=layer.apply_router_weight_on_input, global_num_experts=layer.global_num_experts, @@ -1739,6 +1736,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): workspace=layer.workspace, input_dtype=self.marlin_input_dtype, is_k_full=self.is_k_full, + inplace=not self.moe.disable_inplace, ) @@ -1969,7 +1967,7 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): layer.w2_weight_packed, topk_weights=topk_weights, topk_ids=topk_ids, - inplace=True, + inplace=not self.moe.disable_inplace, activation=layer.activation, apply_router_weight_on_input=layer.apply_router_weight_on_input, global_num_experts=layer.global_num_experts, @@ -2605,6 +2603,7 @@ class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod): s_strides1=self.s_strides1, s_strides2=self.s_strides2, group_size=self.group_size, + apply_router_weight_on_input=layer.apply_router_weight_on_input, ) @property diff --git a/vllm/model_executor/layers/quantization/experts_int8.py b/vllm/model_executor/layers/quantization/experts_int8.py index 5a0bb5d30f9..176bfe04095 100644 --- a/vllm/model_executor/layers/quantization/experts_int8.py +++ b/vllm/model_executor/layers/quantization/experts_int8.py @@ -149,7 +149,7 @@ class ExpertsInt8MoEMethod(FusedMoEMethodBase): layer.w2_weight, topk_weights=topk_weights, topk_ids=topk_ids, - inplace=True, + inplace=not self.moe.disable_inplace, activation=layer.activation, apply_router_weight_on_input=layer.apply_router_weight_on_input, global_num_experts=layer.global_num_experts, diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 8b9fe0f3e93..a612397066e 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -854,7 +854,7 @@ class Fp8MoEMethod(FusedMoEMethodBase): self.moe_quant_config = self.get_fused_moe_quant_config(layer) if self.moe_quant_config: assert self.experts_cls is not None - self.moe_mk, self.use_inplace = make_fp8_moe_kernel( + self.moe_mk = make_fp8_moe_kernel( moe_quant_config=self.moe_quant_config, moe_config=self.moe, fp8_backend=self.fp8_backend, @@ -958,10 +958,6 @@ class Fp8MoEMethod(FusedMoEMethodBase): def supports_eplb(self) -> bool: return True - @property - def allow_inplace(self) -> bool: - return True - @property def is_monolithic(self) -> bool: return self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM @@ -1032,7 +1028,6 @@ class Fp8MoEMethod(FusedMoEMethodBase): layer.w2_weight, topk_weights, topk_ids, - inplace=self.use_inplace, activation=layer.activation, global_num_experts=layer.global_num_experts, expert_map=layer.expert_map, diff --git a/vllm/model_executor/layers/quantization/gptq_marlin.py b/vllm/model_executor/layers/quantization/gptq_marlin.py index 698855c09f4..d18c7207dff 100644 --- a/vllm/model_executor/layers/quantization/gptq_marlin.py +++ b/vllm/model_executor/layers/quantization/gptq_marlin.py @@ -924,4 +924,5 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase): workspace=layer.workspace, is_k_full=self.is_k_full, input_dtype=self.input_dtype, + inplace=not self.moe.disable_inplace, ) diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index e76c109eced..4474e630b76 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -853,7 +853,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase): self.moe_quant_config = self.get_fused_moe_quant_config(layer) if self.moe_quant_config: assert self.experts_cls is not None - self.moe_mk, self.use_inplace = make_fp8_moe_kernel( + self.moe_mk = make_fp8_moe_kernel( moe_quant_config=self.moe_quant_config, moe_config=self.moe, fp8_backend=self.fp8_backend, @@ -967,7 +967,6 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase): layer.w2_weight, topk_weights, topk_ids, - inplace=self.use_inplace, activation=layer.activation, global_num_experts=layer.global_num_experts, expert_map=layer.expert_map, @@ -1538,7 +1537,6 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase): layer.w2_weight, topk_weights, topk_ids, - inplace=False, activation=layer.activation, global_num_experts=layer.global_num_experts, expert_map=layer.expert_map, diff --git a/vllm/model_executor/layers/quantization/moe_wna16.py b/vllm/model_executor/layers/quantization/moe_wna16.py index 34628591ffb..bca2516d4ed 100644 --- a/vllm/model_executor/layers/quantization/moe_wna16.py +++ b/vllm/model_executor/layers/quantization/moe_wna16.py @@ -378,7 +378,7 @@ class MoeWNA16Method(FusedMoEMethodBase): layer.w2_qweight, topk_weights=topk_weights, topk_ids=topk_ids, - inplace=True, + inplace=not self.moe.disable_inplace, apply_router_weight_on_input=layer.apply_router_weight_on_input, global_num_experts=layer.global_num_experts, expert_map=layer.expert_map, diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index a50fa4beea3..50009445d9b 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -881,10 +881,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): f"Incompatible Mxfp4 backend ({self.mxfp4_backend}) for EP" ) - @property - def allow_inplace(self) -> bool: - return True - @property def is_monolithic(self) -> bool: return ( @@ -923,6 +919,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): activation=layer.activation, expert_map=layer.expert_map, input_dtype=self.marlin_input_dtype, + inplace=not self.moe.disable_inplace, ) assert _can_support_mxfp4( diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index d2f0213e809..fc836c56be1 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -388,6 +388,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod): apply_router_weight_on_input=layer.apply_router_weight_on_input, global_num_experts=layer.global_num_experts, expert_map=layer.expert_map, + inplace=not self.moe.disable_inplace, ) else: from vllm.model_executor.layers.fused_moe import fused_experts @@ -398,7 +399,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod): w2=layer.w2_weight, topk_weights=topk_weights, topk_ids=topk_ids, - inplace=True, + inplace=not self.moe.disable_inplace, activation=layer.activation, apply_router_weight_on_input=layer.apply_router_weight_on_input, global_num_experts=layer.global_num_experts, @@ -734,10 +735,6 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod): block_shape=None, ) - @property - def allow_inplace(self) -> bool: - return True - def apply( self, layer: FusedMoE, @@ -769,7 +766,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod): layer.w2_weight, topk_weights=topk_weights, topk_ids=topk_ids, - inplace=True, + inplace=not self.moe.disable_inplace, activation=layer.activation, global_num_experts=layer.global_num_experts, apply_router_weight_on_input=layer.apply_router_weight_on_input, From 87d0d17ab583740bce777f334a6281edf9822e78 Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Fri, 6 Feb 2026 02:29:20 +0800 Subject: [PATCH 107/810] [Models] Consolidate Deepseek-OCR2 processor (#33909) Signed-off-by: Isotr0py --- vllm/model_executor/models/deepencoder2.py | 2 +- vllm/model_executor/models/deepseek_ocr.py | 12 +- vllm/model_executor/models/deepseek_ocr2.py | 16 +- .../processors/deepseek_ocr.py | 38 ++- .../processors/deepseek_ocr2.py | 320 ------------------ 5 files changed, 52 insertions(+), 336 deletions(-) delete mode 100644 vllm/transformers_utils/processors/deepseek_ocr2.py diff --git a/vllm/model_executor/models/deepencoder2.py b/vllm/model_executor/models/deepencoder2.py index b50606d4745..f134249ebfb 100644 --- a/vllm/model_executor/models/deepencoder2.py +++ b/vllm/model_executor/models/deepencoder2.py @@ -31,7 +31,7 @@ class CustomQwen2Decoder(nn.Module): num_key_value_heads: int = 2, intermediate_size: int = 4864, vocab_size: int = 151936, - attn_implementation: str = "sdpa", # ⭐ + attn_implementation: str = "sdpa", rms_norm_eps: float = 1e-06, rope_theta: float = 1000000.0, attention_dropout: float = 0.0, diff --git a/vllm/model_executor/models/deepseek_ocr.py b/vllm/model_executor/models/deepseek_ocr.py index 570ab548442..3425b15709e 100644 --- a/vllm/model_executor/models/deepseek_ocr.py +++ b/vllm/model_executor/models/deepseek_ocr.py @@ -52,7 +52,6 @@ from vllm.transformers_utils.configs.deepseek_vl2 import DeepseekVLV2Config from vllm.transformers_utils.processors.deepseek_ocr import ( BASE_SIZE, CROP_MODE, - IMAGE_SIZE, DeepseekOCRProcessor, count_tiles, ) @@ -66,6 +65,7 @@ from .deepencoder import DeepCLIPVisionTransformer, build_sam_vit_b from .deepseek_vl2 import MlpProjector # The image token id may be various +IMAGE_SIZE = 640 _IMAGE_TOKEN = "" @@ -190,7 +190,15 @@ class DeepseekOCRProcessingInfo(BaseProcessingInfo): return self.ctx.get_hf_config(DeepseekVLV2Config) def get_hf_processor(self, **kwargs: object): - return self.ctx.get_hf_processor(DeepseekOCRProcessor, **kwargs) + v1_processor_config = dict( + image_size=IMAGE_SIZE, + base_size=BASE_SIZE, + crop_mode=CROP_MODE, + strategy="v1", + ) + return self.ctx.get_hf_processor( + DeepseekOCRProcessor, **{**kwargs, **v1_processor_config} + ) def get_supported_mm_limits(self) -> Mapping[str, int | None]: return {"image": None} diff --git a/vllm/model_executor/models/deepseek_ocr2.py b/vllm/model_executor/models/deepseek_ocr2.py index 6541edad20c..cead4368541 100644 --- a/vllm/model_executor/models/deepseek_ocr2.py +++ b/vllm/model_executor/models/deepseek_ocr2.py @@ -48,11 +48,10 @@ from vllm.multimodal.processing import ( from vllm.sequence import IntermediateTensors from vllm.tokenizers import cached_tokenizer_from_config from vllm.transformers_utils.configs.deepseek_vl2 import DeepseekVLV2Config -from vllm.transformers_utils.processors.deepseek_ocr2 import ( +from vllm.transformers_utils.processors.deepseek_ocr import ( BASE_SIZE, CROP_MODE, - IMAGE_SIZE, - DeepseekOCR2Processor, + DeepseekOCRProcessor, ) from ...transformers_utils.processors.deepseek_ocr import count_tiles @@ -62,6 +61,7 @@ from .deepseek_ocr import DeepseekOCRImagePixelInputs from .deepseek_vl2 import MlpProjector # The image token id may be various +IMAGE_SIZE = 768 # different from deepseek-ocr _IMAGE_TOKEN = "" @@ -70,7 +70,15 @@ class DeepseekOCR2ProcessingInfo(BaseProcessingInfo): return self.ctx.get_hf_config(DeepseekVLV2Config) def get_hf_processor(self, **kwargs: object): - return self.ctx.get_hf_processor(DeepseekOCR2Processor, **kwargs) + v2_processor_config = dict( + image_size=IMAGE_SIZE, + base_size=BASE_SIZE, + crop_mode=CROP_MODE, + strategy="v2", + ) + return self.ctx.get_hf_processor( + DeepseekOCRProcessor, **{**kwargs, **v2_processor_config} + ) def get_supported_mm_limits(self) -> Mapping[str, int | None]: return {"image": None} diff --git a/vllm/transformers_utils/processors/deepseek_ocr.py b/vllm/transformers_utils/processors/deepseek_ocr.py index bb7aa0c1748..77e49483640 100644 --- a/vllm/transformers_utils/processors/deepseek_ocr.py +++ b/vllm/transformers_utils/processors/deepseek_ocr.py @@ -1,7 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # adapted from https://github.com/deepseek-ai/DeepSeek-OCR/blob/main/DeepSeek-OCR-master/DeepSeek-OCR-vllm/process/image_process.py +# and https://github.com/deepseek-ai/DeepSeek-OCR-2/blob/main/DeepSeek-OCR2-master/DeepSeek-OCR2-vllm/process/image_process.py import math +from typing import Literal import torch import torchvision.transforms as T @@ -156,10 +158,19 @@ class DeepseekOCRProcessor(ProcessorMixin): sft_format: str = "deepseek", mask_prompt: bool = True, ignore_id: int = -100, + image_size: int = IMAGE_SIZE, + base_size: int = BASE_SIZE, + strategy: Literal["v1", "v2"] = "v1", **kwargs, ): - self.image_size = IMAGE_SIZE - self.base_size = BASE_SIZE + self.image_size = image_size + self.base_size = base_size + + # image token calculation strategy for + # Deepseek-OCR and Deepseek-OCR-2 + self.strategy = strategy + assert strategy in ["v1", "v2"], "Only 'v1' and 'v2' strategies are supported." + self.patch_size = 16 self.image_mean = image_mean self.image_std = image_std @@ -317,16 +328,16 @@ class DeepseekOCRProcessor(ProcessorMixin): image_shapes.append(image.size) images_crop_raw = [] - if image.size[0] <= 640 and image.size[1] <= 640: + if image.size[0] <= self.image_size and image.size[1] <= self.image_size: crop_ratio = [1, 1] elif cropping: images_crop_raw, crop_ratio = dynamic_preprocess( - image, image_size=IMAGE_SIZE + image, image_size=self.image_size ) else: crop_ratio = [1, 1] - if self.image_size <= 640 and not cropping: + if not cropping: image = image.resize((self.image_size, self.image_size)) global_view = ImageOps.pad( @@ -350,12 +361,21 @@ class DeepseekOCRProcessor(ProcessorMixin): (self.base_size // self.patch_size) / self.downsample_ratio ) - tokenized_image = ( - [self.image_token_id] * num_queries_base + [self.image_token_id] - ) * num_queries_base + num_tokens_base = ( + (num_queries_base * (num_queries_base + 1)) + if self.strategy == "v1" + else num_queries_base * num_queries_base + ) + tokenized_image = [self.image_token_id] * num_tokens_base + tokenized_image += [self.image_token_id] if num_width_tiles > 1 or num_height_tiles > 1: - local_row = [self.image_token_id] * (num_queries * num_width_tiles + 1) + num_tokens_per_row = ( + num_queries * num_width_tiles + 1 + if self.strategy == "v1" + else num_queries * num_width_tiles + ) + local_row = [self.image_token_id] * num_tokens_per_row tokenized_image += local_row * (num_queries * num_height_tiles) tokenized_str += tokenized_image images_seq_mask += [True] * len(tokenized_image) diff --git a/vllm/transformers_utils/processors/deepseek_ocr2.py b/vllm/transformers_utils/processors/deepseek_ocr2.py deleted file mode 100644 index 6dbda73d4c7..00000000000 --- a/vllm/transformers_utils/processors/deepseek_ocr2.py +++ /dev/null @@ -1,320 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# adapted from https://github.com/deepseek-ai/DeepSeek-OCR/blob/main/DeepSeek-OCR-master/DeepSeek-OCR-vllm/process/image_process.py -import math - -import torch -from PIL import Image, ImageOps -from transformers import AutoProcessor, BatchFeature, LlamaTokenizerFast -from transformers.processing_utils import ProcessorMixin - -from vllm.transformers_utils.processors.deepseek_ocr import ( - ImageTransform, - dynamic_preprocess, -) - -BASE_SIZE = 1024 -IMAGE_SIZE = 768 -CROP_MODE = True -MIN_CROPS = 2 -MAX_CROPS = 6 - - -class DeepseekOCR2Processor(ProcessorMixin): - tokenizer_class = ("LlamaTokenizer", "LlamaTokenizerFast") - attributes = ["tokenizer"] - - def __init__( - self, - tokenizer: LlamaTokenizerFast, - patch_size: int = 16, - downsample_ratio: int = 4, - image_mean: tuple[float, float, float] = (0.5, 0.5, 0.5), - image_std: tuple[float, float, float] = (0.5, 0.5, 0.5), - normalize: bool = True, - image_token: str = "", - pad_token: str = "<｜▁pad▁｜>", - add_special_token: bool = False, - sft_format: str = "deepseek", - mask_prompt: bool = True, - ignore_id: int = -100, - **kwargs, - ): - self.image_size = IMAGE_SIZE - self.base_size = BASE_SIZE - self.patch_size = 16 - self.image_mean = image_mean - self.image_std = image_std - self.normalize = normalize - self.downsample_ratio = 4 - - self.image_transform = ImageTransform( - mean=image_mean, std=image_std, normalize=normalize - ) - - self.tokenizer = tokenizer - self.tokenizer.padding_side = "left" # must set this，padding side with make a difference in batch inference # noqa: E501 - - # add the pad_token as special token to use 'tokenizer.pad_token' - # and 'tokenizer.pad_token_id' - if self.tokenizer.pad_token is None: - self.tokenizer.add_special_tokens({"pad_token": pad_token}) - - # add image token - self.image_token_id = self.tokenizer.vocab.get(image_token) - self.image_token = image_token - self.pad_token = pad_token - self.add_special_token = add_special_token - self.sft_format = sft_format - self.mask_prompt = mask_prompt - self.ignore_id = ignore_id - - super().__init__( - tokenizer, - **kwargs, - ) - - @property - def bos_id(self): - return self.tokenizer.bos_token_id - - @property - def eos_id(self): - return self.tokenizer.eos_token_id - - @property - def pad_id(self): - return self.tokenizer.pad_token_id - - def encode(self, text: str, bos: bool = True, eos: bool = False): - t = self.tokenizer.encode(text, add_special_tokens=False) - if bos: - t = [self.bos_id] + t - if eos: - t = t + [self.eos_id] - return t - - def decode(self, t: list[int], **kwargs) -> str: - return self.tokenizer.decode(t, **kwargs) - - def process_one( - self, - prompt: str, - images: list[Image.Image], - crop_mode: bool = CROP_MODE, - ): - """ - - Args: - prompt (str): the formatted prompt; - images (List[ImageType]): the list of images; - crop_mode (bool): if True, then crop the image; - - Returns: - outputs (BaseProcessorOutput): the output of the processor, - - input_ids (torch.LongTensor): [N + image tokens] - - target_ids (torch.LongTensor): [N + image tokens] - - pixel_values (torch.FloatTensor): [n_patches, 3, H, W] - - image_id (int): the id of the image token - - num_image_tokens (List[int]): the number of image tokens - """ - - assert prompt is not None and images is not None, ( - "prompt and images must be used at the same time." - ) - - sft_format = prompt - - ( - input_ids, - pixel_values, - images_crop, - images_seq_mask, - images_spatial_crop, - num_image_tokens, - _, - ) = self.tokenize_with_images( - conversation=sft_format, - images=images, - bos=True, - eos=True, - cropping=crop_mode, - ) - - prepare = BatchFeature( - data=dict( - input_ids=input_ids, - pixel_values=pixel_values, - images_crop=images_crop, - images_seq_mask=images_seq_mask, - images_spatial_crop=images_spatial_crop, - num_image_tokens=num_image_tokens, - ), - tensor_type="pt", - ) - return prepare - - def __call__( - self, - *, - prompt: str, - images: list[Image.Image], - crop_mode: bool = CROP_MODE, - **kwargs, - ): - prepare = self.process_one( - prompt=prompt, - images=images, - crop_mode=crop_mode, - ) - - return prepare - - def tokenize_with_images( - self, - conversation: str, - images: list[Image.Image], - bos: bool = True, - eos: bool = True, - cropping: bool = True, - ): - """Tokenize text with tags.""" - - assert conversation.count(self.image_token) == len(images) - text_splits = conversation.split(self.image_token) - images_list, images_crop_list, images_seq_mask, images_spatial_crop = ( - [], - [], - [], - [], - ) - image_shapes = [] - num_image_tokens = [] - tokenized_str = [] - for text_sep, image in zip(text_splits, images): - tokenized_sep = self.encode(text_sep, bos=False, eos=False) - tokenized_str += tokenized_sep - images_seq_mask += [False] * len(tokenized_sep) - - image_shapes.append(image.size) - - images_crop_raw = [] - if image.size[0] <= 768 and image.size[1] <= 768: - crop_ratio = [1, 1] - elif cropping: - images_crop_raw, crop_ratio = dynamic_preprocess( - image, image_size=IMAGE_SIZE - ) - else: - crop_ratio = [1, 1] - - if self.image_size <= 768 and not cropping: - image = image.resize((self.image_size, self.image_size)) - - global_view = ImageOps.pad( - image, - (self.base_size, self.base_size), - color=tuple(int(x * 255) for x in self.image_transform.mean), - ) - images_list.append(self.image_transform(global_view)) - - num_width_tiles, num_height_tiles = crop_ratio - images_spatial_crop.append([num_width_tiles, num_height_tiles]) - - if num_width_tiles > 1 or num_height_tiles > 1: - for cropped_image in images_crop_raw: - images_crop_list.append(self.image_transform(cropped_image)) - - num_queries = math.ceil( - (self.image_size // self.patch_size) / self.downsample_ratio - ) - num_queries_base = math.ceil( - (self.base_size // self.patch_size) / self.downsample_ratio - ) - - tokenized_image = ( - [self.image_token_id] * num_queries_base - ) * num_queries_base - tokenized_image += [self.image_token_id] - if num_width_tiles > 1 or num_height_tiles > 1: - local_row = [self.image_token_id] * (num_queries * num_width_tiles) - tokenized_image += local_row * (num_queries * num_height_tiles) - tokenized_str += tokenized_image - images_seq_mask += [True] * len(tokenized_image) - num_image_tokens.append(len(tokenized_image)) - - """process the last text split""" - tokenized_sep = self.encode(text_splits[-1], bos=False, eos=False) - tokenized_str += tokenized_sep - images_seq_mask += [False] * len(tokenized_sep) - - """add the bos and eos tokens""" - if bos: - tokenized_str = [self.bos_id] + tokenized_str - images_seq_mask = [False] + images_seq_mask - if eos: - tokenized_str = tokenized_str + [self.eos_id] - images_seq_mask = images_seq_mask + [False] - - assert len(tokenized_str) == len(images_seq_mask), ( - f"tokenize_with_images func: tokenized_str's length {len(tokenized_str)} " - f"is not equal to images_seq_mask's length {len(images_seq_mask)}." - ) - - masked_tokenized_str = [] - for token_index in tokenized_str: - if token_index != self.image_token_id: - masked_tokenized_str.append(token_index) - else: - masked_tokenized_str.append(self.ignore_id) - - assert ( - len(tokenized_str) == len(images_seq_mask) == len(masked_tokenized_str) - ), ( - f"tokenized_str's length {len(tokenized_str)}, " - f"input_ids' length {len(masked_tokenized_str)}, " - f"images_seq_mask's length {len(images_seq_mask)}, are not equal." - ) - - input_ids = torch.LongTensor(tokenized_str) - target_ids = torch.LongTensor(masked_tokenized_str) - images_seq_mask = torch.tensor(images_seq_mask, dtype=torch.bool) - - # set input_ids < 0 | input_ids == self.image_token_id as ignore_id - target_ids[(input_ids < 0) | (input_ids == self.image_token_id)] = ( - self.ignore_id - ) - input_ids[input_ids < 0] = self.pad_id - - # Remove the ending eos token - assert input_ids[-1] == self.eos_id - input_ids = input_ids[:-1] - target_ids = target_ids[:-1] - images_seq_mask = images_seq_mask[:-1] - - if len(images_list) == 0: - pixel_values = torch.zeros((0, 3, self.base_size, self.base_size)) - images_spatial_crop = torch.zeros((0, 2), dtype=torch.long) - images_crop = torch.zeros((0, 3, self.image_size, self.image_size)) - else: - pixel_values = torch.stack(images_list, dim=0) - images_spatial_crop = torch.tensor(images_spatial_crop, dtype=torch.long) - if images_crop_list: - images_crop = torch.stack(images_crop_list, dim=0) - else: - images_crop = torch.zeros((0, 3, self.image_size, self.image_size)) - - input_ids = input_ids.unsqueeze(0) - - return ( - input_ids, - pixel_values, - images_crop, - images_seq_mask, - images_spatial_crop, - num_image_tokens, - image_shapes, - ) - - -AutoProcessor.register("DeepseekOCR2Processor", DeepseekOCR2Processor) From 92e7562a994038c904fea859d90462c7e84a3246 Mon Sep 17 00:00:00 2001 From: Tsukasa OI Date: Fri, 6 Feb 2026 03:47:09 +0900 Subject: [PATCH 108/810] [Bugfix] Suppress non-TTY color output on the process name part of the log (#29714) Signed-off-by: Tsukasa OI --- vllm/utils/system_utils.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/vllm/utils/system_utils.py b/vllm/utils/system_utils.py index 180a8d08b73..840056e8bef 100644 --- a/vllm/utils/system_utils.py +++ b/vllm/utils/system_utils.py @@ -179,7 +179,12 @@ def set_process_title( def _add_prefix(file: TextIO, worker_name: str, pid: int) -> None: """Add colored prefix to file output for log decoration.""" - if envs.NO_COLOR: + is_tty = hasattr(file, "isatty") and file.isatty() + if ( + envs.NO_COLOR + or envs.VLLM_LOGGING_COLOR == "0" + or (envs.VLLM_LOGGING_COLOR != "1" and not is_tty) + ): prefix = f"({worker_name} pid={pid}) " else: prefix = f"{CYAN}({worker_name} pid={pid}){RESET} " From 1887acca9e2ceacea8f7b1770bd0a0fd9b6a3b02 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Thu, 5 Feb 2026 19:16:20 +0000 Subject: [PATCH 109/810] Fix tokenizer test for renamed attr on Transformers v5 (#33902) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- tests/entrypoints/openai/test_serving_tokens.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/tests/entrypoints/openai/test_serving_tokens.py b/tests/entrypoints/openai/test_serving_tokens.py index 215de851073..ee3c6055619 100644 --- a/tests/entrypoints/openai/test_serving_tokens.py +++ b/tests/entrypoints/openai/test_serving_tokens.py @@ -7,6 +7,7 @@ import pytest_asyncio from transformers import AutoTokenizer from vllm.config import ModelConfig +from vllm.config.utils import getattr_iter from vllm.v1.engine.detokenizer import check_stop_strings from ...utils import RemoteOpenAIServer @@ -131,7 +132,14 @@ async def test_same_response_as_chat_completions(client, tokenizer, messages): # Post-EOS generation is undefined and may differ eos_tokens = { tokenizer.eos_token_id, - *tokenizer.additional_special_tokens_ids, + *getattr_iter( + tokenizer, + [ + "extra_special_tokens_ids", # Transformers v5 + "additional_special_tokens_ids", # Transformers v4 + ], + [], + ), } # Find first EOS in generated tokens eos_pos = None From 20f5d185a6f570b74ab403577cd4eabe44d14496 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Thu, 5 Feb 2026 20:16:52 +0100 Subject: [PATCH 110/810] [Misc] Rename `translations` to `speech_to_text` for OAI serving component (#33904) Signed-off-by: NickLucche --- vllm/entrypoints/openai/api_server.py | 8 ++++---- vllm/entrypoints/openai/engine/serving.py | 2 +- .../openai/{translations => speech_to_text}/__init__.py | 0 .../openai/{translations => speech_to_text}/api_router.py | 4 ++-- .../openai/{translations => speech_to_text}/protocol.py | 0 .../openai/{translations => speech_to_text}/serving.py | 4 ++-- .../{translations => speech_to_text}/speech_to_text.py | 4 ++-- 7 files changed, 11 insertions(+), 11 deletions(-) rename vllm/entrypoints/openai/{translations => speech_to_text}/__init__.py (100%) rename vllm/entrypoints/openai/{translations => speech_to_text}/api_router.py (97%) rename vllm/entrypoints/openai/{translations => speech_to_text}/protocol.py (100%) rename vllm/entrypoints/openai/{translations => speech_to_text}/serving.py (97%) rename vllm/entrypoints/openai/{translations => speech_to_text}/speech_to_text.py (99%) diff --git a/vllm/entrypoints/openai/api_server.py b/vllm/entrypoints/openai/api_server.py index a1ee3607a05..d1da420f6bd 100644 --- a/vllm/entrypoints/openai/api_server.py +++ b/vllm/entrypoints/openai/api_server.py @@ -190,11 +190,11 @@ def build_app(args: Namespace, supported_tasks: tuple["SupportedTask", ...]) -> register_generate_api_routers(app) if "transcription" in supported_tasks: - from vllm.entrypoints.openai.translations.api_router import ( - attach_router as register_translations_api_router, + from vllm.entrypoints.openai.speech_to_text.api_router import ( + attach_router as register_speech_to_text_api_router, ) - register_translations_api_router(app) + register_speech_to_text_api_router(app) if "realtime" in supported_tasks: from vllm.entrypoints.openai.realtime.api_router import ( @@ -318,7 +318,7 @@ async def init_app_state( ) if "transcription" in supported_tasks: - from vllm.entrypoints.openai.translations.api_router import ( + from vllm.entrypoints.openai.speech_to_text.api_router import ( init_transcription_state, ) diff --git a/vllm/entrypoints/openai/engine/serving.py b/vllm/entrypoints/openai/engine/serving.py index 801c7dcd52a..f87ac5804f5 100644 --- a/vllm/entrypoints/openai/engine/serving.py +++ b/vllm/entrypoints/openai/engine/serving.py @@ -57,7 +57,7 @@ from vllm.entrypoints.openai.responses.protocol import ( from vllm.entrypoints.openai.responses.utils import ( construct_input_messages, ) -from vllm.entrypoints.openai.translations.protocol import ( +from vllm.entrypoints.openai.speech_to_text.protocol import ( TranscriptionRequest, TranscriptionResponse, TranslationRequest, diff --git a/vllm/entrypoints/openai/translations/__init__.py b/vllm/entrypoints/openai/speech_to_text/__init__.py similarity index 100% rename from vllm/entrypoints/openai/translations/__init__.py rename to vllm/entrypoints/openai/speech_to_text/__init__.py diff --git a/vllm/entrypoints/openai/translations/api_router.py b/vllm/entrypoints/openai/speech_to_text/api_router.py similarity index 97% rename from vllm/entrypoints/openai/translations/api_router.py rename to vllm/entrypoints/openai/speech_to_text/api_router.py index 7dd95161f52..7477b79c08b 100644 --- a/vllm/entrypoints/openai/translations/api_router.py +++ b/vllm/entrypoints/openai/speech_to_text/api_router.py @@ -9,13 +9,13 @@ from fastapi import APIRouter, FastAPI, Form, Request from fastapi.responses import JSONResponse, StreamingResponse from vllm.entrypoints.openai.engine.protocol import ErrorResponse -from vllm.entrypoints.openai.translations.protocol import ( +from vllm.entrypoints.openai.speech_to_text.protocol import ( TranscriptionRequest, TranscriptionResponseVariant, TranslationRequest, TranslationResponseVariant, ) -from vllm.entrypoints.openai.translations.serving import ( +from vllm.entrypoints.openai.speech_to_text.serving import ( OpenAIServingTranscription, OpenAIServingTranslation, ) diff --git a/vllm/entrypoints/openai/translations/protocol.py b/vllm/entrypoints/openai/speech_to_text/protocol.py similarity index 100% rename from vllm/entrypoints/openai/translations/protocol.py rename to vllm/entrypoints/openai/speech_to_text/protocol.py diff --git a/vllm/entrypoints/openai/translations/serving.py b/vllm/entrypoints/openai/speech_to_text/serving.py similarity index 97% rename from vllm/entrypoints/openai/translations/serving.py rename to vllm/entrypoints/openai/speech_to_text/serving.py index 646789bba9b..9d18f5aa34c 100644 --- a/vllm/entrypoints/openai/translations/serving.py +++ b/vllm/entrypoints/openai/speech_to_text/serving.py @@ -11,7 +11,7 @@ from vllm.entrypoints.openai.engine.protocol import ( RequestResponseMetadata, ) from vllm.entrypoints.openai.models.serving import OpenAIServingModels -from vllm.entrypoints.openai.translations.protocol import ( +from vllm.entrypoints.openai.speech_to_text.protocol import ( TranscriptionRequest, TranscriptionResponse, TranscriptionResponseStreamChoice, @@ -23,7 +23,7 @@ from vllm.entrypoints.openai.translations.protocol import ( TranslationResponseVerbose, TranslationStreamResponse, ) -from vllm.entrypoints.openai.translations.speech_to_text import OpenAISpeechToText +from vllm.entrypoints.openai.speech_to_text.speech_to_text import OpenAISpeechToText from vllm.logger import init_logger from vllm.outputs import RequestOutput diff --git a/vllm/entrypoints/openai/translations/speech_to_text.py b/vllm/entrypoints/openai/speech_to_text/speech_to_text.py similarity index 99% rename from vllm/entrypoints/openai/translations/speech_to_text.py rename to vllm/entrypoints/openai/speech_to_text/speech_to_text.py index 58bfb3e970f..19dccbb17d4 100644 --- a/vllm/entrypoints/openai/translations/speech_to_text.py +++ b/vllm/entrypoints/openai/speech_to_text/speech_to_text.py @@ -24,7 +24,7 @@ from vllm.entrypoints.openai.engine.protocol import ( ) from vllm.entrypoints.openai.engine.serving import OpenAIServing, SpeechToTextRequest from vllm.entrypoints.openai.models.serving import OpenAIServingModels -from vllm.entrypoints.openai.translations.protocol import ( +from vllm.entrypoints.openai.speech_to_text.protocol import ( TranscriptionResponse, TranscriptionResponseStreamChoice, TranscriptionResponseVerbose, @@ -402,7 +402,7 @@ class OpenAISpeechToText(OpenAIServing): audio_data: bytes, request: SpeechToTextRequest, raw_request: Request, - response_class: type[T | V], + response_class: type[ResponseType], stream_generator_method: Callable[..., AsyncGenerator[str, None]], ) -> T | V | AsyncGenerator[str, None] | ErrorResponse: """Base method for speech-to-text operations like transcription and From 4145e50d854e3182c22bad99ec011c283b9c493f Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Thu, 5 Feb 2026 14:22:19 -0500 Subject: [PATCH 111/810] [Bugfix] Fix DSV3.2 NVFP4 (#33932) Signed-off-by: Matthew Bonanni --- vllm/model_executor/layers/attention/mla_attention.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 862f8493985..4859af43ae4 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -530,7 +530,7 @@ class MLAAttention(nn.Module, AttentionLayerBase): scale=self._k_scale, ) - if fp8_attention: + if fp8_attention and self.kv_cache_dtype != "fp8_ds_mla": kv_cache = kv_cache.view(current_platform.fp8_dtype()) # Sparse MLA impls only support forward_mqa (decode-style attention) @@ -614,7 +614,7 @@ class MLAAttention(nn.Module, AttentionLayerBase): # Convert from (N, B, L) to (B, N, L) mqa_ql_nope = mqa_ql_nope.transpose(0, 1) - if fp8_attention: + if fp8_attention and self.impl.supports_quant_query_input: assert mqa_ql_nope.shape[0] == mqa_q_pe.shape[0] assert mqa_ql_nope.shape[1] == mqa_q_pe.shape[1] mqa_q = self._decode_concat_quant_fp8_op( @@ -1885,6 +1885,8 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]): self.indexer = indexer self.q_pad_num_heads = q_pad_num_heads + self.supports_quant_query_input = True + # Use flashinfer's optimized concat_mla_k kernel when available. # The kernel is optimized for DeepSeek V3 dimensions: # num_heads=128, nope_dim=128, rope_dim=64 From 116880a5a0af3a226e29f4716c484a4cc8422fc1 Mon Sep 17 00:00:00 2001 From: Cyrus Leung Date: Fri, 6 Feb 2026 04:40:58 +0800 Subject: [PATCH 112/810] [Bugfix] Make MM batching more robust (#33817) Signed-off-by: DarkLight1337 --- .buildkite/test-amd.yaml | 3 +- .buildkite/test-pipeline.yaml | 3 +- .buildkite/test_areas/models_basic.yaml | 3 +- tests/models/test_terratorch.py | 12 +- tests/multimodal/media/test_connector.py | 320 +++++++++++++++++++ tests/multimodal/test_hasher.py | 16 + tests/multimodal/test_inputs.py | 46 +++ tests/multimodal/test_utils.py | 386 +++-------------------- vllm/model_executor/models/step3_vl.py | 36 +-- vllm/model_executor/models/terratorch.py | 46 +-- vllm/multimodal/hasher.py | 18 +- vllm/multimodal/inputs.py | 52 +-- vllm/multimodal/utils.py | 112 ++++++- 13 files changed, 625 insertions(+), 428 deletions(-) create mode 100644 tests/multimodal/media/test_connector.py create mode 100644 tests/multimodal/test_inputs.py diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index ca3bebcb0c3..9c3e84af931 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -863,10 +863,11 @@ steps: torch_nightly: true source_file_dependencies: - vllm/ + - tests/models/test_terratorch.py - tests/models/test_transformers.py - tests/models/test_registry.py commands: - - pytest -v -s models/test_transformers.py models/test_registry.py + - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py - label: Basic Models Test (Other CPU) # 5min mirror_hardwares: [amdexperimental, amdproduction] diff --git a/.buildkite/test-pipeline.yaml b/.buildkite/test-pipeline.yaml index b03e4b6d87b..e3146948b37 100644 --- a/.buildkite/test-pipeline.yaml +++ b/.buildkite/test-pipeline.yaml @@ -804,10 +804,11 @@ steps: torch_nightly: true source_file_dependencies: - vllm/ + - tests/models/test_terratorch.py - tests/models/test_transformers.py - tests/models/test_registry.py commands: - - pytest -v -s models/test_transformers.py models/test_registry.py + - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py - label: Basic Models Test (Other CPU) # 5min timeout_in_minutes: 10 diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index aa6161ffa66..df0a98dc9c2 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -33,10 +33,11 @@ steps: timeout_in_minutes: 45 source_file_dependencies: - vllm/ + - tests/models/test_terratorch.py - tests/models/test_transformers.py - tests/models/test_registry.py commands: - - pytest -v -s models/test_transformers.py models/test_registry.py + - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py - label: Basic Models Test (Other CPU) # 5min depends_on: diff --git a/tests/models/test_terratorch.py b/tests/models/test_terratorch.py index 24b624e2695..5de154fa3ab 100644 --- a/tests/models/test_terratorch.py +++ b/tests/models/test_terratorch.py @@ -5,8 +5,10 @@ import pytest import torch from tests.conftest import VllmRunner +from tests.utils import create_new_process_for_each_test +@create_new_process_for_each_test() # Memory is not cleaned up properly otherwise @pytest.mark.parametrize( "model", [ @@ -22,10 +24,14 @@ def test_inference( location_coords = torch.full((1, 2), 1.0, dtype=torch.float16) prompt = dict( prompt_token_ids=[1], - multi_modal_data=dict( - pixel_values=pixel_values, location_coords=location_coords - ), + multi_modal_data={ + "image": { + "pixel_values": pixel_values, + "location_coords": location_coords, + } + }, ) + with vllm_runner( model, runner="pooling", diff --git a/tests/multimodal/media/test_connector.py b/tests/multimodal/media/test_connector.py new file mode 100644 index 00000000000..6ef71fcc06e --- /dev/null +++ b/tests/multimodal/media/test_connector.py @@ -0,0 +1,320 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import asyncio +import base64 +import mimetypes +import os +from tempfile import NamedTemporaryFile, TemporaryDirectory + +import numpy as np +import pytest +import torch +from PIL import Image, ImageChops + +from vllm.multimodal.image import convert_image_mode +from vllm.multimodal.inputs import PlaceholderRange +from vllm.multimodal.media import MediaConnector + +# Test different image extensions (JPG/PNG) and formats (gray/RGB/RGBA) +TEST_IMAGE_ASSETS = [ + "2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg" + "Grayscale_8bits_palette_sample_image.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/Grayscale_8bits_palette_sample_image.png", + "1280px-Venn_diagram_rgb.svg.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/1280px-Venn_diagram_rgb.svg.png", + "RGBA_comp.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/RGBA_comp.png", +] + +TEST_VIDEO_URLS = [ + "https://www.bogotobogo.com/python/OpenCV_Python/images/mean_shift_tracking/slow_traffic_small.mp4", + "https://github.com/opencv/opencv/raw/refs/tags/4.12.0/samples/data/vtest.avi", +] + + +@pytest.fixture(scope="module") +def url_images(local_asset_server) -> dict[str, Image.Image]: + return { + image_url: local_asset_server.get_image_asset(image_url) + for image_url in TEST_IMAGE_ASSETS + } + + +def get_supported_suffixes() -> tuple[str, ...]: + # We should at least test the file types mentioned in GPT-4 with Vision + OPENAI_SUPPORTED_SUFFIXES = (".png", ".jpeg", ".jpg", ".webp", ".gif") + + # Additional file types that are supported by us + EXTRA_SUPPORTED_SUFFIXES = (".bmp", ".tiff") + + return OPENAI_SUPPORTED_SUFFIXES + EXTRA_SUPPORTED_SUFFIXES + + +def _image_equals(a: Image.Image, b: Image.Image) -> bool: + return (np.asarray(a) == np.asarray(convert_image_mode(b, a.mode))).all() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("image_url", TEST_IMAGE_ASSETS, indirect=True) +async def test_fetch_image_http(image_url: str): + connector = MediaConnector() + + image_sync = connector.fetch_image(image_url) + image_async = await connector.fetch_image_async(image_url) + assert _image_equals(image_sync, image_async) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("raw_image_url", TEST_IMAGE_ASSETS) +@pytest.mark.parametrize("suffix", get_supported_suffixes()) +async def test_fetch_image_base64( + url_images: dict[str, Image.Image], raw_image_url: str, suffix: str +): + connector = MediaConnector( + # Domain restriction should not apply to data URLs. + allowed_media_domains=[ + "www.bogotobogo.com", + "github.com", + ] + ) + url_image = url_images[raw_image_url] + + try: + mime_type = Image.MIME[Image.registered_extensions()[suffix]] + except KeyError: + try: + mime_type = mimetypes.types_map[suffix] + except KeyError: + pytest.skip("No MIME type") + + with NamedTemporaryFile(suffix=suffix) as f: + try: + url_image.save(f.name) + except Exception as e: + if e.args[0] == "cannot write mode RGBA as JPEG": + pytest.skip("Conversion not supported") + + raise + + base64_image = base64.b64encode(f.read()).decode("utf-8") + data_url = f"data:{mime_type};base64,{base64_image}" + + data_image_sync = connector.fetch_image(data_url) + if _image_equals(url_image, Image.open(f)): + assert _image_equals(url_image, data_image_sync) + else: + pass # Lossy format; only check that image can be opened + + data_image_async = await connector.fetch_image_async(data_url) + assert _image_equals(data_image_sync, data_image_async) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("image_url", TEST_IMAGE_ASSETS, indirect=True) +async def test_fetch_image_local_files(image_url: str): + connector = MediaConnector() + + with TemporaryDirectory() as temp_dir: + local_connector = MediaConnector(allowed_local_media_path=temp_dir) + + origin_image = connector.fetch_image(image_url) + origin_image.save( + os.path.join(temp_dir, os.path.basename(image_url)), + quality=100, + icc_profile=origin_image.info.get("icc_profile"), + ) + + image_async = await local_connector.fetch_image_async( + f"file://{temp_dir}/{os.path.basename(image_url)}" + ) + image_sync = local_connector.fetch_image( + f"file://{temp_dir}/{os.path.basename(image_url)}" + ) + # Check that the images are equal + assert not ImageChops.difference(image_sync, image_async).getbbox() + + with pytest.raises(ValueError, match="must be a subpath"): + await local_connector.fetch_image_async( + f"file://{temp_dir}/../{os.path.basename(image_url)}" + ) + with pytest.raises(RuntimeError, match="Cannot load local files"): + await connector.fetch_image_async( + f"file://{temp_dir}/../{os.path.basename(image_url)}" + ) + + with pytest.raises(ValueError, match="must be a subpath"): + local_connector.fetch_image( + f"file://{temp_dir}/../{os.path.basename(image_url)}" + ) + with pytest.raises(RuntimeError, match="Cannot load local files"): + connector.fetch_image(f"file://{temp_dir}/../{os.path.basename(image_url)}") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("image_url", [TEST_IMAGE_ASSETS[0]], indirect=True) +async def test_fetch_image_local_files_with_space_in_name(image_url: str): + connector = MediaConnector() + + with TemporaryDirectory() as temp_dir: + local_connector = MediaConnector(allowed_local_media_path=temp_dir) + + origin_image = connector.fetch_image(image_url) + filename = "file name with space.jpg" + origin_image.save( + os.path.join(temp_dir, filename), + quality=100, + icc_profile=origin_image.info.get("icc_profile"), + ) + + try: + image_async = await local_connector.fetch_image_async( + f"file://{temp_dir}/{filename}" + ) + image_sync = local_connector.fetch_image(f"file://{temp_dir}/{filename}") + except FileNotFoundError as e: + pytest.fail("Failed to fetch image with space in name: {}".format(e)) + # Check that the images are equal + assert not ImageChops.difference(image_sync, image_async).getbbox() + + +@pytest.mark.asyncio +async def test_fetch_image_error_conversion(): + connector = MediaConnector() + broken_img = "data:image/png;base64,aGVsbG9fdmxsbV9jb21tdW5pdHkK" + + # PIL.UnidentifiedImageError should be converted to ValueError + with pytest.raises(ValueError): + await connector.fetch_image_async(broken_img) + + with pytest.raises(ValueError): + connector.fetch_image(broken_img) + + +@pytest.mark.flaky(reruns=3, reruns_delay=5) +@pytest.mark.asyncio +@pytest.mark.parametrize("video_url", TEST_VIDEO_URLS) +@pytest.mark.parametrize("num_frames", [-1, 32, 1800]) +async def test_fetch_video_http(video_url: str, num_frames: int): + connector = MediaConnector( + media_io_kwargs={ + "video": { + "num_frames": num_frames, + } + } + ) + + try: + video_sync, metadata_sync = connector.fetch_video(video_url) + video_async, metadata_async = await connector.fetch_video_async(video_url) + except (TimeoutError, asyncio.TimeoutError) as e: + pytest.skip(f"Timeout fetching video (CI network flakiness): {e}") + + assert np.array_equal(video_sync, video_async) + assert metadata_sync == metadata_async + + +@pytest.mark.asyncio +@pytest.mark.parametrize("video_url", TEST_VIDEO_URLS) +@pytest.mark.parametrize("max_duration", [1, 60, 1800]) +@pytest.mark.parametrize("requested_fps", [2, 24]) +async def test_fetch_video_http_with_dynamic_loader( + video_url: str, + max_duration: int, + requested_fps: int, + monkeypatch: pytest.MonkeyPatch, +): + with monkeypatch.context() as m: + m.setenv("VLLM_VIDEO_LOADER_BACKEND", "opencv_dynamic") + connector = MediaConnector( + media_io_kwargs={ + "video": { + "max_duration": max_duration, + "requested_fps": requested_fps, + } + } + ) + + video_sync, metadata_sync = connector.fetch_video(video_url) + video_async, metadata_async = await connector.fetch_video_async(video_url) + + assert np.array_equal(video_sync, video_async) + assert metadata_sync == metadata_async + assert metadata_sync["video_backend"] == "opencv_dynamic" + + +@pytest.mark.parametrize( + "is_embed,start_idx,end_idx,expected", + [ + (None, 2, 4, (2, 4)), + ( + torch.tensor([False, True, False, True, True]), + 3, + 5, + (1, 3), + ), + ( + torch.tensor([False, True, False, True, True]), + 0, + 2, + (0, 1), + ), + ( + torch.tensor([True, False, True, False]), + 2, + 2, + (1, 1), + ), + ], +) +def test_placeholder_range_get_embeds_indices_in_range( + is_embed, start_idx, end_idx, expected +): + length = len(is_embed) if is_embed is not None else 5 + pr = PlaceholderRange(offset=0, length=length, is_embed=is_embed) + assert pr.get_embeds_indices_in_range(start_idx, end_idx) == expected + + +@pytest.mark.parametrize( + "offset,is_embed,expected", + [ + (0, None, [(0, 4)]), + ( + 2, + torch.tensor([False, True, False, True, True]), + [(3, 3), (5, 6)], + ), + (0, torch.tensor([True, True, True, True]), [(0, 3)]), + (0, torch.tensor([False, False, False, False]), []), + ], +) +def test_placeholder_range_extract_embeds_range(offset, is_embed, expected): + length = len(is_embed) if is_embed is not None else 5 + pr = PlaceholderRange(offset=offset, length=length, is_embed=is_embed) + assert pr.extract_embeds_range() == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize("video_url", TEST_VIDEO_URLS) +@pytest.mark.parametrize("num_frames", [-1, 32, 1800]) +async def test_allowed_media_domains(video_url: str, num_frames: int): + connector = MediaConnector( + media_io_kwargs={ + "video": { + "num_frames": num_frames, + } + }, + allowed_media_domains=[ + "www.bogotobogo.com", + "github.com", + ], + ) + + video_sync, metadata_sync = connector.fetch_video(video_url) + video_async, metadata_async = await connector.fetch_video_async(video_url) + assert np.array_equal(video_sync, video_async) + assert metadata_sync == metadata_async + + disallowed_url = "https://upload.wikimedia.org/wikipedia/commons/4/47/PNG_transparency_demonstration_1.png" + with pytest.raises(ValueError): + _, _ = connector.fetch_video(disallowed_url) + + with pytest.raises(ValueError): + _, _ = await connector.fetch_video_async(disallowed_url) diff --git a/tests/multimodal/test_hasher.py b/tests/multimodal/test_hasher.py index 29064f27378..fdedcaea27c 100644 --- a/tests/multimodal/test_hasher.py +++ b/tests/multimodal/test_hasher.py @@ -16,6 +16,22 @@ ASSETS_DIR = Path(__file__).parent / "assets" assert ASSETS_DIR.exists() +def test_hash_single_item_different_shape(): + x1 = torch.zeros(()) + x2 = torch.zeros((1,)) + + hasher = MultiModalHasher + assert hasher.hash_kwargs(x=x1) != hasher.hash_kwargs(x=x2) + + +def test_hash_key_order_invariant(): + x = torch.zeros((5, 10)) + y = torch.ones((5, 10)) + + hasher = MultiModalHasher + assert hasher.hash_kwargs(x=x, y=y) == hasher.hash_kwargs(y=y, x=x) + + # NOTE: Images that are the same visually are allowed to have the same hash @pytest.mark.parametrize("mode_pair", [("1", "L"), ("RGBA", "CMYK")]) def test_hash_collision_image_mode(mode_pair): diff --git a/tests/multimodal/test_inputs.py b/tests/multimodal/test_inputs.py new file mode 100644 index 00000000000..7378c149304 --- /dev/null +++ b/tests/multimodal/test_inputs.py @@ -0,0 +1,46 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest +import torch + +from vllm.multimodal.inputs import PlaceholderRange + + +@pytest.mark.parametrize( + "is_embed,expected", + [ + (None, 5), + (torch.tensor([True, True, True, True, True]), 5), + (torch.tensor([False, False, False, False, False]), 0), + (torch.tensor([True, False, True, False, True]), 3), + (torch.tensor([True]), 1), + ], +) +def test_placeholder_range_get_num_embeds(is_embed, expected): + length = len(is_embed) if is_embed is not None else 5 + pr = PlaceholderRange(offset=0, length=length, is_embed=is_embed) + assert pr.get_num_embeds == expected + + +@pytest.mark.parametrize( + "is_embed,expected", + [ + (None, None), + ( + torch.tensor([False, True, False, True, True]), + torch.tensor([0, 1, 1, 2, 3]), + ), + (torch.tensor([True, True, True]), torch.tensor([1, 2, 3])), + ], +) +def test_placeholder_range_embeds_cumsum(is_embed, expected): + length = len(is_embed) if is_embed is not None else 5 + pr = PlaceholderRange(offset=0, length=length, is_embed=is_embed) + + if expected is None: + assert pr.embeds_cumsum is None + return + + assert torch.equal(pr.embeds_cumsum, expected) + # cached_property should return the same object on repeated access + assert pr.embeds_cumsum is pr.embeds_cumsum diff --git a/tests/multimodal/test_utils.py b/tests/multimodal/test_utils.py index 0a10bc1bb06..4e765ab1b8b 100644 --- a/tests/multimodal/test_utils.py +++ b/tests/multimodal/test_utils.py @@ -1,244 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import asyncio -import base64 -import mimetypes -import os -from tempfile import NamedTemporaryFile, TemporaryDirectory - -import numpy as np import pytest import torch -from PIL import Image, ImageChops -from vllm.multimodal.image import convert_image_mode -from vllm.multimodal.inputs import PlaceholderRange -from vllm.multimodal.media import MediaConnector -from vllm.multimodal.utils import argsort_mm_positions - -# Test different image extensions (JPG/PNG) and formats (gray/RGB/RGBA) -TEST_IMAGE_ASSETS = [ - "2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg" - "Grayscale_8bits_palette_sample_image.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/Grayscale_8bits_palette_sample_image.png", - "1280px-Venn_diagram_rgb.svg.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/1280px-Venn_diagram_rgb.svg.png", - "RGBA_comp.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/RGBA_comp.png", -] - -TEST_VIDEO_URLS = [ - "https://www.bogotobogo.com/python/OpenCV_Python/images/mean_shift_tracking/slow_traffic_small.mp4", - "https://github.com/opencv/opencv/raw/refs/tags/4.12.0/samples/data/vtest.avi", -] - - -@pytest.fixture(scope="module") -def url_images(local_asset_server) -> dict[str, Image.Image]: - return { - image_url: local_asset_server.get_image_asset(image_url) - for image_url in TEST_IMAGE_ASSETS - } - - -def get_supported_suffixes() -> tuple[str, ...]: - # We should at least test the file types mentioned in GPT-4 with Vision - OPENAI_SUPPORTED_SUFFIXES = (".png", ".jpeg", ".jpg", ".webp", ".gif") - - # Additional file types that are supported by us - EXTRA_SUPPORTED_SUFFIXES = (".bmp", ".tiff") - - return OPENAI_SUPPORTED_SUFFIXES + EXTRA_SUPPORTED_SUFFIXES - - -def _image_equals(a: Image.Image, b: Image.Image) -> bool: - return (np.asarray(a) == np.asarray(convert_image_mode(b, a.mode))).all() - - -@pytest.mark.asyncio -@pytest.mark.parametrize("image_url", TEST_IMAGE_ASSETS, indirect=True) -async def test_fetch_image_http(image_url: str): - connector = MediaConnector() - - image_sync = connector.fetch_image(image_url) - image_async = await connector.fetch_image_async(image_url) - assert _image_equals(image_sync, image_async) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("raw_image_url", TEST_IMAGE_ASSETS) -@pytest.mark.parametrize("suffix", get_supported_suffixes()) -async def test_fetch_image_base64( - url_images: dict[str, Image.Image], raw_image_url: str, suffix: str -): - connector = MediaConnector( - # Domain restriction should not apply to data URLs. - allowed_media_domains=[ - "www.bogotobogo.com", - "github.com", - ] - ) - url_image = url_images[raw_image_url] - - try: - mime_type = Image.MIME[Image.registered_extensions()[suffix]] - except KeyError: - try: - mime_type = mimetypes.types_map[suffix] - except KeyError: - pytest.skip("No MIME type") - - with NamedTemporaryFile(suffix=suffix) as f: - try: - url_image.save(f.name) - except Exception as e: - if e.args[0] == "cannot write mode RGBA as JPEG": - pytest.skip("Conversion not supported") - - raise - - base64_image = base64.b64encode(f.read()).decode("utf-8") - data_url = f"data:{mime_type};base64,{base64_image}" - - data_image_sync = connector.fetch_image(data_url) - if _image_equals(url_image, Image.open(f)): - assert _image_equals(url_image, data_image_sync) - else: - pass # Lossy format; only check that image can be opened - - data_image_async = await connector.fetch_image_async(data_url) - assert _image_equals(data_image_sync, data_image_async) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("image_url", TEST_IMAGE_ASSETS, indirect=True) -async def test_fetch_image_local_files(image_url: str): - connector = MediaConnector() - - with TemporaryDirectory() as temp_dir: - local_connector = MediaConnector(allowed_local_media_path=temp_dir) - - origin_image = connector.fetch_image(image_url) - origin_image.save( - os.path.join(temp_dir, os.path.basename(image_url)), - quality=100, - icc_profile=origin_image.info.get("icc_profile"), - ) - - image_async = await local_connector.fetch_image_async( - f"file://{temp_dir}/{os.path.basename(image_url)}" - ) - image_sync = local_connector.fetch_image( - f"file://{temp_dir}/{os.path.basename(image_url)}" - ) - # Check that the images are equal - assert not ImageChops.difference(image_sync, image_async).getbbox() - - with pytest.raises(ValueError, match="must be a subpath"): - await local_connector.fetch_image_async( - f"file://{temp_dir}/../{os.path.basename(image_url)}" - ) - with pytest.raises(RuntimeError, match="Cannot load local files"): - await connector.fetch_image_async( - f"file://{temp_dir}/../{os.path.basename(image_url)}" - ) - - with pytest.raises(ValueError, match="must be a subpath"): - local_connector.fetch_image( - f"file://{temp_dir}/../{os.path.basename(image_url)}" - ) - with pytest.raises(RuntimeError, match="Cannot load local files"): - connector.fetch_image(f"file://{temp_dir}/../{os.path.basename(image_url)}") - - -@pytest.mark.asyncio -@pytest.mark.parametrize("image_url", [TEST_IMAGE_ASSETS[0]], indirect=True) -async def test_fetch_image_local_files_with_space_in_name(image_url: str): - connector = MediaConnector() - - with TemporaryDirectory() as temp_dir: - local_connector = MediaConnector(allowed_local_media_path=temp_dir) - - origin_image = connector.fetch_image(image_url) - filename = "file name with space.jpg" - origin_image.save( - os.path.join(temp_dir, filename), - quality=100, - icc_profile=origin_image.info.get("icc_profile"), - ) - - try: - image_async = await local_connector.fetch_image_async( - f"file://{temp_dir}/{filename}" - ) - image_sync = local_connector.fetch_image(f"file://{temp_dir}/{filename}") - except FileNotFoundError as e: - pytest.fail("Failed to fetch image with space in name: {}".format(e)) - # Check that the images are equal - assert not ImageChops.difference(image_sync, image_async).getbbox() - - -@pytest.mark.asyncio -async def test_fetch_image_error_conversion(): - connector = MediaConnector() - broken_img = "data:image/png;base64,aGVsbG9fdmxsbV9jb21tdW5pdHkK" - - # PIL.UnidentifiedImageError should be converted to ValueError - with pytest.raises(ValueError): - await connector.fetch_image_async(broken_img) - - with pytest.raises(ValueError): - connector.fetch_image(broken_img) - - -@pytest.mark.flaky(reruns=3, reruns_delay=5) -@pytest.mark.asyncio -@pytest.mark.parametrize("video_url", TEST_VIDEO_URLS) -@pytest.mark.parametrize("num_frames", [-1, 32, 1800]) -async def test_fetch_video_http(video_url: str, num_frames: int): - connector = MediaConnector( - media_io_kwargs={ - "video": { - "num_frames": num_frames, - } - } - ) - - try: - video_sync, metadata_sync = connector.fetch_video(video_url) - video_async, metadata_async = await connector.fetch_video_async(video_url) - except (TimeoutError, asyncio.TimeoutError) as e: - pytest.skip(f"Timeout fetching video (CI network flakiness): {e}") - - assert np.array_equal(video_sync, video_async) - assert metadata_sync == metadata_async - - -@pytest.mark.asyncio -@pytest.mark.parametrize("video_url", TEST_VIDEO_URLS) -@pytest.mark.parametrize("max_duration", [1, 60, 1800]) -@pytest.mark.parametrize("requested_fps", [2, 24]) -async def test_fetch_video_http_with_dynamic_loader( - video_url: str, - max_duration: int, - requested_fps: int, - monkeypatch: pytest.MonkeyPatch, -): - with monkeypatch.context() as m: - m.setenv("VLLM_VIDEO_LOADER_BACKEND", "opencv_dynamic") - connector = MediaConnector( - media_io_kwargs={ - "video": { - "max_duration": max_duration, - "requested_fps": requested_fps, - } - } - ) - - video_sync, metadata_sync = connector.fetch_video(video_url) - video_async, metadata_async = await connector.fetch_video_async(video_url) - - assert np.array_equal(video_sync, video_async) - assert metadata_sync == metadata_async - assert metadata_sync["video_backend"] == "opencv_dynamic" +from vllm.multimodal.inputs import ( + MultiModalBatchedField, + MultiModalFieldElem, + MultiModalKwargsItem, + MultiModalSharedField, + PlaceholderRange, +) +from vllm.multimodal.utils import argsort_mm_positions, group_and_batch_mm_items @pytest.mark.parametrize( @@ -412,121 +184,35 @@ def test_argsort_mm_positions(case): assert modality_idxs == expected_modality_idxs -@pytest.mark.parametrize( - "is_embed,expected", - [ - (None, 5), - (torch.tensor([True, True, True, True, True]), 5), - (torch.tensor([False, False, False, False, False]), 0), - (torch.tensor([True, False, True, False, True]), 3), - (torch.tensor([True]), 1), - ], -) -def test_placeholder_range_get_num_embeds(is_embed, expected): - length = len(is_embed) if is_embed is not None else 5 - pr = PlaceholderRange(offset=0, length=length, is_embed=is_embed) - assert pr.get_num_embeds == expected - - -@pytest.mark.parametrize( - "is_embed,expected", - [ - (None, None), - ( - torch.tensor([False, True, False, True, True]), - torch.tensor([0, 1, 1, 2, 3]), - ), - (torch.tensor([True, True, True]), torch.tensor([1, 2, 3])), - ], -) -def test_placeholder_range_embeds_cumsum(is_embed, expected): - length = len(is_embed) if is_embed is not None else 5 - pr = PlaceholderRange(offset=0, length=length, is_embed=is_embed) - - if expected is None: - assert pr.embeds_cumsum is None - return - - assert torch.equal(pr.embeds_cumsum, expected) - # cached_property should return the same object on repeated access - assert pr.embeds_cumsum is pr.embeds_cumsum - - -@pytest.mark.parametrize( - "is_embed,start_idx,end_idx,expected", - [ - (None, 2, 4, (2, 4)), - ( - torch.tensor([False, True, False, True, True]), - 3, - 5, - (1, 3), - ), - ( - torch.tensor([False, True, False, True, True]), - 0, - 2, - (0, 1), - ), - ( - torch.tensor([True, False, True, False]), - 2, - 2, - (1, 1), - ), - ], -) -def test_placeholder_range_get_embeds_indices_in_range( - is_embed, start_idx, end_idx, expected -): - length = len(is_embed) if is_embed is not None else 5 - pr = PlaceholderRange(offset=0, length=length, is_embed=is_embed) - assert pr.get_embeds_indices_in_range(start_idx, end_idx) == expected - - -@pytest.mark.parametrize( - "offset,is_embed,expected", - [ - (0, None, [(0, 4)]), - ( - 2, - torch.tensor([False, True, False, True, True]), - [(3, 3), (5, 6)], - ), - (0, torch.tensor([True, True, True, True]), [(0, 3)]), - (0, torch.tensor([False, False, False, False]), []), - ], -) -def test_placeholder_range_extract_embeds_range(offset, is_embed, expected): - length = len(is_embed) if is_embed is not None else 5 - pr = PlaceholderRange(offset=offset, length=length, is_embed=is_embed) - assert pr.extract_embeds_range() == expected - - -@pytest.mark.asyncio -@pytest.mark.parametrize("video_url", TEST_VIDEO_URLS) -@pytest.mark.parametrize("num_frames", [-1, 32, 1800]) -async def test_allowed_media_domains(video_url: str, num_frames: int): - connector = MediaConnector( - media_io_kwargs={ - "video": { - "num_frames": num_frames, - } - }, - allowed_media_domains=[ - "www.bogotobogo.com", - "github.com", - ], +def test_group_and_batch_mm_items_split_by_fieldset(): + elem = MultiModalFieldElem( + data=torch.empty(1, dtype=torch.uint8), + field=MultiModalBatchedField(), ) + item1 = MultiModalKwargsItem({"x": elem, "y": elem}) + item2 = MultiModalKwargsItem({"y": elem, "x": elem}) + item3 = MultiModalKwargsItem({"x": elem, "y": elem, "z": elem}) + item4 = MultiModalKwargsItem({"x": elem}) + item5 = MultiModalKwargsItem({"x": elem, "y": elem}) - video_sync, metadata_sync = connector.fetch_video(video_url) - video_async, metadata_async = await connector.fetch_video_async(video_url) - assert np.array_equal(video_sync, video_async) - assert metadata_sync == metadata_async + res = group_and_batch_mm_items([item1, item2, item3, item4, item5]) + assert [num_items for num_items, _ in res] == [2, 1, 1, 1] - disallowed_url = "https://upload.wikimedia.org/wikipedia/commons/4/47/PNG_transparency_demonstration_1.png" - with pytest.raises(ValueError): - _, _ = connector.fetch_video(disallowed_url) - with pytest.raises(ValueError): - _, _ = await connector.fetch_video_async(disallowed_url) +def test_group_and_batch_mm_items_split_by_shared_data(): + elem1 = MultiModalFieldElem( + data=torch.zeros(1, dtype=torch.uint8), + field=MultiModalSharedField(batch_size=1), + ) + elem2 = MultiModalFieldElem( + data=torch.zeros(2, dtype=torch.uint8), + field=MultiModalSharedField(batch_size=1), + ) + item1 = MultiModalKwargsItem({"x": elem1}) + item2 = MultiModalKwargsItem({"x": elem1}) + item3 = MultiModalKwargsItem({"x": elem2}) + item4 = MultiModalKwargsItem({"x": elem1}) + item5 = MultiModalKwargsItem({"x": elem2}) + + res = group_and_batch_mm_items([item1, item2, item3, item4, item5]) + assert [num_items for num_items, _ in res] == [2, 1, 1, 1] diff --git a/vllm/model_executor/models/step3_vl.py b/vllm/model_executor/models/step3_vl.py index f3993348b30..11081b04081 100644 --- a/vllm/model_executor/models/step3_vl.py +++ b/vllm/model_executor/models/step3_vl.py @@ -71,9 +71,7 @@ class Step3VLImagePixelInputs(TensorSchema): type: Literal["pixel_values"] pixel_values: Annotated[torch.Tensor, TensorShape("bn", 3, "h", "w")] - patch_pixel_values: Annotated[ - torch.Tensor | None, TensorShape("bnp", 3, "hp", "wp") - ] + patch_pixel_values: Annotated[torch.Tensor, TensorShape("bnp", 3, "hp", "wp")] num_patches: Annotated[torch.Tensor, TensorShape("bn")] @@ -91,7 +89,7 @@ class Step3VLImageEmbeddingInputs(TensorSchema): Step3VLImageInputs: TypeAlias = Step3VLImagePixelInputs | Step3VLImageEmbeddingInputs -ImageWithPatches = tuple[Image.Image, list[Image.Image], list[int] | None] +ImageWithPatches = tuple[Image.Image, list[Image.Image], list[bool] | None] MAX_IMAGE_SIZE: int = 3024 @@ -432,7 +430,7 @@ class Step3VLProcessor: if len(parts) - 1 != len(repls): raise ValueError( - "The number of placeholders does not match the number of replacements." # noqa: E501 + "The number of placeholders does not match the number of replacements." ) result = [parts[0]] @@ -468,7 +466,7 @@ class Step3VLProcessor: image_repl_str_lst = [] image_repl_ids_lst = [] num_patches = [] - for raw_img, img_patches, patch_newline_mask in splitted_images_data: # noqa: E501 + for raw_img, img_patches, patch_newline_mask in splitted_images_data: pixel_values_lst.extend(self._convert_images_to_pixel_values([raw_img])) if len(img_patches) > 0: @@ -486,16 +484,20 @@ class Step3VLProcessor: if patch_newline_mask is not None: patch_newline_mask_lst.extend(patch_newline_mask) + pixel_values = torch.cat(pixel_values_lst) + patch_size = self.patch_size image_inputs = { - "pixel_values": torch.cat(pixel_values_lst), + "pixel_values": pixel_values, "num_patches": num_patches, - } - if patch_pixel_values_lst: - image_inputs["patch_pixel_values"] = torch.cat(patch_pixel_values_lst) - if patch_newline_mask_lst: - image_inputs["patch_newline_mask"] = torch.tensor( + "patch_pixel_values": ( + torch.cat(patch_pixel_values_lst) + if patch_pixel_values_lst + else pixel_values.new_empty((0, 3, patch_size, patch_size)) + ), + "patch_newline_mask": torch.tensor( patch_newline_mask_lst, dtype=torch.bool - ) + ), + } text = [ self.replace_placeholder(t, self.image_token, image_repl_str_lst) @@ -998,13 +1000,11 @@ class Step3VLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP) if pixel_values is None and image_embeds is None: return None - if pixel_values is not None: + if pixel_values is not None and patch_pixel_values is not None: return Step3VLImagePixelInputs( type="pixel_values", pixel_values=pixel_values.to(self.dtype), - patch_pixel_values=patch_pixel_values.to(self.dtype) - if patch_pixel_values is not None - else None, + patch_pixel_values=patch_pixel_values.to(self.dtype), num_patches=num_patches, ) @@ -1039,7 +1039,7 @@ class Step3VLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP) image_features = self._get_vision_model_output(image_input["pixel_values"]) patch_image_features = ( self._get_vision_model_output(image_input["patch_pixel_values"]) - if image_input["patch_pixel_values"] is not None + if len(image_input["patch_pixel_values"]) > 0 else None ) num_patches = image_input["num_patches"] diff --git a/vllm/model_executor/models/terratorch.py b/vllm/model_executor/models/terratorch.py index a4fc3a10b26..b817383ab1e 100644 --- a/vllm/model_executor/models/terratorch.py +++ b/vllm/model_executor/models/terratorch.py @@ -62,7 +62,6 @@ from vllm.multimodal.processing import ( PromptUpdate, ) from vllm.sequence import IntermediateTensors -from vllm.utils import length_from_prompt_token_ids_or_embeds from .interfaces import IsAttentionFree, MultiModalEmbeddings, SupportsMultiModal from .interfaces_base import attn_type @@ -74,7 +73,11 @@ def _terratorch_field_names(input_definition: InputDefinition): return set(input_definition.data.keys()) -def _terratorch_field_factory(input_definition: InputDefinition): +def _terratorch_field_factory( + input_definition: InputDefinition, + *, + is_shared: bool = True, # True for unprocessed data, False for processed data +): def _terratorch_field_config( hf_inputs: Mapping[str, torch.Tensor], ) -> Mapping[str, MultiModalFieldConfig]: @@ -82,7 +85,11 @@ def _terratorch_field_factory(input_definition: InputDefinition): for name, input in input_definition.data.items(): modality = "image" if input.type == InputTypeEnum.tensor: - fields[name] = MultiModalFieldConfig.shared(modality, batch_size=1) + fields[name] = ( + MultiModalFieldConfig.shared(modality, batch_size=1) + if is_shared + else MultiModalFieldConfig.batched(modality) + ) return fields @@ -166,8 +173,14 @@ class TerratorchMultiModalProcessor(BaseMultiModalProcessor[TerratorchProcessing self, hf_inputs: BatchFeature, hf_processor_mm_kwargs: Mapping[str, object], + *, + is_shared: bool = True, ) -> Mapping[str, MultiModalFieldConfig]: - return _terratorch_field_factory(self.info.input_definition)(hf_inputs) + factory = _terratorch_field_factory( + self.info.input_definition, + is_shared=is_shared, + ) + return factory(hf_inputs) def _get_prompt_updates( self, @@ -193,12 +206,19 @@ class TerratorchMultiModalProcessor(BaseMultiModalProcessor[TerratorchProcessing ) _, passthrough_data = self._get_hf_mm_data(mm_items) - mm_processed_data = BatchFeature(dict(passthrough_data), tensor_type="pt") + mm_processed_data = BatchFeature( + {k: torch.tensor(v).unsqueeze(0) for k, v in passthrough_data.items()}, + tensor_type="pt", + ) mm_placeholders = {"image": [PlaceholderRange(offset=0, length=0)]} mm_kwargs = MultiModalKwargsItems.from_hf_inputs( mm_processed_data, - self._get_mm_fields_config(mm_processed_data, hf_processor_mm_kwargs), + self._get_mm_fields_config( + mm_processed_data, + hf_processor_mm_kwargs, + is_shared=False, + ), ) return MultiModalInputs( @@ -235,9 +255,6 @@ class Terratorch(nn.Module, IsAttentionFree, SupportsMultiModal): self.inference_runner = InferenceRunner(config) self.model = self.inference_runner.model - pooler_config = vllm_config.model_config.pooler_config - assert pooler_config is not None - self.pooler = IdentityPooler() def embed_input_ids( @@ -262,15 +279,8 @@ class Terratorch(nn.Module, IsAttentionFree, SupportsMultiModal): inputs_embeds: torch.Tensor | None = None, **kwargs: object, ): - input_len = length_from_prompt_token_ids_or_embeds(input_ids, inputs_embeds) - - batched_kwargs = {k: v.unsqueeze(0) for k, v in kwargs.items()} - model_output = self.inference_runner.forward(**batched_kwargs).output - - # The leading dimension of hidden states needs to equal input length - return model_output.expand( - input_len, *(-1 for _ in range(model_output.ndim - 1)) - ) + model_output = self.inference_runner.forward(**kwargs) + return model_output.output def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: params_list = [] diff --git a/vllm/multimodal/hasher.py b/vllm/multimodal/hasher.py index 009d0bc4495..6caf9c11427 100644 --- a/vllm/multimodal/hasher.py +++ b/vllm/multimodal/hasher.py @@ -102,12 +102,19 @@ class MultiModalHasher: "data": tensor_obj.numpy(), }, ) + return cls.iter_item_to_bytes("tensor", tensor_obj.numpy()) + if isinstance(obj, np.ndarray): - # If the array is non-contiguous, we need to copy it first - arr_data = ( - obj.view(np.uint8).data if obj.flags.c_contiguous else obj.tobytes() - ) + if obj.ndim == 0: + arr_data = obj.item() + elif obj.flags.c_contiguous: + # Not valid for 0-D arrays + arr_data = obj.view(np.uint8).data + else: + # If the array is non-contiguous, we need to copy it first + arr_data = obj.tobytes() + return cls.iter_item_to_bytes( "ndarray", { @@ -116,6 +123,7 @@ class MultiModalHasher: "data": arr_data, }, ) + logger.warning( "No serialization method found for %s. Falling back to pickle.", type(obj) ) @@ -147,7 +155,7 @@ class MultiModalHasher: hasher_factory = _get_hasher_factory(envs.VLLM_MM_HASHER_ALGORITHM) hasher = hasher_factory() - for k, v in kwargs.items(): + for k, v in sorted(kwargs.items(), key=lambda kv: kv[0]): for bytes_ in cls.iter_item_to_bytes(k, v): hasher.update(bytes_) diff --git a/vllm/multimodal/inputs.py b/vllm/multimodal/inputs.py index 262def71220..2cc7900eb67 100644 --- a/vllm/multimodal/inputs.py +++ b/vllm/multimodal/inputs.py @@ -424,8 +424,9 @@ class BaseMultiModalField(ABC): keep_on_cpu: bool = False """ - If `True`, then this field is excluded from being moved to the accelerator - when `MultiModalKwargsItems.get_data()` is called to batch the data. + If `True`, then this field is excluded from being moved to the accelerator when + [`group_and_batch_mm_items`][vllm.multimodal.utils.group_and_batch_mm_items] + is called to batch the data. """ def _field_factory(self): @@ -1006,27 +1007,38 @@ class MultiModalKwargsItems(UserDict[str, Sequence[_I]]): pin_memory: bool = False, ) -> BatchedTensorInputs: """Construct a dictionary of keyword arguments to pass to the model.""" - elems_by_key = defaultdict[str, list[MultiModalFieldElem]](list) - for modality, items in self.items(): - for i, item in enumerate(items): - if item is None: - raise RuntimeError( - f"Cannot build data from empty mm_items[{modality}][{i}]" - ) + from .utils import group_and_batch_mm_items - for key, elem in item.items(): - elems_by_key[key].append(elem) - - data = { - key: elems[0].field.reduce_data( - elems, - device=device, - pin_memory=pin_memory, - ) - for key, elems in elems_by_key.items() + items_by_modality = self.require_data() + batches_by_modality = { + modality: [ + data + for _, data in group_and_batch_mm_items( + items, + device=device, + pin_memory=pin_memory, + ) + ] + for modality, items in items_by_modality.items() + if len(items) > 0 } - return data + out_data: BatchedTensorInputs = {} + for _, batches in batches_by_modality.items(): + if len(batches) != 1: + num_batches_by_modality = { + modality: len(batches) + for modality, batches in batches_by_modality.items() + } + + raise RuntimeError( + f"Some modalities cannot be merged into a single batch " + f"({num_batches_by_modality=})" + ) + + out_data.update(batches[0]) + + return out_data MultiModalKwargsOptionalItems: TypeAlias = ( diff --git a/vllm/multimodal/utils.py b/vllm/multimodal/utils.py index cd116b9b8bc..d94faa67557 100644 --- a/vllm/multimodal/utils.py +++ b/vllm/multimodal/utils.py @@ -3,7 +3,8 @@ import mimetypes import warnings -from collections.abc import Generator +from collections import defaultdict +from collections.abc import Generator, Sequence from itertools import groupby from typing import TYPE_CHECKING, Any @@ -13,11 +14,13 @@ from PIL import Image from vllm.utils.import_utils import LazyLoader +from .hasher import MultiModalHasher from .inputs import ( BatchedTensorInputs, + MultiModalFieldElem, MultiModalKwargsItem, - MultiModalKwargsItems, MultiModalPlaceholderDict, + MultiModalSharedField, ) from .media import AudioMediaIO, ImageMediaIO, MediaConnector, VideoMediaIO @@ -146,32 +149,119 @@ def argsort_mm_positions( return [(modality, idx) for modality, idx, _ in sorted_flat_items] +def _get_group_hash(elem: MultiModalFieldElem): + if not isinstance(elem.field, MultiModalSharedField): + return None + + return MultiModalHasher.hash_kwargs(data=elem.data) + + +def _batch_mm_items( + items: Sequence[MultiModalKwargsItem], + *, + device: torch.types.Device = None, + pin_memory: bool = False, +): + elems = defaultdict[str, list[MultiModalFieldElem]](list) + for item in items: + for key, elem in item.items(): + elems[key].append(elem) + + return { + key: elems[0].field.reduce_data( + elems, + device=device, + pin_memory=pin_memory, + ) + for key, elems in elems.items() + } + + +def group_and_batch_mm_items( + items: Sequence[MultiModalKwargsItem], + *, + device: torch.types.Device = None, + pin_memory: bool = False, +) -> Generator[tuple[int, BatchedTensorInputs]]: + """ + Group consecutive items (possibly from different requests) into batches. + + Items must be split across groups if any of the following occurs, + as the batch would otherwise be invalid: + - They have different fields (e.g. mixed image and embedding inputs). + - They have different values in `MultiModalSharedField`. + + Args: + items: List of `MultiModalKwargsItem`. + device: The device to place the grouped tensors on. + pin_memory: Whether to pin memory for faster host-to-device transfer. + + Yields: + A tuple `(num_items, grouped_kwargs)`, where: + - `kwargs` is a dictionary of keyword arguments to pass to the model; + - `num_items` is the corresponding number of items. + """ + group_ids = [ + tuple( + (key, _get_group_hash(elem)) + for key, elem in sorted(item.items(), key=lambda kv: kv[0]) + ) + for item in items + ] + group_sizes = [sum(1 for _ in group) for _, group in groupby(group_ids)] + + start_idx = 0 + for group_size in group_sizes: + group_data = _batch_mm_items( + items[start_idx : start_idx + group_size], + device=device, + pin_memory=pin_memory, + ) + + yield group_size, group_data + + start_idx += group_size + + assert start_idx == len(items) + + def group_mm_kwargs_by_modality( mm_kwargs: list[tuple[str, MultiModalKwargsItem]], *, device: torch.types.Device = None, pin_memory: bool = False, ) -> Generator[tuple[str, int, BatchedTensorInputs], None, None]: - """Group consecutive `MultiModalKwargsItem`s from `mm_kwargs` with the same - modality together into the same `MultiModalKwargs` instance. + """ + Group consecutive items (possibly from different requests) into batches. + + Items must be split across groups if any of the following occurs, + as the batch would otherwise be invalid: + - They have different fields (e.g. mixed image and embedding inputs). + - They have different values in `MultiModalSharedField`. + + To simplify the implementation of `embed_multimodal`, we add another + restriction that the items in a batch must belong to the same modality. Args: - mm_kwargs: List of `MultiModalKwargsItem`. + mm_kwargs: List of `(modality, item)`. device: The device to place the grouped tensors on. pin_memory: Whether to pin memory for faster host-to-device transfer. Yields: - A tuple `(modality, num_items, grouped_kwargs)`. + A tuple `(modality, num_items, grouped_kwargs)`, where: + - `modality` is the modality of the batch; + - `kwargs` is a dictionary of keyword arguments to pass to the model; + - `num_items` is the corresponding number of items. """ for modality, group in groupby(mm_kwargs, key=lambda x: x[0]): items_lst = [item for _, item in group] - mm_kwargs_items = MultiModalKwargsItems({modality: items_lst}) - mm_kwargs_data = mm_kwargs_items.get_data( + + for num_items, mm_kwargs_batch in group_and_batch_mm_items( + items_lst, device=device, pin_memory=pin_memory, - ) - - yield modality, len(items_lst), mm_kwargs_data + ): + yield modality, num_items, mm_kwargs_batch def fetch_audio( From 42d5d705f93b254179e062003e8504fbe04f1b30 Mon Sep 17 00:00:00 2001 From: Lumosis <30372757+Lumosis@users.noreply.github.com> Date: Thu, 5 Feb 2026 14:05:09 -0800 Subject: [PATCH 113/810] [Minor] Sort safetensors files to ensure deterministic loading order (#33491) Signed-off-by: Lihao Ran Signed-off-by: mgoin Co-authored-by: mgoin --- vllm/model_executor/model_loader/weight_utils.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py index 15fd4423943..02998290b0a 100644 --- a/vllm/model_executor/model_loader/weight_utils.py +++ b/vllm/model_executor/model_loader/weight_utils.py @@ -19,6 +19,7 @@ from typing import IO, Any import filelock import huggingface_hub.constants import numpy as np +import regex as re import torch from huggingface_hub import HfFileSystem, hf_hub_download, snapshot_download from safetensors.torch import load, load_file, safe_open, save_file @@ -143,6 +144,15 @@ def atomic_writer( os.remove(temp_path) +def _natural_sort_key(filepath: str) -> list: + """Natural sort key for filenames with numeric components, such as + model-00001-of-00005.safetensors -> ['model-', 1, '-of-', 5, '.safetensors']""" + return [ + int(s) if s.isdigit() else s + for s in re.split(r"(\d+)", os.path.basename(filepath)) + ] + + def maybe_download_from_modelscope( model: str, revision: str | None = None, @@ -682,9 +692,8 @@ def safetensors_weights_iterator( loading_desc += " (eager)" leftover_state_dict: dict[str, torch.Tensor] = {} - for st_file in tqdm( - hf_weights_files, + sorted(hf_weights_files, key=_natural_sort_key), desc=loading_desc, disable=not enable_tqdm(use_tqdm_on_load), bar_format=_BAR_FORMAT, From d5c4800112c12bbcd4955858ef1b415c16ae16e7 Mon Sep 17 00:00:00 2001 From: Hashem Hashemi <159079214+amd-hhashemi@users.noreply.github.com> Date: Thu, 5 Feb 2026 14:16:02 -0800 Subject: [PATCH 114/810] Adds padding and perf improvements to wvSplitK_fp8 (#33527) Signed-off-by: Hashem Hashemi --- csrc/rocm/skinny_gemms.cu | 300 ++++++++---------- .../quantization/test_rocm_skinny_gemms.py | 92 +++--- .../quantization/kernels/scaled_mm/rocm.py | 6 +- 3 files changed, 169 insertions(+), 229 deletions(-) diff --git a/csrc/rocm/skinny_gemms.cu b/csrc/rocm/skinny_gemms.cu index a6cf63f22fd..770c94c5be5 100644 --- a/csrc/rocm/skinny_gemms.cu +++ b/csrc/rocm/skinny_gemms.cu @@ -1899,8 +1899,9 @@ torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b, template __global__ void __launch_bounds__(WvPrGrp* THRDS) - wvSplitKQ_hf_sml_(const int K, const int Kp, const int M, const int Bx, - const int By, const fp8_t* B, const fp8_t* __restrict__ A, + wvSplitKQ_hf_sml_(const int K, const int Kap, const int Kbp, const int M, + const int Bx, const int By, const fp8_t* B, + const fp8_t* __restrict__ A, const scalar_t* __restrict__ BIAS, scalar_t* C, const float* __restrict__ s_A, const float* __restrict__ s_B, const int _WvPrGrp, @@ -1924,9 +1925,14 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) __shared__ fp8_t s[max_lds_len]; for (uint32_t k = (threadIdx.y * THRDS + threadIdx.x) * A_CHUNK; - k < min__(K * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) { + k < min__(Kap * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) { + #if defined(__gfx950__) + __builtin_amdgcn_global_load_lds((int*)(&A[k]), (int*)(&s[k]), 16, 0, 0); + #else *((bigType*)(&s[k])) = *((bigType*)(&A[k])); + #endif } + asm volatile("s_waitcnt vmcnt(0)"); __syncthreads(); if (threadIdx.y >= _WvPrGrp) return; @@ -1934,37 +1940,24 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) uint32_t m = (blockIdx.x * _WvPrGrp + (threadIdx.y % _WvPrGrp)) * YTILE; using floatx16 = __attribute__((__vector_size__(16 * sizeof(float)))) float; - floatx16 sum[N][YTILE]; float sA = *s_A; float sB = *s_B; while (m < M) { - for (int i = 0; i < YTILE; i++) - for (int n = 0; n < N; n++) sum[n][i] = {0.f}; - - bigType bigA[N][UNRL]; - bigType bigB[YTILE][UNRL]; - + floatx16 sum[N][YTILE] = {}; for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) { - #pragma unroll - for (uint32_t k2 = 0; k2 < UNRL; k2++) { - #pragma unroll - for (uint32_t n = 0; n < N; ++n) bigA[n][k2].h8 = {0.f}; - #pragma unroll - for (uint32_t y = 0; y < YTILE; ++y) bigB[y][k2].h8 = {0.f}; - } + bigType bigA[N][UNRL] = {}; + bigType bigB[YTILE][UNRL]; // Fetch the weight matrix from memory! #pragma unroll for (uint32_t k2 = 0; k2 < UNRL; k2++) { uint32_t k = k1 + k2 * THRDS * A_CHUNK; uint32_t k_ = k + threadIdx.x * A_CHUNK; - if (k_ >= K) break; - - const fp8_t* B_ = &B[(m + 0) * Kp + k_]; + const fp8_t* B_ = &B[min__(k_, K - A_CHUNK)]; #pragma unroll for (uint32_t y = 0; y < YTILE; ++y) { - bigB[y][k2].h8 = (loadnt((scalar8*)(&B_[y * Kp]))); + bigB[y][k2].h8 = (loadnt((scalar8*)(&B_[min__(y + m, M - 1) * Kbp]))); } } @@ -1975,16 +1968,13 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) uint32_t k_ = k + threadIdx.x * A_CHUNK; if (k_ >= K) break; for (int n = 0; n < N; n++) { - bigA[n][k2] = *((const bigType*)(&(s[k_ + K * n]))); + bigA[n][k2] = *((const bigType*)(&(s[k_ + Kap * n]))); } } // Do the matrix multiplication in interleaved manner #pragma unroll for (uint32_t k2 = 0; k2 < UNRL; k2++) { - uint32_t k = k1 + k2 * THRDS * A_CHUNK; - if (k >= K) break; - for (uint32_t n = 0; n < N; n++) { for (int i = 0; i < A_CHUNK; i += 8) { for (int y = 0; y < YTILE; ++y) { @@ -2002,48 +1992,27 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) for (int y = 0; y < YTILE; y++) { float accm0 = sum[n][y][0]; float accm16 = sum[n][y][8]; - asm("v_add_f32 %0, %2, %3 row_shl:1 bound_ctrl:0 " - : "=v"(accm0) - : "0"(accm0), "v"(sum[n][y][1]), "v"(accm0)); - asm("v_add_f32 %0, %2, %3 row_shl:1 bound_ctrl:0 " - : "=v"(accm16) - : "0"(accm16), "v"(sum[n][y][9]), "v"(accm16)); - asm("v_add_f32 %0, %2, %3 row_shl:2 bound_ctrl:0 " - : "=v"(accm0) - : "0"(accm0), "v"(sum[n][y][2]), "v"(accm0)); - asm("v_add_f32 %0, %2, %3 row_shl:2 bound_ctrl:0 " - : "=v"(accm16) - : "0"(accm16), "v"(sum[n][y][10]), "v"(accm16)); - asm("v_add_f32 %0, %2, %3 row_shl:3 bound_ctrl:0 " - : "=v"(accm0) - : "0"(accm0), "v"(sum[n][y][3]), "v"(accm0)); - asm("v_add_f32 %0, %2, %3 row_shl:3 bound_ctrl:0 " - : "=v"(accm16) - : "0"(accm16), "v"(sum[n][y][11]), "v"(accm16)); - asm("v_add_f32 %0, %2, %3 row_shl:8 bound_ctrl:0 " - : "=v"(accm0) - : "0"(accm0), "v"(sum[n][y][4]), "v"(accm0)); - asm("v_add_f32 %0, %2, %3 row_shl:8 bound_ctrl:0 " - : "=v"(accm16) - : "0"(accm16), "v"(sum[n][y][12]), "v"(accm16)); - asm("v_add_f32 %0, %2, %3 row_shl:9 bound_ctrl:0 " - : "=v"(accm0) - : "0"(accm0), "v"(sum[n][y][5]), "v"(accm0)); - asm("v_add_f32 %0, %2, %3 row_shl:9 bound_ctrl:0 " - : "=v"(accm16) - : "0"(accm16), "v"(sum[n][y][13]), "v"(accm16)); - asm("v_add_f32 %0, %2, %3 row_shl:10 bound_ctrl:0 " - : "=v"(accm0) - : "0"(accm0), "v"(sum[n][y][6]), "v"(accm0)); - asm("v_add_f32 %0, %2, %3 row_shl:10 bound_ctrl:0 " - : "=v"(accm16) - : "0"(accm16), "v"(sum[n][y][14]), "v"(accm16)); - asm("v_add_f32 %0, %2, %3 row_shl:11 bound_ctrl:0 " - : "=v"(accm0) - : "0"(accm0), "v"(sum[n][y][7]), "v"(accm0)); - asm("v_add_f32 %0, %2, %3 row_shl:11 bound_ctrl:0 " - : "=v"(accm16) - : "0"(accm16), "v"(sum[n][y][15]), "v"(accm16)); + accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][1], 0x101, 0xf, 0xf, + 1); // row_shl1 + accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][9], 0x101, 0xf, 0xf, 1); + accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][2], 0x102, 0xf, 0xf, + 1); // row_shl2 + accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][10], 0x102, 0xf, 0xf, 1); + accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][3], 0x103, 0xf, 0xf, + 1); // row_shl3 + accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][11], 0x103, 0xf, 0xf, 1); + accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][4], 0x108, 0xf, 0xf, + 1); // row_shl8 + accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][12], 0x108, 0xf, 0xf, 1); + accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][5], 0x109, 0xf, 0xf, + 1); // row_shl9 + accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][13], 0x109, 0xf, 0xf, 1); + accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][6], 0x10a, 0xf, 0xf, + 1); // row_shl10 + accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][14], 0x10a, 0xf, 0xf, 1); + accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][7], 0x10b, 0xf, 0xf, + 1); // row_shl11 + accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][15], 0x10b, 0xf, 0xf, 1); accm0 += __shfl(accm0, 36); accm16 += __shfl(accm16, 52); sum[n][y][0] = accm0 + __shfl(accm16, 16); @@ -2051,19 +2020,23 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) } if (threadIdx.x == 0) { + scalar_t biases[N][YTILE] = {}; + if (BIAS) + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx]; + } + } for (int n = 0; n < N; n++) { for (int y = 0; y < YTILE; y++) { if (y + m >= M) break; // To avoid mem access fault. sum[n][y][0] *= sA * sB; if constexpr (std::is_same_v) { - if (BIAS) - sum[n][y][0] += __half2float(BIAS[(m + y) % Bx + (n % By) * M]); + sum[n][y][0] += __half2float(biases[n][y]); } else if constexpr (std::is_same_v) { - if (BIAS) - sum[n][y][0] += - __bfloat162float(BIAS[(m + y) % Bx + (n % By) * M]); + sum[n][y][0] += __bfloat162float(biases[n][y]); } - C[m + y + n * M] = __float2s(sum[n][y][0]); // * sA * sB); + C[m + y + n * M] = __float2s(sum[n][y][0]); } } } @@ -2074,9 +2047,9 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) #else // !defined(__HIP__MI3XX__) TODO: Add NAVI support template -__global__ void wvSplitKQ_hf_sml_(const int K, const int Kp, const int M, - const int Bx, const int By, const fp8_t* B, - const fp8_t* __restrict__ A, +__global__ void wvSplitKQ_hf_sml_(const int K, const int Kap, const int Kbp, + const int M, const int Bx, const int By, + const fp8_t* B, const fp8_t* __restrict__ A, const scalar_t* __restrict__ BIAS, scalar_t* C, const float* __restrict__ s_A, const float* __restrict__ s_B, @@ -2089,8 +2062,9 @@ __global__ void wvSplitKQ_hf_sml_(const int K, const int Kp, const int M, template __global__ void __launch_bounds__(WvPrGrp* THRDS) - wvSplitKQ_hf_(const int K, const int Kp, const int M, const int Bx, - const int By, const fp8_t* B, const fp8_t* __restrict__ A, + wvSplitKQ_hf_(const int K, const int Kap, const int Kbp, const int M, + const int Bx, const int By, const fp8_t* B, + const fp8_t* __restrict__ A, const scalar_t* __restrict__ BIAS, scalar_t* C, const float* __restrict__ s_A, const float* __restrict__ s_B, const int _WvPrGrp, const int CuCount) { @@ -2113,9 +2087,14 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) __shared__ fp8_t s[max_lds_len]; for (uint32_t k = (threadIdx.y * THRDS + threadIdx.x) * A_CHUNK; - k < min__(K * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) { + k < min__(Kap * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) { + #if defined(__gfx950__) + __builtin_amdgcn_global_load_lds((int*)(&A[k]), (int*)(&s[k]), 16, 0, 0); + #else *((bigType*)(&s[k])) = *((bigType*)(&A[k])); + #endif } + asm volatile("s_waitcnt vmcnt(0)"); __syncthreads(); if (threadIdx.y >= _WvPrGrp) return; @@ -2123,29 +2102,23 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) uint32_t m = (blockIdx.x * _WvPrGrp + (threadIdx.y % _WvPrGrp)) * YTILE; using floatx16 = __attribute__((__vector_size__(16 * sizeof(float)))) float; - floatx16 sum[N][YTILE]; float sA = *s_A; float sB = *s_B; while (m < M) { - for (int i = 0; i < YTILE; i++) - for (int n = 0; n < N; n++) sum[n][i] = {0}; - - bigType bigA[N][UNRL]; - bigType bigB[YTILE][UNRL]; - + floatx16 sum[N][YTILE] = {}; for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) { + bigType bigA[N][UNRL] = {}; + bigType bigB[YTILE][UNRL]; + // Fetch the weight matrix from memory! #pragma unroll for (uint32_t k2 = 0; k2 < UNRL; k2++) { uint32_t k = k1 + k2 * THRDS * A_CHUNK; uint32_t k_ = k + threadIdx.x * A_CHUNK; - if (k_ >= K) break; - - const fp8_t* B_ = &B[(m + 0) * Kp + k_]; + const fp8_t* B_ = &B[min__(k_, K - A_CHUNK)]; for (int y = 0; y < YTILE; ++y) { - if (y + m >= M) break; // To avoid mem access fault. - bigB[y][k2].h8 = (loadnt((scalar8*)(&B_[y * Kp]))); + bigB[y][k2].h8 = (loadnt((scalar8*)(&B_[min__(y + m, M - 1) * Kbp]))); } } @@ -2156,20 +2129,16 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) uint32_t k_ = k + threadIdx.x * A_CHUNK; if (k_ >= K) break; for (int n = 0; n < N; n++) { - if (k_ + K * n < max_lds_len) - bigA[n][k2] = *((const bigType*)(&(s[k_ + K * n]))); + if (k_ + Kap * n < max_lds_len) + bigA[n][k2] = *((const bigType*)(&(s[k_ + Kap * n]))); else - bigA[n][k2] = *((const bigType*)(&(A[k_ + K * n]))); + bigA[n][k2] = *((const bigType*)(&(A[k_ + Kap * n]))); } } // Do the matrix multiplication in interleaved manner #pragma unroll for (uint32_t k2 = 0; k2 < UNRL; k2++) { - uint32_t k = k1 + k2 * THRDS * A_CHUNK; - uint32_t k_ = k + threadIdx.x * A_CHUNK; - if (k_ >= K) break; - for (uint32_t n = 0; n < N; n++) { for (int i = 0; i < A_CHUNK; i += 8) { for (int y = 0; y < YTILE; ++y) { @@ -2187,48 +2156,27 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) for (int y = 0; y < YTILE; y++) { float accm0 = sum[n][y][0]; float accm16 = sum[n][y][8]; - asm("v_add_f32 %0, %2, %3 row_shl:1 bound_ctrl:0 " - : "=v"(accm0) - : "0"(accm0), "v"(sum[n][y][1]), "v"(accm0)); - asm("v_add_f32 %0, %2, %3 row_shl:1 bound_ctrl:0 " - : "=v"(accm16) - : "0"(accm16), "v"(sum[n][y][9]), "v"(accm16)); - asm("v_add_f32 %0, %2, %3 row_shl:2 bound_ctrl:0 " - : "=v"(accm0) - : "0"(accm0), "v"(sum[n][y][2]), "v"(accm0)); - asm("v_add_f32 %0, %2, %3 row_shl:2 bound_ctrl:0 " - : "=v"(accm16) - : "0"(accm16), "v"(sum[n][y][10]), "v"(accm16)); - asm("v_add_f32 %0, %2, %3 row_shl:3 bound_ctrl:0 " - : "=v"(accm0) - : "0"(accm0), "v"(sum[n][y][3]), "v"(accm0)); - asm("v_add_f32 %0, %2, %3 row_shl:3 bound_ctrl:0 " - : "=v"(accm16) - : "0"(accm16), "v"(sum[n][y][11]), "v"(accm16)); - asm("v_add_f32 %0, %2, %3 row_shl:8 bound_ctrl:0 " - : "=v"(accm0) - : "0"(accm0), "v"(sum[n][y][4]), "v"(accm0)); - asm("v_add_f32 %0, %2, %3 row_shl:8 bound_ctrl:0 " - : "=v"(accm16) - : "0"(accm16), "v"(sum[n][y][12]), "v"(accm16)); - asm("v_add_f32 %0, %2, %3 row_shl:9 bound_ctrl:0 " - : "=v"(accm0) - : "0"(accm0), "v"(sum[n][y][5]), "v"(accm0)); - asm("v_add_f32 %0, %2, %3 row_shl:9 bound_ctrl:0 " - : "=v"(accm16) - : "0"(accm16), "v"(sum[n][y][13]), "v"(accm16)); - asm("v_add_f32 %0, %2, %3 row_shl:10 bound_ctrl:0 " - : "=v"(accm0) - : "0"(accm0), "v"(sum[n][y][6]), "v"(accm0)); - asm("v_add_f32 %0, %2, %3 row_shl:10 bound_ctrl:0 " - : "=v"(accm16) - : "0"(accm16), "v"(sum[n][y][14]), "v"(accm16)); - asm("v_add_f32 %0, %2, %3 row_shl:11 bound_ctrl:0 " - : "=v"(accm0) - : "0"(accm0), "v"(sum[n][y][7]), "v"(accm0)); - asm("v_add_f32 %0, %2, %3 row_shl:11 bound_ctrl:0 " - : "=v"(accm16) - : "0"(accm16), "v"(sum[n][y][15]), "v"(accm16)); + accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][1], 0x101, 0xf, 0xf, + 1); // row_shl1 + accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][9], 0x101, 0xf, 0xf, 1); + accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][2], 0x102, 0xf, 0xf, + 1); // row_shl2 + accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][10], 0x102, 0xf, 0xf, 1); + accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][3], 0x103, 0xf, 0xf, + 1); // row_shl3 + accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][11], 0x103, 0xf, 0xf, 1); + accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][4], 0x108, 0xf, 0xf, + 1); // row_shl8 + accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][12], 0x108, 0xf, 0xf, 1); + accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][5], 0x109, 0xf, 0xf, + 1); // row_shl9 + accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][13], 0x109, 0xf, 0xf, 1); + accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][6], 0x10a, 0xf, 0xf, + 1); // row_shl10 + accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][14], 0x10a, 0xf, 0xf, 1); + accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][7], 0x10b, 0xf, 0xf, + 1); // row_shl11 + accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][15], 0x10b, 0xf, 0xf, 1); accm0 += __shfl(accm0, 36); accm16 += __shfl(accm16, 52); sum[n][y][0] = accm0 + __shfl(accm16, 16); @@ -2236,17 +2184,21 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) } if (threadIdx.x == 0) { + scalar_t biases[N][YTILE] = {}; + if (BIAS) + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx]; + } + } for (int n = 0; n < N; n++) { for (int y = 0; y < YTILE; y++) { if (y + m >= M) break; // To avoid mem access fault. sum[n][y][0] *= sA * sB; if constexpr (std::is_same_v) { - if (BIAS) - sum[n][y][0] += __half2float(BIAS[(m + y) % Bx + (n % By) * M]); + sum[n][y][0] += __half2float(biases[n][y]); } else if constexpr (std::is_same_v) { - if (BIAS) - sum[n][y][0] += - __bfloat162float(BIAS[(m + y) % Bx + (n % By) * M]); + sum[n][y][0] += __bfloat162float(biases[n][y]); } C[m + y + n * M] = __float2s(sum[n][y][0]); } @@ -2259,9 +2211,9 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) #else // !defined(__HIP__MI3XX__) TODO: Add NAVI support template -__global__ void wvSplitKQ_hf_(const int K, const int Kp, const int M, - const int Bx, const int By, const fp8_t* B, - const fp8_t* __restrict__ A, +__global__ void wvSplitKQ_hf_(const int K, const int Kap, const int Kbp, + const int M, const int Bx, const int By, + const fp8_t* B, const fp8_t* __restrict__ A, const scalar_t* __restrict__ BIAS, scalar_t* C, const float* __restrict__ s_A, const float* __restrict__ s_B, const int _WvPrGrp, @@ -2270,17 +2222,18 @@ __global__ void wvSplitKQ_hf_(const int K, const int Kp, const int M, } #endif // defined(__HIP__MI3XX__) TODO: Add NAVI support -void wvSplitKQ(const at::Tensor& in_a, const at::Tensor& in_b, +void wvSplitKQ(const at::Tensor& in_b, const at::Tensor& in_a, const std::optional& in_bias, at::Tensor& out_c, const at::Tensor& scale_a, const at::Tensor& scale_b, const int64_t CuCount) { static c10::ScalarType kFp8Type = is_fp8_ocp() ? c10::ScalarType::Float8_e4m3fn : c10::ScalarType::Float8_e4m3fnuz; - auto M_in = in_a.size(0); - auto K_in = in_a.size(1); - auto N_in = in_b.size(0); - auto Kp_in = in_a.stride(0); + auto M_in = in_b.size(0); + auto K_in = in_b.size(1); + auto N_in = in_a.size(0); + auto Kap_in = in_a.stride(0); + auto Kbp_in = in_b.stride(0); auto Bx_in = (in_bias.has_value() && in_bias->numel() > 0) ? (in_bias->sizes().size() == 2) ? in_bias->size(1) : in_bias->size(0) @@ -2300,23 +2253,22 @@ void wvSplitKQ(const at::Tensor& in_a, const at::Tensor& in_b, const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); const int max_lds_len = get_lds_size(); -#define WVSPLITKQ(_WvPrGrp, _YTILEs, _YTILEm, _YTILEb, _UNRLs, _UNRLm, _UNRLb, \ - _N) \ - { \ - dim3 block(64, _WvPrGrp); \ - if ((K_in * N_in <= max_lds_len) && (M_in % _YTILEs == 0)) { \ - int __wvPrGrp = mindiv(M_in, CuCount * _YTILEs, _WvPrGrp); \ - wvSplitKQ_hf_sml_ \ - <<>>(K_in, Kp_in, M_in, Bx_in, By_in, a_ptr, \ - b_ptr, bias_ptr, c_ptr, s_a, s_b, \ - __wvPrGrp, CuCount); \ - } else { \ - int __wvPrGrp = mindiv(M_in, CuCount * _YTILEm, _WvPrGrp); \ - wvSplitKQ_hf_ \ - <<>>(K_in, Kp_in, M_in, Bx_in, By_in, a_ptr, \ - b_ptr, bias_ptr, c_ptr, s_a, s_b, \ - __wvPrGrp, CuCount); \ - } \ +#define WVSPLITKQ(_WvPrGrp, _YTILEs, _YTILEm, _UNRLs, _UNRLm, _N) \ + { \ + dim3 block(64, _WvPrGrp); \ + if ((Kap_in * N_in <= max_lds_len) && (M_in % _YTILEs == 0)) { \ + int __wvPrGrp = min(_WvPrGrp, mindiv(M_in, CuCount * _YTILEs, 16)); \ + wvSplitKQ_hf_sml_ \ + <<>>(K_in, Kap_in, Kbp_in, M_in, Bx_in, \ + By_in, b_ptr, a_ptr, bias_ptr, c_ptr, \ + s_a, s_b, __wvPrGrp, CuCount); \ + } else { \ + int __wvPrGrp = min(_WvPrGrp, mindiv(M_in, CuCount * _YTILEm, 16)); \ + wvSplitKQ_hf_ \ + <<>>(K_in, Kap_in, Kbp_in, M_in, Bx_in, \ + By_in, b_ptr, a_ptr, bias_ptr, c_ptr, \ + s_a, s_b, __wvPrGrp, CuCount); \ + } \ } AT_DISPATCH_REDUCED_FLOATING_TYPES(out_c.scalar_type(), "wvSplitKQ", [&] { @@ -2332,16 +2284,16 @@ void wvSplitKQ(const at::Tensor& in_a, const at::Tensor& in_b, : nullptr; switch (N_in) { case 1: - WVSPLITKQ(16, 2, 2, 2, 2, 2, 2, 1) + WVSPLITKQ(12, 2, 2, 2, 2, 1) break; case 2: - WVSPLITKQ(16, 2, 2, 2, 2, 2, 2, 2) + WVSPLITKQ(12, 2, 2, 2, 2, 2) break; case 3: - WVSPLITKQ(16, 4, 7, 7, 1, 1, 1, 3) + WVSPLITKQ(8, 2, 2, 1, 1, 3) break; case 4: - WVSPLITKQ(16, 4, 7, 7, 1, 1, 1, 4) + WVSPLITKQ(4, 2, 2, 1, 1, 4) break; default: throw std::runtime_error( diff --git a/tests/kernels/quantization/test_rocm_skinny_gemms.py b/tests/kernels/quantization/test_rocm_skinny_gemms.py index 1505604a691..474339ce4fa 100644 --- a/tests/kernels/quantization/test_rocm_skinny_gemms.py +++ b/tests/kernels/quantization/test_rocm_skinny_gemms.py @@ -73,21 +73,40 @@ NKM_FACTORS_WVSPLITKRC = [ NKM_FACTORS_WVSPLITK_FP8 = [ # FP8-specific cases with K % 16 == 0 (1, 16, 16), + (1, 32, 16 + 16), (1, 64, 64), + (1, 64, 64 + 16), + (1, 64 + 16, 64), + (1, 64 + 16, 64 + 16), + (4, 64, 64), + (4, 64, 64 + 16), + (4, 64 + 16, 64), + (4, 64 + 16, 64 + 16), (2, 512, 512), + (3, 512, 512), + (3, 512, 512 + 16), + (4, 512, 512), (3, 2048, 2048), + (3, 2048, 2048 + 16), + (4, 2048 + 16, 2048), + (4, 2048 + 16, 2048 + 16), (4, 4096, 4096), (4, 16400, 2048), + (4, 16400, 2048 + 16), # Extended FP8 dimensions not covered by WVSPLITK (1, 14336, 1024), (2, 24576, 2048), (4, 32768, 28672), + (4, 32768 * 2, 28672), + (4, 32768 * 2, 28672 + 16), + (4, 32768 * 2 + 16, 28672), + (4, 32768 * 2 + 16, 28672 + 16), ] SEEDS = [0] -def pad_weights_fp8(weight): +def pad_fp8(weight): num_pad = 256 // weight.element_size() import torch.nn.functional as F @@ -195,72 +214,41 @@ def test_rocm_wvsplitk_bias2D_kernel(n, k, m, dtype, seed): assert torch.allclose(out, ref_out, rtol=0.01) +@pytest.mark.parametrize("xnorm", [False, True]) @pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK_FP8) @pytest.mark.parametrize("dtype", DTYPES) @pytest.mark.parametrize("seed", SEEDS) -@pytest.mark.parametrize("padded", [False, True]) +@pytest.mark.parametrize("padded_a", [False, True]) +@pytest.mark.parametrize("padded_b", [False, True]) +@pytest.mark.parametrize("biased", [False, True]) @pytest.mark.skipif( not (current_platform.is_rocm() and current_platform.supports_fp8()), reason="only test for rocm fp8", ) -def test_rocm_wvsplitk_fp8_kernel(n, k, m, dtype, seed, padded): +def test_rocm_wvsplitk_fp8_kernel( + xnorm, n, k, m, dtype, seed, padded_a, padded_b, biased +): torch.manual_seed(seed) - A = torch.rand(n, k, device="cuda") - 0.5 - B = torch.rand(m, k, device="cuda") - 0.5 + xavier = math.sqrt(2 / k) if xnorm else 1 # normalize to avoid large deltas + A = (torch.rand(n, k, device="cuda") * 2 - 1) * xavier + B = (torch.rand(m, k, device="cuda") * 2 - 1) * xavier A, scale_a = ref_dynamic_per_tensor_fp8_quant(A) B, scale_b = ref_dynamic_per_tensor_fp8_quant(B) - if padded: - B = pad_weights_fp8(B) + if padded_b: + B = pad_fp8(B) + if padded_a: + A = pad_fp8(A) - ref_out = torch._scaled_mm( - A, B.t(), out_dtype=dtype, scale_a=scale_a, scale_b=scale_b - ) - out = ops.wvSplitKQ( - B, - A, - dtype, - scale_a, - scale_b, - get_cu_count(), - ) - - assert torch.allclose(out, ref_out, rtol=0.01) - - -@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK_FP8) -@pytest.mark.parametrize("dtype", DTYPES) -@pytest.mark.parametrize("seed", SEEDS) -@pytest.mark.parametrize("padded", [False, True]) -@pytest.mark.skipif( - not (current_platform.is_rocm() and current_platform.supports_fp8()), - reason="only test for rocm fp8", -) -def test_rocm_wvsplitk_fp8_bias1D_kernel(n, k, m, dtype, seed, padded): - torch.manual_seed(seed) - - xavier = math.sqrt(2 / k) # normalize to avoid large output-bias deltas - A = (torch.rand(n, k, device="cuda") - 0.5) * xavier - B = (torch.rand(m, k, device="cuda") - 0.5) * xavier - BIAS = torch.rand(m, dtype=dtype, device="cuda") - 0.5 - - A, scale_a = ref_dynamic_per_tensor_fp8_quant(A) - B, scale_b = ref_dynamic_per_tensor_fp8_quant(B) - if padded: - B = pad_weights_fp8(B) + BIAS = None if (not biased) else (torch.rand(m, dtype=dtype, device="cuda") * 2 - 1) ref_out = torch._scaled_mm( A, B.t(), out_dtype=dtype, scale_a=scale_a, scale_b=scale_b, bias=BIAS ) - out = ops.wvSplitKQ( - B, - A, - dtype, - scale_a, - scale_b, - get_cu_count(), - BIAS, - ) + out = ops.wvSplitKQ(B, A, dtype, scale_a, scale_b, get_cu_count(), BIAS) - assert torch.allclose(out, ref_out, rtol=0.01) + if xnorm: + assert torch.allclose(out, ref_out, atol=1e-3, rtol=1e-8) + else: + assert torch.allclose(out, ref_out, 0.01) diff --git a/vllm/model_executor/layers/quantization/kernels/scaled_mm/rocm.py b/vllm/model_executor/layers/quantization/kernels/scaled_mm/rocm.py index ee660812e9f..7a95296245b 100644 --- a/vllm/model_executor/layers/quantization/kernels/scaled_mm/rocm.py +++ b/vllm/model_executor/layers/quantization/kernels/scaled_mm/rocm.py @@ -25,10 +25,10 @@ def rocm_per_tensor_float_w8a8_scaled_mm_impl( bias: torch.Tensor, ) -> torch.Tensor: if ( - A.shape[0] == 1 - and B.shape[1] % 16 == 0 + A.shape[0] <= 4 + and B.shape[0] % 16 == 0 # M TODO: needed? + and B.shape[1] % 16 == 0 # K and ((bias is None) or (bias.dtype == out_dtype)) - and A.is_contiguous() ): output = ops.wvSplitKQ( B.t(), From 91a07ff6187e7308794b2a4863ab9e1f821ed464 Mon Sep 17 00:00:00 2001 From: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Date: Thu, 5 Feb 2026 18:50:49 -0500 Subject: [PATCH 115/810] [Bugfix] Fix DeepSeek v3.2 tokenizer outputting None issue (#33832) Signed-off-by: wzhao18 --- vllm/tokenizers/detokenizer_utils.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/vllm/tokenizers/detokenizer_utils.py b/vllm/tokenizers/detokenizer_utils.py index e586a5d46cb..8e73d5dc537 100644 --- a/vllm/tokenizers/detokenizer_utils.py +++ b/vllm/tokenizers/detokenizer_utils.py @@ -157,6 +157,10 @@ def detokenize_incrementally( ) if isinstance(new_tokens, str): new_tokens = [new_tokens] + else: + # This is required to guard against out-of-vocab prompt token ids + # (for example when using dummy weights) + _replace_none_with_empty(new_tokens) # type: ignore[arg-type] else: new_tokens = [""] output_tokens = prev_tokens + new_tokens From 325ab6b0a896f4b483b26d99afdfc6d75ea61074 Mon Sep 17 00:00:00 2001 From: emricksini-h Date: Fri, 6 Feb 2026 01:59:28 +0100 Subject: [PATCH 116/810] [Feature] OTEL tracing during loading (#31162) --- requirements/common.txt | 2 +- setup.py | 7 + tests/tracing/__init__.py | 0 tests/tracing/conftest.py | 127 +++++++++ tests/tracing/test_loading_tracing.py | 87 ++++++ tests/v1/tracing/test_tracing.py | 107 ++----- vllm/compilation/backends.py | 8 + vllm/config/observability.py | 4 +- vllm/entrypoints/openai/api_server.py | 2 + .../model_loader/base_loader.py | 2 + .../model_loader/default_loader.py | 2 + vllm/model_executor/model_loader/utils.py | 2 + .../model_loader/weight_utils.py | 2 + .../model_executor/warmup/deep_gemm_warmup.py | 2 + vllm/tracing.py | 135 --------- vllm/tracing/__init__.py | 157 +++++++++++ vllm/tracing/otel.py | 265 ++++++++++++++++++ vllm/tracing/utils.py | 72 +++++ vllm/v1/engine/async_llm.py | 10 +- vllm/v1/engine/core.py | 21 ++ vllm/v1/engine/core_client.py | 4 + vllm/v1/engine/llm_engine.py | 4 +- vllm/v1/engine/output_processor.py | 106 +++---- vllm/v1/executor/abstract.py | 2 + vllm/v1/executor/multiproc_executor.py | 11 + vllm/v1/worker/cpu_model_runner.py | 3 + vllm/v1/worker/gpu_model_runner.py | 3 + vllm/v1/worker/gpu_worker.py | 4 + vllm/v1/worker/worker_base.py | 2 + 29 files changed, 873 insertions(+), 280 deletions(-) create mode 100644 tests/tracing/__init__.py create mode 100644 tests/tracing/conftest.py create mode 100644 tests/tracing/test_loading_tracing.py delete mode 100644 vllm/tracing.py create mode 100644 vllm/tracing/__init__.py create mode 100644 vllm/tracing/otel.py create mode 100644 vllm/tracing/utils.py diff --git a/requirements/common.txt b/requirements/common.txt index bc170d90b6d..f8402410bc9 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -52,4 +52,4 @@ anthropic >= 0.71.0 model-hosting-container-standards >= 0.1.13, < 1.0.0 mcp grpcio -grpcio-reflection +grpcio-reflection \ No newline at end of file diff --git a/setup.py b/setup.py index 7caaa684688..14325cdfcda 100644 --- a/setup.py +++ b/setup.py @@ -1049,6 +1049,13 @@ setup( "petit-kernel": ["petit-kernel"], # Optional deps for Helion kernel development "helion": ["helion"], + # Optional deps for OpenTelemetry tracing + "otel": [ + "opentelemetry-sdk>=1.26.0", + "opentelemetry-api>=1.26.0", + "opentelemetry-exporter-otlp>=1.26.0", + "opentelemetry-semantic-conventions-ai>=0.4.1", + ], }, cmdclass=cmdclass, package_data=package_data, diff --git a/tests/tracing/__init__.py b/tests/tracing/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/tracing/conftest.py b/tests/tracing/conftest.py new file mode 100644 index 00000000000..d29933ba8d5 --- /dev/null +++ b/tests/tracing/conftest.py @@ -0,0 +1,127 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import threading +from collections.abc import Callable, Generator, Iterable +from concurrent import futures +from typing import Any, Literal + +import grpc +import pytest +from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ( + ExportTraceServiceRequest, + ExportTraceServiceResponse, +) +from opentelemetry.proto.collector.trace.v1.trace_service_pb2_grpc import ( + TraceServiceServicer, + add_TraceServiceServicer_to_server, +) +from opentelemetry.proto.common.v1.common_pb2 import AnyValue, KeyValue + +FAKE_TRACE_SERVER_ADDRESS = "localhost:4317" + +FieldName = Literal[ + "bool_value", "string_value", "int_value", "double_value", "array_value" +] + + +def decode_value(value: AnyValue): + """Decode an OpenTelemetry AnyValue protobuf message to a Python value.""" + field_decoders: dict[FieldName, Callable] = { + "bool_value": (lambda v: v.bool_value), + "string_value": (lambda v: v.string_value), + "int_value": (lambda v: v.int_value), + "double_value": (lambda v: v.double_value), + "array_value": ( + lambda v: [decode_value(item) for item in v.array_value.values] + ), + } + for field, decoder in field_decoders.items(): + if value.HasField(field): + return decoder(value) + raise ValueError(f"Couldn't decode value: {value}") + + +def decode_attributes(attributes: Iterable[KeyValue]) -> dict[str, Any]: + """Decode OpenTelemetry KeyValue attributes to a Python dictionary.""" + return {kv.key: decode_value(kv.value) for kv in attributes} + + +class FakeTraceService(TraceServiceServicer): + """A fake gRPC trace service for testing OpenTelemetry trace exports.""" + + def __init__(self): + self.requests: list[ExportTraceServiceRequest] = [] + self.evt = threading.Event() + self._lock = threading.Lock() + + def Export(self, request, context): + with self._lock: + self.requests.append(request) + self.evt.set() + return ExportTraceServiceResponse() + + @property + def request(self) -> ExportTraceServiceRequest | None: + """Returns the first request received (for backward compatibility).""" + with self._lock: + return self.requests[0] if self.requests else None + + def get_all_spans(self) -> list[dict]: + """Returns all spans from all received requests as decoded dicts.""" + spans = [] + with self._lock: + for request in self.requests: + for resource_span in request.resource_spans: + for scope_span in resource_span.scope_spans: + for span in scope_span.spans: + spans.append( + { + "name": span.name, + "attributes": decode_attributes(span.attributes), + "trace_id": span.trace_id.hex(), + "span_id": span.span_id.hex(), + "parent_span_id": span.parent_span_id.hex() + if span.parent_span_id + else None, + "start_time_unix_nano": span.start_time_unix_nano, + "end_time_unix_nano": span.end_time_unix_nano, + } + ) + return spans + + def wait_for_spans(self, count: int = 1, timeout: float = 10) -> bool: + """Wait until at least `count` spans have been received.""" + import time + + deadline = time.time() + timeout + while time.time() < deadline: + if len(self.get_all_spans()) >= count: + return True + time.sleep(0.1) + return False + + def clear(self): + """Clear all received requests.""" + with self._lock: + self.requests.clear() + self.evt.clear() + + +@pytest.fixture +def trace_service() -> Generator[FakeTraceService, None, None]: + """Fixture to set up a fake gRPC trace service.""" + server = grpc.server(futures.ThreadPoolExecutor(max_workers=2)) + service = FakeTraceService() + add_TraceServiceServicer_to_server(service, server) + server.add_insecure_port(FAKE_TRACE_SERVER_ADDRESS) + server.start() + + yield service + + server.stop(grace=None) + + +@pytest.fixture +def trace_server_address() -> str: + """Returns the address of the fake trace server.""" + return FAKE_TRACE_SERVER_ADDRESS diff --git a/tests/tracing/test_loading_tracing.py b/tests/tracing/test_loading_tracing.py new file mode 100644 index 00000000000..e7cb3c838ff --- /dev/null +++ b/tests/tracing/test_loading_tracing.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import asyncio + +import pytest +from opentelemetry.sdk.environment_variables import OTEL_EXPORTER_OTLP_TRACES_INSECURE + +from tests.tracing.conftest import FAKE_TRACE_SERVER_ADDRESS, FakeTraceService +from vllm.tracing import init_tracer, instrument, is_otel_available + +# Skip everything if OTel is missing +pytestmark = pytest.mark.skipif(not is_otel_available(), reason="OTel required") + + +class TestCoreInstrumentation: + """Focuses on the @instrument decorator's ability to capture execution data.""" + + @pytest.fixture(autouse=True) + def setup_tracing(self, monkeypatch): + monkeypatch.setenv(OTEL_EXPORTER_OTLP_TRACES_INSECURE, "true") + init_tracer("test.core", FAKE_TRACE_SERVER_ADDRESS) + + def test_decorator_captures_sync_and_async(self, trace_service: FakeTraceService): + """Verify basic span creation for both sync and async functions.""" + + @instrument(span_name="sync_task") + def sync_task(): + return True + + @instrument(span_name="async_task") + async def async_task(): + return True + + sync_task() + asyncio.run(async_task()) + + assert trace_service.wait_for_spans(count=2) + span_names = [s["name"] for s in trace_service.get_all_spans()] + assert "sync_task" in span_names + assert "async_task" in span_names + + def test_nested_spans_hierarchy(self, trace_service: FakeTraceService): + """Verify that nested calls create a parent-child relationship.""" + + @instrument(span_name="child") + def child(): + pass + + @instrument(span_name="parent") + def parent(): + child() + + parent() + + assert trace_service.wait_for_spans(count=2) + spans = trace_service.get_all_spans() + parent_span = next(s for s in spans if s["name"] == "parent") + child_span = next(s for s in spans if s["name"] == "child") + + assert child_span["parent_span_id"] == parent_span["span_id"] + + +class TestInterProcessPropagation: + """Test the propagation of trace context between processes.""" + + def test_pickup_external_context(self, monkeypatch, trace_service): + """Test that vLLM attaches to an existing trace ID if in environment.""" + monkeypatch.setenv(OTEL_EXPORTER_OTLP_TRACES_INSECURE, "true") + + # Manually simulate an external parent trace ID + fake_trace_id = "4bf92f3577b34da6a3ce929d0e0e4736" + fake_parent_id = "00f067aa0ba902b7" + monkeypatch.setenv("traceparent", f"00-{fake_trace_id}-{fake_parent_id}-01") + + init_tracer("test.external", FAKE_TRACE_SERVER_ADDRESS) + + @instrument(span_name="follower") + def follower_func(): + pass + + follower_func() + + assert trace_service.wait_for_spans(count=1) + span = trace_service.get_all_spans()[0] + + assert span["trace_id"] == fake_trace_id + assert span["parent_span_id"] == fake_parent_id diff --git a/tests/v1/tracing/test_tracing.py b/tests/v1/tracing/test_tracing.py index 11d9d18ead7..2b450a6299c 100644 --- a/tests/v1/tracing/test_tracing.py +++ b/tests/v1/tracing/test_tracing.py @@ -2,76 +2,19 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # ruff: noqa # type: ignore -import threading -from collections.abc import Iterable -from concurrent import futures -from typing import Callable, Generator, Literal - -import grpc import pytest -from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ( - ExportTraceServiceResponse, -) -from opentelemetry.proto.collector.trace.v1.trace_service_pb2_grpc import ( - TraceServiceServicer, - add_TraceServiceServicer_to_server, -) -from opentelemetry.proto.common.v1.common_pb2 import AnyValue, KeyValue +import time from opentelemetry.sdk.environment_variables import OTEL_EXPORTER_OTLP_TRACES_INSECURE from vllm import LLM, SamplingParams from vllm.tracing import SpanAttributes -FAKE_TRACE_SERVER_ADDRESS = "localhost:4317" - -FieldName = Literal[ - "bool_value", "string_value", "int_value", "double_value", "array_value" -] - - -def decode_value(value: AnyValue): - field_decoders: dict[FieldName, Callable] = { - "bool_value": (lambda v: v.bool_value), - "string_value": (lambda v: v.string_value), - "int_value": (lambda v: v.int_value), - "double_value": (lambda v: v.double_value), - "array_value": ( - lambda v: [decode_value(item) for item in v.array_value.values] - ), - } - for field, decoder in field_decoders.items(): - if value.HasField(field): - return decoder(value) - raise ValueError(f"Couldn't decode value: {value}") - - -def decode_attributes(attributes: Iterable[KeyValue]): - return {kv.key: decode_value(kv.value) for kv in attributes} - - -class FakeTraceService(TraceServiceServicer): - def __init__(self): - self.request = None - self.evt = threading.Event() - - def Export(self, request, context): - self.request = request - self.evt.set() - return ExportTraceServiceResponse() - - -@pytest.fixture -def trace_service() -> Generator[FakeTraceService, None, None]: - """Fixture to set up a fake gRPC trace service""" - server = grpc.server(futures.ThreadPoolExecutor(max_workers=1)) - service = FakeTraceService() - add_TraceServiceServicer_to_server(service, server) - server.add_insecure_port(FAKE_TRACE_SERVER_ADDRESS) - server.start() - - yield service - - server.stop(None) +# Import shared fixtures from the tracing conftest +from tests.tracing.conftest import ( # noqa: F401 + FAKE_TRACE_SERVER_ADDRESS, + FakeTraceService, + trace_service, +) def test_traces( @@ -97,29 +40,25 @@ def test_traces( outputs = llm.generate(prompts, sampling_params=sampling_params) print(f"test_traces outputs is : {outputs}") - timeout = 10 - if not trace_service.evt.wait(timeout): - raise TimeoutError( - f"The fake trace service didn't receive a trace within " - f"the {timeout} seconds timeout" - ) + # Wait for the "llm_request" span to be exported. + # The BatchSpanProcessor batches spans and exports them periodically, + # so we need to wait specifically for the llm_request span to appear. + timeout = 15 + deadline = time.time() + timeout + llm_request_spans = [] + while time.time() < deadline: + all_spans = trace_service.get_all_spans() + llm_request_spans = [s for s in all_spans if s["name"] == "llm_request"] + if llm_request_spans: + break + time.sleep(0.5) - request = trace_service.request - assert len(request.resource_spans) == 1, ( - f"Expected 1 resource span, but got {len(request.resource_spans)}" - ) - assert len(request.resource_spans[0].scope_spans) == 1, ( - f"Expected 1 scope span, " - f"but got {len(request.resource_spans[0].scope_spans)}" - ) - assert len(request.resource_spans[0].scope_spans[0].spans) == 1, ( - f"Expected 1 span, " - f"but got {len(request.resource_spans[0].scope_spans[0].spans)}" + assert len(llm_request_spans) == 1, ( + f"Expected exactly 1 'llm_request' span, but got {len(llm_request_spans)}. " + f"All span names: {[s['name'] for s in all_spans]}" ) - attributes = decode_attributes( - request.resource_spans[0].scope_spans[0].spans[0].attributes - ) + attributes = llm_request_spans[0]["attributes"] # assert attributes.get(SpanAttributes.GEN_AI_RESPONSE_MODEL) == model assert attributes.get(SpanAttributes.GEN_AI_REQUEST_ID) == outputs[0].request_id assert ( diff --git a/vllm/compilation/backends.py b/vllm/compilation/backends.py index 89981fc2996..85833a7a858 100644 --- a/vllm/compilation/backends.py +++ b/vllm/compilation/backends.py @@ -33,6 +33,7 @@ from vllm.config.utils import Range, hash_factors from vllm.logger import init_logger from vllm.logging_utils import lazy from vllm.platforms import current_platform +from vllm.tracing import instrument, instrument_manual from vllm.utils.import_utils import resolve_obj_by_qualname from .compiler_interface import ( @@ -234,6 +235,7 @@ class CompilerManager: ) return compiled_graph + @instrument(span_name="Compile graph") def compile( self, graph: fx.GraphModule, @@ -497,6 +499,7 @@ class PiecewiseCompileInterpreter(torch.fx.Interpreter): # type: ignore[misc] # When True, it annoyingly dumps the torch.fx.Graph on errors. self.extra_traceback = False + @instrument(span_name="Inductor compilation") def run(self, *args: Any) -> Any: # maybe instead just assert inputs are fake? fake_args = [ @@ -922,6 +925,11 @@ class VllmBackend: ) self.compilation_config.compilation_time += dynamo_time + # Record Dynamo time in tracing if available + start_time = int(torch_compile_start_time * 1e9) + attributes = {"dynamo.time_seconds": dynamo_time} + instrument_manual("Dynamo bytecode transform", start_time, None, attributes) + # we control the compilation process, each instance can only be # called once assert not self._called, "VllmBackend can only be called once" diff --git a/vllm/config/observability.py b/vllm/config/observability.py index 3871759125c..7293cf11ca2 100644 --- a/vllm/config/observability.py +++ b/vllm/config/observability.py @@ -122,9 +122,9 @@ class ObservabilityConfig: @classmethod def _validate_otlp_traces_endpoint(cls, value: str | None) -> str | None: if value is not None: - from vllm.tracing import is_otel_available, otel_import_error_traceback + from vllm.tracing import is_tracing_available, otel_import_error_traceback - if not is_otel_available(): + if not is_tracing_available(): raise ValueError( "OpenTelemetry is not available. Unable to configure " "'otlp_traces_endpoint'. Ensure OpenTelemetry packages are " diff --git a/vllm/entrypoints/openai/api_server.py b/vllm/entrypoints/openai/api_server.py index d1da420f6bd..2170765e87f 100644 --- a/vllm/entrypoints/openai/api_server.py +++ b/vllm/entrypoints/openai/api_server.py @@ -50,6 +50,7 @@ from vllm.logger import init_logger from vllm.reasoning import ReasoningParserManager 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 @@ -377,6 +378,7 @@ def validate_api_server_args(args): ) +@instrument(span_name="API server setup") def setup_server(args): """Validate API server args, set up signal handler, create socket ready to serve.""" diff --git a/vllm/model_executor/model_loader/base_loader.py b/vllm/model_executor/model_loader/base_loader.py index 2c55ee68e25..77fbb41f037 100644 --- a/vllm/model_executor/model_loader/base_loader.py +++ b/vllm/model_executor/model_loader/base_loader.py @@ -14,6 +14,7 @@ from vllm.model_executor.model_loader.utils import ( process_weights_after_loading, ) from vllm.platforms import current_platform +from vllm.tracing import instrument from vllm.utils.mem_utils import format_gib from vllm.utils.torch_utils import set_default_torch_dtype @@ -37,6 +38,7 @@ class BaseModelLoader(ABC): inplace weights loading for an already-initialized model""" raise NotImplementedError + @instrument(span_name="Load model") def load_model( self, vllm_config: VllmConfig, model_config: ModelConfig, prefix: str = "" ) -> nn.Module: diff --git a/vllm/model_executor/model_loader/default_loader.py b/vllm/model_executor/model_loader/default_loader.py index c4e961581ef..7064998af86 100644 --- a/vllm/model_executor/model_loader/default_loader.py +++ b/vllm/model_executor/model_loader/default_loader.py @@ -30,6 +30,7 @@ from vllm.model_executor.model_loader.weight_utils import ( pt_weights_iterator, safetensors_weights_iterator, ) +from vllm.tracing import instrument from vllm.transformers_utils.repo_utils import list_filtered_repo_files logger = init_logger(__name__) @@ -274,6 +275,7 @@ class DefaultModelLoader(BaseModelLoader): allow_patterns_overrides=None, ) + @instrument(span_name="Load weights") def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None: if model_config.quantization == "torchao": quant_config = get_quant_config(model_config, self.load_config) diff --git a/vllm/model_executor/model_loader/utils.py b/vllm/model_executor/model_loader/utils.py index 3e46ccc5ca4..51f62c15b30 100644 --- a/vllm/model_executor/model_loader/utils.py +++ b/vllm/model_executor/model_loader/utils.py @@ -23,11 +23,13 @@ from vllm.model_executor.model_loader.reload import ( set_torchao_reload_attrs, ) from vllm.model_executor.models.interfaces import SupportsQuant +from vllm.tracing import instrument from vllm.utils.platform_utils import is_pin_memory_available logger = init_logger(__name__) +@instrument(span_name="Initialize model") def initialize_model( vllm_config: VllmConfig, *, diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py index 02998290b0a..13a60c7b772 100644 --- a/vllm/model_executor/model_loader/weight_utils.py +++ b/vllm/model_executor/model_loader/weight_utils.py @@ -36,6 +36,7 @@ from vllm.model_executor.layers.quantization import ( get_quantization_config, ) from vllm.platforms import current_platform +from vllm.tracing import instrument from vllm.utils.import_utils import PlaceholderModule try: @@ -443,6 +444,7 @@ def download_gguf( return local_files[0] +@instrument(span_name="Download weights - HF") def download_weights_from_hf( model_name_or_path: str, cache_dir: str | None, diff --git a/vllm/model_executor/warmup/deep_gemm_warmup.py b/vllm/model_executor/warmup/deep_gemm_warmup.py index cd4efe1ca77..a445c0aaf1d 100644 --- a/vllm/model_executor/warmup/deep_gemm_warmup.py +++ b/vllm/model_executor/warmup/deep_gemm_warmup.py @@ -19,6 +19,7 @@ from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import ( ) from vllm.model_executor.layers.linear import LinearBase from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod +from vllm.tracing import instrument from vllm.utils.deep_gemm import ( fp8_gemm_nt, get_mk_alignment_for_contiguous_layout, @@ -358,6 +359,7 @@ def _count_warmup_iterations(model: torch.nn.Module, max_tokens: int) -> int: return total +@instrument(span_name="DeepGemm warmup") def deep_gemm_warmup(model: torch.nn.Module, max_tokens: int): total = _count_warmup_iterations(model, max_tokens) if total == 0: diff --git a/vllm/tracing.py b/vllm/tracing.py deleted file mode 100644 index 01bbebf35cf..00000000000 --- a/vllm/tracing.py +++ /dev/null @@ -1,135 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import os -from collections.abc import Mapping - -from vllm.logger import init_logger -from vllm.utils.func_utils import run_once - -TRACE_HEADERS = ["traceparent", "tracestate"] - -logger = init_logger(__name__) - -_is_otel_imported = False -otel_import_error_traceback: str | None = None -try: - from opentelemetry.context.context import Context - from opentelemetry.sdk.environment_variables import ( - OTEL_EXPORTER_OTLP_TRACES_PROTOCOL, - ) - from opentelemetry.sdk.trace import TracerProvider - from opentelemetry.sdk.trace.export import BatchSpanProcessor - from opentelemetry.trace import SpanKind, Tracer, set_tracer_provider - from opentelemetry.trace.propagation.tracecontext import ( - TraceContextTextMapPropagator, - ) - - _is_otel_imported = True -except ImportError: - # Capture and format traceback to provide detailed context for the import - # error. Only the string representation of the error is retained to avoid - # memory leaks. - # See https://github.com/vllm-project/vllm/pull/7266#discussion_r1707395458 - import traceback - - otel_import_error_traceback = traceback.format_exc() - - class Context: # type: ignore - pass - - class BaseSpanAttributes: # type: ignore - pass - - class SpanKind: # type: ignore - pass - - class Tracer: # type: ignore - pass - - -def is_otel_available() -> bool: - return _is_otel_imported - - -def init_tracer( - instrumenting_module_name: str, otlp_traces_endpoint: str -) -> Tracer | None: - if not is_otel_available(): - raise ValueError( - "OpenTelemetry is not available. Unable to initialize " - "a tracer. Ensure OpenTelemetry packages are installed. " - f"Original error:\n{otel_import_error_traceback}" - ) - trace_provider = TracerProvider() - - span_exporter = get_span_exporter(otlp_traces_endpoint) - trace_provider.add_span_processor(BatchSpanProcessor(span_exporter)) - set_tracer_provider(trace_provider) - - tracer = trace_provider.get_tracer(instrumenting_module_name) - return tracer - - -def get_span_exporter(endpoint): - protocol = os.environ.get(OTEL_EXPORTER_OTLP_TRACES_PROTOCOL, "grpc") - if protocol == "grpc": - from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import ( - OTLPSpanExporter, - ) - elif protocol == "http/protobuf": - from opentelemetry.exporter.otlp.proto.http.trace_exporter import ( - OTLPSpanExporter, # type: ignore - ) - else: - raise ValueError(f"Unsupported OTLP protocol '{protocol}' is configured") - - return OTLPSpanExporter(endpoint=endpoint) - - -def extract_trace_context(headers: Mapping[str, str] | None) -> Context | None: - if is_otel_available(): - headers = headers or {} - return TraceContextTextMapPropagator().extract(headers) - else: - return None - - -def extract_trace_headers(headers: Mapping[str, str]) -> Mapping[str, str]: - return {h: headers[h] for h in TRACE_HEADERS if h in headers} - - -class SpanAttributes: - # Attribute names copied from here to avoid version conflicts: - # https://github.com/open-telemetry/semantic-conventions/blob/main/docs/gen-ai/gen-ai-spans.md - GEN_AI_USAGE_COMPLETION_TOKENS = "gen_ai.usage.completion_tokens" - GEN_AI_USAGE_PROMPT_TOKENS = "gen_ai.usage.prompt_tokens" - GEN_AI_REQUEST_MAX_TOKENS = "gen_ai.request.max_tokens" - GEN_AI_REQUEST_TOP_P = "gen_ai.request.top_p" - GEN_AI_REQUEST_TEMPERATURE = "gen_ai.request.temperature" - GEN_AI_RESPONSE_MODEL = "gen_ai.response.model" - # Attribute names added until they are added to the semantic conventions: - GEN_AI_REQUEST_ID = "gen_ai.request.id" - GEN_AI_REQUEST_N = "gen_ai.request.n" - GEN_AI_USAGE_NUM_SEQUENCES = "gen_ai.usage.num_sequences" - GEN_AI_LATENCY_TIME_IN_QUEUE = "gen_ai.latency.time_in_queue" - GEN_AI_LATENCY_TIME_TO_FIRST_TOKEN = "gen_ai.latency.time_to_first_token" - GEN_AI_LATENCY_E2E = "gen_ai.latency.e2e" - GEN_AI_LATENCY_TIME_IN_SCHEDULER = "gen_ai.latency.time_in_scheduler" - # Time taken in the forward pass for this across all workers - GEN_AI_LATENCY_TIME_IN_MODEL_FORWARD = "gen_ai.latency.time_in_model_forward" - # Time taken in the model execute function. This will include model - # forward, block/sync across workers, cpu-gpu sync time and sampling time. - GEN_AI_LATENCY_TIME_IN_MODEL_EXECUTE = "gen_ai.latency.time_in_model_execute" - GEN_AI_LATENCY_TIME_IN_MODEL_PREFILL = "gen_ai.latency.time_in_model_prefill" - GEN_AI_LATENCY_TIME_IN_MODEL_DECODE = "gen_ai.latency.time_in_model_decode" - GEN_AI_LATENCY_TIME_IN_MODEL_INFERENCE = "gen_ai.latency.time_in_model_inference" - - -def contains_trace_headers(headers: Mapping[str, str]) -> bool: - return any(h in headers for h in TRACE_HEADERS) - - -@run_once -def log_tracing_disabled_warning() -> None: - logger.warning("Received a request with trace context but tracing is disabled") diff --git a/vllm/tracing/__init__.py b/vllm/tracing/__init__.py new file mode 100644 index 00000000000..4f025970ecb --- /dev/null +++ b/vllm/tracing/__init__.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import functools +from collections.abc import Callable +from typing import Any, TypeAlias + +# Import the implementation details +from .otel import ( + SpanKind, + extract_trace_context, + init_otel_tracer, + init_otel_worker_tracer, + instrument_otel, + is_otel_available, + manual_instrument_otel, + otel_import_error_traceback, +) +from .utils import ( + SpanAttributes, + contains_trace_headers, + extract_trace_headers, + log_tracing_disabled_warning, +) + +__all__ = [ + "instrument", + "instrument_manual", + "init_tracer", + "maybe_init_worker_tracer", + "is_tracing_available", + "SpanAttributes", + "SpanKind", + "extract_trace_context", + "extract_trace_headers", + "log_tracing_disabled_warning", + "contains_trace_headers", + "otel_import_error_traceback", +] + +BackendAvailableFunc: TypeAlias = Callable[[], bool] +InstrumentFunc: TypeAlias = Callable[..., Any] +InstrumentManualFunc: TypeAlias = Callable[..., Any] +InitTracerFunc: TypeAlias = Callable[..., Any] +InitWorkerTracerFunc: TypeAlias = Callable[..., Any] +_REGISTERED_TRACING_BACKENDS: dict[ + str, + tuple[ + BackendAvailableFunc, + InitTracerFunc, + InitWorkerTracerFunc, + InstrumentFunc, + InstrumentManualFunc, + ], +] = { + "otel": ( + is_otel_available, + init_otel_tracer, + init_otel_worker_tracer, + instrument_otel, + manual_instrument_otel, + ), +} + + +def init_tracer( + instrumenting_module_name: str, + otlp_traces_endpoint: str, + extra_attributes: dict[str, str] | None = None, +): + is_available, init_tracer_fn, _, _, _ = _REGISTERED_TRACING_BACKENDS["otel"] + if is_available(): + return init_tracer_fn( + instrumenting_module_name, otlp_traces_endpoint, extra_attributes + ) + + +def maybe_init_worker_tracer( + instrumenting_module_name: str, + process_kind: str, + process_name: str, +): + is_available, _, init_worker_tracer_fn, _, _ = _REGISTERED_TRACING_BACKENDS["otel"] + if is_available(): + return init_worker_tracer_fn( + instrumenting_module_name, process_kind, process_name + ) + + +def instrument( + obj: Callable | None = None, + *, + span_name: str = "", + attributes: dict[str, str] | None = None, + record_exception: bool = True, +): + """ + Generic decorator to instrument functions. + """ + if obj is None: + return functools.partial( + instrument, + span_name=span_name, + attributes=attributes, + record_exception=record_exception, + ) + + # Dispatch to OTel (and potentially others later) + is_available, _, _, otel_instrument, _ = _REGISTERED_TRACING_BACKENDS["otel"] + if is_available(): + return otel_instrument( + func=obj, + span_name=span_name, + attributes=attributes, + record_exception=record_exception, + ) + else: + return obj + + +def instrument_manual( + span_name: str, + start_time: int, + end_time: int | None = None, + attributes: dict[str, Any] | None = None, + context: Any = None, + kind: Any = None, +): + """Manually create a span with explicit timestamps. + + Args: + span_name: Name of the span to create. + start_time: Start time in nanoseconds since epoch. + end_time: Optional end time in nanoseconds. If None, ends immediately. + attributes: Optional dict of span attributes. + context: Optional trace context (e.g., from extract_trace_context). + kind: Optional SpanKind (e.g., SpanKind.SERVER). + """ + is_available, _, _, _, manual_instrument_fn = _REGISTERED_TRACING_BACKENDS["otel"] + if is_available(): + return manual_instrument_fn( + span_name, start_time, end_time, attributes, context, kind + ) + else: + return None + + +def is_tracing_available() -> bool: + """ + Returns True if any tracing backend (OTel, Profiler, etc.) is available. + Use this to guard expensive tracing logic in the main code. + """ + check_available = [ + is_available + for is_available, _, _, _, _ in _REGISTERED_TRACING_BACKENDS.values() + ] + return any(check_available) diff --git a/vllm/tracing/otel.py b/vllm/tracing/otel.py new file mode 100644 index 00000000000..ac06ae97255 --- /dev/null +++ b/vllm/tracing/otel.py @@ -0,0 +1,265 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import atexit +import functools +import inspect +import os +import traceback +from collections.abc import Mapping +from contextlib import contextmanager +from typing import Any + +from vllm.logger import init_logger +from vllm.tracing.utils import TRACE_HEADERS, LoadingSpanAttributes + +logger = init_logger(__name__) + +try: + from opentelemetry import trace + from opentelemetry.context.context import Context + from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import ( + OTLPSpanExporter as OTLPGrpcExporter, + ) + from opentelemetry.exporter.otlp.proto.http.trace_exporter import ( + OTLPSpanExporter as OTLPHttpExporter, + ) + from opentelemetry.propagate import inject + from opentelemetry.sdk.environment_variables import ( + OTEL_EXPORTER_OTLP_TRACES_PROTOCOL, + ) + from opentelemetry.sdk.resources import Resource + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import BatchSpanProcessor + from opentelemetry.trace import ( + SpanKind, # noqa: F401 + Tracer, + set_tracer_provider, + ) + from opentelemetry.trace.propagation.tracecontext import ( + TraceContextTextMapPropagator, + ) + + _IS_OTEL_AVAILABLE = True + otel_import_error_traceback = None +except ImportError: + _IS_OTEL_AVAILABLE = False + otel_import_error_traceback = traceback.format_exc() + trace = None # type: ignore + Context = Any # type: ignore + Tracer = Any # type: ignore + inject = None # type: ignore + Resource = None # type: ignore + SpanKind = Any # type: ignore + + +def is_otel_available() -> bool: + return _IS_OTEL_AVAILABLE + + +def init_otel_tracer( + instrumenting_module_name: str, + otlp_traces_endpoint: str, + extra_attributes: dict[str, str] | None = None, +) -> Tracer: + """Initializes the OpenTelemetry tracer provider.""" + if not _IS_OTEL_AVAILABLE: + raise ValueError( + "OpenTelemetry is not available. Unable to initialize " + "a tracer. Ensure OpenTelemetry packages are installed. " + f"Original error:\n{otel_import_error_traceback}" + ) + + # Store the endpoint in environment so child processes can inherit it + os.environ["OTEL_EXPORTER_OTLP_TRACES_ENDPOINT"] = otlp_traces_endpoint + + resource_attrs = {} + resource_attrs["vllm.instrumenting_module_name"] = instrumenting_module_name + resource_attrs["vllm.process_id"] = str(os.getpid()) + if extra_attributes: + resource_attrs.update(extra_attributes) + resource = Resource.create(resource_attrs) + + trace_provider = TracerProvider(resource=resource) + span_exporter = get_span_exporter(otlp_traces_endpoint) + trace_provider.add_span_processor(BatchSpanProcessor(span_exporter)) + set_tracer_provider(trace_provider) + + atexit.register(trace_provider.shutdown) + + tracer = trace_provider.get_tracer(instrumenting_module_name) + return tracer + + +def get_span_exporter(endpoint): + protocol = os.environ.get(OTEL_EXPORTER_OTLP_TRACES_PROTOCOL, "grpc") + if protocol == "grpc": + exporter = OTLPGrpcExporter(endpoint=endpoint, insecure=True) + elif protocol == "http/protobuf": + exporter = OTLPHttpExporter(endpoint=endpoint) + else: + raise ValueError(f"Unsupported OTLP protocol '{protocol}' is configured") + return exporter + + +def init_otel_worker_tracer( + instrumenting_module_name: str, + process_kind: str, + process_name: str, +) -> Tracer: + """ + Backend-specific initialization for OpenTelemetry in a worker process. + """ + # Initialize the tracer if an OTLP endpoint is configured. + # The endpoint is propagated via environment variable from the main process. + otlp_endpoint = os.environ.get("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT") + if not otlp_endpoint: + return None + + extra_attrs = { + "vllm.process_kind": process_kind, + "vllm.process_name": process_name, + } + + return init_otel_tracer(instrumenting_module_name, otlp_endpoint, extra_attrs) + + +def extract_trace_context(headers: Mapping[str, str] | None) -> Context | None: + """Extracts context from HTTP headers.""" + if _IS_OTEL_AVAILABLE and headers: + return TraceContextTextMapPropagator().extract(headers) + return None + + +def instrument_otel(func, span_name, attributes, record_exception): + """Internal wrapper logic for sync and async functions.""" + + # Pre-calculate static code attributes once (these don't change) + code_attrs = { + LoadingSpanAttributes.CODE_FUNCTION: func.__qualname__, + LoadingSpanAttributes.CODE_NAMESPACE: func.__module__, + LoadingSpanAttributes.CODE_FILEPATH: func.__code__.co_filename, + LoadingSpanAttributes.CODE_LINENO: str(func.__code__.co_firstlineno), + } + if attributes: + code_attrs.update(attributes) + + final_span_name = span_name or func.__qualname__ + module_name = func.__module__ + + @functools.wraps(func) + async def async_wrapper(*args, **kwargs): + tracer = trace.get_tracer(module_name) + ctx = _get_smart_context() + with ( + tracer.start_as_current_span( + final_span_name, + context=ctx, + attributes=code_attrs, + record_exception=record_exception, + ), + propagate_trace_to_env(), + ): + return await func(*args, **kwargs) + + @functools.wraps(func) + def sync_wrapper(*args, **kwargs): + tracer = trace.get_tracer(module_name) + ctx = _get_smart_context() + with ( + tracer.start_as_current_span( + final_span_name, + context=ctx, + attributes=code_attrs, + record_exception=record_exception, + ), + propagate_trace_to_env(), + ): + return func(*args, **kwargs) + + return async_wrapper if inspect.iscoroutinefunction(func) else sync_wrapper + + +def manual_instrument_otel( + span_name: str, + start_time: int, + end_time: int | None = None, + attributes: dict[str, Any] | None = None, + context: Context | None = None, + kind: Any = None, # SpanKind, but typed as Any for when OTEL unavailable +): + """Manually create and end a span with explicit timestamps.""" + if not _IS_OTEL_AVAILABLE: + return + + tracer = trace.get_tracer(__name__) + # Use provided context, or fall back to smart context detection + ctx = context if context is not None else _get_smart_context() + + span_kwargs: dict[str, Any] = { + "name": span_name, + "context": ctx, + "start_time": start_time, + } + if kind is not None: + span_kwargs["kind"] = kind + + span = tracer.start_span(**span_kwargs) + if attributes: + span.set_attributes(attributes) + if end_time is not None: + span.end(end_time=end_time) + else: + span.end() + + +def _get_smart_context() -> Context | None: + """ + Determines the parent context. + 1. If a Span is already active in this process, use it. + 2. If not, extract from os.environ, handling the case-sensitivity mismatch. + """ + current_span = trace.get_current_span() + if current_span.get_span_context().is_valid: + return None + + carrier = {} + + if tp := os.environ.get("traceparent", os.environ.get("TRACEPARENT")): # noqa: SIM112 + carrier["traceparent"] = tp + + if ts := os.environ.get("tracestate", os.environ.get("TRACESTATE")): # noqa: SIM112 + carrier["tracestate"] = ts + + if not carrier: + carrier = dict(os.environ) + + return TraceContextTextMapPropagator().extract(carrier) + + +@contextmanager +def propagate_trace_to_env(): + """ + Temporarily injects the current OTel context into os.environ. + This ensures that any subprocesses (like vLLM workers) spawned + within this context inherit the correct traceparent. + """ + if not _IS_OTEL_AVAILABLE: + yield + return + + # Capture original state of relevant keys + original_state = {k: os.environ.get(k) for k in TRACE_HEADERS} + + try: + # inject() writes 'traceparent' and 'tracestate' to os.environ + inject(os.environ) + yield + + finally: + # Restore original environment + for key, original_value in original_state.items(): + if original_value is None: + os.environ.pop(key, None) + else: + os.environ[key] = original_value diff --git a/vllm/tracing/utils.py b/vllm/tracing/utils.py new file mode 100644 index 00000000000..0e11f850b44 --- /dev/null +++ b/vllm/tracing/utils.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Mapping + +from vllm.logger import init_logger +from vllm.utils.func_utils import run_once + +logger = init_logger(__name__) + +# Standard W3C headers used for context propagation +TRACE_HEADERS = ["traceparent", "tracestate"] + + +class SpanAttributes: + """ + Standard attributes for spans. + + These are largely based on OpenTelemetry Semantic Conventions but are defined + here as constants so they can be used by any backend or logger. + """ + + # Attribute names copied from OTel semantic conventions to avoid version conflicts + GEN_AI_USAGE_COMPLETION_TOKENS = "gen_ai.usage.completion_tokens" + GEN_AI_USAGE_PROMPT_TOKENS = "gen_ai.usage.prompt_tokens" + GEN_AI_REQUEST_MAX_TOKENS = "gen_ai.request.max_tokens" + GEN_AI_REQUEST_TOP_P = "gen_ai.request.top_p" + GEN_AI_REQUEST_TEMPERATURE = "gen_ai.request.temperature" + GEN_AI_RESPONSE_MODEL = "gen_ai.response.model" + + # Custom attributes added until they are standardized + GEN_AI_REQUEST_ID = "gen_ai.request.id" + GEN_AI_REQUEST_N = "gen_ai.request.n" + GEN_AI_USAGE_NUM_SEQUENCES = "gen_ai.usage.num_sequences" + GEN_AI_LATENCY_TIME_IN_QUEUE = "gen_ai.latency.time_in_queue" + GEN_AI_LATENCY_TIME_TO_FIRST_TOKEN = "gen_ai.latency.time_to_first_token" + GEN_AI_LATENCY_E2E = "gen_ai.latency.e2e" + GEN_AI_LATENCY_TIME_IN_SCHEDULER = "gen_ai.latency.time_in_scheduler" + + # Latency breakdowns + GEN_AI_LATENCY_TIME_IN_MODEL_FORWARD = "gen_ai.latency.time_in_model_forward" + GEN_AI_LATENCY_TIME_IN_MODEL_EXECUTE = "gen_ai.latency.time_in_model_execute" + GEN_AI_LATENCY_TIME_IN_MODEL_PREFILL = "gen_ai.latency.time_in_model_prefill" + GEN_AI_LATENCY_TIME_IN_MODEL_DECODE = "gen_ai.latency.time_in_model_decode" + GEN_AI_LATENCY_TIME_IN_MODEL_INFERENCE = "gen_ai.latency.time_in_model_inference" + + +class LoadingSpanAttributes: + """Custom attributes for code-level tracing (file, line number).""" + + CODE_NAMESPACE = "code.namespace" + CODE_FUNCTION = "code.function" + CODE_FILEPATH = "code.filepath" + CODE_LINENO = "code.lineno" + + +def contains_trace_headers(headers: Mapping[str, str]) -> bool: + """Check if the provided headers dictionary contains trace context.""" + return any(h in headers for h in TRACE_HEADERS) + + +def extract_trace_headers(headers: Mapping[str, str]) -> Mapping[str, str]: + """ + Extract only trace-related headers from a larger header dictionary. + Useful for logging or passing context to a non-OTel client. + """ + return {h: headers[h] for h in TRACE_HEADERS if h in headers} + + +@run_once +def log_tracing_disabled_warning() -> None: + logger.warning("Received a request with trace context but tracing is disabled") diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 43d63bcff25..28c95777c1c 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -110,6 +110,10 @@ class AsyncLLM(EngineClient): self.model_config = vllm_config.model_config self.vllm_config = vllm_config self.observability_config = vllm_config.observability_config + tracing_endpoint = self.observability_config.otlp_traces_endpoint + if tracing_endpoint is not None: + init_tracer("vllm.llm_engine", tracing_endpoint) + self.log_requests = log_requests custom_stat_loggers = list(stat_loggers or []) @@ -136,10 +140,8 @@ class AsyncLLM(EngineClient): log_stats=self.log_stats, stream_interval=self.vllm_config.scheduler_config.stream_interval, ) - endpoint = self.observability_config.otlp_traces_endpoint - if endpoint is not None: - tracer = init_tracer("vllm.llm_engine", endpoint) - self.output_processor.tracer = tracer + if tracing_endpoint is not None: + self.output_processor.tracing_enabled = True # EngineCore (starts the engine in background process). self.engine_core = EngineCoreClient.make_async_mp_client( diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 216d610b4b8..ebc1febdd56 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -24,6 +24,7 @@ from vllm.logging_utils.dump_input import dump_engine_exception from vllm.lora.request import LoRARequest from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.tasks import POOLING_TASKS, SupportedTask +from vllm.tracing import instrument, maybe_init_worker_tracer from vllm.transformers_utils.config import maybe_register_config_serialize_by_value from vllm.utils.gc_utils import ( freeze_gc_heap, @@ -217,6 +218,7 @@ class EngineCore: # environment variable overrides after this point) enable_envs_cache() + @instrument(span_name="Prepare model") def _initialize_kv_caches( self, vllm_config: VllmConfig ) -> tuple[int, int, KVCacheConfig]: @@ -658,6 +660,7 @@ class EngineCoreProc(EngineCore): ENGINE_CORE_DEAD = b"ENGINE_CORE_DEAD" + @instrument(span_name="EngineCoreProc init") def __init__( self, vllm_config: VllmConfig, @@ -926,8 +929,18 @@ class EngineCoreProc(EngineCore): data_parallel = parallel_config.data_parallel_size > 1 or dp_rank > 0 if data_parallel: parallel_config.data_parallel_rank_local = local_dp_rank + maybe_init_worker_tracer( + instrumenting_module_name="vllm.engine_core", + process_kind="engine_core", + process_name=f"EngineCore_DP{dp_rank}", + ) set_process_title("EngineCore", f"DP{dp_rank}") else: + maybe_init_worker_tracer( + instrumenting_module_name="vllm.engine_core", + process_kind="engine_core", + process_name="EngineCore", + ) set_process_title("EngineCore") decorate_logs() @@ -956,6 +969,7 @@ class EngineCoreProc(EngineCore): parallel_config.data_parallel_rank = 0 engine_core = EngineCoreProc(*args, engine_index=dp_rank, **kwargs) + assert engine_core is not None engine_core.run_busy_loop() except SystemExit: @@ -1485,6 +1499,13 @@ class EngineCoreActorMixin: dp_rank: int = 0, local_dp_rank: int = 0, ): + # Initialize tracer for distributed tracing if configured. + maybe_init_worker_tracer( + instrumenting_module_name="vllm.engine_core", + process_kind="engine_core", + process_name=f"DPEngineCoreActor_DP{dp_rank}", + ) + self.addresses = addresses vllm_config.parallel_config.data_parallel_index = dp_rank vllm_config.parallel_config.data_parallel_rank_local = local_dp_rank diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index 30808619803..f303db2dffe 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -24,6 +24,7 @@ from vllm.envs import VLLM_ENGINE_READY_TIMEOUT_S from vllm.logger import init_logger from vllm.lora.request import LoRARequest from vllm.tasks import SupportedTask +from vllm.tracing import instrument from vllm.utils.async_utils import in_loop from vllm.utils.network_utils import ( close_sockets, @@ -96,6 +97,7 @@ class EngineCoreClient(ABC): return InprocClient(vllm_config, executor_class, log_stats) @staticmethod + @instrument(span_name="Overall Loading") def make_async_mp_client( vllm_config: VllmConfig, executor_class: type[Executor], @@ -650,6 +652,7 @@ def _process_utility_output( class SyncMPClient(MPClient): """Synchronous client for multi-proc EngineCore.""" + @instrument(span_name="SyncMPClient init") def __init__( self, vllm_config: VllmConfig, executor_class: type[Executor], log_stats: bool ): @@ -819,6 +822,7 @@ class SyncMPClient(MPClient): class AsyncMPClient(MPClient): """Asyncio-compatible client for multi-proc EngineCore.""" + @instrument(span_name="AsyncMPClient init") def __init__( self, vllm_config: VllmConfig, diff --git a/vllm/v1/engine/llm_engine.py b/vllm/v1/engine/llm_engine.py index 9cae71a4348..4f44e7101b2 100644 --- a/vllm/v1/engine/llm_engine.py +++ b/vllm/v1/engine/llm_engine.py @@ -100,8 +100,8 @@ class LLMEngine: ) endpoint = self.observability_config.otlp_traces_endpoint if endpoint is not None: - tracer = init_tracer("vllm.llm_engine", endpoint) - self.output_processor.tracer = tracer + init_tracer("vllm.llm_engine", endpoint) + self.output_processor.tracing_enabled = True # EngineCore (gets EngineCoreRequests and gives EngineCoreOutputs) self.engine_core = EngineCoreClient.make_client( diff --git a/vllm/v1/engine/output_processor.py b/vllm/v1/engine/output_processor.py index c497468bbd3..00a5355e022 100644 --- a/vllm/v1/engine/output_processor.py +++ b/vllm/v1/engine/output_processor.py @@ -20,7 +20,12 @@ from vllm.outputs import ( ) from vllm.sampling_params import RequestOutputKind from vllm.tokenizers import TokenizerLike -from vllm.tracing import SpanAttributes, SpanKind, Tracer, extract_trace_context +from vllm.tracing import ( + SpanAttributes, + SpanKind, + extract_trace_context, + instrument_manual, +) from vllm.utils import length_from_prompt_token_ids_or_embeds from vllm.v1.engine import EngineCoreOutput, EngineCoreRequest, FinishReason from vllm.v1.engine.detokenizer import IncrementalDetokenizer @@ -422,7 +427,7 @@ class OutputProcessor: self.parent_requests: dict[str, ParentRequest] = {} self.external_req_ids: defaultdict[str, list[str]] = defaultdict(list) self.lora_states = LoRARequestStates(log_stats) - self.tracer: Tracer | None = None + self.tracing_enabled: bool = False self._requests_drained = asyncio.Event() self._requests_drained.set() @@ -678,7 +683,7 @@ class OutputProcessor: self._update_stats_from_finished( req_state, finish_reason, iteration_stats ) - if self.tracer: + if self.tracing_enabled: self.do_tracing(engine_core_output, req_state, iteration_stats) return OutputProcessorOutput( @@ -714,62 +719,59 @@ class OutputProcessor: ) -> None: assert req_state.stats is not None assert iteration_stats is not None - assert self.tracer is not None - arrival_time_nano_seconds = int(req_state.stats.arrival_time * 1e9) + metrics = req_state.stats + arrival_time_ns = int(metrics.arrival_time * 1e9) trace_context = extract_trace_context(engine_core_output.trace_headers) prompt_length = length_from_prompt_token_ids_or_embeds( req_state.prompt_token_ids, req_state.prompt_embeds ) - with self.tracer.start_as_current_span( - "llm_request", - kind=SpanKind.SERVER, - context=trace_context, - start_time=arrival_time_nano_seconds, - ) as span: - metrics = req_state.stats - e2e_time = iteration_stats.iteration_timestamp - metrics.arrival_time - queued_time = metrics.scheduled_ts - metrics.queued_ts - prefill_time = metrics.first_token_ts - metrics.scheduled_ts - decode_time = metrics.last_token_ts - metrics.first_token_ts - inference_time = metrics.last_token_ts - metrics.scheduled_ts - span.set_attribute( - SpanAttributes.GEN_AI_LATENCY_TIME_TO_FIRST_TOKEN, - metrics.first_token_latency, - ) - span.set_attribute(SpanAttributes.GEN_AI_LATENCY_E2E, e2e_time) - span.set_attribute(SpanAttributes.GEN_AI_LATENCY_TIME_IN_QUEUE, queued_time) - span.set_attribute(SpanAttributes.GEN_AI_USAGE_PROMPT_TOKENS, prompt_length) - span.set_attribute( - SpanAttributes.GEN_AI_USAGE_COMPLETION_TOKENS, - metrics.num_generation_tokens, - ) - span.set_attribute( - SpanAttributes.GEN_AI_LATENCY_TIME_IN_MODEL_PREFILL, prefill_time - ) - span.set_attribute( - SpanAttributes.GEN_AI_LATENCY_TIME_IN_MODEL_DECODE, decode_time - ) - span.set_attribute( - SpanAttributes.GEN_AI_LATENCY_TIME_IN_MODEL_INFERENCE, inference_time - ) - # meta - span.set_attribute( - SpanAttributes.GEN_AI_REQUEST_ID, req_state.external_req_id + # Calculate timing metrics + e2e_time = iteration_stats.iteration_timestamp - metrics.arrival_time + queued_time = metrics.scheduled_ts - metrics.queued_ts + prefill_time = metrics.first_token_ts - metrics.scheduled_ts + decode_time = metrics.last_token_ts - metrics.first_token_ts + inference_time = metrics.last_token_ts - metrics.scheduled_ts + + # Build attributes dict + attributes: dict[str, Any] = { + SpanAttributes.GEN_AI_LATENCY_TIME_TO_FIRST_TOKEN: ( + metrics.first_token_latency + ), + SpanAttributes.GEN_AI_LATENCY_E2E: e2e_time, + SpanAttributes.GEN_AI_LATENCY_TIME_IN_QUEUE: queued_time, + SpanAttributes.GEN_AI_USAGE_PROMPT_TOKENS: prompt_length, + SpanAttributes.GEN_AI_USAGE_COMPLETION_TOKENS: ( + metrics.num_generation_tokens + ), + SpanAttributes.GEN_AI_LATENCY_TIME_IN_MODEL_PREFILL: prefill_time, + SpanAttributes.GEN_AI_LATENCY_TIME_IN_MODEL_DECODE: decode_time, + SpanAttributes.GEN_AI_LATENCY_TIME_IN_MODEL_INFERENCE: inference_time, + SpanAttributes.GEN_AI_REQUEST_ID: req_state.external_req_id, + } + + # Add optional request parameters + if req_state.top_p: + attributes[SpanAttributes.GEN_AI_REQUEST_TOP_P] = req_state.top_p + if req_state.max_tokens_param: + attributes[SpanAttributes.GEN_AI_REQUEST_MAX_TOKENS] = ( + req_state.max_tokens_param ) - if req_state.top_p: - span.set_attribute(SpanAttributes.GEN_AI_REQUEST_TOP_P, req_state.top_p) - if req_state.max_tokens_param: - span.set_attribute( - SpanAttributes.GEN_AI_REQUEST_MAX_TOKENS, req_state.max_tokens_param - ) - if req_state.temperature: - span.set_attribute( - SpanAttributes.GEN_AI_REQUEST_TEMPERATURE, req_state.temperature - ) - if req_state.n: - span.set_attribute(SpanAttributes.GEN_AI_REQUEST_N, req_state.n) + if req_state.temperature: + attributes[SpanAttributes.GEN_AI_REQUEST_TEMPERATURE] = ( + req_state.temperature + ) + if req_state.n: + attributes[SpanAttributes.GEN_AI_REQUEST_N] = req_state.n + + instrument_manual( + span_name="llm_request", + start_time=arrival_time_ns, + attributes=attributes, + context=trace_context, + kind=SpanKind.SERVER, + ) def _update_stats_from_output( self, diff --git a/vllm/v1/executor/abstract.py b/vllm/v1/executor/abstract.py index 0fef6c1d149..32fa87e9d3a 100644 --- a/vllm/v1/executor/abstract.py +++ b/vllm/v1/executor/abstract.py @@ -15,6 +15,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import ( from vllm.logger import init_logger from vllm.lora.request import LoRARequest from vllm.tasks import SupportedTask +from vllm.tracing import instrument from vllm.utils.import_utils import resolve_obj_by_qualname from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.engine import ReconfigureDistributedRequest @@ -84,6 +85,7 @@ class Executor(ABC): ) return executor_class + @instrument(span_name="Executor init") def __init__( self, vllm_config: VllmConfig, diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py index df92259f50e..b63cbd6586f 100644 --- a/vllm/v1/executor/multiproc_executor.py +++ b/vllm/v1/executor/multiproc_executor.py @@ -41,6 +41,7 @@ from vllm.distributed.parallel_state import ( ) from vllm.envs import enable_envs_cache from vllm.logger import init_logger +from vllm.tracing import instrument, maybe_init_worker_tracer from vllm.utils.network_utils import ( get_distributed_init_method, get_loopback_ip, @@ -527,6 +528,7 @@ class WorkerProc: ) ) + @instrument(span_name="Worker init") def __init__( self, vllm_config: VllmConfig, @@ -740,6 +742,15 @@ class WorkerProc: try: reader.close() + + # Initialize tracer + rank = kwargs.get("rank", 0) + maybe_init_worker_tracer( + instrumenting_module_name="vllm.worker", + process_kind="worker", + process_name=f"Worker_{rank}", + ) + worker = WorkerProc(*args, **kwargs) assert worker.worker_response_mq is not None diff --git a/vllm/v1/worker/cpu_model_runner.py b/vllm/v1/worker/cpu_model_runner.py index 6bfbc32d598..8ee758353b1 100644 --- a/vllm/v1/worker/cpu_model_runner.py +++ b/vllm/v1/worker/cpu_model_runner.py @@ -9,6 +9,7 @@ import torch.nn as nn from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.model_executor.model_loader import get_model +from vllm.tracing import instrument from vllm.v1.utils import CpuGpuBuffer from vllm.v1.worker.gpu_model_runner import GPUModelRunner @@ -51,6 +52,7 @@ class CPUModelRunner(GPUModelRunner): if isinstance(v, CpuGpuBuffer): v.gpu = v.cpu + @instrument(span_name="Loading (CPU)") def load_model(self, eep_scale_up: bool = False) -> None: logger.info("Starting to load model %s...", self.model_config.model) self.model = get_model(vllm_config=self.vllm_config) @@ -61,6 +63,7 @@ class CPUModelRunner(GPUModelRunner): def get_model(self) -> nn.Module: return self.model + @instrument(span_name="Warmup (CPU)") def warming_up_model(self) -> None: logger.info("Warming up model for the compilation...") # Only generate graph for the generic shape diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 6b04774a8b6..ec36e159147 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -93,6 +93,7 @@ from vllm.pooling_params import PoolingParams from vllm.sampling_params import SamplingType from vllm.sequence import IntermediateTensors from vllm.tasks import GenerationTask, PoolingTask, SupportedTask +from vllm.tracing import instrument from vllm.utils import length_from_prompt_token_ids_or_embeds from vllm.utils.jsontree import json_map_leaves from vllm.utils.math_utils import cdiv, round_up @@ -4111,6 +4112,7 @@ class GPUModelRunner( new_config = update_config(config, config_overrides) setattr(self, config_name, new_config) + @instrument(span_name="Loading (GPU)") def load_model(self, eep_scale_up: bool = False) -> None: """ Args: @@ -5165,6 +5167,7 @@ class GPUModelRunner( self.encoder_cache.clear() gc.collect() + @instrument(span_name="Capture model") def capture_model(self) -> int: if self.compilation_config.cudagraph_mode == CUDAGraphMode.NONE: logger.warning( diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 09880f79bf1..f1eb4b2bc21 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -42,6 +42,7 @@ from vllm.platforms import current_platform from vllm.profiler.wrapper import CudaProfilerWrapper, TorchProfilerWrapper from vllm.sequence import IntermediateTensors from vllm.tasks import SupportedTask +from vllm.tracing import instrument from vllm.utils.mem_utils import MemorySnapshot, format_gib, memory_profiling from vllm.utils.torch_utils import set_random_seed from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput @@ -186,6 +187,7 @@ class Worker(WorkerBase): self.cache_config.num_gpu_blocks = num_gpu_blocks self.cache_config.num_cpu_blocks = num_cpu_blocks + @instrument(span_name="Init device") def init_device(self): if self.device_config.device_type == "cuda": # This env var set by Ray causes exceptions with graph building. @@ -407,6 +409,7 @@ class Worker(WorkerBase): self.model_runner.update_max_model_len(max_model_len) logger.debug("Updated max_model_len to %d", max_model_len) + @instrument(span_name="Allocate KV cache") def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None: """Allocate GPU KV cache with the specified kv_cache_config.""" @@ -426,6 +429,7 @@ class Worker(WorkerBase): else: self.model_runner.initialize_kv_cache(kv_cache_config) + @instrument(span_name="Warmup (GPU)") def compile_or_warm_up_model(self) -> None: warmup_sizes = [] diff --git a/vllm/v1/worker/worker_base.py b/vllm/v1/worker/worker_base.py index eed371e988b..b4454589d7e 100644 --- a/vllm/v1/worker/worker_base.py +++ b/vllm/v1/worker/worker_base.py @@ -11,6 +11,7 @@ from vllm.config import VllmConfig, set_current_vllm_config from vllm.logger import init_logger from vllm.lora.request import LoRARequest from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.tracing import instrument from vllm.utils.import_utils import resolve_obj_by_qualname from vllm.utils.system_utils import update_environment_variables from vllm.v1.kv_cache_interface import KVCacheSpec @@ -222,6 +223,7 @@ class WorkerWrapperBase: envs = envs_list[self.rpc_rank] update_environment_variables(envs) + @instrument(span_name="Worker init") def init_worker(self, all_kwargs: list[dict[str, Any]]) -> None: """ Here we inject some common logic before initializing the worker. From 79028d438859162841d35bbf2a91eaa8236733fa Mon Sep 17 00:00:00 2001 From: Xin Yang <105740670+xyang16@users.noreply.github.com> Date: Thu, 5 Feb 2026 17:34:00 -0800 Subject: [PATCH 117/810] [Perf] Disable clean_logits in deepgemm fp8_mqa_logits kernel (#33568) --- .../attention/test_deepgemm_attention.py | 20 ++++--- tests/kernels/test_top_k_per_row.py | 56 +++++++++++++------ .../layers/sparse_attn_indexer.py | 2 + vllm/utils/deep_gemm.py | 10 +++- 4 files changed, 61 insertions(+), 27 deletions(-) diff --git a/tests/kernels/attention/test_deepgemm_attention.py b/tests/kernels/attention/test_deepgemm_attention.py index e2ae3b833b2..2dc522598e4 100644 --- a/tests/kernels/attention/test_deepgemm_attention.py +++ b/tests/kernels/attention/test_deepgemm_attention.py @@ -95,7 +95,8 @@ def _ref_fp8_mqa_logits( @pytest.mark.skipif( not current_platform.has_device_capability(90), reason="SM90 and SM100 only" ) -def test_deepgemm_fp8_mqa_logits(): +@pytest.mark.parametrize("clean_logits", [True, False]) +def test_deepgemm_fp8_mqa_logits(clean_logits: bool): torch.manual_seed(0) random.seed(0) num_heads, head_dim = 32, 128 @@ -126,7 +127,9 @@ def test_deepgemm_fp8_mqa_logits(): q_fp8 = q.to(torch.float8_e4m3fn) kv_fp8 = per_custom_dims_cast_to_fp8(kv, (0,), False) - logits = fp8_mqa_logits(q_fp8, kv_fp8, weights, ks, ke) + logits = fp8_mqa_logits( + q_fp8, kv_fp8, weights, ks, ke, clean_logits=clean_logits + ) ref_logits = _ref_fp8_mqa_logits( q=q, @@ -135,13 +138,14 @@ def test_deepgemm_fp8_mqa_logits(): cu_seqlen_ks=ks, cu_seqlen_ke=ke, ) - ref_neginf_mask = ref_logits == float("-inf") - neginf_mask = logits == float("-inf") - assert torch.equal(neginf_mask, ref_neginf_mask) + + if clean_logits: + neginf_mask = logits == float("-inf") + assert torch.equal(neginf_mask, ref_neginf_mask) ref_logits = ref_logits.masked_fill(ref_neginf_mask, 0) - logits = logits.masked_fill(neginf_mask, 0) + logits = logits.masked_fill(ref_neginf_mask, 0) diff = calc_diff(logits, ref_logits) assert diff < 1e-3, f"{diff=}" @@ -201,7 +205,8 @@ def _ref_fp8_paged_mqa_logits( @pytest.mark.skipif( not current_platform.has_device_capability(90), reason="SM90 and SM100 only" ) -def test_deepgemm_fp8_paged_mqa_logits(): +@pytest.mark.parametrize("clean_logits", [True, False]) +def test_deepgemm_fp8_paged_mqa_logits(clean_logits: bool): torch.manual_seed(0) random.seed(0) @@ -264,6 +269,7 @@ def test_deepgemm_fp8_paged_mqa_logits(): block_tables, schedule_metadata, max_model_len, + clean_logits=clean_logits, ) ref_logits = _ref_fp8_paged_mqa_logits( diff --git a/tests/kernels/test_top_k_per_row.py b/tests/kernels/test_top_k_per_row.py index 3bf69389753..2d9dd2a0461 100644 --- a/tests/kernels/test_top_k_per_row.py +++ b/tests/kernels/test_top_k_per_row.py @@ -6,6 +6,7 @@ import pytest import torch from vllm.platforms import current_platform +from vllm.utils.torch_utils import set_random_seed # Test parameters NUM_ROWS = [1, 32, 2050] @@ -20,6 +21,7 @@ def create_random_logits( row_ends: torch.Tensor, dtype: torch.dtype, seed: int, + clean_logits: bool, data_generation: str, ) -> torch.Tensor: """Create random logits tensor for testing.""" @@ -48,8 +50,9 @@ def create_random_logits( ) logits = logits_bits.view(dtype) - for i, end in enumerate(row_ends): - logits[i, end:] = float("-inf") + if clean_logits: + for i, end in enumerate(row_ends): + logits[i, end:] = float("-inf") return logits @@ -121,21 +124,26 @@ def compare_top_k_results( @pytest.mark.parametrize("num_rows", NUM_ROWS) @pytest.mark.parametrize("top_k", TOP_K_VALUES) +@pytest.mark.parametrize("clean_logits", [True, False]) @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") @torch.inference_mode() def test_top_k_per_row( num_rows: int, top_k: int, + clean_logits: bool, ) -> None: """ Test top_k_per_row. """ + set_random_seed(0) torch.set_default_device("cuda:0") # Create test data vocab_size = 20000 row_starts, row_ends = create_row_boundaries(num_rows, vocab_size) - logits = create_random_logits(row_starts, row_ends, torch.float32, 42, "random") + logits = create_random_logits( + row_starts, row_ends, torch.float32, 42, clean_logits, "random" + ) # Create output tensors indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda") @@ -153,11 +161,12 @@ def test_top_k_per_row( ) # Run reference implementation - torch_indices = logits.topk(min(top_k, max(row_ends)), dim=-1)[1] - mask_lo = torch_indices >= 0 - mask_hi = (torch_indices - (row_ends - row_starts)[:, None]) < 0 - mask = mask_lo & mask_hi - torch_indices = torch_indices.masked_fill(~mask, -1) + torch_indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda") + for i in range(num_rows): + row_end = int(row_ends[i]) + k_i = min(top_k, row_end) + idx = logits[i, :row_end].topk(k_i, dim=-1)[1] + torch_indices[i, :k_i] = idx # Compare results assert compare_top_k_results( @@ -170,6 +179,7 @@ def _run_top_k_per_row_decode_test( batch_size: int, next_n: int, vocab_size: int, + clean_logits: bool, data_generation: str, ) -> None: """ @@ -180,14 +190,18 @@ def _run_top_k_per_row_decode_test( # Create test data num_rows = batch_size * next_n seq_lens = torch.randint( - vocab_size, (batch_size,), dtype=torch.int32, device="cuda" + low=next_n, + high=vocab_size, + size=(batch_size,), + dtype=torch.int32, + device="cuda", ) row_starts = torch.zeros(num_rows, dtype=torch.int32, device="cuda") row_indices = torch.arange(num_rows, device="cuda") // next_n next_n_offset = torch.arange(num_rows, device="cuda") % next_n row_ends = seq_lens[row_indices] - next_n + next_n_offset + 1 logits = create_random_logits( - row_starts, row_ends, torch.float32, 42, data_generation + row_starts, row_ends, torch.float32, 42, clean_logits, data_generation ) # Create output tensors @@ -208,11 +222,12 @@ def _run_top_k_per_row_decode_test( torch.cuda.synchronize() # Run reference implementation - torch_indices = logits.topk(min(top_k, max(row_ends)), dim=-1)[1] - mask_lo = torch_indices >= 0 - mask_hi = (torch_indices - (row_ends - row_starts)[:, None]) < 0 - mask = mask_lo & mask_hi - torch_indices = torch_indices.masked_fill(~mask, -1) + torch_indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda") + for i in range(num_rows): + row_end = int(row_ends[i]) + k_i = min(top_k, row_end) + idx = logits[i, :row_end].topk(k_i, dim=-1)[1] + torch_indices[i, :k_i] = idx # Compare results assert compare_top_k_results( @@ -223,6 +238,7 @@ def _run_top_k_per_row_decode_test( @pytest.mark.parametrize("top_k", TOP_K_VALUES) @pytest.mark.parametrize("batch_size", BATCH_SIZE) @pytest.mark.parametrize("next_n", NEXT_N) +@pytest.mark.parametrize("clean_logits", [True, False]) @pytest.mark.parametrize("data_generation", DATA_GENERATION) @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") @torch.inference_mode() @@ -230,28 +246,32 @@ def test_top_k_per_row_decode( top_k: int, batch_size: int, next_n: int, + clean_logits: bool, data_generation: str, ) -> None: """ Test top_k_per_row with seq_lens tensor. """ + set_random_seed(0) vocab_size = 20000 _run_top_k_per_row_decode_test( - top_k, batch_size, next_n, vocab_size, data_generation + top_k, batch_size, next_n, vocab_size, clean_logits, data_generation ) @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") +@pytest.mark.parametrize("clean_logits", [True, False]) @torch.inference_mode() -def test_top_k_per_row_decode_large_vocab_size() -> None: +def test_top_k_per_row_decode_large_vocab_size(clean_logits: bool) -> None: """ Test top_k_per_row_decode with large vocabulary size. """ + set_random_seed(0) top_k = 2048 batch_size = 2 next_n = 2 vocab_size = 300000 data_generation = "random" _run_top_k_per_row_decode_test( - top_k, batch_size, next_n, vocab_size, data_generation + top_k, batch_size, next_n, vocab_size, clean_logits, data_generation ) diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 77fe4c063ac..9ca7a42b70b 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -108,6 +108,7 @@ def sparse_attn_indexer( weights[chunk.token_start : chunk.token_end], chunk.cu_seqlen_ks, chunk.cu_seqlen_ke, + clean_logits=False, ) num_rows = logits.shape[0] @@ -157,6 +158,7 @@ def sparse_attn_indexer( decode_metadata.block_table, decode_metadata.schedule_metadata, max_model_len=max_model_len, + clean_logits=False, ) num_rows = logits.shape[0] diff --git a/vllm/utils/deep_gemm.py b/vllm/utils/deep_gemm.py index 129e9c9fa64..19e85ff6239 100644 --- a/vllm/utils/deep_gemm.py +++ b/vllm/utils/deep_gemm.py @@ -242,6 +242,7 @@ def fp8_mqa_logits( weights: torch.Tensor, cu_seqlen_ks: torch.Tensor, cu_seqlen_ke: torch.Tensor, + clean_logits: bool, ) -> torch.Tensor: """Compute FP8 MQA logits for a single sequence without KV paging. @@ -256,6 +257,7 @@ def fp8_mqa_logits( shape [M], dtype int32. cu_seqlen_ke: End indices (exclusive) for valid K per query position, shape [M], dtype int32. + clean_logits: Whether to clean the unfilled logits into `-inf`. Returns: Logits tensor of shape [M, N], dtype `torch.float32`. @@ -263,7 +265,9 @@ def fp8_mqa_logits( _lazy_init() if _fp8_mqa_logits_impl is None: return _missing() - return _fp8_mqa_logits_impl(q, kv, weights, cu_seqlen_ks, cu_seqlen_ke) + return _fp8_mqa_logits_impl( + q, kv, weights, cu_seqlen_ks, cu_seqlen_ke, clean_logits=clean_logits + ) def get_paged_mqa_logits_metadata( @@ -295,6 +299,7 @@ def fp8_paged_mqa_logits( block_tables: torch.Tensor, schedule_metadata: torch.Tensor, max_model_len: int, + clean_logits: bool, ) -> torch.Tensor: """Compute FP8 MQA logits using paged KV-cache. @@ -312,6 +317,7 @@ def fp8_paged_mqa_logits( schedule_metadata: Returned by `get_paged_mqa_logits_metadata`; used to distribute work across SMs. max_model_len: Maximum sequence length used to size the logits output. + clean_logits: Whether to clean the unfilled logits into `-inf`. Returns: Logits tensor of shape [B * next_n, max_model_len], dtype @@ -328,7 +334,7 @@ def fp8_paged_mqa_logits( block_tables, schedule_metadata, max_model_len, - clean_logits=True, + clean_logits=clean_logits, ) From 5819ca8944af4f7dcbac3c6b73179f760e05910d Mon Sep 17 00:00:00 2001 From: Simon Mo Date: Thu, 5 Feb 2026 17:42:22 -0800 Subject: [PATCH 118/810] [Docs] Add reo analytics (#33957) Signed-off-by: simon-mo --- docs/mkdocs/javascript/reo.js | 3 +++ mkdocs.yaml | 1 + 2 files changed, 4 insertions(+) create mode 100644 docs/mkdocs/javascript/reo.js diff --git a/docs/mkdocs/javascript/reo.js b/docs/mkdocs/javascript/reo.js new file mode 100644 index 00000000000..13350abdc1e --- /dev/null +++ b/docs/mkdocs/javascript/reo.js @@ -0,0 +1,3 @@ +// Reo.Dev documentation tracking +// https://docs.reo.dev/integrations/tracking-beacon/install-javascript-for-documentation +!function(){var e,t,n;e="d5c4337961ef0ac",t=function(){Reo.init({clientID:"d5c4337961ef0ac"})},(n=document.createElement("script")).src="https://static.reo.dev/"+e+"/reo.js",n.defer=!0,n.onload=t,document.head.appendChild(n)}(); diff --git a/mkdocs.yaml b/mkdocs.yaml index 2797f6dee0d..d5d6852f31d 100644 --- a/mkdocs.yaml +++ b/mkdocs.yaml @@ -141,6 +141,7 @@ extra_css: - mkdocs/stylesheets/extra.css extra_javascript: + - mkdocs/javascript/reo.js - mkdocs/javascript/run_llm_widget.js - mkdocs/javascript/mathjax.js - https://unpkg.com/mathjax@3.2.2/es5/tex-mml-chtml.js From 20d7454c9bb0c2de7f59863f5030e5f494cab178 Mon Sep 17 00:00:00 2001 From: Rabi Mishra Date: Fri, 6 Feb 2026 07:52:53 +0530 Subject: [PATCH 119/810] fix(ROCm): Make flash_attn import optional in MLA attention (#33511) Signed-off-by: rabi --- .../layers/attention/mla_attention.py | 22 ++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 4859af43ae4..c31aa7b41d0 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -919,10 +919,20 @@ try: is_vllm_fa = True except ImportError: - # For rocm use upstream flash attention - if current_platform.is_rocm(): - from flash_attn import flash_attn_varlen_func # type: ignore[no-redef] is_vllm_fa = False + flash_attn_varlen_func = None # type: ignore[assignment] + # On ROCm, vllm_flash_attn is not available, try upstream flash_attn instead. + # On CUDA, vllm_flash_attn should always be available (built with vLLM), + # so we don't attempt the fallback there. + if current_platform.is_rocm(): + try: + from flash_attn import flash_attn_varlen_func # type: ignore[no-redef] + except ImportError: + logger.debug( + "flash_attn not available on ROCm; " + "MLA models using TRITON_MLA will require flash_attn. " + "AITER_MLA backends use aiter kernels instead." + ) def dynamic_per_batched_tensor_quant( @@ -1917,6 +1927,12 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]): self._run_prefill_new_tokens = self._run_prefill_new_tokens_cudnn self._pad_v = False else: # Use FlashAttention + if flash_attn_varlen_func is None: + raise RuntimeError( + "MLA attention requires FlashAttention but it is not " + "available. Please install flash_attn or use " + "--attention-backend ROCM_AITER_MLA." + ) logger.info_once("Using FlashAttention prefill for MLA", scope="local") self._run_prefill_context_chunk = self._run_prefill_context_chunk_fa self._run_prefill_new_tokens = self._run_prefill_new_tokens_fa From a32cb49b60688fb64a6d3d7f86378b4d2fad06e6 Mon Sep 17 00:00:00 2001 From: Mingliang Li Date: Fri, 6 Feb 2026 11:38:02 +0800 Subject: [PATCH 120/810] feat(frontend): early-fail tokenization guard for user requests (#31366) Signed-off-by: limingliang Signed-off-by: DarkLight1337 Co-authored-by: limingliang Co-authored-by: DarkLight1337 --- tests/renderers/test_completions.py | 425 ++++++++++++++++------------ vllm/renderers/params.py | 74 +++-- vllm/tokenizers/deepseek_v32.py | 4 + vllm/tokenizers/grok2.py | 6 + vllm/tokenizers/hf.py | 6 + vllm/tokenizers/mistral.py | 5 + vllm/tokenizers/protocol.py | 4 + 7 files changed, 315 insertions(+), 209 deletions(-) diff --git a/tests/renderers/test_completions.py b/tests/renderers/test_completions.py index 17274ccd7c3..84b7230d99f 100644 --- a/tests/renderers/test_completions.py +++ b/tests/renderers/test_completions.py @@ -4,7 +4,6 @@ import io from dataclasses import dataclass from typing import Any -from unittest.mock import AsyncMock import pybase64 import pytest @@ -28,7 +27,6 @@ class MockModelConfig: model: str = MODEL_NAME tokenizer: str = MODEL_NAME trust_remote_code: bool = False - max_model_len: int = 100 tokenizer_revision = None tokenizer_mode = "auto" hf_config = MockHFConfig() @@ -37,25 +35,50 @@ class MockModelConfig: skip_tokenizer_init: bool = False -@pytest.fixture -def mock_model_config(): - return MockModelConfig() +@dataclass +class DummyTokenizer: + truncation_side: str = "left" + max_chars_per_token: int = 1 + + def __post_init__(self) -> None: + self._captured_encode_kwargs: dict = {} + + def decode(self, tokens: list[int]): + return str(tokens) + + def encode(self, text: str, **kwargs): + self._captured_encode_kwargs = kwargs + + in_length = len(text) + truncation = kwargs.get("truncation") + max_length = kwargs.get("max_length") + if truncation and max_length is not None: + return list(range(min(in_length, max_length))) + + return list(range(in_length)) -@pytest.fixture -def mock_async_tokenizer(): - return AsyncMock() +def _build_renderer( + model_config: MockModelConfig, + *, + truncation_side: str = "left", + max_chars_per_token: int = 1, +): + _, tokenizer_name, _, kwargs = tokenizer_args_from_config(model_config) - -@pytest.fixture -def renderer(mock_model_config): - _, tokenizer_name, _, kwargs = tokenizer_args_from_config(mock_model_config) - - return HfRenderer( - mock_model_config, + renderer = HfRenderer( + model_config, tokenizer_kwargs={**kwargs, "tokenizer_name": tokenizer_name}, ) + if not model_config.skip_tokenizer_init: + renderer._tokenizer = DummyTokenizer( + truncation_side=truncation_side, + max_chars_per_token=max_chars_per_token, + ) + + return renderer + class TestValidatePrompt: STRING_INPUTS = [ @@ -81,39 +104,50 @@ class TestValidatePrompt: ] # Test that a nested mixed-type list of lists raises a TypeError. - def test_empty_input(self, renderer): + def test_empty_input(self): + renderer = _build_renderer(MockModelConfig()) + with pytest.raises(ValueError, match="at least one prompt"): renderer.render_completions([]) - def test_invalid_type(self, renderer): + def test_invalid_type(self): + renderer = _build_renderer(MockModelConfig()) + with pytest.raises(TypeError, match="string or an array of tokens"): renderer.render_completions([[1, 2], ["foo", "bar"]]) @pytest.mark.parametrize("string_input", STRING_INPUTS) - def test_string_consistent(self, renderer, string_input: str): + def test_string_consistent(self, string_input: str): + renderer = _build_renderer(MockModelConfig()) + assert renderer.render_completions(string_input) == renderer.render_completions( [string_input] ) @pytest.mark.parametrize("token_input", TOKEN_INPUTS) - def test_token_consistent(self, renderer, token_input: list[int]): + def test_token_consistent(self, token_input: list[int]): + renderer = _build_renderer(MockModelConfig()) + assert renderer.render_completions(token_input) == renderer.render_completions( [token_input] ) @pytest.mark.parametrize("inputs_slice", INPUTS_SLICES) - def test_string_slice(self, renderer, inputs_slice: slice): + def test_string_slice(self, inputs_slice: slice): + renderer = _build_renderer(MockModelConfig()) + assert renderer.render_completions(self.STRING_INPUTS)[ inputs_slice ] == renderer.render_completions(self.STRING_INPUTS[inputs_slice]) class TestRenderPrompt: - @pytest.mark.asyncio - async def test_token_input(self, renderer): + def test_token_input(self): + renderer = _build_renderer(MockModelConfig()) + tokens = [101, 7592, 2088] - prompts = await renderer.render_completions_async(tokens) - results = await renderer.tokenize_prompts_async( + prompts = renderer.render_completions(tokens) + results = renderer.tokenize_prompts( prompts, TokenizeParams(max_total_tokens=100), ) @@ -121,11 +155,12 @@ class TestRenderPrompt: assert len(results) == 1 assert results[0]["prompt_token_ids"] == tokens - @pytest.mark.asyncio - async def test_token_list_input(self, renderer): + def test_token_list_input(self): + renderer = _build_renderer(MockModelConfig()) + token_lists = [[101, 7592, 2088], [102, 1234, 5678, 9012], [103, 4567]] - prompts = await renderer.render_completions_async(token_lists) - results = await renderer.tokenize_prompts_async( + prompts = renderer.render_completions(token_lists) + results = renderer.tokenize_prompts( prompts, TokenizeParams(max_total_tokens=100), ) @@ -135,167 +170,178 @@ class TestRenderPrompt: assert results[1]["prompt_token_ids"] == [102, 1234, 5678, 9012] assert results[2]["prompt_token_ids"] == [103, 4567] - @pytest.mark.asyncio - async def test_text_input(self, renderer, mock_async_tokenizer): - mock_async_tokenizer.encode.return_value = [101, 7592, 2088] - renderer._async_tokenizer = mock_async_tokenizer + def test_text_input(self): + renderer = _build_renderer(MockModelConfig()) - prompts = await renderer.render_completions_async("Hello world") - results = await renderer.tokenize_prompts_async( + text_input = "x" * 10 + prompts = renderer.render_completions(text_input) + results = renderer.tokenize_prompts( prompts, TokenizeParams(max_total_tokens=100), ) assert len(results) == 1 - assert results[0]["prompt_token_ids"] == [101, 7592, 2088] - mock_async_tokenizer.encode.assert_called_once() + assert len(results[0]["prompt_token_ids"]) == 10 - @pytest.mark.asyncio - async def test_text_list_input(self, renderer, mock_async_tokenizer): - mock_async_tokenizer.encode.return_value = [101, 7592, 2088] - renderer._async_tokenizer = mock_async_tokenizer + def test_text_list_input(self): + renderer = _build_renderer(MockModelConfig()) - text_list_input = ["Hello world", "How are you?", "Good morning"] - prompts = await renderer.render_completions_async(text_list_input) - results = await renderer.tokenize_prompts_async( + text_list_input = ["x" * 10, "x" * 12, "x" * 14] + prompts = renderer.render_completions(text_list_input) + results = renderer.tokenize_prompts( prompts, TokenizeParams(max_total_tokens=100), ) assert len(results) == 3 - for result in results: - assert result["prompt_token_ids"] == [101, 7592, 2088] - assert mock_async_tokenizer.encode.call_count == 3 + for text_input, result in zip(text_list_input, results): + assert len(result["prompt_token_ids"]) == len(text_input) - @pytest.mark.asyncio - async def test_no_truncation(self, renderer, mock_async_tokenizer): - mock_async_tokenizer.encode.return_value = [101, 7592, 2088] - renderer._async_tokenizer = mock_async_tokenizer + def test_zero_truncation(self): + renderer = _build_renderer(MockModelConfig()) - prompts = await renderer.render_completions_async("Hello world") - results = await renderer.tokenize_prompts_async( + prompts = renderer.render_completions("x" * 200) + results = renderer.tokenize_prompts( prompts, - TokenizeParams(max_total_tokens=100), + TokenizeParams(max_total_tokens=100, truncate_prompt_tokens=0), ) assert len(results) == 1 - call_args = mock_async_tokenizer.encode.call_args - assert ( - "truncation" not in call_args.kwargs - or call_args.kwargs["truncation"] is False - ) + assert len(results[0]["prompt_token_ids"]) == 0 - @pytest.mark.asyncio - async def test_truncation_positive(self, renderer, mock_async_tokenizer): - mock_async_tokenizer.encode.return_value = [101, 7592, 2088] # Truncated - renderer._async_tokenizer = mock_async_tokenizer + def test_pos_truncation(self): + renderer = _build_renderer(MockModelConfig()) - prompts = await renderer.render_completions_async("Hello world") - results = await renderer.tokenize_prompts_async( + prompts = renderer.render_completions("x" * 200) + results = renderer.tokenize_prompts( prompts, - TokenizeParams( - max_total_tokens=200, - truncate_prompt_tokens=50, - ), + TokenizeParams(max_total_tokens=100, truncate_prompt_tokens=50), ) assert len(results) == 1 - call_args = mock_async_tokenizer.encode.call_args - assert call_args.kwargs["truncation"] is True - assert call_args.kwargs["max_length"] == 50 + assert len(results[0]["prompt_token_ids"]) == 50 - @pytest.mark.asyncio - async def test_truncation_negative(self, renderer, mock_async_tokenizer): - # Test that negative truncation uses model's max_model_len - mock_async_tokenizer.encode.return_value = [ - 101, - 7592, - 2088, - ] # Truncated to max_model_len - renderer._async_tokenizer = mock_async_tokenizer + def test_neg_truncation(self): + renderer = _build_renderer(MockModelConfig()) - prompts = await renderer.render_completions_async("Hello world") - results = await renderer.tokenize_prompts_async( + prompts = renderer.render_completions("x" * 200) + results = renderer.tokenize_prompts( prompts, - TokenizeParams( - max_total_tokens=200, - truncate_prompt_tokens=-1, - ), + TokenizeParams(max_total_tokens=100, truncate_prompt_tokens=-1), ) assert len(results) == 1 - call_args = mock_async_tokenizer.encode.call_args - assert call_args.kwargs["truncation"] is True - assert call_args.kwargs["max_length"] == 200 + assert len(results[0]["prompt_token_ids"]) == 100 # max_total_tokens + + def test_truncation_left(self): + renderer = _build_renderer(MockModelConfig(), truncation_side="left") - @pytest.mark.asyncio - async def test_token_truncation_last_elements(self, renderer): - # Test that token truncation keeps the last N elements long_tokens = [100, 101, 102, 103, 104, 105, 106, 107, 108, 109] # 10 tokens - prompts = await renderer.render_completions_async(long_tokens) - results = await renderer.tokenize_prompts_async( + prompts = renderer.render_completions(long_tokens) + results = renderer.tokenize_prompts( prompts, - TokenizeParams( - max_total_tokens=100, - truncate_prompt_tokens=5, - ), + TokenizeParams(max_total_tokens=100, truncate_prompt_tokens=5), ) assert len(results) == 1 # Should keep the last 5 tokens: [105, 106, 107, 108, 109] assert results[0]["prompt_token_ids"] == [105, 106, 107, 108, 109] - @pytest.mark.asyncio - async def test_max_length_exceeded(self, renderer): - long_tokens = list(range(150)) # Exceeds max_model_len=100 + def test_truncation_right(self): + renderer = _build_renderer(MockModelConfig(), truncation_side="right") - prompts = await renderer.render_completions_async(long_tokens) - - with pytest.raises(ValueError, match="context length is only"): - await renderer.tokenize_prompts_async( - prompts, - TokenizeParams(max_total_tokens=100), - ) - - @pytest.mark.asyncio - async def test_no_tokenizer_for_text(self, renderer): - renderer_no_tokenizer = HfRenderer.from_config( - MockModelConfig(skip_tokenizer_init=True), - tokenizer_kwargs={}, + long_tokens = [100, 101, 102, 103, 104, 105, 106, 107, 108, 109] # 10 tokens + prompts = renderer.render_completions(long_tokens) + results = renderer.tokenize_prompts( + prompts, + TokenizeParams(max_total_tokens=100, truncate_prompt_tokens=5), ) - prompts = await renderer_no_tokenizer.render_completions_async("Hello world") + assert len(results) == 1 + # Should keep the first 5 tokens: [100, 101, 102, 103, 104] + assert results[0]["prompt_token_ids"] == [100, 101, 102, 103, 104] - with pytest.raises(ValueError, match="`skip_tokenizer_init=True`"): - await renderer_no_tokenizer.tokenize_prompts_async( + def test_text_max_length_exceeded_obvious(self): + renderer = _build_renderer(MockModelConfig(), max_chars_per_token=1) + + # Exceeds max_total_tokens and max_total_tokens * VLLM_MAX_CHARS_PER_TOKEN + long_tokens = "x" * 150 + prompts = renderer.render_completions(long_tokens) + + with pytest.raises( + ValueError, + match="input characters and requested .* context length is only", + ): + renderer.tokenize_prompts( prompts, TokenizeParams(max_total_tokens=100), ) - @pytest.mark.asyncio - async def test_token_input_with_needs_detokenization( - self, renderer, mock_async_tokenizer - ): - # When needs_detokenization=True for token inputs, renderer should - # use the async tokenizer to decode and include the original text - # in the returned prompt object. - mock_async_tokenizer.decode = AsyncMock(return_value="decoded text") - renderer._async_tokenizer = mock_async_tokenizer + # Should not even attempt tokenization + assert renderer._tokenizer._captured_encode_kwargs == {} + + def test_text_max_length_exceeded_nonobvious(self): + renderer = _build_renderer(MockModelConfig(), max_chars_per_token=2) + + # Exceeds max_total_tokens but not max_total_tokens * VLLM_MAX_CHARS_PER_TOKEN + long_tokens = "x" * 150 + prompts = renderer.render_completions(long_tokens) + + with pytest.raises( + ValueError, + match="input tokens and requested .* context length is only", + ): + renderer.tokenize_prompts( + prompts, + TokenizeParams(max_total_tokens=100), + ) + + # Should only tokenize the first max_total_tokens + 1 tokens + assert renderer._tokenizer._captured_encode_kwargs["truncation"] is True + assert renderer._tokenizer._captured_encode_kwargs["max_length"] == 101 + + def test_token_max_length_exceeded(self): + renderer = _build_renderer(MockModelConfig()) + + long_tokens = list(range(150)) # Exceeds max_total_tokens=100 + prompts = renderer.render_completions(long_tokens) + + with pytest.raises( + ValueError, + match="input tokens and requested .* context length is only", + ): + renderer.tokenize_prompts( + prompts, + TokenizeParams(max_total_tokens=100, truncate_prompt_tokens=None), + ) + + def test_no_tokenizer_for_text(self): + renderer = _build_renderer(MockModelConfig(skip_tokenizer_init=True)) + + prompts = renderer.render_completions("Hello world") + + with pytest.raises(ValueError, match="`skip_tokenizer_init=True`"): + renderer.tokenize_prompts( + prompts, + TokenizeParams(max_total_tokens=100), + ) + + def test_token_input_with_needs_detokenization(self): + renderer = _build_renderer(MockModelConfig()) tokens = [1, 2, 3, 4] - prompts = await renderer.render_completions_async(tokens) - results = await renderer.tokenize_prompts_async( + prompts = renderer.render_completions(tokens) + results = renderer.tokenize_prompts( prompts, TokenizeParams( - max_total_tokens=renderer.config.max_model_len, + max_total_tokens=100, needs_detokenization=True, ), ) assert len(results) == 1 assert results[0]["prompt_token_ids"] == tokens - assert results[0]["prompt"] == "decoded text" - mock_async_tokenizer.decode.assert_awaited_once() + assert results[0]["prompt"] == "[1, 2, 3, 4]" class TestRenderEmbedPrompt: @@ -306,118 +352,121 @@ class TestRenderEmbedPrompt: buffer.seek(0) return pybase64.b64encode(buffer.read()) - @pytest.mark.asyncio - async def test_single_prompt_embed(self, renderer): - # Create a test tensor - test_tensor = torch.randn(10, 768, dtype=torch.float32) - embed_bytes = self._create_test_embed_bytes(test_tensor) + def test_single_prompt_embed(self): + renderer = _build_renderer(MockModelConfig()) - prompts = await renderer.render_completions_async(prompt_embeds=embed_bytes) - results = await renderer.tokenize_prompts_async( + # Create a test tensor + tensor_input = torch.randn(10, 768, dtype=torch.float32) + embed_bytes = self._create_test_embed_bytes(tensor_input) + + prompts = renderer.render_completions(prompt_embeds=embed_bytes) + results = renderer.tokenize_prompts( prompts, - TokenizeParams(max_total_tokens=renderer.config.max_model_len), + TokenizeParams(max_total_tokens=100), ) assert len(results) == 1 - assert torch.allclose(results[0]["prompt_embeds"], test_tensor) + assert torch.equal(results[0]["prompt_embeds"], tensor_input) + + def test_multiple_prompt_embeds(self): + renderer = _build_renderer(MockModelConfig()) - @pytest.mark.asyncio - async def test_multiple_prompt_embeds(self, renderer): # Create multiple test tensors - test_tensors = [ + tensor_inputs = [ torch.randn(8, 512, dtype=torch.float32), torch.randn(12, 512, dtype=torch.float32), ] - embed_bytes_list = [self._create_test_embed_bytes(t) for t in test_tensors] - prompts = await renderer.render_completions_async( - prompt_embeds=embed_bytes_list + prompts = renderer.render_completions( + prompt_embeds=[self._create_test_embed_bytes(t) for t in tensor_inputs], ) - results = await renderer.tokenize_prompts_async( + results = renderer.tokenize_prompts( prompts, - TokenizeParams(max_total_tokens=renderer.config.max_model_len), + TokenizeParams(max_total_tokens=100), ) assert len(results) == 2 for i, result in enumerate(results): - assert torch.allclose(result["prompt_embeds"], test_tensors[i]) + assert torch.allclose(result["prompt_embeds"], tensor_inputs[i]) + + def test_prompt_embed_truncation(self): + renderer = _build_renderer(MockModelConfig()) - @pytest.mark.asyncio - async def test_prompt_embed_truncation(self, renderer): # Create tensor with more tokens than truncation limit - test_tensor = torch.randn(20, 768, dtype=torch.float32) - embed_bytes = self._create_test_embed_bytes(test_tensor) + tensor_input = torch.randn(20, 768, dtype=torch.float32) - prompts = await renderer.render_completions_async(prompt_embeds=embed_bytes) - results = await renderer.tokenize_prompts_async( + prompts = renderer.render_completions( + prompt_embeds=self._create_test_embed_bytes(tensor_input), + ) + results = renderer.tokenize_prompts( prompts, TokenizeParams( - max_total_tokens=renderer.config.max_model_len, + max_total_tokens=100, truncate_prompt_tokens=10, ), ) assert len(results) == 1 # Should keep last 10 tokens - expected = test_tensor[-10:] - assert torch.allclose(results[0]["prompt_embeds"], expected) + expected = tensor_input[-10:] + assert torch.equal(results[0]["prompt_embeds"], expected) + + def test_prompt_embed_different_dtypes(self): + renderer = _build_renderer(MockModelConfig()) - @pytest.mark.asyncio - async def test_prompt_embed_different_dtypes(self, renderer): # Test different supported dtypes dtypes = [torch.float32, torch.float16, torch.bfloat16] for dtype in dtypes: - test_tensor = torch.randn(5, 256, dtype=dtype) - embed_bytes = self._create_test_embed_bytes(test_tensor) + tensor_input = torch.randn(5, 256, dtype=dtype) - prompts = await renderer.render_completions_async(prompt_embeds=embed_bytes) - results = await renderer.tokenize_prompts_async( + prompts = renderer.render_completions( + prompt_embeds=self._create_test_embed_bytes(tensor_input), + ) + results = renderer.tokenize_prompts( prompts, - TokenizeParams(max_total_tokens=renderer.config.max_model_len), + TokenizeParams(max_total_tokens=100), ) assert len(results) == 1 assert results[0]["prompt_embeds"].dtype == dtype - @pytest.mark.asyncio - async def test_prompt_embed_squeeze_batch_dim(self, renderer): - # Test tensor with batch dimension gets squeezed - test_tensor = torch.randn(1, 10, 768, dtype=torch.float32) - embed_bytes = self._create_test_embed_bytes(test_tensor) + def test_prompt_embed_squeeze_batch_dim(self): + renderer = _build_renderer(MockModelConfig()) - prompts = await renderer.render_completions_async(prompt_embeds=embed_bytes) - results = await renderer.tokenize_prompts_async( + # Test tensor with batch dimension gets squeezed + tensor_input = torch.randn(1, 10, 768, dtype=torch.float32) + + prompts = renderer.render_completions( + prompt_embeds=self._create_test_embed_bytes(tensor_input), + ) + results = renderer.tokenize_prompts( prompts, - TokenizeParams(max_total_tokens=renderer.config.max_model_len), + TokenizeParams(max_total_tokens=100), ) assert len(results) == 1 # Should be squeezed to 2D assert results[0]["prompt_embeds"].shape == (10, 768) - @pytest.mark.asyncio - async def test_both_prompts_and_embeds(self, renderer, mock_async_tokenizer): - # Set up text tokenization - mock_async_tokenizer.encode.return_value = [101, 102, 103] - renderer._async_tokenizer = mock_async_tokenizer + def test_both_prompts_and_embeds(self): + renderer = _build_renderer(MockModelConfig()) - # Create embed - test_tensor = torch.randn(5, 256, dtype=torch.float32) - embed_bytes = self._create_test_embed_bytes(test_tensor) + text_input = "Hello world" + tensor_input = torch.randn(5, 256, dtype=torch.float32) - prompts = await renderer.render_completions_async( - "Hello world", - prompt_embeds=embed_bytes, + prompts = renderer.render_completions( + text_input, + prompt_embeds=self._create_test_embed_bytes(tensor_input), ) - results = await renderer.tokenize_prompts_async( + results = renderer.tokenize_prompts( prompts, - TokenizeParams(max_total_tokens=renderer.config.max_model_len), + TokenizeParams(max_total_tokens=100), ) assert len(results) == 2 # First should be embed prompt - assert torch.allclose(results[0]["prompt_embeds"], test_tensor) + assert torch.equal(results[0]["prompt_embeds"], tensor_input) # Second should be tokens prompt assert "prompt_token_ids" in results[1] - assert results[1]["prompt_token_ids"] == [101, 102, 103] + assert len(results[1]["prompt_token_ids"]) == len(text_input) diff --git a/vllm/renderers/params.py b/vllm/renderers/params.py index d20038478de..a860fcd951f 100644 --- a/vllm/renderers/params.py +++ b/vllm/renderers/params.py @@ -229,23 +229,53 @@ class TokenizeParams: max_length = self.truncate_prompt_tokens if max_length is not None and max_length < 0: max_length = self.max_input_tokens + elif max_length is None and self.max_input_tokens is not None: + # This prevents tokenization from taking up more resources than necessary + # while still failing `self._token_len_check` as expected by users + max_length = self.max_input_tokens + 1 return dict( - truncation=self.truncate_prompt_tokens is not None, + truncation=max_length is not None, max_length=max_length, add_special_tokens=self.add_special_tokens, ) - def _apply_lowercase(self, tokenizer: TokenizerLike | None, text: str) -> str: - if self.do_lower_case: - text = text.lower() + def _text_len_check(self, tokenizer: TokenizerLike | None, text: str) -> str: + """Apply length checks to prompt text if necessary.""" + max_input_tokens = self.max_input_tokens + if max_input_tokens is None: + return text + + if self.truncate_prompt_tokens is None and tokenizer is not None: + max_input_chars = max_input_tokens * tokenizer.max_chars_per_token + + if len(text) > max_input_chars: + # To save resources, fail the request outright without even + # attempting tokenization + raise VLLMValidationError( + f"You passed {len(text)} input characters " + f"and requested {self.max_output_tokens} output tokens. " + f"However, the model's context length is only " + f"{self.max_total_tokens} tokens, resulting in a maximum " + f"input length of {max_input_tokens} tokens " + f"(at most {max_input_chars} characters). " + f"Please reduce the length of the input prompt.", + parameter="input_text", + value=len(text), + ) return text + def _text_lowercase(self, tokenizer: TokenizerLike | None, text: str) -> str: + """Apply lowercase to prompt text if necessary.""" + return text.lower() if self.do_lower_case else text + def _validate_text(self, tokenizer: TokenizerLike | None, text: str) -> str: """Apply all validators to prompt text.""" - # TODO: Implement https://github.com/vllm-project/vllm/pull/31366 - for validator in (self._apply_lowercase,): + for validator in ( + self._text_len_check, + self._text_lowercase, + ): text = validator(tokenizer, text) return text @@ -265,8 +295,8 @@ class TokenizeParams: return prompt - def _apply_padding(self, tokenizer: TokenizerLike | None, tokens: _S) -> _S: - """Apply padding to a token sequence.""" + def _token_padding(self, tokenizer: TokenizerLike | None, tokens: _S) -> _S: + """Apply padding to prompt tokens if necessary.""" pad_length = self.pad_prompt_tokens if pad_length is not None and pad_length < 0: pad_length = self.max_input_tokens @@ -281,8 +311,8 @@ class TokenizeParams: return tokens + [tokenizer.pad_token_id] * (pad_length - len(tokens)) - def _apply_truncation(self, tokenizer: TokenizerLike | None, tokens: _S) -> _S: - """Apply truncation to a token sequence.""" + def _token_truncation(self, tokenizer: TokenizerLike | None, tokens: _S) -> _S: + """Apply truncation to prompt tokens if necessary.""" max_length = self.truncate_prompt_tokens if max_length is not None and max_length < 0: max_length = self.max_input_tokens @@ -297,18 +327,20 @@ class TokenizeParams: return tokens[:max_length] - def _apply_length_check(self, tokenizer: TokenizerLike | None, tokens: _S) -> _S: - """Apply length checks to a token sequence.""" + def _token_len_check(self, tokenizer: TokenizerLike | None, tokens: _S) -> _S: + """Apply length checks to prompt tokens if necessary.""" max_input_tokens = self.max_input_tokens + if max_input_tokens is None: + return tokens - if max_input_tokens is not None and len(tokens) > max_input_tokens: + if len(tokens) > max_input_tokens: raise VLLMValidationError( - f"You passed {len(tokens)} input tokens and " - f"requested {self.max_output_tokens} output tokens. " + f"You passed {len(tokens)} input tokens " + f"and requested {self.max_output_tokens} output tokens. " f"However, the model's context length is only " - f"{self.max_total_tokens}, resulting in a maximum " - f"input length of {max_input_tokens}. " - f"Please reduce the length of the input messages.", + f"{self.max_total_tokens} tokens, resulting in a maximum " + f"input length of {max_input_tokens} tokens. " + f"Please reduce the length of the input prompt.", parameter="input_tokens", value=len(tokens), ) @@ -318,9 +350,9 @@ class TokenizeParams: def _validate_tokens(self, tokenizer: TokenizerLike | None, tokens: _S) -> _S: """Apply all validators to a token sequence.""" for validator in ( - self._apply_padding, - self._apply_truncation, - self._apply_length_check, + self._token_padding, + self._token_truncation, + self._token_len_check, ): tokens = validator(tokenizer, tokens) diff --git a/vllm/tokenizers/deepseek_v32.py b/vllm/tokenizers/deepseek_v32.py index 4402054c9a5..cb0ffe73a0b 100644 --- a/vllm/tokenizers/deepseek_v32.py +++ b/vllm/tokenizers/deepseek_v32.py @@ -115,6 +115,10 @@ class DeepseekV32Tokenizer(CachedHfTokenizer): def max_token_id(self) -> int: return self.tokenizer.max_token_id + @property + def max_chars_per_token(self) -> int: + return self.tokenizer.max_chars_per_token + @property def truncation_side(self) -> str: return self.tokenizer.truncation_side diff --git a/vllm/tokenizers/grok2.py b/vllm/tokenizers/grok2.py index fe00f5e56c5..3b984152ef7 100644 --- a/vllm/tokenizers/grok2.py +++ b/vllm/tokenizers/grok2.py @@ -277,6 +277,8 @@ class Grok2Tokenizer(TokenizerLike): self._pad_token_id = self._special_tokens.get(PAD, self._eos_token_id) self._unk_token_id = self._pad_token_id + self._max_chars_per_token = max(len(tok) for tok in self._token_to_id) + def num_special_tokens_to_add(self) -> int: return 0 @@ -312,6 +314,10 @@ class Grok2Tokenizer(TokenizerLike): def max_token_id(self) -> int: return self._tokenizer.n_vocab - 1 + @property + def max_chars_per_token(self) -> int: + return self._max_chars_per_token + @property def truncation_side(self) -> str: return self._truncation_side diff --git a/vllm/tokenizers/hf.py b/vllm/tokenizers/hf.py index a7b565dca5d..85c81239852 100644 --- a/vllm/tokenizers/hf.py +++ b/vllm/tokenizers/hf.py @@ -28,6 +28,8 @@ def get_cached_tokenizer(tokenizer: HfTokenizer) -> HfTokenizer: tokenizer_len = len(tokenizer) max_token_id = max(tokenizer_vocab.values()) + max_chars_per_token = max(len(tok) for tok in tokenizer_vocab) + # Some tokenizers (e.g., QwenTokenizer) have special tokens that # are added and included in the implementation of the vocab_size # property, but not in get_vocab(); if there is an implementation @@ -49,6 +51,10 @@ def get_cached_tokenizer(tokenizer: HfTokenizer) -> HfTokenizer: def max_token_id(self) -> int: return max_token_id + @property + def max_chars_per_token(self) -> int: + return max_chars_per_token + def get_vocab(self) -> dict[str, int]: return tokenizer_vocab diff --git a/vllm/tokenizers/mistral.py b/vllm/tokenizers/mistral.py index bb85052dba8..b56b2718c74 100644 --- a/vllm/tokenizers/mistral.py +++ b/vllm/tokenizers/mistral.py @@ -272,6 +272,7 @@ class MistralTokenizer(TokenizerLike): # Vocab sorted by token id. self._vocab = self.tokenizer.vocab() self._max_token_id = self.vocab_size - 1 + self._max_chars_per_token = max(len(tok) for tok in self._vocab) # Cache special tokens for faster access. self._special_token_ids = self._get_special_token_ids() @@ -325,6 +326,10 @@ class MistralTokenizer(TokenizerLike): def max_token_id(self) -> int: return self._max_token_id + @property + def max_chars_per_token(self) -> int: + return self._max_chars_per_token + @property def truncation_side(self) -> str: return self.transformers_tokenizer.truncation_side diff --git a/vllm/tokenizers/protocol.py b/vllm/tokenizers/protocol.py index 21e5b3a7bbd..6f091379e11 100644 --- a/vllm/tokenizers/protocol.py +++ b/vllm/tokenizers/protocol.py @@ -57,6 +57,10 @@ class TokenizerLike(Protocol): def max_token_id(self) -> int: raise NotImplementedError + @property + def max_chars_per_token(self) -> int: + raise NotImplementedError + @property def truncation_side(self) -> str: raise NotImplementedError From 035a6cb09a3f7076d2e79db1af3070325230bb59 Mon Sep 17 00:00:00 2001 From: Cyrus Leung Date: Fri, 6 Feb 2026 11:38:39 +0800 Subject: [PATCH 121/810] [Misc] Update code for encoder-decoder models (#33900) Signed-off-by: DarkLight1337 --- vllm/multimodal/inputs.py | 2 +- vllm/v1/core/sched/scheduler.py | 10 ++++++++-- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/vllm/multimodal/inputs.py b/vllm/multimodal/inputs.py index 2cc7900eb67..9f01af9a8ab 100644 --- a/vllm/multimodal/inputs.py +++ b/vllm/multimodal/inputs.py @@ -1098,7 +1098,7 @@ class MultiModalEncDecInputs(MultiModalInputs): Note: Even text-only encoder-decoder models are currently implemented as multi-modal models for convenience. - (Example: https://github.com/neuralmagic/bart-plugin) + (Example: https://github.com/vllm-project/bart-plugin) """ encoder_prompt_token_ids: list[int] diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 745d9ffec77..9f0643e4fda 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -185,7 +185,13 @@ class Scheduler(SchedulerInterface): # NOTE: Text-only encoder-decoder models are implemented as # multi-modal models for convenience - # Example: https://github.com/neuralmagic/bart-plugin + # Example: https://github.com/vllm-project/bart-plugin + if self.is_encoder_decoder: + assert mm_budget and len(mm_budget.mm_max_toks_per_item) <= 1, ( + "Encoder-decoder models are expected to implement the " + "multimodal interface with at most one modality." + ) + self.max_num_encoder_input_tokens = ( mm_budget.encoder_compute_budget if mm_budget else 0 ) @@ -200,7 +206,7 @@ class Scheduler(SchedulerInterface): # TODO (NickLucche): Generalize to models with variable-length encoder inputs. self._num_encoder_max_input_tokens = ( mm_budget.mm_max_toks_per_item[mm_budget.get_modality_with_max_tokens()] - if mm_budget + if mm_budget and mm_budget.mm_max_toks_per_item else 0 ) From ac04dd374f996f8df960933fa076bb5ea53c0c2a Mon Sep 17 00:00:00 2001 From: R3hankhan Date: Fri, 6 Feb 2026 10:27:02 +0530 Subject: [PATCH 122/810] [CPU] Add BF16 Kernel type for s390x (#33788) Signed-off-by: Rehan Khan --- csrc/cpu/mla_decode.cpp | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/csrc/cpu/mla_decode.cpp b/csrc/cpu/mla_decode.cpp index bc0ac5bc5ce..bd489b463d0 100644 --- a/csrc/cpu/mla_decode.cpp +++ b/csrc/cpu/mla_decode.cpp @@ -38,6 +38,15 @@ struct KernelVecType { using qk_vec_type = vec_op::BF16Vec32; using v_load_vec_type = vec_op::BF16Vec16; }; + +#elif defined(__s390x__) +template <> +struct KernelVecType { + using qk_load_vec_type = vec_op::BF16Vec16; + using qk_vec_type = vec_op::FP32Vec16; + using v_load_vec_type = vec_op::BF16Vec16; +}; + #elif defined(__aarch64__) template <> struct KernelVecType { From 7439e4f41b1f877e088d1b8c9ce4f2847c59423f Mon Sep 17 00:00:00 2001 From: Kunshang Ji Date: Fri, 6 Feb 2026 13:03:59 +0800 Subject: [PATCH 123/810] [XPU][4/N] add mxfp4 moe model support (#33679) Signed-off-by: Kunshang Ji --- .../layers/quantization/mxfp4.py | 84 ++++++++++++------- 1 file changed, 53 insertions(+), 31 deletions(-) diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index 50009445d9b..b9dec453056 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -215,7 +215,7 @@ class Mxfp4Config(QuantizationConfig): return UnquantizedLinearMethod() elif isinstance(layer, FusedMoE): if current_platform.is_xpu(): - return IpexMxfp4MoEMethod(layer.moe_config) + return XpuMxfp4MoEMethod(layer.moe_config) else: quant_method = Mxfp4MoEMethod(layer.moe_config) quant_method.marlin_input_dtype = get_marlin_input_dtype(prefix) @@ -1096,7 +1096,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): raise ValueError(f"Unsupported backend: {self.mxfp4_backend}") -class IpexMxfp4MoEMethod(Mxfp4MoEMethod): +class XpuMxfp4MoEMethod(Mxfp4MoEMethod): def __init__(self, moe_config: FusedMoEConfig): super().__init__(moe_config) self.moe_config = moe_config @@ -1121,21 +1121,7 @@ class IpexMxfp4MoEMethod(Mxfp4MoEMethod): self.original_hidden_size = hidden_size def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - import intel_extension_for_pytorch as ipex - - layer.w13_weight.data = layer.w13_weight.data.view(torch.int32) - layer.w2_weight.data = layer.w2_weight.data.view(torch.int32) - ep_rank_start = self.moe_config.ep_rank * self.moe_config.num_local_experts - layer.ipex_fusion = ipex.llm.modules.GatedMLPMOE( - layer.w13_weight, - layer.w2_weight, - w1_scale_inv=layer.w13_weight_scale, - w2_scale_inv=layer.w2_weight_scale, - w13_bias=layer.w13_bias, - w2_bias=layer.w2_bias, - is_mxfp4=True, - experts_start_id=ep_rank_start, - ) + pass @property def is_monolithic(self) -> bool: @@ -1148,19 +1134,55 @@ class IpexMxfp4MoEMethod(Mxfp4MoEMethod): router_logits: torch.Tensor, ) -> torch.Tensor: assert layer.activation == "swigluoai", ( - "Only swiglu_oai activation is supported for IPEX MXFP4 MoE" + "Only swiglu_oai activation is supported for XPU MXFP4 MoE" ) - hidden_size_pad = round_up(self.original_hidden_size, 128) - x_pad = torch.nn.functional.pad(x, (0, hidden_size_pad - x.size(-1))) - hidden_states = layer.ipex_fusion( - x_pad, - layer.use_grouped_topk, - layer.top_k, - router_logits, - layer.renormalize, - layer.topk_group, - layer.num_expert_group, - activation="swiglu_oai", + from vllm_xpu_kernels.fused_moe_interface import xpu_fused_moe + + M, _ = x.size() + routing_weights = torch.empty( + M, layer.top_k, dtype=torch.float32, device=x.device + ) + selected_experts = torch.empty( + M, layer.top_k, dtype=torch.int32, device=x.device + ) + token_expert_indices = torch.empty( + M, layer.top_k, dtype=torch.int32, device=x.device + ) + + if layer.use_grouped_topk: + routing_weights, selected_experts = torch.ops._moe_C.fused_grouped_topk( + x, + router_logits, + layer.top_k, + layer.renormalize, + n_expert_group=layer.num_expert_group, + n_topk_group=layer.topk_group, + scoring_func=layer.scoring_func, + routed_scaling_factor=layer.routed_scaling_factor, + bias=layer.e_score_correction_bias, + ) + else: + torch.ops._moe_C.topk_softmax( + routing_weights, + selected_experts, + token_expert_indices, + router_logits, + layer.renormalize, + layer.e_score_correction_bias, + ) + + return xpu_fused_moe( + hidden_states=x, + w13=layer.w13_weight, + w13_bias=layer.w13_bias if self.moe.has_bias else None, + w13_scales=layer.w13_weight_scale, + w2=layer.w2_weight, + w2_bias=layer.w2_bias if self.moe.has_bias else None, + w2_scales=layer.w2_weight_scale, + topk_weights=routing_weights, + topk_ids=selected_experts, + n_experts_per_token=layer.top_k, + activation=layer.activation, + num_experts=layer.local_num_experts, + is_mxfp4=True, ) - hidden_states = hidden_states[..., : self.original_hidden_size].contiguous() - return hidden_states From 6550815c3ad5fc20e6944483dd8a7a47c18b7c7d Mon Sep 17 00:00:00 2001 From: sihao_li <165983188+1643661061leo@users.noreply.github.com> Date: Fri, 6 Feb 2026 14:02:33 +0800 Subject: [PATCH 124/810] [XPU]Replace pip in docker.xpu with uv pip (#31112) Signed-off-by: sihao.li --- docker/Dockerfile.xpu | 70 +++++++++++++++++++++++++++---------------- 1 file changed, 44 insertions(+), 26 deletions(-) diff --git a/docker/Dockerfile.xpu b/docker/Dockerfile.xpu index 04051827ba4..ba7dd848bdf 100644 --- a/docker/Dockerfile.xpu +++ b/docker/Dockerfile.xpu @@ -1,5 +1,10 @@ FROM intel/deep-learning-essentials:2025.3.2-0-devel-ubuntu24.04 AS vllm-base +WORKDIR /workspace/ + +ARG PYTHON_VERSION=3.12 +ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/xpu" + RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \ echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list && \ add-apt-repository -y ppa:kobuk-team/intel-graphics @@ -22,13 +27,16 @@ RUN apt clean && apt-get update -y && \ python3.12-dev \ python3-pip -RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.12 1 -RUN update-alternatives --install /usr/bin/python python /usr/bin/python3.12 1 - RUN apt update && apt upgrade -y && \ apt install -y libze1 libze-dev libze-intel-gpu1 intel-opencl-icd libze-intel-gpu-raytracing intel-ocloc && \ apt install -y intel-oneapi-compiler-dpcpp-cpp-2025.3 +ENV PATH="/root/.local/bin:$PATH" +ENV VIRTUAL_ENV="/opt/venv" +ENV UV_PYTHON_INSTALL_DIR=/opt/uv/python +RUN curl -LsSf https://astral.sh/uv/install.sh | sh +RUN uv venv --python ${PYTHON_VERSION} --seed ${VIRTUAL_ENV} +ENV PATH="$VIRTUAL_ENV/bin:$PATH" # This oneccl contains the BMG support which is not the case for default version of oneapi 2025.2. ARG ONECCL_INSTALLER="intel-oneccl-2021.15.7.8_offline.sh" @@ -44,20 +52,31 @@ SHELL ["bash", "-c"] CMD ["bash", "-c", "source /root/.bashrc && exec bash"] WORKDIR /workspace/vllm -COPY requirements/xpu.txt /workspace/vllm/requirements/xpu.txt -COPY requirements/common.txt /workspace/vllm/requirements/common.txt -# suppress the python externally managed environment error -RUN python3 -m pip config set global.break-system-packages true +ENV UV_HTTP_TIMEOUT=500 -RUN --mount=type=cache,target=/root/.cache/pip \ - pip install --no-cache-dir \ - -r requirements/xpu.txt +# Configure package index for XPU +ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL} +ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL} +ENV UV_INDEX_STRATEGY="unsafe-best-match" +ENV UV_LINK_MODE="copy" -# arctic-inference is built from source which needs torch-xpu properly installed -# used for suffix method speculative decoding -RUN --mount=type=cache,target=/root/.cache/pip \ - pip install --no-cache-dir arctic-inference==0.1.1 +RUN --mount=type=cache,target=/root/.cache/uv \ + --mount=type=bind,src=requirements/common.txt,target=/workspace/vllm/requirements/common.txt \ + --mount=type=bind,src=requirements/xpu.txt,target=/workspace/vllm/requirements/xpu.txt \ + uv pip install --upgrade pip && \ + uv pip install -r requirements/xpu.txt + + # used for suffix method speculative decoding + # build deps for proto + nanobind-based extensions to set up the build environment +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install grpcio-tools protobuf nanobind + # arctic-inference is built from source which needs torch-xpu properly installed first +RUN --mount=type=cache,target=/root/.cache/uv \ + source /opt/intel/oneapi/setvars.sh --force && \ + source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force && \ + export CMAKE_PREFIX_PATH="$(python -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}" && \ + uv pip install --no-build-isolation arctic-inference==0.1.1 ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/lib/" @@ -69,33 +88,32 @@ RUN --mount=type=bind,source=.git,target=.git \ ENV VLLM_TARGET_DEVICE=xpu ENV VLLM_WORKER_MULTIPROC_METHOD=spawn -RUN --mount=type=cache,target=/root/.cache/pip \ +RUN --mount=type=cache,target=/root/.cache/uv \ --mount=type=bind,source=.git,target=.git \ - pip install --no-build-isolation . + uv pip install --no-build-isolation . CMD ["/bin/bash"] FROM vllm-base AS vllm-openai # install additional dependencies for openai api server -RUN --mount=type=cache,target=/root/.cache/pip \ - pip install accelerate hf_transfer pytest pytest_asyncio lm_eval[api] modelscope +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install accelerate hf_transfer pytest pytest_asyncio lm_eval[api] modelscope # install development dependencies (for testing) -RUN python3 -m pip install -e tests/vllm_test_utils +RUN uv pip install -e tests/vllm_test_utils # install nixl from source code ENV NIXL_VERSION=0.7.0 -RUN python3 /workspace/vllm/tools/install_nixl_from_source_ubuntu.py +RUN python /workspace/vllm/tools/install_nixl_from_source_ubuntu.py # FIX triton -RUN --mount=type=cache,target=/root/.cache/pip pip uninstall triton triton-xpu -y && pip install triton-xpu==3.6.0 --extra-index-url=https://download.pytorch.org/whl/xpu - -# PyJWT-2.7.0 will influence some wheel behaviors, remove its dist-info to avoid conflicts -RUN rm /usr/lib/python3/dist-packages/PyJWT-2.7.0.dist-info/ -rf +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip uninstall triton triton-xpu && \ + uv pip install triton-xpu==3.6.0 # remove torch bundled oneccl to avoid conflicts -RUN --mount=type=cache,target=/root/.cache/pip \ - pip uninstall oneccl oneccl-devel -y +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip uninstall oneccl oneccl-devel ENTRYPOINT ["vllm", "serve"] From 965525667b70dc23463d57295dce792eba1ac452 Mon Sep 17 00:00:00 2001 From: chengchengpei <5881383+chengchengpei@users.noreply.github.com> Date: Thu, 5 Feb 2026 22:23:34 -0800 Subject: [PATCH 125/810] Onboard voyage-4-nano (#33720) Signed-off-by: Chengcheng Pei Signed-off-by: chengchengpei <5881383+chengchengpei@users.noreply.github.com> Co-authored-by: chengchengpei <5881383+chengchengpei@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- docs/models/supported_models.md | 1 + .../language/pooling_mteb_test/test_voyage.py | 56 ++++++++ tests/models/registry.py | 3 + vllm/config/model.py | 4 + vllm/model_executor/models/config.py | 8 ++ vllm/model_executor/models/qwen3.py | 12 +- vllm/model_executor/models/registry.py | 4 + vllm/model_executor/models/voyage.py | 130 ++++++++++++++++++ 8 files changed, 216 insertions(+), 2 deletions(-) create mode 100644 tests/models/language/pooling_mteb_test/test_voyage.py create mode 100644 vllm/model_executor/models/voyage.py diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index e69f68feedc..1e6776faa9c 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -519,6 +519,7 @@ These models primarily support the [`LLM.embed`](./pooling_models.md#llmembed) A | `LlamaModel`C, `LlamaForCausalLM`C, `MistralModel`C, etc. | Llama-based | `intfloat/e5-mistral-7b-instruct`, etc. | ✅︎ | ✅︎ | | `Qwen2Model`C, `Qwen2ForCausalLM`C | Qwen2-based | `ssmits/Qwen2-7B-Instruct-embed-base` (see note), `Alibaba-NLP/gte-Qwen2-7B-instruct` (see note), etc. | ✅︎ | ✅︎ | | `Qwen3Model`C, `Qwen3ForCausalLM`C | Qwen3-based | `Qwen/Qwen3-Embedding-0.6B`, etc. | ✅︎ | ✅︎ | +| `VoyageQwen3BidirectionalEmbedModel`C | Voyage Qwen3-based with bidirectional attention | `voyageai/voyage-4-nano`, etc. | ✅︎ | ✅︎ | | `RobertaModel`, `RobertaForMaskedLM` | RoBERTa-based | `sentence-transformers/all-roberta-large-v1`, etc. | | | | `*Model`C, `*ForCausalLM`C, etc. | Generative models | N/A | \* | \* | diff --git a/tests/models/language/pooling_mteb_test/test_voyage.py b/tests/models/language/pooling_mteb_test/test_voyage.py new file mode 100644 index 00000000000..99ef1de9adf --- /dev/null +++ b/tests/models/language/pooling_mteb_test/test_voyage.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from tests.models.language.pooling.embed_utils import correctness_test_embed_models +from tests.models.utils import EmbedModelInfo + +from .mteb_embed_utils import mteb_test_embed_models + +MODELS = [ + EmbedModelInfo( + "voyageai/voyage-4-nano", + architecture="VoyageQwen3BidirectionalEmbedModel", + enable_test=True, + seq_pooling_type="MEAN", + attn_type="encoder_only", + is_prefix_caching_supported=False, + is_chunked_prefill_supported=False, + hf_overrides={ + "architectures": ["VoyageQwen3BidirectionalEmbedModel"], + "num_labels": 2048, + }, + mteb_score=0.7054, + # === MTEB Results === + # STS12: 0.6613 + # STS13: 0.6906 + # STS14: 0.6556 + # STS15: 0.7843 + # STS16: 0.7340 + # STSBenchmark: 0.7063 + # Average score: 0.7054 + ), +] + + +@pytest.mark.parametrize("model_info", MODELS) +def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None: + # Encoder-only attention models need enforce_eager=True to avoid + # CUDA graph capture issues with piecewise compilation + mteb_test_embed_models( + hf_runner, vllm_runner, model_info, vllm_extra_kwargs={"enforce_eager": True} + ) + + +@pytest.mark.parametrize("model_info", MODELS) +def test_embed_models_correctness( + hf_runner, vllm_runner, model_info: EmbedModelInfo, example_prompts +) -> None: + correctness_test_embed_models( + hf_runner, + vllm_runner, + model_info, + example_prompts, + vllm_extra_kwargs={"enforce_eager": True}, + ) diff --git a/tests/models/registry.py b/tests/models/registry.py index ffa4f52f138..69da8c7afeb 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -565,6 +565,9 @@ _EMBEDDING_EXAMPLE_MODELS = { ), "RobertaModel": _HfExamplesInfo("sentence-transformers/stsb-roberta-base-v2"), "RobertaForMaskedLM": _HfExamplesInfo("sentence-transformers/all-roberta-large-v1"), + "VoyageQwen3BidirectionalEmbedModel": _HfExamplesInfo( + "voyageai/voyage-4-nano", trust_remote_code=True + ), "XLMRobertaModel": _HfExamplesInfo("intfloat/multilingual-e5-small"), "BertSpladeSparseEmbeddingModel": _HfExamplesInfo( "naver/splade-v3", diff --git a/vllm/config/model.py b/vllm/config/model.py index 86b48418180..7c0b3344319 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -1513,6 +1513,10 @@ class ModelConfig: @property def embedding_size(self): + # Check for embedding_size set by model config (e.g., Voyage models) + override = getattr(self.hf_config, "embedding_size", None) + if override is not None: + return override dense_modules = try_get_dense_modules(self.model, revision=self.revision) if dense_modules is not None: return dense_modules[-1]["out_features"] diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index c41f5e18b1c..a6c244b6e18 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -582,6 +582,13 @@ class NemotronHForCausalLMConfig(VerifyAndUpdateConfig): cache_config.mamba_ssm_cache_dtype = mamba_ssm_cache_dtype +class VoyageQwen3BidirectionalEmbedModelConfig(VerifyAndUpdateConfig): + @staticmethod + def verify_and_update_model_config(model_config: "ModelConfig") -> None: + model_config.hf_config.is_causal = False + model_config.hf_config.embedding_size = model_config.hf_config.num_labels + + MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = { "GteModel": SnowflakeGteNewModelConfig, "GteNewModel": GteNewModelConfig, @@ -604,4 +611,5 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = { "DeepseekV32ForCausalLM": DeepseekV32ForCausalLM, "NemotronHForCausalLM": NemotronHForCausalLMConfig, "NemotronHPuzzleForCausalLM": NemotronHForCausalLMConfig, + "VoyageQwen3BidirectionalEmbedModel": VoyageQwen3BidirectionalEmbedModelConfig, } diff --git a/vllm/model_executor/models/qwen3.py b/vllm/model_executor/models/qwen3.py index 06df051446a..43f330eb0fc 100644 --- a/vllm/model_executor/models/qwen3.py +++ b/vllm/model_executor/models/qwen3.py @@ -34,7 +34,10 @@ 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.logger import init_logger -from vllm.model_executor.layers.attention import Attention +from vllm.model_executor.layers.attention.encoder_only_attention import ( + Attention, + EncoderOnlyAttention, +) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor @@ -115,7 +118,12 @@ class Qwen3Attention(nn.Module): rope_parameters=rope_parameters, dual_chunk_attention_config=dual_chunk_attention_config, ) - self.attn = Attention( + attn_cls = ( + EncoderOnlyAttention + if attn_type == AttentionType.ENCODER_ONLY + else Attention + ) + self.attn = attn_cls( self.num_heads, self.head_dim, self.scaling, diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 830a615ce0e..c310f6f177d 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -237,6 +237,10 @@ _EMBEDDING_MODELS = { "RobertaModel": ("roberta", "RobertaEmbeddingModel"), "TeleChatForCausalLM": ("telechat2", "TeleChat2ForCausalLM"), "TeleChat2ForCausalLM": ("telechat2", "TeleChat2ForCausalLM"), + "VoyageQwen3BidirectionalEmbedModel": ( + "voyage", + "VoyageQwen3BidirectionalEmbedModel", + ), "XLMRobertaModel": ("roberta", "RobertaEmbeddingModel"), "BgeM3EmbeddingModel": ("roberta", "BgeM3EmbeddingModel"), # [Multimodal] diff --git a/vllm/model_executor/models/voyage.py b/vllm/model_executor/models/voyage.py new file mode 100644 index 00000000000..0713b128ce9 --- /dev/null +++ b/vllm/model_executor/models/voyage.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from __future__ import annotations + +import re +from collections import defaultdict +from collections.abc import Iterable + +import torch +import torch.nn as nn + +from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.models.qwen3 import Qwen3Model +from vllm.model_executor.models.utils import WeightsMapper + +WeightItem = tuple[str, torch.Tensor] + +_LAYER_RE = re.compile(r"^layers\.(\d+)\.(.+)$") + + +class VoyageQwen3BidirectionalEmbedModel(Qwen3Model): + """ + Qwen3Model + Voyage embedding head + bidirectional attention. + + Checkpoint conventions (HF): + - MLP: gate_proj + up_proj (unfused) + - Attn: q_proj + k_proj + v_proj (unfused) + - Linear head: linear.weight + - Weights prefixed with "model." (e.g., model.layers.0...) + + vLLM Qwen3Model expects: + - mlp.gate_up_proj (fused) + - self_attn.qkv_proj (fused) + - No "model." prefix + + We remap/fuse weights using generator pipeline and load directly + (bypassing parent's stacked_params_mapping which would cause + double-transformation like qkv_proj -> qkqkv_proj). + """ + + hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"model.": ""}) + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # Embedding head (hidden_size -> num_labels, bias=False) + self.linear = nn.Linear( + self.config.hidden_size, + self.config.num_labels, + bias=False, + ) + + def forward(self, *args, **kwargs): + out = super().forward(*args, **kwargs) + return self.linear(out) + + def _fuse_qkv_proj(self, weights: Iterable[WeightItem]) -> Iterable[WeightItem]: + """Fuse q_proj, k_proj, v_proj into qkv_proj.""" + qkv_buf: dict[int, dict[str, torch.Tensor]] = defaultdict(dict) + qkv_suffixes = { + "self_attn.q_proj.weight": "q", + "self_attn.k_proj.weight": "k", + "self_attn.v_proj.weight": "v", + } + + for name, tensor in weights: + m = _LAYER_RE.match(name) + if m and m.group(2) in qkv_suffixes: + layer_idx = int(m.group(1)) + qkv_buf[layer_idx][qkv_suffixes[m.group(2)]] = tensor + else: + yield name, tensor + + # Yield fused QKV weights + for layer_idx in sorted(qkv_buf.keys()): + parts = qkv_buf[layer_idx] + if all(p in parts for p in ("q", "k", "v")): + fused = torch.cat([parts["q"], parts["k"], parts["v"]], dim=0) + yield f"layers.{layer_idx}.self_attn.qkv_proj.weight", fused + elif parts: + missing = [p for p in ("q", "k", "v") if p not in parts] + raise ValueError(f"Layer {layer_idx} missing QKV parts: {missing}") + + def _fuse_gate_up_proj(self, weights: Iterable[WeightItem]) -> Iterable[WeightItem]: + """Fuse gate_proj and up_proj into gate_up_proj.""" + mlp_buf: dict[int, dict[str, torch.Tensor]] = defaultdict(dict) + mlp_suffixes = { + "mlp.gate_proj.weight": "gate", + "mlp.up_proj.weight": "up", + } + + for name, tensor in weights: + m = _LAYER_RE.match(name) + if m and m.group(2) in mlp_suffixes: + layer_idx = int(m.group(1)) + mlp_buf[layer_idx][mlp_suffixes[m.group(2)]] = tensor + else: + yield name, tensor + + # Yield fused gate_up weights + for layer_idx in sorted(mlp_buf.keys()): + parts = mlp_buf[layer_idx] + if all(p in parts for p in ("gate", "up")): + fused = torch.cat([parts["gate"], parts["up"]], dim=0) + yield f"layers.{layer_idx}.mlp.gate_up_proj.weight", fused + elif parts: + missing = [p for p in ("gate", "up") if p not in parts] + raise ValueError(f"Layer {layer_idx} missing MLP parts: {missing}") + + def load_weights(self, weights: Iterable[WeightItem]) -> set[str]: + """Remap, fuse, and load weights using generator pipeline.""" + # Chain weight transformations + weights = self.hf_to_vllm_mapper.apply(weights) + weights = self._fuse_qkv_proj(weights) + weights = self._fuse_gate_up_proj(weights) + + # Load weights directly into model parameters + # (bypass parent's stacked_params_mapping) + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + + for name, loaded_weight in weights: + if name not in params_dict: + 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 From 1363e3d6d5659b58376fa5284afc2c8be548cc9d Mon Sep 17 00:00:00 2001 From: Gassan Salama Date: Fri, 6 Feb 2026 07:01:48 +0000 Subject: [PATCH 126/810] [cpu][performance] CPU Paged Attention NEON BFMMLA BF16 Implementation (#32263) Signed-off-by: Gassan --- csrc/cpu/cpu_attn_impl.hpp | 3 +- csrc/cpu/cpu_attn_neon.hpp | 19 +- csrc/cpu/cpu_attn_neon_bfmmla.hpp | 682 +++++++++++++++++++++++++ vllm/v1/attention/backends/cpu_attn.py | 4 +- 4 files changed, 704 insertions(+), 4 deletions(-) create mode 100644 csrc/cpu/cpu_attn_neon_bfmmla.hpp diff --git a/csrc/cpu/cpu_attn_impl.hpp b/csrc/cpu/cpu_attn_impl.hpp index d2479e11844..89cf2dc3a4f 100644 --- a/csrc/cpu/cpu_attn_impl.hpp +++ b/csrc/cpu/cpu_attn_impl.hpp @@ -1107,7 +1107,8 @@ class AttentionMainLoop { if (sliding_window_left != -1) { pos = std::max(pos, curr_token_pos - sliding_window_left); } - return pos; + // Clamp to tile end to avoid OOB when window starts past the tile + return std::min(pos, kv_tile_end_pos); }(); int32_t right_kv_pos = [&]() { diff --git a/csrc/cpu/cpu_attn_neon.hpp b/csrc/cpu/cpu_attn_neon.hpp index 827f0cfbc71..3523893c38c 100644 --- a/csrc/cpu/cpu_attn_neon.hpp +++ b/csrc/cpu/cpu_attn_neon.hpp @@ -4,6 +4,9 @@ #include "cpu_attn_impl.hpp" #include #include +#ifdef ARM_BF16_SUPPORT + #include "cpu_attn_neon_bfmmla.hpp" +#endif namespace cpu_attention { namespace { @@ -57,7 +60,7 @@ FORCE_INLINE void load_row8_B_as_f32(const c10::BFloat16* p, #endif } -// Mx8, with 1 <= M <= 8 , K streamed, unroll-by-4 with NEON FMLAs +// Mx8, with 1 <= M <= 8 , K streamed, unroll-by-4 with ASIMD FMLAs // #Loads = (K // 4) * (M + 4 * sizeof(kv_cache_t) / 2) // #FMLAs = (K // 4) * (4 * 2 * M) // We have (4 * 2 * M) FMLAs for (M + 4 * sizeof(kv_cache_t) / 2) loads @@ -381,6 +384,18 @@ class AttentionImpl { } } }; + +#ifdef ARM_BF16_SUPPORT +// For BF16 on Arm, reuse the BFMMLA kernels with 32-token alignment. +template +class AttentionImpl + : public AttentionImplNEONBFMMLA {}; +#endif } // namespace cpu_attention -#endif // #ifndef CPU_ATTN_NEON_HPP +#undef BLOCK_SIZE_ALIGNMENT +#undef HEAD_SIZE_ALIGNMENT +#undef MAX_Q_HEAD_NUM_PER_ITER + +#endif // #ifndef CPU_ATTN_ASIMD_HPP diff --git a/csrc/cpu/cpu_attn_neon_bfmmla.hpp b/csrc/cpu/cpu_attn_neon_bfmmla.hpp new file mode 100644 index 00000000000..fb133aa1309 --- /dev/null +++ b/csrc/cpu/cpu_attn_neon_bfmmla.hpp @@ -0,0 +1,682 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +#ifndef CPU_ATTN_NEON_BFMMLA_HPP +#define CPU_ATTN_NEON_BFMMLA_HPP + +#include "cpu_attn_impl.hpp" + +#include + +#include +#include + +namespace cpu_attention { + +namespace { + +// BFMMLA tile dimensions +constexpr int32_t TILE_ROWS = 2; // M dimension +constexpr int32_t TILE_K = 4; // K reduction +constexpr int32_t TILE_COLS = 2; // N dimension (column-pair) + +// Derived constants +constexpr int32_t OUTPUT_COLS_PER_BLOCK = 8; // 4 column-pairs +constexpr int32_t K_TOKENS_PER_GROUP = 8; // Tokens grouped in K cache +constexpr int32_t V_TOKENS_PER_ROW_BLOCK = 4; // Tokens per V cache row block +constexpr int32_t K_INNER_STRIDE = K_TOKENS_PER_GROUP * TILE_K; +constexpr int32_t V_INNER_STRIDE = V_TOKENS_PER_ROW_BLOCK * TILE_COLS; +constexpr int32_t PACK_ELEMENTS_PER_K_CHUNK = TILE_ROWS * TILE_K; // A packing + +// Matrix Packing and Accumulator +// Reshape two rows of Q into BFMMLA-friendly interleaved +// Input: row0 = [a0,a1,a2,a3], row1 = [b0,b1,b2,b3] +// Output: [a0,a1,a2,a3,b0,b1,b2,b3, a4,a5,a6,a7,b4,b5,b6,b7] +// For K tail (K % TILE_K != 0): pads with zeros to complete the final chunk +FORCE_INLINE void reshape_Q_2xK_for_bfmmla(const c10::BFloat16* __restrict r0, + const c10::BFloat16* __restrict r1, + c10::BFloat16* __restrict dst, + int32_t K) { + const uint16_t* s0 = reinterpret_cast(r0); + const uint16_t* s1 = reinterpret_cast(r1); + uint16_t* d = reinterpret_cast(dst); + + // Process TILE_K elements at a time (PACK_ELEMENTS_PER_K_CHUNK output) + int32_t k = 0; + for (; k + TILE_K <= K; k += TILE_K, d += PACK_ELEMENTS_PER_K_CHUNK) { + vst1q_u16(d, vcombine_u16(vld1_u16(s0 + k), vld1_u16(s1 + k))); + } + + // Handle K tail: pack remaining elements with zero-padding + const int32_t tail = K - k; + if (tail > 0) { + // Pack remaining tail elements: [r0[k..k+tail-1], pad, r1[k..k+tail-1], + // pad] + for (int32_t t = 0; t < tail; ++t) { + d[t] = s0[k + t]; + d[t + TILE_K] = s1[k + t]; + } + // Zero-pad the rest + for (int32_t t = tail; t < TILE_K; ++t) { + d[t] = 0; + d[t + TILE_K] = 0; + } + } +} + +// 2x2 accumulator load/store with compile-time row count +template +FORCE_INLINE float32x4_t load_acc_2x2(float* base, int64_t ldc, int col_off) { + static_assert(m_rows == 1 || m_rows == 2); + float32x2_t row0 = vld1_f32(base + col_off); + float32x2_t row1 = + (m_rows == 2) ? vld1_f32(base + ldc + col_off) : vdup_n_f32(0.f); + return vcombine_f32(row0, row1); +} + +template +FORCE_INLINE void store_acc_2x2(float32x4_t acc, float* base, int64_t ldc, + int col_off) { + static_assert(m_rows == 1 || m_rows == 2); + vst1_f32(base + col_off, vget_low_f32(acc)); + if constexpr (m_rows == 2) { + vst1_f32(base + ldc + col_off, vget_high_f32(acc)); + } +} + +// Initialize 4 column-pair accumulators for 2 rows (8 columns total) +#define INIT_ACC_ROWPAIR_4(a0, a1, a2, a3, Crow, ldc, m_rows, accum) \ + do { \ + if (accum) { \ + if (m_rows == 2) { \ + a0 = load_acc_2x2<2>(Crow, ldc, 0); \ + a1 = load_acc_2x2<2>(Crow, ldc, 2); \ + a2 = load_acc_2x2<2>(Crow, ldc, 4); \ + a3 = load_acc_2x2<2>(Crow, ldc, 6); \ + } else { \ + a0 = load_acc_2x2<1>(Crow, ldc, 0); \ + a1 = load_acc_2x2<1>(Crow, ldc, 2); \ + a2 = load_acc_2x2<1>(Crow, ldc, 4); \ + a3 = load_acc_2x2<1>(Crow, ldc, 6); \ + } \ + } else { \ + a0 = a1 = a2 = a3 = vdupq_n_f32(0.f); \ + } \ + } while (0) + +// Store 4 column-pair accumulators back to C matrix +#define STORE_ACC_ROWPAIR_4(a0, a1, a2, a3, Crow, ldc, m_rows) \ + do { \ + if (m_rows == 2) { \ + store_acc_2x2<2>(a0, Crow, ldc, 0); \ + store_acc_2x2<2>(a1, Crow, ldc, 2); \ + store_acc_2x2<2>(a2, Crow, ldc, 4); \ + store_acc_2x2<2>(a3, Crow, ldc, 6); \ + } else { \ + store_acc_2x2<1>(a0, Crow, ldc, 0); \ + store_acc_2x2<1>(a1, Crow, ldc, 2); \ + store_acc_2x2<1>(a2, Crow, ldc, 4); \ + store_acc_2x2<1>(a3, Crow, ldc, 6); \ + } \ + } while (0) + +// Perform 4 BFMMLA operations: acc += A @ B for 4 column-pairs +#define BFMMLA_COMPUTE_4(r0, r1, r2, r3, a, b0, b1, b2, b3) \ + do { \ + r0 = vbfmmlaq_f32(r0, a, b0); \ + r1 = vbfmmlaq_f32(r1, a, b1); \ + r2 = vbfmmlaq_f32(r2, a, b2); \ + r3 = vbfmmlaq_f32(r3, a, b3); \ + } while (0) + +// Micro-kernel: updates a small fixed tile using BFMMLA. +// RP = number of row-pairs (1,2,4) +// Computes C[TILE_ROWS*RP, OUTPUT_COLS_PER_BLOCK] += A_packed @ B. +// A_packed interleaves RP row-pairs; B layout is driven by the attention phase: +// - AttentionGemmPhase::QK -> token-column layout (Q @ K^T) +// - AttentionGemmPhase::PV -> token-row layout (P @ V) +// K_static < 0 enables runtime K (PV only) +template +FORCE_INLINE void gemm_rowpairs_x8_bfmmla_neon( + const bfloat16_t* const* __restrict A_packed_rp, + const int32_t* __restrict m_rows_rp, const bfloat16_t* __restrict B_blk, + float* __restrict C, int64_t ldc, bool accumulate, int64_t b_stride, + int32_t K_runtime = 0) { + static_assert(RP == 1 || RP == 2 || RP == 4, "RP must be 1,2,4"); + static_assert(K_static < 0 || K_static % TILE_K == 0, + "K must be divisible by TILE_K"); + static_assert(K_static >= 0 || phase == AttentionGemmPhase::PV, + "Runtime K only supported for PV"); + + constexpr bool runtime_k = (K_static < 0); + const int32_t K_iters = + runtime_k ? (K_runtime / TILE_K) : (K_static / TILE_K); + const int32_t K_tail = runtime_k ? (K_runtime % TILE_K) : 0; + + if (!runtime_k) { + // Help the compiler fold away unused K_runtime when K is compile-time + (void)K_runtime; + } + + auto* C_al = C; + const auto* B_al = B_blk; + + // Setup A pointers + const bfloat16_t* a_ptr[4] = { + A_packed_rp[0], + (RP >= 2) ? A_packed_rp[1] : nullptr, + (RP >= 4) ? A_packed_rp[2] : nullptr, + (RP >= 4) ? A_packed_rp[3] : nullptr, + }; + + // Setup B pointers based on layout + const bfloat16_t* b_ptr[4]; + if constexpr (phase == AttentionGemmPhase::PV) { + b_ptr[0] = B_blk + 0 * b_stride; + b_ptr[1] = B_blk + 1 * b_stride; + b_ptr[2] = B_blk + 2 * b_stride; + b_ptr[3] = B_blk + 3 * b_stride; + } + + float32x4_t acc[4][4]; + +// Initialize accumulators +#define INIT_RP(rp) \ + if constexpr (RP > rp) { \ + INIT_ACC_ROWPAIR_4(acc[rp][0], acc[rp][1], acc[rp][2], acc[rp][3], \ + C_al + (rp * 2) * ldc, ldc, m_rows_rp[rp], accumulate); \ + } + INIT_RP(0); + INIT_RP(1); + INIT_RP(2); + INIT_RP(3); +#undef INIT_RP + + // Main compute loop + for (int32_t ki = 0; ki < K_iters; ++ki) { + bfloat16x8_t b0, b1, b2, b3; + if constexpr (phase == AttentionGemmPhase::PV) { + b0 = vld1q_bf16(b_ptr[0] + ki * V_INNER_STRIDE); + b1 = vld1q_bf16(b_ptr[1] + ki * V_INNER_STRIDE); + b2 = vld1q_bf16(b_ptr[2] + ki * V_INNER_STRIDE); + b3 = vld1q_bf16(b_ptr[3] + ki * V_INNER_STRIDE); + } else { + const bfloat16_t* b_base = B_al + ki * b_stride; + b0 = vld1q_bf16(b_base + 0 * V_INNER_STRIDE); + b1 = vld1q_bf16(b_base + 1 * V_INNER_STRIDE); + b2 = vld1q_bf16(b_base + 2 * V_INNER_STRIDE); + b3 = vld1q_bf16(b_base + 3 * V_INNER_STRIDE); + } + +#define COMPUTE_RP(rp) \ + if constexpr (RP > rp) { \ + bfloat16x8_t a = vld1q_bf16(a_ptr[rp] + ki * PACK_ELEMENTS_PER_K_CHUNK); \ + BFMMLA_COMPUTE_4(acc[rp][0], acc[rp][1], acc[rp][2], acc[rp][3], a, b0, \ + b1, b2, b3); \ + } + COMPUTE_RP(0); + COMPUTE_RP(1); + COMPUTE_RP(2); + COMPUTE_RP(3); +#undef COMPUTE_RP + } + + // K tail for runtime PV: fallback path + if constexpr (runtime_k) { + if (K_tail > 0) { + const int32_t tail_offset = K_iters * V_INNER_STRIDE; + const int32_t a_tail_offset = K_iters * PACK_ELEMENTS_PER_K_CHUNK; + for (int32_t kt = 0; kt < K_tail; ++kt) { + float32x4_t b_vecs[4]; + for (int32_t p = 0; p < 4; ++p) { + const bfloat16_t* bp = b_ptr[p] + tail_offset + kt * TILE_COLS; + const float b0 = vcvtah_f32_bf16(bp[0]); + const float b1 = vcvtah_f32_bf16(bp[1]); + const float32x2_t b_pair = vset_lane_f32(b1, vdup_n_f32(b0), 1); + b_vecs[p] = vcombine_f32(b_pair, b_pair); + } + +#define TAIL_RP(rp) \ + if constexpr (RP > rp) { \ + const bfloat16_t* ap = A_packed_rp[rp] + a_tail_offset; \ + float a_row0 = vcvtah_f32_bf16(ap[kt]); \ + float a_row1 = \ + (m_rows_rp[rp] == 2) ? vcvtah_f32_bf16(ap[kt + TILE_K]) : 0.0f; \ + const float32x4_t a_vec = \ + vcombine_f32(vdup_n_f32(a_row0), vdup_n_f32(a_row1)); \ + for (int32_t p = 0; p < 4; ++p) { \ + acc[rp][p] = vmlaq_f32(acc[rp][p], a_vec, b_vecs[p]); \ + } \ + } + TAIL_RP(0); + TAIL_RP(1); + TAIL_RP(2); + TAIL_RP(3); +#undef TAIL_RP + } + } + } + + // Store results +#define STORE_RP(rp) \ + if constexpr (RP > rp) { \ + STORE_ACC_ROWPAIR_4(acc[rp][0], acc[rp][1], acc[rp][2], acc[rp][3], \ + C_al + (rp * 2) * ldc, ldc, m_rows_rp[rp]); \ + } + STORE_RP(0); + STORE_RP(1); + STORE_RP(2); + STORE_RP(3); +#undef STORE_RP +} + +// Meso-kernel: packs a small MBxK slice of A, then tiles over N and calls the +// micro-kernel for each OUTPUT_COLS_PER_BLOCK chunk. K_static < 0 enables +// runtime K (PV only). +template +FORCE_INLINE void gemm_packA_compute_MB_xN( + const c10::BFloat16* __restrict A, const c10::BFloat16* __restrict B, + float* __restrict C, int32_t K_runtime, int64_t lda, int64_t ldc, + int64_t b_layout_stride, int64_t b_reduction_stride, bool accumulate) { + static_assert(MB >= 1 && MB <= 8, "MB must be in [1,8]"); + static_assert(N % OUTPUT_COLS_PER_BLOCK == 0, + "N must be a multiple of OUTPUT_COLS_PER_BLOCK"); + static_assert(K_static < 0 || K_static % TILE_K == 0, + "K must be divisible by TILE_K"); + static_assert(K_static >= 0 || phase == AttentionGemmPhase::PV, + "Runtime K only supported for PV"); + + constexpr bool runtime_k = (K_static < 0); + const int32_t K_val = runtime_k ? K_runtime : K_static; + + // Keep small packs on-stack to avoid heap churn + constexpr int32_t STACK_PACK_STRIDE = + (1024 / TILE_K) * PACK_ELEMENTS_PER_K_CHUNK; + + constexpr int32_t ROW_PAIRS = (MB + 1) / TILE_ROWS; + const int32_t pack_stride = + runtime_k ? ((K_val + TILE_K - 1) / TILE_K) * PACK_ELEMENTS_PER_K_CHUNK + : (K_static / TILE_K) * PACK_ELEMENTS_PER_K_CHUNK; + + alignas(64) c10::BFloat16 A_packed_stack[ROW_PAIRS * STACK_PACK_STRIDE]; + std::vector A_packed_heap; + c10::BFloat16* A_packed = + (pack_stride <= STACK_PACK_STRIDE) + ? A_packed_stack + : (A_packed_heap.resize(ROW_PAIRS * pack_stride), + A_packed_heap.data()); + + for (int32_t rp = 0; rp < ROW_PAIRS; ++rp) { + const int32_t m = rp * TILE_ROWS; + const int32_t m_rows = (m + 1 < MB) ? TILE_ROWS : 1; + const c10::BFloat16* A0 = A + m * lda; + const c10::BFloat16* A1 = (m_rows == TILE_ROWS) ? (A + (m + 1) * lda) : A0; + reshape_Q_2xK_for_bfmmla(A0, A1, A_packed + rp * pack_stride, K_val); + } + + for (int32_t n = 0; n < N; n += OUTPUT_COLS_PER_BLOCK) { + const c10::BFloat16* B_blk_c10 = + (phase == AttentionGemmPhase::PV) + ? (B + (n / TILE_COLS) * b_layout_stride) + : (B + (n / OUTPUT_COLS_PER_BLOCK) * b_layout_stride); + const bfloat16_t* B_blk = reinterpret_cast(B_blk_c10); + + // Process row-pairs in groups of 4, 2, then 1 + int32_t row_pair_idx = 0; + +#define PROCESS_RP_GROUP(group_size) \ + for (; row_pair_idx + (group_size - 1) < ROW_PAIRS; \ + row_pair_idx += group_size) { \ + const bfloat16_t* Ap[group_size]; \ + int32_t mr[group_size]; \ + for (int32_t i = 0; i < group_size; ++i) { \ + Ap[i] = reinterpret_cast( \ + A_packed + (row_pair_idx + i) * pack_stride); \ + mr[i] = (((row_pair_idx + i) * TILE_ROWS + 1) < MB) ? TILE_ROWS : 1; \ + } \ + float* C_blk = C + (row_pair_idx * TILE_ROWS) * ldc + n; \ + if constexpr (runtime_k) { \ + gemm_rowpairs_x8_bfmmla_neon( \ + Ap, mr, B_blk, C_blk, ldc, accumulate, b_layout_stride, K_val); \ + } else { \ + gemm_rowpairs_x8_bfmmla_neon( \ + Ap, mr, B_blk, C_blk, ldc, accumulate, \ + (phase == AttentionGemmPhase::PV) ? b_layout_stride \ + : b_reduction_stride); \ + } \ + } + + PROCESS_RP_GROUP(4); + PROCESS_RP_GROUP(2); + PROCESS_RP_GROUP(1); +#undef PROCESS_RP_GROUP + } +} + +// Macro-kernel: iterates over M in MB={8,4,2,1} chunks. +// Supports compile-time K specialization when K >= 0; otherwise uses runtime K +// (runtime K path is only supported for PV). +template +FORCE_INLINE void gemm_macro_neon_bfmmla( + const c10::BFloat16* __restrict A, const c10::BFloat16* __restrict B, + float* __restrict C, int32_t M, int32_t K_runtime, int64_t lda, int64_t ldc, + int64_t b_layout_stride, int64_t b_reduction_stride, bool accumulate) { + static_assert(N % OUTPUT_COLS_PER_BLOCK == 0, + "N must be a multiple of OUTPUT_COLS_PER_BLOCK"); + + if constexpr (K >= 0) { + static_assert(K % TILE_K == 0, "K must be divisible by TILE_K"); + for (int32_t m = 0; m < M;) { + const int32_t rem = M - m; + const c10::BFloat16* A_blk = A + m * lda; + float* C_blk = C + m * ldc; + +#define DISPATCH_MB(mb) \ + gemm_packA_compute_MB_xN(A_blk, B, C_blk, 0, lda, ldc, \ + b_layout_stride, \ + b_reduction_stride, accumulate) + + if (rem >= 8) { + DISPATCH_MB(8); + m += 8; + } else if (rem >= 4) { + DISPATCH_MB(4); + m += 4; + } else if (rem >= 2) { + DISPATCH_MB(2); + m += 2; + } else { + DISPATCH_MB(1); + m += 1; + } +#undef DISPATCH_MB + } + } else { + static_assert(phase == AttentionGemmPhase::PV, + "Runtime K specialization only supported for PV."); + const int32_t K_val = K_runtime; + + for (int32_t m = 0; m < M;) { + const int32_t rem = M - m; + const c10::BFloat16* A_blk = A + m * lda; + float* C_blk = C + m * ldc; + +#define DISPATCH_MB_RUNTIME(mb) \ + gemm_packA_compute_MB_xN(A_blk, B, C_blk, K_val, lda, ldc, \ + b_layout_stride, \ + b_reduction_stride, accumulate) + + if (rem >= 8) { + DISPATCH_MB_RUNTIME(8); + m += 8; + } else if (rem >= 4) { + DISPATCH_MB_RUNTIME(4); + m += 4; + } else if (rem >= 2) { + DISPATCH_MB_RUNTIME(2); + m += 2; + } else { + DISPATCH_MB_RUNTIME(1); + m += 1; + } +#undef DISPATCH_MB_RUNTIME + } + } +} + +#undef INIT_ACC_ROWPAIR_4 +#undef STORE_ACC_ROWPAIR_4 +#undef BFMMLA_COMPUTE_4 + +} // namespace + +// TileGemm Adapter for Attention + +template +class TileGemmNEONBFMMLA { + public: + template + FORCE_INLINE static void gemm(const int32_t m_size, void* __restrict__ a_tile, + kv_cache_t* __restrict__ b_tile, + float* __restrict__ c_tile, const int64_t lda, + [[maybe_unused]] const int64_t ldb, + const int64_t ldc, + [[maybe_unused]] const int32_t block_size, + [[maybe_unused]] const int32_t dynamic_k_size, + const bool accum_c) { + static_assert(BlockTokens % OUTPUT_COLS_PER_BLOCK == 0); + // BFMMLA kernels require compile-time head_dim; keep head_dim_ct only for + // API parity with other tile_gemm implementations. + if constexpr (head_dim_ct >= 0) { + static_assert(head_dim_ct == HeadDim, + "BFMMLA expects head_dim_ct to match HeadDim; PV passes " + "-1 for API parity."); + } + + if constexpr (phase == AttentionGemmPhase::QK) { + const int64_t b_reduction_stride = K_INNER_STRIDE; + const int64_t b_token_block_stride = (HeadDim / TILE_K) * K_INNER_STRIDE; + + gemm_macro_neon_bfmmla( + reinterpret_cast(a_tile), b_tile, c_tile, + m_size, 0, lda, ldc, b_token_block_stride, b_reduction_stride, + accum_c); + } else { + const int64_t b_pair_stride = + (block_size / V_TOKENS_PER_ROW_BLOCK) * V_INNER_STRIDE; + + // PV gemm with runtime K specialization + switch (dynamic_k_size) { + case 32: + gemm_macro_neon_bfmmla( + reinterpret_cast(a_tile), b_tile, c_tile, + m_size, 32, lda, ldc, b_pair_stride, 0, accum_c); + break; + case 128: + gemm_macro_neon_bfmmla( + reinterpret_cast(a_tile), b_tile, c_tile, + m_size, 128, lda, ldc, b_pair_stride, 0, accum_c); + break; + case 256: + gemm_macro_neon_bfmmla( + reinterpret_cast(a_tile), b_tile, c_tile, + m_size, 256, lda, ldc, b_pair_stride, 0, accum_c); + break; + default: + gemm_macro_neon_bfmmla( + reinterpret_cast(a_tile), b_tile, c_tile, + m_size, dynamic_k_size, lda, ldc, b_pair_stride, 0, accum_c); + break; + } + } + } +}; + +// Shared ASIMD BFMMLA implementation (BF16 only). The block size alignment and +// ISA tag are template parameters so we can reuse the same kernels for +// different NEON configurations. +template +class AttentionImplNEONBFMMLA { + public: + using query_t = c10::BFloat16; + using q_buffer_t = c10::BFloat16; + using kv_cache_t = c10::BFloat16; + using logits_buffer_t = float; + using partial_output_buffer_t = float; + using prob_buffer_t = c10::BFloat16; + + static constexpr int64_t BlockSizeAlignment = block_size_alignment; + // HeadDimAlignment equals head_dim so that the PV phase processes + // the full head dimension in a single gemm call. + static constexpr int64_t HeadDimAlignment = head_dim; + static constexpr int64_t MaxQHeadNumPerIteration = 16; + static constexpr int64_t HeadDim = head_dim; + static constexpr ISA ISAType = isa_type; + static constexpr bool scale_on_logits = false; + + static_assert(HeadDim % OUTPUT_COLS_PER_BLOCK == 0); + static_assert(BlockSizeAlignment % OUTPUT_COLS_PER_BLOCK == 0); + static_assert(HeadDim % TILE_K == 0, "HeadDim must be a multiple of TILE_K"); + + public: + template