From 837db7605e240202c43577cfa4da65f3c8f506fb Mon Sep 17 00:00:00 2001 From: Ashish Patel Date: Thu, 18 Jun 2026 21:30:20 +0530 Subject: [PATCH 01/75] [Bugfix][Tool Parser] Handle non-finite numbers in coerce_to_schema_type (#43984) Signed-off-by: ashishpatel26 Co-authored-by: Ben Browning --- tests/tool_parsers/test_utils.py | 67 ++++++++++++++++++++++++++++++++ vllm/tool_parsers/utils.py | 38 ++++++++++++++++-- 2 files changed, 102 insertions(+), 3 deletions(-) diff --git a/tests/tool_parsers/test_utils.py b/tests/tool_parsers/test_utils.py index 592ef580a2b..3276fa9ddd2 100644 --- a/tests/tool_parsers/test_utils.py +++ b/tests/tool_parsers/test_utils.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json + import pytest from vllm.tool_parsers.utils import ( @@ -91,6 +93,71 @@ class TestCoerceToSchemaType: def test_invalid_number_fallback(self): assert coerce_to_schema_type("abc", "number") == "abc" + class TestNonFiniteNumbers: + """Non-finite numeric strings must not crash and must coerce to a + JSON-serializable value. + + Regression: ``int(float("inf"))`` raised an uncaught ``OverflowError`` + (only ``ValueError``/``TypeError`` were handled), and ``"1e999"`` + round-tripped through ``json.loads`` to a float ``inf`` that + ``json.dumps`` renders as invalid JSON ``Infinity``. + """ + + @pytest.mark.parametrize( + "value", ["inf", "-inf", "Infinity", "1e999", "nan", "-nan"] + ) + def test_non_finite_number_does_not_crash(self, value): + # Must not raise (previously OverflowError for inf/1e999/Infinity). + result = coerce_to_schema_type(value, "number") + # Result must serialize to valid, finite JSON and round-trip. + assert json.loads(json.dumps(result)) == result + + @pytest.mark.parametrize("value", ["inf", "-inf", "1e999"]) + def test_non_finite_number_preserved_as_string(self, value): + assert coerce_to_schema_type(value, "number") == value + + @pytest.mark.parametrize("value", ["inf", "1e999", "Infinity"]) + def test_non_finite_integer_not_float_inf(self, value): + result = coerce_to_schema_type(value, "integer") + assert isinstance(result, str) + assert result == value + + class TestNonFiniteContainers: + """Non-finite floats nested in object/array values must not produce + invalid JSON. + + Regression: the ``object``/``array`` branch returned + ``json.loads(value)`` directly, so ``"[1e999]"`` became ``[inf]`` and + ``'{"x": Infinity}'`` became ``{"x": inf}`` -- values that + ``json.dumps`` later renders as invalid JSON (``Infinity``/``NaN``). + """ + + @pytest.mark.parametrize( + "value", ["[1e999]", "[1, 2, 1e999]", "[NaN]", "[-Infinity]"] + ) + def test_array_with_non_finite_preserved_as_string(self, value): + result = coerce_to_schema_type(value, "array") + assert result == value + assert json.loads(json.dumps(result)) == result + + @pytest.mark.parametrize( + "value", ['{"x": 1e999}', '{"x": Infinity}', '{"a": [1e999, 2]}'] + ) + def test_object_with_non_finite_preserved_as_string(self, value): + result = coerce_to_schema_type(value, "object") + assert result == value + assert json.loads(json.dumps(result)) == result + + def test_finite_array_still_coerced(self): + assert coerce_to_schema_type("[1, 2, 3]", "array") == [1, 2, 3] + + def test_finite_object_still_coerced(self): + assert coerce_to_schema_type('{"a": 1}', "object") == {"a": 1} + + def test_unknown_type_non_finite_falls_back_to_string(self): + # Exercises the final json.loads fallback path. + assert coerce_to_schema_type("1e999", "unknown_type") == "1e999" + class TestBooleanType: def test_true(self): assert coerce_to_schema_type("true", "boolean") is True diff --git a/vllm/tool_parsers/utils.py b/vllm/tool_parsers/utils.py index 82cb16233fd..a31420cf1cd 100644 --- a/vllm/tool_parsers/utils.py +++ b/vllm/tool_parsers/utils.py @@ -3,6 +3,7 @@ import ast import json +import math import warnings from json import JSONDecodeError, JSONDecoder from typing import Any, TypeAlias @@ -145,6 +146,20 @@ def is_complete_json(input_str: str) -> bool: return False +def _is_json_finite(obj: Any) -> bool: + """Whether *obj* can be serialized to valid JSON. + + ``json.dumps(..., allow_nan=False)`` raises ``ValueError`` on any + non-finite float (``inf``/``-inf``/``nan``) anywhere in the value, so this + detects non-finite floats nested inside parsed lists/dicts too. + """ + try: + json.dumps(obj, allow_nan=False) + return True + except (ValueError, TypeError): + return False + + def consume_space(i: int, s: str) -> int: while i < len(s) and s[i].isspace(): i += 1 @@ -601,9 +616,15 @@ def coerce_to_schema_type(value: str, schema_type: str | list[str]) -> Any: if candidate_type == "number": try: val = float(value) - return val if val != int(val) else int(val) except (ValueError, TypeError): continue + if not math.isfinite(val): + # inf/-inf/nan are not valid JSON numbers. Fall through so + # the value is preserved as a string instead of crashing + # (int(float("inf")) raises OverflowError) or emitting + # invalid JSON (json.dumps(inf) -> "Infinity"). + continue + return val if val != int(val) else int(val) if candidate_type == "boolean": lower_val = value.lower().strip() if lower_val in ("true", "1"): @@ -613,14 +634,25 @@ def coerce_to_schema_type(value: str, schema_type: str | list[str]) -> Any: continue if candidate_type in ("object", "array"): try: - return json.loads(value) + parsed = json.loads(value) except (json.JSONDecodeError, ValueError, TypeError): continue + if _is_json_finite(parsed): + return parsed + # Non-finite floats (e.g. "[1e999]" -> [inf]) cannot be + # serialized back to valid JSON; preserve the raw string. + continue try: - return json.loads(value) + parsed = json.loads(value) except (json.JSONDecodeError, ValueError): return value + # Reject non-finite results (e.g. json.loads("1e999") -> inf, or nested + # inf/nan inside a parsed list/dict) which json.dumps would render as + # invalid JSON (Infinity/NaN). Preserve the raw string instead. + if not _is_json_finite(parsed): + return value + return parsed def compute_tool_delta( From 058cc0a8b6e33523b1ed75db933726959df43791 Mon Sep 17 00:00:00 2001 From: Yuwen Zhou Date: Fri, 19 Jun 2026 00:20:29 +0800 Subject: [PATCH 02/75] [Bugfix] Restore is_sym guard for zp in GPTQ/CT MoE to fix symmetric quant regression (#45656) Signed-off-by: yuwenzho --- vllm/model_executor/layers/quantization/auto_gptq.py | 10 ++++++++-- .../compressed_tensors_moe_wna16_marlin.py | 4 ++-- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/layers/quantization/auto_gptq.py b/vllm/model_executor/layers/quantization/auto_gptq.py index 459a6158327..f7fe7f6e9e4 100644 --- a/vllm/model_executor/layers/quantization/auto_gptq.py +++ b/vllm/model_executor/layers/quantization/auto_gptq.py @@ -25,6 +25,7 @@ from vllm.model_executor.layers.fused_moe import ( UnquantizedFusedMoEMethod, ) from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import ( + WNA16MoEBackend, convert_to_wna16_moe_kernel_format, make_wna16_moe_kernel, select_wna16_moe_backend, @@ -753,13 +754,18 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase): gptq_marlin_moe_quant_config, ) + # CPU fused_experts_cpu requires zero points even for symmetric quant + use_zp = ( + not self.quant_config.is_sym + or self.wna16_moe_backend == WNA16MoEBackend.CPU + ) return gptq_marlin_moe_quant_config( w1_scale=layer.w13_scales, w2_scale=layer.w2_scales, weight_bits=self.quant_config.weight_bits, group_size=self.quant_config.group_size, - w1_zp=getattr(layer, "w13_qzeros", None), - w2_zp=getattr(layer, "w2_qzeros", None), + w1_zp=getattr(layer, "w13_qzeros", None) if use_zp else None, + w2_zp=getattr(layer, "w2_qzeros", None) if use_zp else None, w1_bias=getattr(layer, "w13_bias", None), w2_bias=getattr(layer, "w2_bias", None), ) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py index a69d2a594ad..82734103917 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py @@ -415,9 +415,9 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): replace_parameter(layer, "w13_weight_scale", w13_scales) replace_parameter(layer, "w2_weight_scale", w2_scales) - if w13_qzeros is not None: + # CPU fused_experts_cpu requires zero points even for symmetric quant + if not self.symmetric or self.wna16_backend == WNA16MoEBackend.CPU: replace_parameter(layer, "w13_weight_zero_point", w13_qzeros) - if w2_qzeros is not None: replace_parameter(layer, "w2_weight_zero_point", w2_qzeros) # Marlin-specific parameters (not needed for Flashinfer) From 509947463375cc27e2a60d05ce5463f6dd059171 Mon Sep 17 00:00:00 2001 From: Rohan Potdar <66227218+Rohan138@users.noreply.github.com> Date: Thu, 18 Jun 2026 11:30:21 -0500 Subject: [PATCH 03/75] [Bugfix][ROCm] Fix rocm_aiter_per_tensor_quant custom op aliasing (#45747) Signed-off-by: Rohan138 --- tests/rocm/aiter/test_quant_op_schema.py | 145 +++++++++++++++++++++++ vllm/_aiter_ops.py | 34 ++++-- 2 files changed, 166 insertions(+), 13 deletions(-) create mode 100644 tests/rocm/aiter/test_quant_op_schema.py diff --git a/tests/rocm/aiter/test_quant_op_schema.py b/tests/rocm/aiter/test_quant_op_schema.py new file mode 100644 index 00000000000..9b2fac6e017 --- /dev/null +++ b/tests/rocm/aiter/test_quant_op_schema.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Schema/aliasing tests for the AITER FP8 quantization custom ops. +# +# These use torch.library.opcheck, whose test_schema check catches custom ops +# whose implementation aliases an input that the registered schema declares as +# non-aliasing -- the failure mode behind the rocm_aiter_per_tensor_quant +# regression (a returned scale that aliased the input scale). +# +# Skipped if AITER is not installed or the platform is not ROCm. + +import importlib.util + +import pytest +import torch + +# this import statement is needed to ensure the ops are registered +from vllm._aiter_ops import rocm_aiter_ops +from vllm.platforms import current_platform + +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", +) + +FP8_DTYPE = current_platform.fp8_dtype() + + +def _x(M=128, N=4096): + return torch.randn((M, N), dtype=torch.float16, device="cuda") + + +# The in-place per-tensor op takes the fp8 output buffer as an input, which +# opcheck's test_schema cannot exercise ("mul_cuda" is unimplemented for fp8), +# so restrict to the utils that run on fp8 inputs. The aliasing contract for +# this op is instead covered by test_per_tensor_quant_torch_compile below. +_INPLACE_OPCHECK_UTILS = ( + "test_faketensor", + "test_aot_dispatch_dynamic", + "test_autograd_registration", +) + + +def test_per_tensor_quant_static_schema(): + """Static per-tensor: caller provides scale (the aliasing regression).""" + x = _x() + out = torch.empty_like(x, dtype=FP8_DTYPE) + scale = torch.ones(1, dtype=torch.float32, device="cuda") + torch.library.opcheck( + torch.ops.vllm.rocm_aiter_per_tensor_quant, + (out, x, scale, False), + test_utils=_INPLACE_OPCHECK_UTILS, + ) + + +def test_per_tensor_quant_dynamic_schema(): + """Dynamic per-tensor: op computes scale into the caller's buffer.""" + x = _x() + out = torch.empty_like(x, dtype=FP8_DTYPE) + scale = torch.empty(1, dtype=torch.float32, device="cuda") + torch.library.opcheck( + torch.ops.vllm.rocm_aiter_per_tensor_quant, + (out, x, scale, True), + test_utils=_INPLACE_OPCHECK_UTILS, + ) + + +def test_per_token_quant_dynamic_schema(): + """Dynamic per-token: op computes scale into a freshly allocated buffer.""" + x = _x() + torch.library.opcheck( + torch.ops.vllm.rocm_aiter_per_token_quant, + (x, FP8_DTYPE, None), + ) + + +def test_group_fp8_quant_schema(): + """Dynamic per-token-group quant.""" + x = _x() + torch.library.opcheck( + torch.ops.vllm.rocm_aiter_group_fp8_quant, + (x, 128), + ) + + +@pytest.mark.parametrize("dynamic", [True, False]) +def test_per_tensor_quant_matches_native(dynamic): + """Wrapper output matches the native scaled_fp8_quant reference.""" + from vllm import _custom_ops as ops + + torch.manual_seed(0) + x = _x() + if dynamic: + scale_in = None + else: + scale_in = torch.tensor([0.5], dtype=torch.float32, device="cuda") + + out, scale = rocm_aiter_ops.per_tensor_quant(x, FP8_DTYPE, scale_in) + ref_out, ref_scale = ops.scaled_fp8_quant(x, scale_in) + + assert out.shape == x.shape + assert out.dtype == FP8_DTYPE + assert scale.shape == ref_scale.shape + if not dynamic: + # static scale is passed through unchanged + assert torch.equal(scale, scale_in) + # Compare dequantized values to be robust to 1-ULP fp8 boundary flips. + deq = out.to(torch.float32) * scale + ref_deq = ref_out.to(torch.float32) * ref_scale + torch.testing.assert_close(deq, ref_deq, rtol=2e-2, atol=2e-2) + + +@pytest.mark.parametrize("dynamic", [True, False]) +def test_per_tensor_quant_torch_compile(monkeypatch, dynamic): + """per_tensor_quant compiles under inductor without an aliasing error. + + Forces the custom-op aliasing check to error (it is otherwise only a + warning outside CI), so a regression that returns an input-aliasing + scale fails here regardless of the CI env var. + """ + aliasing_cfg = pytest.importorskip("torch._functorch.config") + monkeypatch.setattr( + aliasing_cfg, "error_on_custom_op_aliasing", True, raising=False + ) + + x = _x() + scale = None if dynamic else torch.tensor([0.5], dtype=torch.float32, device="cuda") + + def fn(x, s): + return rocm_aiter_ops.per_tensor_quant(x, FP8_DTYPE, s) + + compiled = torch.compile(fn, fullgraph=True, backend="inductor", dynamic=False) + + out_eager, scale_eager = fn(x, scale) + out_compiled, scale_compiled = compiled(x, scale) + + assert out_compiled.shape == out_eager.shape + torch.testing.assert_close( + out_compiled.to(torch.float32) * scale_compiled, + out_eager.to(torch.float32) * scale_eager, + rtol=2e-2, + atol=2e-2, + ) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index d744da0b89b..95a5361032f 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -1019,23 +1019,26 @@ def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_fake( def _rocm_aiter_per_tensor_quant_impl( + out: torch.Tensor, x: torch.Tensor, - quant_dtype: torch.dtype, - scale: torch.Tensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - from aiter.ops.quant import per_tensor_quant_hip + scale: torch.Tensor, + is_dynamic: bool, +) -> None: + from aiter.ops.quant import dynamic_per_tensor_quant, static_per_tensor_quant - return per_tensor_quant_hip(x, scale, quant_dtype) + if is_dynamic: + dynamic_per_tensor_quant(out, x, scale) + else: + static_per_tensor_quant(out, x, scale) def _rocm_aiter_per_tensor_quant_fake( + out: torch.Tensor, x: torch.Tensor, - quant_dtype: torch.dtype, - scale: torch.Tensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - return torch.empty_like(x, dtype=quant_dtype), torch.empty( - 1, dtype=torch.float32, device=x.device - ) + scale: torch.Tensor, + is_dynamic: bool, +) -> None: + pass def _rocm_aiter_per_token_quant_impl( @@ -1979,7 +1982,7 @@ class rocm_aiter_ops: direct_register_custom_op( op_name="rocm_aiter_per_tensor_quant", op_func=_rocm_aiter_per_tensor_quant_impl, - mutates_args=[], + mutates_args=["out", "scale"], fake_impl=_rocm_aiter_per_tensor_quant_fake, dispatch_key=current_platform.dispatch_key, ) @@ -2392,7 +2395,12 @@ class rocm_aiter_ops: quant_dtype: torch.dtype, scale: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - return torch.ops.vllm.rocm_aiter_per_tensor_quant(x, quant_dtype, scale) + out = torch.empty_like(x, dtype=quant_dtype) + is_dynamic = scale is None + if is_dynamic: + scale = torch.empty(1, dtype=torch.float32, device=x.device) + torch.ops.vllm.rocm_aiter_per_tensor_quant(out, x, scale, is_dynamic) + return out, scale @staticmethod def per_token_quant( From 6c379b9e5439ae305913e4a87ebf2b2e816072b4 Mon Sep 17 00:00:00 2001 From: Chauncey Date: Fri, 19 Jun 2026 00:42:10 +0800 Subject: [PATCH 04/75] [Frontend] Add Streaming Parser Engine and new GLM4.7/GLM5.1/GLM5.2 Parser (#45915) Signed-off-by: chaunceyjiang Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/parser/engine/trace_builder.py | 76 + .../test_glm4_moe_reasoning_parser.py | 38 +- .../test_glm47_moe_tool_parser.py | 36 +- .../tool_parsers/test_glm4_moe_tool_parser.py | 1567 ++--------------- vllm/parser/engine/registered_adapters.py | 6 + vllm/parser/glm47_moe.py | 226 +++ vllm/reasoning/__init__.py | 8 +- vllm/reasoning/glm47_moe_reasoning_parser.py | 6 + vllm/tool_parsers/__init__.py | 4 +- vllm/tool_parsers/glm47_moe_tool_parser.py | 36 +- vllm/tool_parsers/glm4_moe_tool_parser.py | 495 ------ 11 files changed, 542 insertions(+), 1956 deletions(-) create mode 100644 vllm/parser/glm47_moe.py create mode 100644 vllm/reasoning/glm47_moe_reasoning_parser.py delete mode 100644 vllm/tool_parsers/glm4_moe_tool_parser.py diff --git a/tests/parser/engine/trace_builder.py b/tests/parser/engine/trace_builder.py index 128e511e690..4817d3b9005 100644 --- a/tests/parser/engine/trace_builder.py +++ b/tests/parser/engine/trace_builder.py @@ -30,6 +30,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import ( ) from vllm.parser.engine.registered_adapters import ( Gemma4Parser, + Glm47MoeParser, MinimaxM2Parser, NemotronV3Parser, Qwen3Parser, @@ -571,6 +572,80 @@ def _build_nemotron_v3(scenario: Scenario, validate: bool = True) -> Sample: ) +# ── GLM-4.7 MoE (XML tool format, starts in REASONING) ────────────── + +_GLM47_MOE_VOCAB: dict[str, int] = { + "": 50, + "": 51, + "": 60, + "": 61, + "": 62, + "": 63, + "": 64, + "": 65, +} + + +def _glm47_moe_arg_value(value: Any) -> str: + if isinstance(value, bool): + return "true" if value else "false" + if isinstance(value, (int, float)): + return str(value) + if isinstance(value, str): + return value + return json.dumps(value, ensure_ascii=False) + + +def _glm47_moe_tool_segments(tc: ToolCallSpec) -> list[tuple[str, bool]]: + segs: list[tuple[str, bool]] = [ + ("", True), + (tc.name, False), + ] + for key, value in tc.arguments.items(): + segs.extend( + [ + ("", True), + (key, False), + ("", True), + ("", True), + (_glm47_moe_arg_value(value), False), + ("", True), + ] + ) + segs.append(("", True)) + return segs + + +def _glm47_moe_segments(scenario: Scenario) -> list[tuple[str, bool]]: + segs: list[tuple[str, bool]] = [] + if scenario.reasoning is not None: + segs.append((scenario.reasoning, False)) + if scenario.content is not None or scenario.tool_calls: + segs.append(("", True)) + if scenario.content is not None: + segs.append((scenario.content, False)) + if scenario.tool_calls: + for tc in scenario.tool_calls: + segs.extend(_glm47_moe_tool_segments(tc)) + return segs + + +def _build_glm47_moe(scenario: Scenario, validate: bool = True) -> Sample: + sample = _make_sample( + sample_id=f"glm47_moe-{scenario.id}", + description=scenario.description, + vocab=_GLM47_MOE_VOCAB, + segments=_glm47_moe_segments(scenario), + expected_reasoning=scenario.reasoning if scenario.reasoning is not None else "", + expected_content=_qwen3_expected_content(scenario), + expected_tool_calls=_expected_tc(scenario), + tools=_expected_tools(scenario), + ) + if validate: + _validate_sample(sample, Glm47MoeParser) + return sample + + # ── Registry and public API ────────────────────────────────────────── _BUILDERS: dict[str, Any] = { @@ -578,6 +653,7 @@ _BUILDERS: dict[str, Any] = { "gemma4": _build_gemma4, "minimax_m2": _build_minimax_m2, "nemotron_v3": _build_nemotron_v3, + "glm47_moe": _build_glm47_moe, } diff --git a/tests/reasoning/test_glm4_moe_reasoning_parser.py b/tests/reasoning/test_glm4_moe_reasoning_parser.py index 6f7827e5b82..3d6f21b5e17 100644 --- a/tests/reasoning/test_glm4_moe_reasoning_parser.py +++ b/tests/reasoning/test_glm4_moe_reasoning_parser.py @@ -11,7 +11,7 @@ parser_name = "glm45" start_token = "" end_token = "" -REASONING_MODEL_NAME = "zai-org/GLM-4.5" +REASONING_MODEL_NAME = "zai-org/GLM-4.7" @pytest.fixture(scope="module") @@ -35,18 +35,32 @@ WITH_THINK_STREAM = { WITHOUT_THINK = { "output": "This is the rest", - "reasoning": None, - "content": "This is the rest", + "reasoning": "This is the rest", + "content": None, "is_reasoning_end": False, } WITHOUT_THINK_STREAM = { "output": "This is the rest", - "reasoning": None, - "content": "This is the rest", + "reasoning": "This is the rest", + "content": None, "is_reasoning_end": False, } +WITHOUT_OPEN_THINK = { + "output": "This is a reasoning sectionThis is the rest", + "reasoning": "This is a reasoning section", + "content": "This is the rest", + "is_reasoning_end": True, +} + +WITHOUT_OPEN_THINK_STREAM = { + "output": "This is a reasoning sectionThis is the rest", + "reasoning": "This is a reasoning section", + "content": "This is the rest", + "is_reasoning_end": True, +} + COMPLETE_REASONING = { "output": "This is a reasoning section", "reasoning": "This is a reasoning section", @@ -61,8 +75,8 @@ MULTILINE_REASONING = { } ONLY_OPEN_TAG = { "output": "This is a reasoning section", - "reasoning": None, - "content": "This is a reasoning section", + "reasoning": "This is a reasoning section", + "content": None, "is_reasoning_end": False, } @@ -94,6 +108,16 @@ TEST_CASES = [ WITHOUT_THINK_STREAM, id="without_think_stream", ), + pytest.param( + False, + WITHOUT_OPEN_THINK, + id="without_open_think", + ), + pytest.param( + True, + WITHOUT_OPEN_THINK_STREAM, + id="without_open_think_stream", + ), pytest.param( False, COMPLETE_REASONING, diff --git a/tests/tool_parsers/test_glm47_moe_tool_parser.py b/tests/tool_parsers/test_glm47_moe_tool_parser.py index 51696c95478..c9767f6f62f 100644 --- a/tests/tool_parsers/test_glm47_moe_tool_parser.py +++ b/tests/tool_parsers/test_glm47_moe_tool_parser.py @@ -16,7 +16,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import ( from vllm.tokenizers import get_tokenizer from vllm.tool_parsers.glm47_moe_tool_parser import Glm47MoeModelToolParser -MODEL = "zai-org/GLM-4.5" +MODEL = "zai-org/GLM-4.7" @pytest.fixture(scope="module") @@ -136,9 +136,10 @@ class TestGlm47Streaming: _reset(glm47_tool_parser) chunks = ["", "get_current_date", ""] current_text = "" + deltas = [] for chunk in chunks: current_text += chunk - glm47_tool_parser.extract_tool_calls_streaming( + delta = glm47_tool_parser.extract_tool_calls_streaming( previous_text="", current_text=current_text, delta_text=chunk, @@ -147,7 +148,23 @@ class TestGlm47Streaming: delta_token_ids=[], request=mock_request, ) - assert len(glm47_tool_parser.prev_tool_call_arr) >= 1 + if delta: + deltas.append(delta) + tool_calls = [ + tool_call for delta in deltas for tool_call in (delta.tool_calls or []) + ] + names = [ + tool_call.function.name + for tool_call in tool_calls + if tool_call.function and tool_call.function.name + ] + arguments = [ + tool_call.function.arguments + for tool_call in tool_calls + if tool_call.function and tool_call.function.arguments + ] + assert names == ["get_current_date"] + assert "".join(arguments) == "{}" def test_with_args(self, glm47_tool_parser, mock_request): _reset(glm47_tool_parser) @@ -161,9 +178,10 @@ class TestGlm47Streaming: "", ] current_text = "" + deltas = [] for chunk in chunks: current_text += chunk - glm47_tool_parser.extract_tool_calls_streaming( + delta = glm47_tool_parser.extract_tool_calls_streaming( previous_text="", current_text=current_text, delta_text=chunk, @@ -172,5 +190,13 @@ class TestGlm47Streaming: delta_token_ids=[], request=mock_request, ) - args = json.loads(glm47_tool_parser.prev_tool_call_arr[0]["arguments"]) + if delta: + deltas.append(delta) + arguments = [ + tool_call.function.arguments + for delta in deltas + for tool_call in (delta.tool_calls or []) + if tool_call.function and tool_call.function.arguments + ] + args = json.loads("".join(arguments)) assert args["city"] == "Beijing" diff --git a/tests/tool_parsers/test_glm4_moe_tool_parser.py b/tests/tool_parsers/test_glm4_moe_tool_parser.py index b0300297ddc..ca110adac0d 100644 --- a/tests/tool_parsers/test_glm4_moe_tool_parser.py +++ b/tests/tool_parsers/test_glm4_moe_tool_parser.py @@ -1,1067 +1,57 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Compatibility tests for GLM-4.5 using the shared GLM XML parser.""" import json -from unittest.mock import Mock - -import pytest -from openai.types.responses import FunctionTool +from typing import Any, TypedDict +from tests.parser.engine.replay_harness import MockTokenizer from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ChatCompletionToolsParam, FunctionDefinition, ) -from vllm.entrypoints.openai.engine.protocol import FunctionCall, ToolCall -from vllm.tokenizers import get_tokenizer -from vllm.tool_parsers.glm4_moe_tool_parser import ( - Glm4MoeModelToolParser, -) +from vllm.tool_parsers import ToolParserManager +from vllm.tool_parsers.glm47_moe_tool_parser import Glm47MoeModelToolParser -# Use a common model that is likely to be available MODEL = "zai-org/GLM-4.5" - -@pytest.fixture(scope="module") -def glm4_moe_tokenizer(): - return get_tokenizer(tokenizer_name=MODEL) +_GLM_VOCAB = { + "": 50, + "": 51, + "": 60, + "": 61, + "": 62, + "": 63, + "": 64, + "": 65, +} -@pytest.fixture -def sample_tools(): +class _CollectedToolDelta(TypedDict): + name: str | None + args_fragments: list[str] + + +def _mock_tokenizer() -> MockTokenizer: + return MockTokenizer(vocab=_GLM_VOCAB, tokens=[]) + + +def _tools() -> list[ChatCompletionToolsParam]: return [ ChatCompletionToolsParam( function=FunctionDefinition( - name="get_weather", - parameters={"city": {"type": "string"}}, - ), - ), - ] - - -@pytest.fixture -def glm4_moe_tool_parser(glm4_moe_tokenizer, sample_tools): - return Glm4MoeModelToolParser(glm4_moe_tokenizer, tools=sample_tools) - - -@pytest.fixture -def mock_request(sample_tools) -> ChatCompletionRequest: - request = Mock(spec=ChatCompletionRequest) - request.tools = sample_tools - return request - - -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 isinstance(actual_tool_call.id, str) - assert len(actual_tool_call.id) > 0 - - assert actual_tool_call.type == "function" - assert actual_tool_call.function.name == expected_tool_call.function.name - # Compare arguments as JSON objects to handle formatting differences - actual_args = json.loads(actual_tool_call.function.arguments) - expected_args = json.loads(expected_tool_call.function.arguments) - assert actual_args == expected_args - - -def test_extract_tool_calls_no_tools(glm4_moe_tool_parser, mock_request): - model_output = "This is a test" - extracted_tool_calls = glm4_moe_tool_parser.extract_tool_calls( - model_output, request=mock_request - ) # 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_call", - "multiple_tool_calls", - "tool_call_with_content_before", - "tool_call_with_mixed_args", - "tool_call_with_chinese_content", - ], - argnames=["model_output", "expected_tool_calls", "expected_content"], - argvalues=[ - ( - """get_current_weather - city - Dallas - state - TX - unit - fahrenheit - """, - [ - ToolCall( - function=FunctionCall( - name="get_current_weather", - arguments=json.dumps( - { - "city": "Dallas", - "state": "TX", - "unit": "fahrenheit", - } - ), - ) - ) - ], - None, - ), - ( - """get_current_weather - city - Dallas - state - TX - unit - fahrenheit - - get_current_weather - city - Orlando - state - FL - unit - 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, - ), - ( - """I'll help you check the weather. get_current_weather - city - Seattle - state - WA - unit - celsius - """, - [ - ToolCall( - function=FunctionCall( - name="get_current_weather", - arguments=json.dumps( - { - "city": "Seattle", - "state": "WA", - "unit": "celsius", - } - ), - ) - ) - ], - "I'll help you check the weather. ", - ), - ( - """get_current_weather - city - New York - state - NY - unit - celsius - """, - [ - ToolCall( - function=FunctionCall( - name="get_current_weather", - arguments=json.dumps( - { - "city": "New York", - "state": "NY", - "unit": "celsius", - } - ), - ) - ) - ], - None, - ), - ( - """I will help you get the weather.get_weather - city - Beijing - date - 2025-08-01 - """, - [ - ToolCall( - function=FunctionCall( - name="get_weather", - arguments=json.dumps( - { - "city": "Beijing", - "date": "2025-08-01", - } - ), - ) - ) - ], - "I will help you get the weather.", - ), - ], -) -def test_extract_tool_calls( - glm4_moe_tool_parser, - mock_request, - model_output, - expected_tool_calls, - expected_content, -): - extracted_tool_calls = glm4_moe_tool_parser.extract_tool_calls( - model_output, request=mock_request - ) # type: ignore[arg-type] - 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_with_thinking_tags(glm4_moe_tool_parser, mock_request): - """Test tool extraction when thinking tags are present.""" - model_output = """I want to get the weather. - -I will help you get the weather. -get_weather -city -Beijing -date -2025-08-01 -""" - - extracted_tool_calls = glm4_moe_tool_parser.extract_tool_calls( - model_output, request=mock_request - ) # type: ignore[arg-type] - - assert extracted_tool_calls.tools_called - assert len(extracted_tool_calls.tool_calls) == 1 - assert extracted_tool_calls.tool_calls[0].function.name == "get_weather" - - expected_content = """I want to get the weather. - -I will help you get the weather. -""" - assert extracted_tool_calls.content == expected_content - - -def test_extract_tool_calls_malformed_xml(glm4_moe_tool_parser, mock_request): - """Test that malformed XML is handled gracefully.""" - model_output = """get_weather -city -Seattle -incomplete_arg -value -""" - - extracted_tool_calls = glm4_moe_tool_parser.extract_tool_calls( - model_output, request=mock_request - ) # type: ignore[arg-type] - - # Should handle malformed XML gracefully - # The parser should either extract what it can or return no tool calls - # depending on how robust we want the parsing to be - assert isinstance(extracted_tool_calls.tools_called, bool) - assert isinstance(extracted_tool_calls.tool_calls, list) - - -def test_extract_tool_calls_empty_arguments(glm4_moe_tool_parser, mock_request): - """Test tool calls with no arguments.""" - model_output = """get_current_time -""" - - extracted_tool_calls = glm4_moe_tool_parser.extract_tool_calls( - model_output, request=mock_request - ) # type: ignore[arg-type] - - 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_time" - # Empty arguments should result in empty JSON object - assert extracted_tool_calls.tool_calls[0].function.arguments == "{}" - - -def test_extract_tool_calls_mixed_content(glm4_moe_tool_parser, mock_request): - """Test extraction with mixed content and multiple tool calls.""" - model_output = """I will help you get the weather info. - -get_weather -city -Beijing -date -2025-08-01 - - -meaningwhile, I will also check the weather in Shanghai. - -get_weather -city -Shanghai -date -2025-08-01 -""" - - extracted_tool_calls = glm4_moe_tool_parser.extract_tool_calls( - model_output, request=mock_request - ) # type: ignore[arg-type] - - assert extracted_tool_calls.tools_called - assert len(extracted_tool_calls.tool_calls) == 2 - - # Check first tool call - assert extracted_tool_calls.tool_calls[0].function.name == "get_weather" - args1 = json.loads(extracted_tool_calls.tool_calls[0].function.arguments) - assert args1["city"] == "Beijing" - assert args1["date"] == "2025-08-01" - - # Check second tool call - assert extracted_tool_calls.tool_calls[1].function.name == "get_weather" - args2 = json.loads(extracted_tool_calls.tool_calls[1].function.arguments) - assert args2["city"] == "Shanghai" - assert args2["date"] == "2025-08-01" - - # Content should be everything before the first tool call - assert extracted_tool_calls.content == "I will help you get the weather info.\n\n" - - -def test_streaming_basic_functionality(glm4_moe_tool_parser, mock_request): - """Test basic streaming functionality.""" - _reset_streaming_state(glm4_moe_tool_parser) - - current_text = """get_weather -city -Beijing -""" - - result = glm4_moe_tool_parser.extract_tool_calls_streaming( - previous_text="", - current_text=current_text, - delta_text=current_text, - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=mock_request, - ) - - # Should return tool call with name and arguments in one shot - assert result is not None - assert result.tool_calls is not None - assert len(result.tool_calls) >= 1 - - -def test_streaming_no_tool_calls(glm4_moe_tool_parser, mock_request): - """Test streaming when there are no tool calls.""" - _reset_streaming_state(glm4_moe_tool_parser) - - current_text = "This is just regular text without any tool calls." - - result = glm4_moe_tool_parser.extract_tool_calls_streaming( - previous_text="", - current_text=current_text, - delta_text=current_text, - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=mock_request, - ) - - # Should return content - assert result is not None - assert result.content == current_text - - -def test_streaming_with_content_before_tool_calls(glm4_moe_tool_parser, mock_request): - """Test streaming when there's content before tool calls.""" - _reset_streaming_state(glm4_moe_tool_parser) - - current_text = "I will help you get the weather." - - result = glm4_moe_tool_parser.extract_tool_calls_streaming( - previous_text="", - current_text=current_text, - delta_text=current_text, - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=mock_request, - ) - - # Should return content before the tag - assert result is not None - assert result.content == "I will help you get the weather." - - -def test_extract_tool_calls_special_characters(glm4_moe_tool_parser, mock_request): - """Test tool calls with special characters and unicode.""" - model_output = """send_message -recipient -Amy -message -It is a nice day -priority -high -""" - - extracted_tool_calls = glm4_moe_tool_parser.extract_tool_calls( - model_output, request=mock_request - ) # type: ignore[arg-type] - - assert extracted_tool_calls.tools_called - assert len(extracted_tool_calls.tool_calls) == 1 - assert extracted_tool_calls.tool_calls[0].function.name == "send_message" - - args = json.loads(extracted_tool_calls.tool_calls[0].function.arguments) - assert args["recipient"] == "Amy" - assert args["message"] == "It is a nice day" - assert args["priority"] == "high" - - -def test_extract_tool_calls_incomplete_tool_call(glm4_moe_tool_parser, mock_request): - """Test incomplete tool calls (missing closing tag).""" - model_output = """get_weather -city -Beijing -date -2025-08-01""" - - extracted_tool_calls = glm4_moe_tool_parser.extract_tool_calls( - model_output, request=mock_request - ) # type: ignore[arg-type] - - # Incomplete tool calls should not be extracted - assert not extracted_tool_calls.tools_called - assert extracted_tool_calls.tool_calls == [] - assert extracted_tool_calls.content == model_output - - -def _reset_streaming_state(parser): - """Helper to reset parser streaming state.""" - parser.current_tool_name_sent = False - parser.prev_tool_call_arr = [] - parser.current_tool_id = -1 - parser.streamed_args_for_tool = [] - parser._tool_call_ids = [] - parser._sent_content_idx = 0 - - -def test_streaming_incremental_string_value(glm4_moe_tool_parser, mock_request): - """Test incremental streaming of string argument values.""" - _reset_streaming_state(glm4_moe_tool_parser) - - # Simulate streaming a tool call chunk by chunk - chunks = [ - "", - "get_weather\n", - "city", - "", - "Bei", - "jing", - "", - "", - ] - - collected_fragments = [] - current_text = "" - for chunk in chunks: - current_text += chunk - result = glm4_moe_tool_parser.extract_tool_calls_streaming( - previous_text="", - current_text=current_text, - delta_text=chunk, - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=mock_request, - ) - if result is not None and result.tool_calls: - for tc in result.tool_calls: - func = tc.function - if isinstance(func, dict): - if func.get("arguments"): - collected_fragments.append(func["arguments"]) - if func.get("name"): - collected_fragments.append(f"name:{func['name']}") - else: - if func.arguments: - collected_fragments.append(func.arguments) - if func.name: - collected_fragments.append(f"name:{func.name}") - - # Verify we got incremental streaming of the argument value - assert len(collected_fragments) > 0 - # The fragments should include the tool name and argument pieces - combined = "".join(collected_fragments) - assert "get_weather" in combined or "name:get_weather" in combined - - -def test_streaming_empty_tool_call(glm4_moe_tool_parser, mock_request): - """Test that empty tool calls don't cause infinite loops.""" - _reset_streaming_state(glm4_moe_tool_parser) - - current_text = "" - result = glm4_moe_tool_parser.extract_tool_calls_streaming( - previous_text="", - current_text=current_text, - delta_text=current_text, - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=mock_request, - ) - - # Should not hang and should return something (None or content) - # The key is that this completes without hanging - assert result is None or hasattr(result, "content") or hasattr(result, "tool_calls") - - -def test_streaming_prev_tool_call_arr_updates(glm4_moe_tool_parser, mock_request): - """Test that prev_tool_call_arr is populated incrementally.""" - _reset_streaming_state(glm4_moe_tool_parser) - - chunks = [ - "get_weather\n", - "city", - "Beijing", - "", - ] - - current_text = "" - for chunk in chunks: - current_text += chunk - glm4_moe_tool_parser.extract_tool_calls_streaming( - previous_text="", - current_text=current_text, - delta_text=chunk, - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=mock_request, - ) - - # After the tool call completes, prev_tool_call_arr should be populated - assert len(glm4_moe_tool_parser.prev_tool_call_arr) == 1 - tool_entry = glm4_moe_tool_parser.prev_tool_call_arr[0] - assert tool_entry.get("name") == "get_weather" - - # arguments is a JSON string in the re-parse approach - args_str = tool_entry.get("arguments") - assert isinstance(args_str, str), f"Expected str, got {type(args_str)}" - parsed = json.loads(args_str) - assert parsed["city"] == "Beijing" - - # streamed_args_for_tool should match prev_tool_call_arr arguments - streamed = glm4_moe_tool_parser.streamed_args_for_tool[0] - assert streamed == args_str - - -def test_streaming_multiple_tool_calls_sequential(glm4_moe_tool_parser, mock_request): - """Test streaming multiple sequential tool calls.""" - _reset_streaming_state(glm4_moe_tool_parser) - - chunks = [ - "get_weather\n", - "city", - "Beijing", - "", - "get_weather\n", - "city", - "Shanghai", - "", - ] - - current_text = "" - for chunk in chunks: - current_text += chunk - glm4_moe_tool_parser.extract_tool_calls_streaming( - previous_text="", - current_text=current_text, - delta_text=chunk, - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=mock_request, - ) - - # Should have two tool calls in prev_tool_call_arr - assert len(glm4_moe_tool_parser.prev_tool_call_arr) == 2 - args0 = json.loads(glm4_moe_tool_parser.prev_tool_call_arr[0]["arguments"]) - args1 = json.loads(glm4_moe_tool_parser.prev_tool_call_arr[1]["arguments"]) - assert args0["city"] == "Beijing" - assert args1["city"] == "Shanghai" - - -def test_streaming_json_escape_in_string(glm4_moe_tool_parser, mock_request): - """Test that special characters in string values are properly escaped.""" - _reset_streaming_state(glm4_moe_tool_parser) - - chunks = [ - "send_message\n", - "message", - 'Hello "world"\nNew line', - "", - ] - - current_text = "" - for chunk in chunks: - current_text += chunk - glm4_moe_tool_parser.extract_tool_calls_streaming( - previous_text="", - current_text=current_text, - delta_text=chunk, - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=mock_request, - ) - - # The streamed_args_for_tool should contain valid JSON - assert len(glm4_moe_tool_parser.streamed_args_for_tool) == 1 - args_json = glm4_moe_tool_parser.streamed_args_for_tool[0] - parsed = json.loads(args_json) - assert "message" in parsed - assert '"' in parsed["message"] or "world" in parsed["message"] - - -def test_streaming_long_content_incremental(glm4_moe_tokenizer): - """Test incremental streaming of long content (Issue #32829). - - This is the core fix: for long string values like code (4000+ chars), - the parser should stream incrementally rather than buffering until - complete. This test verifies we get many fragments, not just 1-3. - """ - - # Bubble sort example from Issue #32829 - realistic long content - bubble_sort_code = '''#!/usr/bin/env python3 -# -*- coding: utf-8 -*- -""" -Bubble Sort Implementation -""" - -def bubble_sort(arr): - n = len(arr) - for i in range(n): - swapped = False - for j in range(0, n - i - 1): - if arr[j] > arr[j + 1]: - arr[j], arr[j + 1] = arr[j + 1], arr[j] - swapped = True - if not swapped: - break - return arr - -if __name__ == "__main__": - test_arr = [64, 34, 25, 12, 22, 11, 90] - print(f"Original: {test_arr}") - sorted_arr = bubble_sort(test_arr.copy()) - print(f"Sorted: {sorted_arr}")''' - - # Create tools with schema to enable string type detection - # This is required for incremental streaming of string values - tools = [ - ChatCompletionToolsParam( - function=FunctionDefinition( - name="write_to_file", + name="get_current_weather", parameters={ "type": "object", "properties": { - "file_path": {"type": "string"}, - "content": {"type": "string"}, + "city": {"type": "string"}, + "state": {"type": "string"}, + "unit": {"type": "string"}, }, }, ), ), - ] - glm4_moe_tool_parser = Glm4MoeModelToolParser(glm4_moe_tokenizer, tools=tools) - request = ChatCompletionRequest( - model=MODEL, - messages=[], - tools=tools, - ) - - # Simulate token-based streaming (special tags as single tokens) - chunks = [ - "", - "write_to_file\n", - "file_path", - "/tmp/bubble_sort.py", - "content", - "", - ] - # Add content line by line (realistic token streaming) - for line in bubble_sort_code.split("\n"): - chunks.append(line + "\n") - chunks.append("") - chunks.append("") - - # Count argument fragments - fragment_count = 0 - current_text = "" - for chunk in chunks: - current_text += chunk - result = glm4_moe_tool_parser.extract_tool_calls_streaming( - previous_text="", - current_text=current_text, - delta_text=chunk, - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=request, - ) - if result is not None and result.tool_calls: - for tc in result.tool_calls: - func = tc.function - if isinstance(func, dict): - args = func.get("arguments") - else: - args = getattr(func, "arguments", None) - if args: - fragment_count += 1 - - # For true incremental streaming, we expect many fragments (10+) - # Old buffered implementation would give only 1-3 fragments - assert fragment_count >= 10, ( - f"Expected >=10 fragments for incremental streaming, got {fragment_count}" - ) - - # Verify final result is valid JSON - assert len(glm4_moe_tool_parser.streamed_args_for_tool) == 1 - args_json = glm4_moe_tool_parser.streamed_args_for_tool[0] - parsed = json.loads(args_json) - assert parsed["file_path"] == "/tmp/bubble_sort.py" - assert "def bubble_sort" in parsed["content"] - - -def test_extract_tool_calls_numeric_deserialization(glm4_moe_tool_parser, mock_request): - """Test that numeric arguments are deserialized as numbers, not strings.""" - model_output = """calculate -operation -add -a -42 -b -3.14 -enabled -true -""" - - extracted_tool_calls = glm4_moe_tool_parser.extract_tool_calls( - model_output, request=mock_request - ) # type: ignore[arg-type] - - assert extracted_tool_calls.tools_called - assert len(extracted_tool_calls.tool_calls) == 1 - - args = json.loads(extracted_tool_calls.tool_calls[0].function.arguments) - - # String should remain string - assert args["operation"] == "add" - assert isinstance(args["operation"], str) - - # Integer should be deserialized as int - assert args["a"] == 42 - assert isinstance(args["a"], int) - - # Float should be deserialized as float - assert args["b"] == 3.14 - assert isinstance(args["b"], float) - - # Boolean should be deserialized as bool - assert args["enabled"] is True - assert isinstance(args["enabled"], bool) - - -def test_whitespace_preserved_in_arg_values(glm4_moe_tokenizer): - """Test that string arguments preserve leading and trailing whitespace.""" - tools = [ - ChatCompletionToolsParam( - function=FunctionDefinition( - name="apply_diff", - parameters={ - "type": "object", - "properties": { - "s": {"type": "string"}, - }, - "required": ["s"], - }, - ), - ), - ] - parser = Glm4MoeModelToolParser(glm4_moe_tokenizer, tools=tools) - request = ChatCompletionRequest(model=MODEL, messages=[], tools=tools) - - model_output = """apply_diff -s - indented code -""" - - extracted_tool_calls = parser.extract_tool_calls(model_output, request=request) - args = json.loads(extracted_tool_calls.tool_calls[0].function.arguments) - - assert args["s"] == " indented code " - - -def test_zero_argument_tool_call(glm4_moe_tool_parser, mock_request): - """Regression: zero-argument tool call crash (PR #32321).""" - model_output = """get_time -""" - - extracted = glm4_moe_tool_parser.extract_tool_calls( - model_output, request=mock_request - ) # type: ignore[arg-type] - - assert extracted.tools_called - assert len(extracted.tool_calls) == 1 - assert extracted.tool_calls[0].function.name == "get_time" - args = json.loads(extracted.tool_calls[0].function.arguments) - assert args == {} - - -def test_malformed_tool_call_no_regex_match(glm4_moe_tool_parser, mock_request): - """Regression: malformed tool_call with no regex match (PR #32321).""" - model_output = " " - - extracted = glm4_moe_tool_parser.extract_tool_calls( - model_output, request=mock_request - ) # type: ignore[arg-type] - - assert extracted.tools_called is False - assert extracted.tool_calls == [] - - -def test_delimiter_preserved_transformers_5x(glm4_moe_tool_parser): - """Regression: adjust_request sets skip_special_tokens=False (PR #31622).""" - # Tools enabled - request_with_tools = ChatCompletionRequest( - model=MODEL, - messages=[], - tools=[ - { - "type": "function", - "function": { - "name": "get_weather", - "parameters": { - "type": "object", - "properties": {"city": {"type": "string"}}, - }, - }, - } - ], - ) # type: ignore - adjusted = glm4_moe_tool_parser.adjust_request(request_with_tools) - assert adjusted.skip_special_tokens is False - - # tool_choice="none" - request_no_choice = ChatCompletionRequest( - model=MODEL, - messages=[], - tools=[ - { - "type": "function", - "function": { - "name": "get_weather", - "parameters": { - "type": "object", - "properties": {"city": {"type": "string"}}, - }, - }, - } - ], - tool_choice="none", - ) # type: ignore - adjusted_none = glm4_moe_tool_parser.adjust_request(request_no_choice) - assert adjusted_none.skip_special_tokens is True - - # No tools at all - request_no_tools = ChatCompletionRequest( - model=MODEL, - messages=[], - ) # type: ignore - adjusted_empty = glm4_moe_tool_parser.adjust_request(request_no_tools) - assert adjusted_empty.skip_special_tokens is True - - -def test_unicode_characters_preserved(glm4_moe_tool_parser, mock_request): - """Regression: Unicode chars must not be escaped to \\uXXXX (PR #30920).""" - model_output = """send_message -greeting -你好世界 -emoji -🎉 -""" - - extracted = glm4_moe_tool_parser.extract_tool_calls( - model_output, request=mock_request - ) # type: ignore[arg-type] - - assert extracted.tools_called - assert len(extracted.tool_calls) == 1 - - raw_args = extracted.tool_calls[0].function.arguments - assert "你好世界" in raw_args - assert "🎉" in raw_args - assert "\\u4f60" not in raw_args - parsed_args = json.loads(raw_args) - assert parsed_args["greeting"] == "你好世界" - assert parsed_args["emoji"] == "🎉" - - -def test_streaming_multi_token_chunks(glm4_moe_tool_parser, mock_request): - """Test that multi-token chunks (stream_interval > 1) are handled correctly. - - With stream_interval > 1 or MTP, multiple XML tags arrive in one delta. - The old buffer-based parser could only return one delta per call, losing - data on the final output. The re-parse approach handles this correctly. - """ - _reset_streaming_state(glm4_moe_tool_parser) - - # Simulate stream_interval=3: chunks contain multiple XML tags - chunks = [ - "get_weather\ncityBei", - "jing", - "", - ] - - current_text = "" - for chunk in chunks: - current_text += chunk - glm4_moe_tool_parser.extract_tool_calls_streaming( - previous_text="", - current_text=current_text, - delta_text=chunk, - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=mock_request, - ) - - # All data should be captured despite multi-token chunks - assert len(glm4_moe_tool_parser.prev_tool_call_arr) == 1 - args = json.loads(glm4_moe_tool_parser.streamed_args_for_tool[0]) - assert args["city"] == "Beijing" - - -def test_streaming_entire_tool_call_at_once(glm4_moe_tool_parser, mock_request): - """Test that a complete tool call arriving in one delta works. - - This simulates the extreme MTP case where all tokens arrive at once. - """ - _reset_streaming_state(glm4_moe_tool_parser) - - full_text = ( - "get_weather\n" - "city" - "Beijing" - "" - ) - - result = glm4_moe_tool_parser.extract_tool_calls_streaming( - previous_text="", - current_text=full_text, - delta_text=full_text, - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=mock_request, - ) - - # Should emit tool call with complete arguments in one shot - assert result is not None - assert result.tool_calls is not None - - # Verify final state - assert len(glm4_moe_tool_parser.prev_tool_call_arr) == 1 - args = json.loads(glm4_moe_tool_parser.streamed_args_for_tool[0]) - assert args["city"] == "Beijing" - - -def test_streaming_content_between_tool_calls_multi_token( - glm4_moe_tool_parser, mock_request -): - """Test content between tool calls with multi-token chunks.""" - _reset_streaming_state(glm4_moe_tool_parser) - - # Deliver everything at once — worst case for the old buffer parser - full_text = ( - "I will check.\n" - "get_weather\n" - "city" - "Beijing" - "" - "\nAlso Shanghai.\n" - "get_weather\n" - "city" - "Shanghai" - "" - ) - - # First call with partial text (content only) - partial = "I will check.\n" - result1 = glm4_moe_tool_parser.extract_tool_calls_streaming( - previous_text="", - current_text=partial, - delta_text=partial, - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=mock_request, - ) - assert result1 is not None - assert result1.content == "I will check.\n" - - # Second call with everything - glm4_moe_tool_parser.extract_tool_calls_streaming( - previous_text="", - current_text=full_text, - delta_text=full_text[len(partial) :], - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=mock_request, - ) - - # Should have both tool calls - assert len(glm4_moe_tool_parser.prev_tool_call_arr) == 2 - args0 = json.loads(glm4_moe_tool_parser.prev_tool_call_arr[0]["arguments"]) - args1 = json.loads(glm4_moe_tool_parser.prev_tool_call_arr[1]["arguments"]) - assert args0["city"] == "Beijing" - assert args1["city"] == "Shanghai" - - -def test_streaming_multi_token_with_multiple_args(glm4_moe_tokenizer): - """Test multi-token streaming with multiple arguments of mixed types.""" - tools = [ ChatCompletionToolsParam( function=FunctionDefinition( name="calculate", @@ -1071,415 +61,168 @@ def test_streaming_multi_token_with_multiple_args(glm4_moe_tokenizer): "operation": {"type": "string"}, "a": {"type": "number"}, "b": {"type": "number"}, + "enabled": {"type": "boolean"}, }, }, ), ), - ] - parser = Glm4MoeModelToolParser(glm4_moe_tokenizer, tools=tools) - request = ChatCompletionRequest( - model=MODEL, - messages=[], - tools=tools, - ) - - # All arguments arrive in two big chunks (simulates stream_interval=5) - chunks = [ - "calculate\noperationadda", - "42b3.14", - ] - - current_text = "" - for chunk in chunks: - current_text += chunk - parser.extract_tool_calls_streaming( - previous_text="", - current_text=current_text, - delta_text=chunk, - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=request, - ) - - args = json.loads(parser.streamed_args_for_tool[0]) - assert args["operation"] == "add" - assert args["a"] == 42 - assert args["b"] == 3.14 - - -def _simulate_streaming(tokenizer, parser, request, text, stream_interval=1): - """Simulate streaming with a given stream_interval. - - Tokens are batched into chunks of ``stream_interval`` tokens, - mimicking how the output processor delivers them. - Returns a list of non-None DeltaMessages. - """ - tokens = tokenizer.encode(text) - previous_text = "" - deltas = [] - for i in range(0, len(tokens), stream_interval): - chunk_ids = tokens[i : i + stream_interval] - delta_text = tokenizer.decode(chunk_ids) - current_text = previous_text + delta_text - delta = parser.extract_tool_calls_streaming( - previous_text=previous_text, - current_text=current_text, - delta_text=delta_text, - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=chunk_ids, - request=request, - ) - previous_text = current_text - if delta is not None: - deltas.append(delta) - return deltas - - -def _collect_from_deltas(deltas): - """Reconstruct tool call names/args and content from a delta stream.""" - tools: dict[int, dict] = {} - content_parts: list[str] = [] - for d in deltas: - if d.content: - content_parts.append(d.content) - if d.tool_calls: - for tc in d.tool_calls: - func = tc.function - if isinstance(func, dict): - name = func.get("name") - args = func.get("arguments") - else: - name = getattr(func, "name", None) - args = getattr(func, "arguments", None) - idx = tc.index - if idx not in tools: - tools[idx] = {"name": None, "args_fragments": []} - if name: - tools[idx]["name"] = name - if args: - tools[idx]["args_fragments"].append(args) - return content_parts, tools - - -@pytest.mark.parametrize("stream_interval", [1, 2, 3, 5, 8]) -def test_stream_interval_single_tool_call(glm4_moe_tokenizer, stream_interval): - """Tool call streaming produces correct name + args at any interval.""" - tools = [ ChatCompletionToolsParam( - function=FunctionDefinition( - name="get_weather", - parameters={ - "type": "object", - "properties": {"city": {"type": "string"}}, - }, - ), - ), - ] - parser = Glm4MoeModelToolParser(glm4_moe_tokenizer, tools=tools) - request = ChatCompletionRequest(model=MODEL, messages=[], tools=tools) - - text = ( - "get_weather\n" - "city" - "Beijing" - "" - ) - - deltas = _simulate_streaming( - glm4_moe_tokenizer, parser, request, text, stream_interval - ) - _, tools_found = _collect_from_deltas(deltas) - - assert 0 in tools_found - assert tools_found[0]["name"] == "get_weather" - args_json = "".join(tools_found[0]["args_fragments"]) - parsed = json.loads(args_json) - assert parsed == {"city": "Beijing"} - - -@pytest.mark.parametrize("stream_interval", [1, 2, 3, 5, 8]) -def test_stream_interval_multiple_tool_calls(glm4_moe_tokenizer, stream_interval): - """Multiple sequential tool calls with correct indices at any interval.""" - tools = [ - ChatCompletionToolsParam( - function=FunctionDefinition( - name="get_weather", - parameters={ - "type": "object", - "properties": {"city": {"type": "string"}}, - }, - ), - ), - ] - parser = Glm4MoeModelToolParser(glm4_moe_tokenizer, tools=tools) - request = ChatCompletionRequest(model=MODEL, messages=[], tools=tools) - - text = ( - "get_weather\n" - "city" - "Beijing" - "" - "get_weather\n" - "city" - "Shanghai" - "" - ) - - deltas = _simulate_streaming( - glm4_moe_tokenizer, parser, request, text, stream_interval - ) - _, tools_found = _collect_from_deltas(deltas) - - assert 0 in tools_found and 1 in tools_found - args0 = json.loads("".join(tools_found[0]["args_fragments"])) - args1 = json.loads("".join(tools_found[1]["args_fragments"])) - assert args0 == {"city": "Beijing"} - assert args1 == {"city": "Shanghai"} - - -@pytest.mark.parametrize("stream_interval", [1, 2, 3, 5, 8]) -def test_stream_interval_content_then_tool_call(glm4_moe_tokenizer, stream_interval): - """Content before a tool call is fully emitted before tool deltas.""" - tools = [ - ChatCompletionToolsParam( - function=FunctionDefinition( - name="get_weather", - parameters={ - "type": "object", - "properties": {"city": {"type": "string"}}, - }, - ), - ), - ] - parser = Glm4MoeModelToolParser(glm4_moe_tokenizer, tools=tools) - request = ChatCompletionRequest(model=MODEL, messages=[], tools=tools) - - text = ( - "I will check the weather for you.\n" - "get_weather\n" - "city" - "Beijing" - "" - ) - - deltas = _simulate_streaming( - glm4_moe_tokenizer, parser, request, text, stream_interval - ) - content_parts, tools_found = _collect_from_deltas(deltas) - - # Content must be present and precede tool calls - full_content = "".join(content_parts) - assert "I will check the weather" in full_content - - # Tool call must be correct - assert 0 in tools_found - assert tools_found[0]["name"] == "get_weather" - args = json.loads("".join(tools_found[0]["args_fragments"])) - assert args == {"city": "Beijing"} - - -def test_stream_interval_extreme_single_chunk(glm4_moe_tokenizer): - """Extreme MTP: entire output arrives in one chunk (interval=9999).""" - tools = [ - ChatCompletionToolsParam( - function=FunctionDefinition( - name="get_weather", - parameters={ - "type": "object", - "properties": {"city": {"type": "string"}}, - }, - ), - ), - ] - parser = Glm4MoeModelToolParser(glm4_moe_tokenizer, tools=tools) - request = ChatCompletionRequest(model=MODEL, messages=[], tools=tools) - - text = ( - "Here is the weather.\n" - "get_weather\n" - "city" - "Beijing" - "" - ) - - deltas = _simulate_streaming( - glm4_moe_tokenizer, parser, request, text, stream_interval=9999 - ) - content_parts, tools_found = _collect_from_deltas(deltas) - - assert "Here is the weather" in "".join(content_parts) - assert 0 in tools_found - assert tools_found[0]["name"] == "get_weather" - args = json.loads("".join(tools_found[0]["args_fragments"])) - assert args == {"city": "Beijing"} - - -@pytest.mark.parametrize("stream_interval", [1, 2, 5]) -def test_stream_interval_content_between_tool_calls( - glm4_moe_tokenizer, stream_interval -): - """Content between tool calls must be emitted, not silently dropped.""" - tools = [ - ChatCompletionToolsParam( - function=FunctionDefinition( - name="get_weather", - parameters={ - "type": "object", - "properties": {"city": {"type": "string"}}, - }, - ), - ), - ] - parser = Glm4MoeModelToolParser(glm4_moe_tokenizer, tools=tools) - request = ChatCompletionRequest(model=MODEL, messages=[], tools=tools) - - text = ( - "Checking Beijing.\n" - "get_weather\n" - "city" - "Beijing" - "" - "\nAlso Shanghai.\n" - "get_weather\n" - "city" - "Shanghai" - "" - ) - - deltas = _simulate_streaming( - glm4_moe_tokenizer, parser, request, text, stream_interval - ) - content_parts, tools_found = _collect_from_deltas(deltas) - - full_content = "".join(content_parts) - # Both prefix and inter-tool-call content must appear - assert "Checking Beijing" in full_content - assert "Also Shanghai" in full_content - - # Both tool calls must be correct - assert 0 in tools_found and 1 in tools_found - args0 = json.loads("".join(tools_found[0]["args_fragments"])) - args1 = json.loads("".join(tools_found[1]["args_fragments"])) - assert args0 == {"city": "Beijing"} - assert args1 == {"city": "Shanghai"} - - -# ── FunctionTool (Responses API) tests ────────────────────────────── - - -@pytest.fixture -def function_tools(): - return [ - FunctionTool( - type="function", - name="get_weather", - parameters={ - "type": "object", - "properties": { - "city": {"type": "string"}, - "unit": {"type": "string"}, - }, - }, - ), - FunctionTool( - type="function", - name="calculate", - parameters={ - "type": "object", - "properties": { - "operation": {"type": "string"}, - "a": {"type": "number"}, - "b": {"type": "number"}, - }, - }, + function=FunctionDefinition(name="get_time", parameters={}), ), ] -@pytest.fixture -def glm4_moe_parser_function_tools(glm4_moe_tokenizer, function_tools): - return Glm4MoeModelToolParser(glm4_moe_tokenizer, tools=function_tools) +def _request(tools: list[ChatCompletionToolsParam]) -> ChatCompletionRequest: + return ChatCompletionRequest(model=MODEL, messages=[], tools=tools) -@pytest.fixture -def mock_request_function_tools(function_tools) -> ChatCompletionRequest: - request = Mock(spec=ChatCompletionRequest) - request.tools = function_tools - return request +def _parser(tools: list[ChatCompletionToolsParam] | None = None): + return Glm47MoeModelToolParser(_mock_tokenizer(), tools=tools) -def test_extract_tool_calls_with_function_tool( - glm4_moe_parser_function_tools, mock_request_function_tools -): - model_output = """get_weather +def _collect_tool_deltas(deltas: Any) -> dict[int, _CollectedToolDelta]: + calls: dict[int, _CollectedToolDelta] = {} + for delta in deltas: + if delta is None or not delta.tool_calls: + continue + for tool_call in delta.tool_calls: + entry = calls.setdefault( + tool_call.index, + {"name": None, "args_fragments": []}, + ) + function = tool_call.function + if function is None: + continue + if isinstance(function, dict): + name = function.get("name") + arguments = function.get("arguments") + else: + name = function.name + arguments = function.arguments + if isinstance(name, str) and name: + entry["name"] = name + if isinstance(arguments, str) and arguments: + entry["args_fragments"].append(arguments) + return calls + + +def test_glm45_uses_shared_glm47_parser(): + assert ToolParserManager.get_tool_parser("glm45") is Glm47MoeModelToolParser + assert ToolParserManager.get_tool_parser("glm47") is Glm47MoeModelToolParser + + +def test_extract_tool_calls_with_glm45_newline_format(): + tools = _tools() + parser = _parser(tools) + model_output = """I'll check it. get_current_weather city Dallas +state +TX unit fahrenheit """ - extracted = glm4_moe_parser_function_tools.extract_tool_calls( - model_output, request=mock_request_function_tools - ) + extracted = parser.extract_tool_calls(model_output, request=_request(tools)) + assert extracted.tools_called + assert extracted.content == "I'll check it." assert len(extracted.tool_calls) == 1 - assert extracted.tool_calls[0].function.name == "get_weather" - args = json.loads(extracted.tool_calls[0].function.arguments) - assert args["city"] == "Dallas" - assert args["unit"] == "fahrenheit" + tool_call = extracted.tool_calls[0] + assert tool_call.function.name == "get_current_weather" + assert json.loads(tool_call.function.arguments) == { + "city": "Dallas", + "state": "TX", + "unit": "fahrenheit", + } -def test_extract_tool_calls_with_function_tool_mixed_types( - glm4_moe_parser_function_tools, mock_request_function_tools -): - model_output = """calculate -operation -add -a -42 -b -3.14 +def test_extract_multiple_tool_calls_with_glm45_newline_format(): + tools = _tools() + parser = _parser(tools) + model_output = """get_current_weather +cityDallas + +get_current_weather +cityOrlando """ - extracted = glm4_moe_parser_function_tools.extract_tool_calls( - model_output, request=mock_request_function_tools - ) + extracted = parser.extract_tool_calls(model_output, request=_request(tools)) + assert extracted.tools_called - args = json.loads(extracted.tool_calls[0].function.arguments) - assert args["operation"] == "add" - assert isinstance(args["a"], (int, float)) - assert isinstance(args["b"], float) + assert [tc.function.name for tc in extracted.tool_calls] == [ + "get_current_weather", + "get_current_weather", + ] + assert [ + json.loads(tc.function.arguments)["city"] for tc in extracted.tool_calls + ] == ["Dallas", "Orlando"] -def test_streaming_with_function_tool( - glm4_moe_parser_function_tools, mock_request_function_tools -): - _reset_streaming_state(glm4_moe_parser_function_tools) +def test_extract_tool_calls_coerces_schema_types(): + tools = _tools() + parser = _parser(tools) + model_output = """calculate +operationadd +a42 +b3.14 +enabledtrue +""" + extracted = parser.extract_tool_calls(model_output, request=_request(tools)) + + assert extracted.tools_called + assert json.loads(extracted.tool_calls[0].function.arguments) == { + "operation": "add", + "a": 42, + "b": 3.14, + "enabled": True, + } + + +def test_extract_zero_argument_tool_call_with_glm45_newline_format(): + tools = _tools() + parser = _parser(tools) + + extracted = parser.extract_tool_calls( + "get_time\n", + request=_request(tools), + ) + + assert extracted.tools_called + assert extracted.tool_calls[0].function.name == "get_time" + assert json.loads(extracted.tool_calls[0].function.arguments) == {} + + +def test_streaming_tool_call_with_glm45_newline_format(): + tools = _tools() + parser = _parser(tools) + request = _request(tools) chunks = [ - "get_weather\n", + "", + "get_current_weather\n", "city", "Bei", - "jing", - "", + "jing", "", ] - + deltas = [] current_text = "" + for chunk in chunks: current_text += chunk - glm4_moe_parser_function_tools.extract_tool_calls_streaming( - previous_text="", - current_text=current_text, - delta_text=chunk, - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=mock_request_function_tools, + deltas.append( + parser.extract_tool_calls_streaming( + previous_text="", + current_text=current_text, + delta_text=chunk, + previous_token_ids=[], + current_token_ids=[], + delta_token_ids=[], + request=request, + ) ) - assert len(glm4_moe_parser_function_tools.prev_tool_call_arr) == 1 - args = json.loads(glm4_moe_parser_function_tools.prev_tool_call_arr[0]["arguments"]) - assert args["city"] == "Beijing" + calls = _collect_tool_deltas(deltas) + assert calls[0]["name"] == "get_current_weather" + assert json.loads("".join(calls[0]["args_fragments"])) == {"city": "Beijing"} diff --git a/vllm/parser/engine/registered_adapters.py b/vllm/parser/engine/registered_adapters.py index d45a82879fa..9d670f30564 100644 --- a/vllm/parser/engine/registered_adapters.py +++ b/vllm/parser/engine/registered_adapters.py @@ -9,6 +9,7 @@ names so that :class:`ReasoningParserManager` and from vllm.parser.engine.adapters import make_adapters from vllm.parser.gemma4 import Gemma4Parser +from vllm.parser.glm47_moe import Glm47MoeParser from vllm.parser.minimax_m2 import MinimaxM2Parser from vllm.parser.nemotron_v3 import NemotronV3Parser from vllm.parser.qwen3 import Qwen3Parser @@ -32,3 +33,8 @@ from vllm.parser.qwen3 import Qwen3Parser Qwen3ParserReasoningAdapter, Qwen3ParserToolAdapter, ) = make_adapters(Qwen3Parser) + +( + Glm47MoeParserReasoningAdapter, + Glm47MoeParserToolAdapter, +) = make_adapters(Glm47MoeParser) diff --git a/vllm/parser/glm47_moe.py b/vllm/parser/glm47_moe.py new file mode 100644 index 00000000000..8aa4feef259 --- /dev/null +++ b/vllm/parser/glm47_moe.py @@ -0,0 +1,226 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GLM-4.7 parser for reasoning and tool calls. + +GLM-4.7 uses XML-like tool calls:: + + func_namekeyvalue + +The function name can be followed directly by the first ```` tag, +and tool calls may have no arguments. +""" + +from __future__ import annotations + +import functools +import json +from typing import TYPE_CHECKING + +import regex as re + +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.parser.engine.events import EventType +from vllm.parser.engine.parser_engine import ParserEngine +from vllm.parser.engine.parser_engine_config import ( + ParserEngineConfig, + ParserState, + Transition, +) + +if TYPE_CHECKING: + from vllm.tokenizers import TokenizerLike + from vllm.tool_parsers.abstract_tool_parser import Tool + +THINK_START = "" +THINK_END = "" +TOOL_CALL_START = "" +TOOL_CALL_END = "" +ARG_KEY_START = "" +ARG_KEY_END = "" +ARG_VALUE_START = "" +ARG_VALUE_END = "" + +_ARG_RE = re.compile( + r"(?P.*?)\s*" + r"(?P.*?)", + re.DOTALL, +) +_PARTIAL_ARG_RE = re.compile( + r"(?P.*?)\s*" + r"(?P.*)$", + re.DOTALL, +) + + +def _glm47_arg_converter(raw_args: str, partial: bool) -> str: + params: dict[str, object] = {} + + for match in _ARG_RE.finditer(raw_args): + params[match.group("key").strip()] = match.group("value") + + if partial: + remaining = _ARG_RE.sub("", raw_args) + match = _PARTIAL_ARG_RE.search(remaining) + if match: + key = match.group("key").strip() + if key: + params[key] = match.group("value") + + return json.dumps(params, ensure_ascii=False) + + +@functools.cache +def glm47_moe_config(thinking: bool = True) -> ParserEngineConfig: + arg_tag_transitions = { + (ParserState.TOOL_ARGS, terminal): Transition( + ParserState.TOOL_ARGS, + (EventType.ARG_VALUE_CHUNK,), + ) + for terminal in ( + "ARG_KEY_START", + "ARG_KEY_END", + "ARG_VALUE_START", + "ARG_VALUE_END", + ) + } + + reasoning_terminals = ( + { + "THINK_START": THINK_START, + "THINK_END": THINK_END, + } + if thinking + else {} + ) + reasoning_token_id_terminals = ( + { + "THINK_START": THINK_START, + "THINK_END": THINK_END, + } + if thinking + else {} + ) + reasoning_transitions = ( + { + (ParserState.CONTENT, "THINK_START"): Transition( + ParserState.REASONING, + (EventType.REASONING_START,), + ), + (ParserState.REASONING, "THINK_END"): Transition( + ParserState.CONTENT, + (EventType.REASONING_END,), + ), + (ParserState.CONTENT, "THINK_END"): Transition( + ParserState.CONTENT, + (), + ), + } + if thinking + else {} + ) + + return ParserEngineConfig( + name="glm47_moe", + initial_state=ParserState.REASONING if thinking else ParserState.CONTENT, + terminals={ + **reasoning_terminals, + "TOOL_START": TOOL_CALL_START, + "TOOL_END": TOOL_CALL_END, + "ARG_KEY_START": ARG_KEY_START, + "ARG_KEY_END": ARG_KEY_END, + "ARG_VALUE_START": ARG_VALUE_START, + "ARG_VALUE_END": ARG_VALUE_END, + }, + token_id_terminals={ + **reasoning_token_id_terminals, + "TOOL_START": TOOL_CALL_START, + "TOOL_END": TOOL_CALL_END, + }, + transitions={ + **reasoning_transitions, + (ParserState.REASONING, "THINK_START"): Transition( + ParserState.REASONING, + (), + ), + (ParserState.REASONING, "TOOL_START"): Transition( + ParserState.TOOL_NAME, + (EventType.REASONING_END, EventType.TOOL_CALL_START), + ), + (ParserState.CONTENT, "TOOL_START"): Transition( + ParserState.TOOL_NAME, + (EventType.TOOL_CALL_START,), + ), + (ParserState.TOOL_NAME, "ARG_KEY_START"): Transition( + ParserState.TOOL_ARGS, + (EventType.ARG_VALUE_CHUNK,), + ), + (ParserState.TOOL_NAME, "TOOL_END"): Transition( + ParserState.CONTENT, + (EventType.TOOL_CALL_END,), + ), + (ParserState.TOOL_ARGS, "TOOL_END"): Transition( + ParserState.CONTENT, + (EventType.TOOL_CALL_END,), + ), + **arg_tag_transitions, + }, + arg_converter=_glm47_arg_converter, + stream_arg_deltas=True, + tool_args_json=False, + validate_tool_names=True, + ) + + +class Glm47MoeParser(ParserEngine): + """GLM-4.7 parser backed by the declarative parser engine.""" + + def __init__( + self, + tokenizer: TokenizerLike, + tools: list[Tool] | None = None, + **kwargs, + ) -> None: + chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {} + thinking = chat_kwargs.get("thinking", None) + enable_thinking = chat_kwargs.get("enable_thinking", None) + self.thinking_enabled = ( + True + if thinking is None and enable_thinking is None + else bool(thinking) or bool(enable_thinking) + ) + kwargs.setdefault( + "parser_engine_config", + glm47_moe_config(thinking=self.thinking_enabled), + ) + super().__init__(tokenizer, tools, **kwargs) + + def _emit_name_delta(self, idx: int, deltas, name: str | None) -> None: + if name is not None: + name = name.strip() + super()._emit_name_delta(idx, deltas, name) + + def _handle_tool_end(self, event, deltas) -> None: + idx = event.tool_index + if 0 <= idx < len(self._tool_slots): + self._tool_slots[idx].name = self._tool_slots[idx].name.strip() + super()._handle_tool_end(event, deltas) + + def is_reasoning_end(self, input_ids: list[int]) -> bool: + if not self.thinking_enabled: + return True + return super().is_reasoning_end(input_ids) + + def extract_content_ids(self, input_ids: list[int]) -> list[int]: + if not self.thinking_enabled: + return input_ids + return super().extract_content_ids(input_ids) + + def extract_reasoning( + self, + model_output: str, + request: ChatCompletionRequest | ResponsesRequest, + ) -> tuple[str | None, str | None]: + if not self.thinking_enabled: + return None, model_output + return super().extract_reasoning(model_output, request) diff --git a/vllm/reasoning/__init__.py b/vllm/reasoning/__init__.py index 7d46faa6de8..cbb1fa350f5 100644 --- a/vllm/reasoning/__init__.py +++ b/vllm/reasoning/__init__.py @@ -53,8 +53,12 @@ _REASONING_PARSERS_TO_REGISTER = { "Gemma4ParserReasoningAdapter", ), "glm45": ( - "deepseek_v3_reasoning_parser", - "DeepSeekV3ReasoningWithThinkingParser", + "glm47_moe_reasoning_parser", + "Glm47MoeParserReasoningAdapter", + ), + "glm47": ( + "glm47_moe_reasoning_parser", + "Glm47MoeParserReasoningAdapter", ), "openai_gptoss": ( "gptoss_reasoning_parser", diff --git a/vllm/reasoning/glm47_moe_reasoning_parser.py b/vllm/reasoning/glm47_moe_reasoning_parser.py new file mode 100644 index 00000000000..8e963f88b09 --- /dev/null +++ b/vllm/reasoning/glm47_moe_reasoning_parser.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm.parser.engine.registered_adapters import Glm47MoeParserReasoningAdapter + +__all__ = ["Glm47MoeParserReasoningAdapter"] diff --git a/vllm/tool_parsers/__init__.py b/vllm/tool_parsers/__init__.py index 407e57ca2f9..bbc4d2edb19 100644 --- a/vllm/tool_parsers/__init__.py +++ b/vllm/tool_parsers/__init__.py @@ -51,8 +51,8 @@ _TOOL_PARSERS_TO_REGISTER = { "Ernie45ToolParser", ), "glm45": ( - "glm4_moe_tool_parser", - "Glm4MoeModelToolParser", + "glm47_moe_tool_parser", + "Glm47MoeModelToolParser", ), "glm47": ( "glm47_moe_tool_parser", diff --git a/vllm/tool_parsers/glm47_moe_tool_parser.py b/vllm/tool_parsers/glm47_moe_tool_parser.py index 80068264b70..70275a6ac03 100644 --- a/vllm/tool_parsers/glm47_moe_tool_parser.py +++ b/vllm/tool_parsers/glm47_moe_tool_parser.py @@ -1,41 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -GLM-4.7 Tool Call Parser. -GLM-4.7 uses a slightly different tool call format compared to GLM-4.5: - - The function name may appear on the same line as ```` without - a newline separator before the first ````. - - Tool calls may have zero arguments - (e.g. ``func``). +from __future__ import annotations -This parser overrides the parent regex patterns to handle both formats. -""" - -import regex as re - -from vllm.logger import init_logger -from vllm.tokenizers import TokenizerLike -from vllm.tool_parsers.abstract_tool_parser import Tool -from vllm.tool_parsers.glm4_moe_tool_parser import Glm4MoeModelToolParser - -logger = init_logger(__name__) +from vllm.parser.engine.registered_adapters import Glm47MoeParserToolAdapter -class Glm47MoeModelToolParser(Glm4MoeModelToolParser): +class Glm47MoeModelToolParser(Glm47MoeParserToolAdapter): # type: ignore[valid-type, misc] supports_required_and_named = False structural_tag_model = "glm_4_7" - - def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): - super().__init__(tokenizer, tools) - # GLM-4.7 format: func_name[...]* - # The function name can be followed by a newline, whitespace, or - # directly by tags (no separator). The arg section is - # optional so that zero-argument calls are supported. - self.func_detail_regex = re.compile( - r"\s*(\S+?)\s*(.*)?", re.DOTALL - ) - self.func_arg_regex = re.compile( - r"(.*?)\s*(.*?)", - re.DOTALL, - ) diff --git a/vllm/tool_parsers/glm4_moe_tool_parser.py b/vllm/tool_parsers/glm4_moe_tool_parser.py deleted file mode 100644 index 213a774535b..00000000000 --- a/vllm/tool_parsers/glm4_moe_tool_parser.py +++ /dev/null @@ -1,495 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -GLM-4 Tool Call Parser with incremental string streaming support. - -This parser fixes the streaming issue reported in Issue #32829 where long string -parameters (e.g., file content with 4000+ characters of code) are buffered until -complete, causing multi-second delays before the user sees any content. - -The fix streams string values incrementally as they arrive, providing a true -streaming experience for long content. -""" - -import json -from collections.abc import Sequence -from typing import Any - -import regex as re - -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 ( - DeltaFunctionCall, - DeltaMessage, - DeltaToolCall, - ExtractedToolCallInformation, - FunctionCall, - ToolCall, -) -from vllm.entrypoints.openai.responses.protocol import ResponsesRequest -from vllm.logger import init_logger -from vllm.tokenizers import TokenizerLike -from vllm.tool_parsers.abstract_tool_parser import ( - Tool, - ToolParser, -) -from vllm.tool_parsers.utils import ( - extract_types_from_schema, - find_tool_properties, - partial_tag_overlap, - safe_literal_eval, -) - -logger = init_logger(__name__) - - -class Glm4MoeModelToolParser(ToolParser): - """Tool parser for GLM-4 models with incremental string streaming. - - On every streaming call the parser re-parses ``current_text`` to find - ```` regions, builds the JSON arguments string for each tool - call, and diffs against what was previously sent to emit only new content. - """ - - supports_required_and_named = False - - def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): - super().__init__(tokenizer, tools) - # Stateful streaming fields - self.current_tool_name_sent: bool = False - self.prev_tool_call_arr: list[dict[str, Any]] = [] - self.current_tool_id: int = -1 - self.streamed_args_for_tool: list[str] = [] - - self.tool_call_start_token: str = "" - self.tool_call_end_token: str = "" - self.arg_key_start: str = "" - self.arg_key_end: str = "" - self.arg_val_start: str = "" - self.arg_val_end: str = "" - - self.tool_calls_start_token = self.tool_call_start_token - - self.func_call_regex = re.compile(r".*?", re.DOTALL) - self.func_detail_regex = re.compile( - r"([^\n]*)\n(.*)", re.DOTALL - ) - self.func_arg_regex = re.compile( - r"(.*?)\s*(.*?)", re.DOTALL - ) - - if not self.model_tokenizer: - raise ValueError( - "The model tokenizer must be passed to the ToolParser " - "constructor during construction." - ) - - self.tool_call_start_token_id = self.vocab.get(self.tool_call_start_token) - self.tool_call_end_token_id = self.vocab.get(self.tool_call_end_token) - - # Pre-compiled pattern for finding the last ... - # before a partial (used in _build_args_json_so_far). - self._arg_key_pattern = re.compile( - re.escape(self.arg_key_start) + r"(.*?)" + re.escape(self.arg_key_end), - re.DOTALL, - ) - - # Streaming state for re-parse-and-diff approach - self._sent_content_idx: int = 0 - self._tool_call_ids: list[str] = [] - - @staticmethod - def _deserialize(value: str) -> Any: - try: - return json.loads(value) - except json.JSONDecodeError: - pass - - try: - return safe_literal_eval(value) - except (ValueError, SyntaxError): - pass - - return value - - @staticmethod - def _json_escape_string_content(s: str) -> str: - """JSON-escape string content for incremental streaming. - - This escapes the content that goes INSIDE a JSON string (between quotes), - not including the surrounding quotes themselves. - """ - if not s: - return "" - return json.dumps(s, ensure_ascii=False)[1:-1] - - def _is_string_type(self, tool_name: str, arg_name: str) -> bool: - tool_properties = find_tool_properties(self.tools, tool_name) - param_schema = tool_properties.get(arg_name) - if param_schema is None: - return False - param_types = extract_types_from_schema(param_schema) - return set(param_types) - {"null"} == {"string"} - - @staticmethod - def _tools_enabled(request: ChatCompletionRequest) -> bool: - """Return whether tool parsing should be applied for this request.""" - try: - tools = getattr(request, "tools", None) - tool_choice = getattr(request, "tool_choice", None) - return bool(tools) and tool_choice != "none" - except Exception: - logger.exception("Failed to determine if tools are enabled.") - return False - - def adjust_request( - self, request: ChatCompletionRequest | ResponsesRequest - ) -> ChatCompletionRequest | ResponsesRequest: - """Adjust request parameters for tool call token handling. - - For required/named tool_choice, skip setting structured_outputs - because GLM models output tool calls in XML format (per chat - template). Guided decoding would force JSON output, conflicting - with the XML format and causing parsing failures. - """ - if request.tools: - tc = request.tool_choice - if tc == "required" or isinstance(tc, ChatCompletionNamedToolChoiceParam): - # Do NOT call super().adjust_request() for required/named, - # because it would set structured_outputs and force JSON - # output via guided decoding. GLM models use XML tool-call - # syntax (defined in the chat template), so guided decoding - # must be skipped to let the model output XML freely. - # The tool_parser handles extraction from XML output. - if request.tool_choice != "none": - request.skip_special_tokens = False - return request - request = super().adjust_request(request) - if request.tools and request.tool_choice != "none": - # Ensure tool call tokens (, ) are not skipped - # during decoding. Even though they are not marked as special tokens, - # setting skip_special_tokens=False ensures proper handling in - # transformers 5.x where decoding behavior may have changed. - request.skip_special_tokens = False - return request - - def extract_tool_calls( - self, - model_output: str, - request: ChatCompletionRequest, - ) -> ExtractedToolCallInformation: - matched_tool_calls = self.func_call_regex.findall(model_output) - logger.debug("model_output: %s", model_output) - try: - tool_calls: list[ToolCall] = [] - for match in matched_tool_calls: - tc_detail = self.func_detail_regex.search(match) - if not tc_detail: - logger.warning( - "Failed to parse tool call details from: %s", - match, - ) - continue - tc_name = tc_detail.group(1).strip() - tc_args = tc_detail.group(2) - pairs = self.func_arg_regex.findall(tc_args) if tc_args else [] - arg_dct: dict[str, Any] = {} - for key, value in pairs: - arg_key = key.strip() - if self._is_string_type(tc_name, arg_key): - arg_val = value - else: - arg_val = self._deserialize(value.strip()) - logger.debug("arg_key = %s, arg_val = %s", arg_key, arg_val) - arg_dct[arg_key] = arg_val - tool_calls.append( - ToolCall( - type="function", - function=FunctionCall( - name=tc_name, - arguments=json.dumps(arg_dct, ensure_ascii=False), - ), - ) - ) - except Exception: - logger.exception("Failed to extract tool call spec") - return ExtractedToolCallInformation( - tools_called=False, tool_calls=[], content=model_output - ) - else: - if len(tool_calls) > 0: - content: str | None = model_output[ - : model_output.find(self.tool_calls_start_token) - ] - # Normalize empty/whitespace-only content to None - if not content or not content.strip(): - content = None - return ExtractedToolCallInformation( - tools_called=True, tool_calls=tool_calls, content=content - ) - return ExtractedToolCallInformation( - tools_called=False, tool_calls=[], content=model_output - ) - - def _extract_content(self, current_text: str) -> str | None: - """Return unsent non-tool-call text, or None. - - Collects all text outside ``...`` regions, - including text between consecutive tool calls. Holds back any - suffix that could be a partial ```` tag. - """ - # Build the "sendable index" — the furthest point we can send - # content up to. We scan through the text collecting segments - # that are outside tool-call regions. - content_segments: list[str] = [] - pos = self._sent_content_idx - - while pos < len(current_text): - start = current_text.find(self.tool_call_start_token, pos) - if start == -1: - # No more tool calls — send up to (len - partial-tag overlap) - tail = current_text[pos:] - overlap = partial_tag_overlap(tail, self.tool_call_start_token) - sendable = tail[: len(tail) - overlap] if overlap else tail - if sendable: - content_segments.append(sendable) - pos = len(current_text) - overlap - break - - # Text before this - if start > pos: - content_segments.append(current_text[pos:start]) - - # Skip past the (or to end if incomplete) - end = current_text.find(self.tool_call_end_token, start) - if end != -1: - pos = end + len(self.tool_call_end_token) - else: - # Incomplete tool call — nothing more to send - pos = start - break - - if content_segments: - self._sent_content_idx = pos - return "".join(content_segments) - # Even if no content, advance past completed tool-call regions - if pos > self._sent_content_idx: - self._sent_content_idx = pos - return None - - def _extract_tool_call_regions(self, text: str) -> list[tuple[str, bool]]: - """Extract ``(inner_text, is_complete)`` for each ```` region.""" - results: list[tuple[str, bool]] = [] - pos = 0 - while True: - start = text.find(self.tool_call_start_token, pos) - if start == -1: - break - inner_start = start + len(self.tool_call_start_token) - end = text.find(self.tool_call_end_token, inner_start) - if end != -1: - results.append((text[inner_start:end], True)) - pos = end + len(self.tool_call_end_token) - else: - # Incomplete tool call — strip partial suffix - raw = text[inner_start:] - overlap = partial_tag_overlap(raw, self.tool_call_end_token) - if overlap: - raw = raw[:-overlap] - results.append((raw, False)) - break - return results - - def _extract_tool_name_from_region(self, inner_text: str) -> str | None: - """Extract the tool name from the beginning of a tool-call region. - - The name is everything before the first ``\\n`` or ````. - Returns ``None`` if the name hasn't fully arrived yet. - """ - nl = inner_text.find("\n") - ak = inner_text.find(self.arg_key_start) - candidates = [i for i in [nl, ak] if i != -1] - if not candidates: - return None - cut = min(candidates) - name = inner_text[:cut].strip() - return name if name else None - - def _build_args_json_so_far( - self, - tool_name: str, - inner_text: str, - is_complete: bool, - ) -> str: - """Build the JSON arguments string from the XML pairs seen so far. - - For complete ``/`` pairs the value is fully - formatted. For the last argument whose ```` has been - opened but not closed, the partial string content is included - (JSON-escaped, with an opening ``"`` but no closing ``"``). - - The closing ``}`` is only appended when ``is_complete`` is True - (i.e. the ```` tag has arrived). - """ - # Find all complete arg pairs - pairs = self.func_arg_regex.findall(inner_text) - - parts: list[str] = [] - for key, value in pairs: - key = key.strip() - key_json = json.dumps(key, ensure_ascii=False) - if self._is_string_type(tool_name, key): - # Don't strip string values — whitespace is significant - # and must match the partial-value path for diffing. - val_json = json.dumps(value, ensure_ascii=False) - else: - val_json = json.dumps( - self._deserialize(value.strip()), ensure_ascii=False - ) - parts.append(f"{key_json}: {val_json}") - - # Check for a partial (incomplete) arg value - # Find the last that isn't closed - last_val_start = inner_text.rfind(self.arg_val_start) - last_val_end = inner_text.rfind(self.arg_val_end) - has_partial_value = last_val_start != -1 and ( - last_val_end == -1 or last_val_end < last_val_start - ) - - if has_partial_value: - # Find the key for this partial value - # Look for the last ... before this - last_key_match = None - for m in self._arg_key_pattern.finditer(inner_text[:last_val_start]): - last_key_match = m - - if last_key_match: - partial_key = last_key_match.group(1).strip() - partial_content_start = last_val_start + len(self.arg_val_start) - partial_content = inner_text[partial_content_start:] - - # Hold back any partial suffix - overlap = partial_tag_overlap(partial_content, self.arg_val_end) - if overlap: - partial_content = partial_content[:-overlap] - - key_json = json.dumps(partial_key, ensure_ascii=False) - if is_complete: - # Tool call finished but is missing - # (malformed output). Treat partial as complete value - # so the diff naturally closes any open quotes. - if self._is_string_type(tool_name, partial_key): - val_json = json.dumps(partial_content, ensure_ascii=False) - else: - val_json = json.dumps( - self._deserialize(partial_content.strip()), - ensure_ascii=False, - ) - parts.append(f"{key_json}: {val_json}") - elif self._is_string_type(tool_name, partial_key): - escaped = self._json_escape_string_content(partial_content) - # Open quote but no close — more content may arrive - parts.append(f'{key_json}: "{escaped}') - else: - # Non-string partial: include raw content, no wrapping - parts.append(f"{key_json}: {partial_content}") - - if not parts: - return "{}" if is_complete else "" - - joined = "{" + ", ".join(parts) - if is_complete: - joined += "}" - return joined - - def _compute_args_diff(self, index: int, args_so_far: str) -> str | None: - """Return new argument text not yet sent for tool *index*, or None.""" - if not args_so_far or len(args_so_far) <= len( - self.streamed_args_for_tool[index] - ): - return None - diff = args_so_far[len(self.streamed_args_for_tool[index]) :] - self.streamed_args_for_tool[index] = args_so_far - self.prev_tool_call_arr[index]["arguments"] = args_so_far - return diff - - def _ensure_tool_state_for(self, index: int) -> None: - """Grow state arrays so that *index* is valid.""" - while len(self._tool_call_ids) <= index: - self._tool_call_ids.append( - make_tool_call_id(id_type="random", func_name=None, idx=None) - ) - while len(self.streamed_args_for_tool) <= index: - self.streamed_args_for_tool.append("") - while len(self.prev_tool_call_arr) <= index: - self.prev_tool_call_arr.append({}) - - 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 not self._tools_enabled(request): - return DeltaMessage(content=delta_text) if delta_text else None - - content = self._extract_content(current_text) - regions = self._extract_tool_call_regions(current_text) - tool_call_deltas: list[DeltaToolCall] = [] - - for i, (inner_text, is_complete) in enumerate(regions): - self._ensure_tool_state_for(i) - - # Extract tool name - tool_name = self._extract_tool_name_from_region(inner_text) - if not tool_name: - break - - # Emit tool name (once per tool call) - if "name" not in self.prev_tool_call_arr[i]: - self.prev_tool_call_arr[i]["name"] = tool_name - tool_call_deltas.append( - DeltaToolCall( - index=i, - id=self._tool_call_ids[i], - type="function", - function=DeltaFunctionCall( - name=tool_name, - arguments="", - ).model_dump(exclude_none=True), - ) - ) - - # Build args JSON so far, diff, emit - args_so_far = self._build_args_json_so_far( - tool_name, inner_text, is_complete - ) - diff = self._compute_args_diff(i, args_so_far) - if diff: - tool_call_deltas.append( - DeltaToolCall( - index=i, - function=DeltaFunctionCall(arguments=diff).model_dump( - exclude_none=True - ), - ) - ) - - # Update current_tool_id for serving layer compatibility - if regions: - self.current_tool_id = len(regions) - 1 - - if content or tool_call_deltas: - return DeltaMessage( - content=content, - tool_calls=tool_call_deltas, - ) - return None From 21da47dabe27559bf46b80ff6caacafd9dde6035 Mon Sep 17 00:00:00 2001 From: Divakar Verma <137818590+divakar-amd@users.noreply.github.com> Date: Thu, 18 Jun 2026 12:50:32 -0400 Subject: [PATCH 05/75] [ROCm][CI] move lora%N test to mi300 and gate (#45970) Signed-off-by: Divakar Verma --- .buildkite/test-amd.yaml | 30 ++++++++++++++---------------- .buildkite/test_areas/lora.yaml | 11 +++++++++++ 2 files changed, 25 insertions(+), 16 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index e8c2d57fd97..a7f3d67e79f 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -415,22 +415,6 @@ steps: commands: - pytest -v -s kernels/mamba -#----------------------------------------------------------- mi250 · lora ------------------------------------------------------------# - -- label: LoRA %N # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - parallelism: 4 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/lora - - tests/lora - - vllm/platforms/rocm.py - commands: - - pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_qwen3_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py - #------------------------------------------------------ mi250 · models / basic -------------------------------------------------------# - label: Basic Models Test (Other CPU) # TBD @@ -1699,6 +1683,20 @@ steps: #----------------------------------------------------------- mi300 · lora ------------------------------------------------------------# +- label: LoRA %N # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + parallelism: 4 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/lora + - tests/lora + - vllm/platforms/rocm.py + commands: + - pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_qwen3_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py + - label: LoRA TP (Distributed) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] diff --git a/.buildkite/test_areas/lora.yaml b/.buildkite/test_areas/lora.yaml index 3ccf92f9a7a..bd437c52265 100644 --- a/.buildkite/test_areas/lora.yaml +++ b/.buildkite/test_areas/lora.yaml @@ -12,6 +12,17 @@ steps: commands: - pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_qwen3_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py parallelism: 4 + mirror: + amd: + device: mi325_1 + working_dir: "/vllm-workspace/tests" + timeout_in_minutes: 60 + source_file_dependencies: + - vllm/lora + - tests/lora + - vllm/platforms/rocm.py + depends_on: + - image-build-amd - label: LoRA TP (Distributed) From 4583630b562124c033551e5630a7ab3d607a6f03 Mon Sep 17 00:00:00 2001 From: Humphrey Date: Thu, 18 Jun 2026 11:58:22 -0500 Subject: [PATCH 06/75] [Bugfix][Kernel] Check output alignment in vectorize_with_alignment (fixes misaligned-address crash for non-multiple-of-8 head sizes) (#45466) Signed-off-by: HumphreySun98 --- .../quantization/vectorization_utils.cuh | 28 +++++++++++--- tests/kernels/attention/test_cache.py | 37 +++++++++++++++++++ 2 files changed, 60 insertions(+), 5 deletions(-) diff --git a/csrc/libtorch_stable/quantization/vectorization_utils.cuh b/csrc/libtorch_stable/quantization/vectorization_utils.cuh index 98b491b7e23..0cc89bf289d 100644 --- a/csrc/libtorch_stable/quantization/vectorization_utils.cuh +++ b/csrc/libtorch_stable/quantization/vectorization_utils.cuh @@ -24,13 +24,21 @@ __device__ inline void vectorize_with_alignment( ScaOp&& scalar_op) { // InT -> OutT static_assert(VEC_SIZE > 0 && (VEC_SIZE & (VEC_SIZE - 1)) == 0, "VEC_SIZE must be a positive power-of-two"); - constexpr int WIDTH = VEC_SIZE * sizeof(InT); // eg: 64 B + constexpr int WIDTH = VEC_SIZE * sizeof(InT); // eg: 16 B + constexpr int OUT_WIDTH = VEC_SIZE * sizeof(OutT); // eg: 16 B uintptr_t addr = reinterpret_cast(in); + uintptr_t out_addr = reinterpret_cast(out); - // fast path when the whole region is already aligned - // Note: currently the output is guaranteed to be same as the input, so we - // don't check it here, comments here just for future reference. - bool can_vec = ((addr & (WIDTH - 1)) == 0) && ((len & (VEC_SIZE - 1)) == 0); + // fast path when input and output are both fully aligned. The vector + // load/store below go through vec_n_t, declared + // __align__(VEC_SIZE * sizeof(T)), so each side must be aligned to its + // own vector width. out is NOT generally co-aligned with in: e.g. + // reshape_and_cache_flash writes KV-cache rows whose byte offset is a + // multiple of head_size, which for head sizes that are not a multiple + // of VEC_SIZE puts some rows off the vector-width boundary. + bool can_vec = ((addr & (WIDTH - 1)) == 0) && + ((out_addr & (OUT_WIDTH - 1)) == 0) && + ((len & (VEC_SIZE - 1)) == 0); if (can_vec) { int num_vec = len / VEC_SIZE; @@ -55,6 +63,16 @@ __device__ inline void vectorize_with_alignment( prefix_elems /= sizeof(InT); prefix_elems = min(prefix_elems, len); // 0 ≤ prefix < 16 + // the prefix below aligns in; if that does not also align out (their + // addresses differ modulo the vector width), vectorizing is impossible + // and the whole copy must stay scalar. + if (((out_addr + prefix_elems * sizeof(OutT)) & (OUT_WIDTH - 1)) != 0) { + for (int i = tid; i < len; i += stride) { + scalar_op(out[i], in[i]); + } + return; + } + // 1. prefill the when it is unsafe to vectorize for (int i = tid; i < prefix_elems; i += stride) { scalar_op(out[i], in[i]); diff --git a/tests/kernels/attention/test_cache.py b/tests/kernels/attention/test_cache.py index 9b022a042c8..4cbeb7a0b97 100644 --- a/tests/kernels/attention/test_cache.py +++ b/tests/kernels/attention/test_cache.py @@ -428,6 +428,43 @@ def test_reshape_and_cache_flash( torch.testing.assert_close(value_cache_compact, cloned_value_cache) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE) +@pytest.mark.parametrize("kv_cache_layout", CACHE_LAYOUTS) +@pytest.mark.parametrize("implementation", RESHAPE_FLASH_IMPLEMENTATIONS) +@torch.inference_mode() +def test_reshape_and_cache_flash_unaligned_rows( + kv_cache_factory_flashinfer, + dtype: torch.dtype, + kv_cache_dtype: str, + kv_cache_layout: str, + implementation: str, +) -> None: + """Regression test for https://github.com/vllm-project/vllm/issues/41257. + + head_size=46 with num_heads=13 places KV-cache rows at byte offsets + that are not a multiple of the vector width (NHD row pitch + 13*46*itemsize, HND head pitch 46*itemsize), unlike HEAD_SIZES above + which are all 16-byte multiples. The CUDA kernel used to issue + vectorized stores to those rows -> CUDA misaligned address. + """ + test_reshape_and_cache_flash( + kv_cache_factory_flashinfer, + num_tokens=42, + num_heads=13, + head_size=46, + block_size=16, + num_blocks=128, + dtype=dtype, + seed=0, + device=CUDA_DEVICES[0], + kv_cache_dtype=kv_cache_dtype, + kv_cache_layout=kv_cache_layout, + kv_scale_type="tensor", + implementation=implementation, + ) + + @pytest.mark.parametrize("direction", COPYING_DIRECTION) @pytest.mark.parametrize("num_mappings", NUM_MAPPINGS) @pytest.mark.parametrize("num_heads", NUM_HEADS) From 25faa1f4cc2ec5d0db50b2b2b04c43f58d8a0931 Mon Sep 17 00:00:00 2001 From: qli88 Date: Thu, 18 Jun 2026 11:59:09 -0500 Subject: [PATCH 07/75] [CI]Enable mxfp4 lora test for ROCm platform (#43802) Signed-off-by: Qiang Li --- tests/lora/test_gptoss_tp.py | 36 +++++++++++++++++++++++++----------- 1 file changed, 25 insertions(+), 11 deletions(-) diff --git a/tests/lora/test_gptoss_tp.py b/tests/lora/test_gptoss_tp.py index 7aa8643cd9c..838c3ab7dd9 100644 --- a/tests/lora/test_gptoss_tp.py +++ b/tests/lora/test_gptoss_tp.py @@ -70,17 +70,20 @@ def generate_and_test(llm: vllm.LLM, lora_path: str, lora_id: int) -> None: assert generated_texts[i].startswith(EXPECTED_LORA_OUTPUT[i]) -@pytest.mark.skipif( - not current_platform.is_cuda(), - reason=( - "Mxfp4 LoRA on ROCm is blocked by a spawn compatibility issue. " - "The fused_moe_lora Triton kernel crashes in spawned subprocesses, " - "and vLLM forces spawn mode when HIP is initialized before " - "multiprocessing. Fixing this requires either making the LoRA " - "Triton kernel spawn-safe or pre-warming the kernel cache." - ), +# TODO: make the Mxfp4MoeBackend.TRITON spawn-safe. +# For now just use TRITON_UNFUSED kernel +@pytest.mark.parametrize( + "mxfp4_use_marlin", + [ + False, + pytest.param( + True, + marks=pytest.mark.skipif( + current_platform.is_rocm(), reason="marlin not supported" + ), + ), + ], ) -@pytest.mark.parametrize("mxfp4_use_marlin", [True, False]) @pytest.mark.parametrize("specialize_active_lora", [True, False]) def test_gpt_oss_lora( gptoss20b_lora_files, @@ -109,7 +112,18 @@ def test_gpt_oss_lora( @multi_gpu_test(num_gpus=2) @pytest.mark.parametrize("fully_sharded_loras", [False, True]) -@pytest.mark.parametrize("mxfp4_use_marlin", [True, False]) +@pytest.mark.parametrize( + "mxfp4_use_marlin", + [ + False, + pytest.param( + True, + marks=pytest.mark.skipif( + current_platform.is_rocm(), reason="marlin not supported" + ), + ), + ], +) def test_gpt_oss_lora_tp2( gptoss20b_lora_files, fully_sharded_loras, From e2352c29743aeec4a2dafc66c4fdd0e10b37072e Mon Sep 17 00:00:00 2001 From: stefankoncarevic Date: Thu, 18 Jun 2026 18:59:37 +0200 Subject: [PATCH 08/75] [ROCm][Spec Decode] Fix probabilistic draft probs test attention backend (#45706) Signed-off-by: Stefan Koncarevic --- .buildkite/test_areas/misc.yaml | 6 ++++++ tests/v1/spec_decode/test_eagle.py | 8 ++++++-- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 67fecf06df3..7db72be7b52 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -21,6 +21,12 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn # TODO: create another `optional` test group for slow tests - pytest -v -s -m 'not slow_test' v1/spec_decode + mirror: + amd: + device: mi300_1 + timeout_in_minutes: 65 + depends_on: + - image-build-amd - label: V1 Sample + Logits key: v1-sample-logits diff --git a/tests/v1/spec_decode/test_eagle.py b/tests/v1/spec_decode/test_eagle.py index 848130725ac..fecb72800e0 100644 --- a/tests/v1/spec_decode/test_eagle.py +++ b/tests/v1/spec_decode/test_eagle.py @@ -1002,7 +1002,11 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch): assert torch.equal(result, expected_tokens) -def test_propose_stores_probabilistic_draft_probs(monkeypatch): +@pytest.mark.parametrize( + "attn_backend", + ["ROCM_ATTN", "TRITON_ATTN"] if current_platform.is_rocm() else ["FLASH_ATTN"], +) +def test_propose_stores_probabilistic_draft_probs(attn_backend, monkeypatch): device = torch.device(DEVICE_TYPE) batch_size = 2 seq_lens = [5, 3] @@ -1053,7 +1057,7 @@ def test_propose_stores_probabilistic_draft_probs(monkeypatch): ) attn_metadata_builder_cls, _ = try_get_attention_backend( - AttentionBackendEnum.FLASH_ATTN + AttentionBackendEnum[attn_backend] ) attn_metadata_builder = attn_metadata_builder_cls( kv_cache_spec=create_standard_kv_cache_spec(proposer.vllm_config), From a0df04e4775efbfebd65c997259d63af0ec548ce Mon Sep 17 00:00:00 2001 From: Palaiologos1453 <2260891073@qq.com> Date: Fri, 19 Jun 2026 01:37:39 +0800 Subject: [PATCH 09/75] [Tests] Add Qwen3 streaming parser delta boundary cases (#45708) Signed-off-by: test test <2260891073@qq.com> --- .../test_qwen3coder_tool_parser.py | 67 +++++++++++++++++++ 1 file changed, 67 insertions(+) diff --git a/tests/tool_parsers/test_qwen3coder_tool_parser.py b/tests/tool_parsers/test_qwen3coder_tool_parser.py index ac770ff8e5b..1f5e51412b9 100644 --- a/tests/tool_parsers/test_qwen3coder_tool_parser.py +++ b/tests/tool_parsers/test_qwen3coder_tool_parser.py @@ -1300,6 +1300,73 @@ def test_streaming_multi_param_single_chunk(qwen3_tool_parser, qwen3_tokenizer): assert args["unit"] == "fahrenheit" +def test_streaming_complete_tool_call_single_delta(qwen3_tool_parser): + """Regression: one delta may contain a complete tool call.""" + request = ChatCompletionRequest(model=MODEL, messages=[]) + + from tests.tool_parsers.utils import ( + run_tool_extraction_streaming, + ) + + reconstructor = run_tool_extraction_streaming( + qwen3_tool_parser, + [ + ( + "\n" + "\n" + "\nDallas\n\n" + "\nTX\n\n" + "\n" + "" + ) + ], + request, + assert_one_tool_per_delta=False, + ) + + assert len(reconstructor.tool_calls) == 1 + assert reconstructor.tool_calls[0].function.name == "get_current_weather" + args = json.loads(reconstructor.tool_calls[0].function.arguments) + assert args == {"city": "Dallas", "state": "TX"} + + +def test_streaming_next_tool_call_starts_in_close_delta(qwen3_tool_parser): + """Regression: a close delta may also contain the next tool call.""" + request = ChatCompletionRequest(model=MODEL, messages=[]) + + from tests.tool_parsers.utils import ( + run_tool_extraction_streaming, + ) + + reconstructor = run_tool_extraction_streaming( + qwen3_tool_parser, + [ + "\n", + "\n", + "\nDallas\n\n", + "\nTX\n\n", + "", + ( + "\n\n" + "\n" + "\n" + "\nOrlando\n\n" + "\nFL\n\n" + "\n" + "" + ), + ], + request, + assert_one_tool_per_delta=False, + ) + + assert len(reconstructor.tool_calls) == 2 + first_args = json.loads(reconstructor.tool_calls[0].function.arguments) + second_args = json.loads(reconstructor.tool_calls[1].function.arguments) + assert first_args == {"city": "Dallas", "state": "TX"} + assert second_args == {"city": "Orlando", "state": "FL"} + + def test_no_double_serialization_string_args(qwen3_tool_parser): """Regression: string arguments must not be double-serialized (PR #35615).""" tools = [ From ea6078fe6a7242e7a5a89798e617b807d2540466 Mon Sep 17 00:00:00 2001 From: Itay Etelis <92247226+Etelis@users.noreply.github.com> Date: Thu, 18 Jun 2026 20:43:35 +0300 Subject: [PATCH 10/75] [KV Connector][Offloading] Disable parallel-agnostic fs-tier cache on V2 model runner (#46044) Signed-off-by: Itay Etelis Co-authored-by: Itay Etelis --- tests/v1/kv_offload/test_file_mapper.py | 14 ++++++++++++++ tests/v1/kv_offload/tiering/test_obj_tier.py | 1 + vllm/v1/kv_offload/file_mapper.py | 3 +++ 3 files changed, 18 insertions(+) diff --git a/tests/v1/kv_offload/test_file_mapper.py b/tests/v1/kv_offload/test_file_mapper.py index 0e462f8de2b..6f6e0d66196 100644 --- a/tests/v1/kv_offload/test_file_mapper.py +++ b/tests/v1/kv_offload/test_file_mapper.py @@ -64,6 +64,7 @@ def make_mapper_from_offloading_spec(**kwargs) -> FileMapper: "dcp_size", 1 ) mock_vllm_config.parallel_config.rank = kwargs.get("rank", 0) + mock_vllm_config.use_v2_model_runner = kwargs.get("use_v2_model_runner", False) mock_kv_cache_config = MagicMock() mock_kv_cache_config.kv_cache_groups = kwargs.get("kv_cache_groups", []) @@ -210,3 +211,16 @@ def test_parallel_agnostic_excludes_mla(): ) assert fm.fields["tp_size"] == 2 assert fm.rank == 1 + + +def test_parallel_agnostic_disabled_on_v2_model_runner(): + # V2's KV layout is not known to be parallelism-invariant: don't collapse. + fm = make_mapper_from_offloading_spec( + tp_size=2, + rank=1, + kv_cache_groups=[_full_attention_group()], + use_v2_model_runner=True, + parallel_agnostic=True, + ) + assert fm.fields["tp_size"] == 2 + assert fm.rank == 1 diff --git a/tests/v1/kv_offload/tiering/test_obj_tier.py b/tests/v1/kv_offload/tiering/test_obj_tier.py index bac5729eafb..aae3c60c539 100644 --- a/tests/v1/kv_offload/tiering/test_obj_tier.py +++ b/tests/v1/kv_offload/tiering/test_obj_tier.py @@ -37,6 +37,7 @@ def _make_vllm_config(): decode_context_parallel_size=1, rank=0, ), + use_v2_model_runner=False, ) diff --git a/vllm/v1/kv_offload/file_mapper.py b/vllm/v1/kv_offload/file_mapper.py index c19f07ff514..d8fadb09988 100644 --- a/vllm/v1/kv_offload/file_mapper.py +++ b/vllm/v1/kv_offload/file_mapper.py @@ -84,10 +84,13 @@ class FileMapper: ] # Only a single full-attention group is parallelism-invariant. MLA is # excluded: its latent KV is replicated per rank, never head-sharded. + # The V2 model runner is excluded: its KV layout is not known to be + # parallelism-invariant. groups = kv_cache_config.kv_cache_groups spec = groups[0].kv_cache_spec if len(groups) == 1 else None parallel_agnostic = ( parallel_agnostic + and not vllm_config.use_v2_model_runner and isinstance(spec, FullAttentionSpec) and not isinstance(spec, MLAAttentionSpec) ) From 09f3cd5c1080de42c9001803f638852b7f6a4310 Mon Sep 17 00:00:00 2001 From: Ben Browning Date: Thu, 18 Jun 2026 14:04:06 -0400 Subject: [PATCH 11/75] [Bugfix] [Parser] Fix Qwen3 latent bug in partial params dropping values containing `<` (#46047) Signed-off-by: Ben Browning --- tests/parser/engine/test_qwen3.py | 18 ++++++++++++++++++ vllm/parser/qwen3.py | 2 +- 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/tests/parser/engine/test_qwen3.py b/tests/parser/engine/test_qwen3.py index 7c2255ac7b2..06784212e1b 100644 --- a/tests/parser/engine/test_qwen3.py +++ b/tests/parser/engine/test_qwen3.py @@ -615,6 +615,24 @@ class TestArgConverter: assert result["command"] == "ls -la" assert result["desc"] == "\npartial value" + def test_partial_value_with_angle_bracket(self): + from vllm.parser.qwen3 import ( + _qwen3_arg_converter, + ) + + raw = "x<5" + result = json.loads(_qwen3_arg_converter(raw, partial=True)) + assert result == {"expr": "x<5"} + + def test_partial_value_with_angle_bracket_and_complete_param(self): + from vllm.parser.qwen3 import ( + _qwen3_arg_converter, + ) + + raw = "Tokyo\nx<5" + result = json.loads(_qwen3_arg_converter(raw, partial=True)) + assert result == {"city": "Tokyo", "expr": "x<5"} + class TestSchemaAwareTypeCoercion: """Verify that _fix_arg_types corrects miscoerced values using the diff --git a/vllm/parser/qwen3.py b/vllm/parser/qwen3.py index ed47b1b9254..45e3c7e4325 100644 --- a/vllm/parser/qwen3.py +++ b/vllm/parser/qwen3.py @@ -49,7 +49,7 @@ _PARAM_RE = re.compile( r"(?:<\s*/\s*parameter\s*>|(?=<\s*parameter\s*=))", re.DOTALL, ) -_PARTIAL_PARAM_RE = re.compile(r"<\s*parameter\s*=\s*([^>]+)>([^<]*)$", re.DOTALL) +_PARTIAL_PARAM_RE = re.compile(r"<\s*parameter\s*=\s*([^>]+)>(.*)$", re.DOTALL) def _qwen3_arg_converter(raw_args: str, partial: bool) -> str: From 79ca54d2215b22d9a4fc17378eb7aa2b2eb9dbd1 Mon Sep 17 00:00:00 2001 From: Ting SUN Date: Fri, 19 Jun 2026 02:18:25 +0800 Subject: [PATCH 12/75] [Bugfix][Quantization] Don't reject fp8_e5m2 KV cache for non-fp8 quantized checkpoints (#45040) Signed-off-by: Ting Sun Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../layers/attention/attention.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/attention/attention.py b/vllm/model_executor/layers/attention/attention.py index 5974e09624d..cdfe9fa1bce 100644 --- a/vllm/model_executor/layers/attention/attention.py +++ b/vllm/model_executor/layers/attention/attention.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast import torch import torch.nn as nn @@ -166,7 +166,21 @@ def _init_kv_cache_quant( # TODO (mgoin): kv cache dtype should be specified in the FP8 # checkpoint config and become the "auto" behavior if layer.kv_cache_dtype == "fp8_e5m2": - raise ValueError("fp8_e5m2 kv-cache is not supported with fp8 checkpoints.") + # A compressed-tensors checkpoint stores fp8 KV scales only when it + # declares a kv_cache_scheme; weight-only ones declare none and must + # keep fp8_e5m2, the only fp8 KV dtype usable on Ampere. + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501 + CompressedTensorsConfig, + CompressedTensorsKVCacheMethod, + ) + + if not isinstance(quant_method, CompressedTensorsKVCacheMethod) or ( + cast(CompressedTensorsConfig, quant_method.quant_config).kv_cache_scheme + is not None + ): + raise ValueError( + "fp8_e5m2 kv-cache is not supported with fp8 checkpoints." + ) # If quantization is enabled, we make "k_scale" and "v_scale" # parameters so that it can be loaded from the model checkpoint. # The k/v_scale will then be converted back to native float32 From b53b1c7ffe7aebdafd0876350f30e51d1226c92a Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Thu, 18 Jun 2026 15:20:44 -0400 Subject: [PATCH 13/75] [Model Runner V2] Migration to support quantized model by default [5/N] (#44446) Signed-off-by: yewentao256 --- tests/test_config.py | 2 +- vllm/config/vllm.py | 3 --- 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/tests/test_config.py b/tests/test_config.py index d992ac29696..eb9b11535b8 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -188,7 +188,7 @@ def test_v2_model_runner_env_tri_state(monkeypatch, env_value, expected): is_moe=False, is_quantized=True, ), - False, + True, ), ( SimpleNamespace( diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index ba20d75fa11..ba7d26c93b2 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -555,9 +555,6 @@ class VllmConfig: if model_config.runner_type != "generate": return False - if model_config.is_quantized: - return False - architectures = getattr(model_config, "architectures", []) return any( arch in DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES for arch in architectures From f6ba7209632936d4908499afc799e96f6eee2725 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Juan=20P=C3=A9rez=20de=20Algaba?= <124347725+jperezdealgaba@users.noreply.github.com> Date: Thu, 18 Jun 2026 21:35:13 +0200 Subject: [PATCH 14/75] (security) Upgrade Starlette to >= 1.0.1 to fix CVE-2026-48710 (#45675) Signed-off-by: jperezde Co-authored-by: Isotr0py --- requirements/common.txt | 5 +-- requirements/test/cuda.txt | 71 ++++++++++++-------------------------- requirements/test/rocm.txt | 68 +++++++++++------------------------- requirements/test/xpu.txt | 3 +- 4 files changed, 49 insertions(+), 98 deletions(-) diff --git a/requirements/common.txt b/requirements/common.txt index fde1ba4f0c9..a5d74e14e64 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -11,13 +11,14 @@ transformers >= 5.5.3 tokenizers >= 0.21.1 # Required for fast incremental detokenization. safetensors >= 0.6.2 # MXFP4/MXFP6 dtype support (F8_E8M0, F4) added in 0.6.0: https://github.com/huggingface/safetensors/pull/611 protobuf >= 5.29.6, !=6.30.*, !=6.31.*, !=6.32.*, !=6.33.0.*, !=6.33.1.*, !=6.33.2.*, !=6.33.3.*, !=6.33.4.* # Required by LlamaTokenizer, gRPC. CVE-2026-0994 -fastapi[standard] >= 0.115.0 # Required by FastAPI's form models in the OpenAI API server's audio transcriptions endpoint. +fastapi[standard] >= 0.133.0, < 0.137.0 # First version supporting Starlette 1.0; < 0.137.0 avoids route-tree change that breaks model-hosting-container-standards handler overrides. +starlette >= 1.0.1 # CVE-2026-48710: Host header injection in < 1.0.1 aiohttp >= 3.13.3 openai >= 2.0.0 # For Responses API with reasoning content pydantic >= 2.12.0 prometheus_client >= 0.18.0 pillow # Required for image processing -prometheus-fastapi-instrumentator >= 7.0.0 +prometheus-fastapi-instrumentator >= 8.0.0 # v8 unblocks starlette >= 1.0 tiktoken >= 0.6.0 # Required for DBRX tokenizer lm-format-enforcer == 0.11.3 llguidance >= 1.7.0, < 1.8.0; platform_machine == "x86_64" or platform_machine == "arm64" or platform_machine == "aarch64" or platform_machine == "ppc64le" diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index c6d9ed24adb..76c343b91b1 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -35,14 +35,11 @@ arctic-inference==0.1.1 # via -r requirements/test/cuda.in argcomplete==3.5.1 # via datamodel-code-generator -arrow==1.3.0 - # via isoduration attrs==24.2.0 # via # aiohttp # hypothesis # jsonschema - # pytest-subtests # referencing audioread==3.0.1 # via librosa @@ -57,9 +54,7 @@ azure-identity==1.25.2 azure-storage-blob==12.28.0 # via runai-model-streamer-azure backoff==2.2.1 - # via - # -r requirements/test/cuda.in - # schemathesis + # via -r requirements/test/cuda.in bitsandbytes==0.49.2 # via -r requirements/test/cuda.in black==24.10.0 @@ -110,7 +105,6 @@ colorama==0.4.6 # via # perceptron # sacrebleu - # schemathesis colorful==0.5.6 # via ray colorlog==6.10.1 @@ -183,7 +177,7 @@ et-xmlfile==2.0.0 # via openpyxl evaluate==0.4.3 # via lm-eval -fastapi==0.128.0 +fastapi==0.136.3 # via # -c requirements/common.txt # gpt-oss @@ -206,8 +200,6 @@ filelock==3.16.1 # virtualenv fonttools==4.55.0 # via matplotlib -fqdn==1.5.1 - # via jsonschema frozendict==2.4.6 # via einx frozenlist==1.5.0 @@ -269,7 +261,7 @@ h11==0.14.0 # uvicorn h2==4.3.0 # via httpx -harfile==0.3.0 +harfile==0.5.0 # via schemathesis hf-xet==1.4.3 # via huggingface-hub @@ -309,7 +301,7 @@ hypothesis==6.131.0 # hypothesis-graphql # hypothesis-jsonschema # schemathesis -hypothesis-graphql==0.11.1 +hypothesis-graphql==0.13.0 # via schemathesis hypothesis-jsonschema==0.23.1 # via schemathesis @@ -318,7 +310,6 @@ idna==3.10 # anyio # email-validator # httpx - # jsonschema # requests # yarl imagehash==4.3.2 @@ -335,8 +326,6 @@ instanttensor==0.1.5 # via -r requirements/test/cuda.in isodate==0.7.2 # via azure-storage-blob -isoduration==20.11.0 - # via jsonschema isort==5.13.2 # via datamodel-code-generator jinja2==3.1.6 @@ -356,15 +345,14 @@ joblib==1.4.2 # librosa # nltk # scikit-learn -jsonpointer==3.0.0 - # via jsonschema jsonschema==4.23.0 # via # -c requirements/common.txt # hypothesis-jsonschema # mistral-common # ray - # schemathesis +jsonschema-rs==0.46.5 + # via schemathesis jsonschema-specifications==2024.10.1 # via jsonschema junit-xml==1.9 @@ -715,18 +703,20 @@ pydantic-core==2.41.1 pydantic-extra-types==2.10.5 # via mistral-common pygments==2.18.0 - # via rich + # via + # pytest + # rich pyjwt==2.11.0 # via msal pyparsing==3.2.0 # via matplotlib -pyrate-limiter==3.7.0 +pyrate-limiter==4.4.0 # via schemathesis pystemmer==3.0.0 # via mteb pytablewriter==1.2.0 # via lm-eval -pytest==8.3.5 +pytest==9.1.0 # via # -r requirements/test/cuda.in # buildkite-test-collector @@ -737,10 +727,9 @@ pytest==8.3.5 # pytest-mock # pytest-rerunfailures # pytest-shard - # pytest-subtests # pytest-timeout # schemathesis -pytest-asyncio==0.24.0 +pytest-asyncio==1.4.0 # via -r requirements/test/cuda.in pytest-cov==6.3.0 # via -r requirements/test/cuda.in @@ -752,13 +741,10 @@ pytest-rerunfailures==14.0 # via -r requirements/test/cuda.in pytest-shard==0.1.2 # via -r requirements/test/cuda.in -pytest-subtests==0.14.1 - # via schemathesis pytest-timeout==2.3.1 # via -r requirements/test/cuda.in python-dateutil==2.9.0.post0 # via - # arrow # botocore # matplotlib # pandas @@ -829,15 +815,12 @@ requests==2.32.3 # tiktoken responses==0.25.3 # via genai-perf -rfc3339-validator==0.1.4 - # via jsonschema -rfc3987==1.3.8 - # via jsonschema rich==13.9.4 # via # genai-perf # mteb # perceptron + # schemathesis # typer rouge-score==0.1.2 # via lm-eval @@ -868,7 +851,7 @@ safetensors==0.7.0 # segmentation-models-pytorch # timm # transformers -schemathesis==3.39.15 +schemathesis==4.21.6 # via -r requirements/test/cuda.in scikit-image==0.25.2 # via albumentations @@ -912,7 +895,6 @@ six==1.16.0 # junit-xml # opencensus # python-dateutil - # rfc3339-validator # rouge-score smart-open==7.1.0 # via ray @@ -938,10 +920,10 @@ sqlalchemy==2.0.41 # optuna sqlitedict==2.1.0 # via lm-eval -starlette==0.50.0 +starlette==1.3.1 # via + # -c requirements/common.txt # fastapi - # schemathesis # starlette-testclient starlette-testclient==0.4.1 # via schemathesis @@ -966,6 +948,7 @@ tenacity==9.1.2 # gpt-oss # lm-eval # plotly + # schemathesis tensorizer==2.10.1 # via -r requirements/test/cuda.in termcolor==3.1.0 @@ -990,10 +973,6 @@ tokenizers==0.22.2 # -c requirements/common.txt # -r requirements/test/cuda.in # transformers -tomli==2.2.1 - # via schemathesis -tomli-w==1.2.0 - # via schemathesis torch==2.11.0+cu130 # via # -c requirements/cuda.txt @@ -1066,8 +1045,6 @@ typer==0.15.2 # huggingface-hub # perceptron # transformers -types-python-dateutil==2.9.0.20241206 - # via arrow typing-extensions==4.15.0 # via # -c requirements/common.txt @@ -1092,6 +1069,8 @@ typing-extensions==4.15.0 # pydantic # pydantic-core # pydantic-extra-types + # pytest-asyncio + # schemathesis # sentence-transformers # sqlalchemy # starlette @@ -1099,11 +1078,11 @@ typing-extensions==4.15.0 # typer # typing-inspection typing-inspection==0.4.2 - # via pydantic + # via + # fastapi + # pydantic tzdata==2024.2 # via pandas -uri-template==1.3.0 - # via jsonschema urllib3==2.2.3 # via # blobfile @@ -1122,8 +1101,6 @@ vocos==0.1.0 # via -r requirements/test/cuda.in wcwidth==0.2.13 # via ftfy -webcolors==24.11.1 - # via jsonschema werkzeug==3.1.3 # via schemathesis word2number==1.1 @@ -1135,8 +1112,6 @@ xxhash==3.5.0 # datasets # evaluate yarl==1.17.1 - # via - # aiohttp - # schemathesis + # via aiohttp zipp==3.23.0 # via importlib-metadata diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt index 879a3286444..842d2ff3188 100644 --- a/requirements/test/rocm.txt +++ b/requirements/test/rocm.txt @@ -51,15 +51,12 @@ arctic-inference==0.1.1 # via -r requirements/test/rocm.in argcomplete==3.6.3 # via datamodel-code-generator -arrow==1.4.0 - # via isoduration astor==0.8.1 # via depyf attrs==26.1.0 # via # aiohttp # jsonschema - # pytest-subtests # referencing audioread==3.0.1 # via librosa @@ -74,9 +71,7 @@ azure-identity==1.25.3 azure-storage-blob==12.28.0 # via runai-model-streamer-azure backoff==2.2.1 - # via - # -r requirements/test/rocm.in - # schemathesis + # via -r requirements/test/rocm.in bitsandbytes==0.49.2 # via -r requirements/test/rocm.in black==26.3.1 @@ -139,7 +134,6 @@ colorama==0.4.6 # via # perceptron # sacrebleu - # schemathesis colorful==0.5.8 # via ray colorlog==6.10.1 @@ -258,8 +252,6 @@ filelock==3.25.2 # virtualenv fonttools==4.62.1 # via matplotlib -fqdn==1.5.1 - # via jsonschema frozendict==2.4.7 # via einx frozenlist==1.8.0 @@ -328,7 +320,7 @@ h11==0.16.0 # uvicorn h2==4.3.0 # via httpx -harfile==0.4.0 +harfile==0.5.0 # via schemathesis hf-xet==1.4.3 # via huggingface-hub @@ -378,7 +370,7 @@ hypothesis==6.151.9 # hypothesis-graphql # hypothesis-jsonschema # schemathesis -hypothesis-graphql==0.12.0 +hypothesis-graphql==0.13.0 # via schemathesis hypothesis-jsonschema==0.23.1 # via schemathesis @@ -387,7 +379,6 @@ idna==3.11 # anyio # email-validator # httpx - # jsonschema # requests # yarl ijson==3.5.0 @@ -408,8 +399,6 @@ interegular==0.3.3 # via lm-format-enforcer isodate==0.7.2 # via azure-storage-blob -isoduration==20.11.0 - # via jsonschema isort==8.0.1 # via datamodel-code-generator jinja2==3.1.6 @@ -435,8 +424,6 @@ joblib==1.5.3 # librosa # nltk # scikit-learn -jsonpointer==3.1.0 - # via jsonschema jsonschema==4.26.0 # via # -c requirements/common.txt @@ -445,7 +432,8 @@ jsonschema==4.26.0 # mcp # mistral-common # ray - # schemathesis +jsonschema-rs==0.46.5 + # via schemathesis jsonschema-specifications==2025.9.1 # via jsonschema junit-xml==1.9 @@ -792,7 +780,7 @@ prometheus-client==0.24.1 # opentelemetry-exporter-prometheus # prometheus-fastapi-instrumentator # ray -prometheus-fastapi-instrumentator==7.1.0 +prometheus-fastapi-instrumentator==8.0.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -876,20 +864,22 @@ pydantic-settings==2.13.1 # fastapi # mcp pygments==2.19.2 - # via rich + # via + # pytest + # rich pyjwt==2.12.1 # via # mcp # msal pyparsing==3.3.2 # via matplotlib -pyrate-limiter==3.9.0 +pyrate-limiter==4.4.0 # via schemathesis pystemmer==3.0.0 # via mteb pytablewriter==1.2.1 # via lm-eval -pytest==8.3.5 +pytest==9.1.0 # via # -r requirements/test/rocm.in # buildkite-test-collector @@ -900,10 +890,9 @@ pytest==8.3.5 # pytest-mock # pytest-rerunfailures # pytest-shard - # pytest-subtests # pytest-timeout # schemathesis -pytest-asyncio==0.24.0 +pytest-asyncio==1.4.0 # via -r requirements/test/rocm.in pytest-cov==6.3.0 # via -r requirements/test/rocm.in @@ -915,13 +904,10 @@ pytest-rerunfailures==14.0 # via -r requirements/test/rocm.in pytest-shard==0.1.2 # via -r requirements/test/rocm.in -pytest-subtests==0.14.2 - # via schemathesis pytest-timeout==2.3.1 # via -r requirements/test/rocm.in python-dateutil==2.9.0.post0 # via - # arrow # botocore # matplotlib # pandas @@ -1016,16 +1002,13 @@ requests==2.32.5 # tiktoken responses==0.26.0 # via genai-perf -rfc3339-validator==0.1.4 - # via jsonschema -rfc3987==1.3.8 - # via jsonschema rich==14.3.3 # via # genai-perf # mteb # perceptron # rich-toolkit + # schemathesis # typer rich-toolkit==0.19.7 # via @@ -1063,7 +1046,7 @@ safetensors==0.7.0 # segmentation-models-pytorch # timm # transformers -schemathesis==3.39.15 +schemathesis==4.21.6 # via -r requirements/test/rocm.in scikit-image==0.26.0 # via albumentations @@ -1120,7 +1103,6 @@ six==1.17.0 # junit-xml # opencensus # python-dateutil - # rfc3339-validator # rouge-score smart-open==7.5.1 # via ray @@ -1149,13 +1131,14 @@ sqlitedict==2.1.0 # via lm-eval sse-starlette==3.3.4 # via mcp -starlette==0.52.1 +starlette==1.3.1 # via + # -c requirements/common.txt + # -r requirements/test/../common.txt # fastapi # mcp # model-hosting-container-standards # prometheus-fastapi-instrumentator - # schemathesis # sse-starlette # starlette-testclient starlette-testclient==0.4.1 @@ -1182,6 +1165,7 @@ tenacity==9.1.4 # via # gpt-oss # lm-eval + # schemathesis tensorizer==2.10.1 # via # -c requirements/rocm.txt @@ -1215,10 +1199,6 @@ tokenizers==0.22.2 # -r requirements/test/../common.txt # -r requirements/test/rocm.in # transformers -tomli==2.4.0 - # via schemathesis -tomli-w==1.2.0 - # via schemathesis torch-c-dlpack-ext==0.1.5 # via tilelang tqdm==4.67.3 @@ -1301,8 +1281,10 @@ typing-extensions==4.15.0 # pydantic # pydantic-core # pydantic-extra-types + # pytest-asyncio # referencing # rich-toolkit + # schemathesis # sentence-transformers # sqlalchemy # starlette @@ -1317,10 +1299,6 @@ typing-inspection==0.4.2 # mcp # pydantic # pydantic-settings -tzdata==2025.3 - # via arrow -uri-template==1.3.0 - # via jsonschema urllib3==2.6.3 # via # blobfile @@ -1351,8 +1329,6 @@ watchfiles==1.1.1 # uvicorn wcwidth==0.6.0 # via ftfy -webcolors==25.10.0 - # via jsonschema websockets==16.0 # via uvicorn werkzeug==3.1.6 @@ -1370,9 +1346,7 @@ xxhash==3.6.0 # datasets # evaluate yarl==1.23.0 - # via - # aiohttp - # schemathesis + # via aiohttp z3-solver==4.15.4.0 # via tilelang zipp==3.23.0 diff --git a/requirements/test/xpu.txt b/requirements/test/xpu.txt index 1b1f3c91c5e..40f23b95d10 100644 --- a/requirements/test/xpu.txt +++ b/requirements/test/xpu.txt @@ -593,8 +593,9 @@ soxr==0.5.0.post1 # mistral-common sqlitedict==2.1.0 # via lm-eval -starlette==1.0.0 +starlette==1.3.1 # via + # -c requirements/common.txt # fastapi # starlette-testclient starlette-testclient==0.4.1 From 225936a1dd10586798c0181696d628e7b609ea90 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Thu, 18 Jun 2026 15:37:39 -0400 Subject: [PATCH 15/75] [CI Bug] Revert #42379 to fix CI `Multi-Modal Models (Extended Generation 1)` (#46070) Signed-off-by: yewentao256 --- csrc/libtorch_stable/layernorm_kernels.cu | 13 ++++++----- .../layernorm_quant_kernels.cu | 23 ++++++++++++++----- 2 files changed, 24 insertions(+), 12 deletions(-) diff --git a/csrc/libtorch_stable/layernorm_kernels.cu b/csrc/libtorch_stable/layernorm_kernels.cu index eb121b0b880..f29734fc265 100644 --- a/csrc/libtorch_stable/layernorm_kernels.cu +++ b/csrc/libtorch_stable/layernorm_kernels.cu @@ -81,11 +81,11 @@ __global__ void rms_norm_kernel( #pragma unroll for (int j = 0; j < VEC_SIZE; j++) { float x = static_cast(src1.val[j]); - scalar_t normalized = static_cast(x * s_variance); if constexpr (HasWeight) { - dst.val[j] = normalized * src2.val[j]; + float w = static_cast(src2.val[j]); + dst.val[j] = static_cast(x * s_variance * w); } else { - dst.val[j] = normalized; + dst.val[j] = static_cast(x * s_variance); } } v_out[i] = dst; @@ -151,7 +151,8 @@ fused_add_rms_norm_kernel( #pragma unroll for (int j = 0; j < width; ++j) { float x = Converter::convert(res.data[j]); - out.data[j] = Converter::convert(x * s_variance) * w.data[j]; + float wf = Converter::convert(w.data[j]); + out.data[j] = Converter::convert(x * s_variance * wf); } } else { #pragma unroll @@ -198,8 +199,8 @@ fused_add_rms_norm_kernel( for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) { float x = (float)residual[blockIdx.x * hidden_size + idx]; if constexpr (HasWeight) { - input[blockIdx.x * input_stride + idx] = - (scalar_t)(x * s_variance) * weight[idx]; + float w = (float)weight[idx]; + input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance * w); } else { input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance); } diff --git a/csrc/libtorch_stable/layernorm_quant_kernels.cu b/csrc/libtorch_stable/layernorm_quant_kernels.cu index 32f3495f4e9..26ffa76d6e1 100644 --- a/csrc/libtorch_stable/layernorm_quant_kernels.cu +++ b/csrc/libtorch_stable/layernorm_quant_kernels.cu @@ -66,8 +66,13 @@ __global__ void rms_norm_static_fp8_quant_kernel( #pragma unroll for (int j = 0; j < VEC_SIZE; j++) { float x = static_cast(src1.val[j]); - // Multiply in weight's native dtype to match rms_norm_kernel. - scalar_t out_norm = static_cast(x * s_variance) * src2.val[j]; + float w = static_cast(src2.val[j]); + // Round normalized result through scalar_t to match the precision of the + // unfused composite (rms_norm writes scalar_t, then + // static_scaled_fp8_quant re-loads it as float before FP8 conversion). + // Without this round, the fused path is strictly more accurate and + // disagrees with the composite at exact E4M3 quantization tie boundaries. + scalar_t out_norm = static_cast(x * s_variance * w); out[blockIdx.x * hidden_size + idx * VEC_SIZE + j] = scaled_fp8_conversion(static_cast(out_norm), scale_inv); @@ -137,8 +142,12 @@ fused_add_rms_norm_static_fp8_quant_kernel( #pragma unroll for (int i = 0; i < width; ++i) { float x = Converter::convert(res.data[i]); - // Multiply in weight's native dtype to match fused_add_rms_norm_kernel. - HipT out_norm_h = Converter::convert(x * s_variance) * w.data[i]; + float wf = Converter::convert(w.data[i]); + // See note in rms_norm_static_fp8_quant_kernel: round through scalar_t + // to match the unfused composite path at FP8 boundaries. We use the + // backend's hip_type for the intermediate since c10::Half/BFloat16 has + // ambiguous conversions on CUDA and no implicit conversion on ROCm. + HipT out_norm_h = Converter::convert(x * s_variance * wf); out[id * width + i] = scaled_fp8_conversion( Converter::convert(out_norm_h), scale_inv); } @@ -183,8 +192,10 @@ fused_add_rms_norm_static_fp8_quant_kernel( for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) { float x = (float)residual[blockIdx.x * hidden_size + idx]; - // Multiply in weight's native dtype to match fused_add_rms_norm_kernel. - scalar_t out_norm = static_cast(x * s_variance) * weight[idx]; + float w = (float)weight[idx]; + // See note in rms_norm_static_fp8_quant_kernel: round through scalar_t + // to match the unfused composite path at FP8 boundaries. + scalar_t out_norm = static_cast(x * s_variance * w); out[blockIdx.x * hidden_size + idx] = scaled_fp8_conversion( static_cast(out_norm), scale_inv); } From 16908e132e10f75af93049e865130f8987573f5d Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Thu, 18 Jun 2026 12:42:09 -0700 Subject: [PATCH 16/75] [MRV2] Make FP32 Gumbel sampling more accurate (#45996) Signed-off-by: Woosuk Kwon --- tests/v1/worker/test_gpu_gumbel_sample.py | 227 ++++++++++++++++++++++ vllm/v1/worker/gpu/sample/gumbel.py | 21 +- 2 files changed, 240 insertions(+), 8 deletions(-) create mode 100644 tests/v1/worker/test_gpu_gumbel_sample.py diff --git a/tests/v1/worker/test_gpu_gumbel_sample.py b/tests/v1/worker/test_gpu_gumbel_sample.py new file mode 100644 index 00000000000..9db175113ce --- /dev/null +++ b/tests/v1/worker/test_gpu_gumbel_sample.py @@ -0,0 +1,227 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the Model Runner V2 Gumbel-max sampling kernel. + +Accuracy: define a target categorical distribution as a non-negative int64 +count tensor summing to N, turn it into logits (= log(count)), sample many +times with `gumbel_sample`, and check the empirical distribution matches. + +The count tensor is deliberately heavy-tailed (one dominant token, the rest +~18 logits below). That tail is the sensitive part: the fp32 Gumbel noise must +reach ~18 to ever sample it. A flat distribution would keep every token within +a few logits of the top and would not exercise the noise tail at all. +""" + +import math + +import pytest +import torch + +pytest.importorskip("triton") +if not torch.cuda.is_available(): + pytest.skip("CUDA required for Gumbel sampler tests", allow_module_level=True) + +from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample + +DEVICE = "cuda" +VOCAB_SIZE = 200_000 +NUM_SAMPLES = 500_000 +# Dominant token is exp(HEAD_LOG_GAP)x larger than the unit-count tail, so the +# tail sits ~HEAD_LOG_GAP logits below the top. +HEAD_LOG_GAP = 18.0 +# 10-sigma band: a correct sampler effectively never trips it. +Z_TOLERANCE = 10.0 + + +def _make_heavy_tailed_counts(seed: int = 1234) -> torch.Tensor: + """Non-negative int64 counts of shape [VOCAB_SIZE]; target prob = counts/N.""" + gen = torch.Generator(device=DEVICE).manual_seed(seed) + counts = torch.randint( + 1, 4, (VOCAB_SIZE,), generator=gen, dtype=torch.int64, device=DEVICE + ) + counts[0] = round(math.exp(HEAD_LOG_GAP)) # dominant token + return counts + + +def _counts_to_logits(counts: torch.Tensor) -> torch.Tensor: + # softmax(log(count)) == count / sum(count); count 0 -> logit -inf -> prob 0. + return counts.double().log().to(torch.float32) + + +def _sample( + logits_1d: torch.Tensor, + num_samples: int, + *, + use_fp64: bool = False, + temperature: float = 1.0, +) -> torch.Tensor: + """Sample `num_samples` tokens from one logit vector. + + Fixed seed with a distinct `pos` per sample gives independent draws; the + logits are broadcast with a 0-stride view to avoid materializing + [num_samples, vocab_size]. + """ + vocab_size = logits_1d.shape[0] + logits = logits_1d.unsqueeze(0).expand(num_samples, vocab_size) + idx_mapping = torch.zeros(num_samples, dtype=torch.int32, device=DEVICE) + temp = torch.tensor([temperature], dtype=torch.float32, device=DEVICE) + seed = torch.tensor([0xABCD], dtype=torch.int64, device=DEVICE) + pos = torch.arange(num_samples, dtype=torch.int64, device=DEVICE) + return gumbel_sample( + logits, + idx_mapping, + temp, + seed, + pos, + apply_temperature=True, + use_fp64=use_fp64, + ) + + +def _z_score(observed: int, expected: float, num_trials: int) -> float: + p = expected / num_trials + return (observed - expected) / math.sqrt(num_trials * p * (1 - p)) + + +def _sample_histogram( + logits_1d: torch.Tensor, num_samples: int, *, chunk: int = 1_000_000 +) -> torch.Tensor: + """Histogram of `num_samples` draws, accumulated in chunks. + + Chunking keeps the kernel's per-sample scratch ([chunk, num_blocks]) bounded + so a large sample count does not blow up memory. + """ + vocab_size = logits_1d.shape[0] + hist = torch.zeros(vocab_size, dtype=torch.float64, device=DEVICE) + for start in range(0, num_samples, chunk): + size = min(chunk, num_samples - start) + logits = logits_1d.unsqueeze(0).expand(size, vocab_size) + idx_mapping = torch.zeros(size, dtype=torch.int32, device=DEVICE) + temp = torch.tensor([1.0], dtype=torch.float32, device=DEVICE) + seed = torch.tensor([0xABCD], dtype=torch.int64, device=DEVICE) + pos = torch.arange(start, start + size, dtype=torch.int64, device=DEVICE) + out = gumbel_sample( + logits, idx_mapping, temp, seed, pos, apply_temperature=True + ) + hist += torch.bincount(out, minlength=vocab_size).double() + return hist + + +# ----------------------------- Accuracy ------------------------------------ + + +@pytest.mark.parametrize("use_fp64", [False, True]) +def test_sampling_matches_target_distribution(use_fp64: bool): + counts = _make_heavy_tailed_counts() + total = counts.sum().item() + logits = _counts_to_logits(counts) + + sampled = _sample(logits, NUM_SAMPLES, use_fp64=use_fp64) + assert sampled.min() >= 0 and sampled.max() < VOCAB_SIZE + + # The dominant token (index 0) and the aggregate tail are the two + # statistically resolvable bins (individual tail tokens are far below the + # ~5/N detectability floor). The tail mass is small but well above noise, + # and it lives beyond the fp32 Gumbel cap -- the regime sensitive to noise + # precision -- so matching it is the meaningful check. + tail_prob = (total - counts[0].item()) / total + tail_count = (sampled != 0).sum().item() + z = _z_score(tail_count, NUM_SAMPLES * tail_prob, NUM_SAMPLES) + assert abs(z) < Z_TOLERANCE, ( + f"sampled tail mass {tail_count / NUM_SAMPLES:.3e} != target " + f"{tail_prob:.3e} (z={z:.2f})" + ) + + +def test_full_vocab_distribution_fidelity(): + """The sampled distribution matches the target across the WHOLE vocab. + + A near-flat count tensor makes every one of the 200K bins individually + measurable. With ~20 samples/bin, a goodness-of-fit over all bins checks + that no part of the vocab is over- or under-represented (the heavy-tailed + test above only resolves head vs aggregate tail). Empirically the fp32 + sampler is as faithful here as torch.multinomial; the residual error is the + multinomial sampling-noise floor, not the kernel. + """ + gen = torch.Generator(device=DEVICE).manual_seed(2024) + counts = torch.randint( + 500, 1500, (VOCAB_SIZE,), generator=gen, dtype=torch.int64, device=DEVICE + ) + total = counts.sum().item() + logits = _counts_to_logits(counts) + + num_samples = 4_000_000 + hist = _sample_histogram(logits, num_samples) + + # Diversity: essentially every token must be reachable (no starved region). + coverage = (hist > 0).sum().item() / VOCAB_SIZE + assert coverage > 0.99, f"only {coverage:.4f} of the vocab was ever sampled" + + # Goodness-of-fit across all bins (each has expected count >= ~10). + expected = (counts.double() / total) * num_samples + chi2 = (((hist - expected) ** 2) / expected).sum().item() + df = VOCAB_SIZE - 1 + assert chi2 < df + 10 * math.sqrt(2 * df), f"chi2={chi2:.0f}, df={df}" + + +# ----------------------------- Edge cases ---------------------------------- + + +def test_greedy_temperature_zero_returns_argmax(): + """temperature == 0 skips Gumbel noise and returns the exact argmax.""" + torch.manual_seed(0) + num_reqs = 128 + logits = torch.randn(num_reqs, VOCAB_SIZE, device=DEVICE, dtype=torch.float32) + idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=DEVICE) + temp = torch.zeros(num_reqs, dtype=torch.float32, device=DEVICE) + seed = torch.arange(num_reqs, dtype=torch.int64, device=DEVICE) + pos = torch.arange(num_reqs, dtype=torch.int64, device=DEVICE) + + sampled = gumbel_sample( + logits, idx_mapping, temp, seed, pos, apply_temperature=True + ) + assert torch.equal(sampled, logits.argmax(dim=-1)) + + +def test_zero_count_tokens_are_never_sampled(): + """Count 0 -> -inf logit -> probability 0; must never be selected.""" + counts = _make_heavy_tailed_counts(seed=7) + zeroed = torch.arange(1, VOCAB_SIZE, 2, device=DEVICE) # odd indices (not head) + counts[zeroed] = 0 + logits = _counts_to_logits(counts) + + sampled = _sample(logits, NUM_SAMPLES) + assert sampled.min() >= 0 and sampled.max() < VOCAB_SIZE + assert not torch.isin(sampled, zeroed).any(), "sampled a zero-probability token" + + +def test_single_nonzero_token_is_always_sampled(): + """A lone finite logit must win every draw, regardless of its index.""" + counts = torch.zeros(VOCAB_SIZE, dtype=torch.int64, device=DEVICE) + counts[123_456] = 1000 + logits = _counts_to_logits(counts) + + sampled = _sample(logits, 10_000) + assert (sampled == 123_456).all() + + +@pytest.mark.parametrize("vocab_size", [1, 999, 1024, 4097]) +def test_vocab_size_not_multiple_of_block(vocab_size: int): + """Per-block tail masking for non-block-aligned vocab; all bins measurable.""" + gen = torch.Generator(device=DEVICE).manual_seed(vocab_size) + counts = torch.randint( + 20, 200, (vocab_size,), generator=gen, dtype=torch.int64, device=DEVICE + ) + total = counts.sum().item() + logits = _counts_to_logits(counts) + num_samples = max(40 * vocab_size, 50_000) + + sampled = _sample(logits, num_samples) + assert sampled.min() >= 0 and sampled.max() < vocab_size + + observed = torch.bincount(sampled, minlength=vocab_size).double() + expected = (counts.double() / total) * num_samples + chi2 = (((observed - expected) ** 2) / expected).sum().item() + df = vocab_size - 1 + if df >= 1: + assert chi2 < df + 10 * math.sqrt(2 * df), f"chi2={chi2:.1f}, df={df}" diff --git a/vllm/v1/worker/gpu/sample/gumbel.py b/vllm/v1/worker/gpu/sample/gumbel.py index 44d12738cca..fab53fef7ee 100644 --- a/vllm/v1/worker/gpu/sample/gumbel.py +++ b/vllm/v1/worker/gpu/sample/gumbel.py @@ -2,18 +2,16 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch -from vllm.triton_utils import HAS_TRITON, tl, triton +from vllm.triton_utils import HAS_TRITON, tl, tldevice, triton -# Smallest positive normal fp32 value. Used to clamp the uniform draw so that -# `log(u)` cannot produce -inf (and thus `-log(-log(u))` stays finite). +# Smallest positive value produced by Triton's fp32 `tl.rand`. Used to clamp +# zero draws before the flipped Gumbel transform below. # # Triton requires globals accessed from `@triton.jit` functions to be wrapped # in `tl.constexpr(...)`. We can only do that when Triton is actually # available — on the CPU worker path `tl` is a placeholder whose `constexpr` # attribute is `None`, and `tl.constexpr(...)` would crash at import time. -_FP32_TINY = ( - tl.constexpr(float.fromhex("0x1p-126")) if HAS_TRITON else float.fromhex("0x1p-126") -) +_TL_RAND_MIN = tl.constexpr(4.6566127342e-10) if HAS_TRITON else 4.6566127342e-10 @triton.jit @@ -131,10 +129,17 @@ def gumbel_block_argmax( if USE_FP64: u = tl_rand64(gumbel_seed, block, includes_zero=False) + gumbel_noise = -tl.log(-tl.log(u)) else: u = tl.rand(gumbel_seed, block) - u = tl.maximum(u, _FP32_TINY) - gumbel_noise = -tl.log(-tl.log(u)) + u = tl.maximum(u, _TL_RAND_MIN) + # Draw the large-noise tail (which decides the argmax winner) from u -> 0, + # where fp32 has fine resolution, instead of u -> 1, where fp32 spacing is + # ~2**-24. The naive `-log(-log(u))` puts the winning tail at u -> 1, + # hard-capping the noise at ~16.6 and coarsely quantizing it; using + # `log1p(-u)` == `log(1 - u)` keeps the tail in the well-resolved region. + # Note `1 - u` would lose precision for small u, so `log1p` is required. + gumbel_noise = -tl.log(-tldevice.log1p(-u)) # Apply gumbel noise. logits = tl.where(mask, logits + gumbel_noise, float("-inf")) From 4ce2d0145312809ef6122ccb7be8ae7cafa462a9 Mon Sep 17 00:00:00 2001 From: MrFan <642664360@qq.com> Date: Fri, 19 Jun 2026 04:19:11 +0800 Subject: [PATCH 17/75] fix(anthropic): auto-detect template support for mid-conversation system messages (#46025) Signed-off-by: felix0080 Signed-off-by: Ben Browning Co-authored-by: felix0080 Co-authored-by: Ben Browning --- .../test_anthropic_messages_conversion.py | 47 +++++++++++ vllm/entrypoints/anthropic/serving.py | 79 +++++++++++++++++-- 2 files changed, 120 insertions(+), 6 deletions(-) diff --git a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py index 2fb0f21c877..4663a6565d6 100644 --- a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py +++ b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py @@ -1096,3 +1096,50 @@ class TestMessageStartIncludesTypeAndRole: message = events[0][1]["message"] assert message["type"] == "message" assert message["role"] == "assistant" + + +# ====================================================================== +# Auto-detection of system-first template requirement +# ====================================================================== + + +Q35_TEMPLATE = ( + "{%- for message in messages %}" + "{%- if message.role == 'system' %}" + "{%- if not loop.first %}" + "{{- raise_exception('System message must be at the beginning.') }}" + "{%- endif %}" + "{%- endif %}" + "{%- endfor %}" +) + + +class TestDetectMergeInlineSystem: + """Verify _detect_merge_inline_system auto-detection. + + Tests three scenarios: + 1. Template with system-first guard (e.g. Qwen) → merge needed + 2. Template without restrictions → no merge, cache-friendly + 3. No template provided → safe default: merge + """ + + def test_qwen_template_requires_merge(self): + """Template with loop.first guard rejects mid-conversation system.""" + assert ( + AnthropicServingMessages._detect_merge_inline_system(Q35_TEMPLATE) is True + ) + + def test_no_restriction_no_merge(self): + """Template without restriction accepts mid-conversation system.""" + assert ( + AnthropicServingMessages._detect_merge_inline_system( + "{%- for message in messages %}" + "{{- message.role }}: {{ message.content }}\n" + "{%- endfor %}" + ) + is False + ) + + def test_no_template_defaults_merge(self): + """No chat_template → conservative default: merge.""" + assert AnthropicServingMessages._detect_merge_inline_system(None) is True diff --git a/vllm/entrypoints/anthropic/serving.py b/vllm/entrypoints/anthropic/serving.py index 5a7e8ae95ea..9d5852428df 100644 --- a/vllm/entrypoints/anthropic/serving.py +++ b/vllm/entrypoints/anthropic/serving.py @@ -12,6 +12,7 @@ import uuid from collections.abc import AsyncGenerator from typing import TYPE_CHECKING, Any +import jinja2 from fastapi import Request from vllm.engine.protocol import EngineClient @@ -99,6 +100,36 @@ class AnthropicServingMessages(OpenAIServingChat): "length": "max_tokens", "tool_calls": "tool_use", } + self._merge_inline_system = self._detect_merge_inline_system(chat_template) + + @staticmethod + def _detect_merge_inline_system(chat_template: str | None) -> bool: + """Auto-detect whether the chat template requires system-first ordering. + + Renders a [system, user, system, user] conversation against the + template; if it raises (e.g. Qwen's ``loop.first`` guard), the + model needs inline system messages merged into the leading block. + """ + if not chat_template: + return True + try: + env = jinja2.sandbox.ImmutableSandboxedEnvironment( + trim_blocks=True, + lstrip_blocks=True, + extensions=[jinja2.ext.loopcontrols], + ) + env.from_string(chat_template).render( + messages=[ + {"role": "system", "content": "t"}, + {"role": "user", "content": "t"}, + {"role": "system", "content": "t"}, + {"role": "user", "content": "t"}, + ], + add_generation_prompt=False, + ) + return False + except jinja2.TemplateError: + return True @staticmethod def _convert_image_source_to_url(source: dict[str, Any]) -> str: @@ -123,13 +154,24 @@ class AnthropicServingMessages(OpenAIServingChat): @classmethod def _convert_anthropic_to_openai_request( - cls, anthropic_request: AnthropicMessagesRequest | AnthropicCountTokensRequest + cls, + anthropic_request: AnthropicMessagesRequest | AnthropicCountTokensRequest, + *, + merge_inline_system: bool = False, ) -> ChatCompletionRequest: """Convert Anthropic message format to OpenAI format""" openai_messages: list[dict[str, Any]] = [] - cls._convert_system_message(anthropic_request, openai_messages) - cls._convert_messages(anthropic_request.messages, openai_messages) + cls._convert_system_message( + anthropic_request, + openai_messages, + merge_inline_system=merge_inline_system, + ) + cls._convert_messages( + anthropic_request.messages, + openai_messages, + merge_inline_system=merge_inline_system, + ) req = cls._build_base_request(anthropic_request, openai_messages) cls._handle_streaming_options(req, anthropic_request) cls._handle_output_config(req, anthropic_request) @@ -142,6 +184,8 @@ class AnthropicServingMessages(OpenAIServingChat): cls, anthropic_request: AnthropicMessagesRequest | AnthropicCountTokensRequest, openai_messages: list[dict[str, Any]], + *, + merge_inline_system: bool = False, ) -> None: """Convert Anthropic system message to OpenAI format""" system_parts: list[str] = [] @@ -159,6 +203,17 @@ class AnthropicServingMessages(OpenAIServingChat): continue system_parts.append(block.text) + # When the template requires system-first ordering, extract inline + # system messages from the messages array and merge them into the + # top-level block so the template doesn't reject them. + if merge_inline_system: + for msg in anthropic_request.messages: + if msg.role != "system": + continue + text = cls._extract_system_text(msg) + if text: + system_parts.append(text) + if system_parts: openai_messages.append({"role": "system", "content": "".join(system_parts)}) @@ -180,7 +235,11 @@ class AnthropicServingMessages(OpenAIServingChat): @classmethod def _convert_messages( - cls, messages: list, openai_messages: list[dict[str, Any]] + cls, + messages: list, + openai_messages: list[dict[str, Any]], + *, + merge_inline_system: bool = False, ) -> None: """Convert Anthropic messages to OpenAI format""" for msg in messages: @@ -190,6 +249,8 @@ class AnthropicServingMessages(OpenAIServingChat): # doesn't strip billing headers and may produce messages with # no "content" key. if msg.role == "system": + if merge_inline_system: + continue # already merged into top-level by _convert_system_message text = cls._extract_system_text(msg) if text: openai_messages.append({"role": "system", "content": text}) @@ -497,7 +558,10 @@ class AnthropicServingMessages(OpenAIServingChat): """ if logger.isEnabledFor(logging.DEBUG): logger.debug("Received messages request %s", request.model_dump_json()) - chat_req = self._convert_anthropic_to_openai_request(request) + chat_req = self._convert_anthropic_to_openai_request( + request, + merge_inline_system=self._merge_inline_system, + ) if logger.isEnabledFor(logging.DEBUG): logger.debug("Convert to OpenAI request %s", chat_req.model_dump_json()) generator = await self.create_chat_completion(chat_req, raw_request) @@ -905,7 +969,10 @@ class AnthropicServingMessages(OpenAIServingChat): raw_request: Request | None = None, ) -> AnthropicCountTokensResponse | ErrorResponse: """Implements Anthropic's messages.count_tokens endpoint.""" - chat_req = self._convert_anthropic_to_openai_request(request) + chat_req = self._convert_anthropic_to_openai_request( + request, + merge_inline_system=self._merge_inline_system, + ) result = await self.render_chat_request(chat_req) if isinstance(result, ErrorResponse): return result From 35e4dd4a69b6b95feb74866341daa46c3836aed0 Mon Sep 17 00:00:00 2001 From: Yifan Qiao Date: Thu, 18 Jun 2026 14:44:02 -0700 Subject: [PATCH 18/75] [KV Connector][Mooncake] Async lookup to reduce scheduler overhead (#45659) Signed-off-by: Yifan Qiao Signed-off-by: Nick Hill Co-authored-by: Nick Hill --- .../mooncake_store_connector_usage.md | 1 + .../unit/test_mooncake_store_connector.py | 127 +++++++++++++++++- .../unit/test_mooncake_store_scheduler.py | 9 +- .../v1/mooncake/store/connector.py | 2 +- .../v1/mooncake/store/scheduler.py | 25 +++- .../kv_connector/v1/mooncake/store/worker.py | 44 +++++- 6 files changed, 197 insertions(+), 11 deletions(-) diff --git a/docs/features/mooncake_store_connector_usage.md b/docs/features/mooncake_store_connector_usage.md index bab69410978..cb857856b78 100644 --- a/docs/features/mooncake_store_connector_usage.md +++ b/docs/features/mooncake_store_connector_usage.md @@ -203,6 +203,7 @@ the vLLM JSON config. ### kv_connector_extra_config - `load_async` (bool): Enable asynchronous loading for better compute-I/O overlap. Default: `true`. +- `lookup_async` (bool): Run the external prefix-cache lookup on a background thread so it never blocks the scheduler step. The request is held until the in-flight lookup completes, then resumed on a later step. Default: `false`. - `enable_cross_layers_blocks` (bool): Enable cross-layer block packing for reduced store operations. Default: `false`. - `lookup_rpc_port` (int): Custom port for the ZMQ lookup RPC socket. Default: `0`. - `cache_prefix` (str): Namespace prepended to every store key. Lets separate deployments share one Mooncake master without polluting each other — instances configured with different prefixes never see each other's cached blocks, even for identical prompts. All instances that should share a prefix cache must use the same value. Default: `""` (no prefix; keys are byte-identical to the unprefixed format). diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_connector.py b/tests/v1/kv_connector/unit/test_mooncake_store_connector.py index d3992b02b68..951b447fd6b 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_connector.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_connector.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import threading +import time from unittest.mock import MagicMock, patch from vllm.config import set_current_vllm_config @@ -406,7 +408,9 @@ def test_lookup_key_client_lookup_prepends_typed_tag(): fake_socket = mock_make_socket.return_value fake_socket.recv.return_value = (5).to_bytes(4, "big") - assert client.lookup(token_len=128, block_hashes=[]) == 5 + # Blocking lookup (non_block defaults to False) runs on the executor and + # returns the resolved hit length. + assert client.lookup("req0", token_len=128, block_hashes=[]) == 5 sent_frames = fake_socket.send_multipart.call_args[0][0] assert sent_frames[0] == protocol.LOOKUP_MSG @@ -435,6 +439,127 @@ def test_lookup_key_client_reset_uses_typed_protocol(): assert client.reset() is False +def _poll_lookup(client, req_id, token_len=128, block_hashes=(), timeout=5.0): + """Drive non-blocking lookup until the executor completes it.""" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + result = client.lookup(req_id, token_len, list(block_hashes), non_block=True) + if result is not None: + return result + time.sleep(0.005) + return None + + +def _gated_recv(gate: threading.Event, value: int): + """Mock recv side-effect that blocks until ``gate`` is set, so the + executor's lookup can be held pending deterministically.""" + + def recv(): + gate.wait() + return value.to_bytes(4, "big") + + return recv + + +def test_lookup_key_client_non_block_lookup_async(): + """Non-blocking lookup defers to the executor: None first, hit once the + Future resolves.""" + vllm_config = _make_vllm_config() + + with patch( + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store." + "worker.make_zmq_socket" + ) as mock_make_socket: + client = worker.LookupKeyClient(vllm_config) + + fake_socket = mock_make_socket.return_value + # Hold the executor's lookup pending until we release the gate. + gate = threading.Event() + fake_socket.recv.side_effect = _gated_recv(gate, 7) + + # First query submits the lookup and returns None while it is in flight. + assert client.lookup("req1", 128, [], non_block=True) is None + # Release the executor; a later poll returns the hit length. + gate.set() + assert _poll_lookup(client, "req1") == 7 + # Future is consumed (popped) on read. + assert "req1" not in client.futures + + +def test_lookup_key_client_discard_clears_state(): + """discard() drops a completed lookup Future so it is not served stale.""" + vllm_config = _make_vllm_config() + + with patch( + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store." + "worker.make_zmq_socket" + ) as mock_make_socket: + client = worker.LookupKeyClient(vllm_config) + + fake_socket = mock_make_socket.return_value + gate = threading.Event() + fake_socket.recv.side_effect = _gated_recv(gate, 9) + + # Submit while gated so the call returns None and the Future stays in + # `futures` (unconsumed) once it resolves. + assert client.lookup("req2", 128, [], non_block=True) is None + gate.set() + deadline = time.monotonic() + 5.0 + while time.monotonic() < deadline: + if client.futures["req2"].done(): + break + time.sleep(0.005) + # discard() drops the completed result before any lookup consumes it. + client.discard("req2") + assert "req2" not in client.futures + # A fresh query re-submits rather than returning a stale value: hold the + # gate so the resubmitted lookup stays in flight. + gate.clear() + assert client.lookup("req2", 128, [], non_block=True) is None + gate.set() # release the executor so the worker thread can drain + + +def test_get_num_new_matched_tokens_async_defers_then_reports(): + """Async lookup returns (None, False) until ready, then the hit count.""" + vllm_config = create_vllm_config( + kv_connector="MooncakeStoreConnector", + kv_role="kv_both", + kv_connector_extra_config={"lookup_async": True}, + ) + kv_cache_config = _make_kv_cache_config() + + with ( + set_current_vllm_config(vllm_config), + patch( + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store." + "scheduler.LookupKeyClient" + ) as mock_client_cls, + ): + sched = scheduler.MooncakeStoreScheduler(vllm_config, kv_cache_config) + + assert sched.lookup_async is True + mock_client = mock_client_cls.return_value + + block_size = sched._block_size + request = MagicMock() + request.request_id = "r1" + request.num_tokens = 4 * block_size + request.block_hashes = [] + + # Lookup not ready -> defer. + mock_client.lookup.return_value = None + assert sched.get_num_new_matched_tokens(request, 0) == (None, False) + assert "r1" not in sched.load_specs + + # Lookup ready with a hit -> report need_to_allocate + async-load flag. + hit = 3 * block_size + mock_client.lookup.return_value = hit + need, load_async = sched.get_num_new_matched_tokens(request, 0) + assert need == hit + assert load_async == sched.load_async + assert sched.load_specs["r1"].kvpool_cached_tokens == hit + + def test_protocol_tags_are_distinct_and_non_empty(): """Protocol tags must be unique and non-empty to avoid collision.""" tags = {protocol.LOOKUP_MSG, protocol.RESET_MSG} diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py b/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py index ac36005c63e..8ef1277bb39 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py @@ -16,6 +16,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.scheduler impor def _make_bare_scheduler() -> MooncakeStoreScheduler: scheduler = object.__new__(MooncakeStoreScheduler) scheduler.kv_role = "kv_both" + scheduler.lookup_async = False scheduler._block_size = 16 scheduler.load_specs = {} scheduler._preempted_req_ids = set() @@ -405,7 +406,13 @@ class _StubLookupClient: def __init__(self, hit_tokens: int) -> None: self._hit_tokens = hit_tokens - def lookup(self, token_len: int, block_hashes: list[bytes]) -> int: + def lookup( + self, + req_id: str, + token_len: int, + block_hashes: list[bytes], + non_block: bool = False, + ) -> int: return self._hit_tokens diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py index d53cd13c2e4..bf6038a897a 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py @@ -176,7 +176,7 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA): self, request: Request, num_computed_tokens: int, - ) -> tuple[int, bool]: + ) -> tuple[int | None, bool]: assert self.connector_scheduler is not None return self.connector_scheduler.get_num_new_matched_tokens( request, num_computed_tokens diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py index 4c4d55df3e1..620fa2f5ba1 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py @@ -54,9 +54,9 @@ class MooncakeStoreScheduler: ): assert vllm_config.kv_transfer_config is not None self.kv_role = vllm_config.kv_transfer_config.kv_role - self.load_async = vllm_config.kv_transfer_config.kv_connector_extra_config.get( - "load_async", True - ) + kvc_extra_config = vllm_config.kv_transfer_config.kv_connector_extra_config + self.load_async = kvc_extra_config.get("load_async", True) + self.lookup_async = kvc_extra_config.get("lookup_async", False) self.client = LookupKeyClient(vllm_config) # Align with the engine's own scheduler_block_size and hash_block_size. @@ -75,14 +75,26 @@ class MooncakeStoreScheduler: self, request: Request, num_computed_tokens: int, - ) -> tuple[int, bool]: - """Check for external KV cache hit.""" + ) -> tuple[int | None, bool]: + """Check for external KV cache hit. + + Returns ``(None, False)`` when an async lookup is still in flight, + signaling the scheduler to retry this request on a later step. + """ # Look up against the full prefill range, not just the prompt. token_len = request.num_tokens // self._block_size * self._block_size if token_len < self._block_size: return 0, False - num_external_hit_tokens = self.client.lookup(token_len, request.block_hashes) + num_external_hit_tokens = self.client.lookup( + request.request_id, + token_len, + request.block_hashes, + non_block=self.lookup_async, + ) + if num_external_hit_tokens is None: + # Lookup not ready yet; scheduler will retry on a later step. + return None, False if num_external_hit_tokens == request.num_tokens: # Leave a sub-block tail uncomputed for sampling, on a block @@ -158,6 +170,7 @@ class MooncakeStoreScheduler: force_skip_save = self.kv_role == "kv_consumer" for finished_req_id in scheduler_output.finished_req_ids: + self.client.discard(finished_req_id) self.load_specs.pop(finished_req_id, None) self._request_trackers.pop(finished_req_id, None) self._unfinished_requests.pop(finished_req_id, None) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index f5a55b54c75..0d9633f7596 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -19,6 +19,7 @@ import threading import time from collections import defaultdict from collections.abc import Callable +from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import dataclass from typing import Any, Literal, TypeVar @@ -1560,7 +1561,13 @@ class LookupKeyClient: bind=False, ) - def lookup(self, token_len: int, block_hashes: list[BlockHash]) -> int: + # Async lookup support + self.executor = ThreadPoolExecutor( + max_workers=1, thread_name_prefix="MooncakeLookupClient" + ) + self.futures: dict[str, Future[int]] = {} + + def _lookup(self, token_len: int, block_hashes: list[BlockHash]) -> int: hash_strs = [h.hex() for h in block_hashes] hash_frames = self.encoder.encode(hash_strs) token_len_bytes = token_len.to_bytes(4, byteorder="big") @@ -1570,7 +1577,36 @@ class LookupKeyClient: result = int.from_bytes(resp, "big") return result - def reset(self) -> bool: + def lookup( + self, + req_id: str, + token_len: int, + block_hashes: list[BlockHash], + non_block: bool = False, + ) -> int | None: + """If non_block is True, will return None until the result is ready, + so the caller retries on a later step.""" + future = self.futures.get(req_id) + if future is None: + future = self.executor.submit(self._lookup, token_len, list(block_hashes)) + self.futures[req_id] = future + if non_block and not future.done(): + return None + try: + return future.result() + except Exception as e: + logger.error("Async Mooncake lookup failed for %s: %s", req_id, e) + return 0 + finally: + del self.futures[req_id] + + def discard(self, req_id: str) -> None: + """Drop any cached/in-flight lookup for ``req_id`` (e.g. on abort).""" + future = self.futures.pop(req_id, None) + if future is not None: + future.cancel() + + def _reset(self) -> bool: """Trigger ``store.remove_all(force=True)`` on worker rank 0. Ordering assumption: caller MUST ensure no in-flight Mooncake @@ -1582,7 +1618,11 @@ class LookupKeyClient: resp = self.socket.recv() return bytes(resp) == RESP_OK + def reset(self) -> bool: + return self.executor.submit(self._reset).result() + def close(self): + self.executor.shutdown(wait=False, cancel_futures=True) self.socket.close(linger=0) From 41dcf49ca52ab25178ca8869298275b1787f328a Mon Sep 17 00:00:00 2001 From: Yifan Qiao Date: Thu, 18 Jun 2026 15:13:44 -0700 Subject: [PATCH 19/75] [Bugfix][KV Connector] Disable Mooncake TP put-striding when DCP > 1 (#45371) Signed-off-by: Yifan Qiao Co-authored-by: Jingyi Yang Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../unit/test_mooncake_store_worker.py | 87 +++++++++++++++++-- .../kv_connector/v1/mooncake/store/worker.py | 8 +- 2 files changed, 87 insertions(+), 8 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index aa5d7d1ff3b..5213805115e 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -175,14 +175,17 @@ class _FakeModelConfig: def _make_vllm_config( - *, extra_config: dict[str, object] | None = None + *, + extra_config: dict[str, object] | None = None, + rank: int = 0, + decode_context_parallel_size: int = 1, ) -> SimpleNamespace: return SimpleNamespace( model_config=_FakeModelConfig(), parallel_config=SimpleNamespace( pipeline_parallel_size=1, - rank=0, - decode_context_parallel_size=1, + rank=rank, + decode_context_parallel_size=decode_context_parallel_size, prefill_context_parallel_size=1, ), kv_transfer_config=_FakeKVTransferConfig(extra_config=extra_config), @@ -231,13 +234,23 @@ def _install_fake_mooncake(monkeypatch, store_instance: MagicMock): return FakeReplicateConfig -def _patch_worker_runtime(monkeypatch, *, local_ip: str = "10.0.0.7") -> None: +def _patch_worker_runtime( + monkeypatch, + *, + local_ip: str = "10.0.0.7", + tp_rank: int = 0, + tp_size: int = 1, + dcp_size: int = 1, +) -> None: single_rank_group = SimpleNamespace(world_size=1, rank_in_group=0) + # DCP groups are contiguous splits of the TP group (see + # parallel_state.py), so dcp_rank == tp_rank % dcp_size. + dcp_group = SimpleNamespace(world_size=dcp_size, rank_in_group=tp_rank % dcp_size) monkeypatch.setattr(worker, "get_mooncake_dp_engine_index", lambda _: 0) - monkeypatch.setattr(worker, "get_tensor_model_parallel_rank", lambda: 0) - monkeypatch.setattr(worker, "get_tensor_model_parallel_world_size", lambda: 1) + monkeypatch.setattr(worker, "get_tensor_model_parallel_rank", lambda: tp_rank) + monkeypatch.setattr(worker, "get_tensor_model_parallel_world_size", lambda: tp_size) monkeypatch.setattr(worker, "get_pcp_group", lambda: single_rank_group) - monkeypatch.setattr(worker, "get_dcp_group", lambda: single_rank_group) + monkeypatch.setattr(worker, "get_dcp_group", lambda: dcp_group) monkeypatch.setattr(worker, "get_ip", lambda: local_ip) @@ -884,6 +897,66 @@ def test_requester_worker_init_builds_replicate_config_for_preferred_segment( assert w.store_replicate_config.preferred_segment == "10.0.0.7:50053" +@pytest.mark.parametrize("dcp_size", [1, 4]) +def test_worker_put_striding_covers_every_rank_get_namespace( + tmp_path, monkeypatch, dcp_size +): + """Every key a rank GETs must have been PUT by some rank. + + When num_kv_head < tp_size, ranks holding the same KV heads stripe + their PUTs across one shared key namespace. That dedup is only valid + when those ranks really share a namespace: with DCP > 1 each rank GETs + every key from its own ``@dcpN`` namespace, so striding must be + disabled. + """ + tp_size = 4 + store = MagicMock() + store.setup.return_value = 0 + _install_fake_mooncake(monkeypatch, store) + monkeypatch.setenv( + "MOONCAKE_CONFIG_PATH", + _write_mooncake_config( + tmp_path, + { + "metadata_server": "http://metadata/endpoint", + "protocol": "tcp", + "device_name": "", + "master_server_address": "10.0.0.7:50051", + }, + ), + ) + + # _FakeModelConfig has num_kv_head=1 < tp_size, which enables striding. + block_hashes = [f"hash-{i}".encode() for i in range(4)] + put_keys: set[str] = set() + get_keys_per_rank: dict[int, set[str]] = {} + for tp_rank in range(tp_size): + _patch_worker_runtime( + monkeypatch, tp_rank=tp_rank, tp_size=tp_size, dcp_size=dcp_size + ) + w = worker.MooncakeStoreWorker( + _make_vllm_config(rank=tp_rank, decode_context_parallel_size=dcp_size), + _make_kv_cache_config(), + ) + db = w.token_dbs[0] + token_len = len(block_hashes) * db.block_size + keys = [ + key.to_string() for _, _, key in db.process_tokens(token_len, block_hashes) + ] + assert len(keys) == len(block_hashes) + # PUT side: mirrors KVCacheStoreSendingThread's striding slice. + put_keys.update(keys[w.tp_rank % w.put_step :: w.put_step]) + # GET side: KVCacheStoreRecvingThread fetches every key. + get_keys_per_rank[tp_rank] = set(keys) + + for tp_rank, rank_keys in get_keys_per_rank.items(): + missing = rank_keys - put_keys + assert not missing, ( + f"tp_rank={tp_rank} would GET {len(missing)}/{len(rank_keys)} keys " + f"that no rank PUT (Mooncake OBJECT_NOT_FOUND): {sorted(missing)}" + ) + + # --------------------------------------------------------------------------- # Helpers for register_kv_caches tests # --------------------------------------------------------------------------- diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index 0d9633f7596..62c2d30c9c4 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -972,7 +972,13 @@ class MooncakeStoreWorker: else: self.num_kv_head = model_config.get_total_num_kv_heads() - if self.num_kv_head < self.tp_size: + if self.num_kv_head < self.tp_size and self.dcp_size <= 1: + # Dedup: TP ranks holding the same KV heads stripe PUTs across + # one shared key namespace. DCP splits the TP group, so with + # DCP>1 those ranks have different `@dcpN` namespaces and + # striping would leave keys unwritten (OBJECT_NOT_FOUND on + # GET). PCP is outer to TP (pcp_rank is constant within a TP + # group), so it needs no guard. self.put_step = self.tp_size // self.num_kv_head self.head_or_tp_rank = self.tp_rank // self.put_step else: From c3c6d723fdd1c315322e5d5a51c479eb2bc017a2 Mon Sep 17 00:00:00 2001 From: Ivy Xu Date: Fri, 19 Jun 2026 06:24:29 +0800 Subject: [PATCH 20/75] [Perf] Remove unused loggers in `reasoning/` (#45988) Signed-off-by: Ivy --- vllm/reasoning/deepseek_v3_reasoning_parser.py | 3 --- vllm/reasoning/ernie45_reasoning_parser.py | 3 --- vllm/reasoning/granite_reasoning_parser.py | 3 --- vllm/reasoning/hunyuan_a13b_reasoning_parser.py | 3 --- vllm/reasoning/identity_reasoning_parser.py | 3 --- vllm/reasoning/minimax_m2_reasoning_parser.py | 3 --- vllm/reasoning/mistral_reasoning_parser.py | 3 --- vllm/reasoning/olmo3_reasoning_parser.py | 3 --- vllm/reasoning/step3_reasoning_parser.py | 3 --- 9 files changed, 27 deletions(-) diff --git a/vllm/reasoning/deepseek_v3_reasoning_parser.py b/vllm/reasoning/deepseek_v3_reasoning_parser.py index bb79afd8ded..dbaf0b1cf89 100644 --- a/vllm/reasoning/deepseek_v3_reasoning_parser.py +++ b/vllm/reasoning/deepseek_v3_reasoning_parser.py @@ -6,7 +6,6 @@ from typing import TYPE_CHECKING from transformers import PreTrainedTokenizerBase -from vllm.logger import init_logger from vllm.reasoning import ReasoningParser from vllm.reasoning.deepseek_r1_reasoning_parser import DeepSeekR1ReasoningParser @@ -17,8 +16,6 @@ if TYPE_CHECKING: from vllm.entrypoints.openai.engine.protocol import DeltaMessage from vllm.entrypoints.openai.responses.protocol import ResponsesRequest -logger = init_logger(__name__) - class DeepSeekV3ReasoningParser(ReasoningParser): """ diff --git a/vllm/reasoning/ernie45_reasoning_parser.py b/vllm/reasoning/ernie45_reasoning_parser.py index 593eba4ecb4..a755c72a1e3 100644 --- a/vllm/reasoning/ernie45_reasoning_parser.py +++ b/vllm/reasoning/ernie45_reasoning_parser.py @@ -7,15 +7,12 @@ from typing import TYPE_CHECKING from transformers import PreTrainedTokenizerBase from vllm.entrypoints.openai.engine.protocol import DeltaMessage -from vllm.logger import init_logger from vllm.reasoning.basic_parsers import BaseThinkingReasoningParser if TYPE_CHECKING: from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.entrypoints.openai.responses.protocol import ResponsesRequest -logger = init_logger(__name__) - class Ernie45ReasoningParser(BaseThinkingReasoningParser): """ diff --git a/vllm/reasoning/granite_reasoning_parser.py b/vllm/reasoning/granite_reasoning_parser.py index 2d8052f614d..c6d63fc3614 100644 --- a/vllm/reasoning/granite_reasoning_parser.py +++ b/vllm/reasoning/granite_reasoning_parser.py @@ -8,15 +8,12 @@ import regex as re from transformers import PreTrainedTokenizerBase from vllm.entrypoints.openai.engine.protocol import DeltaMessage -from vllm.logger import init_logger from vllm.reasoning import ReasoningParser if TYPE_CHECKING: from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.entrypoints.openai.responses.protocol import ResponsesRequest -logger = init_logger(__name__) - class GraniteReasoningParser(ReasoningParser): """ diff --git a/vllm/reasoning/hunyuan_a13b_reasoning_parser.py b/vllm/reasoning/hunyuan_a13b_reasoning_parser.py index f833f8f32f6..257dc0f9540 100644 --- a/vllm/reasoning/hunyuan_a13b_reasoning_parser.py +++ b/vllm/reasoning/hunyuan_a13b_reasoning_parser.py @@ -8,15 +8,12 @@ import regex as re from transformers import PreTrainedTokenizerBase from vllm.entrypoints.openai.engine.protocol import DeltaMessage -from vllm.logger import init_logger from vllm.reasoning import ReasoningParser if TYPE_CHECKING: from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.entrypoints.openai.responses.protocol import ResponsesRequest -logger = init_logger(__name__) - class HunyuanA13BReasoningParser(ReasoningParser): """ diff --git a/vllm/reasoning/identity_reasoning_parser.py b/vllm/reasoning/identity_reasoning_parser.py index c6f117e2f98..ee35360ea6c 100644 --- a/vllm/reasoning/identity_reasoning_parser.py +++ b/vllm/reasoning/identity_reasoning_parser.py @@ -7,15 +7,12 @@ from typing import TYPE_CHECKING from transformers import PreTrainedTokenizerBase from vllm.entrypoints.openai.engine.protocol import DeltaMessage -from vllm.logger import init_logger from vllm.reasoning import ReasoningParser if TYPE_CHECKING: from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.entrypoints.openai.responses.protocol import ResponsesRequest -logger = init_logger(__name__) - class IdentityReasoningParser(ReasoningParser): """ diff --git a/vllm/reasoning/minimax_m2_reasoning_parser.py b/vllm/reasoning/minimax_m2_reasoning_parser.py index 935a3b26aa5..9c3a502e4f8 100644 --- a/vllm/reasoning/minimax_m2_reasoning_parser.py +++ b/vllm/reasoning/minimax_m2_reasoning_parser.py @@ -7,7 +7,6 @@ from typing import TYPE_CHECKING from vllm.entrypoints.openai.engine.protocol import ( DeltaMessage, ) -from vllm.logger import init_logger from vllm.parser.engine.registered_adapters import MinimaxM2ParserReasoningAdapter from vllm.reasoning.abs_reasoning_parsers import ReasoningParser from vllm.tokenizers import TokenizerLike @@ -16,8 +15,6 @@ if TYPE_CHECKING: from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.entrypoints.openai.responses.protocol import ResponsesRequest -logger = init_logger(__name__) - class MiniMaxM2ReasoningParser(MinimaxM2ParserReasoningAdapter): # type: ignore[valid-type, misc] """ diff --git a/vllm/reasoning/mistral_reasoning_parser.py b/vllm/reasoning/mistral_reasoning_parser.py index 74e32cfd163..c224c3c165c 100644 --- a/vllm/reasoning/mistral_reasoning_parser.py +++ b/vllm/reasoning/mistral_reasoning_parser.py @@ -5,7 +5,6 @@ from collections.abc import Iterable, Sequence from functools import cached_property from typing import TYPE_CHECKING -from vllm.logger import init_logger from vllm.reasoning import ReasoningParser from vllm.reasoning.basic_parsers import BaseThinkingReasoningParser from vllm.tokenizers.mistral import MistralTokenizer @@ -14,8 +13,6 @@ if TYPE_CHECKING: from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.entrypoints.openai.responses.protocol import ResponsesRequest -logger = init_logger(__name__) - class MistralReasoningParser(BaseThinkingReasoningParser): """ diff --git a/vllm/reasoning/olmo3_reasoning_parser.py b/vllm/reasoning/olmo3_reasoning_parser.py index 102508b9ac1..dd323501dfb 100644 --- a/vllm/reasoning/olmo3_reasoning_parser.py +++ b/vllm/reasoning/olmo3_reasoning_parser.py @@ -9,7 +9,6 @@ from typing import TYPE_CHECKING import regex as re from vllm.entrypoints.openai.engine.protocol import DeltaMessage -from vllm.logger import init_logger from vllm.reasoning import ReasoningParser if TYPE_CHECKING: @@ -17,8 +16,6 @@ if TYPE_CHECKING: from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.tokenizers import TokenizerLike -logger = init_logger(__name__) - class Olmo3ReasoningState(enum.Enum): REASONING = 1 diff --git a/vllm/reasoning/step3_reasoning_parser.py b/vllm/reasoning/step3_reasoning_parser.py index a50fcf02db4..bc80003edc3 100644 --- a/vllm/reasoning/step3_reasoning_parser.py +++ b/vllm/reasoning/step3_reasoning_parser.py @@ -9,15 +9,12 @@ import regex as re from transformers import PreTrainedTokenizerBase from vllm.entrypoints.openai.engine.protocol import DeltaMessage -from vllm.logger import init_logger from vllm.reasoning import ReasoningParser if TYPE_CHECKING: from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.entrypoints.openai.responses.protocol import ResponsesRequest -logger = init_logger(__name__) - class Step3ReasoningParser(ReasoningParser): """ From 7f616c327d24a259dd81605e513c42ce2b9dc204 Mon Sep 17 00:00:00 2001 From: Ben Browning Date: Thu, 18 Jun 2026 19:17:18 -0400 Subject: [PATCH 21/75] [Bugfix] [Parser] Fix empty tool block silently dropping subsequent content (#46091) Signed-off-by: Ben Browning Co-authored-by: Flora Feng <4florafeng@gmail.com> --- tests/parser/engine/trace_builder.py | 19 +++++++++++++++++-- vllm/parser/engine/parser_engine.py | 2 +- vllm/parser/gemma4.py | 4 ++++ vllm/parser/qwen3.py | 4 ++++ 4 files changed, 26 insertions(+), 3 deletions(-) diff --git a/tests/parser/engine/trace_builder.py b/tests/parser/engine/trace_builder.py index 4817d3b9005..bee3d5d8b28 100644 --- a/tests/parser/engine/trace_builder.py +++ b/tests/parser/engine/trace_builder.py @@ -143,6 +143,12 @@ SCENARIOS: list[Scenario] = [ tool_calls=[_READ_TOOL], after_tool_response=True, ), + Scenario( + id="empty-tool-block", + description="Empty tool block followed by content (edge case recovery)", + content="Content after empty tools.", + tool_calls=[], + ), ] @@ -344,8 +350,11 @@ def _qwen3_segments(scenario: Scenario) -> list[tuple[str, bool]]: segs: list[tuple[str, bool]] = [] if scenario.reasoning is not None: segs.append((scenario.reasoning, False)) - if scenario.content is not None or scenario.tool_calls: + if scenario.content is not None or scenario.tool_calls is not None: segs.append(("", True)) + if scenario.tool_calls is not None and not scenario.tool_calls: + segs.append(("", True)) + segs.append(("", True)) if scenario.content is not None: segs.append((scenario.content, False)) if scenario.tool_calls: @@ -437,8 +446,11 @@ def _minimax_m2_segments(scenario: Scenario) -> list[tuple[str, bool]]: segs: list[tuple[str, bool]] = [] if scenario.reasoning is not None: segs.append((scenario.reasoning, False)) - if scenario.content is not None or scenario.tool_calls: + if scenario.content is not None or scenario.tool_calls is not None: segs.append(("", True)) + if scenario.tool_calls is not None and not scenario.tool_calls: + segs.append(("", True)) + segs.append(("", True)) if scenario.content is not None: segs.append((scenario.content, False)) if scenario.tool_calls: @@ -534,6 +546,9 @@ def _gemma4_segments(scenario: Scenario) -> list[tuple[str, bool]]: segs.append((_GEMMA4_THOUGHT_PREFIX, False)) segs.append((scenario.reasoning, False)) segs.append(("", True)) + if scenario.tool_calls is not None and not scenario.tool_calls: + segs.append(("<|tool_call>", True)) + segs.append(("", True)) if scenario.content is not None: segs.append((scenario.content, False)) if scenario.tool_calls: diff --git a/vllm/parser/engine/parser_engine.py b/vllm/parser/engine/parser_engine.py index 237e2745632..dafb26fc48d 100644 --- a/vllm/parser/engine/parser_engine.py +++ b/vllm/parser/engine/parser_engine.py @@ -672,7 +672,7 @@ class ParserEngine(Parser): if len(tool_call_deltas) > 1: tool_call_deltas = self._coalesce_tool_call_deltas(tool_call_deltas) - if self._deferred_content and not seen_tool_event: + if self._deferred_content and (not seen_tool_event or not tool_call_deltas): content_parts.insert(0, self._deferred_content) self._deferred_content = "" diff --git a/vllm/parser/gemma4.py b/vllm/parser/gemma4.py index 5dd07e44e3e..e9223ee72f7 100644 --- a/vllm/parser/gemma4.py +++ b/vllm/parser/gemma4.py @@ -375,6 +375,10 @@ def gemma4_config() -> ParserEngineConfig: ParserState.TOOL_PREAMBLE, (EventType.REASONING_END, EventType.TOOL_CALL_START), ), + (ParserState.TOOL_PREAMBLE, "TOOL_END"): Transition( + ParserState.CONTENT, + (EventType.TOOL_CALL_END,), + ), (ParserState.TOOL_PREAMBLE, "CALL_PREFIX"): Transition( ParserState.TOOL_NAME, (), diff --git a/vllm/parser/qwen3.py b/vllm/parser/qwen3.py index 45e3c7e4325..583d3481bd8 100644 --- a/vllm/parser/qwen3.py +++ b/vllm/parser/qwen3.py @@ -125,6 +125,10 @@ def qwen3_config(thinking: bool = True) -> ParserEngineConfig: ParserState.TOOL_NAME, (EventType.TOOL_CALL_START,), ), + (ParserState.TOOL_PREAMBLE, "TOOL_END"): Transition( + ParserState.CONTENT, + (EventType.TOOL_CALL_END,), + ), (ParserState.TOOL_PREAMBLE, "FUNC_PREFIX"): Transition( ParserState.TOOL_NAME, (), From 675cd5d228869d152eba17526f1bef0b97f58ed8 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Thu, 18 Jun 2026 20:36:40 -0400 Subject: [PATCH 22/75] [Model Runner V2] Fix MRv2 memory leak test (#46095) Signed-off-by: yewentao256 --- tests/models/multimodal/generation/test_memory_leak.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/models/multimodal/generation/test_memory_leak.py b/tests/models/multimodal/generation/test_memory_leak.py index 743a71f928f..5ee505257c1 100644 --- a/tests/models/multimodal/generation/test_memory_leak.py +++ b/tests/models/multimodal/generation/test_memory_leak.py @@ -25,7 +25,7 @@ TEST_IMAGE_NAMES = [ ] MAX_MODEL_LEN = 8192 REQUESTS_PER_ROUND = 4 -WARMUP_ROUNDS = 1 +WARMUP_ROUNDS = 2 MEASURED_ROUNDS = 16 GPU_GROWTH_THRESHOLD_MIB = 0 CPU_PEAK_GROWTH_THRESHOLD_MIB = 0 From 560fb8b867aaa444d471b35fd846368ebacf12b9 Mon Sep 17 00:00:00 2001 From: Flora Feng <4florafeng@gmail.com> Date: Thu, 18 Jun 2026 21:02:11 -0400 Subject: [PATCH 23/75] [Cohere] Remove dead prepare_structured_tag override in Cohere parser (#46099) Signed-off-by: sfeng33 <4florafeng@gmail.com> --- vllm/reasoning/cohere_command_reasoning_parser.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/vllm/reasoning/cohere_command_reasoning_parser.py b/vllm/reasoning/cohere_command_reasoning_parser.py index 949c9ff5d99..34066ef2d92 100644 --- a/vllm/reasoning/cohere_command_reasoning_parser.py +++ b/vllm/reasoning/cohere_command_reasoning_parser.py @@ -20,7 +20,6 @@ except ImportError as e: ) from e -from vllm.entrypoints.mcp.tool_server import ToolServer from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ) @@ -481,15 +480,6 @@ class BaseCohereCommandReasoningParser(ReasoningParser): def is_reasoning_end(self, input_ids: Sequence[int]) -> bool: return any(tid == self.end_token_id for tid in reversed(input_ids)) - def prepare_structured_tag( - self, original_tag: str | None, tool_server: ToolServer | None - ) -> str | None: - # Responses API replaces ``structural_tag`` via the reasoning parser. - # Default ``ReasoningParser.prepare_structured_tag`` returns None, which - # would clear a Cohere tag produced in ``adjust_request`` and break - # ``StructuredOutputsParams`` validation. Preserve the existing tag. - return original_tag - def adjust_request( self, request: ChatCompletionRequest | ResponsesRequest ) -> ChatCompletionRequest | ResponsesRequest: From 9ea3a4015b412d146d38ee1b697aafe92979c6ae Mon Sep 17 00:00:00 2001 From: nv-nedelman-1 <49536618+nv-nedelman-1@users.noreply.github.com> Date: Thu, 18 Jun 2026 20:26:09 -0500 Subject: [PATCH 24/75] [Bugfix] Fix corrupt outputs in MoE FP8 LoRA responses and MoE base model responses when LoRAs are loaded (#42120) Signed-off-by: Nicholas Edelman Signed-off-by: Jee Jee Li Co-authored-by: Jee Jee Li Co-authored-by: Jee Jee Li --- tests/lora/test_punica_ops.py | 124 ++++++++++++++++++ vllm/lora/punica_wrapper/punica_gpu.py | 8 +- .../layers/fused_moe/experts/lora_context.py | 7 + .../layers/fused_moe/experts/triton_moe.py | 51 ++++++- .../layers/fused_moe/modular_kernel.py | 10 ++ 5 files changed, 196 insertions(+), 4 deletions(-) diff --git a/tests/lora/test_punica_ops.py b/tests/lora/test_punica_ops.py index 7706d0e2aab..be878472620 100644 --- a/tests/lora/test_punica_ops.py +++ b/tests/lora/test_punica_ops.py @@ -482,3 +482,127 @@ def test_kernels_hidden_size( seq_length=128, add_inputs=True, ) + + +@pytest.mark.parametrize("device", DEVICES) +def test_add_lora_fused_moe_early_exit(device): + """ + Ensures add_lora_fused_moe does not invoke the LoRA kernel or + modify the output tensor when no_lora_flag_cpu is True + """ + from types import SimpleNamespace + + from vllm.lora.punica_wrapper.punica_gpu import PunicaWrapperGPU + + torch.set_default_device(device) + torch.accelerator.set_device_index(device) + + max_loras, num_tokens = 4, 16 + num_experts, top_k, max_lora_rank = 8, 2, 16 + K, N = 256, 128 + + # build PunicaWrapperGPU with minimal lora_config mock + lora_config = SimpleNamespace( + max_loras=max_loras, + specialize_active_lora=False, + ) + wrapper = PunicaWrapperGPU( + max_num_batched_tokens=num_tokens, + max_batches=num_tokens, + device=device, + lora_config=lora_config, + ) + + # simulate a prior LoRA batch so the internal mapping is + # populated with stale LoRA IDs + lora_mapping = torch.zeros( + num_tokens, + dtype=torch.int32, + device=device, + ) + lora_mapping[:8] = 1 + lora_mapping[8:] = 2 + wrapper.token_mapping_meta.prepare_tensors(lora_mapping) + + # simulate a base-model batch (all -1) + base_mapping = torch.full( + (num_tokens,), + -1, + dtype=torch.int32, + device=device, + ) + wrapper.token_mapping_meta.prepare_tensors(base_mapping) + + assert wrapper.token_mapping_meta.no_lora_flag_cpu[0].item() is True + + # dummy tensors for add_lora_fused_moe + y = torch.rand(num_tokens, top_k, N, dtype=torch.bfloat16, device=device) + y_snapshot = y.clone() + x = torch.rand(num_tokens, K, dtype=torch.bfloat16, device=device) + + lora_a_stacked = ( + torch.rand( + max_loras, + num_experts, + max_lora_rank, + K, + dtype=torch.bfloat16, + device=device, + ), + ) + lora_b_stacked = ( + torch.rand( + max_loras, + num_experts, + N, + max_lora_rank, + dtype=torch.bfloat16, + device=device, + ), + ) + topk_weights = torch.ones( + num_tokens, + top_k, + dtype=torch.float32, + device=device, + ) + adapter_enabled = torch.ones( + max_loras + 1, + dtype=torch.int32, + device=device, + ) + shrink_config = expand_config = { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NUM_WARPS": 4, + "NUM_STAGES": 3, + "SPLIT_K": 1, + } + + # call add_lora_fused_moe - the early exit should prevent any + # modification to the output + wrapper.add_lora_fused_moe( + y=y, + x=x, + lora_a_stacked=lora_a_stacked, + lora_b_stacked=lora_b_stacked, + topk_weights=topk_weights, + sorted_token_ids=None, + expert_ids=torch.zeros( + num_tokens * top_k, + dtype=torch.int32, + device=device, + ), + num_tokens_post_padded=None, + max_lora_rank=max_lora_rank, + top_k_num=top_k, + shrink_config=shrink_config, + expand_config=expand_config, + adapter_enabled=adapter_enabled, + ) + + assert torch.equal(y, y_snapshot), ( + "add_lora_fused_moe modified output tensor despite no_lora_flag_cpu=True" + ) diff --git a/vllm/lora/punica_wrapper/punica_gpu.py b/vllm/lora/punica_wrapper/punica_gpu.py index ccf95eb6847..18272354b47 100644 --- a/vllm/lora/punica_wrapper/punica_gpu.py +++ b/vllm/lora/punica_wrapper/punica_gpu.py @@ -446,11 +446,17 @@ class PunicaWrapperGPU(PunicaWrapperBase): _, _, lora_ids, - _, + no_lora_flag, num_active_loras, ) = self.token_mapping_meta.meta_args( x.size(0), self.lora_config.specialize_active_lora ) + + assert no_lora_flag.numel() == 1 + if no_lora_flag.item(): + # None of the inputs require LoRA. + return + if token_lora_mapping is None: token_lora_mapping = token_lora_mapping_meta fused_moe_lora( diff --git a/vllm/model_executor/layers/fused_moe/experts/lora_context.py b/vllm/model_executor/layers/fused_moe/experts/lora_context.py index 404457bb34b..117f744aeea 100644 --- a/vllm/model_executor/layers/fused_moe/experts/lora_context.py +++ b/vllm/model_executor/layers/fused_moe/experts/lora_context.py @@ -59,3 +59,10 @@ class MoELoRAContext: # None means no dispatch happened (non-EP path), in which case callers # fall back to punica_wrapper.token_mapping_meta. local_token_lora_mapping: torch.Tensor | None = None + + # Original unquantized hidden states, stashed by the modular kernel + # before the prepare step potentially quantizes them. Used by + # apply_w13_lora so the LoRA kernel sees correct-magnitude activations + # instead of raw quantized values that are missing the activation scale. + # Set per forward pass; None until the modular kernel writes it. + original_hidden_states: torch.Tensor | None = None diff --git a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py index d81458b3751..0d9b43658f9 100644 --- a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py @@ -77,6 +77,16 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular): def activation_format() -> mk.FusedMoEActivationFormat: return mk.FusedMoEActivationFormat.Standard + @property + def expects_unquantized_inputs(self) -> bool: + # Defer activation quantization to apply() only when LoRA is active AND + # tokens are dispatched across ranks (DP+EP all2all). + return ( + self._lora_context is not None + and self.quant_dtype is not None + and self.moe_config.moe_parallel_config.use_all2all_kernels + ) + @staticmethod def _supports_current_device() -> bool: return current_platform.is_cuda_alike() or current_platform.is_xpu() @@ -223,6 +233,25 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular): torch.float8_e4m3fnuz, ] + # We declared expects_unquantized_inputs (LoRA + DP/EP all2all), so the + # prepare step deferred activation quantization to this kernel: + # `hidden_states` arrives unquantized. Keep the unquantized tensor for + # the LoRA shrink input and quantize a copy here for the base GEMM + # (mirrors what the prepare step would have done, but after the + # all-gather so the layout matches the gathered topk_ids / token map). + lora_unquantized_hidden_states: torch.Tensor | None = None + if self.expects_unquantized_inputs: + assert a1q_scale is None + lora_unquantized_hidden_states = hidden_states + hidden_states, a1q_scale = moe_kernel_quantize_input( + hidden_states, + self.a1_scale, + self.quant_dtype, + self.per_act_token_quant, + self.block_shape, + quantization_emulation=self.quantization_emulation, + ) + E, num_tokens, N, K, top_k_num = self.moe_problem_size( hidden_states, w1, w2, topk_ids ) @@ -280,12 +309,28 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular): # GEMM on the default stream and the LoRA fast-path on aux_stream; # the LoRA writes its delta into a fresh zero buffer (add_inputs= # False) and we sum it into intermediate_cache1 after both finish. - + # + # The LoRA shrink kernel needs unquantized, gathered-layout + # activations. When activation quant was deferred to this kernel + # (expects_unquantized_inputs), the input we quantized above is exactly + # that, so use it directly. Otherwise fall back to the context stash + # (e.g. weight-only quant), guarding on a row-count match so a + # DP-gathered layout never indexes a local stash out of bounds. sorted_token_ids_lora = None expert_ids_lora = None num_tokens_post_padded_lora = None token_lora_mapping = None lora_context = self._lora_context + if lora_unquantized_hidden_states is not None: + lora_x = lora_unquantized_hidden_states + elif ( + lora_context is not None + and lora_context.original_hidden_states is not None + and lora_context.original_hidden_states.shape[0] == hidden_states.shape[0] + ): + lora_x = lora_context.original_hidden_states + else: + lora_x = hidden_states def _base_w13_fn(): invoke_fused_moe_triton_kernel( @@ -322,7 +367,7 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular): return self.apply_w13_lora( lora_context, y=lora_delta_w13, - x=hidden_states, + x=lora_x, topk_ids=topk_ids, topk_weights=topk_weights, expert_map=expert_map, @@ -359,7 +404,7 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular): ) = self.apply_w13_lora( lora_context, y=intermediate_cache1, - x=hidden_states, + x=lora_x, topk_ids=topk_ids, topk_weights=topk_weights, expert_map=expert_map, diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index e80224be70f..0e55e827c20 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -1407,6 +1407,13 @@ class FusedMoEKernelModularImpl: apply_router_weight_on_input, ) + # Stash the original unquantized hidden states on the LoRA context + # so apply_w13_lora sees correct-magnitude activations instead of + # the potentially quantized values produced by _prepare(). + lora_ctx = getattr(self.fused_experts, "_lora_context", None) + if lora_ctx is not None: + lora_ctx.original_hidden_states = hidden_states + fused_out = self._fused_experts( in_dtype=hidden_states.dtype, a1q=a1q, @@ -1424,6 +1431,9 @@ class FusedMoEKernelModularImpl: output_alias=output, ) + if lora_ctx is not None: + lora_ctx.original_hidden_states = None + return self._finalize( output, fused_out, From ab666069935c1f23e8ef56038b4659ac9e8f19f8 Mon Sep 17 00:00:00 2001 From: Jared Wen Date: Fri, 19 Jun 2026 09:57:51 +0800 Subject: [PATCH 25/75] [bugfix]Indexer init skip and MTP TopK share for iteration (#45895) Signed-off-by: JaredforReal --- .../layers/attention/mla_attention.py | 6 +++ vllm/model_executor/layers/mla.py | 1 + vllm/model_executor/models/deepseek_mtp.py | 8 +++- vllm/model_executor/models/deepseek_v2.py | 39 +++++++++++-------- .../backends/mla/flashinfer_mla_sparse.py | 10 +++-- .../attention/backends/mla/flashmla_sparse.py | 8 +++- .../backends/mla/rocm_aiter_mla_sparse.py | 10 +++-- .../attention/backends/mla/xpu_mla_sparse.py | 10 +++-- vllm/v1/spec_decode/llm_base_proposer.py | 7 ++++ 9 files changed, 69 insertions(+), 30 deletions(-) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 21e3215479f..ab3874c5dad 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -349,6 +349,7 @@ class MLAAttention(nn.Module, AttentionLayerBase): attn_backend: type[AttentionBackend] | None = None, use_sparse: bool = False, indexer: object | None = None, + topk_indices_buffer: torch.Tensor | None = None, **extra_impl_args, ): super().__init__() @@ -437,6 +438,11 @@ class MLAAttention(nn.Module, AttentionLayerBase): ) cache_config.enable_prefix_caching = False + # Sparse MLA reads top-k indices from a shared buffer. Pass it + # explicitly so backbone "skip" layers (indexer=None) still find it. + if use_sparse: + extra_impl_args["topk_indices_buffer"] = topk_indices_buffer + impl_cls = cast(type[MLAAttentionImpl], self.attn_backend.get_impl_cls()) self.impl = impl_cls( # type: ignore[assignment] # impl_cls always returns an MLAAttentionImpl subclass num_heads=self.num_heads, diff --git a/vllm/model_executor/layers/mla.py b/vllm/model_executor/layers/mla.py index 856f6bb8a3c..66a95b43c71 100644 --- a/vllm/model_executor/layers/mla.py +++ b/vllm/model_executor/layers/mla.py @@ -112,6 +112,7 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer): kv_b_proj=self.kv_b_proj, use_sparse=self.is_sparse, indexer=self.indexer, + topk_indices_buffer=mla_modules.topk_indices_buffer, ) self.prefix = prefix diff --git a/vllm/model_executor/models/deepseek_mtp.py b/vllm/model_executor/models/deepseek_mtp.py index d46eb67c5ea..88f33ac021b 100644 --- a/vllm/model_executor/models/deepseek_mtp.py +++ b/vllm/model_executor/models/deepseek_mtp.py @@ -119,8 +119,12 @@ class DeepSeekMultiTokenPredictorLayer(nn.Module): hidden_states=hidden_states, residual=None, ) - hidden_states = residual + hidden_states - return hidden_states + hidden_states = residual + hidden_states # pre-final-norm (logits hidden) + # Recycle the post-final-norm hidden into the next draft step. + # compute_logits applies shared_head (== final norm) to the pre-norm + # element, so logits and the recycle each get exactly one final-norm. + # Matches SGLang's deepseek_nextn. + return hidden_states, self.shared_head(hidden_states) class DeepSeekMultiTokenPredictor(nn.Module): diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 80d518dacbd..22c4003d3fa 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -998,8 +998,29 @@ class DeepseekV2MLAAttention(nn.Module): self.is_v32 = hasattr(config, "index_topk") + # IndexCache config + # Refer: https://arxiv.org/abs/2603.12201 for more details. _skip_topk = False - if self.is_v32: + _index_topk_freq = getattr(config, "index_topk_freq", 1) + _index_topk_pattern = getattr(config, "index_topk_pattern", None) + _index_skip_topk_offset = getattr(config, "index_skip_topk_offset", 2) + layer_id = extract_layer_index(prefix) + + if _index_topk_pattern is None: + _skip_topk = ( + max(layer_id - _index_skip_topk_offset + 1, 0) % _index_topk_freq != 0 + ) + elif 0 <= layer_id < len(_index_topk_pattern): + _skip_topk = _index_topk_pattern[layer_id] == "S" + + # The skip pattern only governs backbone layers. MTP/nextn layers + # (layer_id >= num_hidden_layers) always build a full indexer: they + # compute indices at draft step 0 and toggle at runtime via + # set_skip_topk (index_share_for_mtp_iteration). + _num_hidden_layers = getattr(config, "num_hidden_layers", None) + is_mtp_layer = _num_hidden_layers is not None and layer_id >= _num_hidden_layers + + if self.is_v32 and (not _skip_topk or is_mtp_layer): self.indexer_rope_emb = get_rope( qk_rope_head_dim, max_position=max_position_embeddings, @@ -1017,22 +1038,6 @@ class DeepseekV2MLAAttention(nn.Module): f"{prefix}.indexer", is_inplace_rope=self.indexer_rope_emb.enabled(), ) - - # IndexCache config - # Refer: https://arxiv.org/abs/2603.12201 for more details. - _index_topk_freq = getattr(config, "index_topk_freq", 1) - _index_topk_pattern = getattr(config, "index_topk_pattern", None) - _index_skip_topk_offset = getattr(config, "index_skip_topk_offset", 2) - layer_id = extract_layer_index(prefix) - - if _index_topk_pattern is None: - _skip_topk = ( - max(layer_id - _index_skip_topk_offset + 1, 0) % _index_topk_freq - != 0 - ) - elif 0 <= layer_id < len(_index_topk_pattern): - _skip_topk = _index_topk_pattern[layer_id] == "S" - else: self.indexer_rope_emb = None self.indexer = None diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index aa6301c13bf..01716f567d0 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -271,7 +271,7 @@ class FlashInferMLASparseImpl(SparseMLAAttentionImpl[FlashInferMLASparseMetadata attn_type: str, kv_sharing_target_layer_name: str | None, # MLA Specific Arguments - topk_indice_buffer: torch.Tensor | None = None, + topk_indices_buffer: torch.Tensor | None = None, indexer: "Indexer | None" = None, **mla_args, ) -> None: @@ -301,8 +301,12 @@ class FlashInferMLASparseImpl(SparseMLAAttentionImpl[FlashInferMLASparseMetadata self.qk_nope_head_dim: int = mla_args["qk_nope_head_dim"] self.qk_rope_head_dim: int = mla_args["qk_rope_head_dim"] - assert indexer is not None, "Indexer required for sparse MLA" - self.topk_indices_buffer: torch.Tensor | None = indexer.topk_indices_buffer + # The indexer carries the shared buffer for normal layers and tests; + # the explicitly-passed buffer covers backbone skip layers, whose + # indexer is not constructed (see deepseek_v2.py). + self.topk_indices_buffer: torch.Tensor | None = ( + indexer.topk_indices_buffer if indexer is not None else topk_indices_buffer + ) self._workspace_buffer: torch.Tensor | None = None self.bmm1_scale: float | None = None diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index 2da71f9d2c3..6d8dfe13128 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -568,8 +568,12 @@ class FlashMLASparseImpl(SparseMLAAttentionImpl[FlashMLASparseMetadata]): self.kv_cache_dtype = kv_cache_dtype self.kv_lora_rank: int = mla_args["kv_lora_rank"] self.softmax_scale = scale - assert indexer is not None - self.topk_indices_buffer: torch.Tensor | None = indexer.topk_indices_buffer + # The indexer carries the shared buffer for normal layers and tests; + # the explicitly-passed buffer covers backbone skip layers, whose + # indexer is not constructed (see deepseek_v2.py). + self.topk_indices_buffer: torch.Tensor | None = ( + indexer.topk_indices_buffer if indexer is not None else topk_indices_buffer + ) # Prefill BF16 kernel requires 64 on Hopper, 128 on Blackwell self.prefill_padding = ( 128 if current_platform.is_device_capability_family(100) else 64 diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py index 705ac167f20..1225352acee 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py @@ -629,7 +629,7 @@ class ROCMAiterMLASparseImpl(SparseMLAAttentionImpl[ROCMAiterMLASparseMetadata]) attn_type: str, kv_sharing_target_layer_name: str | None, # MLA Specific Arguments - topk_indice_buffer: torch.Tensor | None = None, + topk_indices_buffer: torch.Tensor | None = None, indexer: "Indexer | None" = None, **mla_args, ) -> None: @@ -642,8 +642,12 @@ class ROCMAiterMLASparseImpl(SparseMLAAttentionImpl[ROCMAiterMLASparseMetadata]) self.kv_cache_dtype = kv_cache_dtype self.kv_lora_rank: int = mla_args["kv_lora_rank"] self.softmax_scale = scale - assert indexer is not None - self.topk_indices_buffer: torch.Tensor | None = indexer.topk_indices_buffer + # The indexer carries the shared buffer for normal layers and tests; + # the explicitly-passed buffer covers backbone skip layers, whose + # indexer is not constructed (see deepseek_v2.py). + self.topk_indices_buffer: torch.Tensor | None = ( + indexer.topk_indices_buffer if indexer is not None else topk_indices_buffer + ) vllm_config = get_current_vllm_config() max_tokens = vllm_config.scheduler_config.max_num_batched_tokens diff --git a/vllm/v1/attention/backends/mla/xpu_mla_sparse.py b/vllm/v1/attention/backends/mla/xpu_mla_sparse.py index 2fa91d01838..9aad4532103 100644 --- a/vllm/v1/attention/backends/mla/xpu_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/xpu_mla_sparse.py @@ -184,7 +184,7 @@ class XPUMLASparseImpl(SparseMLAAttentionImpl[XPUMLASparseMetadata]): attn_type: str, kv_sharing_target_layer_name: str | None, # MLA Specific Arguments - topk_indice_buffer: torch.Tensor | None = None, + topk_indices_buffer: torch.Tensor | None = None, indexer: Optional["Indexer"] = None, **mla_args, ) -> None: @@ -195,8 +195,12 @@ class XPUMLASparseImpl(SparseMLAAttentionImpl[XPUMLASparseMetadata]): self.kv_cache_dtype = kv_cache_dtype self.kv_lora_rank: int = mla_args["kv_lora_rank"] self.softmax_scale = scale - assert indexer is not None - self.topk_indices_buffer: torch.Tensor | None = indexer.topk_indices_buffer + # The indexer carries the shared buffer for normal layers and tests; + # the explicitly-passed buffer covers backbone skip layers, whose + # indexer is not constructed (see deepseek_v2.py). + self.topk_indices_buffer: torch.Tensor | None = ( + indexer.topk_indices_buffer if indexer is not None else topk_indices_buffer + ) def _forward_bf16_kv( self, diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index d4f2c1007b0..b7c01d3ec1c 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -918,6 +918,13 @@ class SpecDecodeBaseProposer: return per_group_attn_metadata, per_layer_attn_metadata def model_returns_tuple(self) -> bool: + if self.method == "mtp": + # DeepSeek-family MTP (deepseek_mtp.py) recycles the post-final- + # norm hidden, so its forward returns (logit_hidden, + # recycle_hidden). Other MTP families return a single tensor. + return "DeepSeekMTPModel" in ( + self.draft_model_config.hf_config.architectures or [] + ) return self.method not in ("mtp", "draft_model", "dflash") def prepare_next_token_ids_cpu( From 2a6c6b94293edb54bff8088a5d64b703aac187ff Mon Sep 17 00:00:00 2001 From: "Jeff (Junze) Ma" <93145857+majunze2001@users.noreply.github.com> Date: Thu, 18 Jun 2026 20:10:12 -0700 Subject: [PATCH 26/75] [DeepSeek-V4] Support TEP=16 for the block-FP8 shared expert (#46001) Signed-off-by: Jeff Ma Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- vllm/models/deepseek_v4/nvidia/model.py | 44 +++++++++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 364754f9d77..868fc3f5fdb 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -64,6 +64,7 @@ from vllm.models.deepseek_v4.nvidia.flashinfer_sparse import ( from vllm.models.deepseek_v4.nvidia.flashmla import DeepseekV4FlashMLAAttention from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs from vllm.sequence import IntermediateTensors +from vllm.utils.math_utils import cdiv from vllm.v1.attention.backends.registry import AttentionBackendEnum @@ -85,6 +86,15 @@ class DeepseekV4MLP(nn.Module): # across the ranks within the tp_group. In this case the weights are # replicated and no collective ops are needed. # Otherwise we use standard TP with an allreduce at the end. + # + # Block-FP8 shards in whole 128-blocks; cdiv rounds the per-rank block + # count up so the linear's even TP split stays block-aligned, with the + # trailing ranks zero-filled by load_weights. + block_size = getattr(quant_config, "weight_block_size", None) + if block_size is not None and not is_sequence_parallel: + tp_size = get_tensor_model_parallel_world_size() + n_local = cdiv(intermediate_size // block_size[0], tp_size) + intermediate_size = n_local * block_size[0] * tp_size self.gate_up_proj = MergedColumnParallelLinear( hidden_size, [intermediate_size] * 2, @@ -892,6 +902,8 @@ class DeepseekV4Model(nn.Module): config = vllm_config.model_config.hf_config quant_config = vllm_config.quant_config self.config = config + self.quant_config = quant_config + self.parallel_config = vllm_config.parallel_config self.use_mega_moe = ( vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe" ) @@ -1080,7 +1092,17 @@ class DeepseekV4Model(nn.Module): # Pre-compute expert mapping ONCE. expert_mapping = self.get_expert_mapping() + # Block-FP8 shared experts: pad the intermediate up to the TP-uniform + # block count so the standard loaders below slice it evenly (trailing + # ranks land on the zero pad). SP / unquantized ones need no padding. + pad_shared_expert = ( + getattr(self.quant_config, "weight_block_size", None) is not None + and not self.parallel_config.use_sequence_parallel_moe + ) + for name, loaded_weight in weights: + if pad_shared_expert and ".shared_experts." in name: + loaded_weight = self._pad_shared_expert_weight(name, loaded_weight) for param_name, weight_name, shard_id in stacked_params_mapping: # Skip non-stacked layers and experts (experts handled below). if ".experts." in name: @@ -1155,6 +1177,28 @@ class DeepseekV4Model(nn.Module): return loaded_params + def _pad_shared_expert_weight( + self, name: str, loaded_weight: torch.Tensor + ) -> torch.Tensor: + """Zero-pad a block-FP8 shared-expert weight/scale on its intermediate + axis so the standard TP loaders split it into even, block-aligned shards + (trailing ranks get the zero pad). gate (w1)/up (w3) [I, H] pad dim 0; + down (w2 -> down_proj) [H, I] pads dim 1. + """ + block_size = getattr(self.quant_config, "weight_block_size", None) + assert block_size is not None + # Round the intermediate axis up to a whole number of TP shards. The axis + # is in elements for weights (step = block) and in blocks for scales. + step = 1 if name.endswith("weight_scale_inv") else block_size[0] + dim = 1 if ".down_proj." in name else 0 + mult = get_tensor_model_parallel_world_size() * step + pad = cdiv(loaded_weight.shape[dim], mult) * mult - loaded_weight.shape[dim] + if pad == 0: + return loaded_weight + pad_shape = list(loaded_weight.shape) + pad_shape[dim] = pad + return torch.cat([loaded_weight, loaded_weight.new_zeros(pad_shape)], dim=dim) + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: first_layer = next(iter(islice(self.layers, self.start_layer, self.end_layer))) if first_layer.ffn.use_mega_moe: From c9135db27cafb853af5e2cb86c1a0b3c6b5b8c91 Mon Sep 17 00:00:00 2001 From: Samuel Shen Date: Thu, 18 Jun 2026 20:21:36 -0700 Subject: [PATCH 27/75] [Docs] Update stale LMCache examples (#45762) Signed-off-by: Samuel Shen --- .../integrations/production-stack.md | 2 +- docs/features/disagg_prefill.md | 2 +- examples/disaggregated/lmcache/README.md | 48 ++++-- .../lmcache/cpu_offload_lmcache.py | 38 +---- .../lmcache/cpu_offload_lmcache_mp.sh | 43 ++++++ .../lmcache/disagg_prefill_lmcache_v0.py | 144 ------------------ .../disagg_vllm_launcher.sh | 2 - .../lmcache/kv_cache_sharing_lmcache_v1.py | 2 - 8 files changed, 84 insertions(+), 197 deletions(-) create mode 100755 examples/disaggregated/lmcache/cpu_offload_lmcache_mp.sh delete mode 100644 examples/disaggregated/lmcache/disagg_prefill_lmcache_v0.py diff --git a/docs/deployment/integrations/production-stack.md b/docs/deployment/integrations/production-stack.md index 4db595164e3..d93300a2b06 100644 --- a/docs/deployment/integrations/production-stack.md +++ b/docs/deployment/integrations/production-stack.md @@ -4,7 +4,7 @@ Deploying vLLM on Kubernetes is a scalable and efficient way to serve machine le * **Upstream vLLM compatibility** – It wraps around upstream vLLM without modifying its code. * **Ease of use** – Simplified deployment via Helm charts and observability through Grafana dashboards. -* **High performance** – Optimized for LLM workloads with features like multimodel support, model-aware and prefix-aware routing, fast vLLM bootstrapping, and KV cache offloading with [LMCache](https://github.com/LMCache/LMCache), among others. +* **High performance** – Optimized for LLM workloads with features like multimodel support, model-aware and prefix-aware routing, fast vLLM bootstrapping, and KV cache offloading with [LMCache](https://github.com/LMCache/LMCache) (wired up in vLLM via `--kv-offloading-backend lmcache`; see the [LMCache examples](https://github.com/vllm-project/vllm/tree/main/examples/disaggregated/lmcache) and [docs.lmcache.ai](https://docs.lmcache.ai)), among others. If you are new to Kubernetes, don't worry: in the vLLM production stack [repo](https://github.com/vllm-project/production-stack), we provide a step-by-step [guide](https://github.com/vllm-project/production-stack/blob/main/tutorials/00-install-kubernetes-env.md) and a [short video](https://www.youtube.com/watch?v=EsTJbQtzj0g) to set up everything and get started in **4 minutes**! diff --git a/docs/features/disagg_prefill.md b/docs/features/disagg_prefill.md index 8352d2f20e0..578343096df 100644 --- a/docs/features/disagg_prefill.md +++ b/docs/features/disagg_prefill.md @@ -20,7 +20,7 @@ Two main reasons: Now supports 9 types of connectors: - **ExampleConnector**: refer to [examples/disaggregated/example_connector/run.sh](../../examples/disaggregated/example_connector/run.sh) for the example usage of ExampleConnector disaggregated prefilling. -- **LMCacheConnectorV1**: refer to [examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh](../../examples/disaggregated/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. +- **LMCacheConnectorV1**: refer to [examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh](../../examples/disaggregated/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. LMCache also offers a multi-process (MP) mode via `LMCacheMPConnector`, where a standalone `lmcache server` holds the KV cache shared by one or more vLLM instances; see the [LMCache examples](../../examples/disaggregated/lmcache/README.md) and the [LMCache docs](https://docs.lmcache.ai) for setup. - **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). For feature compatibility details, see [NixlConnector Compatibility Matrix](nixl_connector_compatibility.md). You may specify one or multiple NIXL transfer backends, such as: ```bash diff --git a/examples/disaggregated/lmcache/README.md b/examples/disaggregated/lmcache/README.md index 759be55d6f1..87fec826842 100644 --- a/examples/disaggregated/lmcache/README.md +++ b/examples/disaggregated/lmcache/README.md @@ -1,10 +1,38 @@ # LMCache Examples -This folder demonstrates how to use LMCache for disaggregated prefilling, CPU offloading and KV cache sharing. +This folder demonstrates how to use LMCache with vLLM v1 for KV cache +offloading, disaggregated prefilling, and KV cache sharing. -## 1. Disaggregated Prefill in vLLM v1 +## Integration modes -This example demonstrates how to run LMCache with disaggregated prefill using NIXL on a single node. +LMCache integrates with vLLM v1 in two ways: + +- **In-process mode** (`LMCacheConnectorV1`): LMCache runs inside the vLLM + process and is configured through environment variables or a YAML config + file (`LMCACHE_CONFIG_FILE`). This is the simplest way to add single-node + CPU/disk offloading. +- **Multi-process (MP) mode** (`LMCacheMPConnector`): LMCache runs as a + standalone server (`lmcache server`) that owns the KV cache storage; one or + more vLLM instances connect to it. This is the recommended mode for + distributed KV storage and for sharing KV cache across instances. See the + [LMCache docs](https://docs.lmcache.ai) for the full MP setup. + +## 1. CPU offload (in-process) + +- `python cpu_offload_lmcache.py` - CPU offloading with `LMCacheConnectorV1` + for vLLM v1. + +## 2. CPU offload (multi-process) + +- `bash cpu_offload_lmcache_mp.sh` - CPU offloading with `LMCacheMPConnector`, + using a standalone `lmcache server`. vLLM provides a built-in shortcut for + this setup via `--kv-offloading-backend lmcache` and + `--kv-offloading-size `. + +## 3. Disaggregated Prefill in vLLM v1 + +This example demonstrates how to run LMCache with disaggregated prefill using +NIXL on a single node. ### Prerequisites @@ -46,15 +74,7 @@ The main script generates several log files: - `decoder.log` - Logs from the decode server - `proxy.log` - Logs from the proxy server -## 2. CPU Offload Examples +## 4. KV Cache Sharing -- `python cpu_offload_lmcache.py -v v0` - CPU offloading implementation for vLLM v0 -- `python cpu_offload_lmcache.py -v v1` - CPU offloading implementation for vLLM v1 - -## 3. KV Cache Sharing - -The `kv_cache_sharing_lmcache_v1.py` example demonstrates how to share KV caches between vLLM v1 instances. - -## 4. Disaggregated Prefill in vLLM v0 - -The `disaggregated_prefill_lmcache_v0.py` provides an example of how to run disaggregated prefill in vLLM v0. +The `kv_cache_sharing_lmcache_v1.py` example demonstrates how to share KV +caches between vLLM v1 instances through a centralized LMCache server. diff --git a/examples/disaggregated/lmcache/cpu_offload_lmcache.py b/examples/disaggregated/lmcache/cpu_offload_lmcache.py index 53036b3eb0f..b67a929e5d9 100644 --- a/examples/disaggregated/lmcache/cpu_offload_lmcache.py +++ b/examples/disaggregated/lmcache/cpu_offload_lmcache.py @@ -1,20 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ -This file demonstrates the example usage of cpu offloading -with LMCache in vLLM v1 or v0. - -Usage: - - Specify vLLM version - - -v v0 : Use LMCacheConnector - model = mistralai/Mistral-7B-Instruct-v0.2 - (Includes enable_chunked_prefill = True) - - -v v1 : Use LMCacheConnectorV1 (default) - model = meta-llama/Meta-Llama-3.1-8B-Instruct - (Without enable_chunked_prefill) +This file demonstrates the example usage of CPU offloading +with LMCache in vLLM v1. Note that `lmcache` is needed to run this example. Requirements: @@ -23,7 +11,6 @@ Learn more about LMCache environment setup, please refer to: https://docs.lmcache.ai/getting_started/installation.html """ -import argparse import contextlib import os import time @@ -39,8 +26,6 @@ from vllm.engine.arg_utils import EngineArgs def setup_environment_variables(): # LMCache-related environment variables - # Use experimental features in LMCache - os.environ["LMCACHE_USE_EXPERIMENTAL"] = "True" # LMCache is set to use 256 tokens per chunk os.environ["LMCACHE_CHUNK_SIZE"] = "256" # Enable local CPU backend in LMCache @@ -50,9 +35,9 @@ def setup_environment_variables(): @contextlib.contextmanager -def build_llm_with_lmcache(lmcache_connector: str, model: str): +def build_llm_with_lmcache(model: str): ktc = KVTransferConfig( - kv_connector=lmcache_connector, + kv_connector="LMCacheConnectorV1", kv_role="kv_both", ) # Set GPU memory utilization to 0.8 for an A40 GPU with 40GB @@ -92,23 +77,10 @@ def print_output( print("-" * 50) -def parse_args(): - parser = argparse.ArgumentParser() - parser.add_argument( - "-v", - "--version", - choices=["v0", "v1"], - default="v1", - help="Specify vLLM version (default: v1)", - ) - return parser.parse_args() - - def main(): - lmcache_connector = "LMCacheConnectorV1" model = "meta-llama/Meta-Llama-3.1-8B-Instruct" setup_environment_variables() - with build_llm_with_lmcache(lmcache_connector, model) as llm: + with build_llm_with_lmcache(model) as llm: # This example script runs two requests with a shared prefix. # Define the shared prompt and specific prompts shared_prompt = "Hello, how are you?" * 1000 diff --git a/examples/disaggregated/lmcache/cpu_offload_lmcache_mp.sh b/examples/disaggregated/lmcache/cpu_offload_lmcache_mp.sh new file mode 100755 index 00000000000..2372eabe1a8 --- /dev/null +++ b/examples/disaggregated/lmcache/cpu_offload_lmcache_mp.sh @@ -0,0 +1,43 @@ +#!/bin/bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# CPU offloading with LMCache in multi-process (MP) mode. +# +# In MP mode, LMCache runs as a standalone server process (`lmcache server`) +# that owns the KV cache storage. One or more vLLM instances connect to it via +# the `LMCacheMPConnector`. This is the recommended way to run LMCache for +# distributed KV storage and for sharing KV cache across vLLM instances. +# +# vLLM ships a built-in shortcut for this setup: pass `--kv-offloading-backend +# lmcache` together with `--kv-offloading-size ` and vLLM wires up the +# `LMCacheMPConnector` for you (it defaults to the LMCache server at +# tcp://localhost:5555, matching the `lmcache server` default). +# +# Requires `lmcache` to be installed (`pip install lmcache`). +# Learn more: https://docs.lmcache.ai +set -euo pipefail + +MODEL=${MODEL:-meta-llama/Meta-Llama-3.1-8B-Instruct} + +# 1. Launch the standalone LMCache server (binds tcp://localhost:5555 by +# default). `--l1-size-gb` sets the CPU memory budget for the L1 cache. +echo "Starting LMCache server..." +lmcache server --host localhost --port 5555 --l1-size-gb 5 & +LMCACHE_SERVER_PID=$! +trap 'kill $LMCACHE_SERVER_PID 2>/dev/null || true' EXIT + +# 2. Launch vLLM and offload KV cache to the LMCache server. +# The MP connector currently requires the non-hybrid KV cache manager. +echo "Starting vLLM server with LMCache MP offloading..." +vllm serve "$MODEL" \ + --port 8000 \ + --kv-offloading-size 5 \ + --kv-offloading-backend lmcache \ + --disable-hybrid-kv-cache-manager + +# Equivalent explicit configuration (instead of the two flags above): +# --kv-transfer-config \ +# '{"kv_connector":"LMCacheMPConnector","kv_role":"kv_both", +# "kv_connector_extra_config":{"lmcache.mp.host":"tcp://localhost", +# "lmcache.mp.port":5555}}' diff --git a/examples/disaggregated/lmcache/disagg_prefill_lmcache_v0.py b/examples/disaggregated/lmcache/disagg_prefill_lmcache_v0.py deleted file mode 100644 index 6669eb3fb3d..00000000000 --- a/examples/disaggregated/lmcache/disagg_prefill_lmcache_v0.py +++ /dev/null @@ -1,144 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This file demonstrates the example usage of disaggregated prefilling -with LMCache. -We will launch 2 vllm instances (GPU 0 for prefill and GPU 1 for decode), -and launch an additional LMCache server. -KV cache is transferred in the following manner: -vLLM prefill node -> LMCache server -> vLLM decode node. - -Note that `pip install lmcache` is needed to run this example. -Learn more about LMCache in https://github.com/LMCache/LMCache. -""" - -import os -import subprocess -import time -from multiprocessing import Event, Process - -from lmcache.experimental.cache_engine import LMCacheEngineBuilder -from lmcache.integration.vllm.utils import ENGINE_NAME - -from vllm import LLM, SamplingParams -from vllm.config import KVTransferConfig - -# LMCache-related environment variables -# The port to start LMCache server -port = 8100 -# Use experimental features in LMCache -os.environ["LMCACHE_USE_EXPERIMENTAL"] = "True" -# LMCache is set to use 256 tokens per chunk -os.environ["LMCACHE_CHUNK_SIZE"] = "256" -# Disable local CPU backend in LMCache -os.environ["LMCACHE_LOCAL_CPU"] = "False" -# Set local CPU memory buffer limit to 5.0 GB -os.environ["LMCACHE_MAX_LOCAL_CPU_SIZE"] = "5.0" -# Set the remote URL for LMCache server -os.environ["LMCACHE_REMOTE_URL"] = f"lm://localhost:{port}" -# Set the serializer/deserializer between vllm and LMCache server -# `naive` indicates using raw bytes of the tensor without any compression -os.environ["LMCACHE_REMOTE_SERDE"] = "naive" - -prompts = [ - "Hello, how are you?" * 1000, -] - - -def run_prefill(prefill_done, prompts): - # We use GPU 0 for prefill node. - os.environ["CUDA_VISIBLE_DEVICES"] = "0" - - sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=1) - - ktc = KVTransferConfig( - kv_connector="LMCacheConnector", - kv_role="kv_producer", - kv_rank=0, - kv_parallel_size=2, - ) - # Set GPU memory utilization to 0.8 for an A40 GPU with 40GB - # memory. Reduce the value if your GPU has less memory. - llm = LLM( - model="mistralai/Mistral-7B-Instruct-v0.2", - kv_transfer_config=ktc, - max_model_len=8000, - gpu_memory_utilization=0.8, - enforce_eager=True, - ) - - # llm.generate(prompts, sampling_params) - outputs = llm.generate(prompts, sampling_params) - for output in outputs: - generated_text = output.outputs[0].text - print(f"Generated text: {generated_text!r}") - print("Prefill node is finished.") - prefill_done.set() - - # Clean up lmcache backend - LMCacheEngineBuilder.destroy(ENGINE_NAME) - - -def run_decode(prefill_done, prompts, timeout=1): - # We use GPU 1 for decode node. - os.environ["CUDA_VISIBLE_DEVICES"] = "1" - - sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10) - - ktc = KVTransferConfig( - kv_connector="LMCacheConnector", - kv_role="kv_consumer", - kv_rank=1, - kv_parallel_size=2, - ) - # Set GPU memory utilization to 0.8 for an A40 GPU with 40GB - # of memory. Reduce the value if your GPU has less memory. - llm = LLM( - model="mistralai/Mistral-7B-Instruct-v0.2", - kv_transfer_config=ktc, - max_model_len=8000, - gpu_memory_utilization=0.8, - enforce_eager=True, - ) - - print("Waiting for prefill node to finish...") - prefill_done.wait() - time.sleep(timeout) - - outputs = llm.generate(prompts, sampling_params) - for output in outputs: - generated_text = output.outputs[0].text - print(f"Generated text: {generated_text!r}") - - # Clean up lmcache backend - LMCacheEngineBuilder.destroy(ENGINE_NAME) - - -def run_lmcache_server(port): - server_proc = subprocess.Popen( - ["python", "-m", "lmcache.experimental.server", "localhost", str(port)] - ) - return server_proc - - -def main(): - prefill_done = Event() - prefill_process = Process(target=run_prefill, args=(prefill_done, prompts)) - decode_process = Process(target=run_decode, args=(prefill_done, prompts)) - lmcache_server_process = run_lmcache_server(port) - - # Start prefill node - prefill_process.start() - - # Start decode node - decode_process.start() - - # Clean up the processes - decode_process.join() - prefill_process.terminate() - lmcache_server_process.terminate() - lmcache_server_process.wait() - - -if __name__ == "__main__": - main() diff --git a/examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh b/examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh index 363c35028aa..61e578460c4 100644 --- a/examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh +++ b/examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh @@ -30,7 +30,6 @@ if [[ $1 == "prefiller" ]]; then UCX_TLS=cuda_ipc,cuda_copy,tcp \ LMCACHE_CONFIG_FILE=$prefill_config_file \ - LMCACHE_USE_EXPERIMENTAL=True \ VLLM_ENABLE_V1_MULTIPROCESSING=1 \ VLLM_WORKER_MULTIPROC_METHOD=spawn \ CUDA_VISIBLE_DEVICES=0 \ @@ -47,7 +46,6 @@ elif [[ $1 == "decoder" ]]; then UCX_TLS=cuda_ipc,cuda_copy,tcp \ LMCACHE_CONFIG_FILE=$decode_config_file \ - LMCACHE_USE_EXPERIMENTAL=True \ VLLM_ENABLE_V1_MULTIPROCESSING=1 \ VLLM_WORKER_MULTIPROC_METHOD=spawn \ CUDA_VISIBLE_DEVICES=1 \ diff --git a/examples/disaggregated/lmcache/kv_cache_sharing_lmcache_v1.py b/examples/disaggregated/lmcache/kv_cache_sharing_lmcache_v1.py index 46e2d903d4b..489ff132122 100644 --- a/examples/disaggregated/lmcache/kv_cache_sharing_lmcache_v1.py +++ b/examples/disaggregated/lmcache/kv_cache_sharing_lmcache_v1.py @@ -26,8 +26,6 @@ from vllm.config import KVTransferConfig # LMCache-related environment variables # The port to start LMCache server port = 8100 -# Use experimental features in LMCache -os.environ["LMCACHE_USE_EXPERIMENTAL"] = "True" # LMCache is set to use 256 tokens per chunk os.environ["LMCACHE_CHUNK_SIZE"] = "256" # Disable local CPU backend in LMCache From ecf9d83520eb217401b47d8a5451a27c5231b8c2 Mon Sep 17 00:00:00 2001 From: Oxana Korzh Date: Thu, 18 Jun 2026 22:06:56 -0600 Subject: [PATCH 28/75] [AMD][CI] Fix Language Models Test (Extended Generation) failures (#45509) Signed-off-by: Oxana Korzh Co-authored-by: Claude Co-authored-by: Cursor --- tests/models/language/generation/test_common.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/models/language/generation/test_common.py b/tests/models/language/generation/test_common.py index 1b6c8ef5583..50c87d7729e 100644 --- a/tests/models/language/generation/test_common.py +++ b/tests/models/language/generation/test_common.py @@ -130,8 +130,12 @@ def test_models( monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") if model == "TitanML/tiny-mixtral": # Untrained model: near-uniform logits make argmax sensitive to - # AITER's bfloat16 rounding error in plain rms_norm. + # AITER's bfloat16 rounding error. Route the plain rms_norm and the + # fused MoE (whose near-uniform router logits flip expert selection + # under ~1 ULP drift) through the native kernels for this model. + # See ROCm/aiter#3806 for the tracking issue and minimal repro. monkeypatch.setenv("VLLM_ROCM_USE_AITER_RMSNORM", "0") + monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "0") elif use_rocm_aiter and model not in AITER_MODEL_LIST: # Skip model that are not using AITER tests. # When more AITER kernels are added, this list will not be From ec67d7ae619435f5f27279081f40e3a733ca9ab7 Mon Sep 17 00:00:00 2001 From: Kunshang Ji Date: Fri, 19 Jun 2026 15:37:20 +0800 Subject: [PATCH 29/75] [xpu] bump up vllm-xpu-kernels v0.1.10 and upgrade 2618 umd (#40367) Signed-off-by: Kunshang Ji Signed-off-by: Kunshang Ji --- .buildkite/intel_jobs/test-intel.yaml | 2 +- docker/Dockerfile.xpu | 14 +++++++------- docs/getting_started/installation/gpu.xpu.inc.md | 1 + requirements/xpu.txt | 2 +- 4 files changed, 10 insertions(+), 9 deletions(-) diff --git a/.buildkite/intel_jobs/test-intel.yaml b/.buildkite/intel_jobs/test-intel.yaml index 7ca48e6841f..f365bf76512 100644 --- a/.buildkite/intel_jobs/test-intel.yaml +++ b/.buildkite/intel_jobs/test-intel.yaml @@ -67,7 +67,7 @@ steps: pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py && pytest -v -s v1/structured_output && pytest -v -s v1/test_serial_utils.py && - pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py && + pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py --ignore=v1/spec_decode/test_speculators_correctness.py && pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py --ignore=v1/kv_connector/unit/test_offloading_connector.py' - label: "XPU server test" depends_on: diff --git a/docker/Dockerfile.xpu b/docker/Dockerfile.xpu index 529388f0c68..ed8a347005c 100644 --- a/docker/Dockerfile.xpu +++ b/docker/Dockerfile.xpu @@ -75,13 +75,13 @@ RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRO # Install UMD RUN mkdir neo && \ cd neo && \ - wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-core-2_2.24.8+20344_amd64.deb && \ - wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-opencl-2_2.24.8+20344_amd64.deb && \ - wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-ocloc_25.48.36300.8-0_amd64.deb && \ - wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-opencl-icd_25.48.36300.8-0_amd64.deb && \ - wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libigdgmm12_22.8.2_amd64.deb && \ - wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libze-intel-gpu1_25.48.36300.8-0_amd64.deb && \ - wget https://github.com/oneapi-src/level-zero/releases/download/v1.26.0/level-zero_1.26.0+u24.04_amd64.deb && \ + wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.34.4/intel-igc-core-2_2.34.4+21428_amd64.deb && \ + wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.34.4/intel-igc-opencl-2_2.34.4+21428_amd64.deb && \ + wget https://github.com/intel/compute-runtime/releases/download/26.18.38308.1/intel-ocloc_26.18.38308.1-0_amd64.deb && \ + wget https://github.com/intel/compute-runtime/releases/download/26.18.38308.1/intel-opencl-icd_26.18.38308.1-0_amd64.deb && \ + wget https://github.com/intel/compute-runtime/releases/download/26.18.38308.1/libigdgmm12_22.10.0_amd64.deb && \ + wget https://github.com/intel/compute-runtime/releases/download/26.18.38308.1/libze-intel-gpu1_26.18.38308.1-0_amd64.deb && \ + wget https://github.com/oneapi-src/level-zero/releases/download/v1.28.2/level-zero_1.28.2+u24.04_amd64.deb && \ dpkg -i *.deb && \ cd .. && \ rm -rf neo diff --git a/docs/getting_started/installation/gpu.xpu.inc.md b/docs/getting_started/installation/gpu.xpu.inc.md index f22f5159473..8564f2a7265 100644 --- a/docs/getting_started/installation/gpu.xpu.inc.md +++ b/docs/getting_started/installation/gpu.xpu.inc.md @@ -27,6 +27,7 @@ Currently, there are no pre-built XPU wheels. - First, install required [driver](https://dgpu-docs.intel.com/driver/installation.html#installing-gpu-drivers). - Second, install Python packages for vLLM XPU backend building (Intel OneAPI dependencies are installed automatically as part of `torch-xpu`, see [PyTorch XPU get started](https://docs.pytorch.org/docs/stable/notes/get_start_xpu.html)): +- Start from vllm-xpu-kernels v0.1.10, we recommend user upgrade driver to [compute runtime 26.18](https://github.com/intel/compute-runtime/releases/tag/26.14.37833.4) release, to avoid potential compatibility issue. ```bash git clone https://github.com/vllm-project/vllm.git diff --git a/requirements/xpu.txt b/requirements/xpu.txt index f17e2281f7a..a24ac9ae534 100644 --- a/requirements/xpu.txt +++ b/requirements/xpu.txt @@ -17,4 +17,4 @@ torchaudio torchvision auto_round_lib>=0.13.3 -vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.9.1/vllm_xpu_kernels-0.1.9.1-cp38-abi3-manylinux_2_28_x86_64.whl +vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.10/vllm_xpu_kernels-0.1.10-cp38-abi3-manylinux_2_28_x86_64.whl From 69bdd345428408a2fdf745e225c87defbc2c07d0 Mon Sep 17 00:00:00 2001 From: Muhammad Fawaz <135441198+professorsab@users.noreply.github.com> Date: Fri, 19 Jun 2026 16:11:11 +0500 Subject: [PATCH 30/75] [Bugfix] Fall back to Pydantic loc for param in validation errors (#46038) Signed-off-by: professorsab <135441198+professorsab@users.noreply.github.com> Co-authored-by: Mahad Durrani <114791389+mahadrehmann@users.noreply.github.com> --- .../serve/utils/test_server_utils.py | 65 +++++++++++++++++++ vllm/entrypoints/serve/utils/server_utils.py | 6 ++ 2 files changed, 71 insertions(+) create mode 100644 tests/entrypoints/serve/utils/test_server_utils.py diff --git a/tests/entrypoints/serve/utils/test_server_utils.py b/tests/entrypoints/serve/utils/test_server_utils.py new file mode 100644 index 00000000000..91896986137 --- /dev/null +++ b/tests/entrypoints/serve/utils/test_server_utils.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests that validation_exception_handler populates the `param` field +in its error response using the Pydantic error's `loc`, even when no +custom VLLMValidationError context is present. + +Previously, `param` was only populated for errors carrying a custom +VLLMValidationError in their Pydantic `ctx`. Plain validation failures +(missing fields, wrong types) left `param` as None, even though the +field name was readily available from `error['loc']`. +""" + +import json +from types import SimpleNamespace + +import pytest +from fastapi.exceptions import RequestValidationError + +from vllm.entrypoints.serve.utils.server_utils import validation_exception_handler + + +def _fake_request(log_error_stack: bool = False) -> SimpleNamespace: + """Minimal stand-in for a FastAPI Request - just enough for the + handler to read req.app.state.args.log_error_stack.""" + return SimpleNamespace( + app=SimpleNamespace( + state=SimpleNamespace(args=SimpleNamespace(log_error_stack=log_error_stack)) + ), + state=SimpleNamespace(), # no request_metadata -> hasattr(...) is False + ) + + +class TestValidationErrorParamFallback: + """Ensure `param` falls back to the Pydantic error's `loc` when no + custom VLLMValidationError context is present.""" + + @pytest.mark.parametrize( + ("error_type", "msg"), + [ + ("missing", "Field required"), + ("list_type", "Input should be a valid list"), + ], + ids=["missing-field", "wrong-type"], + ) + @pytest.mark.asyncio + async def test_param_falls_back_to_loc(self, error_type: str, msg: str): + errors = [{"type": error_type, "loc": ("body", "messages"), "msg": msg}] + exc = RequestValidationError(errors) + + response = await validation_exception_handler(_fake_request(), exc) + body = json.loads(response.body) + + assert body["error"]["param"] == "body.messages" + + @pytest.mark.asyncio + async def test_param_fallback_does_not_crash_on_non_dict_error(self): + """Schemathesis fuzzing found that errors[0] isn't always a dict. + The fallback must not crash in that case - it should just leave + param as None instead of raising.""" + exc = RequestValidationError(["some unexpected non-dict error"]) + + response = await validation_exception_handler(_fake_request(), exc) + body = json.loads(response.body) + + assert body["error"]["param"] is None diff --git a/vllm/entrypoints/serve/utils/server_utils.py b/vllm/entrypoints/serve/utils/server_utils.py index d24d492b61e..93a8ef757d4 100644 --- a/vllm/entrypoints/serve/utils/server_utils.py +++ b/vllm/entrypoints/serve/utils/server_utils.py @@ -427,6 +427,12 @@ async def validation_exception_handler(req: Request, exc: RequestValidationError param = ctx_error.parameter break + if param is None and errors: + first_error = errors[0] + loc = first_error.get("loc") if isinstance(first_error, dict) else None + if loc: + param = ".".join(str(part) for part in loc) + exc_str = str(exc) errors_str = str(errors) From b9a7cd464c9ae9b1b450f8982b76d7be4de73724 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Fri, 19 Jun 2026 09:57:26 -0400 Subject: [PATCH 31/75] [12/n] final _C library kernel migration (#45415) --- CMakeLists.txt | 128 ++++++++---------- cmake/external_projects/qutlass.cmake | 34 ++++- csrc/{ => libtorch_stable}/core/math.hpp | 0 .../moe/moe_align_sum_kernels.cu | 2 +- csrc/libtorch_stable/ops.h | 28 ++++ .../quantization/activation_kernels.cu | 118 ++++++++-------- .../fp4/nvfp4_scaled_mm_kernels.cu | 2 +- .../fp4/nvfp4_scaled_mm_sm120_kernels.cu | 2 +- .../w8a8/cutlass/c3x/cutlass_gemm_caller.cuh | 2 +- .../w8a8/cutlass/c3x/scaled_mm.cuh | 2 +- .../w8a8/cutlass/scaled_mm_c2x.cuh | 2 +- csrc/libtorch_stable/torch_bindings.cpp | 27 ++++ csrc/ops.h | 32 ----- csrc/qutlass_registration.cpp | 5 + csrc/torch_bindings.cpp | 40 ------ setup.py | 5 +- vllm/platforms/cuda.py | 22 ++- 17 files changed, 239 insertions(+), 212 deletions(-) rename csrc/{ => libtorch_stable}/core/math.hpp (100%) rename csrc/{ => libtorch_stable}/quantization/activation_kernels.cu (87%) create mode 100644 csrc/qutlass_registration.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index a2651ab344c..e95fe38d329 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -319,82 +319,35 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() # -# _C extension +# Legacy _C extension (ROCm only — CUDA ops migrated to _C_stable_libtorch) # -set(VLLM_EXT_SRC - "csrc/quantization/activation_kernels.cu" - "csrc/torch_bindings.cpp") - -if(VLLM_GPU_LANG STREQUAL "CUDA") - SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library") - - # Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building. - set(CUTLASS_REVISION "v4.4.2") - - # Use the specified CUTLASS source directory for compilation if VLLM_CUTLASS_SRC_DIR is provided - if (DEFINED ENV{VLLM_CUTLASS_SRC_DIR}) - set(VLLM_CUTLASS_SRC_DIR $ENV{VLLM_CUTLASS_SRC_DIR}) - endif() - - if(VLLM_CUTLASS_SRC_DIR) - if(NOT IS_ABSOLUTE VLLM_CUTLASS_SRC_DIR) - get_filename_component(VLLM_CUTLASS_SRC_DIR "${VLLM_CUTLASS_SRC_DIR}" ABSOLUTE) - endif() - message(STATUS "The VLLM_CUTLASS_SRC_DIR is set, using ${VLLM_CUTLASS_SRC_DIR} for compilation") - FetchContent_Declare(cutlass SOURCE_DIR ${VLLM_CUTLASS_SRC_DIR}) - else() - FetchContent_Declare( - cutlass - GIT_REPOSITORY https://github.com/nvidia/cutlass.git - # Please keep this in sync with CUTLASS_REVISION line above. - GIT_TAG ${CUTLASS_REVISION} - GIT_PROGRESS TRUE - - # Speed up CUTLASS download by retrieving only the specified GIT_TAG instead of the history. - # Important: If GIT_SHALLOW is enabled then GIT_TAG works only with branch names and tags. - # So if the GIT_TAG above is updated to a commit hash, GIT_SHALLOW must be set to FALSE - GIT_SHALLOW TRUE - ) - endif() - FetchContent_MakeAvailable(cutlass) - - set_gencode_flags_for_srcs( - SRCS "${VLLM_EXT_SRC}" - CUDA_ARCHS "${CUDA_ARCHS}") - -# if CUDA endif -endif() - -if (VLLM_GPU_LANG STREQUAL "HIP") - # Add QuickReduce kernels (ROCm-only; not part of stable ABI migration). - # TODO: Remove the cuda_view when ROCm upgrade to torch 2.11. - list(APPEND VLLM_EXT_SRC +if(VLLM_GPU_LANG STREQUAL "HIP") + set(VLLM_EXT_SRC + "csrc/torch_bindings.cpp" "csrc/custom_quickreduce.cu" "csrc/cuda_view.cu" - "csrc/libtorch_stable/cuda_utils_kernels.cu" - ) -# if ROCM endif -endif() + "csrc/libtorch_stable/cuda_utils_kernels.cu") -message(STATUS "Enabling C extension.") -define_extension_target( - _C - DESTINATION vllm - LANGUAGE ${VLLM_GPU_LANG} - SOURCES ${VLLM_EXT_SRC} - COMPILE_FLAGS ${VLLM_GPU_FLAGS} - ARCHITECTURES ${VLLM_GPU_ARCHES} - INCLUDE_DIRECTORIES ${CUTLASS_INCLUDE_DIR} - INCLUDE_DIRECTORIES ${CUTLASS_TOOLS_UTIL_INCLUDE_DIR} - USE_SABI 3 - WITH_SOABI) + message(STATUS "Enabling C extension.") + define_extension_target( + _C + DESTINATION vllm + LANGUAGE ${VLLM_GPU_LANG} + SOURCES ${VLLM_EXT_SRC} + COMPILE_FLAGS ${VLLM_GPU_FLAGS} + ARCHITECTURES ${VLLM_GPU_ARCHES} + INCLUDE_DIRECTORIES ${CUTLASS_INCLUDE_DIR} + INCLUDE_DIRECTORIES ${CUTLASS_TOOLS_UTIL_INCLUDE_DIR} + USE_SABI 3 + WITH_SOABI) -# If CUTLASS is compiled on NVCC >= 12.5, it by default uses -# cudaGetDriverEntryPointByVersion as a wrapper to avoid directly calling the -# driver API. This causes problems when linking with earlier versions of CUDA. -# Setting this variable sidesteps the issue by calling the driver directly. -target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1) + # If CUTLASS is compiled on NVCC >= 12.5, it by default uses + # cudaGetDriverEntryPointByVersion as a wrapper to avoid directly calling the + # driver API. This causes problems when linking with earlier versions of CUDA. + # Setting this variable sidesteps the issue by calling the driver directly. + target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1) +endif() # _C HIP endif if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # @@ -403,6 +356,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") set(VLLM_STABLE_EXT_SRC "csrc/libtorch_stable/torch_bindings.cpp" "csrc/libtorch_stable/activation_kernels.cu" + "csrc/libtorch_stable/quantization/activation_kernels.cu" "csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu" "csrc/libtorch_stable/quantization/w8a8/fp8/common.cu" "csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu" @@ -429,6 +383,38 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu") if(VLLM_GPU_LANG STREQUAL "CUDA") + SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library") + + # Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building. + set(CUTLASS_REVISION "v4.4.2") + + # Use the specified CUTLASS source directory for compilation if VLLM_CUTLASS_SRC_DIR is provided + if (DEFINED ENV{VLLM_CUTLASS_SRC_DIR}) + set(VLLM_CUTLASS_SRC_DIR $ENV{VLLM_CUTLASS_SRC_DIR}) + endif() + + if(VLLM_CUTLASS_SRC_DIR) + if(NOT IS_ABSOLUTE VLLM_CUTLASS_SRC_DIR) + get_filename_component(VLLM_CUTLASS_SRC_DIR "${VLLM_CUTLASS_SRC_DIR}" ABSOLUTE) + endif() + message(STATUS "The VLLM_CUTLASS_SRC_DIR is set, using ${VLLM_CUTLASS_SRC_DIR} for compilation") + FetchContent_Declare(cutlass SOURCE_DIR ${VLLM_CUTLASS_SRC_DIR}) + else() + FetchContent_Declare( + cutlass + GIT_REPOSITORY https://github.com/nvidia/cutlass.git + # Please keep this in sync with CUTLASS_REVISION line above. + GIT_TAG ${CUTLASS_REVISION} + GIT_PROGRESS TRUE + + # Speed up CUTLASS download by retrieving only the specified GIT_TAG instead of the history. + # Important: If GIT_SHALLOW is enabled then GIT_TAG works only with branch names and tags. + # So if the GIT_TAG above is updated to a commit hash, GIT_SHALLOW must be set to FALSE + GIT_SHALLOW TRUE + ) + endif() + FetchContent_MakeAvailable(cutlass) + list(APPEND VLLM_STABLE_EXT_SRC "csrc/libtorch_stable/cuda_view.cu" "csrc/libtorch_stable/cuda_utils_kernels.cu" @@ -929,7 +915,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") SRCS "${FP4_SM120_SRCS}" CUDA_ARCHS "${FP4_SM120_ARCHS}") list(APPEND VLLM_STABLE_EXT_SRC "${FP4_SM120_SRCS}") - target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM120=1) list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1") list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1") message(STATUS "Building SM12x NVFP4 for archs: ${FP4_SM120_ARCHS}") @@ -962,7 +947,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") SRCS "${FP4_SM100_SRCS}" CUDA_ARCHS "${FP4_SM100_ARCHS}") list(APPEND VLLM_STABLE_EXT_SRC "${FP4_SM100_SRCS}") - target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM100=1) list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1") list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1") message(STATUS "Building SM10x/11x NVFP4/MXFP4 for archs: ${FP4_SM100_ARCHS}") diff --git a/cmake/external_projects/qutlass.cmake b/cmake/external_projects/qutlass.cmake index 66c001919b0..b653bbfce7b 100644 --- a/cmake/external_projects/qutlass.cmake +++ b/cmake/external_projects/qutlass.cmake @@ -60,6 +60,7 @@ endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS) set(QUTLASS_SOURCES + csrc/qutlass_registration.cpp ${qutlass_SOURCE_DIR}/qutlass/csrc/bindings.cpp ${qutlass_SOURCE_DIR}/qutlass/csrc/gemm.cu ${qutlass_SOURCE_DIR}/qutlass/csrc/gemm_ada.cu @@ -78,8 +79,19 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS) if(CUTLASS_INCLUDE_DIR AND EXISTS "${CUTLASS_INCLUDE_DIR}/cutlass/cutlass.h") list(APPEND QUTLASS_INCLUDES "${CUTLASS_INCLUDE_DIR}") + if(CUTLASS_TOOLS_UTIL_INCLUDE_DIR AND + EXISTS "${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}/cutlass/util/packed_stride.hpp") + list(APPEND QUTLASS_INCLUDES "${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}") + else() + get_filename_component(_qutlass_cutlass_root "${CUTLASS_INCLUDE_DIR}" DIRECTORY) + if(EXISTS "${_qutlass_cutlass_root}/tools/util/include/cutlass/util/packed_stride.hpp") + list(APPEND QUTLASS_INCLUDES "${_qutlass_cutlass_root}/tools/util/include") + endif() + endif() elseif(EXISTS "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include/cutlass/cutlass.h") - list(APPEND QUTLASS_INCLUDES "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include") + list(APPEND QUTLASS_INCLUDES + "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include" + "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/tools/util/include") message(STATUS "[QUTLASS] Using QuTLASS vendored CUTLASS headers (no vLLM CUTLASS detected).") else() message(FATAL_ERROR "[QUTLASS] CUTLASS headers not found. " @@ -91,12 +103,23 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS) CUDA_ARCHS "${QUTLASS_ARCHS}" ) - target_sources(_C PRIVATE ${QUTLASS_SOURCES}) - target_include_directories(_C PRIVATE ${QUTLASS_INCLUDES}) - target_compile_definitions(_C PRIVATE + # QuTLASS uses legacy ATen headers and cannot be built with TORCH_TARGET_VERSION. + # Keep it as its own extension (registers torch.ops._qutlass_C). + define_extension_target( + _qutlass_C + DESTINATION vllm + LANGUAGE ${VLLM_GPU_LANG} + SOURCES ${QUTLASS_SOURCES} + COMPILE_FLAGS ${VLLM_GPU_FLAGS} + ARCHITECTURES ${VLLM_GPU_ARCHES} + INCLUDE_DIRECTORIES ${QUTLASS_INCLUDES} + USE_SABI 3 + WITH_SOABI) + + target_compile_definitions(_qutlass_C PRIVATE QUTLASS_DISABLE_PYBIND=1 TARGET_CUDA_ARCH=${QUTLASS_TARGET_CC} - ) + CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1) set_property(SOURCE ${QUTLASS_SOURCES} APPEND PROPERTY COMPILE_OPTIONS $<$:--expt-relaxed-constexpr --use_fast_math -O3> @@ -111,4 +134,5 @@ else() "[QUTLASS] Skipping build: no supported arch (12.0f / 10.0f) found in " "CUDA_ARCHS='${CUDA_ARCHS}'.") endif() + add_custom_target(_qutlass_C) endif() diff --git a/csrc/core/math.hpp b/csrc/libtorch_stable/core/math.hpp similarity index 100% rename from csrc/core/math.hpp rename to csrc/libtorch_stable/core/math.hpp diff --git a/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu b/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu index d7c68ff25a6..1e842381349 100644 --- a/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu +++ b/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu @@ -9,7 +9,7 @@ #include #include "../../cuda_compat.h" -#include "core/math.hpp" +#include "libtorch_stable/core/math.hpp" #include "libtorch_stable/dispatch_utils.h" #include "libtorch_stable/torch_utils.h" diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 9efc12e9f49..1cc8e8167a6 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -2,9 +2,25 @@ #include #include +#include #include #include +#include + +#include + +inline torch::stable::Tensor weak_ref_tensor(torch::stable::Tensor& tensor) { + // Ensure tensor is on CUDA + STD_TORCH_CHECK(tensor.device().is_cuda(), "Tensor must be on CUDA device"); + + // Get the raw data pointer + void* data_ptr = tensor.mutable_data_ptr(); + + /// Create a new tensor from the raw data pointer + return torch::stable::from_blob(data_ptr, tensor.sizes(), tensor.strides(), + tensor.device(), tensor.scalar_type()); +} void per_token_group_quant_fp8(const torch::stable::Tensor& input, torch::stable::Tensor& output_q, @@ -371,6 +387,18 @@ void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input); void silu_and_mul_clamp(torch::stable::Tensor& out, torch::stable::Tensor& input, double limit, double alpha = 1.0, double beta = 0.0); + +void silu_and_mul_quant(torch::stable::Tensor& out, + torch::stable::Tensor& input, + torch::stable::Tensor& scale); + +void persistent_masked_m_silu_mul_quant( + const torch::stable::Tensor& input, // (E, T, 2*H) + const torch::stable::Tensor& tokens_per_expert, // (E) + torch::stable::Tensor& y_q, // (E, T, H) [OUT] + torch::stable::Tensor& y_s, // (E, T, H//group_size) [OUT] + bool use_ue8m0); + void mul_and_silu(torch::stable::Tensor& out, torch::stable::Tensor& input); void gelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input); void gelu_tanh_and_mul(torch::stable::Tensor& out, diff --git a/csrc/quantization/activation_kernels.cu b/csrc/libtorch_stable/quantization/activation_kernels.cu similarity index 87% rename from csrc/quantization/activation_kernels.cu rename to csrc/libtorch_stable/quantization/activation_kernels.cu index 8cc645c33e2..822a41969e7 100644 --- a/csrc/quantization/activation_kernels.cu +++ b/csrc/libtorch_stable/quantization/activation_kernels.cu @@ -1,16 +1,12 @@ -#include -#include -#include +#include "libtorch_stable/torch_utils.h" #include -#include "core/math.hpp" -#include "../cuda_compat.h" -#include "dispatch_utils.h" +#include "libtorch_stable/core/math.hpp" +#include "cuda_compat.h" +#include "libtorch_stable/dispatch_utils.h" #include "quantization/w8a8/fp8/common.cuh" -#include - #ifndef USE_ROCM #include #include @@ -33,7 +29,6 @@ typedef __hip_fp8x4_e4m3_fnuz __nv_fp8x4_e4m3; #endif #endif -#include "core/registration.h" namespace vllm { template @@ -564,41 +559,47 @@ __global__ void silu_mul_fp8_quant_deep_gemm_kernel( } // namespace vllm // Launch activation, gating, and quantize kernel. -#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL) \ - int d = input.size(-1) / 2; \ - int64_t num_tokens = input.numel() / input.size(-1); \ - dim3 grid(num_tokens, num_tokens > 16 ? num_tokens > 32 ? 1 : 2 : 4); \ - dim3 block(std::min(d, 512)); \ - const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \ - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); \ - VLLM_DISPATCH_FLOATING_TYPES( \ - input.scalar_type(), "act_and_mul_kernel", [&] { \ - VLLM_DISPATCH_FP8_TYPES( \ - out.scalar_type(), "fused_add_rms_norm_kernel_fp8_type", [&] { \ - vllm::act_and_mul_quant_kernel, \ - fp8_t> \ - <<>>(out.data_ptr(), \ - input.data_ptr(), \ - scale.data_ptr(), d); \ - }); \ +#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL) \ + int d = input.size(-1) / 2; \ + int64_t num_tokens = input.numel() / input.size(-1); \ + dim3 grid(num_tokens, num_tokens > 16 ? num_tokens > 32 ? 1 : 2 : 4); \ + dim3 block(std::min(d, 512)); \ + const torch::stable::accelerator::DeviceGuard device_guard( \ + input.get_device_index()); \ + const cudaStream_t stream = \ + get_current_cuda_stream(input.get_device_index()); \ + VLLM_STABLE_DISPATCH_FLOATING_TYPES( \ + input.scalar_type(), "act_and_mul_kernel", [&] { \ + VLLM_STABLE_DISPATCH_FP8_TYPES( \ + out.scalar_type(), "act_and_mul_quant_kernel_fp8_type", [&] { \ + vllm::act_and_mul_quant_kernel, \ + fp8_t> \ + <<>>( \ + out.mutable_data_ptr(), \ + input.const_data_ptr(), \ + scale.const_data_ptr(), d); \ + }); \ }); -void silu_and_mul_quant(torch::Tensor& out, // [..., d] - torch::Tensor& input, // [..., 2 * d] - torch::Tensor& scale) { - TORCH_CHECK(out.dtype() == torch::kFloat8_e4m3fn || - out.dtype() == torch::kFloat8_e4m3fnuz); - TORCH_CHECK(input.dtype() == torch::kFloat16 || - input.dtype() == torch::kBFloat16); - TORCH_CHECK(input.size(-1) % 2 == 0); +void silu_and_mul_quant(torch::stable::Tensor& out, // [..., d] + torch::stable::Tensor& input, // [..., 2 * d] + torch::stable::Tensor& scale) { + STD_TORCH_CHECK( + out.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn || + out.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fnuz); + STD_TORCH_CHECK( + input.scalar_type() == torch::headeronly::ScalarType::Half || + input.scalar_type() == torch::headeronly::ScalarType::BFloat16, + "Input must be FP16 or BF16"); + STD_TORCH_CHECK(input.size(-1) % 2 == 0); LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel); } void persistent_masked_m_silu_mul_quant( - const at::Tensor& input, // (E, T, 2*H) - const at::Tensor& tokens_per_expert, // (E) - at::Tensor& y_q, // (E, T, H) [OUT] - at::Tensor& y_s, // (E, T, H//group_size) [OUT] + const torch::stable::Tensor& input, // (E, T, 2*H) + const torch::stable::Tensor& tokens_per_expert, // (E) + torch::stable::Tensor& y_q, // (E, T, H) [OUT] + torch::stable::Tensor& y_s, // (E, T, H//group_size) [OUT] bool cast_scale_ue8m0) { #ifndef USE_ROCM @@ -606,14 +607,18 @@ void persistent_masked_m_silu_mul_quant( // fixed GROUP_SIZE of 128. static constexpr int GROUP_SIZE = 128; - TORCH_CHECK(input.dtype() == torch::kBFloat16); - TORCH_CHECK(y_q.dtype() == torch::kFloat8_e4m3fn || - y_q.dtype() == torch::kFloat8_e4m3fnuz); - TORCH_CHECK(input.size(-1) % (GROUP_SIZE * 2) == 0); + STD_TORCH_CHECK(input.scalar_type() == + torch::headeronly::ScalarType::BFloat16); + STD_TORCH_CHECK( + y_q.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn || + y_q.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fnuz); + STD_TORCH_CHECK(input.size(-1) % (GROUP_SIZE * 2) == 0); bool const is_packed_ue8m0 = - (y_s.dtype() == torch::kInt32 && cast_scale_ue8m0); - TORCH_CHECK(y_s.dtype() == torch::kFloat32 || is_packed_ue8m0); + (y_s.scalar_type() == torch::headeronly::ScalarType::Int && + cast_scale_ue8m0); + STD_TORCH_CHECK(y_s.scalar_type() == torch::headeronly::ScalarType::Float || + is_packed_ue8m0); using Idx_t = int64_t; @@ -631,7 +636,7 @@ void persistent_masked_m_silu_mul_quant( int const NUM_GROUPS = H / GROUP_SIZE; - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const cudaStream_t stream = get_current_cuda_stream(input.get_device_index()); // TODO: Get this from cuda_arch ? static constexpr int SILU_V2_BLOCK_COUNT = 132 * 32; @@ -643,18 +648,21 @@ void persistent_masked_m_silu_mul_quant( static constexpr int max_shared_mem_bytes = \ GROUP_SIZE * 2 * STAGES * NUM_WARPS * 2; \ dim3 grid(sms), block(THREAD_COUNT); \ - const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \ - VLLM_DISPATCH_FP8_TYPES( \ + const torch::stable::accelerator::DeviceGuard device_guard( \ + input.get_device_index()); \ + VLLM_STABLE_DISPATCH_FP8_TYPES( \ y_q.scalar_type(), "silu_mul_fp8_quant_deep_gemm_kernel", [&] { \ vllm::silu_mul_fp8_quant_deep_gemm_kernel< \ BLOCK_COUNT, max_shared_mem_bytes, fp8_t, scale_t, THREAD_COUNT, \ Idx_t, CEIL_UE8M0, GROUP_SIZE, STAGES> \ <<>>( \ - reinterpret_cast<__nv_bfloat16*>(input.data_ptr()), \ - (fp8_t*)y_q.data_ptr(), \ - reinterpret_cast(y_s.data_ptr()), \ - reinterpret_cast(tokens_per_expert.data_ptr()), E, \ - T, H, stride_i_e, stride_i_t, stride_i_h, stride_yq_e, \ + reinterpret_cast( \ + input.const_data_ptr()), \ + y_q.mutable_data_ptr(), \ + reinterpret_cast(y_s.mutable_data_ptr()), \ + reinterpret_cast( \ + tokens_per_expert.const_data_ptr()), \ + E, T, H, stride_i_e, stride_i_t, stride_i_h, stride_yq_e, \ stride_yq_t, stride_yq_h, STRIDE_YS_E, STRIDE_YS_T, \ STRIDE_YS_G, STRIDE_YS_P, stride_counts_e); \ }); @@ -679,7 +687,7 @@ void persistent_masked_m_silu_mul_quant( Idx_t stride_ys_g = y_s.stride(2); Idx_t stride_ys_p = 0; if (!cast_scale_ue8m0) { - TORCH_CHECK(!is_packed_ue8m0); + STD_TORCH_CHECK(!is_packed_ue8m0); LAUNCH_ON_H(float, stride_ys_e, stride_ys_t, stride_ys_g, stride_ys_p, false); return; @@ -692,8 +700,8 @@ void persistent_masked_m_silu_mul_quant( return; } - TORCH_CHECK(cast_scale_ue8m0 && is_packed_ue8m0); - TORCH_CHECK(y_s.dtype() == torch::kInt32); + STD_TORCH_CHECK(cast_scale_ue8m0 && is_packed_ue8m0); + STD_TORCH_CHECK(y_s.scalar_type() == torch::headeronly::ScalarType::Int); // Int32 packed ue8m0 scales tensor. // Let E, T, G be the number to experts, number of tokens and number of groups diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu b/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu index 86355bf7060..af9f24a70e0 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu +++ b/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu @@ -31,7 +31,7 @@ #include "cutlass/util/packed_stride.hpp" -#include "core/math.hpp" +#include "libtorch_stable/core/math.hpp" #include "core/batch_invariant.hpp" using namespace cute; diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu b/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu index 7adba6308fa..3a45ede8dfd 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu +++ b/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu @@ -31,7 +31,7 @@ #include "cutlass/util/packed_stride.hpp" -#include "core/math.hpp" +#include "libtorch_stable/core/math.hpp" #include "core/batch_invariant.hpp" using namespace cute; diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh index 1eed7579924..1d9023484fa 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh +++ b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh @@ -19,7 +19,7 @@ #include "cutlass/gemm/collective/collective_builder.hpp" #include "cutlass/util/packed_stride.hpp" -#include "core/math.hpp" +#include "libtorch_stable/core/math.hpp" #include "libtorch_stable/cutlass_extensions/common.hpp" // clang-format on diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm.cuh b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm.cuh index 4cb591be056..7b7d4d71473 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm.cuh +++ b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm.cuh @@ -14,7 +14,7 @@ #include "cutlass/epilogue/collective/collective_builder.hpp" #include "cutlass/gemm/collective/collective_builder.hpp" -#include "core/math.hpp" +#include "libtorch_stable/core/math.hpp" #include "libtorch_stable/cutlass_extensions/common.hpp" // clang-format on diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cuh b/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cuh index 7846e609fe7..d2b54cb911b 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cuh +++ b/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cuh @@ -22,7 +22,7 @@ #include "cutlass/epilogue/threadblock/fusion/visitors.hpp" #include "cutlass/gemm/kernel/default_gemm_universal_with_visitor.h" -#include "core/math.hpp" +#include "libtorch_stable/core/math.hpp" #include "libtorch_stable/cutlass_extensions/common.hpp" // clang-format on diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index c1d2d26fcd8..d55c12d382a 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -34,6 +34,20 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { // TODO: Remove this once ROCm upgrade to torch 2.11. ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor"); + // Note about marlin kernel 'workspace' arguments: + // Technically these should be mutable since they are modified by the kernel. + // But since they are set back to zero once the kernel is finished we can + // hand wave and say that they have no net effect. + // + // The reason to mark 'workspace' as immutable is so that they don't interfere + // with using ScalarType arguments in the ops. If they are marked as mutable, + // pytorch throws an assert in + // 'torch._higher_order_ops._register_effectful_op' that prevents these + // kernels from being torch.compile'd. + // See the following document for more info on custom types and ops that use + // custom types: + // https://docs.google.com/document/d/18fBMPuOJ0fY5ZQ6YyrHUppw9FA332CpNtgB6SOIgyuA + // Machete (Dense) Optimized Mixed Precision GEMM for Hopper. ops.def( "machete_supported_schedules(" @@ -480,6 +494,11 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "Tensor workspace, int k, int max_seq_len) -> ()"); // Activation ops + ops.def( + "persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! " + "y_q, Tensor! y_s, bool use_ue8m0) -> ()"); + ops.def("weak_ref_tensor(Tensor input) -> Tensor"); + // Activation function used in SwiGLU. ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()"); @@ -492,6 +511,10 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit, " "float alpha=1.0, float beta=0.0) -> ()"); + // SwiGLU activation with FP8 quantization. + ops.def( + "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()"); + // Activation function used in GeGLU with `none` approximation. ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()"); @@ -690,6 +713,10 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { ops.impl("persistent_topk", TORCH_BOX(&persistent_topk)); // Activation kernels (shared CUDA/ROCm) + ops.impl("persistent_masked_m_silu_mul_quant", + TORCH_BOX(&persistent_masked_m_silu_mul_quant)); + ops.impl("weak_ref_tensor", TORCH_BOX(&weak_ref_tensor)); + ops.impl("silu_and_mul_quant", TORCH_BOX(&silu_and_mul_quant)); ops.impl("silu_and_mul", TORCH_BOX(&silu_and_mul)); ops.impl("mul_and_silu", TORCH_BOX(&mul_and_silu)); ops.impl("gelu_and_mul", TORCH_BOX(&gelu_and_mul)); diff --git a/csrc/ops.h b/csrc/ops.h index ec3f5e187cc..398ae1016f3 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -9,28 +9,6 @@ #include -torch::Tensor weak_ref_tensor(torch::Tensor& tensor) { - // Ensure tensor is on CUDA - if (!tensor.is_cuda()) { - throw std::runtime_error("Tensor must be on CUDA device"); - } - - // Get the raw data pointer - void* data_ptr = tensor.data_ptr(); - - // Get tensor sizes and strides - std::vector sizes = tensor.sizes().vec(); - std::vector strides = tensor.strides().vec(); - - // Get tensor options (dtype, device) - auto options = tensor.options(); - - // Create a new tensor from the raw data pointer - auto new_tensor = torch::from_blob(data_ptr, sizes, strides, options); - - return new_tensor; -} - // rms_norm and fused_add_rms_norm declarations also exist in // csrc/libtorch_stable/ops.h (torch::stable ABI for CUDA). They remain here // because the CPU build still uses these torch::Tensor declarations. @@ -53,16 +31,6 @@ void silu_and_mul(torch::Tensor& out, torch::Tensor& input); void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit, double alpha = 1.0, double beta = 0.0); -void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input, - torch::Tensor& scale); - -void persistent_masked_m_silu_mul_quant( - const at::Tensor& input, // (E, T, 2*H) - const at::Tensor& counts, // (E) - at::Tensor& y_q, // (E, T, H) [OUT] - at::Tensor& y_s, // (E, T, H//group_size) [OUT] - bool use_ue8m0); - void gelu_and_mul(torch::Tensor& out, torch::Tensor& input); void gelu_tanh_and_mul(torch::Tensor& out, torch::Tensor& input); diff --git a/csrc/qutlass_registration.cpp b/csrc/qutlass_registration.cpp new file mode 100644 index 00000000000..effb4404135 --- /dev/null +++ b/csrc/qutlass_registration.cpp @@ -0,0 +1,5 @@ +#include "core/registration.h" + +// QuTLASS registers torch.ops._qutlass_C via TORCH_LIBRARY in bindings.cpp. +// This stub lets Python import vllm._qutlass_C to trigger op registration. +REGISTER_EXTENSION(_qutlass_C) diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index cfd185394a4..e1430c08d3a 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -20,17 +20,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { // vLLM custom ops - // - - ops.def( - "persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! " - "y_q, Tensor! y_s," - "bool use_ue8m0) -> ()"); - ops.impl("persistent_masked_m_silu_mul_quant", torch::kCUDA, - &persistent_masked_m_silu_mul_quant); - - ops.def("weak_ref_tensor(Tensor input) -> Tensor"); - ops.impl("weak_ref_tensor", torch::kCUDA, &weak_ref_tensor); #ifdef USE_ROCM // TODO: Remove this once we upgrade to torch 2.11. @@ -39,35 +28,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor"); ops.impl("get_cuda_view_from_cpu_tensor", torch::kCPU, &get_cuda_view_from_cpu_tensor); -#endif - - // Activation ops (quantized only — basic ops moved to _C_stable_libtorch) - ops.def( - "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()"); - ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant); - - // Horizontally-fused DeepseekV4-MLA: per-head RMSNorm + GPT-J RoPE for Q, and - // GPT-J RoPE + UE8M0 FP8 quant + paged cache insert for KV, all in one - // kernel launch. Registered in _C_stable_libtorch (incl. the FlashInfer V4 - // full-cache bf16/fp8 variants). - - // Quantization ops -#ifndef USE_ROCM - - // Note about marlin kernel 'workspace' arguments: - // Technically these should be mutable since they are modified by the kernel. - // But since they are set back to zero once the kernel is finished we can - // hand wave and say that they have no net effect. - // - // The reason to mark 'workspace' as immutable is so that they don't interfere - // with using ScalarType arguments in the ops. If they are marked as mutable, - // pytorch throws an assert in - // 'torch._higher_order_ops._register_effectful_op' that prevents these - // kernels from being torch.compile'd. - // See the following document for more info on custom types and ops that use - // custom types: - // https://docs.google.com/document/d/18fBMPuOJ0fY5ZQ6YyrHUppw9FA332CpNtgB6SOIgyuA - #endif } diff --git a/setup.py b/setup.py index 2aaa7dfc49c..b807b2215db 100644 --- a/setup.py +++ b/setup.py @@ -769,6 +769,7 @@ class precompiled_wheel_utils: "vllm/_C.abi3.so", "vllm/_C_stable_libtorch.abi3.so", "vllm/_moe_C_stable_libtorch.abi3.so", + "vllm/_qutlass_C.abi3.so", "vllm/_flashmla_C.abi3.so", "vllm/_flashmla_extension_C.abi3.so", "vllm/_sparse_flashmla_C.abi3.so", @@ -1135,6 +1136,7 @@ if _is_cuda(): # DeepGEMM requires CUDA 12.3+ (SM90/SM100) # Optional since it won't build on unsupported architectures ext_modules.append(CMakeExtension(name="vllm._deep_gemm_C", optional=True)) + ext_modules.append(CMakeExtension(name="vllm._qutlass_C", optional=True)) # fmha_sm100 is a Python/CuTe-DSL package installed into vllm.third_party. ext_modules.append(CMakeExtension(name="vllm.fmha_sm100", optional=True)) @@ -1149,7 +1151,8 @@ if _is_cpu(): ext_modules.append(CMakeExtension(name="vllm._C")) if _build_custom_ops(): - ext_modules.append(CMakeExtension(name="vllm._C")) + if _is_hip(): + ext_modules.append(CMakeExtension(name="vllm._C")) if _is_cuda() or _is_hip(): ext_modules.append(CMakeExtension(name="vllm._C_stable_libtorch")) ext_modules.append(CMakeExtension(name="vllm._moe_C_stable_libtorch")) diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 49181eaec6c..30a16e27469 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -19,7 +19,6 @@ from torch.distributed.distributed_c10d import is_nccl_available from typing_extensions import ParamSpec # import custom ops, trigger op registration -import vllm._C # noqa import vllm._C_stable_libtorch # noqa import vllm.envs as envs from vllm.logger import init_logger @@ -40,6 +39,11 @@ else: logger = init_logger(__name__) +try: + import vllm._qutlass_C # noqa: F401 +except ImportError as e: + logger.warning("Failed to import from vllm._qutlass_C: %r", e) + _P = ParamSpec("_P") _R = TypeVar("_R") @@ -187,6 +191,22 @@ class CudaPlatformBase(Platform): "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES", ] + @classmethod + def import_kernels(cls) -> None: + """Import CUDA kernel extensions (_C_stable_libtorch, optional _qutlass_C).""" + try: + import vllm._C_stable_libtorch # noqa: F401 + except ImportError as e: + logger.warning("Failed to import from vllm._C_stable_libtorch: %r", e) + try: + import vllm._moe_C_stable_libtorch # noqa: F401 + except ImportError as e: + logger.warning("Failed to import from vllm._moe_C_stable_libtorch: %r", e) + try: + import vllm._qutlass_C # noqa: F401 + except ImportError as e: + logger.warning("Failed to import from vllm._qutlass_C: %r", e) + @property def supported_dtypes(self) -> list[torch.dtype]: if self.has_device_capability(80): From 01192139bf022bec84e2cca3a3e36e8bb5293b5c Mon Sep 17 00:00:00 2001 From: Tyler Michael Smith Date: Fri, 19 Jun 2026 12:55:42 -0400 Subject: [PATCH 32/75] [DSv4] Pack KV caches into contiguous per-block allocations for DeepSeek V4 (#44577) Signed-off-by: Tyler Michael Smith Signed-off-by: Matthew Bonanni Signed-off-by: Lucas Wilkinson Co-authored-by: Claude Opus 4.6 (1M context) Co-authored-by: Matthew Bonanni Co-authored-by: Lucas Wilkinson Co-authored-by: Lucas Wilkinson Co-authored-by: OpenAI Codex --- tests/v1/core/test_contiguous_kv_packing.py | 135 ++++++++++++++++++ .../kv_connector/v1/nixl/base_worker.py | 98 +++++++++++++ .../kv_connector/v1/offloading/worker.py | 14 +- vllm/v1/core/kv_cache_utils.py | 13 +- vllm/v1/kv_cache_interface.py | 2 + vllm/v1/worker/gpu/attn_utils.py | 48 ++++++- vllm/v1/worker/gpu_model_runner.py | 56 ++++++-- 7 files changed, 344 insertions(+), 22 deletions(-) create mode 100644 tests/v1/core/test_contiguous_kv_packing.py diff --git a/tests/v1/core/test_contiguous_kv_packing.py b/tests/v1/core/test_contiguous_kv_packing.py new file mode 100644 index 00000000000..79f8937c637 --- /dev/null +++ b/tests/v1/core/test_contiguous_kv_packing.py @@ -0,0 +1,135 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for contiguous KV cache packing in _get_kv_cache_config_deepseek_v4.""" + +from unittest.mock import MagicMock + +import pytest +import torch + +from vllm.v1.core.kv_cache_utils import _get_kv_cache_config_deepseek_v4 +from vllm.v1.kv_cache_interface import ( + KVCacheGroupSpec, + MLAAttentionSpec, + UniformTypeKVCacheSpecs, +) + + +def _make_mla_spec(page_size: int, block_size: int = 256) -> MLAAttentionSpec: + return MLAAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=512, + dtype=torch.uint8, + page_size_padded=page_size, + cache_dtype_str="fp8_ds_mla", + model_version="deepseek_v4", + alignment=576, + ) + + +def _make_groups(n_c4, n_c128, n_swa): + PS_C4_MLA = 37440 + PS_C4_IDX = 8640 + PS_C128 = 1728 + PS_SWA = 37440 + + mla_specs = {} + for i in range(n_c4): + mla_specs[f"c4_mla.{i}"] = _make_mla_spec(PS_C4_MLA) + mla_specs[f"c4_idx.{i}"] = _make_mla_spec(PS_C4_IDX) + for i in range(n_c128): + mla_specs[f"c128_mla.{i}"] = _make_mla_spec(PS_C128) + + mla_group = KVCacheGroupSpec( + layer_names=list(mla_specs.keys()), + kv_cache_spec=UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=mla_specs), + ) + + swa_specs = {} + for i in range(n_swa): + swa_specs[f"swa.{i}"] = _make_mla_spec(PS_SWA) + + swa_group = KVCacheGroupSpec( + layer_names=list(swa_specs.keys()), + kv_cache_spec=UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=swa_specs), + ) + + return [mla_group, swa_group] + + +def _mock_vllm_config(): + config = MagicMock() + config.cache_config.num_gpu_blocks_override = None + return config + + +def _run(n_c4=3, n_c128=2, n_swa=5, mem=100 * 1024 * 1024): + groups = _make_groups(n_c4, n_c128, n_swa) + return _get_kv_cache_config_deepseek_v4(_mock_vllm_config(), groups, mem) + + +def _page_sizes_by_layer( + groups: list[KVCacheGroupSpec], +) -> dict[str, int]: + page_sizes = {} + for group in groups: + specs = group.kv_cache_spec.kv_cache_specs + for layer_name in group.layer_names: + page_sizes[layer_name] = specs[layer_name].page_size_bytes + return page_sizes + + +class TestInterleavedPacking: + def test_all_tensors_have_block_stride(self): + _, tensors = _run() + for t in tensors: + assert t.block_stride > 0 + + def test_all_tensors_share_same_size(self): + _, tensors = _run() + sizes = set(t.size for t in tensors) + assert len(sizes) == 1 + assert sizes.pop() > 0 + + def test_offsets_within_one_block(self): + _, tensors = _run() + for t in tensors: + assert t.offset < t.block_stride + + def test_all_layers_accounted_for(self): + n_c4, n_c128, n_swa = 5, 4, 7 + _, tensors = _run(n_c4=n_c4, n_c128=n_c128, n_swa=n_swa) + all_names = set() + for t in tensors: + all_names.update(t.shared_by) + expected = n_c4 * 2 + n_c128 + n_swa + assert len(all_names) == expected + + def test_strided_views_are_independent(self): + groups = _make_groups(n_c4=3, n_c128=2, n_swa=5) + page_sizes = _page_sizes_by_layer(groups) + num_blocks, tensors = _get_kv_cache_config_deepseek_v4( + _mock_vllm_config(), groups, 100 * 1024 * 1024 + ) + backing = torch.zeros(tensors[0].size, dtype=torch.uint8) + views = [] + for t in tensors: + page_size = page_sizes[t.shared_by[0]] + v = torch.as_strided( + backing, + size=(num_blocks, page_size), + stride=(t.block_stride, 1), + storage_offset=t.offset, + ) + views.append(v) + + for i, v in enumerate(views): + v.fill_(i + 1) + + for i, v in enumerate(views): + assert (v == i + 1).all(), f"View {i} was corrupted" + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 5804732f80f..7ee072ceaf1 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -841,8 +841,106 @@ class NixlBaseConnectorWorker: # Forwarding a real layer name rather than a synthetic key self.register_kv_caches({first_layer: kv_cache}) + def _register_packed_kv_cache( + self, + storage: torch.UntypedStorage, + ) -> None: + """Register a packed KV cache as a single NIXL region. + + The packed allocation interleaves all layers per block, so each + block_stride-byte chunk is one logical block. We register 1 + NIXL region and create 1 descriptor per block. + """ + self.transfer_topo = TransferTopology( + tp_rank=self.tp_rank, + tp_size=self.world_size, + block_size=self.block_size, + engine_id=self.engine_id, + is_mla=self.use_mla, + total_num_kv_heads=self.model_config.get_total_num_kv_heads(), + attn_backends=self.attn_backends, + tensor_shape=None, + is_mamba=self._has_mamba, + ) + self.compat_hash = compute_nixl_compatibility_hash( + self.vllm_config, + self.backend_name, + self.transfer_topo.cross_layers_blocks, + ) + + total_size = storage.nbytes() + block_stride = total_size // self.num_blocks + base_addr = storage.data_ptr() + device_id = storage.device.index + assert device_id is not None + + logger.info( + "Registering packed KV cache: total_size=%s, block_stride=%s, " + "num_blocks=%s, num_regions=1", + total_size, + block_stride, + self.num_blocks, + ) + + self.device_id = device_id + caches_data = [(base_addr, total_size, self.device_id, "")] + + self.block_len_per_layer = [block_stride] + self.num_regions = 1 + self.num_descs = self.num_blocks + self.kv_caches_base_addr[self.engine_id][self.tp_rank] = [base_addr] + + descs = self.nixl_wrapper.get_reg_descs(caches_data, self.nixl_memory_type) + self.nixl_wrapper.register_memory(descs, backends=self.nixl_backends) + self._registered_descs.append(descs) + + self.dst_num_blocks[self.engine_id] = self.num_blocks + + self.src_xfer_handles_by_block_size[self.block_size], (self.src_blocks_data) = ( + self.register_local_xfer_handler(self.block_size) + ) + + agent_metadata = NixlAgentMetadata( + engine_id=self.engine_id, + agent_metadata=self.nixl_wrapper.get_agent_metadata(), + device_id=self.device_id, + kv_caches_base_addr=( + self.kv_caches_base_addr[self.engine_id][self.tp_rank] + ), + num_blocks=self.num_blocks, + block_lens=self.block_len_per_layer, + kv_cache_layout=self.kv_cache_layout, + block_size=self.block_size, + ssm_sizes=self._mamba_ssm_size, + attn_backend_name=self.backend_name, + physical_blocks_per_logical_kv_block=( + self._physical_blocks_per_logical_kv_block + ), + ) + assert self.compat_hash is not None + encoder = msgspec.msgpack.Encoder() + self.xfer_handshake_metadata = NixlHandshakePayload( + compatibility_hash=self.compat_hash, + agent_metadata_bytes=encoder.encode(agent_metadata), + ) + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): """Register the KV Cache data in nixl.""" + + # Detect packed allocation: all tensors are strided views into the + # same backing storage (different data_ptr but same storage). + # This happens with DSv4-style contiguous per-block packing. + if len(kv_caches) > 1 and not self._has_mamba: + storage = next(iter(kv_caches.values())).untyped_storage() + storage_ptrs = { + cache.untyped_storage().data_ptr() for cache in kv_caches.values() + } + data_ptrs = {cache.data_ptr() for cache in kv_caches.values()} + if len(storage_ptrs) == 1 and len(data_ptrs) > 1: + self._register_packed_kv_cache(storage) + self.device_kv_caches = kv_caches + return + self.transfer_topo = TransferTopology( tp_rank=self.tp_rank, tp_size=self.world_size, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py index 744a0c74294..8583bb4b1e0 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py @@ -72,18 +72,22 @@ class OffloadingConnectorWorker: if isinstance(layer_kv_cache_spec, AttentionSpec): layer_kv_cache = kv_caches[layer_name] assert isinstance(layer_kv_cache, torch.Tensor) - assert layer_kv_cache.storage_offset() == 0 - storage = layer_kv_cache.untyped_storage() page = layer_kv_cache_spec.page_size_bytes + elem_size = layer_kv_cache.element_size() + byte_offset = layer_kv_cache.storage_offset() * elem_size + block_stride_bytes = layer_kv_cache.stride(0) * elem_size tensors_per_block[layer_name] = ( torch.tensor( [], dtype=torch.int8, device=layer_kv_cache.device, - ) - .set_(storage) - .view(num_blocks, page), + ).set_( + layer_kv_cache.untyped_storage(), + byte_offset, + (num_blocks, page), + (block_stride_bytes, 1), + ), ) page_size_bytes[layer_name] = layer_kv_cache_spec.page_size_bytes unpadded_page_size_bytes[layer_name] = ( diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 72ca6a2fa67..a1ebe08c078 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -1236,10 +1236,21 @@ def _get_kv_cache_config_deepseek_v4( num_blocks = available_memory // total_num_bytes_per_block num_blocks = may_override_num_blocks(vllm_config, num_blocks) + total_size = total_num_bytes_per_block * num_blocks + kv_cache_tensors: list[KVCacheTensor] = [] + byte_offset = 0 for ps, slots in buckets.items(): for slot in slots: - kv_cache_tensors.append(KVCacheTensor(size=ps * num_blocks, shared_by=slot)) + kv_cache_tensors.append( + KVCacheTensor( + size=total_size, + shared_by=slot, + offset=byte_offset, + block_stride=total_num_bytes_per_block, + ) + ) + byte_offset += ps return num_blocks, kv_cache_tensors diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 9528fb65af1..2e779b2c2a4 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -847,6 +847,8 @@ class KVCacheTensor: size: int # size of the KV cache tensor in bytes shared_by: list[str] # layer names that share the same KV cache tensor + offset: int = 0 # byte offset of this layer within a contiguous block + block_stride: int = 0 # total bytes per block in a packed layout (0 = not packed) @dataclass diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 74158f92bf8..7b85e6fa316 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Iterable, Sequence from dataclasses import dataclass +from math import prod from typing import Any, cast import torch @@ -155,8 +156,17 @@ def _allocate_kv_cache( kv_cache_config: KVCacheConfig, shared_layers: dict[str, str], device: torch.device ): kv_cache_raw_tensors: dict[str, torch.Tensor] = {} + packed_backing: torch.Tensor | None = None for kv_cache_tensor in kv_cache_config.kv_cache_tensors: - tensor = torch.zeros(kv_cache_tensor.size, dtype=torch.int8, device=device) + if kv_cache_tensor.block_stride > 0: + # Allocate once; all packed tensors alias the same backing. + if packed_backing is None: + packed_backing = torch.zeros( + kv_cache_tensor.size, dtype=torch.int8, device=device + ) + tensor = packed_backing + else: + tensor = torch.zeros(kv_cache_tensor.size, dtype=torch.int8, device=device) for layer_name in kv_cache_tensor.shared_by: kv_cache_raw_tensors[layer_name] = tensor @@ -176,10 +186,18 @@ def _reshape_kv_cache( cache_dtype: str, kernel_block_sizes: list[int], shared_kv_cache_layers: dict[str, str], + kv_cache_config: "KVCacheConfig | None" = None, ) -> dict[str, Any]: kv_caches: dict[str, Any] = {} has_attn, has_mamba = False, False + layer_packing: dict[str, tuple[int, int]] = {} + if kv_cache_config is not None: + for kv_tensor in kv_cache_config.kv_cache_tensors: + if kv_tensor.block_stride > 0: + for ln in kv_tensor.shared_by: + layer_packing[ln] = (kv_tensor.offset, kv_tensor.block_stride) + for group in attn_groups: if group.kv_cache_group_id >= len(kernel_block_sizes): continue @@ -198,8 +216,13 @@ def _reshape_kv_cache( continue kv_raw_tensor = kv_cache_raw_tensors[layer_name] - assert kv_raw_tensor.numel() % kv_cache_spec.page_size_bytes == 0 - num_blocks = kv_raw_tensor.numel() // kv_cache_spec.page_size_bytes + packing = layer_packing.get(layer_name) + if packing is not None: + _, blk_stride = packing + num_blocks = kv_raw_tensor.numel() // blk_stride + else: + assert kv_raw_tensor.numel() % kv_cache_spec.page_size_bytes == 0 + num_blocks = kv_raw_tensor.numel() // kv_cache_spec.page_size_bytes if isinstance(kv_cache_spec, AttentionSpec): has_attn = True @@ -232,8 +255,18 @@ def _reshape_kv_cache( ] dtype = kv_cache_spec.dtype - kv_tensor = kv_raw_tensor.view(dtype) - if kv_cache_spec.page_size_padded is not None: + if packing is not None: + offset, block_stride = packing + assert inv_order[0] == 0 + page_bytes = prod(kv_cache_shape[1:]) * get_dtype_size(dtype) + kv_cache = ( + kv_raw_tensor.view(-1, block_stride)[ + :, offset : offset + page_bytes + ] + .view(dtype) + .view(kv_cache_shape) + ) + elif kv_cache_spec.page_size_padded is not None: # Use strided view to handle page_size_bytes that # include padding. This follows the same pattern as # MambaSpec handling in gpu_model_runner.py. @@ -246,13 +279,13 @@ def _reshape_kv_cache( strides = list(torch.empty(kv_cache_shape).stride()) strides[inv_order[0]] = page_stride kv_cache = torch.as_strided( - kv_tensor, + kv_raw_tensor.view(dtype), size=kv_cache_shape, stride=tuple(strides), ) else: # No padding — safe to use a contiguous view. - kv_cache = kv_tensor.view(kv_cache_shape) + kv_cache = kv_raw_tensor.view(dtype).view(kv_cache_shape) kv_caches[layer_name] = kv_cache.permute(*inv_order) elif isinstance(kv_cache_spec, MambaSpec): @@ -365,6 +398,7 @@ def init_kv_cache( kernel_block_sizes=kernel_block_sizes, cache_dtype=cache_dtype, shared_kv_cache_layers=shared_kv_cache_layers, + kv_cache_config=kv_cache_config, ) bind_kv_cache(kv_caches, forward_context, runner_kv_caches) return kv_caches diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index b958ef79d07..3221dc46c63 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -12,6 +12,7 @@ from contextlib import contextmanager from copy import copy, deepcopy from dataclasses import dataclass, replace from functools import reduce +from math import prod from typing import TYPE_CHECKING, Any, NamedTuple, TypeAlias, cast import numpy as np @@ -7029,10 +7030,21 @@ class GPUModelRunner( corresponding memory buffer for KV cache. """ kv_cache_raw_tensors: dict[str, torch.Tensor] = {} + packed_backing: torch.Tensor | None = None for kv_cache_tensor in kv_cache_config.kv_cache_tensors: - tensor = torch.zeros( - kv_cache_tensor.size, dtype=torch.int8, device=self.device - ) + if kv_cache_tensor.block_stride > 0: + # Allocate once; all packed tensors alias the same backing. + if packed_backing is None: + packed_backing = torch.zeros( + kv_cache_tensor.size, + dtype=torch.int8, + device=self.device, + ) + tensor = packed_backing + else: + tensor = torch.zeros( + kv_cache_tensor.size, dtype=torch.int8, device=self.device + ) for layer_name in kv_cache_tensor.shared_by: kv_cache_raw_tensors[layer_name] = tensor @@ -7074,6 +7086,14 @@ class GPUModelRunner( """ kv_caches: dict[str, torch.Tensor] = {} has_attn, has_mamba = False, False + + # Map layer names to (offset, block_stride) within the packed + # backing tensor so we can create strided views per layer. + layer_packing: dict[str, tuple[int, int]] = {} + for kv_tensor in self.kv_cache_config.kv_cache_tensors: + if kv_tensor.block_stride > 0: + for ln in kv_tensor.shared_by: + layer_packing[ln] = (kv_tensor.offset, kv_tensor.block_stride) for group in self._kv_cache_spec_attn_group_iterator(): kv_cache_spec = group.kv_cache_spec attn_backend = group.backend @@ -7085,8 +7105,13 @@ class GPUModelRunner( if layer_name in self.runner_only_attn_layers: continue raw_tensor = kv_cache_raw_tensors[layer_name] - assert raw_tensor.numel() % kv_cache_spec.page_size_bytes == 0 - num_blocks = raw_tensor.numel() // kv_cache_spec.page_size_bytes + packing = layer_packing.get(layer_name) + if packing is not None: + _, blk_stride = packing + num_blocks = raw_tensor.numel() // blk_stride + else: + assert raw_tensor.numel() % kv_cache_spec.page_size_bytes == 0 + num_blocks = raw_tensor.numel() // kv_cache_spec.page_size_bytes if isinstance(kv_cache_spec, AttentionSpec): has_attn = True num_blocks_per_kv_block = ( @@ -7127,8 +7152,17 @@ class GPUModelRunner( for i in range(len(kv_cache_stride_order)) ] - raw_tensor = kv_cache_raw_tensors[layer_name].view(dtype) - if kv_cache_spec.page_size_padded is not None: + if packing is not None: + offset, block_stride = packing + assert inv_order[0] == 0 + page_bytes = prod(kv_cache_shape[1:]) * get_dtype_size(dtype) + kv_cache = ( + kv_cache_raw_tensors[layer_name] + .view(-1, block_stride)[:, offset : offset + page_bytes] + .view(dtype) + .view(kv_cache_shape) + ) + elif kv_cache_spec.page_size_padded is not None: # Use strided view to handle page_size_bytes that # include padding. This follows # the same pattern as MambaSpec handling below. @@ -7142,13 +7176,17 @@ class GPUModelRunner( strides = list(torch.empty(kv_cache_shape).stride()) strides[inv_order[0]] = page_stride kv_cache = torch.as_strided( - raw_tensor, + kv_cache_raw_tensors[layer_name].view(dtype), size=kv_cache_shape, stride=tuple(strides), ) else: # No padding — safe to use a contiguous view. - kv_cache = raw_tensor.view(kv_cache_shape) + kv_cache = ( + kv_cache_raw_tensors[layer_name] + .view(dtype) + .view(kv_cache_shape) + ) kv_caches[layer_name] = kv_cache.permute(*inv_order) elif isinstance(kv_cache_spec, MambaSpec): From 4a8abf37c75b4a2587bfdad48bc6b442dc71332a Mon Sep 17 00:00:00 2001 From: Ben Browning Date: Fri, 19 Jun 2026 14:05:18 -0400 Subject: [PATCH 33/75] [Test] Migrate test_openai_schema.py to schemathesis 4.x (#46173) Signed-off-by: Ben Browning --- requirements/test/cuda.in | 2 +- requirements/test/nightly-torch.txt | 2 +- requirements/test/rocm.in | 2 +- .../entrypoints/openai/test_openai_schema.py | 101 ++++++++++-------- 4 files changed, 60 insertions(+), 47 deletions(-) diff --git a/requirements/test/cuda.in b/requirements/test/cuda.in index 8d7ad7d0aa2..a7fc65def8e 100644 --- a/requirements/test/cuda.in +++ b/requirements/test/cuda.in @@ -40,7 +40,7 @@ lm-eval[api]>=0.4.12 # required for model evaluation test mteb[bm25s]>=2, <3 # required for mteb test transformers==5.5.3 tokenizers==0.22.2 -schemathesis>=3.39.15 # Required for openai schema test. +schemathesis>=4.0.0 # Required for openai schema test. # quantization bitsandbytes==0.49.2 buildkite-test-collector==0.1.9 diff --git a/requirements/test/nightly-torch.txt b/requirements/test/nightly-torch.txt index 10eb7a62191..a58e0fa248f 100644 --- a/requirements/test/nightly-torch.txt +++ b/requirements/test/nightly-torch.txt @@ -31,7 +31,7 @@ lm-eval[api]>=0.4.12 # required for model evaluation test mteb[bm25s]>=2, <3 # required for mteb test transformers==5.5.3 tokenizers==0.22.2 -schemathesis>=3.39.15 # Required for openai schema test. +schemathesis>=4.0.0 # Required for openai schema test. # quantization bitsandbytes>=0.49.2 buildkite-test-collector==0.1.9 diff --git a/requirements/test/rocm.in b/requirements/test/rocm.in index ed10270f565..046ca09ff7f 100644 --- a/requirements/test/rocm.in +++ b/requirements/test/rocm.in @@ -39,7 +39,7 @@ lm-eval[api]>=0.4.12 # required for model evaluation test mteb[bm25s]>=2, <3 # required for mteb test transformers==5.5.3 tokenizers==0.22.2 -schemathesis>=3.39.15 # Required for openai schema test +schemathesis>=4.0.0 # Required for openai schema test # quantization bitsandbytes==0.49.2 buildkite-test-collector==0.1.9 diff --git a/tests/entrypoints/openai/test_openai_schema.py b/tests/entrypoints/openai/test_openai_schema.py index 56e4e9baf2e..38ea2661c86 100644 --- a/tests/entrypoints/openai/test_openai_schema.py +++ b/tests/entrypoints/openai/test_openai_schema.py @@ -6,15 +6,22 @@ from typing import Final import pytest import schemathesis from hypothesis import HealthCheck, settings -from schemathesis import GenerationConfig -from schemathesis.models import Case +from schemathesis import GenerationMode +from schemathesis.config import ( + ChecksConfig, + CoveragePhaseConfig, + GenerationConfig, + PhasesConfig, + PositiveDataAcceptanceConfig, + ProjectConfig, + ProjectsConfig, + SchemathesisConfig, +) from vllm.platforms import current_platform from ...utils import RemoteOpenAIServer -schemathesis.experimental.OPEN_API_3_1.enable() - MODEL_NAME = "HuggingFaceTB/SmolVLM-256M-Instruct" MAXIMUM_IMAGES = 2 _ROCM_TIMEOUT_MULTIPLIER = 3 if current_platform.is_rocm() else 1 @@ -44,21 +51,38 @@ def server(): @pytest.fixture(scope="module") def get_schema(server): # avoid generating null (\x00) bytes in strings during test case generation - return schemathesis.openapi.from_uri( + return schemathesis.openapi.from_url( f"{server.url_root}/openapi.json", - generation_config=GenerationConfig(allow_x00=False), + config=SchemathesisConfig( + projects=ProjectsConfig( + default=ProjectConfig( + generation=GenerationConfig( + allow_x00=False, + modes=[GenerationMode.POSITIVE], + ), + checks=ChecksConfig( + positive_data_acceptance=PositiveDataAcceptanceConfig( + enabled=False, + ), + ), + phases=PhasesConfig( + coverage=CoveragePhaseConfig(enabled=False), + ), + ), + ), + ), ) -schema = schemathesis.from_pytest_fixture("get_schema") +schema = schemathesis.pytest.from_fixture("get_schema") @schemathesis.hook -def before_generate_case(context: schemathesis.hooks.HookContext, strategy): +def before_generate_case(context: schemathesis.HookContext, strategy): op = context.operation assert op is not None - def no_invalid_types(case: schemathesis.models.Case): + def no_invalid_types(case: schemathesis.Case): """ Skips tool_calls with `"type": "custom"` which schemathesis incorrectly generates instead of the valid `"type": "function"`. @@ -68,39 +92,25 @@ def before_generate_case(context: schemathesis.hooks.HookContext, strategy): -d '{"messages": [{"role": "assistant", "tool_calls": [{"custom": {"input": "", "name": ""}, "id": "", "type": "custom"}]}]}' \ http://localhost:8000/v1/chat/completions """ # noqa: E501 - if hasattr(case, "body") and isinstance(case.body, dict): - if ( - "messages" in case.body - and isinstance(case.body["messages"], list) - and len(case.body["messages"]) > 0 - ): - for message in case.body["messages"]: - if not isinstance(message, dict): - continue + if ( + hasattr(case, "body") + and isinstance(case.body, dict) + and "messages" in case.body + and isinstance(case.body["messages"], list) + and len(case.body["messages"]) > 0 + ): + for message in case.body["messages"]: + if not isinstance(message, dict): + continue - tool_calls = message.get("tool_calls", []) - if isinstance(tool_calls, list): - for tool_call in tool_calls: - if isinstance(tool_call, dict): - if tool_call.get("type") != "function": - return False - if "custom" in tool_call: - return False - - # Sometimes structured_outputs.grammar is generated to be empty - # Causing a server error in EBNF grammar parsing - # https://github.com/vllm-project/vllm/pull/22587#issuecomment-3195253421 - structured_outputs = case.body.get("structured_outputs", {}) - grammar = ( - structured_outputs.get("grammar") - if isinstance(structured_outputs, dict) - else None - ) - - if grammar == "": - # Allow None (will be handled as no grammar) - # But skip empty strings - return False + tool_calls = message.get("tool_calls", []) + if isinstance(tool_calls, list): + for tool_call in tool_calls: + if isinstance(tool_call, dict): + if tool_call.get("type") != "function": + return False + if "custom" in tool_call: + return False return True @@ -108,7 +118,6 @@ def before_generate_case(context: schemathesis.hooks.HookContext, strategy): @schema.parametrize() -@schema.override(headers={"Content-Type": "application/json"}) @settings( deadline=LONG_TIMEOUT_SECONDS * 1000, max_examples=50, @@ -122,7 +131,7 @@ def before_generate_case(context: schemathesis.hooks.HookContext, strategy): # generating large-but-valid request bodies before vLLM is called. suppress_health_check=[HealthCheck.filter_too_much, HealthCheck.data_too_large], ) -def test_openapi_stateless(case: Case): +def test_openapi_stateless(case: schemathesis.Case): key = ( case.operation.method.upper(), case.operation.path, @@ -151,4 +160,8 @@ def test_openapi_stateless(case: Case): }.get(key, DEFAULT_TIMEOUT_SECONDS) # No need to verify SSL certificate for localhost - case.call_and_validate(verify=False, timeout=timeout) + case.call_and_validate( + verify=False, + timeout=timeout, + headers={"Content-Type": "application/json"}, + ) From 0a49fb2b13e474be71723c589cec5f4df1b5341d Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Fri, 19 Jun 2026 19:16:09 +0100 Subject: [PATCH 34/75] Fix dead link in docs (#46181) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- docs/contributing/model/basic.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/contributing/model/basic.md b/docs/contributing/model/basic.md index dceb78f5263..59e57e4ad14 100644 --- a/docs/contributing/model/basic.md +++ b/docs/contributing/model/basic.md @@ -133,7 +133,7 @@ The model should inherit protocol `IsAttentionFree` and also implement class met For the mamba layers themselves, please use the [`MambaMixer`](../../../vllm/model_executor/layers/mamba/mamba_mixer.py) (for Mamba-1) or [`MambaMixer2`](../../../vllm/model_executor/layers/mamba/mamba_mixer2.py) (for Mamba-2) classes. The model should also be added to the `MODELS_CONFIG_MAP` dictionary in [vllm/model_executor/models/config.py](../../../vllm/model_executor/models/config.py) to ensure that the runtime defaults are optimized. -For case (2), we recommend using as a reference the implementation of [`JambaForCausalLM`](../../../vllm/model_executor/models/jamba.py) (for an example of a model that uses Mamba-1 and attention together) or [`BambaForCausalLM`](../../../vllm/model_executor/models/bamba.py) (for an example of a model that uses Mamba-2 and attention together). +For case (2), we recommend using as a reference the implementation of [`JambaForCausalLM`](../../../vllm/model_executor/models/jamba.py) (for an example of a model that uses Mamba-1 and attention together) or [`NemotronHForCausalLM`](../../../vllm/model_executor/models/nemotron_h.py) (for an example of a model that uses Mamba-2 and attention together). These models should follow the same instructions as case (1), but they should inherit protocol `IsHybrid` (instead of `IsAttentionFree`) and it is *not* necessary to add them to the `MODELS_CONFIG_MAP` (their runtime defaults will be inferred from the protocol). 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. From dec860fb19fcd8a39c62a2204c5939feb4781f14 Mon Sep 17 00:00:00 2001 From: djramic Date: Fri, 19 Jun 2026 20:24:02 +0200 Subject: [PATCH 35/75] [ROCm] Use vLLM's fp8 quant max in AITER hipBLASLt accuracy test (#46176) Signed-off-by: Djordje Ramic --- tests/rocm/aiter/test_aiter_hipb_mm_linear_kernel.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/rocm/aiter/test_aiter_hipb_mm_linear_kernel.py b/tests/rocm/aiter/test_aiter_hipb_mm_linear_kernel.py index 92017e95cb7..c855d0e819d 100644 --- a/tests/rocm/aiter/test_aiter_hipb_mm_linear_kernel.py +++ b/tests/rocm/aiter/test_aiter_hipb_mm_linear_kernel.py @@ -18,6 +18,7 @@ from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import ( FP8ScaledMMLinearLayerConfig, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( + get_fp8_min_max, kFp8DynamicTokenSym, kFp8StaticChannelSym, kFp8StaticTensorSym, @@ -309,7 +310,7 @@ def test_hipb_mm_kernel_forward_accuracy(enable_hipb_mm_kernel): _check_bpreshuffle_runtime_support(weight_shape, num_tokens=num_tokens) fp8_dtype = current_platform.fp8_dtype() - fp8_max = torch.finfo(fp8_dtype).max + fp8_max = get_fp8_min_max()[1] device = torch.device("cuda") # Build a bf16 weight and quantize per output channel (one scale per row). From ca7e1f2c43834d1e720b7377e2832097978c1e35 Mon Sep 17 00:00:00 2001 From: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com> Date: Sat, 20 Jun 2026 00:12:40 +0400 Subject: [PATCH 36/75] Move CI failure diagnosis docs into ci-fails-buildkite skill (#45975) Signed-off-by: Vadim Gimpelson --- .claude/skills/ci-fails-buildkite/SKILL.md | 35 ++++++++++++++++++++++ .github/CODEOWNERS | 5 ++-- .gitignore | 4 ++- AGENTS.md | 11 ------- 4 files changed, 40 insertions(+), 15 deletions(-) create mode 100644 .claude/skills/ci-fails-buildkite/SKILL.md diff --git a/.claude/skills/ci-fails-buildkite/SKILL.md b/.claude/skills/ci-fails-buildkite/SKILL.md new file mode 100644 index 00000000000..d195c02f723 --- /dev/null +++ b/.claude/skills/ci-fails-buildkite/SKILL.md @@ -0,0 +1,35 @@ +--- +name: ci-fails-buildkite +description: Fetch and diagnose vLLM Buildkite CI failure logs. Use when investigating failing CI jobs on a PR or build, when the user pastes a buildkite.com URL, or asks to fetch/diagnose CI logs. +--- + +# Diagnosing vLLM Buildkite CI Failures + +Buildkite logs are public; no login needed. + +`.buildkite/scripts/ci-fetch-log.sh` saves each log as `ci--.log`, stripped of timestamps and ANSI codes. Existing files are kept; set `CI_FETCH_LOG_FORCE=1` to refetch. + +## Fetching logs + +```bash +# All failed jobs in a PR's latest build (current branch's PR if omitted): +.buildkite/scripts/ci-fetch-log.sh --pr + +# All failed jobs in a build (--soft also includes soft-failed jobs; +# --all fetches every finished job): +.buildkite/scripts/ci-fetch-log.sh "https://buildkite.com/vllm/ci/builds/" + +# One job — `gh pr checks` URLs (#) and web UI URLs (?sid=) both +# work; pass "-" as a second argument to stream to stdout: +.buildkite/scripts/ci-fetch-log.sh "https://buildkite.com/vllm/ci/builds/#" +``` + +To clean an already-downloaded log with `.buildkite/scripts/ci-clean-log.sh`: + +```bash +./ci-clean-log.sh ci.log +``` + +## Reference + +See [docs/contributing/ci/failures.md](../../../docs/contributing/ci/failures.md) for the full guide: filing CI failure issues, investigating/bisecting, reproducing flaky tests, and daily triage. diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 3a12aa3e6b5..15bd35f80e4 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -2,15 +2,14 @@ # for more info about CODEOWNERS file # This lists cover the "core" components of vLLM that require careful review -/vllm/compilation @zou3519 @youkaichao @ProExpertProg @BoyuanFeng @vadiklyutiy +/vllm/compilation @zou3519 @youkaichao @ProExpertProg @BoyuanFeng /vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery @xuechendi /vllm/lora @jeejeelee /vllm/model_executor/layers/attention @LucasWilkinson @MatthewBonanni /vllm/model_executor/layers/fused_moe @mgoin @pavanimajety @zyongye /vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth @yewentao256 @pavanimajety @zyongye /vllm/model_executor/layers/mamba @tdoublep @tomeras91 -/vllm/model_executor/layers/mamba/gdn_linear_attn.py @tdoublep @ZJY0516 @vadiklyutiy -/vllm/model_executor/layers/rotary_embedding.py @vadiklyutiy +/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @tdoublep @ZJY0516 @vadiklyutiy /vllm/model_executor/model_loader @22quinn /vllm/model_executor/layers/batch_invariant.py @yewentao256 /vllm/ir @ProExpertProg diff --git a/.gitignore b/.gitignore index c70200ed091..26cd21a015d 100644 --- a/.gitignore +++ b/.gitignore @@ -199,7 +199,9 @@ cython_debug/ .vscode/ # Claude -.claude/ +.claude/* +!.claude/skills/ +!.claude/skills/** # Codex .codex/ diff --git a/AGENTS.md b/AGENTS.md index 1f3a083f80c..7d6fd9e0970 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -114,17 +114,6 @@ Follow these rules for all code changes in this repository: - Keep comments and docstrings minimal and concise. - Assume the reader is familiar with vLLM. -### Diagnosing CI failures - -Buildkite logs are public; no login needed. Details: [docs/contributing/ci/failures.md](docs/contributing/ci/failures.md). - -```bash -# All failed-job logs for a PR's latest build (current branch's PR if omitted): -.buildkite/scripts/ci-fetch-log.sh --pr -# Any Buildkite build or job URL also works: -.buildkite/scripts/ci-fetch-log.sh "" -``` - ### Commit messages Add attribution using commit trailers such as `Co-authored-by:` (other projects use `Assisted-by:` or `Generated-by:`). For example: From 4a083cc858f075209dd964ade48c0f8ec87c3393 Mon Sep 17 00:00:00 2001 From: Micah Williamson Date: Fri, 19 Jun 2026 15:20:06 -0500 Subject: [PATCH 37/75] [ROCm][CI] Pin `test_rocm_compressed_tensors_w8a8` to TRITON_ATTN (#46180) Signed-off-by: Micah Williamson --- .buildkite/test_areas/kernels.yaml | 1 + tests/kernels/quantization/test_triton_scaled_mm.py | 6 ++++-- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 159f940530e..ebcb95a9d82 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -104,6 +104,7 @@ steps: source_file_dependencies: - csrc/quantization/ - vllm/model_executor/layers/quantization + - vllm/config/ - tests/kernels/quantization - tests/kernels/quantization/test_rocm_skinny_gemms.py - vllm/_aiter_ops.py diff --git a/tests/kernels/quantization/test_triton_scaled_mm.py b/tests/kernels/quantization/test_triton_scaled_mm.py index 1cef5eb93a5..d857d495f2d 100644 --- a/tests/kernels/quantization/test_triton_scaled_mm.py +++ b/tests/kernels/quantization/test_triton_scaled_mm.py @@ -60,8 +60,10 @@ def test_rocm_compressed_tensors_w8a8( vllm_runner, example_prompts, model_path, max_tokens, num_logprobs ): dtype = "bfloat16" - - with vllm_runner(model_path, dtype=dtype) as vllm_model: + # Pin to TRITON_ATTN, see https://github.com/vllm-project/vllm/issues/46179 + with vllm_runner( + model_path, dtype=dtype, attention_backend="TRITON_ATTN" + ) as vllm_model: vllm_model.generate_greedy_logprobs(example_prompts, max_tokens, num_logprobs) From 859e4d436ba0fb0da8a655a80d5c4fab12adc82e Mon Sep 17 00:00:00 2001 From: Ben Browning Date: Fri, 19 Jun 2026 18:09:28 -0400 Subject: [PATCH 38/75] [Bugfix][Parser] Fix U+FFFD leak at reasoning-to-content transition in engine parsers (#46159) Signed-off-by: Ben Browning --- tests/parser/engine/replay_harness.py | 3 + tests/parser/engine/test_delegating_replay.py | 3 +- .../engine/test_ufffd_reasoning_transition.py | 181 ++++++++++++++++++ vllm/parser/abstract_parser.py | 7 +- vllm/parser/engine/parser_engine.py | 4 +- 5 files changed, 191 insertions(+), 7 deletions(-) create mode 100644 tests/parser/engine/test_ufffd_reasoning_transition.py diff --git a/tests/parser/engine/replay_harness.py b/tests/parser/engine/replay_harness.py index fac643390b1..9abd460f769 100644 --- a/tests/parser/engine/replay_harness.py +++ b/tests/parser/engine/replay_harness.py @@ -96,6 +96,9 @@ class MockTokenizer: return "".join(parts) +CHUNK_SIZES = [1, 2, 3, 5, 11, 23, None] + + def make_mock_tokenizer(sample: Sample) -> MockTokenizer: """Build a mock tokenizer from a sample's vocab and token data.""" return MockTokenizer( diff --git a/tests/parser/engine/test_delegating_replay.py b/tests/parser/engine/test_delegating_replay.py index 7460ab21ec5..5d5d6b3247d 100644 --- a/tests/parser/engine/test_delegating_replay.py +++ b/tests/parser/engine/test_delegating_replay.py @@ -19,6 +19,7 @@ import pytest from pydantic import TypeAdapter from tests.parser.engine.replay_harness import ( + CHUNK_SIZES, MockTokenizer, assert_parse_output, collect_output, @@ -113,8 +114,6 @@ _PAIRINGS = _discover_pairings() _ALL_SAMPLES = [(p.parser_cls, s) for p in _PAIRINGS for s in p.samples] -CHUNK_SIZES = [1, 2, 3, 5, 11, 23, None] - @pytest.mark.parametrize("chunk_size", CHUNK_SIZES, ids=lambda c: f"chunk={c}") @pytest.mark.parametrize( diff --git a/tests/parser/engine/test_ufffd_reasoning_transition.py b/tests/parser/engine/test_ufffd_reasoning_transition.py new file mode 100644 index 00000000000..ffd2ead95d7 --- /dev/null +++ b/tests/parser/engine/test_ufffd_reasoning_transition.py @@ -0,0 +1,181 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression test for U+FFFD leak at reasoning→content transition. + +When byte-fallback tokens span the reasoning/content boundary, +decoding isolated content-side token IDs via tokenizer.decode() +produces U+FFFD (Unicode replacement character). The fix flushes +the reasoning parser's engine lexer instead. + +Reproduces the bug at various chunk sizes and validates that the +fix prevents U+FFFD from leaking into streamed content. +""" + +from __future__ import annotations + +import pytest + +from tests.parser.engine.replay_harness import ( + CHUNK_SIZES, + MockTokenizer, + collect_output, + replay_streaming, +) +from vllm.parser.abstract_parser import DelegatingParser +from vllm.parser.engine.registered_adapters import ( + Glm47MoeParserReasoningAdapter, + Glm47MoeParserToolAdapter, + Qwen3ParserReasoningAdapter, + Qwen3ParserToolAdapter, +) + + +class ByteFallbackMockTokenizer(MockTokenizer): + """MockTokenizer that returns U+FFFD for specified token IDs. + + Simulates byte-fallback tokenizer behavior where isolated + partial-byte tokens decode to the Unicode replacement character. + """ + + def __init__( + self, + vocab: dict[str, int], + tokens: list[tuple[int, str]], + ufffd_token_ids: set[int], + ) -> None: + super().__init__(vocab, tokens) + self._ufffd_token_ids = frozenset(ufffd_token_ids) + + def decode(self, ids: list[int], skip_special_tokens: bool = False) -> str: + parts: list[str] = [] + for tid in ids: + if skip_special_tokens and tid in self._special_ids: + continue + if tid in self._ufffd_token_ids: + parts.append("�") + else: + text = self._token_decode_map.get(tid, f"?{tid}?") + parts.append(text) + return "".join(parts) + + +# ── Model-specific DelegatingParser subclasses ─────────────────────── + + +class _Glm47Delegating(DelegatingParser): + reasoning_parser_cls = Glm47MoeParserReasoningAdapter + tool_parser_cls = Glm47MoeParserToolAdapter + + +class _Qwen3Delegating(DelegatingParser): + reasoning_parser_cls = Qwen3ParserReasoningAdapter + tool_parser_cls = Qwen3ParserToolAdapter + + +# ── Shared test data ───────────────────────────────────────────────── + +_SHARED_TOKENS: list[tuple[int, str]] = [ + (100, "Let me"), + (101, " think"), + (102, " about"), + (103, " Samsung."), + (51, ""), + (200, "삼성"), + (201, "전자의"), + (202, " 주가를"), + (203, " 분석합니다."), +] + +_SHARED_UFFFD_IDS: set[int] = {200} + +EXPECTED_REASONING = "Let me think about Samsung." +EXPECTED_CONTENT = "삼성전자의 주가를 분석합니다." + +_MODEL_CONFIGS = [ + pytest.param( + { + "": 50, + "": 51, + "": 60, + "": 61, + "": 62, + "": 63, + "": 64, + "": 65, + }, + _Glm47Delegating, + id="glm47", + ), + pytest.param( + { + "": 50, + "": 51, + "": 60, + "": 61, + }, + _Qwen3Delegating, + id="qwen3", + ), +] + + +# ── Tests ──────────────────────────────────────────────────────────── + + +class TestUfffdReasoningTransition: + """U+FFFD must not appear at the reasoning→content transition.""" + + @pytest.mark.parametrize("vocab,delegating_cls", _MODEL_CONFIGS) + @pytest.mark.parametrize("chunk_size", CHUNK_SIZES, ids=lambda c: f"chunk={c}") + def test_no_ufffd(self, chunk_size, vocab, delegating_cls): + tokenizer = ByteFallbackMockTokenizer(vocab, _SHARED_TOKENS, _SHARED_UFFFD_IDS) + parser = delegating_cls(tokenizer) + deltas = replay_streaming( + parser, + _SHARED_TOKENS, + chunk_size=chunk_size, + finished_on_last=True, + ) + output = collect_output(deltas) + + assert "�" not in output.content, ( + f"U+FFFD leaked into content: {output.content!r}" + ) + assert output.content == EXPECTED_CONTENT + assert output.reasoning == EXPECTED_REASONING + + def test_byte_fallback_tokenizer_produces_ufffd(self): + """Validate the fixture: decode() returns U+FFFD for isolated + byte-fallback token IDs, proving the old code path would leak.""" + vocab = dict(_MODEL_CONFIGS[0].values[0]) + tokenizer = ByteFallbackMockTokenizer(vocab, _SHARED_TOKENS, _SHARED_UFFFD_IDS) + assert tokenizer.decode([200]) == "�" + + @pytest.mark.parametrize("chunk_size", CHUNK_SIZES, ids=lambda c: f"chunk={c}") + def test_multiple_ufffd_tokens_at_boundary(self, chunk_size): + """Multiple consecutive byte-fallback tokens at the boundary.""" + tokens: list[tuple[int, str]] = [ + (100, "Reasoning."), + (51, ""), + (200, "삼"), + (201, "성"), + (202, "전자"), + ] + ufffd_ids: set[int] = {200, 201} + vocab = dict(_MODEL_CONFIGS[0].values[0]) + + tokenizer = ByteFallbackMockTokenizer(vocab, tokens, ufffd_ids) + parser = _Glm47Delegating(tokenizer) + deltas = replay_streaming( + parser, + tokens, + chunk_size=chunk_size, + finished_on_last=True, + ) + output = collect_output(deltas) + + assert "�" not in output.content, ( + f"U+FFFD leaked into content: {output.content!r}" + ) + assert output.content == "삼성전자" + assert output.reasoning == "Reasoning." diff --git a/vllm/parser/abstract_parser.py b/vllm/parser/abstract_parser.py index 915d401f7bd..11fca8e43ab 100644 --- a/vllm/parser/abstract_parser.py +++ b/vllm/parser/abstract_parser.py @@ -794,11 +794,10 @@ class DelegatingParser(Parser): reasoning_transitioned = True current_token_ids = self.extract_content_ids(delta_token_ids) if self._engine_based: + flush_delta = reasoning_parser.finish_streaming() # type: ignore[union-attr, attr-defined] current_text = ( - self.model_tokenizer.decode(current_token_ids) - if current_token_ids - else "" - ) + (delta_message.content if delta_message else None) or "" + ) + ((flush_delta.content if flush_delta else None) or "") if delta_message and self._tool_parser is not None: delta_message.content = None else: diff --git a/vllm/parser/engine/parser_engine.py b/vllm/parser/engine/parser_engine.py index dafb26fc48d..ba838d31a0b 100644 --- a/vllm/parser/engine/parser_engine.py +++ b/vllm/parser/engine/parser_engine.py @@ -172,7 +172,9 @@ class ParserEngine(Parser): def finish_streaming(self) -> DeltaMessage | None: events = self._engine.finish() - return self._events_to_delta(events) if events else None + if events or self._deferred_content: + return self._events_to_delta(events, finished=True) + return None def _reset(self, initial_state: ParserState | None = None) -> None: self._engine.reset(initial_state=initial_state) From e6cd8913ddfe63b4620e45ff8c2da1d37318dbe5 Mon Sep 17 00:00:00 2001 From: Charlie Fu Date: Fri, 19 Jun 2026 17:20:10 -0500 Subject: [PATCH 39/75] [ROCm][CI] Skip Qwen3.5-35B-A3B-MXFP4-AITER-TP2 for non gfx950 (#46109) Signed-off-by: charlifu --- .buildkite/test_areas/lm_eval.yaml | 9 +++++++++ tests/evals/gsm8k/test_gsm8k_correctness.py | 9 +++++++++ 2 files changed, 18 insertions(+) diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index fc8e72699e4..217fc5665c8 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -101,6 +101,15 @@ steps: num_devices: 8 commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200.txt + mirror: + amd: + device: mi300_8 + timeout_in_minutes: 180 + depends_on: + - image-build-amd + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt - label: MoE Refactor Integration Test (H100 - TEMPORARY) key: moe-refactor-integration-test-h100-temporary diff --git a/tests/evals/gsm8k/test_gsm8k_correctness.py b/tests/evals/gsm8k/test_gsm8k_correctness.py index e7a254e760f..cd90d71669a 100644 --- a/tests/evals/gsm8k/test_gsm8k_correctness.py +++ b/tests/evals/gsm8k/test_gsm8k_correctness.py @@ -78,7 +78,16 @@ def test_gsm8k_correctness(config_filename): "Skipping DeepSeek-V3.2 and DeepSeek-R1 on ROCm platforms " "due to agent pool disk space issues and pod evictions." ) + if current_platform.is_rocm() and ( + "Qwen3.5-35B-A3B-MXFP4-AITER-TP2" in config_filename.name + ): + from vllm.platforms.rocm import on_gfx950 + if not on_gfx950(): + pytest.skip( + "Skipping Qwen3.5-35B-A3B-MXFP4-AITER-TP2 on non-GFX950 platforms. " + "The quantization scheme is not supported on non-GFX950 platforms." + ) # Parse server arguments from config (use shlex to handle quoted strings) server_args_str = eval_config.get("server_args", "") server_args = shlex.split(server_args_str) if server_args_str else [] From 0fbf42af841993ab1c189efca34de6b9799526b7 Mon Sep 17 00:00:00 2001 From: djramic Date: Sat, 20 Jun 2026 00:20:59 +0200 Subject: [PATCH 40/75] [ROCm] Fix VRAM not freed in test_phi3v (#46046) Signed-off-by: Djordje Ramic --- .buildkite/test_areas/models_multimodal.yaml | 1 - tests/models/multimodal/pooling/test_phi3v.py | 2 ++ 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/.buildkite/test_areas/models_multimodal.yaml b/.buildkite/test_areas/models_multimodal.yaml index a7358e8dbd6..27e73e55a3f 100644 --- a/.buildkite/test_areas/models_multimodal.yaml +++ b/.buildkite/test_areas/models_multimodal.yaml @@ -68,7 +68,6 @@ steps: - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model # Otherwise, mp_method="spawn" doesn't work mirror: amd: - soft_fail: true device: mi325_1 depends_on: - image-build-amd diff --git a/tests/models/multimodal/pooling/test_phi3v.py b/tests/models/multimodal/pooling/test_phi3v.py index 285ded375da..ba017f065a1 100644 --- a/tests/models/multimodal/pooling/test_phi3v.py +++ b/tests/models/multimodal/pooling/test_phi3v.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest +import torch import torch.nn.functional as F import transformers.utils from PIL import Image @@ -52,6 +53,7 @@ def _get_cherry_blossom_image() -> Image.Image: ) +@torch.inference_mode() def _run_test( hf_runner: type[HfRunner], vllm_runner: type[VllmRunner], From 93bad119120d0f9bff707dcbf5af5c029158b969 Mon Sep 17 00:00:00 2001 From: JasonLi314 <47095666+JasonLi314@users.noreply.github.com> Date: Fri, 19 Jun 2026 20:27:45 -0700 Subject: [PATCH 41/75] [Bugfix] Fix gridDim.y overflow for large row counts (#45255) Signed-off-by: Jason Li --- .../w8a8/fp8/per_token_group_quant.cu | 26 +++++---- .../test_per_token_group_quant.py | 57 +++++++++++++++++++ 2 files changed, 71 insertions(+), 12 deletions(-) diff --git a/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu b/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu index e3017e6ca21..902391b8f6d 100644 --- a/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu +++ b/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu @@ -301,8 +301,9 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel( const int sf_k_local = local_group_id % kGroupsPerBlockX; const int row_local = local_group_id / kGroupsPerBlockX; - const int sf_k_idx = blockIdx.x * kGroupsPerBlockX + sf_k_local; - const int mn_idx = blockIdx.y * kRowsPerBlock + row_local; + // Rows on grid.x: mn scales with tokens and can exceed the 65535 grid.y cap. + const int sf_k_idx = blockIdx.y * kGroupsPerBlockX + sf_k_local; + const int mn_idx = blockIdx.x * kRowsPerBlock + row_local; #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) asm volatile("griddepcontrol.wait;"); @@ -496,14 +497,15 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input, " is not a multiple of 4."); const int kx = GetGroupsPerBlockX(padded_groups_per_row); const int ry = 16 / kx; - const int64_t blocks_x = padded_groups_per_row / kx; - const int64_t blocks_y = (tma_aligned_mn + ry - 1) / ry; + const int64_t row_blocks = (tma_aligned_mn + ry - 1) / ry; + const int64_t sf_k_blocks = padded_groups_per_row / kx; const int num_threads = (kx * ry) * THREADS_PER_GROUP; - // CUDA caps grid.x and grid.y at 2^31 - 1; guard against pathological inputs. - STD_TORCH_CHECK(blocks_x <= static_cast(INT32_MAX) && - blocks_y <= static_cast(INT32_MAX), + // CUDA caps grid.x at 2^31 - 1 and grid.y at 2^16 - 1 (65535). + constexpr int64_t kMaxGridDimYZ = 65535; + STD_TORCH_CHECK(row_blocks <= static_cast(INT32_MAX) && + sf_k_blocks <= kMaxGridDimYZ, "per_token_group_quant_8bit_packed grid too large: (", - blocks_x, ", ", blocks_y, ")."); + row_blocks, ", ", sf_k_blocks, ")."); auto dst_type = output_q.scalar_type(); @@ -513,8 +515,8 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input, #define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, KX, RY) \ do { \ cudaLaunchConfig_t config = {}; \ - config.gridDim = dim3(static_cast(blocks_x), \ - static_cast(blocks_y)); \ + config.gridDim = dim3(static_cast(row_blocks), \ + static_cast(sf_k_blocks)); \ config.blockDim = dim3(num_threads); \ config.dynamicSmemBytes = 0; \ config.stream = stream; \ @@ -539,8 +541,8 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input, #else #define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, KX, RY) \ do { \ - dim3 grid(static_cast(blocks_x), \ - static_cast(blocks_y)); \ + dim3 grid(static_cast(row_blocks), \ + static_cast(sf_k_blocks)); \ dim3 block(num_threads); \ per_token_group_quant_8bit_packed_register_kernel \ diff --git a/tests/kernels/quantization/test_per_token_group_quant.py b/tests/kernels/quantization/test_per_token_group_quant.py index d957cefed4d..0d9b6c0c3e8 100644 --- a/tests/kernels/quantization/test_per_token_group_quant.py +++ b/tests/kernels/quantization/test_per_token_group_quant.py @@ -345,6 +345,63 @@ def test_per_token_group_quant_fp8_packed_zero_fills_padded_output_q( ) +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="packed FP8 per-token-group quant kernel requires a CUDA-alike GPU", +) +def test_per_token_group_quant_fp8_packed_large_mn(): + """Regression test for https://github.com/vllm-project/vllm/issues/45099. + + Some background: gridDim.x and gridDim.y have different limits of 2^31 - 1 and + 2^16 - 1, respectively. + Prior code introduced a bug where it incorrectly assumed grid.x and y both have + 2^31 - 1 limits and mixed them up, which doesn't surface until the kernel is + launched with a large mn that exceeds grid.y limit (2^16 - 1). + + This issue doesn't surface often because each forward pass only processes a + bounded token batch, not the full context. + Quantizing tensors with more rows than that will fail at launch with + "CUDA error: invalid argument". + This is a differential test that compares fp8 output against Triton output + reference when token size sits just above the gridDim.y 2^16 - 1 limit. + """ + + device = "cuda" + group_size = 128 + # hidden 2048 -> 2048/128 = 16 groups per row -> kx=16, ry=1: one grid row per mn + # row, so any mn > 65535 overflowed grid.y before the fix. + num_tokens, hidden_dim = 65537, 2048 + torch.manual_seed(42) + x = torch.randn((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) * 8 + + out_q, out_s_packed = fp8_utils.per_token_group_quant_fp8_packed_for_deepgemm( + x, + group_size=group_size, + use_ue8m0=True, + ) + + with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False): + ref_q, ref_s = fp8_utils.per_token_group_quant_fp8( + x, group_size, use_ue8m0=True + ) + + assert torch.equal(out_q, ref_q), "Quantized output mismatch" + + # Vectorized packed-scale check; the per-element loop used by the smaller + # tests is too slow at this size. groups_per_row is a multiple of 4 here, + # so there is no K padding and the packed view lines up. + mn = num_tokens + groups_per_row = hidden_dim // group_size + k_num_packed = (groups_per_row + 3) // 4 + assert groups_per_row % 4 == 0 + ref_exponents = (ref_s.reshape(mn, groups_per_row).view(torch.int32) >> 23) & 0xFF + exp = ref_exponents.view(mn, k_num_packed, 4) + expected = ( + exp[..., 0] | (exp[..., 1] << 8) | (exp[..., 2] << 16) | (exp[..., 3] << 24) + ) + assert torch.equal(out_s_packed.cpu(), expected.cpu()), "Packed scale mismatch" + + @pytest.mark.parametrize("shape", [(32, 128), (64, 256), (16, 512)]) @pytest.mark.parametrize("group_size", [64, 128]) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") From dced2907693e3d6bf9eb7168d0a8fecf1cd22dca Mon Sep 17 00:00:00 2001 From: Matt <156021403+mawong-amd@users.noreply.github.com> Date: Sat, 20 Jun 2026 02:04:35 -0500 Subject: [PATCH 42/75] [Hardware][AMD][CI] Fix e2e core test group (#46024) Signed-off-by: Matthew Wong Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .buildkite/test-amd.yaml | 15 +-------------- .buildkite/test_areas/engine.yaml | 10 ++++++++++ tests/v1/e2e/general/test_cascade_attention.py | 7 +++++++ 3 files changed, 18 insertions(+), 14 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index a7f3d67e79f..5550c0a0c18 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -647,7 +647,7 @@ steps: - pytest -v -s v1/cudagraph/test_cudagraph_mode.py - label: e2e Core (1 GPU) # TBD - timeout_in_minutes: 180 + timeout_in_minutes: 35 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 optional: true @@ -2075,19 +2075,6 @@ steps: - export VLLM_ALLOW_INSECURE_SERIALIZATION=1 - pytest -v -s v1/spec_decode/test_acceptance_length.py -m slow_test -- label: e2e Core (1 GPU) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - agent_pool: mi300_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/v1/ - - tests/v1/e2e/ - - vllm/platforms/rocm.py - commands: - - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py - - label: e2e Scheduling (1 GPU) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml index 67ed8e377ae..98c8231831d 100644 --- a/.buildkite/test_areas/engine.yaml +++ b/.buildkite/test_areas/engine.yaml @@ -74,6 +74,16 @@ steps: - tests/v1/e2e/general/ commands: - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py + mirror: + amd: + device: mi250_1 + timeout_in_minutes: 35 + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/v1/ + - tests/v1/e2e/general/ + - vllm/platforms/rocm.py - label: V1 e2e (2 GPUs) key: v1-e2e-2-gpus diff --git a/tests/v1/e2e/general/test_cascade_attention.py b/tests/v1/e2e/general/test_cascade_attention.py index be889b38690..251746271de 100644 --- a/tests/v1/e2e/general/test_cascade_attention.py +++ b/tests/v1/e2e/general/test_cascade_attention.py @@ -4,9 +4,16 @@ import pytest from vllm import LLM, SamplingParams +from vllm.platforms import current_platform from ....utils import create_new_process_for_each_test +if current_platform.is_rocm(): + pytest.skip( + "Cascade attention backends FLASH_ATTN and FLASHINFER are notsupported on ROCm", + allow_module_level=True, + ) + @create_new_process_for_each_test() @pytest.mark.parametrize("attn_backend", ["FLASH_ATTN", "FLASHINFER"]) From 7ff7f5c8eb98354d3776f6c60c90aebc2b41c1da Mon Sep 17 00:00:00 2001 From: Sumanth R Hegde <39546518+SumanthRH@users.noreply.github.com> Date: Sat, 20 Jun 2026 07:09:09 -0700 Subject: [PATCH 43/75] Revert "Fix Stale Encoder Cache After Weight Update" (#46125) --- vllm/entrypoints/llm.py | 6 ------ vllm/v1/engine/async_llm.py | 6 ------ 2 files changed, 12 deletions(-) diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index 349091f4b79..892e5035ab6 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -898,12 +898,6 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin): def finish_weight_update(self) -> None: """Finish the current weight update.""" self.llm_engine.collective_rpc("finish_weight_update") - # Invalidate cached state computed with the old weights so it isn't - # reused for subsequent requests: - # - prefix cache: KV blocks computed with the old weights - # - encoder cache: multimodal embeddings keyed only by mm_hash - self.llm_engine.reset_prefix_cache() - self.llm_engine.reset_encoder_cache() def __repr__(self) -> str: """Return a transformers-style hierarchical view of the model.""" diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 26b3f53d2c4..419e15163a9 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -1109,9 +1109,3 @@ class AsyncLLM(EngineClient): async def finish_weight_update(self) -> None: """Finish the current weight update.""" await self.collective_rpc("finish_weight_update") - # Invalidate cached state computed with the old weights so it isn't - # reused for subsequent requests: - # - prefix cache: KV blocks computed with the old weights - # - encoder cache: multimodal embeddings keyed only by mm_hash - await self.reset_prefix_cache() - await self.reset_encoder_cache() From d272418f459a82e1012b60116ac00659a7017cde Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?L=E4=B8=B6?= <30801931+Sirius29@users.noreply.github.com> Date: Sat, 20 Jun 2026 22:09:18 +0800 Subject: [PATCH 44/75] [Perf] Optimize Qwen3-VL multi-video prompt processing (#46026) Signed-off-by: Sirius29 <422058530@qq.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../multimodal/processing/test_qwen3_vl.py | 46 ++++++++ vllm/model_executor/models/qwen3_vl.py | 104 +++++++++++++----- 2 files changed, 124 insertions(+), 26 deletions(-) diff --git a/tests/models/multimodal/processing/test_qwen3_vl.py b/tests/models/multimodal/processing/test_qwen3_vl.py index d69c31b582a..9155fde5033 100644 --- a/tests/models/multimodal/processing/test_qwen3_vl.py +++ b/tests/models/multimodal/processing/test_qwen3_vl.py @@ -92,3 +92,49 @@ def test_processor_num_frames_timestamp( assert len(video_phs) == 1, ( f"Expected exactly 1 video placeholder, got {len(video_phs)}" ) + + +@pytest.mark.parametrize("model_id", [MODEL_ID]) +@pytest.mark.parametrize("num_videos", [2, 4]) +def test_processor_multi_video( + model_id: str, + num_videos: int, +) -> None: + """Verify that multi-video processing produces correct placeholders. + + This exercises the token-level replacement path in + ``_call_hf_processor`` which avoids the quadratic text-level + prompt expansion. + """ + ctx = build_model_context( + model_id, + limit_mm_per_prompt={"image": 0, "video": num_videos}, + ) + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) + + prompt = "<|vision_start|><|video_pad|><|vision_end|>" * num_videos + mm_data = {"video": [_build_video_mm_data(num_frames=8)["video"][0]] * num_videos} + + processed = processor( + prompt, + mm_items=processor.info.parse_mm_data(mm_data), + hf_processor_mm_kwargs={"num_frames": 8}, + ) + + token_ids = processed["prompt_token_ids"] + assert len(token_ids) > 0 + + video_phs = processed["mm_placeholders"].get("video", []) + assert len(video_phs) == num_videos, ( + f"Expected {num_videos} video placeholders, got {len(video_phs)}" + ) + + # All placeholders should have the same length (same video params) + # and must not overlap. + lengths = {ph.length for ph in video_phs} + assert len(lengths) == 1, f"Placeholder lengths differ: {lengths}" + for i in range(1, len(video_phs)): + prev_end = video_phs[i - 1].offset + video_phs[i - 1].length + assert video_phs[i].offset >= prev_end, ( + f"Placeholder {i} overlaps with placeholder {i - 1}" + ) diff --git a/vllm/model_executor/models/qwen3_vl.py b/vllm/model_executor/models/qwen3_vl.py index 1423770be02..3183a23ffde 100644 --- a/vllm/model_executor/models/qwen3_vl.py +++ b/vllm/model_executor/models/qwen3_vl.py @@ -1202,6 +1202,49 @@ class Qwen3VLDummyInputsBuilder(BaseDummyInputsBuilder[Qwen3VLProcessingInfo]): return video_items +def _replace_video_token_placeholders( + prompt_ids: list[int], + target: list[int], + replacements: list[list[int]], +) -> list[int]: + """Replace each 3-token video placeholder with its expanded sequence. + + Args: + prompt_ids: Token IDs of the original (unexpanded) prompt. + target: 3-element list ``[vision_start_id, video_pad_id, + vision_end_id]`` to search for. + replacements: Per-video expanded token sequences, in prompt order. + + Returns: + Token IDs with every placeholder triplet replaced. + """ + result: list[int] = [] + repl_idx = 0 + i = 0 + n = len(prompt_ids) + t0, t1, t2 = target + num_repl = len(replacements) + + while i < n: + if ( + i + 2 < n + and prompt_ids[i] == t0 + and prompt_ids[i + 1] == t1 + and prompt_ids[i + 2] == t2 + ): + result.extend(replacements[repl_idx]) + repl_idx += 1 + i += 3 + else: + result.append(prompt_ids[i]) + i += 1 + + assert repl_idx == num_repl, ( + f"Found {repl_idx} video placeholders but expected {num_repl}" + ) + return result + + class Qwen3VLMultiModalProcessor(BaseMultiModalProcessor[Qwen3VLProcessingInfo]): def _call_hf_processor( self, @@ -1211,15 +1254,23 @@ class Qwen3VLMultiModalProcessor(BaseMultiModalProcessor[Qwen3VLProcessingInfo]) tok_kwargs: Mapping[str, object], ) -> BatchFeature: mm_data = dict(mm_data) - processor = self.info.get_hf_processor(**mm_kwargs) # Separate video processing from image processing. Because the videos # are processed into several image patches + video_input_ids_lst: list[list[int]] = [] if videos := mm_data.pop("videos", []): video_grid_thw_lst = [] pixel_values_videos_lst = [] timestamps_per_video = [] + hf_config = self.info.get_hf_config() + tokenizer = self.info.get_tokenizer() + merge_size = hf_config.vision_config.spatial_merge_size + video_pruning_rate = self.info.ctx.get_mm_config().video_pruning_rate + vision_start_token_id = hf_config.vision_start_token_id + vision_end_token_id = hf_config.vision_end_token_id + video_token_id = hf_config.video_token_id + for item in videos: video_array, metadata = item @@ -1269,55 +1320,38 @@ class Qwen3VLMultiModalProcessor(BaseMultiModalProcessor[Qwen3VLProcessingInfo]) tok_kwargs=tok_kwargs, ) - merge_size = processor.video_processor.merge_size - # Get video grid info for EVS calculation. + # Discard HF output input_ids — we use get_video_repl below + # to generate the correct (EVS-adjusted) token sequence. + video_outputs.pop("input_ids", None) + video_grid_thw = video_outputs["video_grid_thw"] num_frames = int(video_grid_thw[0, 0]) tokens_per_frame_base = int(video_grid_thw[0, 1:].prod()) // ( merge_size**2 ) - # Apply EVS if enabled. - video_pruning_rate = self.info.ctx.get_mm_config().video_pruning_rate if video_pruning_rate is not None and video_pruning_rate > 0.0: num_tokens = compute_retained_tokens_count( tokens_per_frame=tokens_per_frame_base, num_frames=num_frames, q=video_pruning_rate, ) - # Here we just need placeholders that won't actually be replaced - - # we just need to make sure the total number of tokens is correct - # assign all tokens to the first frame. tokens_per_frame = [num_tokens] + [0] * (num_frames - 1) select_token_id = False else: tokens_per_frame = [tokens_per_frame_base] * num_frames select_token_id = True - # Generate the video replacement with EVS-adjusted token counts - tokenizer = self.info.get_tokenizer() - hf_config = self.info.get_hf_config() video_repl = Qwen3VLMultiModalProcessor.get_video_repl( tokens_per_frame=tokens_per_frame, timestamps=timestamps, tokenizer=tokenizer, - vision_start_token_id=hf_config.vision_start_token_id, - vision_end_token_id=hf_config.vision_end_token_id, - video_token_id=hf_config.video_token_id, + vision_start_token_id=vision_start_token_id, + vision_end_token_id=vision_end_token_id, + video_token_id=video_token_id, select_token_id=select_token_id, ) - - # Convert token IDs to text for the HF processor flow - video_placeholder = tokenizer.decode( - video_repl.full, skip_special_tokens=False - ) - input_ids = video_outputs.pop("input_ids") - video_placeholder = processor.tokenizer.batch_decode(input_ids)[0] - prompt = prompt.replace( - "<|vision_start|><|video_pad|><|vision_end|>", - video_placeholder, - 1, - ) + video_input_ids_lst.append(list(video_repl.full)) video_grid_thw_lst.append(video_outputs["video_grid_thw"]) pixel_values_videos_lst.append(video_outputs["pixel_values_videos"]) @@ -1335,6 +1369,24 @@ class Qwen3VLMultiModalProcessor(BaseMultiModalProcessor[Qwen3VLProcessingInfo]) mm_kwargs=mm_kwargs, tok_kwargs=tok_kwargs, ) + + # Replace each placeholder triplet with pre-computed video tokens. + if video_input_ids_lst: + hf_config = self.info.get_hf_config() + video_target = [ + hf_config.vision_start_token_id, + hf_config.video_token_id, + hf_config.vision_end_token_id, + ] + input_ids = processed_outputs.pop("input_ids") + if not isinstance(input_ids, list): + input_ids = input_ids.tolist() + (prompt_ids,) = input_ids + expanded_ids = _replace_video_token_placeholders( + prompt_ids, video_target, video_input_ids_lst + ) + processed_outputs["input_ids"] = [expanded_ids] + combined_outputs = dict( processed_outputs, **video_outputs, From e9de72fe6c56cfc7117768f671d2a1ff1f3bfb02 Mon Sep 17 00:00:00 2001 From: Tyler Michael Smith Date: Sat, 20 Jun 2026 15:26:38 -0400 Subject: [PATCH 45/75] [Bugfix] Guard model_config access in _log_compilation_config (#46198) Signed-off-by: Tyler Michael Smith Co-authored-by: Claude Opus 4.6 --- vllm/compilation/backends.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/compilation/backends.py b/vllm/compilation/backends.py index 5a67415f103..dc12acbaf4a 100644 --- a/vllm/compilation/backends.py +++ b/vllm/compilation/backends.py @@ -991,7 +991,7 @@ class VllmBackend: }, payload_fn=lambda: json.dumps( { - "model": self.vllm_config.model_config.model, + "model": getattr(self.vllm_config.model_config, "model", "unknown"), "prefix": self.prefix, "mode": str(cc.mode), "backend": cc.backend, From ebfbcfe46aa895d428933244908bae08b1ca6397 Mon Sep 17 00:00:00 2001 From: Tyler Michael Smith Date: Sat, 20 Jun 2026 16:38:10 -0400 Subject: [PATCH 46/75] Stop setting CUDA_VISIBLE_DEVICES internally in vLLM, add device_ids arg (#45026) Signed-off-by: Tyler Michael Smith Co-authored-by: Claude Co-authored-by: Codex Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: kourosh hakhamaneshi --- .../attention/mla/sm100_cutlass_mla_kernel.cu | 11 +- tests/engine/test_arg_utils.py | 193 ++++++++++++++++++ .../entrypoints/openai/test_dp_supervisor.py | 4 +- vllm/config/parallel.py | 9 + .../device_communicators/all2all.py | 9 +- .../device_communicators/all_reduce_utils.py | 33 ++- .../device_communicators/custom_all_reduce.py | 18 +- .../device_communicators/quick_all_reduce.py | 10 +- .../device_communicators/shm_broadcast.py | 8 +- .../v1/lmcache_integration/vllm_v1_adapter.py | 9 +- vllm/distributed/parallel_state.py | 28 ++- vllm/distributed/stateless_coordinator.py | 22 +- vllm/engine/arg_utils.py | 58 ++++++ vllm/entrypoints/openai/dp_supervisor.py | 37 ++-- vllm/platforms/cuda.py | 9 + vllm/platforms/interface.py | 103 +++++++++- vllm/v1/engine/core.py | 25 ++- vllm/v1/engine/utils.py | 109 ++++++---- vllm/v1/executor/multiproc_executor.py | 10 + vllm/v1/executor/ray_executor.py | 53 ++--- vllm/v1/executor/ray_executor_v2.py | 87 +++++--- vllm/v1/executor/ray_utils.py | 8 +- vllm/v1/worker/gpu_worker.py | 48 ++++- vllm/v1/worker/worker_base.py | 6 + 24 files changed, 722 insertions(+), 185 deletions(-) diff --git a/csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu b/csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu index 55d75383476..de62052b4b0 100644 --- a/csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu +++ b/csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu @@ -268,9 +268,14 @@ int64_t sm100_cutlass_mla_get_workspace_size(int64_t max_seq_len, int64_t num_ba using TileShapeD = typename MlaSm100Type::TileShapeD; arguments.problem_shape = cute::make_tuple(TileShapeH{}, static_cast(max_seq_len), TileShapeD{}, static_cast(num_batches)); - // Assumes device 0 when getting sm_count. - arguments.hw_info.sm_count = - sm_count <= 0 ? cutlass::KernelHardwareInfo::query_device_multiprocessor_count(/*device_id=*/0) : sm_count; + if (sm_count <= 0) { + int current_device = 0; + cudaGetDevice(¤t_device); + arguments.hw_info.sm_count = + cutlass::KernelHardwareInfo::query_device_multiprocessor_count(current_device); + } else { + arguments.hw_info.sm_count = sm_count; + } arguments.split_kv = static_cast(num_kv_splits); MlaSm100Type::Fmha::set_split_kv(arguments); diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 9d34975032e..a35f4453027 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -649,3 +649,196 @@ def test_cloud_storage_tokenizer_skips_get_model_path(monkeypatch): args = EngineArgs(model="s3://bucket/model", tokenizer="s3://bucket/tokenizer") assert args.model == "s3://bucket/model" assert args.tokenizer == "s3://bucket/tokenizer" + + +class TestDeviceIds: + def test_device_ids_with_cvd_out_of_range(self, monkeypatch): + """--device-ids index beyond the CVD set raises ValueError.""" + from vllm.platforms import current_platform + + key = current_platform.device_control_env_var + monkeypatch.setenv(key, "4,5") + args = EngineArgs(model="m", device_ids=[0, 2]) + with pytest.raises(ValueError, match="out of range"): + args._resolve_device_ids() + + def test_device_ids_with_cvd_resolve_to_physical_ids(self, monkeypatch): + """--device-ids are CVD-local indices resolved to physical ids.""" + from vllm.platforms import current_platform + + key = current_platform.device_control_env_var + monkeypatch.setenv(key, "4,5") + args = EngineArgs(model="m", device_ids=[0, 1]) + assert args._resolve_device_ids() == [4, 5] + + def test_device_ids_with_uuid_cvd_resolve_to_physical_ids(self, monkeypatch): + """--device-ids support UUID CVD values resolved by the platform.""" + from vllm.platforms import current_platform + + key = current_platform.device_control_env_var + monkeypatch.setenv(key, "GPU-abcd1234,GPU-ef567890") + monkeypatch.setattr( + type(current_platform), + "device_control_id_to_physical_device_id", + classmethod( + lambda cls, device_id: {"GPU-abcd1234": 4, "GPU-ef567890": 5}[device_id] + ), + ) + + args = EngineArgs(model="m", device_ids=[0, 1]) + assert args._resolve_device_ids() == [4, 5] + + def test_device_ids_with_uuid_args_resolve_to_physical_ids(self, monkeypatch): + """UUID --device-ids are resolved to physical IDs immediately.""" + from vllm.platforms import current_platform + + monkeypatch.setattr( + type(current_platform), + "device_control_id_to_physical_device_id", + classmethod(lambda cls, device_id: {"GPU-abcd1234": 4}[device_id]), + ) + + args = EngineArgs(model="m", device_ids=["GPU-abcd1234"]) + assert args._resolve_device_ids() == [4] + + def test_device_ids_reject_mixed_integer_and_uuid_args(self): + """--device-ids must not mix CVD indices and UUIDs.""" + args = EngineArgs(model="m", device_ids=[0, "GPU-abcd1234"]) + with pytest.raises(ValueError, match="must not mix"): + args._resolve_device_ids() + + def test_no_device_ids(self): + """No --device-ids returns None.""" + args = EngineArgs(model="m") + assert args._resolve_device_ids() is None + + def test_cli_parsing(self): + """--device-ids parses comma-separated string from CLI.""" + parser = FlexibleArgumentParser() + EngineArgs.add_cli_args(parser) + parsed = parser.parse_args(["--model", "m", "--device-ids", "0,2,4"]) + assert parsed.device_ids == [0, 2, 4] + + def test_cli_parsing_uuid(self): + """--device-ids parses comma-separated UUID strings from CLI.""" + parser = FlexibleArgumentParser() + EngineArgs.add_cli_args(parser) + parsed = parser.parse_args( + ["--model", "m", "--device-ids", "GPU-abcd1234,GPU-ef567890"] + ) + assert parsed.device_ids == ["GPU-abcd1234", "GPU-ef567890"] + + def test_assigned_physical_gpu_ids_are_physical_with_cvd(self, monkeypatch): + """assigned_physical_gpu_ids are already physical and not composed with CVD.""" + import vllm.platforms.interface as platform_interface + from vllm.platforms import current_platform + + monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [4, 5]) + monkeypatch.setenv(current_platform.device_control_env_var, "4,5") + + assert current_platform.device_id_to_physical_device_id(0) == 4 + assert current_platform.device_id_to_physical_device_id(1) == 5 + assert current_platform.logical_device_id_to_visible_device_id(0) == 0 + assert current_platform.logical_device_id_to_visible_device_id(1) == 1 + + def test_assigned_physical_gpu_ids_map_to_visible_uuid_cvd(self, monkeypatch): + """Physical IDs map back to visible ordinals when CVD uses UUIDs.""" + import vllm.platforms.interface as platform_interface + from vllm.platforms import current_platform + + monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [5]) + monkeypatch.setenv( + current_platform.device_control_env_var, + "GPU-abcd1234,GPU-ef567890", + ) + monkeypatch.setattr( + type(current_platform), + "device_control_id_to_physical_device_id", + classmethod( + lambda cls, device_id: {"GPU-abcd1234": 4, "GPU-ef567890": 5}[device_id] + ), + ) + + assert current_platform.logical_device_id_to_visible_device_id(0) == 1 + + def test_device_ids_reject_duplicates(self): + """--device-ids must not contain duplicate entries.""" + args = EngineArgs(model="m", device_ids=[2, 2]) + with pytest.raises(ValueError, match="duplicates"): + args._resolve_device_ids() + + def test_cli_parsing_strips_whitespace(self): + """--device-ids tolerates whitespace around commas.""" + parser = FlexibleArgumentParser() + EngineArgs.add_cli_args(parser) + parsed = parser.parse_args(["--model", "m", "--device-ids", "0, 2, 4"]) + assert parsed.device_ids == [0, 2, 4] + + def test_visible_ordinal_to_physical_ignores_assigned_ids(self, monkeypatch): + """visible_device_id_to_physical_device_id maps torch device ordinals, + independent of the logical-to-physical mapping. + + Regression test: CustomAllreduce passes device.index (a visible + ordinal) and must not index into assigned_physical_gpu_ids, which + raised IndexError for non-identity --device-ids like [2, 3]. + """ + import vllm.platforms.interface as platform_interface + from vllm.platforms import current_platform + + monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [2, 3]) + monkeypatch.delenv(current_platform.device_control_env_var, raising=False) + + # CVD unset: visible ordinal == physical ID, even beyond the + # assigned list's length. + assert current_platform.visible_device_id_to_physical_device_id(2) == 2 + assert current_platform.visible_device_id_to_physical_device_id(3) == 3 + + monkeypatch.setenv(current_platform.device_control_env_var, "4,5") + assert current_platform.visible_device_id_to_physical_device_id(1) == 5 + with pytest.raises(IndexError, match="out of range"): + current_platform.visible_device_id_to_physical_device_id(2) + + +class TestDpDeviceIdSharding: + def test_dp_supervisor_device_ids_stay_env_relative(self): + """Regression test: the DP supervisor must pass env-relative indices, + not physical IDs, because each child re-resolves --device-ids + against its inherited device-control env var.""" + import argparse + + from vllm.entrypoints.openai.dp_supervisor import _build_device_ids + + args = argparse.Namespace( + tensor_parallel_size=2, pipeline_parallel_size=1, device_ids=None + ) + assert _build_device_ids(args, local_rank=0) == [0, 1] + assert _build_device_ids(args, local_rank=1) == [2, 3] + + def test_dp_supervisor_shards_user_device_ids(self): + """User-provided --device-ids are sharded across DP children.""" + import argparse + + from vllm.entrypoints.openai.dp_supervisor import _build_device_ids + + args = argparse.Namespace( + tensor_parallel_size=2, pipeline_parallel_size=1, device_ids=[4, 5, 6, 7] + ) + assert _build_device_ids(args, local_rank=0) == [4, 5] + assert _build_device_ids(args, local_rank=1) == [6, 7] + with pytest.raises(ValueError, match="needs devices"): + _build_device_ids(args, local_rank=2) + + def test_dp_rank_shards_user_assigned_gpu_ids(self): + """get_physical_gpu_ids_for_local_dp_rank slices the user-provided + --device-ids list instead of recomputing from the env var.""" + from vllm.platforms import current_platform + from vllm.v1.engine.utils import get_physical_gpu_ids_for_local_dp_rank + + evar = current_platform.device_control_env_var + assert get_physical_gpu_ids_for_local_dp_rank( + evar, local_dp_rank=1, world_size=2, user_assigned_gpu_ids=[4, 5, 6, 7] + ) == [6, 7] + with pytest.raises(ValueError, match="needs devices"): + get_physical_gpu_ids_for_local_dp_rank( + evar, local_dp_rank=2, world_size=2, user_assigned_gpu_ids=[4, 5, 6, 7] + ) diff --git a/tests/entrypoints/openai/test_dp_supervisor.py b/tests/entrypoints/openai/test_dp_supervisor.py index 9967e6d86d0..1dd6537f201 100644 --- a/tests/entrypoints/openai/test_dp_supervisor.py +++ b/tests/entrypoints/openai/test_dp_supervisor.py @@ -364,7 +364,7 @@ class MockVLLMServer: await self._serve_task -def launch_mock_vllm(child_args: argparse.Namespace, env_updates: dict[str, str]): +def launch_mock_vllm(child_args: argparse.Namespace): logger.info("Launching mock vLLM on port %s", child_args.port) mock_vllm = MockVLLMServer( port=child_args.port, @@ -375,7 +375,7 @@ def launch_mock_vllm(child_args: argparse.Namespace, env_updates: dict[str, str] def launch_mock_vllm_with_drain( - child_args: argparse.Namespace, env_updates: dict[str, str] + child_args: argparse.Namespace, ): logger.info("Launching mock vLLM with 15s drain on port %s", child_args.port) mock_vllm = MockVLLMServer( diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index a194640f2ec..2ae773d79c7 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -302,6 +302,14 @@ class ParallelConfig: Each entry must use `numactl --physcpubind` CPU-list syntax, for example `"0-3"` or `"0,2,4-7"`. """ + assigned_physical_gpu_ids: list[int] | None = None + """Mapping from vLLM-local logical GPU IDs to physical GPU IDs. + + For example, ``[2, 3]`` means logical GPU 0 maps to physical GPU 2, + and logical GPU 1 maps to physical GPU 3. Physical IDs are used only + at platform/topology boundaries such as NVML, NIC affinity, P2P + checks, and final CUDA device selection when needed. When None, + logical IDs map to visible device IDs in order.""" distributed_timeout_seconds: int | None = None """Timeout in seconds for distributed operations (e.g., init_process_group). @@ -772,6 +780,7 @@ class ParallelConfig: "numa_bind", "numa_bind_nodes", "numa_bind_cpus", + "assigned_physical_gpu_ids", } from vllm.config.utils import get_hash_factors, hash_factors diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index 967ce5d75c3..0066a60dd02 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -704,7 +704,14 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase): self.num_experts = num_experts self.cleanup() - gpus_per_node = torch.accelerator.device_count() + from vllm.platforms.interface import get_assigned_physical_gpu_ids + + assigned_physical_gpu_ids = get_assigned_physical_gpu_ids() + gpus_per_node = ( + len(assigned_physical_gpu_ids) + if assigned_physical_gpu_ids is not None + else torch.accelerator.device_count() + ) logger.debug( "Making One-sided NVLink mapping: rank=%d, world size=%d", self.rank, diff --git a/vllm/distributed/device_communicators/all_reduce_utils.py b/vllm/distributed/device_communicators/all_reduce_utils.py index cebf2c49b44..d50d84fa5ca 100644 --- a/vllm/distributed/device_communicators/all_reduce_utils.py +++ b/vllm/distributed/device_communicators/all_reduce_utils.py @@ -320,13 +320,21 @@ def gpu_p2p_access_check(src: int, tgt: int) -> bool: is_distributed = dist.is_initialized() - num_dev = current_platform.device_count() - cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES - if cuda_visible_devices is None: - cuda_visible_devices = ",".join(str(i) for i in range(num_dev)) + from vllm.platforms.interface import get_assigned_physical_gpu_ids + + assigned_physical_gpu_ids = get_assigned_physical_gpu_ids() + if assigned_physical_gpu_ids is not None: + # Key by the ordered list: the cache stores directed local-index + # pairs, so permutations of the same set are distinct mappings. + cache_key = ",".join(str(i) for i in assigned_physical_gpu_ids) + num_dev = len(assigned_physical_gpu_ids) + else: + num_dev = current_platform.device_count() + cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES + cache_key = cuda_visible_devices or ",".join(str(i) for i in range(num_dev)) path = os.path.join( - envs.VLLM_CACHE_ROOT, f"gpu_p2p_access_cache_for_{cuda_visible_devices}.json" + envs.VLLM_CACHE_ROOT, f"gpu_p2p_access_cache_for_{cache_key}.json" ) os.makedirs(os.path.dirname(path), exist_ok=True) from vllm.distributed.parallel_state import get_world_group @@ -338,7 +346,15 @@ def gpu_p2p_access_check(src: int, tgt: int) -> bool: # enter this block to calculate the cache logger.info("generating GPU P2P access cache in %s", path) cache: dict[str, bool] = {} - ids = list(range(num_dev)) + # The probe subprocesses inherit this process's device-control env + # var, so they must be given visible ordinals, not physical IDs. + if assigned_physical_gpu_ids is not None: + ids = [ + current_platform.logical_device_id_to_visible_device_id(local) + for local in range(num_dev) + ] + else: + ids = list(range(num_dev)) # batch of all pairs of GPUs batch_src, batch_tgt = zip(*list(product(ids, ids))) # NOTE: we use `subprocess` rather than `multiprocessing` here @@ -368,8 +384,11 @@ def gpu_p2p_access_check(src: int, tgt: int) -> bool: ) from e with open(output_file.name, "rb") as f: result = pickle.load(f) + # Cache entries must be keyed by local indices (0..N-1) because + # gpu_p2p_access_check() is called with local ranks. + id_to_local = {device_id: local for local, device_id in enumerate(ids)} for _i, _j, r in zip(batch_src, batch_tgt, result): - cache[f"{_i}->{_j}"] = r + cache[f"{id_to_local[_i]}->{id_to_local[_j]}"] = r with open(path, "w") as f: json.dump(cache, f, indent=4) if is_distributed: diff --git a/vllm/distributed/device_communicators/custom_all_reduce.py b/vllm/distributed/device_communicators/custom_all_reduce.py index c57cc74fc06..95db6cc9245 100644 --- a/vllm/distributed/device_communicators/custom_all_reduce.py +++ b/vllm/distributed/device_communicators/custom_all_reduce.py @@ -34,7 +34,12 @@ def _can_p2p(rank: int, world_size: int) -> bool: continue if envs.VLLM_SKIP_P2P_CHECK: logger.debug("Skipping P2P check and trusting the driver's P2P report.") - return torch.cuda.can_device_access_peer(rank, i) + # can_device_access_peer takes visible device ordinals, while + # rank and i are logical local IDs. + return torch.cuda.can_device_access_peer( + current_platform.logical_device_id_to_visible_device_id(rank), + current_platform.logical_device_id_to_visible_device_id(i), + ) if not gpu_p2p_access_check(rank, i): return False return True @@ -126,13 +131,10 @@ class CustomAllreduce: CUSTOM_ALL_REDUCE_MAX_SIZES[device_capability_str][world_size], max_size, ) - cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES - if cuda_visible_devices: - device_ids = list(map(int, cuda_visible_devices.split(","))) - else: - device_ids = list(range(current_platform.device_count())) - - physical_device_id = device_ids[device.index] + # device.index is a visible ordinal, not a logical local ID. + physical_device_id = current_platform.visible_device_id_to_physical_device_id( + device.index + ) tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu") gather_list = [ torch.tensor([0], dtype=torch.int, device="cpu") for _ in range(world_size) diff --git a/vllm/distributed/device_communicators/quick_all_reduce.py b/vllm/distributed/device_communicators/quick_all_reduce.py index 8c7ee7452f1..c54eaf7555d 100644 --- a/vllm/distributed/device_communicators/quick_all_reduce.py +++ b/vllm/distributed/device_communicators/quick_all_reduce.py @@ -129,12 +129,10 @@ class QuickAllReduce: assert isinstance(device, torch.device) self.device = device - cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES - if cuda_visible_devices: - device_ids = list(map(int, cuda_visible_devices.split(","))) - else: - device_ids = list(range(current_platform.device_count())) - physical_device_id = device_ids[device.index] + # device.index is a visible ordinal, not a logical local ID. + physical_device_id = current_platform.visible_device_id_to_physical_device_id( + device.index + ) tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu") gather_list = [ torch.tensor([0], dtype=torch.int, device="cpu") diff --git a/vllm/distributed/device_communicators/shm_broadcast.py b/vllm/distributed/device_communicators/shm_broadcast.py index 9482568461c..43e066c44b0 100644 --- a/vllm/distributed/device_communicators/shm_broadcast.py +++ b/vllm/distributed/device_communicators/shm_broadcast.py @@ -840,7 +840,13 @@ class MessageQueue: The MessageQueue instance for the calling process, and a list of handles (only non-empty for the reader process). """ - local_size = current_platform.device_count() + from vllm.platforms.interface import get_assigned_physical_gpu_ids + + assigned_physical_gpu_ids = get_assigned_physical_gpu_ids() + if assigned_physical_gpu_ids is not None: + local_size = len(assigned_physical_gpu_ids) + else: + local_size = current_platform.device_count() rank = dist.get_rank() same_node = rank // local_size == reader_rank // local_size buffer_io = MessageQueue( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/vllm_v1_adapter.py b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/vllm_v1_adapter.py index d16fbee585a..d72ebb5cd1e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/vllm_v1_adapter.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/vllm_v1_adapter.py @@ -482,10 +482,11 @@ def _init_lmcache_engine( ) # Change current device. - num_gpus = torch.accelerator.device_count() - local_rank = parallel_config.rank % num_gpus - torch.accelerator.set_device_index(local_rank) - device = torch.device(f"cuda:{local_rank}") + from vllm.distributed.parallel_state import get_world_group + + device_index = get_world_group().device_index + torch.accelerator.set_device_index(device_index) + device = torch.device(f"cuda:{device_index}") metadata = LMCacheEngineMetadata( model_config.model, parallel_config.world_size, diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index 8bd6e92157a..11b9e24e864 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -392,6 +392,14 @@ class GroupCoordinator: self.rank = torch.distributed.get_rank() self.local_rank = local_rank + self.device_index: int + if _WORLD is not None: + self.device_index = _WORLD.device_index + else: + assert local_rank >= 0, ( + "local_rank must be provided when creating the world group" + ) + self.device_index = local_rank self_device_group = None self_cpu_group = None @@ -442,11 +450,18 @@ class GroupCoordinator: from vllm.platforms import current_platform if current_platform.is_cuda_alike(): - self.device = torch.device(f"cuda:{local_rank}") + visible_device_index = ( + current_platform.logical_device_id_to_visible_device_id( + self.device_index + ) + ) + self.device = torch.device(f"cuda:{visible_device_index}") elif current_platform.is_xpu(): - self.device = torch.device(f"xpu:{local_rank}") + self.device = torch.device(f"xpu:{self.device_index}") elif current_platform.is_out_of_tree(): - self.device = torch.device(f"{current_platform.device_name}:{local_rank}") + self.device = torch.device( + f"{current_platform.device_name}:{self.device_index}" + ) else: self.device = torch.device("cpu") @@ -1438,7 +1453,12 @@ def _init_process_group_for_split_group( """ if torch.accelerator.is_available() and backend != "gloo": init_backend = "cpu:gloo,cuda:nccl" - device_id: torch.device | None = torch.device(f"cuda:{local_rank}") + from vllm.platforms import current_platform + + visible_device_index = current_platform.logical_device_id_to_visible_device_id( + local_rank + ) + device_id: torch.device | None = torch.device(f"cuda:{visible_device_index}") else: init_backend = "gloo" device_id = None diff --git a/vllm/distributed/stateless_coordinator.py b/vllm/distributed/stateless_coordinator.py index 549284df32d..5f4597d07cb 100644 --- a/vllm/distributed/stateless_coordinator.py +++ b/vllm/distributed/stateless_coordinator.py @@ -86,6 +86,15 @@ class StatelessGroupCoordinator(GroupCoordinator): self.rank = global_rank self.local_rank = local_rank + from vllm.distributed.parallel_state import _WORLD + + if _WORLD is not None: + self.device_index = _WORLD.device_index + else: + assert local_rank >= 0, ( + "local_rank must be provided when creating the world group" + ) + self.device_index = local_rank self_device_group = None self_cpu_group = None @@ -152,11 +161,18 @@ class StatelessGroupCoordinator(GroupCoordinator): self.tcp_store_group = self_tcp_store_group if current_platform.is_cuda_alike(): - self.device = torch.device(f"cuda:{local_rank}") + visible_device_index = ( + current_platform.logical_device_id_to_visible_device_id( + self.device_index + ) + ) + self.device = torch.device(f"cuda:{visible_device_index}") elif current_platform.is_xpu(): - self.device = torch.device(f"xpu:{local_rank}") + self.device = torch.device(f"xpu:{self.device_index}") elif current_platform.is_out_of_tree(): - self.device = torch.device(f"{current_platform.device_name}:{local_rank}") + self.device = torch.device( + f"{current_platform.device_name}:{self.device_index}" + ) else: self.device = torch.device("cpu") diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 9172a8728a0..921f31466b3 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -6,6 +6,7 @@ import copy import dataclasses import functools import json +import os import sys from collections.abc import Callable from dataclasses import MISSING, asdict, dataclass, fields, is_dataclass @@ -465,6 +466,7 @@ class EngineArgs: numa_bind: bool = ParallelConfig.numa_bind numa_bind_nodes: list[int] | None = ParallelConfig.numa_bind_nodes numa_bind_cpus: list[str] | None = ParallelConfig.numa_bind_cpus + device_ids: list[int | str] | None = None tensor_parallel_size: int = ParallelConfig.tensor_parallel_size prefill_context_parallel_size: int = ParallelConfig.prefill_context_parallel_size decode_context_parallel_size: int = ParallelConfig.decode_context_parallel_size @@ -979,6 +981,20 @@ class EngineArgs: parallel_group.add_argument( "--numa-bind-cpus", **parallel_kwargs["numa_bind_cpus"] ) + parallel_group.add_argument( + "--device-ids", + type=lambda s: [ + int(device_id) if device_id.isdigit() else device_id + for device_id in (part.strip() for part in s.split(",")) + ], + default=None, + help="Comma-separated physical GPU device IDs or UUIDs to use " + '(e.g. --device-ids "2,3,5,7"). Avoids setting ' + "CUDA_VISIBLE_DEVICES, preserving full GPU topology " + "visibility for GPU-NIC affinity and DeepGEMM. " + "Note: has no effect with Ray executors; use Ray " + "placement groups for GPU selection instead.", + ) parallel_group.add_argument( "--tensor-parallel-size", "-tp", **parallel_kwargs["tensor_parallel_size"] ) @@ -1716,6 +1732,47 @@ class EngineArgs: ) return SpeculativeConfig(**self.speculative_config) + def _resolve_device_ids(self) -> list[int] | None: + if not self.device_ids: + return None + if self.distributed_executor_backend == "ray": + logger.warning( + "--device-ids has no effect when using the Ray executor. " + "Use Ray placement groups for GPU selection instead." + ) + ids = self.device_ids + if len(set(ids)) != len(ids): + raise ValueError(f"--device-ids must not contain duplicates: {ids}") + if all(isinstance(i, str) for i in ids): + return [ + current_platform.device_control_id_to_physical_device_id(i) + for i in cast(list[str], ids) + ] + if any(isinstance(i, str) for i in ids): + raise ValueError("--device-ids must not mix integer IDs and UUIDs") + int_ids = cast(list[int], ids) + # Compose with CUDA_VISIBLE_DEVICES: if CVD is set, treat + # --device-ids values as indices into the CVD-visible set. + cvd = getattr( + envs, + current_platform.device_control_env_var, + os.environ.get(current_platform.device_control_env_var), + ) + if cvd: + cvd_ids = [ + current_platform.device_control_id_to_physical_device_id(x) + for x in cvd.split(",") + ] + for i in int_ids: + if i >= len(cvd_ids): + raise ValueError( + f"--device-ids index {i} is out of range for " + f"{current_platform.device_control_env_var}" + f"={cvd} ({len(cvd_ids)} devices visible)" + ) + return [cvd_ids[i] for i in int_ids] + return int_ids + def create_diffusion_config(self) -> DiffusionConfig | None: if self.diffusion_config is None: return None @@ -2029,6 +2086,7 @@ class EngineArgs: cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size, _api_process_count=self._api_process_count, _api_process_rank=self._api_process_rank, + assigned_physical_gpu_ids=self._resolve_device_ids(), numa_bind=self.numa_bind, numa_bind_nodes=self.numa_bind_nodes, numa_bind_cpus=self.numa_bind_cpus, diff --git a/vllm/entrypoints/openai/dp_supervisor.py b/vllm/entrypoints/openai/dp_supervisor.py index 13444015ecc..73b10a04ea5 100644 --- a/vllm/entrypoints/openai/dp_supervisor.py +++ b/vllm/entrypoints/openai/dp_supervisor.py @@ -23,12 +23,10 @@ import uvloop from fastapi import FastAPI, Response from vllm.logger import init_logger -from vllm.platforms import current_platform from vllm.utils.system_utils import ( decorate_logs, kill_process_tree, set_process_title, - update_environment_variables, ) logger = init_logger(__name__) @@ -127,22 +125,29 @@ def _build_vllm_dp_server_args( child_args.data_parallel_multi_port_external_lb = False child_args.data_parallel_supervisor_port = None child_args.api_server_count = 1 + child_args.device_ids = _build_device_ids(args, local_rank) return child_args -def _build_vllm_dp_server_env( - args: argparse.Namespace, local_rank: int -) -> dict[str, str]: - # set visible devices for the child process +def _build_device_ids(args: argparse.Namespace, local_rank: int) -> list[int | str]: + """Build the --device-ids value for a DP child process. + + The child resolves these against its own inherited device-control env + var (e.g. CUDA_VISIBLE_DEVICES), so integer IDs must stay env-relative + here rather than being translated to physical IDs. + """ devices_per_rank = args.tensor_parallel_size * args.pipeline_parallel_size start = local_rank * devices_per_rank stop = start + devices_per_rank - device_env = current_platform.device_control_env_var - visible_devices = ",".join( - str(current_platform.device_id_to_physical_device_id(idx)) - for idx in range(start, stop) - ) - return {device_env: visible_devices} + device_ids = getattr(args, "device_ids", None) + if device_ids is not None: + if stop > len(device_ids): + raise ValueError( + f"--device-ids has {len(device_ids)} entries, but DP rank " + f"{local_rank} needs devices [{start}, {stop})" + ) + return device_ids[start:stop] + return list(range(start, stop)) def _child_base_url(args: argparse.Namespace, port: int) -> str: @@ -228,9 +233,7 @@ def _build_dp_supervisor_app(supervisor: DPSupervisor) -> FastAPI: return app -def _run_vllm_dp_server( - child_args: argparse.Namespace, env_updates: dict[str, str] -) -> None: +def _run_vllm_dp_server(child_args: argparse.Namespace) -> None: """ Entrypoint function for the vLLM DP Server. """ @@ -241,7 +244,6 @@ def _run_vllm_dp_server( os.setpgrp() name = f"APIServer_DP{child_args.data_parallel_rank}" - update_environment_variables(env_updates) set_process_title(name) decorate_logs(name) uvloop.run(run_server(child_args)) @@ -345,11 +347,10 @@ class DPSupervisor: context = multiprocessing.get_context("spawn") for local_rank in range(self.args.data_parallel_size_local): child_args = _build_vllm_dp_server_args(self.args, local_rank) - child_env = _build_vllm_dp_server_env(self.args, local_rank) process = context.Process( target=_run_vllm_dp_server, name=f"APIServer_DPRank_{child_args.data_parallel_rank}", - args=(child_args, child_env), + args=(child_args,), ) process.start() self._processes.append(process) diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 30a16e27469..6bf1793eefd 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -685,6 +685,15 @@ class CudaPlatformBase(Platform): # all the related functions work on real physical device ids. # the major benefit of using NVML is that it will not initialize CUDA class NvmlCudaPlatform(CudaPlatformBase): + @classmethod + @with_nvml_context + def device_control_id_to_physical_device_id(cls, device_id: str) -> int: + try: + return int(device_id) + except ValueError: + handle = pynvml.nvmlDeviceGetHandleByUUID(device_id) + return pynvml.nvmlDeviceGetIndex(handle) + @classmethod @cache @with_nvml_context diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index 7fed06950bd..82c87416093 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -30,6 +30,33 @@ else: logger = init_logger(__name__) +_assigned_physical_gpu_ids: list[int] | None = None + + +def set_assigned_physical_gpu_ids(ids: list[int]) -> None: + """Set the physical GPU IDs assigned to this worker process. + Called during worker init so that device_id_to_physical_device_id() + can map local_rank to the correct physical device without relying + on CUDA_VISIBLE_DEVICES. + + Idempotent: a second call with the same value is a no-op. + Raises RuntimeError if called again with a different value. + + This is expected to run during single-threaded worker initialization.""" + global _assigned_physical_gpu_ids + if _assigned_physical_gpu_ids is not None: + if _assigned_physical_gpu_ids != ids: + raise RuntimeError( + f"set_assigned_physical_gpu_ids called with conflicting values: " + f"existing={_assigned_physical_gpu_ids}, new={ids}" + ) + return + _assigned_physical_gpu_ids = ids + + +def get_assigned_physical_gpu_ids() -> list[int] | None: + return _assigned_physical_gpu_ids + @functools.cache def in_wsl() -> bool: @@ -233,8 +260,34 @@ class Platform: """ import vllm.kernels # noqa: F401 + @classmethod + def device_control_id_to_physical_device_id(cls, device_id: str) -> int: + """Map one device-control env entry to an integer physical device ID.""" + try: + return int(device_id) + except ValueError as e: + raise ValueError( + f"Non-integer device ID {device_id!r} is not supported by " + f"{cls.device_name}." + ) from e + @classmethod def device_id_to_physical_device_id(cls, device_id: int): + """Map a vLLM-local logical device ID to a physical device ID. + + The input is a logical local ID (e.g. a local rank), NOT a visible + device ordinal; for the latter use + visible_device_id_to_physical_device_id(). The two coincide only + when no logical-to-physical mapping is in effect. + """ + if _assigned_physical_gpu_ids is not None: + if device_id >= len(_assigned_physical_gpu_ids): + raise IndexError( + f"device_id {device_id} is out of range for " + f"assigned_physical_gpu_ids {_assigned_physical_gpu_ids} " + f"({len(_assigned_physical_gpu_ids)} devices assigned)" + ) + return _assigned_physical_gpu_ids[device_id] # Treat empty device control env var as unset. This is a valid # configuration in Ray setups where the engine is launched in # a CPU-only placement group located on a GPU node. @@ -244,10 +297,58 @@ class Platform: ): device_ids = os.environ[cls.device_control_env_var].split(",") physical_device_id = device_ids[device_id] - return int(physical_device_id) + return cls.device_control_id_to_physical_device_id(physical_device_id) else: return device_id + @classmethod + def logical_device_id_to_visible_device_id(cls, device_id: int) -> int: + """Map a vLLM-local logical device ID to the current process's + visible accelerator ordinal. + + vLLM internals use logical local IDs. Physical IDs are used only + at platform/topology boundaries. This helper performs the final + translation needed by APIs such as ``torch.device("cuda:N")``. + """ + physical_device_id = cls.device_id_to_physical_device_id(device_id) + device_control_env = os.environ.get(cls.device_control_env_var, "") + if not device_control_env: + return physical_device_id + + visible_physical_device_ids = [ + cls.device_control_id_to_physical_device_id(physical_id) + for physical_id in device_control_env.split(",") + ] + if physical_device_id not in visible_physical_device_ids: + raise RuntimeError( + f"Physical device {physical_device_id} for logical device " + f"{device_id} is not visible in {cls.device_control_env_var}=" + f"{device_control_env}" + ) + return visible_physical_device_ids.index(physical_device_id) + + @classmethod + def visible_device_id_to_physical_device_id(cls, device_id: int) -> int: + """Map a visible accelerator ordinal (e.g. ``torch.device.index``) + to a physical device ID. + + This is the inverse of the env-var translation performed by + logical_device_id_to_visible_device_id() and is independent of any + logical-to-physical mapping set via set_assigned_physical_gpu_ids(). + """ + device_control_env = os.environ.get(cls.device_control_env_var, "") + if not device_control_env: + return device_id + visible_device_ids = device_control_env.split(",") + if device_id >= len(visible_device_ids): + raise IndexError( + f"visible device ordinal {device_id} is out of range for " + f"{cls.device_control_env_var}={device_control_env}" + ) + return cls.device_control_id_to_physical_device_id( + visible_device_ids[device_id] + ) + @classmethod def import_kernels(cls) -> None: """Import any platform-specific C kernels.""" diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index ac7037800a0..8f6baa46936 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -74,7 +74,7 @@ from vllm.v1.engine.utils import ( EngineHandshakeMetadata, EngineZmqAddresses, SignalCallback, - get_device_indices, + get_physical_gpu_ids_for_local_dp_rank, ) from vllm.v1.executor import Executor from vllm.v1.kv_cache_interface import KVCacheConfig, get_kv_cache_spec_kind @@ -2175,23 +2175,30 @@ class EngineCoreActorMixin: pass else: device_control_env_var = current_platform.device_control_env_var - self._set_cuda_visible_devices( + self._set_assigned_physical_gpu_ids( vllm_config, local_dp_rank, device_control_env_var ) - def _set_cuda_visible_devices( - self, vllm_config: VllmConfig, local_dp_rank: int, device_control_env_var: str + def _set_assigned_physical_gpu_ids( + self, + vllm_config: VllmConfig, + local_dp_rank: int, + device_control_env_var: str, ): world_size = vllm_config.parallel_config.world_size - # Set CUDA_VISIBLE_DEVICES or equivalent. try: - value = get_device_indices( - device_control_env_var, local_dp_rank, world_size + physical_gpu_ids = get_physical_gpu_ids_for_local_dp_rank( + device_control_env_var, + local_dp_rank, + world_size, + user_assigned_gpu_ids=( + vllm_config.parallel_config.assigned_physical_gpu_ids + ), ) - os.environ[device_control_env_var] = value + vllm_config.parallel_config.assigned_physical_gpu_ids = physical_gpu_ids except IndexError as e: raise Exception( - f"Error setting {device_control_env_var}: " + f"Error computing assigned_physical_gpu_ids: " f"local range: [{local_dp_rank * world_size}, " f"{(local_dp_rank + 1) * world_size}) " f'base value: "{os.getenv(device_control_env_var)}"' diff --git a/vllm/v1/engine/utils.py b/vllm/v1/engine/utils.py index e13301f03c6..093f065475a 100644 --- a/vllm/v1/engine/utils.py +++ b/vllm/v1/engine/utils.py @@ -12,7 +12,6 @@ from multiprocessing import Process, connection from multiprocessing.process import BaseProcess from multiprocessing.queues import Queue from typing import TYPE_CHECKING, cast -from unittest.mock import patch import msgspec import zmq @@ -175,38 +174,38 @@ class CoreEngineProcManager: self.manager_stopped = threading.Event() self.failed_proc_name: str | None = None + # All ranks share this config object: capture the user-provided + # --device-ids list before the per-rank shard overwrites it. Mutating + # the config before each proc.start() works because the spawn method + # pickles process args at start() time, sequentially per rank. + user_assigned_gpu_ids = vllm_config.parallel_config.assigned_physical_gpu_ids try: for proc, local_dp_rank in zip(self.processes, local_dp_ranks): - # Adjust device control in DP for platforms that cannot rely - # on torch.accelerator.set_device_index(), and for Ray launchers. - device_control_context: contextlib.AbstractContextManager[None] = ( - contextlib.nullcontext() - ) + # Populate the logical-to-physical GPU mapping in DP for + # platforms that cannot rely on + # torch.accelerator.set_device_index(), and for Ray. needs_device_env_isolation = not ( current_platform.is_cuda_alike() or current_platform.is_xpu() ) if is_dp and ( needs_device_env_isolation or vllm_config.parallel_config.use_ray ): - device_control_context = set_device_control_env_var( - vllm_config, local_dp_rank + set_assigned_physical_gpu_ids_for_dp_rank( + vllm_config, local_dp_rank, user_assigned_gpu_ids ) - with ( - device_control_context, - numa_utils.configure_subprocess( - # EngineCore itself does not have a TP/PP-local rank. - # When DP is enabled, set_device_control_env_var() - # narrows visible devices to this DP shard first, so - # local_rank=0 means "the first local GPU in this - # shard". The actual TP/PP worker processes spawned by - # the executor are bound separately with their own - # local_rank values. - vllm_config, - local_rank=0, - dp_local_rank=local_dp_rank, - process_kind="EngineCore", - ), + with numa_utils.configure_subprocess( + # EngineCore itself does not have a TP/PP-local rank. + # When DP is enabled, set_assigned_physical_gpu_ids_for_dp_rank() + # populates the logical-to-physical mapping for this DP + # shard, so local_rank=0 means "the first local GPU in + # this shard". The actual TP/PP worker processes spawned + # by the executor are bound separately with their own + # local_rank values. + vllm_config, + local_rank=0, + dp_local_rank=local_dp_rank, + process_kind="EngineCore", ): proc.start() finally: @@ -281,55 +280,79 @@ class SignalCallback: self._event.set() -@contextlib.contextmanager -def set_device_control_env_var( - vllm_config: VllmConfig, local_dp_rank: int -) -> Iterator[None]: +def set_assigned_physical_gpu_ids_for_dp_rank( + vllm_config: VllmConfig, + local_dp_rank: int, + user_assigned_gpu_ids: list[int] | None = None, +) -> None: """ - Temporarily set CUDA_VISIBLE_DEVICES or equivalent - for engine subprocess. + Populate assigned_physical_gpu_ids on the config for the given DP rank. + + user_assigned_gpu_ids is the full (un-sharded) --device-ids list, if the + user provided one; this DP rank's shard is sliced from it. It is passed + explicitly rather than read from the config because callers may reuse + one config object across DP ranks, overwriting the field each time. """ world_size = vllm_config.parallel_config.world_size local_world_size = vllm_config.parallel_config.local_world_size evar = current_platform.device_control_env_var - value = get_device_indices(evar, local_dp_rank, world_size, local_world_size) - with patch.dict(os.environ, values=((evar, value),)): - yield + physical_gpu_ids = get_physical_gpu_ids_for_local_dp_rank( + evar, + local_dp_rank, + world_size, + local_world_size, + user_assigned_gpu_ids=user_assigned_gpu_ids, + ) + vllm_config.parallel_config.assigned_physical_gpu_ids = physical_gpu_ids -def get_device_indices( +def get_physical_gpu_ids_for_local_dp_rank( device_control_env_var: str, local_dp_rank: int, world_size: int, local_world_size: int | None = None, -): + user_assigned_gpu_ids: list[int] | None = None, +) -> list[int]: """ - Returns a comma-separated string of device indices for the specified + Returns list of physical GPU IDs for the specified data parallel rank. For example, if world_size=2 and local_dp_rank=1, and there are 4 devices, - this will select devices 2 and 3 for local_dp_rank=1. + this will return [2, 3] for local_dp_rank=1. + + If user_assigned_gpu_ids is provided (e.g. from --device-ids), this DP + rank's shard is sliced from it instead of being derived from the + device-control env var. """ if local_world_size is None: local_world_size = world_size + if user_assigned_gpu_ids is not None: + start = local_dp_rank * world_size + stop = start + local_world_size + if stop > len(user_assigned_gpu_ids): + raise ValueError( + f"--device-ids provides {len(user_assigned_gpu_ids)} devices, " + f"but DP rank {local_dp_rank} needs devices [{start}, {stop})" + ) + return user_assigned_gpu_ids[start:stop] try: - value = ",".join( - str(current_platform.device_id_to_physical_device_id(i)) + return [ + current_platform.device_id_to_physical_device_id(i) for i in range( local_dp_rank * world_size, local_dp_rank * world_size + local_world_size, ) - ) + ] except IndexError as e: raise Exception( - f"Error setting {device_control_env_var}: " + f"Error computing device indices for " + f"{device_control_env_var}: " f"local range: [{local_dp_rank * world_size}, " f"{(local_dp_rank + 1) * world_size}) " "base value: " f'"{os.getenv(device_control_env_var)}"' ) from e - return value def _apply_dp_identity_suffix(dp_vllm_config, dp_rank: int) -> None: @@ -453,11 +476,11 @@ class CoreEngineActorManager: # https://github.com/ray-project/ray/blob/master/python/ray/_private/accelerators/intel_gpu.py#L56 # noqa: E501 if current_platform.is_xpu(): device_evar = current_platform.device_control_env_var - device_indices = get_device_indices( + physical_gpu_ids = get_physical_gpu_ids_for_local_dp_rank( device_evar, local_index, world_size ) actor_env_vars = self.env_vars_dict.copy() - actor_env_vars[device_evar] = device_indices + actor_env_vars[device_evar] = ",".join(str(d) for d in physical_gpu_ids) runtime_env = RuntimeEnv(env_vars=actor_env_vars) actor = ( diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py index 7bc81118e6b..9b7581311e8 100644 --- a/vllm/v1/executor/multiproc_executor.py +++ b/vllm/v1/executor/multiproc_executor.py @@ -826,6 +826,16 @@ class WorkerProc: signal.signal(signal.SIGTERM, signal_handler) signal.signal(signal.SIGINT, signal_handler) + # Publish the logical-to-physical mapping early so topology helpers + # work before init_device (needed by set_worker_net_device below). + assigned_physical_gpu_ids = kwargs[ + "vllm_config" + ].parallel_config.assigned_physical_gpu_ids + if assigned_physical_gpu_ids is not None: + from vllm.platforms.interface import set_assigned_physical_gpu_ids + + set_assigned_physical_gpu_ids(assigned_physical_gpu_ids) + # Set net device env vars for the worker if VLLM_GPU_NIC_PCIE_MAPPING is set set_worker_net_device(kwargs.get("local_rank", 0), kwargs["vllm_config"]) diff --git a/vllm/v1/executor/ray_executor.py b/vllm/v1/executor/ray_executor.py index 749e59e04c2..39749ffc257 100644 --- a/vllm/v1/executor/ray_executor.py +++ b/vllm/v1/executor/ray_executor.py @@ -258,30 +258,35 @@ class RayDistributedExecutor(Executor): } self.collective_rpc("adjust_rank", args=(rerank_mapping,)) - # Get the set of GPU IDs used on each node. - worker_node_and_gpu_ids = [] + # Get the set of physical GPU IDs used on each node. + worker_node_and_physical_gpu_ids = [] for worker in [self.driver_dummy_worker] + self.workers: if worker is None: # driver_dummy_worker can be None when using ray spmd worker. continue - worker_node_and_gpu_ids.append( - ray.get(worker.get_node_and_gpu_ids.remote()) # type: ignore[attr-defined] + worker_node_and_physical_gpu_ids.append( + ray.get(worker.get_node_and_physical_gpu_ids.remote()) # type: ignore[attr-defined] ) node_workers = defaultdict(list) # node id -> list of worker ranks - node_gpus = defaultdict(list) # node id -> list of gpu ids + node_physical_gpu_ids = defaultdict(list) # node id -> physical GPU IDs - for i, (node_id, gpu_ids) in enumerate(worker_node_and_gpu_ids): + for i, (node_id, physical_gpu_ids) in enumerate( + worker_node_and_physical_gpu_ids + ): node_workers[node_id].append(i) - # `gpu_ids` can be a list of strings or integers. + # `physical_gpu_ids` can be a list of strings or integers. # convert them to integers for consistency. - # NOTE: gpu_ids can be larger than 9 (e.g. 16 GPUs), + # NOTE: physical GPU IDs can be larger than 9 (e.g. 16 GPUs), # string sorting is not sufficient. # see https://github.com/vllm-project/vllm/issues/5590 - gpu_ids = [int(x) for x in gpu_ids] - node_gpus[node_id].extend(gpu_ids) - for node_id, gpu_ids in node_gpus.items(): - node_gpus[node_id] = sorted(gpu_ids) + physical_gpu_ids = [ + current_platform.device_control_id_to_physical_device_id(str(x)) + for x in physical_gpu_ids + ] + node_physical_gpu_ids[node_id].extend(physical_gpu_ids) + for node_id, physical_gpu_ids in node_physical_gpu_ids.items(): + node_physical_gpu_ids[node_id] = sorted(physical_gpu_ids) all_ips = set(worker_ips + [driver_ip]) n_ips = len(all_ips) @@ -297,23 +302,8 @@ class RayDistributedExecutor(Executor): " each node." ) - # 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( - map(str, node_gpus[node_id]) - ), - } - for (node_id, _) in worker_node_and_gpu_ids + all_args_to_update_environment_variables: list[dict[str, str]] = [ + {} for _ in worker_node_and_physical_gpu_ids ] # Environment variables to copy from driver to workers @@ -336,7 +326,7 @@ class RayDistributedExecutor(Executor): "update_environment_variables", args=(self._get_env_vars_to_be_updated(),) ) - if len(node_gpus) == 1: + if len(node_physical_gpu_ids) == 1: # in single node case, we don't need to get the IP address. # the loopback address is sufficient # NOTE: a node may have several IP addresses, one for each @@ -352,10 +342,11 @@ class RayDistributedExecutor(Executor): # Initialize the actual workers inside worker wrapper. all_kwargs = [] - for rank, (node_id, _) in enumerate(worker_node_and_gpu_ids): + for rank, (node_id, _) in enumerate(worker_node_and_physical_gpu_ids): local_rank = node_workers[node_id].index(rank) kwargs = dict( vllm_config=self.vllm_config, + assigned_physical_gpu_ids=sorted(node_physical_gpu_ids[node_id]), local_rank=local_rank, rank=rank, distributed_init_method=distributed_init_method, diff --git a/vllm/v1/executor/ray_executor_v2.py b/vllm/v1/executor/ray_executor_v2.py index 0665b5fc1b8..d50f06cc620 100644 --- a/vllm/v1/executor/ray_executor_v2.py +++ b/vllm/v1/executor/ray_executor_v2.py @@ -79,24 +79,25 @@ class RayWorkerProc(WorkerProc): 1. __init__: lightweight setup, stores init args (no device/model init) 2. initialize_worker: called after GPU IDs are discovered, completes the full WorkerProc initialization with the correct local_rank and - CUDA_VISIBLE_DEVICES. + logical-to-physical GPU mapping. - CUDA_VISIBLE_DEVICES setup flow: + GPU assignment flow: 1. RayExecutorV2 enables RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES so Ray does not set CUDA_VISIBLE_DEVICES on RayWorkerProc actors at creation time. 2. Each actor is scheduled with a placement group and bundle index; Ray resolves the physical GPU ID for that bundle at placement time. - 3. After placement, the worker discovers that GPU ID and sets - CUDA_VISIBLE_DEVICES before finishing WorkerProc initialization. + 3. After placement, the executor discovers each worker's GPU ID and passes the + node's logical-to-physical mapping (assigned_physical_gpu_ids) to + initialize_worker(); CUDA_VISIBLE_DEVICES is never modified. - There is no workaround for this unset-and-reset sequence when the placement group - is externally managed: scheduling must complete before CUDA_VISIBLE_DEVICES can - match the GPU tied to the worker's bundle. + Scheduling must complete before the mapping is known when the placement + group is externally managed: only then is the GPU tied to the worker's + bundle resolved. This sequence allows multiple vLLM instances to coexist on the same node: each instance is unaware which physical devices others hold, and the - externally managed placement group avoids CUDA_VISIBLE_DEVICES conflicts + externally managed placement group avoids device assignment conflicts by binding workers to specific placement group bundles. """ @@ -120,28 +121,33 @@ class RayWorkerProc(WorkerProc): is_driver_worker=is_driver_worker, ) - def get_node_and_gpu_ids(self) -> tuple[str, list[int]]: - """Return (node_id, gpu_ids) assigned to this actor by Ray.""" + def get_node_and_physical_gpu_ids(self) -> tuple[str, list[int]]: + """Return (node_id, physical_gpu_ids) assigned to this actor by Ray.""" node_id = ray.get_runtime_context().get_node_id() device_key = current_platform.ray_device_key if not device_key: raise RuntimeError( f"current platform {current_platform.device_name} does not support ray." ) - gpu_ids = ray.get_runtime_context().get_accelerator_ids()[device_key] - return node_id, [int(x) for x in gpu_ids] + physical_gpu_ids = ray.get_runtime_context().get_accelerator_ids()[device_key] + return node_id, [ + current_platform.device_control_id_to_physical_device_id(str(x)) + for x in physical_gpu_ids + ] def initialize_worker( self, local_rank: int, env_vars: dict[str, str], driver_env_vars: dict[str, str] | None = None, + assigned_physical_gpu_ids: list[int] | None = None, ) -> None: """Complete initialization after GPU assignment is known. *driver_env_vars* are applied with ``setdefault`` — they fill in missing vars but never overwrite node-local values. - *env_vars* (e.g. CUDA_VISIBLE_DEVICES) always overwrite. + *env_vars* always overwrite. + *assigned_physical_gpu_ids* maps local_rank to physical CUDA device ID. """ if driver_env_vars: for key, value in driver_env_vars.items(): @@ -149,6 +155,13 @@ class RayWorkerProc(WorkerProc): for key, value in env_vars.items(): os.environ[key] = value + if assigned_physical_gpu_ids is not None: + vllm_config = self._init_kwargs["vllm_config"] + assert isinstance(vllm_config, VllmConfig) + vllm_config.parallel_config.assigned_physical_gpu_ids = ( + assigned_physical_gpu_ids + ) + self.local_rank = local_rank super().__init__( local_rank=local_rank, @@ -365,36 +378,48 @@ class RayExecutorV2(MultiprocExecutor): ) self.ray_worker_handles.append(handle) - # Step 6: Discover GPU IDs assigned to each worker via Ray runtime context. - worker_node_and_gpu_ids = ray.get( - [h.actor.get_node_and_gpu_ids.remote() for h in self.ray_worker_handles] + # Step 6: Discover physical GPU IDs assigned to each worker via Ray + # runtime context. + worker_node_and_physical_gpu_ids = ray.get( + [ + h.actor.get_node_and_physical_gpu_ids.remote() + for h in self.ray_worker_handles + ] ) node_workers: dict[str, list[int]] = defaultdict(list) - node_gpus: dict[str, list[int]] = defaultdict(list) - for i, (node_id, gpu_ids) in enumerate(worker_node_and_gpu_ids): + node_physical_gpu_ids: dict[str, list[int]] = defaultdict(list) + for i, (node_id, physical_gpu_ids) in enumerate( + worker_node_and_physical_gpu_ids + ): node_workers[node_id].append(i) - node_gpus[node_id].extend(gpu_ids) - for node_id, gpu_ids in node_gpus.items(): - node_gpus[node_id] = sorted(gpu_ids) + node_physical_gpu_ids[node_id].extend(physical_gpu_ids) + for node_id, physical_gpu_ids in node_physical_gpu_ids.items(): + node_physical_gpu_ids[node_id] = sorted(physical_gpu_ids) - # Step 7: Initialize workers with correct local_rank and - # CUDA_VISIBLE_DEVICES. Each worker sees all GPUs assigned to - # this executor on its node; local_rank indexes into that set. + # Step 7: Initialize workers with local logical ranks and the + # logical-to-physical GPU mapping discovered from Ray placement. init_worker_refs = [] - for i, (node_id, _) in enumerate(worker_node_and_gpu_ids): + for i, (node_id, _) in enumerate(worker_node_and_physical_gpu_ids): local_rank = node_workers[node_id].index(i) - worker_env_vars = { - current_platform.device_control_env_var: ",".join( - map(str, node_gpus[node_id]) - ), - } + assigned_physical_gpu_ids = sorted(node_physical_gpu_ids[node_id]) + worker_env_vars: dict[str, str] = {} self.ray_worker_handles[i].local_rank = local_rank init_worker_refs.append( self.ray_worker_handles[i].actor.initialize_worker.remote( - local_rank, worker_env_vars, self.driver_env_vars + local_rank, + worker_env_vars, + self.driver_env_vars, + assigned_physical_gpu_ids=assigned_physical_gpu_ids, ) ) + # Also set on the executor-side config for consistency. The mapping + # is per-node, so only do this when all workers share one node. + if len(node_physical_gpu_ids) == 1: + node_id_0 = worker_node_and_physical_gpu_ids[0][0] + self.vllm_config.parallel_config.assigned_physical_gpu_ids = sorted( + node_physical_gpu_ids[node_id_0] + ) ray.get(init_worker_refs) # Step 8: Collect response MQ handles diff --git a/vllm/v1/executor/ray_utils.py b/vllm/v1/executor/ray_utils.py index 9083b919591..cc17c39e35f 100644 --- a/vllm/v1/executor/ray_utils.py +++ b/vllm/v1/executor/ray_utils.py @@ -93,7 +93,7 @@ try: def get_node_ip(self) -> str: return get_ip() - def get_node_and_gpu_ids(self) -> tuple[str, list[int]]: + def get_node_and_physical_gpu_ids(self) -> tuple[str, list[int]]: node_id = ray.get_runtime_context().get_node_id() device_key = vllm.platforms.current_platform.ray_device_key if not device_key: @@ -101,8 +101,10 @@ try: "current platform %s does not support ray.", vllm.platforms.current_platform.device_name, ) - gpu_ids = ray.get_runtime_context().get_accelerator_ids()[device_key] - return node_id, gpu_ids + physical_gpu_ids = ray.get_runtime_context().get_accelerator_ids()[ + device_key + ] + return node_id, physical_gpu_ids def setup_device_if_necessary(self): # TODO(swang): This is needed right now because Ray CG executes diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 0291faf1afc..5e266a31354 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -270,19 +270,47 @@ class Worker(WorkerBase): # DP_LOCAL_RANK * TP_PP_WORLD_SIZE + TP_LOCAL_RANK self.local_rank += dp_local_rank * tp_pp_world_size + + # Publish the logical-to-physical mapping for topology queries + # such as NIC affinity and P2P checks. + assigned_physical_gpu_ids = parallel_config.assigned_physical_gpu_ids + if assigned_physical_gpu_ids is not None: + from vllm.platforms.interface import set_assigned_physical_gpu_ids + + set_assigned_physical_gpu_ids(assigned_physical_gpu_ids) + assert self.local_rank < len(assigned_physical_gpu_ids), ( + f"local_rank {self.local_rank} is out of bounds for " + f"assigned_physical_gpu_ids {assigned_physical_gpu_ids}" + ) + # NOTE(patch pr45026): local_world_size is derived from + # parallel_config.nnodes, which is only set for the "mp" + # multi-node backend. With the "ray"/"external_launcher" + # backends nnodes stays 1, so local_world_size collapses to + # the full world_size and this check wrongly fires on + # cross-node deployments. assigned_physical_gpu_ids is already + # per-node and the local_rank bound above fully validates the + # mapping for these backends, so skip the check for them. + if parallel_config.distributed_executor_backend not in ( + "ray", + "external_launcher", + ): + assert self.parallel_config.local_world_size <= len( + assigned_physical_gpu_ids + ), ( + f"local_world_size ({self.parallel_config.local_world_size})" + " exceeds assigned_physical_gpu_ids count " + f"({len(assigned_physical_gpu_ids)})" + ) + else: assert self.local_rank < torch.accelerator.device_count(), ( - f"DP adjusted local rank {self.local_rank} is out of bounds. " - ) - visible_device_count = ( - torch.accelerator.device_count() if torch.cuda.is_available() else 0 - ) - assert self.parallel_config.local_world_size <= visible_device_count, ( - f"local_world_size ({self.parallel_config.local_world_size}) must " - f"be less than or equal to the number of visible devices " - f"({visible_device_count})." + f"DP adjusted local rank {self.local_rank} is out of " + f"bounds for {torch.accelerator.device_count()} devices." ) - self.device = torch.device(f"cuda:{self.local_rank}") + visible_device_index = ( + current_platform.logical_device_id_to_visible_device_id(self.local_rank) + ) + self.device = torch.device(f"cuda:{visible_device_index}") torch.accelerator.set_device_index(self.device) current_platform.check_if_supports_dtype(self.model_config.dtype) diff --git a/vllm/v1/worker/worker_base.py b/vllm/v1/worker/worker_base.py index 19bb18bd39f..9381d71913d 100644 --- a/vllm/v1/worker/worker_base.py +++ b/vllm/v1/worker/worker_base.py @@ -286,6 +286,12 @@ class WorkerWrapperBase: extended_calls, ) + assigned_physical_gpu_ids = kwargs.pop("assigned_physical_gpu_ids", None) + if assigned_physical_gpu_ids is not None: + vllm_config.parallel_config.assigned_physical_gpu_ids = ( + assigned_physical_gpu_ids + ) + shared_worker_lock = kwargs.pop("shared_worker_lock", None) if shared_worker_lock is None: msg = ( From 1bdf9810aae30ae0b7002ac9f98bb9520b34e631 Mon Sep 17 00:00:00 2001 From: TJian Date: Sun, 21 Jun 2026 04:38:42 +0800 Subject: [PATCH 47/75] [ROCm] [Bugfix] Bugfix ROCm Sparse Indexer (#46222) Signed-off-by: tjtanaa --- vllm/v1/attention/ops/rocm_aiter_mla_sparse.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index dbd4d8d1d4c..2153a460f69 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -58,7 +58,9 @@ def _indexer_k_quant_and_cache_kernel( slot_id = tl.load(slot_mapping_ptr + tid) if slot_id < 0: return - block_id = slot_id // block_size + # The packed KV layout makes per-block strides large + # enough that block_id * stride can exceed 32-bit range. + block_id = (slot_id // block_size).to(tl.int64) block_offset = slot_id % block_size tile_block_id = block_offset // BLOCK_TILE_SIZE tile_block_offset = block_offset % BLOCK_TILE_SIZE @@ -179,7 +181,9 @@ def _cp_gather_indexer_quant_cache_kernel( block_table_ptr + block_table_offset, mask=valid_block_table, other=-1 ) valid_block = valid_block_table & (block_id >= 0) & (block_id < NUM_BLOCKS) - safe_block_id = tl.where(valid_block, block_id, 0) + # The packed KV layout makes per-block strides large + # enough that block_id * stride can exceed 32-bit range. + safe_block_id = tl.where(valid_block, block_id, 0).to(tl.int64) safe_block_offset = tl.where(valid_block, block_offset, 0) tiled_block_offset = safe_block_offset % BLOCK_TILE_SIZE if LAYOUT == "SHUFFLE": From 891cc4b9c58fa0ab7e4b29ee1df90724647229fd Mon Sep 17 00:00:00 2001 From: shuoming zhang <48345809+zhangshuoming990105@users.noreply.github.com> Date: Sun, 21 Jun 2026 05:12:48 +0800 Subject: [PATCH 48/75] [Frontend] Report cache usage in Anthropic /v1/messages API (#40912) Signed-off-by: mistral0105 Signed-off-by: Tyler Michael Smith Co-authored-by: Tyler Michael Smith --- .../test_anthropic_messages_conversion.py | 240 +++++++++++++++++- vllm/entrypoints/anthropic/serving.py | 65 ++++- 2 files changed, 295 insertions(+), 10 deletions(-) diff --git a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py index 4663a6565d6..b3447387c8f 100644 --- a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py +++ b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py @@ -8,6 +8,8 @@ AnthropicServingMessages._convert_anthropic_to_openai_request(). Also covers extended-thinking edge cases such as ``redacted_thinking`` blocks echoed back by Anthropic clients, and streaming conversion in ``message_stream_converter``. + +Also covers cache usage computation in ``_build_anthropic_usage``. """ import json @@ -18,7 +20,11 @@ import pytest from vllm.entrypoints.anthropic.protocol import ( AnthropicMessagesRequest, ) -from vllm.entrypoints.anthropic.serving import AnthropicServingMessages +from vllm.entrypoints.anthropic.serving import ( + AnthropicServingMessages, + _build_anthropic_usage, + _get_cached_tokens, +) from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionResponseStreamChoice, ChatCompletionStreamResponse, @@ -27,6 +33,7 @@ from vllm.entrypoints.openai.engine.protocol import ( DeltaFunctionCall, DeltaMessage, DeltaToolCall, + PromptTokenUsageInfo, UsageInfo, ) @@ -653,6 +660,108 @@ class TestThinkingBlockConversion: assert asst.get("content") == "Hi!" +# ====================================================================== +# Cache usage computation +# ====================================================================== + + +class TestGetCachedTokens: + """Tests for _get_cached_tokens helper.""" + + def test_none_usage(self): + assert _get_cached_tokens(None) is None + + def test_no_prompt_tokens_details(self): + usage = UsageInfo(prompt_tokens=100, completion_tokens=10) + assert _get_cached_tokens(usage) is None + + def test_cached_tokens_present(self): + usage = UsageInfo( + prompt_tokens=100, + completion_tokens=10, + prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=80), + ) + assert _get_cached_tokens(usage) == 80 + + def test_cached_tokens_zero(self): + """Zero cached tokens should return 0, not None.""" + usage = UsageInfo( + prompt_tokens=100, + completion_tokens=10, + prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=0), + ) + assert _get_cached_tokens(usage) == 0 + + def test_cached_tokens_none_in_details(self): + usage = UsageInfo( + prompt_tokens=100, + completion_tokens=10, + prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=None), + ) + assert _get_cached_tokens(usage) is None + + +class TestBuildAnthropicUsage: + """Tests for _build_anthropic_usage helper. + + Anthropic defines: total_input = input_tokens + cache_read + cache_creation + vLLM's prompt_tokens is the total. + """ + + def test_no_cache_info(self): + """When cache info is unavailable, return raw prompt_tokens.""" + result = _build_anthropic_usage(100, 10, None) + assert result.input_tokens == 100 + assert result.output_tokens == 10 + assert result.cache_read_input_tokens is None + assert result.cache_creation_input_tokens is None + + def test_cache_hit(self): + """When cache is hit, input_tokens excludes cached tokens.""" + usage = UsageInfo( + prompt_tokens=100, + completion_tokens=10, + prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=80), + ) + result = _build_anthropic_usage(100, 10, usage) + assert result.input_tokens == 20 # 100 - 80 + assert result.output_tokens == 10 + assert result.cache_read_input_tokens == 80 + assert result.cache_creation_input_tokens == 0 + + def test_zero_cached_tokens(self): + """Zero cached tokens should still set cache_creation to 0.""" + usage = UsageInfo( + prompt_tokens=100, + completion_tokens=10, + prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=0), + ) + result = _build_anthropic_usage(100, 10, usage) + assert result.input_tokens == 100 # 100 - 0 + assert result.cache_read_input_tokens == 0 + assert result.cache_creation_input_tokens == 0 + + def test_all_tokens_cached(self): + """When all tokens are cached, input_tokens should be 0.""" + usage = UsageInfo( + prompt_tokens=100, + completion_tokens=10, + prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=100), + ) + result = _build_anthropic_usage(100, 10, usage) + assert result.input_tokens == 0 + assert result.cache_read_input_tokens == 100 + assert result.cache_creation_input_tokens == 0 + + def test_no_prompt_tokens_details(self): + """UsageInfo without prompt_tokens_details returns no cache info.""" + usage = UsageInfo(prompt_tokens=100, completion_tokens=10) + result = _build_anthropic_usage(100, 10, usage) + assert result.input_tokens == 100 + assert result.cache_read_input_tokens is None + assert result.cache_creation_input_tokens is None + + class TestInlineSystemMessageInMessagesArray: """Verify that ``role: system`` messages embedded inside the ``messages`` array are preserved in their original position. @@ -1098,6 +1207,135 @@ class TestMessageStartIncludesTypeAndRole: assert message["role"] == "assistant" +class TestStreamingCacheUsageSemantics: + """Locks in the documented streaming behavior of cache usage fields. + + vLLM's OpenAI chat completion streaming only attaches + ``prompt_tokens_details`` to the terminal usage chunk. The Anthropic layer + mirrors that contract: cache fields are omitted on ``message_start`` (key + absence signals "unknown") and populated on ``message_delta`` (the final + cumulative count). This is intentionally consistent with vLLM's OpenAI + behavior, even though Anthropic's upstream API populates cache fields on + ``message_start``; closing that gap requires plumbing cache info into the + first chunk at the OpenAI layer, which is out of scope here. + """ + + @pytest.mark.asyncio + async def test_streaming_cache_fields_absent_then_populated(self): + """First chunk lacks prompt_tokens_details (vLLM contract); + message_start omits cache fields. The final chunk carries + prompt_tokens_details, so message_delta carries resolved values.""" + + async def sse_input(): + yield _make_stream_chunk( + delta=DeltaMessage(role="assistant", content="hi"), + usage=UsageInfo(prompt_tokens=100, total_tokens=100), + ) + yield _make_stream_chunk(finish_reason="stop") + yield _make_stream_chunk( + choices=[], + usage=UsageInfo( + prompt_tokens=100, + completion_tokens=5, + total_tokens=105, + prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=80), + ), + ) + yield "data: [DONE]" + + converter = _make_stream_converter() + output = [] + async for event in converter.message_stream_converter(sse_input()): + output.append(event) + events = _parse_sse_events(output) + + # message_start: cache fields unknown → omitted from JSON entirely. + start_usage = events[0][1]["message"]["usage"] + assert events[0][0] == "message_start" + assert start_usage["input_tokens"] == 100 + assert "cache_read_input_tokens" not in start_usage + assert "cache_creation_input_tokens" not in start_usage + + # message_delta: authoritative usage with cache fields populated. + delta_usage = next( + data["usage"] for ev, data in events if ev == "message_delta" + ) + assert delta_usage["input_tokens"] == 20 # 100 - 80 + assert delta_usage["cache_read_input_tokens"] == 80 + assert delta_usage["cache_creation_input_tokens"] == 0 + + @pytest.mark.asyncio + async def test_streaming_no_cache_hit(self): + """When the final chunk reports cached_tokens=0, message_delta carries + cache fields = 0 (cache miss); message_start still omits them.""" + + async def sse_input(): + yield _make_stream_chunk( + delta=DeltaMessage(role="assistant"), + usage=UsageInfo(prompt_tokens=50, total_tokens=50), + ) + yield _make_stream_chunk(finish_reason="stop") + yield _make_stream_chunk( + choices=[], + usage=UsageInfo( + prompt_tokens=50, + completion_tokens=5, + total_tokens=55, + prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=0), + ), + ) + yield "data: [DONE]" + + converter = _make_stream_converter() + output = [] + async for event in converter.message_stream_converter(sse_input()): + output.append(event) + events = _parse_sse_events(output) + + start_usage = events[0][1]["message"]["usage"] + delta_usage = next( + data["usage"] for ev, data in events if ev == "message_delta" + ) + assert start_usage["input_tokens"] == 50 + assert "cache_read_input_tokens" not in start_usage + assert "cache_creation_input_tokens" not in start_usage + assert delta_usage["input_tokens"] == 50 # 50 - 0 + assert delta_usage["cache_read_input_tokens"] == 0 + assert delta_usage["cache_creation_input_tokens"] == 0 + + @pytest.mark.asyncio + async def test_streaming_no_prompt_tokens_details_at_all(self): + """If --enable-prompt-tokens-details is off, no chunk carries cache + info; both message_start and message_delta omit cache fields.""" + + async def sse_input(): + yield _make_stream_chunk( + delta=DeltaMessage(role="assistant"), + usage=UsageInfo(prompt_tokens=30, total_tokens=30), + ) + yield _make_stream_chunk(finish_reason="stop") + yield _make_stream_chunk( + choices=[], + usage=UsageInfo(prompt_tokens=30, completion_tokens=2, total_tokens=32), + ) + yield "data: [DONE]" + + converter = _make_stream_converter() + output = [] + async for event in converter.message_stream_converter(sse_input()): + output.append(event) + events = _parse_sse_events(output) + + start_usage = events[0][1]["message"]["usage"] + delta_usage = next( + data["usage"] for ev, data in events if ev == "message_delta" + ) + assert "cache_read_input_tokens" not in start_usage + assert "cache_creation_input_tokens" not in start_usage + assert "cache_read_input_tokens" not in delta_usage + assert "cache_creation_input_tokens" not in delta_usage + + # ====================================================================== # Auto-detection of system-first template requirement # ====================================================================== diff --git a/vllm/entrypoints/anthropic/serving.py b/vllm/entrypoints/anthropic/serving.py index 9d5852428df..3d0151aefa8 100644 --- a/vllm/entrypoints/anthropic/serving.py +++ b/vllm/entrypoints/anthropic/serving.py @@ -43,6 +43,7 @@ from vllm.entrypoints.openai.engine.protocol import ( JsonSchemaResponseFormat, ResponseFormat, StreamOptions, + UsageInfo, ) from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.entrypoints.serve.utils.api_utils import sanitize_message @@ -54,6 +55,49 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) +def _get_cached_tokens(usage: UsageInfo | None) -> int | None: + """Extract cached token count from OpenAI UsageInfo.""" + if usage is None or usage.prompt_tokens_details is None: + return None + return usage.prompt_tokens_details.cached_tokens + + +def _build_anthropic_usage( + prompt_tokens: int, + completion_tokens: int | None, + usage: UsageInfo | None, +) -> AnthropicUsage: + """Build an AnthropicUsage from OpenAI-style token counts. + + Anthropic defines ``total_input == input_tokens + cache_read + + cache_creation``. vLLM's ``prompt_tokens`` is the total, so + ``input_tokens = prompt_tokens - cached_tokens``. + + OpenAI usage only exposes ``cached_tokens`` (hits); there is no + cache-creation analog, so ``cache_creation_input_tokens`` is ``0`` + when cache info is present. When cache info is absent (e.g. + ``--enable-prompt-tokens-details`` off, or a streaming chunk that + hasn't carried it yet), cache fields are left **unset** so + ``exclude_unset=True`` serialization omits them entirely. + + ``completion_tokens`` follows ``UsageInfo`` and may be ``None`` on + intermediate stream chunks; we coerce to ``0`` for the wire format. + """ + output_tokens = completion_tokens or 0 + cached = _get_cached_tokens(usage) + if cached is not None: + return AnthropicUsage( + input_tokens=prompt_tokens - cached, + output_tokens=output_tokens, + cache_read_input_tokens=cached, + cache_creation_input_tokens=0, + ) + return AnthropicUsage( + input_tokens=prompt_tokens, + output_tokens=output_tokens, + ) + + def wrap_data_with_event(data: str, event: str): return f"event: {event}\ndata: {data}\n\n" @@ -582,9 +626,10 @@ class AnthropicServingMessages(OpenAIServingChat): id=generator.id, content=[], model=generator.model, - usage=AnthropicUsage( - input_tokens=generator.usage.prompt_tokens, - output_tokens=generator.usage.completion_tokens, + usage=_build_anthropic_usage( + generator.usage.prompt_tokens, + generator.usage.completion_tokens, + generator.usage, ), kv_transfer_params=generator.kv_transfer_params, ) @@ -765,11 +810,12 @@ class AnthropicServingMessages(OpenAIServingChat): model=origin_chunk.model, stop_reason=None, stop_sequence=None, - usage=AnthropicUsage( - input_tokens=origin_chunk.usage.prompt_tokens + usage=_build_anthropic_usage( + origin_chunk.usage.prompt_tokens if origin_chunk.usage else 0, - output_tokens=0, + 0, + origin_chunk.usage, ), ), ) @@ -788,13 +834,14 @@ class AnthropicServingMessages(OpenAIServingChat): chunk = AnthropicStreamEvent( type="message_delta", delta=AnthropicDelta(stop_reason=stop_reason), - usage=AnthropicUsage( - input_tokens=origin_chunk.usage.prompt_tokens + usage=_build_anthropic_usage( + origin_chunk.usage.prompt_tokens if origin_chunk.usage else 0, - output_tokens=origin_chunk.usage.completion_tokens + origin_chunk.usage.completion_tokens if origin_chunk.usage else 0, + origin_chunk.usage, ), ) data = chunk.model_dump_json(exclude_unset=True) From 77148992cfc905ded5fbd34d746553aa7f099da4 Mon Sep 17 00:00:00 2001 From: Tyler Michael Smith Date: Sat, 20 Jun 2026 17:19:10 -0400 Subject: [PATCH 49/75] [Bugfix] Move extract_layer_index back inside is_v32 guard (#46199) Signed-off-by: Tyler Michael Smith Co-authored-by: Claude Opus 4.6 --- vllm/model_executor/models/deepseek_v2.py | 38 +++++++++++++---------- 1 file changed, 22 insertions(+), 16 deletions(-) diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 22c4003d3fa..2f6a472fe35 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -1001,24 +1001,30 @@ class DeepseekV2MLAAttention(nn.Module): # IndexCache config # Refer: https://arxiv.org/abs/2603.12201 for more details. _skip_topk = False - _index_topk_freq = getattr(config, "index_topk_freq", 1) - _index_topk_pattern = getattr(config, "index_topk_pattern", None) - _index_skip_topk_offset = getattr(config, "index_skip_topk_offset", 2) - layer_id = extract_layer_index(prefix) + is_mtp_layer = False + if self.is_v32: + _index_topk_freq = getattr(config, "index_topk_freq", 1) + _index_topk_pattern = getattr(config, "index_topk_pattern", None) + _index_skip_topk_offset = getattr(config, "index_skip_topk_offset", 2) + layer_id = extract_layer_index(prefix) - if _index_topk_pattern is None: - _skip_topk = ( - max(layer_id - _index_skip_topk_offset + 1, 0) % _index_topk_freq != 0 + if _index_topk_pattern is None: + _skip_topk = ( + max(layer_id - _index_skip_topk_offset + 1, 0) % _index_topk_freq + != 0 + ) + elif 0 <= layer_id < len(_index_topk_pattern): + _skip_topk = _index_topk_pattern[layer_id] == "S" + + # The skip pattern only governs backbone layers. MTP/nextn + # layers (layer_id >= num_hidden_layers) always build a full + # indexer: they compute indices at draft step 0 and toggle + # at runtime via set_skip_topk + # (index_share_for_mtp_iteration). + _num_hidden_layers = getattr(config, "num_hidden_layers", None) + is_mtp_layer = ( + _num_hidden_layers is not None and layer_id >= _num_hidden_layers ) - elif 0 <= layer_id < len(_index_topk_pattern): - _skip_topk = _index_topk_pattern[layer_id] == "S" - - # The skip pattern only governs backbone layers. MTP/nextn layers - # (layer_id >= num_hidden_layers) always build a full indexer: they - # compute indices at draft step 0 and toggle at runtime via - # set_skip_topk (index_share_for_mtp_iteration). - _num_hidden_layers = getattr(config, "num_hidden_layers", None) - is_mtp_layer = _num_hidden_layers is not None and layer_id >= _num_hidden_layers if self.is_v32 and (not _skip_topk or is_mtp_layer): self.indexer_rope_emb = get_rope( From cc22621b51207e1af96269a840108ea654af9b42 Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Sat, 20 Jun 2026 17:19:40 -0400 Subject: [PATCH 50/75] [KV Offload] Support packed HMA KV cache layout (#46205) Signed-off-by: Lucas Wilkinson Co-authored-by: OpenAI Codex Co-authored-by: Tyler Michael Smith --- tests/v1/core/test_contiguous_kv_packing.py | 97 ++++++++++++++++++- .../kv_connector/v1/offloading/worker.py | 35 ++++++- vllm/envs.py | 6 ++ vllm/v1/core/kv_cache_utils.py | 38 +++++--- vllm/v1/kv_offload/cpu/spec.py | 10 +- vllm/v1/simple_kv_offload/manager.py | 10 +- 6 files changed, 175 insertions(+), 21 deletions(-) diff --git a/tests/v1/core/test_contiguous_kv_packing.py b/tests/v1/core/test_contiguous_kv_packing.py index 79f8937c637..f4b7ee520ad 100644 --- a/tests/v1/core/test_contiguous_kv_packing.py +++ b/tests/v1/core/test_contiguous_kv_packing.py @@ -1,16 +1,23 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for contiguous KV cache packing in _get_kv_cache_config_deepseek_v4.""" +"""Tests for contiguous KV cache packing.""" from unittest.mock import MagicMock import pytest import torch -from vllm.v1.core.kv_cache_utils import _get_kv_cache_config_deepseek_v4 +from vllm import envs +from vllm.v1.core.kv_cache_utils import ( + _get_kv_cache_config_deepseek_v4, + get_kv_cache_config_from_groups, +) from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, KVCacheGroupSpec, + KVCacheTensor, MLAAttentionSpec, + SlidingWindowSpec, UniformTypeKVCacheSpecs, ) @@ -28,6 +35,25 @@ def _make_mla_spec(page_size: int, block_size: int = 256) -> MLAAttentionSpec: ) +def _make_full_spec() -> FullAttentionSpec: + return FullAttentionSpec( + block_size=16, + num_kv_heads=2, + head_size=64, + dtype=torch.float16, + ) + + +def _make_sw_spec() -> SlidingWindowSpec: + return SlidingWindowSpec( + block_size=16, + num_kv_heads=2, + head_size=64, + dtype=torch.float16, + sliding_window=128, + ) + + def _make_groups(n_c4, n_c128, n_swa): PS_C4_MLA = 37440 PS_C4_IDX = 8640 @@ -130,6 +156,73 @@ class TestInterleavedPacking: for i, v in enumerate(views): assert (v == i + 1).all(), f"View {i} was corrupted" + def test_hma_attention_groups_keep_default_backing(self, monkeypatch): + monkeypatch.setattr(envs, "VLLM_USE_PACKED_HMA_KV_CACHE", False, raising=False) + full = _make_full_spec() + sw = _make_sw_spec() + page_size = full.page_size_bytes + groups = [ + KVCacheGroupSpec(["full.0", "full.1"], full), + KVCacheGroupSpec(["sw.0", "sw.2"], sw), + KVCacheGroupSpec(["sw.1", "sw.3"], sw), + ] + + config = get_kv_cache_config_from_groups( + _mock_vllm_config(), groups, available_memory=page_size * 2 * 32 + ) + + assert config.num_blocks == 32 + assert sum(t.size for t in config.kv_cache_tensors) == page_size * 2 * 32 + assert config.kv_cache_tensors == [ + KVCacheTensor(size=page_size * 32, shared_by=["full.0", "sw.0", "sw.1"]), + KVCacheTensor(size=page_size * 32, shared_by=["full.1", "sw.2", "sw.3"]), + ] + + def test_hma_attention_groups_use_packed_backing_with_flag(self, monkeypatch): + monkeypatch.setattr(envs, "VLLM_USE_PACKED_HMA_KV_CACHE", True, raising=False) + full = _make_full_spec() + sw = _make_sw_spec() + page_size = full.page_size_bytes + groups = [ + KVCacheGroupSpec(["full.0", "full.1"], full), + KVCacheGroupSpec(["sw.0", "sw.2"], sw), + KVCacheGroupSpec(["sw.1", "sw.3"], sw), + ] + + config = get_kv_cache_config_from_groups( + _mock_vllm_config(), groups, available_memory=page_size * 2 * 32 + ) + + assert config.num_blocks == 32 + assert {t.size for t in config.kv_cache_tensors} == {page_size * 2 * 32} + assert config.kv_cache_tensors == [ + KVCacheTensor( + size=page_size * 2 * 32, + shared_by=["full.0", "sw.0", "sw.1"], + offset=0, + block_stride=page_size * 2, + ), + KVCacheTensor( + size=page_size * 2 * 32, + shared_by=["full.1", "sw.2", "sw.3"], + offset=page_size, + block_stride=page_size * 2, + ), + ] + + def test_single_group_attention_keeps_unpacked_layout(self): + spec = _make_full_spec() + groups = [KVCacheGroupSpec(["full.0", "full.1"], spec)] + + config = get_kv_cache_config_from_groups( + _mock_vllm_config(), groups, available_memory=spec.page_size_bytes * 2 * 32 + ) + + assert sum(t.size for t in config.kv_cache_tensors) == ( + spec.page_size_bytes * 2 * 32 + ) + assert [t.block_stride for t in config.kv_cache_tensors] == [0, 0] + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py index 8583bb4b1e0..f22d6738b4f 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py @@ -50,7 +50,8 @@ class OffloadingConnectorWorker: def register_kv_caches( self, kv_caches: dict[str, torch.Tensor | list[torch.Tensor]] ): - num_blocks = self.spec.kv_cache_config.num_blocks + kv_cache_config = self.spec.kv_cache_config + num_blocks = kv_cache_config.num_blocks # layer_name -> (num_blocks, page_size_bytes) tensor tensors_per_block: dict[str, tuple[torch.Tensor, ...]] = {} @@ -58,7 +59,7 @@ class OffloadingConnectorWorker: unpadded_page_size_bytes: dict[str, int] = {} # layer_name -> size of page in bytes page_size_bytes: dict[str, int] = {} - for kv_cache_group in self.spec.kv_cache_config.kv_cache_groups: + for kv_cache_group in kv_cache_config.kv_cache_groups: group_layer_names = kv_cache_group.layer_names group_kv_cache_spec = kv_cache_group.kv_cache_spec if isinstance(group_kv_cache_spec, UniformTypeKVCacheSpecs): @@ -122,9 +123,35 @@ class OffloadingConnectorWorker: else: raise NotImplementedError + packed_kv_cache_tensor = next( + (t for t in kv_cache_config.kv_cache_tensors if t.block_stride), None + ) + is_dsv4 = all( + isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) + for group in kv_cache_config.kv_cache_groups + ) + if packed_kv_cache_tensor is not None and not is_dsv4: + (tensor,) = tensors_per_block[packed_kv_cache_tensor.shared_by[0]] + block_stride = tensor.stride(0) + packed_tensor = tensor.as_strided( + (num_blocks, block_stride), + (block_stride, 1), + storage_offset=0, + ) + self._register_handlers( + CanonicalKVCaches( + [CanonicalKVCacheTensor(packed_tensor, block_stride)], + [ + [CanonicalKVCacheRef(0, block_stride)] + for _ in kv_cache_config.kv_cache_groups + ], + ) + ) + return + block_tensors: list[CanonicalKVCacheTensor] = [] block_data_refs: dict[str, list[CanonicalKVCacheRef]] = defaultdict(list) - for kv_cache_tensor in self.spec.kv_cache_config.kv_cache_tensors: + for kv_cache_tensor in kv_cache_config.kv_cache_tensors: # Filter to layers that were actually processed above. # _get_kv_cache_config_deepseek_v4 emits KVCacheTensor entries for # every (tuple_idx, page_size) slot; slots where no group has a @@ -166,7 +193,7 @@ class OffloadingConnectorWorker: ) group_data_refs: list[list[CanonicalKVCacheRef]] = [] - for kv_cache_group in self.spec.kv_cache_config.kv_cache_groups: + for kv_cache_group in kv_cache_config.kv_cache_groups: group_refs: list[CanonicalKVCacheRef] = [] for layer_name in kv_cache_group.layer_names: group_refs += block_data_refs[layer_name] diff --git a/vllm/envs.py b/vllm/envs.py index a94e084ab62..d9b10afba20 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -209,6 +209,7 @@ if TYPE_CHECKING: VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: int = 300 VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS: int = 5 VLLM_KV_CACHE_LAYOUT: Literal["NHD", "HND"] | None = None + VLLM_USE_PACKED_HMA_KV_CACHE: bool = False VLLM_SSM_CONV_STATE_LAYOUT: Literal["SD", "DS"] | None = None VLLM_COMPUTE_NANS_IN_LOGITS: bool = False VLLM_ROCM_QUICK_REDUCE_QUANTIZATION: Literal[ @@ -1608,6 +1609,11 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_KV_CACHE_LAYOUT": env_with_choices( "VLLM_KV_CACHE_LAYOUT", None, ["NHD", "HND"] ), + # Opt into packed per-block KV cache allocation for multi-group + # attention-only HMA models (e.g. gpt-oss, Gemma 3/4). + "VLLM_USE_PACKED_HMA_KV_CACHE": lambda: bool( + int(os.getenv("VLLM_USE_PACKED_HMA_KV_CACHE", "0")) + ), # SSM conv state layout used for Mamba models. # - SD: (state_len, dim) — dim contiguous (default) # - DS: (dim, state_len) — TP-sharded dim on dim1, diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index a1ebe08c078..4e1d28d7d5d 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -938,9 +938,7 @@ def _pool_bytes_per_block(kv_cache_groups: list[KVCacheGroupSpec]) -> int: kv_cache_groups[0].kv_cache_spec, UniformTypeKVCacheSpecs ): return kv_cache_groups[0].kv_cache_spec.page_size_bytes - if all( - isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs) for g in kv_cache_groups - ): + if _use_packed_kv_cache_groups(kv_cache_groups): # buckets = {page_size: [[layer_names], [layer_names], ...]} buckets = _bucket_layers_by_page_size(kv_cache_groups) return sum(ps * len(slots) for ps, slots in buckets.items()) @@ -1218,16 +1216,29 @@ def _bucket_layers_by_page_size( return buckets -def _get_kv_cache_config_deepseek_v4( +def _use_packed_kv_cache_groups( + kv_cache_groups: list[KVCacheGroupSpec], +) -> bool: + is_dsv4 = all( + isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) + for group in kv_cache_groups + ) + return is_dsv4 or ( + bool(envs.VLLM_USE_PACKED_HMA_KV_CACHE) and len(kv_cache_groups) > 1 + ) + + +def _get_kv_cache_config_packed( vllm_config: VllmConfig, kv_cache_groups: list[KVCacheGroupSpec], available_memory: int, ) -> tuple[int, list[KVCacheTensor]]: - """DeepseekV4 KV cache tensor layout planning. + """Plan a packed per-block KV cache tensor layout. Emit one KVCacheTensor per (slot_idx, page_size). Layers from different groups at the same slot share a tensor (they have independent block - tables so block-id namespaces never collide). + tables so block-id namespaces never collide). Each emitted tensor aliases + one physical backing allocation, with per-block data laid out contiguously. """ # buckets = {page_size: [[layer_names], [layer_names], ...]} buckets = _bucket_layers_by_page_size(kv_cache_groups) @@ -1255,6 +1266,9 @@ def _get_kv_cache_config_deepseek_v4( return num_blocks, kv_cache_tensors +_get_kv_cache_config_deepseek_v4 = _get_kv_cache_config_packed + + def get_kv_cache_config_from_groups( vllm_config: VllmConfig, kv_cache_groups: list[KVCacheGroupSpec], @@ -1299,13 +1313,11 @@ def get_kv_cache_config_from_groups( ) for layer_name in kv_cache_groups[0].layer_names ] - elif all( - isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) - for group in kv_cache_groups - ): - # DeepseekV4: UniformTypeKVCacheSpecs but multiple groups. - # Delegate to the DeepseekV4-specific allocator. - num_blocks, kv_cache_tensors = _get_kv_cache_config_deepseek_v4( + elif _use_packed_kv_cache_groups(kv_cache_groups): + # DeepSeek V4 keeps the existing packed layout. Other multi-group + # attention-only HMA layouts can opt in with + # VLLM_USE_PACKED_HMA_KV_CACHE=1. + num_blocks, kv_cache_tensors = _get_kv_cache_config_packed( vllm_config, kv_cache_groups, available_memory ) else: diff --git a/vllm/v1/kv_offload/cpu/spec.py b/vllm/v1/kv_offload/cpu/spec.py index d65ba9439e1..b8fb893f14d 100644 --- a/vllm/v1/kv_offload/cpu/spec.py +++ b/vllm/v1/kv_offload/cpu/spec.py @@ -58,7 +58,15 @@ class CPUOffloadingSpec(OffloadingSpec): self.cpu_page_size_per_worker = 0 assert kv_cache_config is not None if kv_cache_config.num_blocks > 0 and world_size > 0: - total_gpu_kv_bytes = sum(t.size for t in kv_cache_config.kv_cache_tensors) + is_packed = any(t.block_stride for t in kv_cache_config.kv_cache_tensors) + assert not is_packed or all( + t.block_stride for t in kv_cache_config.kv_cache_tensors + ) + total_gpu_kv_bytes = ( + kv_cache_config.kv_cache_tensors[0].size + if is_packed + else sum(t.size for t in kv_cache_config.kv_cache_tensors) + ) kv_bytes_per_block = ( total_gpu_kv_bytes // kv_cache_config.num_blocks ) * world_size diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index fe984be96a2..5e431e62388 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/vllm/v1/simple_kv_offload/manager.py @@ -187,7 +187,13 @@ class SimpleCPUOffloadScheduler: assert len(gpu_config.kv_cache_tensors) > 0 - gpu_total_bytes = sum(t.size for t in gpu_config.kv_cache_tensors) + is_packed = any(t.block_stride for t in gpu_config.kv_cache_tensors) + assert not is_packed or all(t.block_stride for t in gpu_config.kv_cache_tensors) + gpu_total_bytes = ( + gpu_config.kv_cache_tensors[0].size + if is_packed + else sum(t.size for t in gpu_config.kv_cache_tensors) + ) num_gpu_blocks = gpu_config.num_blocks num_cpu_blocks = max(1, num_gpu_blocks * cpu_capacity_bytes // gpu_total_bytes) # Create CPU kv_cache_tensors mirroring GPU by scaling size proportionally. @@ -195,6 +201,8 @@ class SimpleCPUOffloadScheduler: KVCacheTensor( size=t.size // num_gpu_blocks * num_cpu_blocks, shared_by=list(t.shared_by), + offset=t.offset, + block_stride=t.block_stride, ) for t in gpu_config.kv_cache_tensors ] From 3b4a76b63fb1a6bbf8641fa87f4ecbc9a229ac94 Mon Sep 17 00:00:00 2001 From: Varun Sundar Rabindranath Date: Sat, 20 Jun 2026 17:21:55 -0400 Subject: [PATCH 51/75] [KV-Offloading] : Expose CPU cache usage metric (#45737) Signed-off-by: Varun Sundar Rabindranath Signed-off-by: <> Co-authored-by: Varun Sundar Rabindranath --- tests/v1/kv_offload/cpu/test_manager.py | 46 ++++++++++++++++++++++--- vllm/v1/kv_offload/cpu/common.py | 5 ++- vllm/v1/kv_offload/cpu/manager.py | 27 ++++++++++----- vllm/v1/kv_offload/cpu/spec.py | 30 +++++++++++----- 4 files changed, 86 insertions(+), 22 deletions(-) diff --git a/tests/v1/kv_offload/cpu/test_manager.py b/tests/v1/kv_offload/cpu/test_manager.py index 6e4cbb1c6b8..aa4fb829597 100644 --- a/tests/v1/kv_offload/cpu/test_manager.py +++ b/tests/v1/kv_offload/cpu/test_manager.py @@ -14,12 +14,13 @@ from vllm.v1.kv_offload.base import ( ReqContext, make_offload_key, ) -from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec +from vllm.v1.kv_offload.cpu.common import ( + CPULoadStoreSpec, + CPUOffloadingMetrics, +) from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy -STORES_SKIPPED = "vllm:kv_offload_stores_skipped" - def make_req_context( req_id: str = "", kv_transfer_params: dict | None = None @@ -181,10 +182,45 @@ def test_filter_reused_manager_reports_stores_skipped_counter(): ) stats = manager.get_stats() assert stats is not None - assert stats.reduce()[STORES_SKIPPED] == 3 + assert stats.reduce()[CPUOffloadingMetrics.STORES_SKIPPED] == 3 stats = manager.get_stats() assert stats is not None - assert stats.reduce()[STORES_SKIPPED] == 0 + assert stats.reduce()[CPUOffloadingMetrics.STORES_SKIPPED] == 0 + + +def test_cpu_manager_reports_cache_usage_gauge(): + def check_usage_stats(manager: CPUOffloadingManager, value: float): + stats = manager.get_stats() + assert stats is not None + assert stats.reduce()[ + CPUOffloadingMetrics.CPU_CACHE_USAGE_PERC + ] == pytest.approx(value) + + # Zero-capacity manager always reports 0.0 + manager = make_cpu_manager(num_blocks=0) + check_usage_stats(manager, 0.0) + + # Empty manager (4 blocks, none allocated): usage = 0.0 + manager = make_cpu_manager(num_blocks=4) + check_usage_stats(manager, 0.0) + + # After allocating 2 of 4 blocks: usage = 0.5 + manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) + check_usage_stats(manager, 0.5) + + # After filling all 4 blocks: usage = 1.0 + manager.prepare_store(to_keys([3, 4]), _EMPTY_REQ_CTX) + check_usage_stats(manager, 1.0) + + # After completing store, the blocks becomes evictable as it is not actively used + # and usage drops. + manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX) + check_usage_stats(manager, 0.5) + + # After completing store, the blocks becomes evictable as it is not actively used + # and usage drops. + manager.complete_store(to_keys([3, 4]), _EMPTY_REQ_CTX) + check_usage_stats(manager, 0.0) def test_cpu_manager(): diff --git a/vllm/v1/kv_offload/cpu/common.py b/vllm/v1/kv_offload/cpu/common.py index 46bca1b9065..14c96680fd0 100644 --- a/vllm/v1/kv_offload/cpu/common.py +++ b/vllm/v1/kv_offload/cpu/common.py @@ -4,7 +4,10 @@ from typing_extensions import override from vllm.v1.kv_offload.base import BlockIDsLoadStoreSpec -METRIC_STORES_SKIPPED = "vllm:kv_offload_stores_skipped" + +class CPUOffloadingMetrics: + STORES_SKIPPED = "vllm:kv_offload_stores_skipped" + CPU_CACHE_USAGE_PERC = "vllm:kv_offload_cpu_cache_usage_perc" class CPULoadStoreSpec(BlockIDsLoadStoreSpec): diff --git a/vllm/v1/kv_offload/cpu/manager.py b/vllm/v1/kv_offload/cpu/manager.py index 7835d35309a..7d92844d1f4 100644 --- a/vllm/v1/kv_offload/cpu/manager.py +++ b/vllm/v1/kv_offload/cpu/manager.py @@ -18,7 +18,10 @@ from vllm.v1.kv_offload.base import ( ReqContext, RequestOffloadingContext, ) -from vllm.v1.kv_offload.cpu.common import METRIC_STORES_SKIPPED, CPULoadStoreSpec +from vllm.v1.kv_offload.cpu.common import ( + CPULoadStoreSpec, + CPUOffloadingMetrics, +) from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy from vllm.v1.kv_offload.cpu.policies.lru import LRUCachePolicy @@ -282,13 +285,21 @@ class CPUOffloadingManager(OffloadingManager): self.events.clear() def get_stats(self) -> OffloadingConnectorStats | None: - if self.store_threshold < 2: - return None - stats = OffloadingConnectorStats() - stats.increase_counter( - METRIC_STORES_SKIPPED, - self.stores_skipped_in_current_batch, + + # Compute cache usage. + num_used = ( + self._num_allocated_blocks + - len(self._free_list) + - self._num_evictable_cache_blocks ) - self.stores_skipped_in_current_batch = 0 + usage = num_used / self._num_blocks if self._num_blocks > 0 else 0.0 + stats.set_gauge(CPUOffloadingMetrics.CPU_CACHE_USAGE_PERC, usage) + + if self.store_threshold >= 2: + stats.increase_counter( + CPUOffloadingMetrics.STORES_SKIPPED, + self.stores_skipped_in_current_batch, + ) + self.stores_skipped_in_current_batch = 0 return stats diff --git a/vllm/v1/kv_offload/cpu/spec.py b/vllm/v1/kv_offload/cpu/spec.py index b8fb893f14d..9b1dff24a87 100644 --- a/vllm/v1/kv_offload/cpu/spec.py +++ b/vllm/v1/kv_offload/cpu/spec.py @@ -14,11 +14,15 @@ from vllm.v1.kv_offload.base import ( GPULoadStoreSpec, LoadStoreSpec, OffloadingCounterMetadata, + OffloadingGaugeMetadata, OffloadingManager, OffloadingMetricMetadata, OffloadingSpec, ) -from vllm.v1.kv_offload.cpu.common import METRIC_STORES_SKIPPED, CPULoadStoreSpec +from vllm.v1.kv_offload.cpu.common import ( + CPULoadStoreSpec, + CPUOffloadingMetrics, +) from vllm.v1.kv_offload.cpu.gpu_worker import CpuGpuOffloadingHandlers from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager from vllm.v1.kv_offload.worker.worker import OffloadingHandler @@ -31,17 +35,27 @@ class CPUOffloadingSpec(OffloadingSpec): def build_metric_definitions( cls, extra_config: dict[str, Any] ) -> dict[str, OffloadingMetricMetadata]: - store_threshold = int(extra_config.get("store_threshold", 0)) - if store_threshold < 2: - return {} - return { - METRIC_STORES_SKIPPED: OffloadingCounterMetadata( + definitions: dict[str, OffloadingMetricMetadata] = { + CPUOffloadingMetrics.CPU_CACHE_USAGE_PERC: OffloadingGaugeMetadata( documentation=( - "Number of KV offload stores skipped because the reuse " - "threshold was not reached." + "Fraction of CPU KV-cache space currently pinned by active " + "transfers (0.0 = idle, 1.0 = saturated). Sustained high " + "values indicate transfers (stores or promotions) may be " + "dropped due to insufficient capacity." ), ) } + store_threshold = int(extra_config.get("store_threshold", 0)) + if store_threshold >= 2: + definitions[CPUOffloadingMetrics.STORES_SKIPPED] = ( + OffloadingCounterMetadata( + documentation=( + "Number of KV offload stores skipped because the reuse " + "threshold was not reached." + ), + ) + ) + return definitions def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig): super().__init__(vllm_config, kv_cache_config) From ab7fcbdd5dbcb457c61f722e0a854de29491cf4d Mon Sep 17 00:00:00 2001 From: Yifan Qiao Date: Sat, 20 Jun 2026 15:00:11 -0700 Subject: [PATCH 52/75] [Perf][KVConnector][Mooncake] Compact chunk-hash keys and zero-copy lookup wire format (#45969) --- .../unit/test_mooncake_store_coordinator.py | 7 +- .../unit/test_mooncake_store_hma_e2e.py | 13 ++- .../unit/test_mooncake_store_worker.py | 38 +++++++- .../v1/mooncake/store/coordinator.py | 24 ++--- .../kv_connector/v1/mooncake/store/data.py | 87 +++++++++++++++++-- .../v1/mooncake/store/protocol.py | 5 +- .../kv_connector/v1/mooncake/store/worker.py | 43 +++++---- 7 files changed, 164 insertions(+), 53 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py b/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py index 8d00345157f..0cddd56a60a 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py @@ -7,7 +7,10 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator imp ExternalCachedBlockPool, MooncakeStoreCoordinator, ) -from vllm.v1.core.kv_cache_utils import BlockHash, BlockHashListWithBlockSize +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( + chunk_hashes_for_block_size, +) +from vllm.v1.core.kv_cache_utils import BlockHash from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheGroupSpec, @@ -182,7 +185,7 @@ def test_coordinator_group_block_size_double_hash(): ] coord = _make_coord(groups, hash_block_size=16) hs = _hashes(4) - big_hashes = list(BlockHashListWithBlockSize(hs, 16, 32)) + big_hashes = list(chunk_hashes_for_block_size(hs, 16, 32)) exists = {(0, bytes(h)) for h in hs} exists |= {(1, bytes(bh)) for bh in big_hashes} cmap = ExternalCachedBlockPool(exists) diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py index 01d4f4821ea..9e9a57cdf74 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py @@ -323,8 +323,8 @@ def test_recv_skips_swa_blocks_before_window(): def test_chunked_token_database_hash_block_size_smaller_than_block_size(): """DSv4-style: hash_block_size=4, group block_size=16 — process_tokens - must merge every 4 fine hashes into one chunk hash via - BlockHashListWithBlockSize.""" + keys each 16-token chunk by its last fine hash, keeping the Mooncake key + at one digest instead of concatenating all 4 fine hashes.""" md = KeyMetadata("m", 0, 0, 0, 0, group_id=3) db = ChunkedTokenDatabase(md, block_size=16, hash_block_size=4) db.set_kv_caches_base_addr([0]) @@ -335,8 +335,7 @@ def test_chunked_token_database_hash_block_size_smaller_than_block_size(): assert len(out) == 2 assert out[0][0] == 0 and out[0][1] == 16 assert out[1][0] == 16 and out[1][1] == 32 - # Each chunk's hash is the concatenation of 4 fine hashes. - expected0 = b"".join(fine_hashes[0:4]).hex() - expected1 = b"".join(fine_hashes[4:8]).hex() - assert out[0][2].chunk_hash == expected0 - assert out[1][2].chunk_hash == expected1 + # Each chunk's hash is its last (4th) fine hash, which already chains the + # prior three. + assert out[0][2].chunk_hash == fine_hashes[3].hex() + assert out[1][2].chunk_hash == fine_hashes[7].hex() diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index 5213805115e..96dd866babe 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -23,6 +23,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import ( worker as mooncake_store_worker, ) from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( + BlobBlockHashes, ChunkedTokenDatabase, KeyMetadata, LoadSpec, @@ -32,6 +33,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.metrics import ( MooncakeStoreConnectorStats, ) +from vllm.v1.core.kv_cache_utils import BlockHash def _default_send_coord() -> mooncake_store_worker.MooncakeStoreCoordinator: @@ -1179,9 +1181,9 @@ def test_store_sending_thread_kv_events_use_group_chunk_metadata(): assert full_event.group_idx == 0 assert full_event.block_size == 32 assert full_event.token_ids == list(range(32)) - assert full_event.block_hashes == [ - maybe_convert_block_hash(BlockHash(b"".join(hs))) - ] + # block_size=32 over hash_block_size=8 (scale 4): the chunk is keyed by its + # last sub-hash, not the concatenation of all four. + assert full_event.block_hashes == [maybe_convert_block_hash(BlockHash(hs[3]))] assert swa_event.group_idx == 1 assert swa_event.block_size == 8 @@ -1749,3 +1751,33 @@ def test_store_worker_close_swallows_store_errors(): worker.close() assert worker.store is None + + +def test_blob_block_hashes_wire_roundtrip(): + """The lookup wire format sends a ``hash_len`` frame plus the raw hashes + concatenated back-to-back; the server rebuilds them through a zero-copy + ``BlobBlockHashes`` view over the frame buffer.""" + hashes = [BlockHash(bytes([i]) * 16) for i in range(5)] + hash_len = len(hashes[0]) + + # Client side (LookupKeyClient._lookup): flat payload frame. + blob = b"".join(hashes) + + # Server side (LookupKeyServer): view over the frame buffer (a memoryview), + # never materializing the full hash list upfront. + view = BlobBlockHashes(memoryview(blob), hash_len) + + assert len(view) == 5 + assert list(view) == hashes # default Sequence iter terminates via IndexError + assert [bytes(h) for h in view] == hashes + assert bytes(view[-1]) == hashes[-1] + assert [bytes(h) for h in view[1:3]] == hashes[1:3] + with pytest.raises(IndexError): + _ = view[5] + + +def test_blob_block_hashes_empty(): + """Empty lookups send hash_len=0 and an empty payload.""" + view = BlobBlockHashes(memoryview(b""), 0) + assert len(view) == 0 + assert list(view) == [] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py index b1513e72699..89ffb560038 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py @@ -2,13 +2,15 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """External-store cache-hit coordinator for MooncakeStoreConnector.""" +from collections.abc import Sequence from typing import cast +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( + chunk_hashes_for_block_size, +) from vllm.v1.core.block_pool import BlockPool from vllm.v1.core.kv_cache_utils import ( BlockHash, - BlockHashList, - BlockHashListWithBlockSize, KVCacheBlock, ) from vllm.v1.core.single_type_kv_cache_manager import ( @@ -120,7 +122,7 @@ class MooncakeStoreCoordinator: def find_longest_cache_hit( self, - block_hashes: list[BlockHash], + block_hashes: Sequence[BlockHash], max_length: int, cached_block_pool: ExternalCachedBlockPool, *, @@ -147,7 +149,7 @@ class MooncakeStoreCoordinator: def load_mask( self, - block_hashes: list[BlockHash], + block_hashes: Sequence[BlockHash], token_len: int, ) -> tuple[list[bool], ...]: """Per-group load masks: ``mask[g][i]`` is True iff group ``g``'s @@ -236,17 +238,15 @@ class MooncakeStoreCoordinator: return tuple(masks) def block_hashes_for_spec( - self, block_hashes: list[BlockHash], spec: KVCacheSpec - ) -> BlockHashList: - if spec.block_size == self.hash_block_size: - return block_hashes - return BlockHashListWithBlockSize( + self, block_hashes: Sequence[BlockHash], spec: KVCacheSpec + ) -> Sequence[BlockHash]: + return chunk_hashes_for_block_size( block_hashes, self.hash_block_size, spec.block_size ) def _find_hit_blocks( self, - block_hashes: list[BlockHash], + block_hashes: Sequence[BlockHash], max_length: int, cached_block_pool: ExternalCachedBlockPool, *, @@ -264,7 +264,7 @@ class MooncakeStoreCoordinator: spec, group_ids, manager_cls = self.attention_groups[0] hashes = self.block_hashes_for_spec(block_hashes, spec) hit_blocks = manager_cls.find_longest_cache_hit( - block_hashes=hashes, + block_hashes=hashes, # type: ignore[arg-type] max_length=max_length, kv_cache_group_ids=group_ids, block_pool=cast(BlockPool, cached_block_pool), @@ -304,7 +304,7 @@ class MooncakeStoreCoordinator: _max_length = min(curr_hit_length + spec.block_size, max_length) hashes = self.block_hashes_for_spec(block_hashes, spec) hit_blocks = manager_cls.find_longest_cache_hit( - block_hashes=hashes, + block_hashes=hashes, # type: ignore[arg-type] max_length=_max_length, kv_cache_group_ids=group_ids, block_pool=cast(BlockPool, cached_block_pool), diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py index 55e2bd0633d..12ad46a8480 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py @@ -5,8 +5,9 @@ # (vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/). """Data classes for MooncakeStoreConnector.""" -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from dataclasses import dataclass +from typing import cast import torch @@ -23,6 +24,77 @@ from vllm.v1.core.kv_cache_utils import ( logger = init_logger(__name__) +class BlobBlockHashes(Sequence[BlockHash]): + """Lazy view over a flat buffer of fixed-size block hashes to avoid the overhead + of materializing all hashes upfront. + """ + + def __init__(self, blob: memoryview, hash_len: int): + self._blob = blob + self._hash_len = hash_len + self._n = len(blob) // hash_len if hash_len else 0 + + def __len__(self) -> int: + return self._n + + def __getitem__(self, idx): + if isinstance(idx, slice): + return [self[i] for i in range(*idx.indices(self._n))] + if idx < 0: + idx += self._n + if not 0 <= idx < self._n: + raise IndexError(idx) + off = idx * self._hash_len + return BlockHash(self._blob[off : off + self._hash_len]) + + +class _CompactChunkHashList(BlockHashListWithBlockSize): + """View that keys each ``block_size`` chunk by the last constituent + ``hash_block_size`` hash instead of concatenating all of them. + + The engine chains block hashes (each hash folds in the previous one), so the + final sub-block hash of a chunk already uniquely identifies the whole chunk + and its prefix. Using it keeps a Mooncake key at a single hash digest + regardless of the ``block_size`` / ``hash_block_size`` ratio, instead of + growing the key linearly with it (e.g. 64x for ``block_size=256``, + ``hash_block_size=4``). + """ + + def __init__( + self, + block_hashes: Sequence[BlockHash], + hash_block_size: int, + target_block_size: int, + ): + # Accept any indexable sequence (e.g. the lazy ``BlobBlockHashes``), not + # just ``list``; the base only indexes/sizes it. + assert target_block_size % hash_block_size == 0 + self.block_hashes = block_hashes # type: ignore[assignment] + self.scale_factor = target_block_size // hash_block_size + + def _get_value_at(self, idx: int) -> BlockHash: + return self.block_hashes[idx * self.scale_factor + self.scale_factor - 1] + + +def chunk_hashes_for_block_size( + block_hashes: Sequence[BlockHash], + hash_block_size: int, + block_size: int, +) -> Sequence[BlockHash]: + """Map ``hash_block_size``-granular block hashes to one compact hash per + ``block_size`` chunk (the chunk's last sub-hash). Returns ``block_hashes`` + unchanged when the two sizes are equal. + """ + if block_size == hash_block_size: + return block_hashes + # Structurally a Sequence[BlockHash] (indexable + sized); the base class + # just isn't declared as one. + return cast( + "Sequence[BlockHash]", + _CompactChunkHashList(block_hashes, hash_block_size, block_size), + ) + + @dataclass class KeyMetadata: """Metadata for constructing pool keys.""" @@ -138,18 +210,15 @@ class ChunkedTokenDatabase: Args: token_len: Total number of tokens. block_hashes: Block hashes computed at ``hash_block_size`` granularity. - When ``block_size > hash_block_size`` consecutive hashes are merged - up to the group's ``block_size`` via ``BlockHashListWithBlockSize``. + When ``block_size > hash_block_size`` each group's ``block_size`` chunk + is keyed by its last sub-hash via ``chunk_hashes_for_block_size``. mask_num: Number of tokens to skip from the beginning. """ if not block_hashes: return - if self.block_size == self.hash_block_size: - chunk_hashes: Iterable[BlockHash] = block_hashes - else: - chunk_hashes = BlockHashListWithBlockSize( - block_hashes, self.hash_block_size, self.block_size - ) + chunk_hashes: Iterable[BlockHash] = chunk_hashes_for_block_size( + block_hashes, self.hash_block_size, self.block_size + ) for chunk_id, h in enumerate(chunk_hashes): start_idx = chunk_id * self.block_size if start_idx >= token_len: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/protocol.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/protocol.py index 1317d781673..fc91b0aeebc 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/protocol.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/protocol.py @@ -11,7 +11,10 @@ Wire format (REQ/REP over IPC): msg_type == LOOKUP_MSG: frame 1: token_len (u32 big-endian, 4 bytes) - frame 2..n: msgpack-encoded list[str] of block-hash hex digests + frame 2: hash_len (u16 big-endian, 2 bytes) — byte length of each + fixed-size block hash (0 when there are no hashes) + frame 3: raw block hashes concatenated back-to-back (each hash_len + bytes); the server splits on hash_len Response: [hit_count: u32 big-endian, 4 bytes] msg_type == RESET_MSG: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index 62c2d30c9c4..e5db88ccbff 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -18,7 +18,7 @@ import socket import threading import time from collections import defaultdict -from collections.abc import Callable +from collections.abc import Callable, Sequence from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import dataclass from typing import Any, Literal, TypeVar @@ -45,6 +45,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator imp MooncakeStoreCoordinator, ) from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( # noqa: E501 + BlobBlockHashes, ChunkedTokenDatabase, KeyMetadata, MooncakeStoreConnectorMetadata, @@ -65,7 +66,6 @@ from vllm.v1.core.kv_cache_utils import ( resolve_kv_cache_block_sizes, ) from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheGroupSpec -from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder from .metrics import MooncakeStoreConnectorStats @@ -1372,7 +1372,7 @@ class MooncakeStoreWorker: return finished_sending - def lookup(self, token_len: int, block_hashes: list[BlockHash]) -> int: + def lookup(self, token_len: int, block_hashes: Sequence[BlockHash]) -> int: """Check how many prefix tokens exist in the store. Checks across all TP ranks and PP ranks. @@ -1392,6 +1392,11 @@ class MooncakeStoreWorker: group_hashes = self.coord.block_hashes_for_spec( block_hashes, self._kv_cache_groups[g_idx].kv_cache_spec ) + metadata_templates = [ + dataclasses.replace(db.metadata, tp_rank=tp, pp_rank=pp) + for tp in range(tp_count) + for pp in range(self.pp_size) + ] for chunk_id, h in enumerate(group_hashes): start_idx = chunk_id * spec_block_size if start_idx >= token_len: @@ -1400,11 +1405,11 @@ class MooncakeStoreWorker: chunk_id >= len(lookup_mask) or not lookup_mask[chunk_id] ): continue - for tp in range(tp_count): - for pp in range(self.pp_size): - md = dataclasses.replace(db.metadata, tp_rank=tp, pp_rank=pp) - candidate_keys.append(PoolKey(md, h.hex()).to_string()) - candidate_meta.append((g_idx, bytes(h))) + h_hex = h.hex() + h_bytes = bytes(h) + for md in metadata_templates: + candidate_keys.append(PoolKey(md, h_hex).to_string()) + candidate_meta.append((g_idx, h_bytes)) if not candidate_keys: return 0 @@ -1483,7 +1488,6 @@ class LookupKeyServer: store_worker: MooncakeStoreWorker, vllm_config: VllmConfig, ): - self.decoder = MsgpackDecoder() self.ctx = zmq.Context() # type: ignore[attr-defined] socket_path = get_zmq_rpc_path_lookup(vllm_config) self._ipc_path = socket_path.removeprefix("ipc://") @@ -1506,9 +1510,9 @@ class LookupKeyServer: if msg_type == LOOKUP_MSG: token_len = int.from_bytes(all_frames[1], byteorder="big") - hash_frames = all_frames[2:] - hashes_str = self.decoder.decode(hash_frames) - block_hashes = [BlockHash(bytes.fromhex(s)) for s in hashes_str] + hash_len = int.from_bytes(all_frames[2], byteorder="big") + blob = all_frames[3].buffer + block_hashes = BlobBlockHashes(blob, hash_len) result = self.store_worker.lookup(token_len, block_hashes) self.socket.send(result.to_bytes(4, "big")) @@ -1557,7 +1561,6 @@ class LookupKeyClient: """ def __init__(self, vllm_config: VllmConfig): - self.encoder = MsgpackEncoder() self.ctx = zmq.Context() # type: ignore[attr-defined] socket_path = get_zmq_rpc_path_lookup(vllm_config) self.socket = make_zmq_socket( @@ -1574,14 +1577,16 @@ class LookupKeyClient: self.futures: dict[str, Future[int]] = {} def _lookup(self, token_len: int, block_hashes: list[BlockHash]) -> int: - hash_strs = [h.hex() for h in block_hashes] - hash_frames = self.encoder.encode(hash_strs) - token_len_bytes = token_len.to_bytes(4, byteorder="big") - all_frames = [LOOKUP_MSG, token_len_bytes] + list(hash_frames) + hash_len = len(block_hashes[0]) if block_hashes else 0 + all_frames = ( + LOOKUP_MSG, + token_len.to_bytes(4, byteorder="big"), + hash_len.to_bytes(2, byteorder="big"), + b"".join(block_hashes), + ) self.socket.send_multipart(all_frames, copy=False) resp = self.socket.recv() - result = int.from_bytes(resp, "big") - return result + return int.from_bytes(resp, "big") def lookup( self, From c88d3d4775c41793543e97f1609d3161a5689905 Mon Sep 17 00:00:00 2001 From: Jonathan Chen Date: Sat, 20 Jun 2026 18:01:06 -0400 Subject: [PATCH 53/75] [SimpleCPUOffloadConnector] PCP + DCP support (#39831) Signed-off-by: Jonathan Chen --- tests/v1/simple_kv_offload/test_scheduler.py | 230 +++++++++++++++++++ vllm/v1/simple_kv_offload/manager.py | 20 +- 2 files changed, 243 insertions(+), 7 deletions(-) diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index cff60ea01d2..1ec986eada6 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -6,6 +6,7 @@ from __future__ import annotations from dataclasses import dataclass +import pytest import torch from vllm import SamplingParams @@ -1528,3 +1529,232 @@ def test_reset_pending_loads() -> None: # All GPU blocks free num_used = gpu_pool.num_gpu_blocks - gpu_pool.get_num_free_blocks() assert num_used == 1, f"Expected only null block in use, got {num_used}" + + +def _make_cp_vllm_config( + dcp_world_size: int = 1, + pcp_world_size: int = 1, +) -> VllmConfig: + """VllmConfig with context-parallel sizes set for scheduler-only tests.""" + cfg = _make_vllm_config() + + cfg.parallel_config.decode_context_parallel_size = dcp_world_size + cfg.parallel_config.prefill_context_parallel_size = pcp_world_size + return cfg + + +def _make_cp_scheduler( + *, + dcp_world_size: int = 1, + pcp_world_size: int = 1, + num_cpu_blocks: int = 8, + num_gpu_blocks: int = 16, + lazy: bool = False, +) -> SchedulerFixture: + """Build a SimpleCPUOffloadScheduler with CP-scaled virtual block size.""" + cp_world_size = dcp_world_size * pcp_world_size + virtual_block_size = BLOCK_SIZE * cp_world_size + + kv_cache_config = _make_kv_cache_config(num_gpu_blocks) + vllm_config = _make_cp_vllm_config(dcp_world_size, pcp_world_size) + cpu_capacity_bytes = _BYTES_PER_BLOCK * num_cpu_blocks + + sched = SimpleCPUOffloadScheduler( + vllm_config=vllm_config, + kv_cache_config=kv_cache_config, + cpu_capacity_bytes=cpu_capacity_bytes, + scheduler_block_size=virtual_block_size, + hash_block_size=virtual_block_size, + lazy_offload=lazy, + ) + + gpu_block_pool = BlockPool( + num_gpu_blocks=num_gpu_blocks, + enable_caching=True, + hash_block_size=virtual_block_size, + ) + sched.bind_gpu_block_pool(gpu_block_pool) + + return SchedulerFixture( + scheduler=sched, + gpu_block_pool=gpu_block_pool, + vllm_config=vllm_config, + kv_cache_config=kv_cache_config, + ) + + +def _make_cp_request( + num_blocks: int, + virtual_block_size: int, + request_id: str | None = None, +) -> Request: + """Create a request whose block hashes are computed at the virtual + (CP-scaled) block size, matching what the real scheduler does. + """ + global _req_counter + _req_counter += 1 + if request_id is None: + request_id = f"req-cp-{_req_counter}" + + num_tokens = num_blocks * virtual_block_size + 1 + start = _req_counter * 10000 + prompt_token_ids = list(range(start, start + num_tokens)) + sampling_params = SamplingParams(max_tokens=1) + + return Request( + request_id=request_id, + prompt_token_ids=prompt_token_ids, + sampling_params=sampling_params, + pooling_params=None, + mm_features=None, + block_hasher=get_request_block_hasher(virtual_block_size, sha256), + ) + + +def _allocate_cp_gpu_blocks( + gpu_block_pool: BlockPool, + request: Request, + num_blocks: int, + virtual_block_size: int, + group_id: int = 0, +) -> list: + """Allocate GPU blocks and cache them using the CP-scaled block size.""" + blocks = gpu_block_pool.get_new_blocks(num_blocks) + num_full = min(num_blocks, len(request.block_hashes)) + if num_full > 0: + gpu_block_pool.cache_full_blocks( + request=request, + blocks=blocks, + num_cached_blocks=0, + num_full_blocks=num_full, + block_size=virtual_block_size, + kv_cache_group_id=group_id, + ) + return blocks + + +# --------------------------------------------------------------------------- +# Test 15: CP block size scaling is correct +# --------------------------------------------------------------------------- +@pytest.mark.parametrize( + "dcp_world_size, pcp_world_size", + [ + (2, 1), # DCP only + (1, 2), # PCP only + (2, 2), # DCP + PCP + ], +) +def test_cp_block_size_scaling(dcp_world_size: int, pcp_world_size: int) -> None: + """Verify that the scheduler's block_size and cp_world_size are correctly + scaled when context parallelism is enabled.""" + fix = _make_cp_scheduler( + dcp_world_size=dcp_world_size, pcp_world_size=pcp_world_size + ) + sched = fix.scheduler + + expected_cp = dcp_world_size * pcp_world_size + assert sched.cp_world_size == expected_cp + assert sched.block_size == BLOCK_SIZE * expected_cp + + +# --------------------------------------------------------------------------- +# Test 16: CP eager store-and-load roundtrip +# --------------------------------------------------------------------------- +@pytest.mark.parametrize( + "dcp_world_size, pcp_world_size", + [ + (2, 1), + (1, 2), + ], +) +def test_cp_eager_store_and_load_roundtrip( + dcp_world_size: int, pcp_world_size: int +) -> None: + """With CP enabled, store blocks to CPU and reload them for a new request + with matching tokens. Verifies that hash matching and transfer-pair + construction work with the virtual block size.""" + fix = _make_cp_scheduler( + dcp_world_size=dcp_world_size, + pcp_world_size=pcp_world_size, + num_cpu_blocks=8, + num_gpu_blocks=16, + lazy=False, + ) + sched = fix.scheduler + cp = dcp_world_size * pcp_world_size + vbs = BLOCK_SIZE * cp + + num_blocks = 2 + req = _make_cp_request(num_blocks, vbs) + + # Allocate GPU blocks and register hashes + gpu_blocks = _allocate_cp_gpu_blocks(fix.gpu_block_pool, req, num_blocks, vbs) + kv_blocks = KVCacheBlocks(blocks=(gpu_blocks,)) + req.num_computed_tokens = num_blocks * vbs + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) + + block_ids = kv_blocks.get_block_ids() + sched_out = make_scheduler_output( + {req.request_id: num_blocks * vbs}, + new_reqs={req.request_id: block_ids}, + ) + + meta = sched.build_connector_meta(sched_out) + assert meta.store_event >= 0, "Expected a store event" + assert len(meta.store_gpu_blocks) == num_blocks + assert len(meta.store_cpu_blocks) == num_blocks + simulate_store_completion(sched, meta.store_event) + + # New request with same tokens — should get a full CPU cache hit. + req2 = Request( + request_id="req-cp-load", + prompt_token_ids=req.prompt_token_ids, + sampling_params=req.sampling_params, + pooling_params=None, + mm_features=None, + block_hasher=req._block_hasher, + ) + + hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0) + assert hit_tokens == num_blocks * vbs + assert is_async is True + + # Allocate fresh GPU blocks for the load. + gpu_blocks2 = fix.gpu_block_pool.get_new_blocks(num_blocks) + kv_blocks2 = KVCacheBlocks(blocks=(gpu_blocks2,)) + sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens) + + sched_out2 = make_scheduler_output( + {req2.request_id: 1}, + new_reqs={req2.request_id: kv_blocks2.get_block_ids()}, + ) + meta2 = sched.build_connector_meta(sched_out2) + assert meta2.load_event >= 0, "Expected a load event" + assert len(meta2.load_gpu_blocks) == num_blocks + assert len(meta2.load_cpu_blocks) == num_blocks + + +# --------------------------------------------------------------------------- +# Test 17: CP lazy target blocks are scaled correctly +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("cp_world_size", [1, 2, 4]) +def test_cp_lazy_target_blocks_scaling(cp_world_size: int) -> None: + """_estimate_lazy_target_blocks returns fewer blocks when cp_world_size > 1 + because each virtual block covers more tokens.""" + kv_cache_config = _make_kv_cache_config(num_blocks=16) + max_batched = 64 + + target_base = SimpleCPUOffloadScheduler._estimate_lazy_target_blocks( + kv_cache_config, max_batched, cp_world_size=1 + ) + target_cp = SimpleCPUOffloadScheduler._estimate_lazy_target_blocks( + kv_cache_config, max_batched, cp_world_size=cp_world_size + ) + + if cp_world_size == 1: + assert target_cp == target_base + else: + assert target_cp < target_base, ( + f"cp_world_size={cp_world_size}: target_cp={target_cp} should be " + f"less than target_base={target_base}" + ) diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index 5e431e62388..07978a9dd61 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/vllm/v1/simple_kv_offload/manager.py @@ -82,6 +82,9 @@ class SimpleCPUOffloadScheduler: vllm_config.kv_events_config is not None and vllm_config.kv_events_config.enable_kv_cache_events ) + dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size + pcp_world_size = vllm_config.parallel_config.prefill_context_parallel_size + self.cp_world_size = dcp_world_size * pcp_world_size self.block_size = scheduler_block_size self.hash_block_size = hash_block_size assert self.block_size % self.hash_block_size == 0 @@ -113,9 +116,6 @@ class SimpleCPUOffloadScheduler: ) # TODO (yifan): maybe need to enable kv_cache_events and metrics_collector here. - dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size - pcp_world_size = vllm_config.parallel_config.prefill_context_parallel_size - assert dcp_world_size == 1 and pcp_world_size == 1 self.cpu_coordinator: KVCacheCoordinator = get_kv_cache_coordinator( kv_cache_config=self.cpu_kv_cache_config, max_model_len=vllm_config.model_config.max_model_len, @@ -155,6 +155,7 @@ class SimpleCPUOffloadScheduler: self._target_free = self._estimate_lazy_target_blocks( kv_cache_config, vllm_config.scheduler_config.max_num_batched_tokens, + self.cp_world_size, ) else: self._target_free = 0 @@ -215,19 +216,22 @@ class SimpleCPUOffloadScheduler: @staticmethod def _estimate_lazy_target_blocks( - kv_cache_config: "KVCacheConfig", max_num_batched_tokens: int + kv_cache_config: "KVCacheConfig", + max_num_batched_tokens: int, + cp_world_size: int = 1, ) -> int: """GPU blocks to keep available (free/offloaded) per step in lazy mode.""" WATERMARK_RATIO = 1.0 # Reserve larger space to avoid running out of GPU blocks target = 0 for g in kv_cache_config.kv_cache_groups: spec = g.kv_cache_spec + block_size = spec.block_size * cp_world_size if isinstance(spec, MambaSpec): target += 2 elif isinstance(spec, SlidingWindowSpec): - target += cdiv(spec.sliding_window, spec.block_size) + 1 + target += cdiv(spec.sliding_window, block_size) + 1 else: - target += cdiv(max_num_batched_tokens, spec.block_size) + target += cdiv(max_num_batched_tokens, block_size) return int(target * (1 + WATERMARK_RATIO)) def bind_gpu_block_pool(self, gpu_block_pool: BlockPool) -> None: @@ -363,7 +367,9 @@ class SimpleCPUOffloadScheduler: continue # Number of blocks in the computed range for this group. - g_block_size = kv_cache_groups[g].kv_cache_spec.block_size + g_block_size = ( + kv_cache_groups[g].kv_cache_spec.block_size * self.cp_world_size + ) n_computed_g = cdiv(total_computed_tokens, g_block_size) # Back-trace: ext blocks sit at the tail of the computed range. From 6e919960af42f79d6811d84b2d4316212fcf59cb Mon Sep 17 00:00:00 2001 From: aman Date: Sat, 20 Jun 2026 23:36:57 +0100 Subject: [PATCH 54/75] [Perf] Skip/shrink all_token_ids copy in scheduler for non-async and V2 runner (#45840) Signed-off-by: amanchugh89 Signed-off-by: Nick Hill Co-authored-by: Claude Co-authored-by: Nick Hill --- tests/v1/core/test_scheduler.py | 37 +++++++++++++++++++++++++++++++++ vllm/v1/core/sched/output.py | 4 ++-- vllm/v1/core/sched/scheduler.py | 19 +++++++++-------- 3 files changed, 49 insertions(+), 11 deletions(-) diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index b2825c34df8..6eb5ff5cc44 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -144,6 +144,43 @@ def test_async_scheduling_pp_allows_rescheduling_with_output_placeholders(): assert req.request_id in output.num_scheduled_tokens +def test_cached_request_data_resumed_all_token_ids_mrv1_only(): + """all_token_ids carries a resumed request's token ids to the connector + for the V1 model runner, but is skipped entirely for the V2 model runner. + """ + from vllm.v1.core.kv_cache_manager import KVCacheBlocks + + scheduler = create_scheduler() + (req,) = create_requests(num_requests=1, num_tokens=8) + req.append_output_token_ids([101, 102, 103]) + + # A resumed request was not scheduled in the previous step. + assert req.request_id not in scheduler.prev_step_scheduled_req_ids + + empty_blocks = KVCacheBlocks(blocks=((),)) + + def make_cached(): + return scheduler._make_cached_request_data( + running_reqs=[], + resumed_reqs=[req], + num_scheduled_tokens={req.request_id: 1}, + spec_decode_tokens={}, + req_to_new_blocks={req.request_id: empty_blocks}, + ) + + # V1 model runner: the full token id list is propagated. + assert not scheduler.use_v2_model_runner + cached = make_cached() + assert req.request_id in cached.resumed_req_ids + assert cached.all_token_ids[req.request_id] == list(req.all_token_ids) + + # V2 model runner: all_token_ids is skipped entirely. + scheduler.use_v2_model_runner = True + cached = make_cached() + assert req.request_id in cached.resumed_req_ids + assert cached.all_token_ids == {} + + def test_schedule_partial_requests(): """Test scheduling behavior with partial requests. diff --git a/vllm/v1/core/sched/output.py b/vllm/v1/core/sched/output.py index 0c1b9d34c55..291e73bc64b 100644 --- a/vllm/v1/core/sched/output.py +++ b/vllm/v1/core/sched/output.py @@ -118,8 +118,8 @@ class CachedRequestData: # NOTE(woosuk): new_token_ids is only used for pipeline parallelism. # When PP is not used, new_token_ids will be empty. new_token_ids: list[list[int]] - # For requests not scheduled in the last step, propagate the token ids to the - # connector. Won't contain requests that were scheduled in the prior step. + # MRV1-only: For requests not scheduled in the last step, propagate the token ids + # to the connector. Won't contain requests scheduled in the prior step. all_token_ids: dict[str, list[int]] new_block_ids: list[tuple[list[int], ...] | None] num_computed_tokens: list[int] diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 25ccf79bc3a..9f7c8d74dea 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -101,6 +101,7 @@ class Scheduler(SchedulerInterface): self.finished_req_ids_dict: dict[int, set[str]] | None = ( defaultdict(set) if include_finished_set else None ) + # Track requests scheduled in prior step (MRV1-only). self.prev_step_scheduled_req_ids: set[str] = set() # Scheduling constraints. @@ -1010,8 +1011,8 @@ class Scheduler(SchedulerInterface): # Construct the scheduler output. if self.use_v2_model_runner: - scheduled_new_reqs = scheduled_new_reqs + scheduled_resumed_reqs - scheduled_resumed_reqs = [] + scheduled_new_reqs.extend(scheduled_resumed_reqs) + scheduled_resumed_reqs.clear() new_reqs_data = [ NewRequestData.from_request( req, @@ -1037,9 +1038,10 @@ class Scheduler(SchedulerInterface): req_to_new_blocks, ) - # Record the request ids that were scheduled in this step. - self.prev_step_scheduled_req_ids.clear() - self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys()) + # Record the request ids that were scheduled in this step (MRV1-only). + if not self.use_v2_model_runner: + self.prev_step_scheduled_req_ids.clear() + self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys()) new_block_ids_to_zero = ( (self.kv_cache_manager.take_new_block_ids() or None) @@ -1252,12 +1254,11 @@ class Scheduler(SchedulerInterface): req.num_computed_tokens : req.num_computed_tokens + num_tokens ] new_token_ids.append(token_ids) - scheduled_in_prev_step = req_id in self.prev_step_scheduled_req_ids if idx >= num_running_reqs: - assert not scheduled_in_prev_step resumed_req_ids.add(req_id) - if not scheduled_in_prev_step: - all_token_ids[req_id] = req.all_token_ids.copy() + if not self.use_v2_model_runner: # noqa: SIM102 + if req_id not in self.prev_step_scheduled_req_ids: + all_token_ids[req_id] = req.all_token_ids.copy() new_block_ids.append( req_to_new_blocks[req_id].get_block_ids(allow_none=True) ) From f57ac274b24cbb5a4e079a529175eef5c0745606 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Sun, 21 Jun 2026 07:43:32 +0800 Subject: [PATCH 55/75] [Render] Add reasoning/tool parsing to /derender + fix byte-fallback FFFD (#45919) Signed-off-by: aoshen524 Co-authored-by: Martin Hickey --- .../entrypoints/serve/render/test_derender.py | 436 ++++++++++++++++++ vllm/entrypoints/openai/api_server.py | 2 +- vllm/entrypoints/serve/disagg/protocol.py | 12 +- vllm/entrypoints/serve/render/serving.py | 188 ++++++-- 4 files changed, 601 insertions(+), 37 deletions(-) diff --git a/tests/entrypoints/serve/render/test_derender.py b/tests/entrypoints/serve/render/test_derender.py index a3006595c19..e452b7367a2 100644 --- a/tests/entrypoints/serve/render/test_derender.py +++ b/tests/entrypoints/serve/render/test_derender.py @@ -8,6 +8,7 @@ import pytest import pytest_asyncio from tests.utils import RemoteLaunchRenderServer +from vllm.tokenizers import get_tokenizer MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM" @@ -486,3 +487,438 @@ async def test_derender_completion_kv_transfer_params_passthrough(client): ) assert response.status_code == 200 assert response.json()["kv_transfer_params"] == kv + + +# --------------------------------------------------------------------------- +# E2E: render -> derender roundtrip with parser (reasoning + tool calls) +# --------------------------------------------------------------------------- + +PARSER_MODEL = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B" + +_E2E_TOOLS = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather for a city", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + }, + } +] + + +@pytest.fixture(scope="module") +def parser_server(): + args = [ + "--enable-auto-tool-choice", + "--tool-call-parser", + "hermes", + "--reasoning-parser", + "deepseek_r1", + ] + with RemoteLaunchRenderServer(PARSER_MODEL, args) as remote_server: + yield remote_server + + +@pytest_asyncio.fixture +async def parser_client(parser_server): + async with httpx.AsyncClient( + base_url=parser_server.url_for(""), timeout=60.0 + ) as http_client: + yield http_client + + +@pytest.fixture(scope="module") +def parser_tokenizer(): + return get_tokenizer(PARSER_MODEL) + + +def _encode(tokenizer, text: str) -> list[int]: + return tokenizer.encode(text, add_special_tokens=False) + + +def _decoded(tokenizer, token_ids: list[int]) -> str: + return tokenizer.decode(token_ids, skip_special_tokens=True) + + +def _require_markers_survive(tokenizer, text: str, *markers: str) -> list[int]: + """Encode text and skip the test if any marker is lost in roundtrip.""" + ids = _encode(tokenizer, text) + decoded = tokenizer.decode(ids, skip_special_tokens=False) + for m in markers: + if m not in decoded: + pytest.skip(f"Marker {m!r} lost in encode->decode roundtrip") + return ids + + +async def _e2e_render_chat( + client: httpx.AsyncClient, + model: str, + messages: list[dict], +) -> dict: + resp = await client.post( + "/v1/chat/completions/render", + json={"model": model, "messages": messages}, + ) + assert resp.status_code == 200, resp.text + return resp.json() + + +def _e2e_generate_response( + token_ids: list[int], + request_id: str = "chatcmpl-e2e-test", +) -> dict: + return { + "request_id": request_id, + "choices": [ + { + "index": 0, + "token_ids": token_ids, + "finish_reason": "stop", + } + ], + } + + +@pytest.mark.asyncio +async def test_e2e_plain_roundtrip(parser_client, parser_tokenizer): + """Plain text without reasoning markers roundtrips correctly.""" + messages = [{"role": "user", "content": "What is 2+2?"}] + gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages) + + answer = "The answer is four." + output_ids = _encode(parser_tokenizer, answer) + expected = _decoded(parser_tokenizer, output_ids) + + resp = await parser_client.post( + "/v1/chat/completions/derender", + json={ + "model": PARSER_MODEL, + "generate_response": _e2e_generate_response(output_ids), + "prompt_tokens": len(gen_req["token_ids"]), + }, + ) + assert resp.status_code == 200, resp.text + content = resp.json()["choices"][0]["message"]["content"] + assert content == expected + + +@pytest.mark.asyncio +async def test_e2e_token_identity(parser_client, parser_tokenizer): + """encode(derender(token_ids)) == token_ids (RL invariant).""" + messages = [{"role": "user", "content": "Hi"}] + gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages) + + answer = "Hello! How can I help?" + output_ids = _encode(parser_tokenizer, answer) + + resp = await parser_client.post( + "/v1/chat/completions/derender", + json={ + "model": PARSER_MODEL, + "generate_response": _e2e_generate_response(output_ids), + "prompt_tokens": len(gen_req["token_ids"]), + }, + ) + assert resp.status_code == 200 + content = resp.json()["choices"][0]["message"]["content"] + re_encoded = _encode(parser_tokenizer, content) + assert output_ids == re_encoded + + +@pytest.mark.asyncio +async def test_e2e_non_ascii_roundtrip(parser_client, parser_tokenizer): + """CJK + emoji roundtrip without U+FFFD.""" + messages = [{"role": "user", "content": "Reply in Chinese"}] + gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages) + + answer = "你好世界 😀" + output_ids = _encode(parser_tokenizer, answer) + + resp = await parser_client.post( + "/v1/chat/completions/derender", + json={ + "model": PARSER_MODEL, + "generate_response": _e2e_generate_response(output_ids), + "prompt_tokens": len(gen_req["token_ids"]), + }, + ) + assert resp.status_code == 200 + content = resp.json()["choices"][0]["message"]["content"] + assert "�" not in content + + +@pytest.mark.asyncio +async def test_e2e_parsed_reasoning(parser_client, parser_tokenizer): + """... splits into reasoning + content.""" + messages = [{"role": "user", "content": "What is 2+3?"}] + gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages) + + reasoning_text = "The user wants 2 plus 3. That is 5." + answer_text = "The answer is 5." + output_text = f"{reasoning_text}{answer_text}" + output_ids = _require_markers_survive(parser_tokenizer, output_text, "") + + resp = await parser_client.post( + "/v1/chat/completions/derender", + json={ + "model": PARSER_MODEL, + "generate_response": _e2e_generate_response(output_ids), + "prompt_tokens": len(gen_req["token_ids"]), + "chat_request": { + "model": PARSER_MODEL, + "messages": messages, + "include_reasoning": True, + }, + }, + ) + assert resp.status_code == 200, resp.text + msg = resp.json()["choices"][0]["message"] + assert msg["reasoning"] is not None + assert reasoning_text in msg["reasoning"] + assert answer_text in msg["content"] + assert "" not in msg["content"] + + +@pytest.mark.asyncio +async def test_e2e_parsed_tool_call(parser_client, parser_tokenizer): + """ extracted into tool_calls field.""" + messages = [{"role": "user", "content": "Weather in Paris?"}] + gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages) + + output_text = ( + "Let me check the weather." + '\n{"name": "get_weather", ' + '"arguments": {"city": "Paris"}}\n' + ) + output_ids = _require_markers_survive( + parser_tokenizer, + output_text, + "", + "", + "", + ) + + resp = await parser_client.post( + "/v1/chat/completions/derender", + json={ + "model": PARSER_MODEL, + "generate_response": _e2e_generate_response(output_ids), + "prompt_tokens": len(gen_req["token_ids"]), + "chat_request": { + "model": PARSER_MODEL, + "messages": messages, + "tools": _E2E_TOOLS, + "tool_choice": "auto", + }, + }, + ) + assert resp.status_code == 200, resp.text + choice = resp.json()["choices"][0] + assert choice["message"]["tool_calls"] + assert choice["message"]["tool_calls"][0]["function"]["name"] == "get_weather" + + +@pytest.mark.asyncio +async def test_e2e_parsed_reasoning_and_tool_call(parser_client, parser_tokenizer): + """Reasoning + tool call in the same output.""" + messages = [{"role": "user", "content": "Weather in Paris?"}] + gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages) + + reasoning_text = "I should look up the weather." + tool_text = ( + '\n{"name": "get_weather", ' + '"arguments": {"city": "Paris"}}\n' + ) + output_text = f"{reasoning_text}{tool_text}" + output_ids = _require_markers_survive( + parser_tokenizer, output_text, "", "" + ) + + resp = await parser_client.post( + "/v1/chat/completions/derender", + json={ + "model": PARSER_MODEL, + "generate_response": _e2e_generate_response(output_ids), + "prompt_tokens": len(gen_req["token_ids"]), + "chat_request": { + "model": PARSER_MODEL, + "messages": messages, + "tools": _E2E_TOOLS, + "tool_choice": "auto", + "include_reasoning": True, + }, + }, + ) + assert resp.status_code == 200, resp.text + choice = resp.json()["choices"][0] + assert choice["message"]["reasoning"] is not None + assert reasoning_text in choice["message"]["reasoning"] + assert choice["message"]["tool_calls"] + + +@pytest.mark.asyncio +async def test_e2e_no_chat_request_fallback(parser_client, parser_tokenizer): + """Without chat_request, derender falls back to plain detokenization.""" + messages = [{"role": "user", "content": "Hello"}] + gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages) + + answer = "Hi there!" + output_ids = _encode(parser_tokenizer, answer) + + resp = await parser_client.post( + "/v1/chat/completions/derender", + json={ + "model": PARSER_MODEL, + "generate_response": _e2e_generate_response(output_ids), + "prompt_tokens": len(gen_req["token_ids"]), + }, + ) + assert resp.status_code == 200 + content = resp.json()["choices"][0]["message"]["content"] + assert "Hi" in content + + +# --------------------------------------------------------------------------- +# E2E: HarmonyParser + GPT-OSS +# --------------------------------------------------------------------------- + +HARMONY_MODEL = "openai/gpt-oss-20b" + + +def _ensure_harmony_vocab(): + """Pre-cache the o200k_base BPE file needed by openai-harmony. + + The Rust tiktoken-rs backend downloads from Azure Blob Storage, which + may be unreachable in some environments. When the cache is cold we + fetch the file ourselves and place it in ``/tmp/tiktoken-rs-cache/`` + using the SHA-1(URL) filename that tiktoken-rs expects. + """ + import hashlib + import urllib.request + from pathlib import Path + + url = "https://openaipublic.blob.core.windows.net/encodings/o200k_base.tiktoken" + cache_dir = Path("/tmp/tiktoken-rs-cache") + cache_key = hashlib.sha1(url.encode()).hexdigest() + cache_file = cache_dir / cache_key + if not cache_file.exists(): + cache_dir.mkdir(parents=True, exist_ok=True) + urllib.request.urlretrieve(url, cache_file) + + +@pytest.fixture(scope="module") +def harmony_server(): + _ensure_harmony_vocab() + args = [ + "--trust-remote-code", + "--enable-auto-tool-choice", + "--tool-call-parser", + "openai", + "--reasoning-parser", + "openai_gptoss", + ] + with RemoteLaunchRenderServer(HARMONY_MODEL, args) as remote_server: + yield remote_server + + +@pytest_asyncio.fixture +async def harmony_client(harmony_server): + async with httpx.AsyncClient( + base_url=harmony_server.url_for(""), timeout=60.0 + ) as http_client: + yield http_client + + +@pytest.fixture(scope="module") +def harmony_tokenizer(): + return get_tokenizer(HARMONY_MODEL, trust_remote_code=True) + + +def _harmony_extract_assistant_ids( + tokenizer, assistant_msg: dict, user_content: str = "test" +) -> list[int]: + """Extract assistant token IDs via apply_chat_template diff.""" + prompt = [{"role": "user", "content": user_content}] + full = prompt + [assistant_msg] + text_prompt = tokenizer.apply_chat_template( + prompt, add_generation_prompt=True, tokenize=False + ) + text_full = tokenizer.apply_chat_template( + full, add_generation_prompt=False, tokenize=False + ) + prompt_ids = tokenizer.encode(text_prompt) + full_ids = tokenizer.encode(text_full) + assistant_ids = list(full_ids[len(prompt_ids) :]) + if not assistant_ids: + pytest.skip("Could not extract assistant tokens for Harmony") + return assistant_ids + + +@pytest.mark.asyncio +async def test_e2e_harmony_plain_roundtrip(harmony_client, harmony_tokenizer): + """GPT-OSS content-only roundtrip.""" + messages = [{"role": "user", "content": "What is 2+2?"}] + gen_req = await _e2e_render_chat(harmony_client, HARMONY_MODEL, messages) + + assistant_msg = {"role": "assistant", "content": "Four."} + output_ids = _harmony_extract_assistant_ids(harmony_tokenizer, assistant_msg) + + resp = await harmony_client.post( + "/v1/chat/completions/derender", + json={ + "model": HARMONY_MODEL, + "generate_response": _e2e_generate_response(output_ids), + "prompt_tokens": len(gen_req["token_ids"]), + "chat_request": { + "model": HARMONY_MODEL, + "messages": messages, + }, + }, + ) + assert resp.status_code == 200, resp.text + content = resp.json()["choices"][0]["message"]["content"] + assert content is not None and len(content) > 0 + assert "Four" in content + + +@pytest.mark.asyncio +async def test_e2e_harmony_reasoning(harmony_client, harmony_tokenizer): + """GPT-OSS reasoning: analysis channel extracted.""" + messages = [{"role": "user", "content": "Add 2 and 3."}] + gen_req = await _e2e_render_chat(harmony_client, HARMONY_MODEL, messages) + + reasoning_text = "The user wants 2 plus 3." + answer_text = "The answer is 5." + assistant_msg = { + "role": "assistant", + "thinking": reasoning_text, + "content": answer_text, + } + output_ids = _harmony_extract_assistant_ids(harmony_tokenizer, assistant_msg) + + decoded = harmony_tokenizer.decode(output_ids) + if reasoning_text not in decoded: + pytest.skip("Harmony template did not render thinking") + + resp = await harmony_client.post( + "/v1/chat/completions/derender", + json={ + "model": HARMONY_MODEL, + "generate_response": _e2e_generate_response(output_ids), + "prompt_tokens": len(gen_req["token_ids"]), + "chat_request": { + "model": HARMONY_MODEL, + "messages": messages, + "include_reasoning": True, + }, + }, + ) + assert resp.status_code == 200, resp.text + msg = resp.json()["choices"][0]["message"] + assert msg["reasoning"] is not None + assert reasoning_text in msg["reasoning"] + assert answer_text in (msg["content"] or "") diff --git a/vllm/entrypoints/openai/api_server.py b/vllm/entrypoints/openai/api_server.py index e1e2ef72bbd..34c4f0ca5d7 100644 --- a/vllm/entrypoints/openai/api_server.py +++ b/vllm/entrypoints/openai/api_server.py @@ -455,7 +455,7 @@ async def init_render_app_state( enable_auto_tools=args.enable_auto_tool_choice, exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none, tool_parser=args.tool_call_parser, - reasoning_parser=args.structured_outputs_config.reasoning_parser, + reasoning_parser=args.reasoning_parser, default_chat_template_kwargs=args.default_chat_template_kwargs, log_error_stack=args.log_error_stack, ) diff --git a/vllm/entrypoints/serve/disagg/protocol.py b/vllm/entrypoints/serve/disagg/protocol.py index c13c4c1705c..7e776ae7178 100644 --- a/vllm/entrypoints/serve/disagg/protocol.py +++ b/vllm/entrypoints/serve/disagg/protocol.py @@ -219,10 +219,14 @@ class GenerateResponse(BaseModel): class DerenderChatRequest(BaseModel): - """Request for the /v1/chat/completions/derender endpoint. + """Request for the /v1/chat/completions/derender endpoint (non-streaming). - Wraps a GenerateResponse and caller-supplied metadata needed to produce - a fully-formed ChatCompletionResponse without a GPU. + Wraps a complete GenerateResponse and caller-supplied metadata needed to + produce a fully-formed ChatCompletionResponse without a GPU. + + Streaming derender would require a separate endpoint design with + incremental token delivery, ``OutputProcessor``-based detokenization, + and ``parser.parse_delta()`` instead of ``parser.parse()``. """ model: str @@ -244,7 +248,7 @@ class DerenderChatRequest(BaseModel): class DerenderCompletionRequest(BaseModel): - """Request for the /v1/completions/derender endpoint. + """Request for the /v1/completions/derender endpoint (non-streaming). Parallel to DerenderChatRequest but handles the multi-prompt completions case: one GenerateResponse per prompt, mirroring the list[GenerateRequest] diff --git a/vllm/entrypoints/serve/render/serving.py b/vllm/entrypoints/serve/render/serving.py index 1f7296cdaa7..612ff6d35e0 100644 --- a/vllm/entrypoints/serve/render/serving.py +++ b/vllm/entrypoints/serve/render/serving.py @@ -27,6 +27,7 @@ from vllm.entrypoints.openai.completion.protocol import ( ) from vllm.entrypoints.openai.engine.protocol import ( ErrorResponse, + ToolCall, UsageInfo, ) from vllm.entrypoints.openai.engine.serving import resolve_token_id_placeholder @@ -43,7 +44,6 @@ from vllm.entrypoints.serve.disagg.protocol import ( DerenderChatRequest, DerenderCompletionRequest, GenerateRequest, - GenerateResponseChoice, MultiModalFeatures, PlaceholderRangeInfo, ) @@ -76,21 +76,83 @@ from vllm.utils.mistral import mt as _mt logger = init_logger(__name__) +def _parse_token_id_placeholder(token: str) -> int | None: + """Extract token ID from a 'token_id:N' placeholder string.""" + if not token.startswith("token_id:"): + return None + try: + return int(token[len("token_id:") :]) + except ValueError: + return None + + +def _correct_decoded_token( + token_id: int, context_token_ids: list[int], tokenizer: TokenizerLike +) -> str: + """Use preceding tokens as context to fix U+FFFD from byte-fallback. + + Mirrors LogprobsProcessor._correct_decoded_token in v1/engine/logprobs.py. + """ + max_ctx = min(len(context_token_ids), 4) + + for num_ctx in range(1, max_ctx + 1): + context = context_token_ids[-num_ctx:] + full_decoded = tokenizer.decode(context + [token_id]) + + if full_decoded.endswith("�"): + continue + + clean_end = len(context) + for j in range(len(context) - 1, -1, -1): + if tokenizer.decode([context[j]]).endswith("�"): + clean_end = j + else: + break + + clean_prefix = tokenizer.decode(context[:clean_end]) if clean_end > 0 else "" + + if full_decoded.startswith(clean_prefix): + return full_decoded[len(clean_prefix) :] + + common_len = 0 + for a, b in zip(clean_prefix, full_decoded): + if a != b: + break + common_len += 1 + return full_decoded[common_len:] + + return "" + + def _resolve_logprobs( logprobs: ChatCompletionLogProbs, tokenizer: TokenizerLike ) -> ChatCompletionLogProbs: - """Resolve all token_id:N placeholders in a ChatCompletionLogProbs object.""" + """Resolve token_id:N placeholders in a ChatCompletionLogProbs object.""" if logprobs.content is None: return logprobs + + context_token_ids: list[int] = [] resolved_content = [] + for entry in logprobs.content: token_str, token_bytes = resolve_token_id_placeholder(entry.token, tokenizer) + sampled_id = _parse_token_id_placeholder(entry.token) + + if token_str.endswith("�") and sampled_id is not None: + token_str = _correct_decoded_token(sampled_id, context_token_ids, tokenizer) + token_bytes = list(token_str.encode("utf-8")) + resolved_top = [] for top in entry.top_logprobs: top_str, top_bytes = resolve_token_id_placeholder(top.token, tokenizer) + top_id = _parse_token_id_placeholder(top.token) + if top_str.endswith("�") and top_id is not None: + top_str = _correct_decoded_token(top_id, context_token_ids, tokenizer) + top_bytes = list(top_str.encode("utf-8")) resolved_top.append( top.model_copy(update={"token": top_str, "bytes": top_bytes}) ) + resolved_content.append( entry.model_copy( update={ @@ -100,6 +162,10 @@ def _resolve_logprobs( } ) ) + + if sampled_id is not None: + context_token_ids.append(sampled_id) + return ChatCompletionLogProbs(content=resolved_content) @@ -136,30 +202,6 @@ def _convert_chat_logprobs_to_completion_logprobs( ) -def _build_chat_choice( - choice: GenerateResponseChoice, tokenizer: TokenizerLike -) -> ChatCompletionResponseChoice: - """Detokenize and resolve logprobs for a single GenerateResponseChoice. - - Raises: - ValueError: if choice.token_ids is empty or None. - """ - if not choice.token_ids: - raise ValueError(f"choice {choice.index} has empty or null token_ids") - decoded_text = tokenizer.decode(choice.token_ids, skip_special_tokens=True) - resolved_logprobs = ( - _resolve_logprobs(choice.logprobs, tokenizer) - if choice.logprobs is not None - else None - ) - return ChatCompletionResponseChoice( - index=choice.index, - message=ChatMessage(role="assistant", content=decoded_text), - logprobs=resolved_logprobs, - finish_reason=choice.finish_reason, - ) - - class OpenAIServingRender: def __init__( self, @@ -536,9 +578,12 @@ class OpenAIServingRender: ) -> ChatCompletionResponse | ErrorResponse: """Postprocess a GenerateResponse into a ChatCompletionResponse. - This is the symmetric inverse of render_chat_request: it detokenizes - output token IDs, resolves token_id:N logprob placeholders, and - formats the result as an OpenAI-compatible chat completion response. + Non-streaming only: expects the complete GenerateResponse with all + token IDs present. Uses ``parser.parse()`` for one-shot extraction. + + When ``request.chat_request`` is provided, the parser splits the + output into (reasoning, content, tool_calls). Otherwise falls + back to plain detokenization. """ error_check_ret = await self._check_model(request) if error_check_ret is not None: @@ -546,11 +591,89 @@ class OpenAIServingRender: tokenizer = self.renderer.get_tokenizer() gen = request.generate_response + chat_request = request.chat_request choices: list[ChatCompletionResponseChoice] = [] try: for choice in gen.choices: - choices.append(_build_chat_choice(choice, tokenizer)) + if not choice.token_ids: + raise ValueError( + f"choice {choice.index} has empty or null token_ids" + ) + + resolved_logprobs = ( + _resolve_logprobs(choice.logprobs, tokenizer) + if choice.logprobs is not None + else None + ) + + if self.parser is not None and chat_request is not None: + # Parser path: decode with special tokens preserved + # so the parser can see markers like , + # , or Harmony channel tokens. + decoded_text = tokenizer.decode( + choice.token_ids, skip_special_tokens=False + ) + + chat_template_kwargs: dict[str, Any] = {} + if not self.use_harmony: + chat_template_kwargs = ( + chat_request.build_chat_params( + self.chat_template, + self.chat_template_content_format, + ) + .with_defaults(self.default_chat_template_kwargs) + .chat_template_kwargs + ) + + parser = self.parser( + tokenizer, + chat_request.tools, + chat_template_kwargs=chat_template_kwargs, + ) + reasoning, content, tool_calls = parser.parse( + decoded_text, + chat_request, + enable_auto_tools=self.enable_auto_tools, + model_output_token_ids=choice.token_ids, + ) + + if not getattr(chat_request, "include_reasoning", True): + reasoning = None + + tc_items = ( + [ + ToolCall( + id=random_uuid(), + function=tc, + ) + for tc in tool_calls + ] + if tool_calls + else [] + ) + + message = ChatMessage( + role="assistant", + reasoning=reasoning, + content=content, + tool_calls=tc_items, + ) + else: + # No parser: plain detokenization. + decoded_text = tokenizer.decode( + choice.token_ids, skip_special_tokens=True + ) + message = ChatMessage(role="assistant", content=decoded_text) + + choices.append( + ChatCompletionResponseChoice( + index=choice.index, + message=message, + logprobs=resolved_logprobs, + finish_reason=choice.finish_reason, + ) + ) except ValueError as exc: return self.create_error_response(str(exc)) @@ -587,8 +710,9 @@ class OpenAIServingRender: ) -> CompletionResponse | ErrorResponse: """Postprocess a list of GenerateResponses into a CompletionResponse. - Mirrors the multi-prompt completions case: one GenerateResponse per - prompt, parallel to the list[GenerateRequest] from /v1/completions/render. + Non-streaming only. Mirrors the multi-prompt completions case: one + GenerateResponse per prompt, parallel to the list[GenerateRequest] + from /v1/completions/render. """ error_check_ret = await self._check_model(request) if error_check_ret is not None: From 8dd1b702f27edeed24a4336f531b01c346e04253 Mon Sep 17 00:00:00 2001 From: Umut Polat <52835619+umut-polat@users.noreply.github.com> Date: Sun, 21 Jun 2026 02:57:01 +0300 Subject: [PATCH 56/75] [Misc] Fix stale doc URL and docstring module path (#35530) Signed-off-by: umut-polat <52835619+umut-polat@users.noreply.github.com> Co-authored-by: Flora Feng <4florafeng@gmail.com> --- vllm/envs.py | 2 +- vllm/tool_parsers/__init__.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/envs.py b/vllm/envs.py index d9b10afba20..190b15667dd 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -486,7 +486,7 @@ def get_vllm_port() -> int | None: raise ValueError( f"VLLM_PORT '{port}' appears to be a URI. " "This may be caused by a Kubernetes service discovery issue," - "check the warning in: https://docs.vllm.ai/en/stable/serving/env_vars.html" + "check the warning in: https://docs.vllm.ai/en/latest/configuration/env_vars.html" ) from None raise ValueError(f"VLLM_PORT '{port}' must be a valid integer") from err diff --git a/vllm/tool_parsers/__init__.py b/vllm/tool_parsers/__init__.py index bbc4d2edb19..7ce1520ffe5 100644 --- a/vllm/tool_parsers/__init__.py +++ b/vllm/tool_parsers/__init__.py @@ -15,7 +15,7 @@ Register a lazy module mapping. Example: ToolParserManager.register_lazy_module( name="kimi_k2", - module_path="vllm.tool_parsers.kimi_k2_parser", + module_path="vllm.tool_parsers.kimi_k2_tool_parser", class_name="KimiK2ToolParser", ) """ From 7df3d7dada840c68b85b26b79de7f59f676d58e3 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Sat, 20 Jun 2026 20:02:24 -0700 Subject: [PATCH 57/75] [Core] Ensure memory is pinned prior to async h2d copy (#45424) Signed-off-by: Nick Hill --- .../v1/logits_processors/test_correctness.py | 2 - .../test_gpu_model_runner_streaming.py | 1 - tests/v1/worker/test_gpu_input_batch.py | 6 -- tests/v1/worker/test_gpu_model_runner.py | 4 - vllm/device_allocator/cumem.py | 4 +- vllm/device_allocator/xpumem.py | 4 +- vllm/lora/lora_model.py | 4 +- vllm/lora/lora_weights.py | 4 +- vllm/lora/model_manager.py | 4 +- .../layers/attention/mla_attention.py | 18 ++-- .../layers/attention/mm_encoder_attention.py | 3 +- .../layers/pooler/seqwise/methods.py | 16 ++-- vllm/model_executor/models/moonvit.py | 5 +- vllm/model_executor/models/qwen2_5_vl.py | 5 +- vllm/models/deepseek_v4/sparse_mla.py | 3 +- vllm/multimodal/inputs.py | 16 +++- vllm/platforms/__init__.py | 3 +- vllm/platforms/cuda.py | 3 +- vllm/platforms/xpu.py | 12 +-- vllm/utils/torch_utils.py | 33 ++++---- vllm/v1/attention/backends/flashinfer.py | 6 +- vllm/v1/attention/backends/flex_attention.py | 8 +- vllm/v1/attention/backends/gdn_attn.py | 19 +++-- vllm/v1/attention/backends/mamba2_attn.py | 27 +++--- .../backends/mla/flashinfer_mla_sparse.py | 4 +- .../attention/backends/mla/flashmla_sparse.py | 4 +- vllm/v1/attention/backends/utils.py | 35 ++++---- vllm/v1/kv_offload/cpu/gpu_worker.py | 6 +- vllm/v1/pool/metadata.py | 6 +- vllm/v1/sample/logits_processor/builtin.py | 15 ++-- vllm/v1/sample/ops/penalties.py | 5 +- vllm/v1/sample/sampler.py | 4 +- vllm/v1/sample/thinking_budget_state.py | 5 +- vllm/v1/serial_utils.py | 4 +- vllm/v1/simple_kv_offload/worker.py | 4 +- vllm/v1/spec_decode/extract_hidden_states.py | 5 +- vllm/v1/spec_decode/llm_base_proposer.py | 16 ++-- vllm/v1/spec_decode/ngram_proposer_gpu.py | 4 +- .../backend_lm_format_enforcer.py | 4 +- vllm/v1/structured_output/backend_outlines.py | 4 +- vllm/v1/structured_output/utils.py | 25 +++--- vllm/v1/utils.py | 3 +- vllm/v1/worker/cpu/shm.py | 11 ++- vllm/v1/worker/gpu/buffer_utils.py | 11 ++- vllm/v1/worker/gpu/mm/encoder_runner.py | 5 +- vllm/v1/worker/gpu/model_runner.py | 5 +- vllm/v1/worker/gpu/model_states/whisper.py | 7 +- vllm/v1/worker/gpu_input_batch.py | 33 ++++---- vllm/v1/worker/gpu_model_runner.py | 83 +++++++++---------- 49 files changed, 254 insertions(+), 264 deletions(-) diff --git a/tests/v1/logits_processors/test_correctness.py b/tests/v1/logits_processors/test_correctness.py index 80083fd57fe..c93593865e0 100644 --- a/tests/v1/logits_processors/test_correctness.py +++ b/tests/v1/logits_processors/test_correctness.py @@ -145,7 +145,6 @@ def _generate_fake_sampling_metadata( vllm_config.scheduler_config.max_num_seqs, num_spec, device, - PIN_MEMORY_AVAILABLE, ) fake_sampling_metadata = SamplingMetadata( temperature=torch.full((batch_size,), 0.0), @@ -880,7 +879,6 @@ def test_maybe_create_thinking_budget_holder_without_reasoning(): cfg.scheduler_config.max_num_seqs, 0, torch.device("cpu"), - False, ) is None ) diff --git a/tests/v1/streaming_input/test_gpu_model_runner_streaming.py b/tests/v1/streaming_input/test_gpu_model_runner_streaming.py index 946ca99507d..9b130e570f6 100644 --- a/tests/v1/streaming_input/test_gpu_model_runner_streaming.py +++ b/tests/v1/streaming_input/test_gpu_model_runner_streaming.py @@ -35,7 +35,6 @@ def mock_model_runner_with_input_batch(): max_model_len=1024, max_num_batched_tokens=1024, device="cpu", - pin_memory=False, vocab_size=32000, block_sizes=[16], kernel_block_sizes=[16], diff --git a/tests/v1/worker/test_gpu_input_batch.py b/tests/v1/worker/test_gpu_input_batch.py index 3a478d21013..bfd4016c9fe 100644 --- a/tests/v1/worker/test_gpu_input_batch.py +++ b/tests/v1/worker/test_gpu_input_batch.py @@ -10,7 +10,6 @@ import torch from vllm.platforms import current_platform from vllm.sampling_params import SamplingParams -from vllm.utils.platform_utils import is_pin_memory_available from vllm.utils.torch_utils import make_tensor_with_pad from vllm.v1.pool.metadata import PoolingMetadata from vllm.v1.sample.logits_processor import LogitsProcessors @@ -236,7 +235,6 @@ def test_sampling_metadata_in_input_batch(device: str, batch_size: int): max_model_len=1024, max_num_batched_tokens=1024, device=torch.device(device), - pin_memory=is_pin_memory_available(), vocab_size=1024, block_sizes=[1], kernel_block_sizes=[1], @@ -331,7 +329,6 @@ def test_swap_states_in_input_batch(device: str, batch_size: int, swap_list: lis max_model_len=1024, max_num_batched_tokens=1024, device=torch.device(device), - pin_memory=is_pin_memory_available(), vocab_size=1024, block_sizes=[1], kernel_block_sizes=[1], @@ -341,7 +338,6 @@ def test_swap_states_in_input_batch(device: str, batch_size: int, swap_list: lis max_model_len=1024, max_num_batched_tokens=1024, device=torch.device(device), - pin_memory=is_pin_memory_available(), vocab_size=1024, block_sizes=[1], kernel_block_sizes=[1], @@ -410,7 +406,6 @@ def test_pooling_prompt_lens_not_aliased(device: str): max_model_len=MAX_PROMPT_SIZE + NUM_OUTPUT_TOKENS, max_num_batched_tokens=batch_size * (MAX_PROMPT_SIZE + NUM_OUTPUT_TOKENS), device=torch.device(device), - pin_memory=is_pin_memory_available(), vocab_size=VOCAB_SIZE, block_sizes=[16], kernel_block_sizes=[16], @@ -459,7 +454,6 @@ def test_pooling_metadata_token_id_buffers( max_model_len=MAX_PROMPT_SIZE + NUM_OUTPUT_TOKENS, max_num_batched_tokens=MAX_PROMPT_SIZE + NUM_OUTPUT_TOKENS, device=torch.device("cpu"), - pin_memory=False, vocab_size=VOCAB_SIZE, block_sizes=[16], kernel_block_sizes=[16], diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 80dd8ee306b..75d8c9c7460 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -85,7 +85,6 @@ def initialize_kv_cache(runner: GPUModelRunner): max_model_len=runner.max_model_len, max_num_batched_tokens=runner.max_num_tokens, device=runner.device, - pin_memory=runner.pin_memory, vocab_size=runner.model_config.get_vocab_size(), block_sizes=[kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size], kernel_block_sizes=[ @@ -1405,7 +1404,6 @@ def test_input_batch_with_kernel_block_sizes(): max_model_len = 512 max_num_batched_tokens = 512 device = torch.device(DEVICE_TYPE) - pin_memory = False vocab_size = 50272 # Test with different kernel block sizes @@ -1417,7 +1415,6 @@ def test_input_batch_with_kernel_block_sizes(): max_model_len=max_model_len, max_num_batched_tokens=max_num_batched_tokens, device=device, - pin_memory=pin_memory, vocab_size=vocab_size, block_sizes=block_sizes, kernel_block_sizes=kernel_block_sizes, @@ -1478,7 +1475,6 @@ def test_hybrid_cache_integration(default_vllm_config, dist_init): max_model_len=runner.max_model_len, max_num_batched_tokens=runner.max_num_tokens, device=runner.device, - pin_memory=runner.pin_memory, vocab_size=runner.model_config.get_vocab_size(), block_sizes=[kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size], kernel_block_sizes=[16], diff --git a/vllm/device_allocator/cumem.py b/vllm/device_allocator/cumem.py index c30790df9ca..59c0cf45f5d 100644 --- a/vllm/device_allocator/cumem.py +++ b/vllm/device_allocator/cumem.py @@ -18,8 +18,8 @@ import torch from vllm.device_allocator import AllocationData, HandleType from vllm.logger import init_logger -from vllm.utils.platform_utils import is_pin_memory_available from vllm.utils.system_utils import find_loaded_library +from vllm.utils.torch_utils import PIN_MEMORY logger = init_logger(__name__) @@ -196,7 +196,7 @@ class CuMemAllocator: size_in_bytes, dtype=torch.uint8, device="cpu", - pin_memory=is_pin_memory_available(), + pin_memory=PIN_MEMORY, ) cpu_ptr = cpu_backup_tensor.data_ptr() libcudart.cudaMemcpy(cpu_ptr, ptr, size_in_bytes) diff --git a/vllm/device_allocator/xpumem.py b/vllm/device_allocator/xpumem.py index 7d99ced7ef5..e0f359b200d 100644 --- a/vllm/device_allocator/xpumem.py +++ b/vllm/device_allocator/xpumem.py @@ -11,7 +11,7 @@ import torch from vllm.device_allocator import AllocationData, HandleType from vllm.logger import init_logger -from vllm.utils.platform_utils import is_pin_memory_available +from vllm.utils.torch_utils import PIN_MEMORY logger = init_logger(__name__) @@ -188,7 +188,7 @@ class XpuMemAllocator: size_in_bytes, dtype=torch.uint8, device="cpu", - pin_memory=is_pin_memory_available(), + pin_memory=PIN_MEMORY, ) cpu_ptr = cpu_backup_tensor.data_ptr() _xpu_memcpy_sync( diff --git a/vllm/lora/lora_model.py b/vllm/lora/lora_model.py index e3cb82e3569..859ed02f871 100644 --- a/vllm/lora/lora_model.py +++ b/vllm/lora/lora_model.py @@ -17,7 +17,7 @@ from vllm.lora.utils import ( ) from vllm.model_executor.model_loader.tensorizer import TensorizerConfig from vllm.model_executor.models.utils import WeightsMapper -from vllm.utils.platform_utils import is_pin_memory_available +from vllm.utils.torch_utils import PIN_MEMORY logger = init_logger(__name__) @@ -126,7 +126,7 @@ class LoRAModel: skip_prefixes: list[str] | None = None, ) -> "LoRAModel": """Create a LoRAModel from a dictionary of tensors.""" - pin_memory = str(device) == "cpu" and is_pin_memory_available() + pin_memory = str(device) == "cpu" and PIN_MEMORY loras: dict[str, LoRALayerWeights] = {} for tensor_name, tensor in tensors.items(): if is_base_embedding_weights(tensor_name): diff --git a/vllm/lora/lora_weights.py b/vllm/lora/lora_weights.py index 90b7df818a8..f90724c5eb5 100644 --- a/vllm/lora/lora_weights.py +++ b/vllm/lora/lora_weights.py @@ -7,7 +7,7 @@ import torch import torch.types from vllm.lora.peft_helper import PEFTHelper -from vllm.utils.platform_utils import is_pin_memory_available +from vllm.utils.torch_utils import PIN_MEMORY class LoRALayerWeights: @@ -79,7 +79,7 @@ class LoRALayerWeights: dtype: torch.dtype, device: torch.types.Device, ) -> "LoRALayerWeights": - pin_memory = str(device) == "cpu" and is_pin_memory_available() + pin_memory = str(device) == "cpu" and PIN_MEMORY lora_a = torch.zeros( [rank, input_dim], dtype=dtype, device=device, pin_memory=pin_memory ) diff --git a/vllm/lora/model_manager.py b/vllm/lora/model_manager.py index 8063f485bf5..39c3bb0ea16 100644 --- a/vllm/lora/model_manager.py +++ b/vllm/lora/model_manager.py @@ -42,7 +42,7 @@ from vllm.model_executor.models.utils import PPMissingLayer from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.encoder_budget import MultiModalBudget from vllm.utils.cache import LRUCache -from vllm.utils.platform_utils import is_pin_memory_available +from vllm.utils.torch_utils import PIN_MEMORY logger = init_logger(__name__) @@ -801,7 +801,7 @@ class LoRAModelManager: # 2. The weight packing above (e.g., pack_moe) may invalidate the # pin_memory allocation, so we execute it after packing. - pin_memory = str(lora_device) == "cpu" and is_pin_memory_available() + pin_memory = str(lora_device) == "cpu" and PIN_MEMORY if pin_memory: for lora in lora_model.loras.values(): if isinstance(lora.lora_a, list): diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index ab3874c5dad..247d6dc3a4b 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -1684,12 +1684,13 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]): # [[0, 0, 0, 0], [256, 256, 256, 256], [512, 512, 512, 512]] # Note(simon): this is done in CPU because of downstream's # of `to_list`. - chunk_starts = ( + chunk_starts = torch.empty( + num_chunks, num_prefills, dtype=torch.int32, pin_memory=True + ).copy_( torch.arange(num_chunks, dtype=torch.int32) + .multiply_(max_context_chunk) .unsqueeze(1) - .expand(-1, num_prefills) - * max_context_chunk - ).pin_memory() + ) chunk_ends = torch.min( context_lens_cpu.unsqueeze(0), chunk_starts + max_context_chunk ) @@ -1746,12 +1747,13 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]): ) * self.dcp_local_block_size ) - local_chunk_starts = ( + local_chunk_starts = torch.empty( + num_chunks, num_prefills, dtype=torch.int32, pin_memory=True + ).copy_( torch.arange(num_chunks, dtype=torch.int32) + .multiply_(padded_local_max_context_chunk_across_ranks) .unsqueeze(1) - .expand(-1, num_prefills) - * padded_local_max_context_chunk_across_ranks - ).pin_memory() + ) local_chunk_ends = torch.min( padded_local_context_lens_cpu.unsqueeze(0), local_chunk_starts diff --git a/vllm/model_executor/layers/attention/mm_encoder_attention.py b/vllm/model_executor/layers/attention/mm_encoder_attention.py index 2ca051ad9e4..bb1c995aeb5 100644 --- a/vllm/model_executor/layers/attention/mm_encoder_attention.py +++ b/vllm/model_executor/layers/attention/mm_encoder_attention.py @@ -28,6 +28,7 @@ from vllm.utils.flashinfer import ( is_flashinfer_cudnn_fp8_prefill_attn_supported, ) from vllm.utils.math_utils import round_up +from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.attention.backends.fa_utils import get_flash_attn_version from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.attention.ops.vit_attn_wrappers import ( @@ -311,7 +312,7 @@ class MMEncoderAttention(CustomOp): ) cu_seqlens = np.concatenate([cu_seqlens_qko, cu_seqlens_v]) - cu_seqlens = torch.from_numpy(cu_seqlens).to(device, non_blocking=True) + cu_seqlens = async_tensor_h2d(cu_seqlens, device=device) return cu_seqlens def __init__( diff --git a/vllm/model_executor/layers/pooler/seqwise/methods.py b/vllm/model_executor/layers/pooler/seqwise/methods.py index 06dddde7deb..5dea5d76273 100644 --- a/vllm/model_executor/layers/pooler/seqwise/methods.py +++ b/vllm/model_executor/layers/pooler/seqwise/methods.py @@ -10,6 +10,7 @@ import torch.nn as nn from vllm.config.pooler import SequencePoolingType from vllm.model_executor.layers.pooler import PoolingParamsUpdate from vllm.tasks import PoolingTask +from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.pool.metadata import PoolingMetadata SequencePoolingMethodOutput: TypeAlias = torch.Tensor | list[torch.Tensor] @@ -74,15 +75,14 @@ class MeanPool(SequencePoolingMethod): # early return for empty batch return hidden_states.new_empty((0, hidden_size), dtype=torch.float32) - # Build segment_ids on CPU so repeat_interleave doesn't need to sync - # GPU->CPU to learn its data-dependent output length, then upload - # non-blocking. eg. [2, 1, 3] -> [0, 0, 1, 2, 2, 2] + prompt_lens = async_tensor_h2d( + prompt_lens_cpu, device=hidden_states.device, dtype=torch.int64 + ) + # eg. [2, 1, 3] -> [0, 0, 1, 2, 2, 2] segment_ids = torch.repeat_interleave( - torch.arange(num_seqs, dtype=torch.long), - prompt_lens_cpu, - ).to(hidden_states.device, non_blocking=True) - prompt_lens = prompt_lens_cpu.to( - hidden_states.device, dtype=torch.int64, non_blocking=True + torch.arange(num_seqs, device=hidden_states.device, dtype=torch.long), + prompt_lens, + output_size=int(prompt_lens_cpu.sum()), ) segment_sums = torch.zeros( (num_seqs, hidden_size), diff --git a/vllm/model_executor/models/moonvit.py b/vllm/model_executor/models/moonvit.py index 73e17cb9fb6..56204dd3c61 100644 --- a/vllm/model_executor/models/moonvit.py +++ b/vllm/model_executor/models/moonvit.py @@ -66,6 +66,7 @@ from vllm.model_executor.models.utils import maybe_prefix from vllm.model_executor.models.vision import is_vit_use_data_parallel from vllm.platforms import current_platform from vllm.transformers_utils.configs.moonvit import MoonViTConfig +from vllm.utils.torch_utils import async_tensor_h2d def _apply_rope_input_validation(x, freqs_cis): @@ -758,7 +759,7 @@ class MoonVitPretrainedModel(PreTrainedModel): ), ] ) - metadata["cu_seqlens"] = torch.from_numpy(cu_seqlens_np).to(device) + metadata["cu_seqlens"] = async_tensor_h2d(cu_seqlens_np, device=device) if max_seqlen_override is not None: max_seqlen_val = int(max_seqlen_override) @@ -770,7 +771,7 @@ class MoonVitPretrainedModel(PreTrainedModel): metadata["max_seqlen"] = torch.tensor(max_seqlen_val, dtype=torch.int32) gather_idx_np = _build_merge_gather_idx(grid_pairs, self.merge_kernel_size) - metadata["merge_gather_idx"] = torch.from_numpy(gather_idx_np).to(device) + metadata["merge_gather_idx"] = async_tensor_h2d(gather_idx_np, device=device) return metadata diff --git a/vllm/model_executor/models/qwen2_5_vl.py b/vllm/model_executor/models/qwen2_5_vl.py index 04c54f1b348..986783fa34d 100644 --- a/vllm/model_executor/models/qwen2_5_vl.py +++ b/vllm/model_executor/models/qwen2_5_vl.py @@ -83,9 +83,8 @@ from vllm.multimodal.parse import MultiModalDataItems from vllm.multimodal.processing import PromptReplacement, PromptUpdate from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors -from vllm.utils.platform_utils import is_pin_memory_available from vllm.utils.tensor_schema import TensorSchema, TensorShape -from vllm.utils.torch_utils import async_tensor_h2d +from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphReplayBuffers @@ -825,7 +824,7 @@ class Qwen2_5_VisionTransformer(nn.Module): @staticmethod def invert_permutation(perm: torch.Tensor) -> torch.Tensor: # building the inverse permutation in O(n) time - inv = torch.empty_like(perm, pin_memory=is_pin_memory_available()) + inv = torch.empty_like(perm, pin_memory=PIN_MEMORY) inv[perm] = torch.arange(perm.numel(), device=perm.device, dtype=perm.dtype) return inv diff --git a/vllm/models/deepseek_v4/sparse_mla.py b/vllm/models/deepseek_v4/sparse_mla.py index ca14fe20b13..2d2a42824c3 100644 --- a/vllm/models/deepseek_v4/sparse_mla.py +++ b/vllm/models/deepseek_v4/sparse_mla.py @@ -13,6 +13,7 @@ from vllm.config.cache import CacheDType from vllm.platforms.interface import DeviceCapability from vllm.triton_utils import tl, triton from vllm.utils.math_utils import cdiv +from vllm.utils.torch_utils import np_to_pinned_tensor from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, @@ -207,7 +208,7 @@ class DeepseekV4FlashMLAMetadataBuilder( # Zero-fill for cudagraphs self.req_id_per_token_buffer.fill_(0) self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_( - torch.from_numpy(req_id_per_token), non_blocking=True + np_to_pinned_tensor(req_id_per_token), non_blocking=True ) req_id_per_token = self.req_id_per_token_buffer[:num_tokens] diff --git a/vllm/multimodal/inputs.py b/vllm/multimodal/inputs.py index d98a1624ac3..c55bbe4623c 100644 --- a/vllm/multimodal/inputs.py +++ b/vllm/multimodal/inputs.py @@ -488,7 +488,13 @@ class MultiModalBatchedField(BaseMultiModalField): # An optimization when `batch` contains only one tensor: # - produce exactly same result as `torch.stack(batch)` # - will achieve zero-copy if the tensor is contiguous - return batch[0].unsqueeze(0).contiguous() + out = batch[0].unsqueeze(0) + if not pin_memory: + return out.contiguous() + # Avoid extra copy - pinning unpinned memory will make it contiguous + if not out.is_contiguous() and out.is_pinned(): + out = out.contiguous() + return out.pin_memory() first_shape = batch[0].shape if all(elem.shape == first_shape for elem in batch): out = torch.empty( @@ -538,7 +544,13 @@ class MultiModalFlatField(BaseMultiModalField): # An optimization when `batch` contains only one tensor: # - produce exactly same result as `torch.concat(batch)` # - will achieve zero-copy if the tensor is contiguous - return batch[0].contiguous() + out = batch[0] + if not pin_memory: + return out.contiguous() + # Avoid extra copy - pinning unpinned memory will make it contiguous + if not out.is_contiguous() and out.is_pinned(): + out = out.contiguous() + return out.pin_memory() dim = self.dim + (self.dim < 0) * len(batch[0].shape) diff --git a/vllm/platforms/__init__.py b/vllm/platforms/__init__.py index 36692c7b76f..ac536aff00c 100644 --- a/vllm/platforms/__init__.py +++ b/vllm/platforms/__init__.py @@ -9,7 +9,6 @@ from typing import TYPE_CHECKING from vllm import envs from vllm.plugins import PLATFORM_PLUGINS_GROUP, load_plugins_by_group from vllm.utils.import_utils import resolve_obj_by_qualname -from vllm.utils.torch_utils import supports_xccl from .interface import CpuArchEnum, Platform, PlatformEnum @@ -135,7 +134,7 @@ def xpu_platform_plugin() -> str | None: try: import torch - if supports_xccl(): + if torch.distributed.is_xccl_available(): dist_backend = "xccl" from vllm.platforms.xpu import XPUPlatform diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 6bf1793eefd..259077da356 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -23,7 +23,6 @@ import vllm._C_stable_libtorch # noqa import vllm.envs as envs from vllm.logger import init_logger from vllm.utils.import_utils import import_pynvml -from vllm.utils.torch_utils import is_quantized_kv_cache from vllm.v1.attention.backends.registry import AttentionBackendEnum from .interface import DeviceCapability, Platform, PlatformEnum, in_wsl @@ -88,6 +87,8 @@ def _get_backend_priorities( kv_cache_dtype: CacheDType | None = None, ) -> list[AttentionBackendEnum]: """Get backend priorities with lazy import to avoid circular dependency.""" + from vllm.utils.torch_utils import is_quantized_kv_cache + if use_mla: if device_capability.major == 10: # Sparse MLA backend priorities diff --git a/vllm/platforms/xpu.py b/vllm/platforms/xpu.py index 3e208688e81..030b4933bb6 100644 --- a/vllm/platforms/xpu.py +++ b/vllm/platforms/xpu.py @@ -14,7 +14,6 @@ import vllm_xpu_kernels._xpu_C # noqa import vllm.envs as envs from vllm.logger import init_logger -from vllm.utils.torch_utils import supports_xpu_graph from vllm.v1.attention.backends.registry import AttentionBackendEnum from .interface import DeviceCapability, Platform, PlatformEnum @@ -178,8 +177,6 @@ class XPUPlatform(Platform): @classmethod def check_and_update_config(cls, vllm_config: VllmConfig) -> None: - parallel_config = vllm_config.parallel_config - # lazy import to avoid circular import from vllm.config import CUDAGraphMode @@ -190,6 +187,10 @@ class XPUPlatform(Platform): attention_config = vllm_config.attention_config if attention_config.backend is None: attention_config.backend = AttentionBackendEnum.FLASH_ATTN + + # lazy import to avoid circular import + from vllm.utils.torch_utils import supports_xpu_graph + if not supports_xpu_graph(): compilation_config.cudagraph_mode = CUDAGraphMode.NONE logger.warning( @@ -324,9 +325,8 @@ class XPUPlatform(Platform): @classmethod def get_device_communicator_cls(cls) -> str: - from vllm.utils.torch_utils import supports_xccl - - if not supports_xccl(): + if not torch.distributed.is_xccl_available(): + # Supports xccl with PyTorch versions >= 2.8.0.dev for XPU platform logger.warning( "xccl is not enabled in this torch build, communication" " is not available." diff --git a/vllm/utils/torch_utils.py b/vllm/utils/torch_utils.py index 12ec5b0fcc6..9269fbb44d7 100644 --- a/vllm/utils/torch_utils.py +++ b/vllm/utils/torch_utils.py @@ -3,7 +3,6 @@ import contextlib import importlib.metadata import os -import platform import random import threading from collections.abc import Callable, Collection @@ -18,6 +17,7 @@ from torch.library import Library, infer_schema import vllm.envs as envs from vllm.logger import init_logger +from vllm.utils.platform_utils import is_pin_memory_available if TYPE_CHECKING: from vllm.config import ModelConfig @@ -68,9 +68,7 @@ MODELOPT_TO_VLLM_KV_CACHE_DTYPE_MAP = { T = TypeVar("T") -# Pin memory in non-WSL case. -# Logic duplicated here for now to avoid circular import. -PIN_MEMORY = "microsoft" not in " ".join(platform.uname()).lower() +PIN_MEMORY = is_pin_memory_available() def is_quantized_kv_cache(kv_cache_dtype: str) -> bool: @@ -606,14 +604,24 @@ def create_kv_caches_with_random( def async_tensor_h2d( - data: list, - dtype: torch.dtype, + data: list | np.ndarray | torch.Tensor, device: str | torch.device, - pin_memory: bool = PIN_MEMORY, + dtype: torch.dtype | None = None, ) -> torch.Tensor: - """Asynchronously create a tensor and copy it from host to device.""" - t = torch.tensor(data, dtype=dtype, pin_memory=pin_memory, device="cpu") - return t.to(device=device, non_blocking=True) + """Copy list/numpy array/tensor async from host to device.""" + if isinstance(data, np.ndarray): + data = torch.from_numpy(data) + if isinstance(data, torch.Tensor): + t = data.pin_memory() if PIN_MEMORY else data + else: + t = torch.tensor(data, dtype=dtype, pin_memory=PIN_MEMORY, device="cpu") + assert t.is_cpu + return t.to(device=device, dtype=dtype, non_blocking=True) + + +def np_to_pinned_tensor(array: np.ndarray) -> torch.Tensor: + t = torch.from_numpy(array) + return t.pin_memory() if PIN_MEMORY else t def make_ndarray_with_pad( @@ -914,11 +922,6 @@ def _encode_layer_name(layer_name: str) -> str | LayerName: return LayerName(layer_name) if _USE_LAYERNAME else layer_name -# Supports xccl with PyTorch versions >= 2.8.0.dev for XPU platform -def supports_xccl() -> bool: - return torch.distributed.is_xccl_available() - - # Supports XPU Graph with PyTorch versions >= 2.11.0.dev for XPU platform def supports_xpu_graph() -> bool: return is_torch_equal_or_newer("2.11.0.dev") diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 486aa7e4054..666c32bca85 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -41,8 +41,8 @@ from vllm.utils.flashinfer import ( use_trtllm_attention, ) from vllm.utils.math_utils import cdiv -from vllm.utils.platform_utils import is_pin_memory_available from vllm.utils.torch_utils import ( + PIN_MEMORY, canonicalize_singleton_dim_strides, is_quantized_kv_cache, is_strictly_contiguous, @@ -708,9 +708,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): # Since we do not have explicit synchronization in ModelRunnerV2, we do not pin # reused CPU buffers to avoid a race condition between step N async copies to # GPU and step N+1 buffer updates. - self.pin_memory = ( - not vllm_config.use_v2_model_runner and is_pin_memory_available() - ) + self.pin_memory = not vllm_config.use_v2_model_runner and PIN_MEMORY self.paged_kv_indptr = self._make_buffer(max_num_reqs + 1) self.paged_kv_indptr_cpu_buffer = torch.zeros_like( self.paged_kv_indptr.cpu, pin_memory=self.pin_memory diff --git a/vllm/v1/attention/backends/flex_attention.py b/vllm/v1/attention/backends/flex_attention.py index 829f3472dd7..983544b5602 100644 --- a/vllm/v1/attention/backends/flex_attention.py +++ b/vllm/v1/attention/backends/flex_attention.py @@ -28,7 +28,11 @@ from vllm.logger import init_logger from vllm.model_executor.layers.attention import Attention from vllm.platforms import current_platform from vllm.utils.math_utils import cdiv -from vllm.utils.torch_utils import is_quantized_kv_cache, is_torch_equal_or_newer +from vllm.utils.torch_utils import ( + async_tensor_h2d, + is_quantized_kv_cache, + is_torch_equal_or_newer, +) from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, @@ -58,7 +62,7 @@ def _offsets_to_doc_ids_tensor( doc_ids = torch.repeat_interleave( torch.arange(len(counts), dtype=torch.int32), counts ) - return doc_ids.to(device, non_blocking=True) + return async_tensor_h2d(doc_ids, device=device) def pad_to_multiple(x: torch.Tensor, multiple: int, dim: int): diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 9323c5d8a46..c615ab62c1c 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -8,6 +8,7 @@ from typing import Literal import torch from vllm.config import VllmConfig +from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, @@ -203,8 +204,8 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] spec_sequence_masks = None spec_sequence_masks_cpu = None else: - spec_sequence_masks = spec_sequence_masks_cpu.to( - query_start_loc.device, non_blocking=True + spec_sequence_masks = async_tensor_h2d( + spec_sequence_masks_cpu, device=query_start_loc.device ) if spec_sequence_masks is None: @@ -376,12 +377,14 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] ) assert prefill_query_start_loc_cpu is not None - chunk_indices = prepare_chunk_indices( - prefill_query_start_loc_cpu, FLA_CHUNK_SIZE - ).to(device=gpu_device, non_blocking=True) - chunk_offsets = prepare_chunk_offsets( - prefill_query_start_loc_cpu, FLA_CHUNK_SIZE - ).to(device=gpu_device, non_blocking=True) + chunk_indices = async_tensor_h2d( + prepare_chunk_indices(prefill_query_start_loc_cpu, FLA_CHUNK_SIZE), + device=gpu_device, + ) + chunk_offsets = async_tensor_h2d( + prepare_chunk_offsets(prefill_query_start_loc_cpu, FLA_CHUNK_SIZE), + device=gpu_device, + ) if num_prefills > 0: has_initial_state = context_lens_tensor > 0 diff --git a/vllm/v1/attention/backends/mamba2_attn.py b/vllm/v1/attention/backends/mamba2_attn.py index 5f25c4a7952..6b4999ab35b 100644 --- a/vllm/v1/attention/backends/mamba2_attn.py +++ b/vllm/v1/attention/backends/mamba2_attn.py @@ -7,6 +7,7 @@ from typing import Any import torch from vllm.config import VllmConfig +from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.attention.backend import ( AttentionBackend, CommonAttentionMetadata, @@ -68,22 +69,22 @@ def compute_varlen_chunk_metadata( # Exclusive prefix sum over logical-chunk lengths if chunk_lens: - cu_chunk_seqlens = torch.tensor( - [0] + list(itertools.accumulate(chunk_lens)), - device=device, - dtype=torch.int32, - ) - # Final boundary must equal total tokens - assert int(cu_chunk_seqlens[-1].item()) == total + cu_chunk_seqlens_list = [0] + list(itertools.accumulate(chunk_lens)) + # Final boundary must equal total tokens (check on host to avoid a sync) + assert cu_chunk_seqlens_list[-1] == total else: - cu_chunk_seqlens = torch.tensor([0], device=device, dtype=torch.int32) + cu_chunk_seqlens_list = [0] + cu_chunk_seqlens = async_tensor_h2d( + cu_chunk_seqlens_list, dtype=torch.int32, device=device + ) - last_chunk_indices_t = ( - torch.tensor(last_chunk_indices, device=device, dtype=torch.int32) - if len(starts) > 0 - else torch.empty((0,), device=device, dtype=torch.int32) + # last_chunk_indices is empty when there are no sequences (len(starts) == 0). + last_chunk_indices_t = async_tensor_h2d( + last_chunk_indices, dtype=torch.int32, device=device + ) + seq_idx_chunks_t = async_tensor_h2d( + seq_idx_chunks, dtype=torch.int32, device=device ) - seq_idx_chunks_t = torch.tensor(seq_idx_chunks, device=device, dtype=torch.int32) return cu_chunk_seqlens, last_chunk_indices_t, seq_idx_chunks_t diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index 01716f567d0..5547c626493 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -26,7 +26,7 @@ from vllm.model_executor.layers.attention.mla_attention import ( get_mla_dims, ) from vllm.platforms.interface import DeviceCapability -from vllm.utils.torch_utils import is_quantized_kv_cache +from vllm.utils.torch_utils import is_quantized_kv_cache, np_to_pinned_tensor from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, @@ -217,7 +217,7 @@ class FlashInferMLASparseMetadataBuilder( # Zero-fill for cudagraphs self.req_id_per_token_buffer.fill_(0) self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_( - torch.from_numpy(req_id_per_token), non_blocking=True + np_to_pinned_tensor(req_id_per_token), non_blocking=True ) req_id_per_token_tensor = self.req_id_per_token_buffer[:num_tokens] diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index 6d8dfe13128..19381efd732 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -16,7 +16,7 @@ from vllm.model_executor.layers.attention.mla_attention import ( from vllm.platforms import current_platform from vllm.platforms.interface import DeviceCapability from vllm.utils.platform_utils import num_compute_units -from vllm.utils.torch_utils import is_quantized_kv_cache +from vllm.utils.torch_utils import is_quantized_kv_cache, np_to_pinned_tensor from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, @@ -503,7 +503,7 @@ class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetad # Zero-fill for cudagraphs self.req_id_per_token_buffer.fill_(0) self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_( - torch.from_numpy(req_id_per_token), non_blocking=True + np_to_pinned_tensor(req_id_per_token), non_blocking=True ) req_id_per_token = self.req_id_per_token_buffer[:num_tokens] diff --git a/vllm/v1/attention/backends/utils.py b/vllm/v1/attention/backends/utils.py index 30db5d5f5a8..1b7b8a01a59 100644 --- a/vllm/v1/attention/backends/utils.py +++ b/vllm/v1/attention/backends/utils.py @@ -17,7 +17,7 @@ from typing_extensions import runtime_checkable from vllm.config import VllmConfig, get_layers_from_vllm_config from vllm.utils.math_utils import cdiv -from vllm.utils.torch_utils import async_tensor_h2d +from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d, np_to_pinned_tensor from vllm.v1.kv_cache_interface import KVCacheSpec, MambaSpec if TYPE_CHECKING: @@ -364,8 +364,8 @@ def make_local_attention_virtual_batches( # tensor first, which recovers perf. # Upload the index tensors to the block_table's device up-front so that the # fancy indexing below doesn't implicitly force a synchronous H2D copy. - batch_indices_torch = torch.from_numpy(batch_indices).to(device, non_blocking=True) - block_indices_torch = torch.from_numpy(block_indices).to(device, non_blocking=True) + batch_indices_torch = async_tensor_h2d(batch_indices, device=device) + block_indices_torch = async_tensor_h2d(block_indices, device=device) # Save as a lambda so we can return this for update_block_table make_block_table = lambda block_table: block_table[ @@ -379,8 +379,8 @@ def make_local_attention_virtual_batches( return CommonAttentionMetadata( query_start_loc_cpu=query_start_loc_cpu, - query_start_loc=query_start_loc_cpu.to(device=device, non_blocking=True), - seq_lens=seq_lens_cpu.to(device=device, non_blocking=True), + query_start_loc=async_tensor_h2d(query_start_loc_cpu, device=device), + seq_lens=async_tensor_h2d(seq_lens_cpu, device=device), num_reqs=len(seq_lens_cpu), num_actual_tokens=common_attn_metadata.num_actual_tokens, max_query_len=seqlens_q_local.max(), @@ -808,14 +808,12 @@ def create_fast_prefill_custom_backend( def compute_causal_conv1d_metadata( - query_start_loc_p_cpu: torch.Tensor, - *, - device: torch.device, -): + query_start_loc_p_cpu: torch.Tensor, *, device: torch.device +) -> tuple[dict[int, dict[str, Any]], torch.Tensor, torch.Tensor]: # Needed for causal_conv1d. Use the CPU query_start_loc to avoid DtoH sync. assert query_start_loc_p_cpu.device.type == "cpu" seqlens = query_start_loc_p_cpu.diff() - nums_dict = {} # type: ignore + nums_dict: dict[int, dict[str, Any]] = {} batch_ptr = None token_chunk_offset_ptr = None for BLOCK_M in [8]: # cover all BLOCK_M values @@ -823,7 +821,7 @@ def compute_causal_conv1d_metadata( nums_dict[BLOCK_M] = {} nums_dict[BLOCK_M]["nums"] = nums nums_dict[BLOCK_M]["tot"] = nums.sum().item() - mlist = torch.from_numpy(np.repeat(np.arange(len(nums)), nums)) + mlist = np_to_pinned_tensor(np.repeat(np.arange(len(nums)), nums)) nums_dict[BLOCK_M]["mlist"] = mlist mlist_len = len(nums_dict[BLOCK_M]["mlist"]) nums_dict[BLOCK_M]["mlist_len"] = mlist_len @@ -831,7 +829,7 @@ def compute_causal_conv1d_metadata( offsetlist = [] # type: ignore for idx, num in enumerate(nums): offsetlist.extend(range(num)) - offsetlist = torch.tensor(offsetlist, dtype=torch.int32) + offsetlist = torch.tensor(offsetlist, dtype=torch.int32, pin_memory=PIN_MEMORY) nums_dict[BLOCK_M]["offsetlist"] = offsetlist if batch_ptr is None: @@ -845,16 +843,15 @@ def compute_causal_conv1d_metadata( else: if batch_ptr.nelement() < MAX_NUM_PROGRAMS: batch_ptr.resize_(MAX_NUM_PROGRAMS).fill_(PAD_SLOT_ID) - token_chunk_offset_ptr.resize_( # type: ignore - MAX_NUM_PROGRAMS - ).fill_(PAD_SLOT_ID) + assert token_chunk_offset_ptr is not None + token_chunk_offset_ptr.resize_(MAX_NUM_PROGRAMS).fill_(PAD_SLOT_ID) + assert batch_ptr is not None batch_ptr[0:mlist_len].copy_(mlist, non_blocking=True) - token_chunk_offset_ptr[ # type: ignore - 0:mlist_len - ].copy_(offsetlist, non_blocking=True) + assert token_chunk_offset_ptr is not None + token_chunk_offset_ptr[0:mlist_len].copy_(offsetlist, non_blocking=True) nums_dict[BLOCK_M]["batch_ptr"] = batch_ptr - nums_dict[BLOCK_M]["token_chunk_offset_ptr"] = token_chunk_offset_ptr # type: ignore + nums_dict[BLOCK_M]["token_chunk_offset_ptr"] = token_chunk_offset_ptr return nums_dict, batch_ptr, token_chunk_offset_ptr diff --git a/vllm/v1/kv_offload/cpu/gpu_worker.py b/vllm/v1/kv_offload/cpu/gpu_worker.py index 81545281b64..f4d3869dc1e 100644 --- a/vllm/v1/kv_offload/cpu/gpu_worker.py +++ b/vllm/v1/kv_offload/cpu/gpu_worker.py @@ -14,7 +14,7 @@ from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.triton_utils import HAS_TRITON, triton from vllm.utils.math_utils import cdiv -from vllm.utils.platform_utils import is_pin_memory_available +from vllm.utils.torch_utils import PIN_MEMORY from vllm.v1.kv_offload.base import ( BlockIDsLoadStoreSpec, CanonicalKVCacheRef, @@ -156,7 +156,7 @@ def pin_mmap_region(region: SharedOffloadRegion) -> None: def _new_descriptor_buffers( num_copy_ops: int, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - pin = is_pin_memory_available() + pin = PIN_MEMORY # CUDA cache_kernels.cu requires int64; XPU DMA engine requires uint64. ptr_dtype = torch.uint64 if current_platform.is_xpu() else torch.int64 return ( @@ -482,7 +482,7 @@ class CpuGpuOffloadingHandlers: num_cpu_blocks: int, mmap_region: SharedOffloadRegion | None = None, ): - pin_memory = is_pin_memory_available() + pin_memory = PIN_MEMORY logger.info("Allocating %d CPU tensors...", len(kv_caches.tensors)) self._mmap_region = mmap_region if mmap_region is not None and pin_memory: diff --git a/vllm/v1/pool/metadata.py b/vllm/v1/pool/metadata.py index f772c850f0d..9a9bb2b0e71 100644 --- a/vllm/v1/pool/metadata.py +++ b/vllm/v1/pool/metadata.py @@ -7,9 +7,7 @@ import torch from vllm.pooling_params import PoolingParams from vllm.tasks import PoolingTask -from vllm.utils.platform_utils import is_pin_memory_available - -pin_memory = is_pin_memory_available() +from vllm.utils.torch_utils import PIN_MEMORY @dataclass @@ -134,7 +132,7 @@ class PoolingMetadata: num_scheduled_tokens_cpu = torch.from_numpy(num_scheduled_tokens_np) if query_start_loc_gpu is None: cumsum = torch.zeros( - n_seq + 1, dtype=torch.int64, pin_memory=pin_memory, device="cpu" + n_seq + 1, dtype=torch.int64, pin_memory=PIN_MEMORY, device="cpu" ) torch.cumsum(num_scheduled_tokens_cpu, dim=0, out=cumsum[1:]) cumsum = cumsum.to(device, non_blocking=True) diff --git a/vllm/v1/sample/logits_processor/builtin.py b/vllm/v1/sample/logits_processor/builtin.py index 11a52711d67..d7c9444380b 100644 --- a/vllm/v1/sample/logits_processor/builtin.py +++ b/vllm/v1/sample/logits_processor/builtin.py @@ -7,6 +7,7 @@ import numpy as np import torch from vllm import SamplingParams +from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.sample.logits_processor.interface import ( BatchUpdate, LogitsProcessor, @@ -118,7 +119,6 @@ class MinPLogitsProcessor(LogitsProcessor): class LogitBiasLogitsProcessor(LogitsProcessor): def __init__(self, _, device: torch.device, is_pin_memory: bool): self.device = device - self.pin_memory = is_pin_memory self.biases: dict[int, dict[int, float]] = {} self.bias_tensor: torch.Tensor = torch.tensor(()) @@ -154,9 +154,7 @@ class LogitBiasLogitsProcessor(LogitsProcessor): ) def _device_tensor(self, data: list, dtype: torch.dtype) -> torch.Tensor: - return torch.tensor( - data, device="cpu", dtype=dtype, pin_memory=self.pin_memory - ).to(device=self.device, non_blocking=True) + return async_tensor_h2d(data, device=self.device, dtype=dtype) def apply(self, logits: torch.Tensor) -> torch.Tensor: if self.biases: @@ -170,7 +168,6 @@ class MinTokensLogitsProcessor(LogitsProcessor): ): # index -> (min_toks, output_token_ids, stop_token_ids) self.device = device - self.pin_memory = is_pin_memory self.min_toks: dict[int, tuple[int, Sequence[int], set[int]]] = {} # (req_idx_tensor,eos_tok_id_tensor) @@ -227,9 +224,7 @@ class MinTokensLogitsProcessor(LogitsProcessor): ) def _device_tensor(self, data: list, dtype: torch.dtype) -> torch.Tensor: - return torch.tensor( - data, device="cpu", dtype=dtype, pin_memory=self.pin_memory - ).to(device=self.device, non_blocking=True) + return async_tensor_h2d(data, device=self.device, dtype=dtype) def apply(self, logits: torch.Tensor) -> torch.Tensor: if self.min_toks: @@ -283,8 +278,8 @@ class MinTokensLogitsProcessor(LogitsProcessor): toks_arr = np.concatenate(all_toks) # (row_indices, token_indices) for index_put_ to set -inf. logits_slice = ( - torch.from_numpy(rows_arr).to(self.device, non_blocking=True), - torch.from_numpy(toks_arr).to(self.device, non_blocking=True), + async_tensor_h2d(rows_arr, device=self.device), + async_tensor_h2d(toks_arr, device=self.device), ) logits.index_put_(logits_slice, self.neg_inf_tensor) diff --git a/vllm/v1/sample/ops/penalties.py b/vllm/v1/sample/ops/penalties.py index 241d9de957e..7bc6ec7ab89 100644 --- a/vllm/v1/sample/ops/penalties.py +++ b/vllm/v1/sample/ops/penalties.py @@ -4,8 +4,7 @@ import torch from vllm.model_executor.layers.utils import apply_penalties -from vllm.utils.platform_utils import is_pin_memory_available -from vllm.utils.torch_utils import make_tensor_with_pad +from vllm.utils.torch_utils import PIN_MEMORY, make_tensor_with_pad def apply_all_penalties( @@ -52,6 +51,6 @@ def _convert_to_tensors( pad=vocab_size, device="cpu", dtype=torch.int64, - pin_memory=is_pin_memory_available(), + pin_memory=PIN_MEMORY, ) return output_tokens_tensor.to(device, non_blocking=True) diff --git a/vllm/v1/sample/sampler.py b/vllm/v1/sample/sampler.py index eadc009c254..bb20432a081 100644 --- a/vllm/v1/sample/sampler.py +++ b/vllm/v1/sample/sampler.py @@ -6,7 +6,7 @@ import torch import torch.nn as nn from vllm.config.model import LogprobsMode -from vllm.utils.platform_utils import is_pin_memory_available +from vllm.utils.torch_utils import PIN_MEMORY from vllm.v1.outputs import LogprobsTensors, SamplerOutput from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.sample.ops.bad_words import apply_bad_words @@ -65,7 +65,7 @@ class Sampler(nn.Module): ): super().__init__() self.topk_topp_sampler = TopKTopPSampler(logprobs_mode, use_fp64_gumbel) - self.pin_memory = is_pin_memory_available() + self.pin_memory = PIN_MEMORY self.logprobs_mode = logprobs_mode self.use_fp64_gumbel = use_fp64_gumbel diff --git a/vllm/v1/sample/thinking_budget_state.py b/vllm/v1/sample/thinking_budget_state.py index 8789e6afdc4..d32d1b30296 100644 --- a/vllm/v1/sample/thinking_budget_state.py +++ b/vllm/v1/sample/thinking_budget_state.py @@ -7,7 +7,7 @@ from typing import TYPE_CHECKING, Any import torch from vllm.platforms import current_platform -from vllm.utils.torch_utils import async_tensor_h2d +from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d from vllm.v1.sample.logits_processor.interface import ( BatchUpdate, MoveDirectionality, @@ -22,12 +22,11 @@ def maybe_create_thinking_budget_state_holder( max_num_seqs: int, num_spec_tokens: int, device: torch.device, - is_pin_memory: bool, ) -> "ThinkingBudgetStateHolder | None": if reasoning_config is None: return None return ThinkingBudgetStateHolder( - reasoning_config, max_num_seqs, num_spec_tokens, device, is_pin_memory + reasoning_config, max_num_seqs, num_spec_tokens, device, PIN_MEMORY ) diff --git a/vllm/v1/serial_utils.py b/vllm/v1/serial_utils.py index 204c8bd0e41..bc4619a7eb3 100644 --- a/vllm/v1/serial_utils.py +++ b/vllm/v1/serial_utils.py @@ -33,7 +33,7 @@ from vllm.multimodal.inputs import ( MultiModalSharedField, NestedTensors, ) -from vllm.utils.platform_utils import is_pin_memory_available +from vllm.utils.torch_utils import PIN_MEMORY from vllm.v1.utils import tensor_data logger = init_logger(__name__) @@ -327,7 +327,7 @@ class MsgpackDecoder: oob_tensor_provider: OOBTensorProvider | None = None, ): self.share_mem = share_mem - self.pin_tensors = is_pin_memory_available() + self.pin_tensors = PIN_MEMORY args = () if t is None else (t,) self.decoder = msgpack.Decoder( *args, ext_hook=self.ext_hook, dec_hook=self.dec_hook diff --git a/vllm/v1/simple_kv_offload/worker.py b/vllm/v1/simple_kv_offload/worker.py index c23b44f2917..d33e5f76204 100644 --- a/vllm/v1/simple_kv_offload/worker.py +++ b/vllm/v1/simple_kv_offload/worker.py @@ -8,7 +8,7 @@ import torch from vllm.config import VllmConfig from vllm.logger import init_logger -from vllm.utils.platform_utils import is_pin_memory_available +from vllm.utils.torch_utils import PIN_MEMORY from vllm.v1.simple_kv_offload.copy_backend import DmaCopyBackend from vllm.v1.simple_kv_offload.cuda_mem_ops import pin_tensor from vllm.v1.simple_kv_offload.metadata import ( @@ -149,7 +149,7 @@ class SimpleCPUOffloadWorker: (self.num_cpu_blocks * total_bytes_per_block) / (1024**3), ) - pin_memory = is_pin_memory_available() + pin_memory = PIN_MEMORY if not pin_memory: logger.warning( "Pinned memory not available. CPU offload performance may be degraded." diff --git a/vllm/v1/spec_decode/extract_hidden_states.py b/vllm/v1/spec_decode/extract_hidden_states.py index a0a1f03c716..b6f9eac4dfa 100644 --- a/vllm/v1/spec_decode/extract_hidden_states.py +++ b/vllm/v1/spec_decode/extract_hidden_states.py @@ -12,7 +12,7 @@ from vllm.config import CUDAGraphMode, VllmConfig, get_layers_from_vllm_config from vllm.forward_context import set_forward_context from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.model_loader import get_model -from vllm.utils.platform_utils import is_pin_memory_available +from vllm.utils.torch_utils import PIN_MEMORY from vllm.v1.attention.backend import AttentionMetadataBuilder, CommonAttentionMetadata from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher from vllm.v1.utils import CpuGpuBuffer @@ -58,7 +58,7 @@ class ExtractHiddenStatesProposer: self.backup_next_token_ids = CpuGpuBuffer( max_batch_size, dtype=torch.int32, - pin_memory=is_pin_memory_available(), + pin_memory=PIN_MEMORY, device=device, with_numpy=True, ) @@ -317,7 +317,6 @@ class ExtractHiddenStatesProposer: (batch_size, 1). For each request we either use the sampled token (if valid and not discarded) or a backup token from the request state. """ - num_reqs = gpu_input_batch.num_reqs # Precompute backup token IDs for discarded requests. num_reqs = gpu_input_batch.num_reqs diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index b7c01d3ec1c..bdc10313f4a 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -26,7 +26,7 @@ from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM from vllm.model_executor.models.qwen3_dflash import DFlashQwen3ForCausalLM from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.platforms import current_platform -from vllm.utils.platform_utils import is_pin_memory_available +from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d from vllm.v1.attention.backend import CommonAttentionMetadata from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.attention.backends.triton_attn import TritonAttentionMetadata @@ -228,7 +228,7 @@ class SpecDecodeBaseProposer: self.backup_next_token_ids = CpuGpuBuffer( self.max_batch_size, dtype=torch.int32, - pin_memory=is_pin_memory_available(), + pin_memory=PIN_MEMORY, device=device, with_numpy=True, ) @@ -239,9 +239,7 @@ class SpecDecodeBaseProposer: self._last_draft_probs: torch.Tensor | None = None self._slot_mapping_buffer = torch.zeros( - self.max_positions, - dtype=torch.int64, - device=device, + self.max_positions, dtype=torch.int64, device=device ) # Determine allowed attention backends once during initialization. @@ -1127,7 +1125,7 @@ class SpecDecodeBaseProposer: new_query_start_loc_cpu = torch.zeros( query_start_loc_cpu.shape, dtype=torch.int32, - pin_memory=is_pin_memory_available(), + pin_memory=PIN_MEMORY, ) new_query_start_loc_np = new_query_start_loc_cpu.numpy() np.cumsum(new_num_tokens_per_req_np, out=new_query_start_loc_np[1:]) @@ -1160,11 +1158,11 @@ class SpecDecodeBaseProposer: # q1 + 0, q1 + 1, q1 + 2, q1 + 3, // req 2 # q1 + q2 + 0, q1 + q2 + 1, q1 + q2 + 2] // req 3 token_indices_np = token_offsets + old_query_start_locs_expanded - token_indices = torch.from_numpy(token_indices_np).to(device, non_blocking=True) + token_indices = async_tensor_h2d(token_indices_np, device=device) spec_common_attn_metadata = CommonAttentionMetadata( - query_start_loc=new_query_start_loc_cpu.to(device, non_blocking=True), - seq_lens=new_seq_lens_cpu.to(device, non_blocking=True), + query_start_loc=async_tensor_h2d(new_query_start_loc_cpu, device=device), + seq_lens=async_tensor_h2d(new_seq_lens_cpu, device=device), query_start_loc_cpu=new_query_start_loc_cpu, _seq_lens_cpu=new_seq_lens_cpu, _num_computed_tokens_cpu=common_attn_metadata._num_computed_tokens_cpu, diff --git a/vllm/v1/spec_decode/ngram_proposer_gpu.py b/vllm/v1/spec_decode/ngram_proposer_gpu.py index b8a0116edee..ed544bb27c1 100644 --- a/vllm/v1/spec_decode/ngram_proposer_gpu.py +++ b/vllm/v1/spec_decode/ngram_proposer_gpu.py @@ -545,7 +545,7 @@ def update_ngram_gpu_tensors_incremental( num_tokens = input_batch.num_tokens_no_spec[idx] if num_tokens > 0: token_ids_gpu_tensor[idx, :num_tokens].copy_( - input_batch.token_ids_cpu_tensor[idx, :num_tokens], + input_batch.token_ids_cpu_tensor[idx, :num_tokens].pin_memory(), non_blocking=True, ) @@ -591,7 +591,7 @@ def update_ngram_gpu_tensors_incremental( num_tokens = input_batch.num_tokens_no_spec[new_req_idx] if num_tokens > 0: token_ids_gpu_tensor[new_req_idx, :num_tokens].copy_( - input_batch.token_ids_cpu_tensor[new_req_idx, :num_tokens], + input_batch.token_ids_cpu_tensor[new_req_idx, :num_tokens].pin_memory(), non_blocking=True, ) diff --git a/vllm/v1/structured_output/backend_lm_format_enforcer.py b/vllm/v1/structured_output/backend_lm_format_enforcer.py index 94568b09a7f..bbda96e60b2 100644 --- a/vllm/v1/structured_output/backend_lm_format_enforcer.py +++ b/vllm/v1/structured_output/backend_lm_format_enforcer.py @@ -11,7 +11,7 @@ from transformers import PreTrainedTokenizerBase from vllm.sampling_params import SamplingParams from vllm.utils.import_utils import LazyLoader -from vllm.utils.platform_utils import is_pin_memory_available +from vllm.utils.torch_utils import PIN_MEMORY from vllm.v1.structured_output.backend_types import ( StructuredOutputBackend, StructuredOutputGrammar, @@ -139,7 +139,7 @@ class LMFormatEnforcerBackend(StructuredOutputBackend): (max_num_seqs, (self.vocab_size + 31) // 32), -1, dtype=torch.int32, - pin_memory=is_pin_memory_available(), + pin_memory=PIN_MEMORY, ) def destroy(self): diff --git a/vllm/v1/structured_output/backend_outlines.py b/vllm/v1/structured_output/backend_outlines.py index 71dd5d80648..91627ff154c 100644 --- a/vllm/v1/structured_output/backend_outlines.py +++ b/vllm/v1/structured_output/backend_outlines.py @@ -15,7 +15,7 @@ from regex import escape as regex_escape from vllm.sampling_params import SamplingParams from vllm.utils.import_utils import LazyLoader -from vllm.utils.platform_utils import is_pin_memory_available +from vllm.utils.torch_utils import PIN_MEMORY from vllm.v1.structured_output.backend_types import ( StructuredOutputBackend, StructuredOutputGrammar, @@ -101,7 +101,7 @@ class OutlinesBackend(StructuredOutputBackend): (max_num_seqs, (self.vocab_size + 31) // 32), -1, dtype=torch.int32, - pin_memory=is_pin_memory_available(), + pin_memory=PIN_MEMORY, ) def destroy(self): diff --git a/vllm/v1/structured_output/utils.py b/vllm/v1/structured_output/utils.py index d30dcf26170..cde31e0fd5c 100644 --- a/vllm/v1/structured_output/utils.py +++ b/vllm/v1/structured_output/utils.py @@ -10,7 +10,6 @@ from collections.abc import Callable from concurrent.futures import ThreadPoolExecutor, TimeoutError from typing import TYPE_CHECKING, TypeVar -import numpy as np import regex as re import torch from cachetools import LRUCache @@ -18,7 +17,7 @@ from cachetools import LRUCache import vllm.envs as envs from vllm.logger import init_logger from vllm.utils.import_utils import LazyLoader -from vllm.utils.platform_utils import is_pin_memory_available +from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput if TYPE_CHECKING: @@ -123,11 +122,13 @@ def apply_grammar_bitmask( out_indices = [] # Reorder the bitmask to match the order of the requests in the batch. - sorted_bitmask = np.full( - shape=(logits.shape[0], grammar_bitmask.shape[1]), - fill_value=-1, - dtype=grammar_bitmask.dtype, + sorted_bitmask_tensor = torch.full( + (logits.shape[0], grammar_bitmask.shape[1]), + -1, + dtype=torch.from_numpy(grammar_bitmask[:0]).dtype, + pin_memory=PIN_MEMORY, ) + sorted_bitmask = sorted_bitmask_tensor.numpy() cumulative_index = 0 for req_id in grammar_output.structured_output_request_ids: num_spec_tokens = len(spec_tokens.get(req_id, ())) @@ -138,10 +139,8 @@ def apply_grammar_bitmask( out_indices.append(bitmask_index) cumulative_index += 1 + num_spec_tokens - # Copy async to device as tensor. - grammar_bitmask = torch.from_numpy(sorted_bitmask).to( - logits.device, non_blocking=True - ) + # Copy async to device. + grammar_bitmask = sorted_bitmask_tensor.to(logits.device, non_blocking=True) # If the length of out indices and the logits have the same shape # we don't need to pass indices to the kernel, @@ -154,11 +153,9 @@ def apply_grammar_bitmask( # xgrammar expects a python list of indices but it will actually work with # a tensor. If we copy the tensor ourselves here we can do it in a # non_blocking manner and there should be no cpu sync within xgrammar. - pin_memory = is_pin_memory_available() - index_tensor = torch.tensor( - out_indices, dtype=torch.int32, device="cpu", pin_memory=pin_memory + index_tensor = async_tensor_h2d( + out_indices, dtype=torch.int32, device=logits.device ) - index_tensor = index_tensor.to(logits.device, non_blocking=True) xgr.apply_token_bitmask_inplace(logits, grammar_bitmask, indices=index_tensor) return diff --git a/vllm/v1/utils.py b/vllm/v1/utils.py index ba66358c66f..71ade9c8607 100644 --- a/vllm/v1/utils.py +++ b/vllm/v1/utils.py @@ -31,6 +31,7 @@ from vllm.logger import init_logger from vllm.usage.usage_lib import UsageContext, is_usage_stats_enabled, usage_message from vllm.utils.network_utils import get_open_zmq_ipc_path, get_tcp_uri from vllm.utils.system_utils import decorate_logs, kill_process_tree, set_process_title +from vllm.utils.torch_utils import PIN_MEMORY from vllm.v1.core.sched.output import SchedulerOutput if TYPE_CHECKING: @@ -114,7 +115,7 @@ class CpuGpuBuffer: *size: int | torch.SymInt, dtype: torch.dtype, device: torch.device, - pin_memory: bool, + pin_memory: bool = PIN_MEMORY, with_numpy: bool = True, ) -> None: # these buffers are mutable runtime state, so allocate them as normal diff --git a/vllm/v1/worker/cpu/shm.py b/vllm/v1/worker/cpu/shm.py index 92aa1b5b95f..9399b823ce6 100644 --- a/vllm/v1/worker/cpu/shm.py +++ b/vllm/v1/worker/cpu/shm.py @@ -7,6 +7,8 @@ from typing import Any +import numpy as np + # Patch torch APIs import torch @@ -45,11 +47,14 @@ import vllm.utils.torch_utils as torch_utils def async_tensor_h2d( - data: list, - dtype: torch.dtype, + data: list | np.ndarray | torch.Tensor, device: str | torch.device, - pin_memory: bool = False, + dtype: torch.dtype | None = None, ) -> torch.Tensor: + if isinstance(data, np.ndarray): + data = torch.from_numpy(data) + if isinstance(data, torch.Tensor): + return data.to(dtype=dtype) return torch.tensor(data, dtype=dtype, device="cpu") diff --git a/vllm/v1/worker/gpu/buffer_utils.py b/vllm/v1/worker/gpu/buffer_utils.py index cf5b2c1a2d4..f8336fa0749 100644 --- a/vllm/v1/worker/gpu/buffer_utils.py +++ b/vllm/v1/worker/gpu/buffer_utils.py @@ -36,10 +36,9 @@ def async_copy_to_gpu( assert device is not None out = torch.empty_like(x, device=device) - # Copy directly to GPU — explicit pin_memory() causes sporadic stalls - # under high concurrency due to CUDA driver contention. The driver - # handles the transfer efficiently without manual pinning. - return out.copy_(x, non_blocking=True) + # pin_memory() is no-op if the memory is already pinned. + pinned = x.pin_memory() + return out.copy_(pinned, non_blocking=True) class UvaBuffer: @@ -183,7 +182,7 @@ class StagedWriteTensor: # Special handling for write_contents write_contents = async_tensor_h2d( - self._staged_write_contents, self.dtype, self.device + self._staged_write_contents, device=self.device, dtype=self.dtype ) # Write diffs to the GPU buffer @@ -255,7 +254,7 @@ class FusedStagedWriter: indices_uva = self.indices.copy_to_uva(indices) starts_uva = self.starts.copy_to_uva(starts) cu_lens_uva = self.cu_lens.copy_to_uva(cu_lens) - contents_gpu = async_tensor_h2d(contents, torch.int32, self.device) + contents_gpu = async_tensor_h2d(contents, device=self.device, dtype=torch.int32) _apply_write_kernel[(len(group_ids),)]( output_ptrs, diff --git a/vllm/v1/worker/gpu/mm/encoder_runner.py b/vllm/v1/worker/gpu/mm/encoder_runner.py index 7c813c9b848..d86a166fbbd 100644 --- a/vllm/v1/worker/gpu/mm/encoder_runner.py +++ b/vllm/v1/worker/gpu/mm/encoder_runner.py @@ -49,12 +49,11 @@ class EncoderRunner: @torch.inference_mode() def execute_mm_encoder( - self, - mm_kwargs: list[tuple[str, MultiModalKwargsItem]], + self, mm_kwargs: list[tuple[str, MultiModalKwargsItem]] ) -> list[torch.Tensor]: encoder_outputs: list[torch.Tensor] = [] for modality, num_items, mm_kwargs_batch in group_and_batch_mm_kwargs( - mm_kwargs, device=self.device, pin_memory=False + mm_kwargs, device=self.device, pin_memory=True ): batch_outputs = self.model.embed_multimodal(**mm_kwargs_batch) sanity_check_mm_encoder_outputs(batch_outputs, expected_num_items=num_items) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index a96068dd913..124eb101862 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -46,8 +46,7 @@ from vllm.sequence import IntermediateTensors from vllm.tasks import SupportedTask from vllm.utils.math_utils import cdiv from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib -from vllm.utils.platform_utils import is_pin_memory_available -from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE +from vllm.utils.torch_utils import PIN_MEMORY, STR_DTYPE_TO_TORCH_DTYPE from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput @@ -498,7 +497,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): """Build KV-block zeroing metadata; invoked from gpu_worker.""" self.kv_block_zeroer = KVBlockZeroer( self.device, - is_pin_memory_available(), + pin_memory=PIN_MEMORY, attn_groups_iter=(g for groups in self.attn_groups for g in groups), kernel_block_sizes=self.kernel_block_sizes, cache_dtype=self.cache_config.cache_dtype, diff --git a/vllm/v1/worker/gpu/model_states/whisper.py b/vllm/v1/worker/gpu/model_states/whisper.py index fc2909de037..df432d82ce7 100644 --- a/vllm/v1/worker/gpu/model_states/whisper.py +++ b/vllm/v1/worker/gpu/model_states/whisper.py @@ -170,7 +170,8 @@ class WhisperModelState(ModelState): for_capture: bool, num_reqs: int, ) -> dict[int, tuple[torch.Tensor, np.ndarray]]: - encoder_seq_lens_np = np.zeros(num_reqs, dtype=np.int32) + encoder_seq_lens = torch.zeros(num_reqs, dtype=torch.int32, pin_memory=True) + encoder_seq_lens_np = encoder_seq_lens.numpy() if not for_capture: # During normal execution, use actual encoder lengths. for i, req_id in enumerate(req_ids): @@ -183,9 +184,7 @@ class WhisperModelState(ModelState): # is captured with the correct value for cross-attention. encoder_seq_lens_np[:] = self.max_encoder_len - self.encoder_seq_lens_gpu[:num_reqs].copy_( - torch.from_numpy(encoder_seq_lens_np), non_blocking=True - ) + self.encoder_seq_lens_gpu[:num_reqs].copy_(encoder_seq_lens, non_blocking=True) self.encoder_seq_lens_gpu[num_reqs:].fill_(0) encoder_seq_lens_gpu = self.encoder_seq_lens_gpu[:num_reqs] diff --git a/vllm/v1/worker/gpu_input_batch.py b/vllm/v1/worker/gpu_input_batch.py index 89d69c0bde6..28d1a04b780 100644 --- a/vllm/v1/worker/gpu_input_batch.py +++ b/vllm/v1/worker/gpu_input_batch.py @@ -15,6 +15,7 @@ from vllm.pooling_params import PoolingParams from vllm.sampling_params import SamplingParams, SamplingType from vllm.utils import length_from_prompt_token_ids_or_embeds from vllm.utils.collection_utils import swap_dict_values +from vllm.utils.torch_utils import PIN_MEMORY from vllm.v1.outputs import LogprobsTensors from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates from vllm.v1.sample.logits_processor import ( @@ -95,7 +96,6 @@ class InputBatch: max_model_len: int, max_num_batched_tokens: int, device: torch.device, - pin_memory: bool, vocab_size: int, block_sizes: list[int], # The block_size of each kv cache group kernel_block_sizes: list[int], @@ -112,7 +112,6 @@ class InputBatch: max_num_reqs, num_spec_tokens, device, - pin_memory, ) self.thinking_token_budget_reqs: set[str] = set() self.is_pooling_model = is_pooling_model @@ -120,7 +119,6 @@ class InputBatch: self.max_model_len = max_model_len self.max_num_batched_tokens = max_num_batched_tokens self.device = device - self.pin_memory = pin_memory self.vocab_size = vocab_size self._req_ids: list[str | None] = [] @@ -138,7 +136,10 @@ class InputBatch: ) self.token_ids_cpu = self.token_ids_cpu_tensor.numpy() self.is_token_ids_tensor = torch.zeros( - (max_num_reqs, max_model_len), device="cpu", dtype=bool, pin_memory=False + (max_num_reqs, max_model_len), + device="cpu", + dtype=bool, + pin_memory=False, ) self.is_token_ids = self.is_token_ids_tensor.numpy() # Store prompt embeddings per request to avoid OOM from large upfront @@ -149,21 +150,21 @@ class InputBatch: (max_num_reqs,), device="cpu", dtype=torch.int32, - pin_memory=pin_memory, + pin_memory=PIN_MEMORY, ) self.num_tokens_no_spec = self.num_tokens_no_spec_cpu_tensor.numpy() self.num_prompt_tokens_cpu_tensor = torch.zeros( (max_num_reqs,), device="cpu", dtype=torch.int32, - pin_memory=pin_memory, + pin_memory=PIN_MEMORY, ) self.num_prompt_tokens = self.num_prompt_tokens_cpu_tensor.numpy() self.num_computed_tokens_cpu_tensor = torch.zeros( (max_num_reqs,), device="cpu", dtype=torch.int32, - pin_memory=pin_memory, + pin_memory=PIN_MEMORY, ) self.num_computed_tokens_cpu = self.num_computed_tokens_cpu_tensor.numpy() @@ -172,7 +173,7 @@ class InputBatch: max_num_reqs=max_num_reqs, max_model_len=max_model_len, max_num_batched_tokens=max_num_batched_tokens, - pin_memory=pin_memory, + pin_memory=PIN_MEMORY, device=device, block_sizes=block_sizes, kernel_block_sizes=kernel_block_sizes, @@ -185,7 +186,7 @@ class InputBatch: (max_num_reqs,), dtype=torch.float32, device=device ) self.temperature_cpu_tensor = torch.empty( - (max_num_reqs,), dtype=torch.float32, device="cpu", pin_memory=pin_memory + (max_num_reqs,), dtype=torch.float32, device="cpu", pin_memory=PIN_MEMORY ) self.temperature_cpu = self.temperature_cpu_tensor.numpy() self.greedy_reqs: set[str] = set() @@ -193,14 +194,14 @@ class InputBatch: self.top_p = torch.empty((max_num_reqs,), dtype=torch.float32, device=device) self.top_p_cpu_tensor = torch.empty( - (max_num_reqs,), dtype=torch.float32, device="cpu", pin_memory=pin_memory + (max_num_reqs,), dtype=torch.float32, device="cpu", pin_memory=PIN_MEMORY ) self.top_p_cpu = self.top_p_cpu_tensor.numpy() self.top_p_reqs: set[str] = set() self.top_k = torch.empty((max_num_reqs,), dtype=torch.int32, device=device) self.top_k_cpu_tensor = torch.empty( - (max_num_reqs,), dtype=torch.int32, device="cpu", pin_memory=pin_memory + (max_num_reqs,), dtype=torch.int32, device="cpu", pin_memory=PIN_MEMORY ) self.top_k_cpu = self.top_k_cpu_tensor.numpy() self.top_k_reqs: set[str] = set() @@ -210,7 +211,7 @@ class InputBatch: (max_num_reqs,), dtype=torch.float, device=device ) self.frequency_penalties_cpu_tensor = torch.empty( - (max_num_reqs,), dtype=torch.float, device="cpu", pin_memory=pin_memory + (max_num_reqs,), dtype=torch.float, device="cpu", pin_memory=PIN_MEMORY ) self.frequency_penalties_cpu = self.frequency_penalties_cpu_tensor.numpy() self.frequency_penalties_reqs: set[str] = set() @@ -220,7 +221,7 @@ class InputBatch: (max_num_reqs,), dtype=torch.float, device=device ) self.presence_penalties_cpu_tensor = torch.empty( - (max_num_reqs,), dtype=torch.float, device="cpu", pin_memory=pin_memory + (max_num_reqs,), dtype=torch.float, device="cpu", pin_memory=PIN_MEMORY ) self.presence_penalties_cpu = self.presence_penalties_cpu_tensor.numpy() self.presence_penalties_reqs: set[str] = set() @@ -230,14 +231,14 @@ class InputBatch: (max_num_reqs,), dtype=torch.float, device=device ) self.repetition_penalties_cpu_tensor = torch.empty( - (max_num_reqs,), dtype=torch.float, device="cpu", pin_memory=pin_memory + (max_num_reqs,), dtype=torch.float, device="cpu", pin_memory=PIN_MEMORY ) self.repetition_penalties_cpu = self.repetition_penalties_cpu_tensor.numpy() self.repetition_penalties_reqs: set[str] = set() # Speculative decoding self.num_accepted_tokens_cpu_tensor = torch.ones( - (max_num_reqs,), dtype=torch.int32, device="cpu", pin_memory=pin_memory + (max_num_reqs,), dtype=torch.int32, device="cpu", pin_memory=PIN_MEMORY ) self.num_accepted_tokens_cpu = self.num_accepted_tokens_cpu_tensor.numpy() @@ -963,7 +964,7 @@ class InputBatch: (self.num_reqs, max_prompt_len), device="cpu", dtype=torch.int64, - pin_memory=self.pin_memory, + pin_memory=PIN_MEMORY, ) prompt_token_ids = prompt_token_ids_cpu_tensor.numpy() prompt_token_ids[:] = self.token_ids_cpu[:num_reqs, :max_prompt_len] diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 3221dc46c63..b554542e65d 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -116,8 +116,10 @@ from vllm.utils import length_from_prompt_token_ids_or_embeds from vllm.utils.math_utils import cdiv, round_up from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib from vllm.utils.nvtx_pytorch_hooks import PytHooks -from vllm.utils.platform_utils import is_pin_memory_available, num_compute_units +from vllm.utils.platform_utils import num_compute_units from vllm.utils.torch_utils import ( + PIN_MEMORY, + async_tensor_h2d, get_dtype_size, is_quantized_kv_cache, kv_cache_dtype_str_to_dtype, @@ -441,7 +443,6 @@ class GPUModelRunner( scheduler_config = self.scheduler_config parallel_config = self.parallel_config self.device = device - self.pin_memory = is_pin_memory_available() self.dtype = self.model_config.dtype self.kv_cache_dtype = kv_cache_dtype_str_to_dtype( @@ -666,7 +667,6 @@ class GPUModelRunner( max_model_len=max(self.max_model_len, self.max_encoder_len), max_num_batched_tokens=self.max_num_tokens, device=self.device, - pin_memory=self.pin_memory, vocab_size=self.model_config.get_vocab_size(), block_sizes=[placeholder_block_size], kernel_block_sizes=[placeholder_block_size], @@ -674,7 +674,7 @@ class GPUModelRunner( logitsprocs=build_logitsprocs( self.vllm_config, self.device, - self.pin_memory, + PIN_MEMORY, self.is_pooling_model, custom_logitsprocs, ), @@ -729,7 +729,7 @@ class GPUModelRunner( self.max_num_reqs, dtype=torch.int32, device=self.device ) self.optimistic_seq_lens_cpu = torch.zeros( - self.max_num_reqs, dtype=torch.int32, pin_memory=self.pin_memory + self.max_num_reqs, dtype=torch.int32, pin_memory=PIN_MEMORY ) self.num_computed_tokens = torch.zeros( self.max_num_reqs, dtype=torch.int32, device=self.device @@ -846,7 +846,7 @@ class GPUModelRunner( and self.speculative_config.use_ngram_gpu() ): self._num_valid_draft_tokens_cpu = torch.empty( - self.max_num_reqs, dtype=torch.int32, pin_memory=self.pin_memory + self.max_num_reqs, dtype=torch.int32, pin_memory=PIN_MEMORY ) self._num_valid_draft_tokens_event = torch.cuda.Event() self._num_valid_draft_tokens_copy_stream = torch.cuda.Stream() @@ -857,7 +857,7 @@ class GPUModelRunner( (self.max_num_reqs, 1), dtype=torch.int64, device="cpu", - pin_memory=self.pin_memory, + pin_memory=PIN_MEMORY, ) # Pre-allocated tensor for copying valid sampled token counts to CPU, @@ -879,7 +879,7 @@ class GPUModelRunner( (self.max_num_reqs, self.num_spec_tokens), dtype=torch.int64, device="cpu", - pin_memory=self.pin_memory, + pin_memory=PIN_MEMORY, ) if self.use_async_scheduling: self.valid_sampled_token_count_event = torch.Event() @@ -888,7 +888,7 @@ class GPUModelRunner( self.max_num_reqs, dtype=torch.int32, device="cpu", - pin_memory=self.pin_memory, + pin_memory=PIN_MEMORY, ) # Model weight offloader @@ -1000,7 +1000,6 @@ class GPUModelRunner( *size, dtype=dtype, device=self.device, - pin_memory=self.pin_memory, with_numpy=numpy, ) @@ -1055,7 +1054,7 @@ class GPUModelRunner( token_type_ids.append(ids) token_type_ids_cpu = torch.empty( - sum(seq_lens_cpu), dtype=torch.int32, pin_memory=self.pin_memory + sum(seq_lens_cpu), dtype=torch.int32, pin_memory=PIN_MEMORY ) torch.cat(token_type_ids, out=token_type_ids_cpu) model_kwargs["token_type_ids"] = token_type_ids_cpu.to( @@ -1095,12 +1094,12 @@ class GPUModelRunner( """ self._kv_block_zeroer = KVBlockZeroer( self.device, - self.pin_memory, + pin_memory=PIN_MEMORY, attn_groups_iter=self._kv_cache_spec_attn_group_iterator(), kernel_block_sizes=self._kernel_block_sizes, cache_dtype=self.cache_config.cache_dtype, runner_only_attn_layers=self.runner_only_attn_layers, - static_forward_context=(self.compilation_config.static_forward_context), + static_forward_context=self.compilation_config.static_forward_context, ) def _zero_block_ids(self, block_ids: list[int]) -> None: @@ -1651,7 +1650,7 @@ class GPUModelRunner( for _, _, mm_kwargs_batch in group_and_batch_mm_kwargs( mm_kwargs, device=self.device, - pin_memory=self.pin_memory, + pin_memory=PIN_MEMORY, ): mm_kwargs_combined.update(mm_kwargs_batch) @@ -1807,10 +1806,10 @@ class GPUModelRunner( return # Upload the index tensors asynchronously so the scatter can be non-blocking. sampled_tokens_index_tensor = torch.tensor( - sample_flattened_indices, dtype=torch.int64, pin_memory=self.pin_memory + sample_flattened_indices, dtype=torch.int64, pin_memory=PIN_MEMORY ).to(self.device, non_blocking=True) prev_common_req_indices_tensor = torch.tensor( - prev_indices, dtype=torch.int64, pin_memory=self.pin_memory + prev_indices, dtype=torch.int64, pin_memory=PIN_MEMORY ).to(self.device, non_blocking=True) self.input_ids.gpu.scatter_( dim=0, @@ -1826,10 +1825,10 @@ class GPUModelRunner( assert isinstance(self._draft_token_ids, torch.Tensor) draft_tokens_index_tensor = torch.tensor( - spec_flattened_indices, dtype=torch.int64, pin_memory=self.pin_memory + spec_flattened_indices, dtype=torch.int64, pin_memory=PIN_MEMORY ).to(self.device, non_blocking=True) prev_draft_token_indices_tensor = torch.tensor( - prev_draft_token_indices, dtype=torch.int64, pin_memory=self.pin_memory + prev_draft_token_indices, dtype=torch.int64, pin_memory=PIN_MEMORY ).to(self.device, non_blocking=True) # because input_ids dtype is torch.int32, @@ -2788,21 +2787,16 @@ class GPUModelRunner( # [0, 1, 2, 5, 6, 9] target_logits_indices += self._arange_scratch[: cu_num_draft_tokens[-1]] - # TODO: Optimize the CPU -> GPU copy. - cu_num_draft_tokens = torch.from_numpy(cu_num_draft_tokens).to( - self.device, non_blocking=True + cu_num_draft_tokens = async_tensor_h2d(cu_num_draft_tokens, device=self.device) + cu_num_sampled_tokens = async_tensor_h2d( + cu_num_sampled_tokens, device=self.device ) - cu_num_sampled_tokens = torch.from_numpy(cu_num_sampled_tokens).to( - self.device, non_blocking=True + logits_indices = async_tensor_h2d(logits_indices, device=self.device) + target_logits_indices = async_tensor_h2d( + target_logits_indices, device=self.device ) - logits_indices = torch.from_numpy(logits_indices).to( - self.device, non_blocking=True - ) - target_logits_indices = torch.from_numpy(target_logits_indices).to( - self.device, non_blocking=True - ) - bonus_logits_indices = torch.from_numpy(bonus_logits_indices).to( - self.device, non_blocking=True + bonus_logits_indices = async_tensor_h2d( + bonus_logits_indices, device=self.device ) # Compute the draft token ids. @@ -3012,9 +3006,7 @@ class GPUModelRunner( # Track the current index in mm_kwargs/mm_lora_refs to map groups to request IDs current_item_idx = 0 for modality, num_items, mm_kwargs_batch in group_and_batch_mm_kwargs( - mm_kwargs, - device=self.device, - pin_memory=self.pin_memory, + mm_kwargs, device=self.device, pin_memory=PIN_MEMORY ): batch_outputs: MultiModalEmbeddings @@ -3048,7 +3040,7 @@ class GPUModelRunner( group_and_batch_mm_kwargs( [video_mm_kwargs_item], device=self.device, - pin_memory=self.pin_memory, + pin_memory=PIN_MEMORY, ) ) @@ -3107,7 +3099,10 @@ class GPUModelRunner( mm_embeds = list[torch.Tensor]() is_mm_embed = torch.zeros( - total_num_scheduled_tokens, dtype=torch.bool, device="cpu" + total_num_scheduled_tokens, + dtype=torch.bool, + device="cpu", + pin_memory=PIN_MEMORY, ) req_start_idx = 0 @@ -3515,8 +3510,7 @@ class GPUModelRunner( token_ids_idx_np = np.nonzero(is_token_ids)[0] # Some tokens ids may need to become embeds if token_ids_idx_np.size > 0: - token_ids_idx = torch.from_numpy(token_ids_idx_np) - token_ids_idx = token_ids_idx.to(self.device, non_blocking=True) + token_ids_idx = async_tensor_h2d(token_ids_idx_np, device=self.device) token_ids = self.input_ids.gpu[token_ids_idx] tokens_to_embeds = self.model.embed_input_ids(input_ids=token_ids) self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds @@ -4953,7 +4947,7 @@ class GPUModelRunner( ): indices.append(offset + len(tokens) - 1) offset += num_draft + 1 - indices = torch.tensor(indices, device=self.device) + indices = async_tensor_h2d(indices, device=self.device) hidden_states = sample_hidden_states[indices] draft_token_ids = self.drafter.propose( @@ -5483,8 +5477,8 @@ class GPUModelRunner( continue num_prompt_tokens = len(request.prompt_token_ids) - prompt_token_ids = torch.tensor(request.prompt_token_ids).to( - self.device, non_blocking=True + prompt_token_ids = async_tensor_h2d( + request.prompt_token_ids, device=self.device ) # Set up target LogprobsTensors object. @@ -5651,7 +5645,7 @@ class GPUModelRunner( for _, _, mm_kwargs_batch in group_and_batch_mm_kwargs( [(modality, dummy_mm_item)] * max_items_per_batch, device=self.device, - pin_memory=self.pin_memory, + pin_memory=PIN_MEMORY, ) ) @@ -6995,7 +6989,6 @@ class GPUModelRunner( max_model_len=max_model_len, max_num_batched_tokens=self.max_num_tokens, device=self.device, - pin_memory=self.pin_memory, vocab_size=self.model_config.get_vocab_size(), block_sizes=block_sizes, kernel_block_sizes=kernel_block_sizes, @@ -7430,7 +7423,7 @@ class GPUModelRunner( self.routed_experts_capturer.device_buffer.shape, dtype=self.routed_experts_capturer.device_buffer.dtype, device="cpu", - pin_memory=self.pin_memory, + pin_memory=PIN_MEMORY, ) # ``slot_mapping`` dtype is fixed to int64 by # ``block_table.slot_mapping``; we mirror that here. @@ -7439,7 +7432,7 @@ class GPUModelRunner( (max_tokens,), dtype=torch.int64, device="cpu", - pin_memory=self.pin_memory, + pin_memory=PIN_MEMORY, ) # Private device buffer so the shared ``block_table.slot_mapping`` # can be overwritten by the next ``_prepare_inputs`` while the From a346d589f5932d4234bf5bf8718f10e26d187021 Mon Sep 17 00:00:00 2001 From: Matt <156021403+mawong-amd@users.noreply.github.com> Date: Sat, 20 Jun 2026 23:13:10 -0500 Subject: [PATCH 58/75] [Bugfix] Fix NVFP4/OCP MX MoE emulation (#46254) Signed-off-by: Matthew Wong --- .buildkite/test_areas/lm_eval.yaml | 1 + .../layers/fused_moe/experts/nvfp4_emulation_moe.py | 9 --------- .../layers/fused_moe/experts/ocp_mx_emulation_moe.py | 11 ----------- .../layers/fused_moe/experts/triton_moe.py | 2 +- vllm/model_executor/layers/fused_moe/utils.py | 1 + 5 files changed, 3 insertions(+), 21 deletions(-) diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index 217fc5665c8..a5beb5ea36c 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -109,6 +109,7 @@ steps: - image-build-amd commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - export PYTORCH_ROCM_ARCH=gfx942 # Limit Quark compilation to save time - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt - label: MoE Refactor Integration Test (H100 - TEMPORARY) diff --git a/vllm/model_executor/layers/fused_moe/experts/nvfp4_emulation_moe.py b/vllm/model_executor/layers/fused_moe/experts/nvfp4_emulation_moe.py index de5b45ccb87..d7ed53612e0 100644 --- a/vllm/model_executor/layers/fused_moe/experts/nvfp4_emulation_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/nvfp4_emulation_moe.py @@ -21,7 +21,6 @@ from vllm.model_executor.layers.fused_moe.config import ( FusedMoEQuantConfig, ) from vllm.model_executor.layers.fused_moe.experts.triton_moe import TritonExperts -from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( dequantize_to_dtype, ) @@ -135,14 +134,6 @@ class Nvfp4QuantizationEmulationTritonExperts(TritonExperts): swizzle=False, ) - hidden_states, _ = moe_kernel_quantize_input( - A=hidden_states, - A_scale=self.quant_config.a1_gscale, - quant_dtype="nvfp4", - per_act_token_quant=False, - quantization_emulation=True, - ) - # Activation quantization/dequantization is deferred to # `moe_kernel_quantize_input` in TritonExperts.apply. super().apply( diff --git a/vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py b/vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py index feb8c2ea769..b29e2fde015 100644 --- a/vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py @@ -21,7 +21,6 @@ from vllm.model_executor.layers.fused_moe.config import ( FusedMoEQuantConfig, ) from vllm.model_executor.layers.fused_moe.experts.triton_moe import TritonExperts -from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input from vllm.model_executor.layers.quantization.utils.mxfp4_utils import dequant_mxfp4 from vllm.model_executor.layers.quantization.utils.mxfp6_utils import dequant_mxfp6 from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import ( @@ -155,16 +154,6 @@ class OCP_MXQuantizationEmulationTritonExperts(TritonExperts): w2, self.w2_scale_val, hidden_states.dtype ) - # Apply activation QDQ if needed by the OCP MX scheme - hidden_states, _ = moe_kernel_quantize_input( - A=hidden_states, - A_scale=None, - quant_dtype=self.quant_config.quant_dtype, - per_act_token_quant=False, - ocp_mx_scheme=self.ocp_mx_scheme, - quantization_emulation=True, - ) - # Activation quantization/dequantization is deferred to # `moe_kernel_quantize_input` in TritonExperts.apply. super().apply( diff --git a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py index 0d9b43658f9..abe31e017d5 100644 --- a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py @@ -245,7 +245,7 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular): lora_unquantized_hidden_states = hidden_states hidden_states, a1q_scale = moe_kernel_quantize_input( hidden_states, - self.a1_scale, + self.a1_scale or self.a1_gscale, self.quant_dtype, self.per_act_token_quant, self.block_shape, diff --git a/vllm/model_executor/layers/fused_moe/utils.py b/vllm/model_executor/layers/fused_moe/utils.py index b8c84ad2af2..8866b4f09f2 100644 --- a/vllm/model_executor/layers/fused_moe/utils.py +++ b/vllm/model_executor/layers/fused_moe/utils.py @@ -296,6 +296,7 @@ def moe_kernel_quantize_input( if not quantization_emulation: return _nvfp4_quantize(A, A_scale, is_sf_swizzled_layout=is_scale_swizzled) else: + assert A_scale is not None A = ref_nvfp4_quant_dequant(A, A_scale, block_size=16) return A, None elif quant_dtype == "mxfp4": From 183a430c137db3d5cd0b9025b816f26ee87328e7 Mon Sep 17 00:00:00 2001 From: Ting SUN Date: Sun, 21 Jun 2026 12:06:49 +0700 Subject: [PATCH 59/75] [Bugfix][Model Runner V2] Fix min_tokens off-by-one in the V2 GPU sampler (#46243) Signed-off-by: Ting Sun --- vllm/v1/worker/gpu/sample/logit_bias.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/v1/worker/gpu/sample/logit_bias.py b/vllm/v1/worker/gpu/sample/logit_bias.py index cabb3fc11f8..396f9f509c6 100644 --- a/vllm/v1/worker/gpu/sample/logit_bias.py +++ b/vllm/v1/worker/gpu/sample/logit_bias.py @@ -222,7 +222,7 @@ def _bias_kernel( num_stop_token_ids = tl.load(num_stop_token_ids_ptr + req_state_idx) pos = tl.load(pos_ptr + token_idx) min_len = tl.load(min_lens_ptr + req_state_idx) - if num_stop_token_ids > 0 and pos < min_len: + if num_stop_token_ids > 0 and pos + 1 < min_len: mask = block < num_stop_token_ids stop_token_ids = tl.load( stop_token_ids_ptr + req_state_idx * stop_token_ids_stride + block, From b5495cc5f9099cf77571524a9af88ad26814f324 Mon Sep 17 00:00:00 2001 From: Shifani Rajabose Date: Sun, 21 Jun 2026 02:00:50 -0400 Subject: [PATCH 60/75] Fix memory pointer overflow in Mamba state buffers (#44665) Signed-off-by: Shifani Rajabose Co-authored-by: Kunshang Ji --- vllm/v1/worker/mamba_utils.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index a2718b72607..45166ef9a3a 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -257,9 +257,10 @@ class MambaCopyBuffers: for gid in mamba_group_ids ) * len(copy_funcs) n = max_num_reqs * entries_per_req + return cls( - src_ptrs=make_buffer(n, dtype=torch.int64), - dst_ptrs=make_buffer(n, dtype=torch.int64), + src_ptrs=make_buffer(n, dtype=torch.uint64), + dst_ptrs=make_buffer(n, dtype=torch.uint64), sizes=make_buffer(n, dtype=torch.int32), mamba_group_ids=mamba_group_ids, mamba_spec=mamba_spec, From b80ce9dd2f30913b3b054308a09bd2d86ec6202f Mon Sep 17 00:00:00 2001 From: xiaolinchen <3400259131@qq.com> Date: Sun, 21 Jun 2026 15:11:19 +0800 Subject: [PATCH 61/75] [CI][test] Replace InternVL2-1B with InternVL3-1B in test_pipeline_parallel.py (#46241) Signed-off-by: wentian-byte <192079369+wentian-byte@users.noreply.github.com> Co-authored-by: wentian-byte <192079369+wentian-byte@users.noreply.github.com> --- tests/distributed/test_pipeline_parallel.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/distributed/test_pipeline_parallel.py b/tests/distributed/test_pipeline_parallel.py index d1196b8e0d5..44dc9089dc2 100644 --- a/tests/distributed/test_pipeline_parallel.py +++ b/tests/distributed/test_pipeline_parallel.py @@ -175,7 +175,7 @@ MULTIMODAL_MODELS = { "facebook/chameleon-7b": PPTestSettings.fast(), "adept/fuyu-8b": PPTestSettings.fast(), "zai-org/glm-4v-9b": PPTestSettings.fast(), - "OpenGVLab/InternVL2-1B": PPTestSettings.fast(), + "OpenGVLab/InternVL3-1B": PPTestSettings.fast(), "llava-hf/llava-1.5-7b-hf": PPTestSettings.fast(), "llava-hf/llava-v1.6-mistral-7b-hf": PPTestSettings.fast(), "llava-hf/LLaVA-NeXT-Video-7B-hf": PPTestSettings.fast(), @@ -203,7 +203,7 @@ TEST_MODELS = [ "intfloat/e5-mistral-7b-instruct", "BAAI/bge-multilingual-gemma2", # [MULTIMODAL GENERATION] - "OpenGVLab/InternVL2-1B", + "OpenGVLab/InternVL3-1B", "microsoft/Phi-3.5-vision-instruct", "fixie-ai/ultravox-v0_5-llama-3_2-1b", # [LANGUAGE GENERATION - HYBRID ARCH] From d3ad8e8bcd1a015026981e479d5537549fba3e97 Mon Sep 17 00:00:00 2001 From: Palaiologos1453 <2260891073@qq.com> Date: Sun, 21 Jun 2026 19:30:13 +0800 Subject: [PATCH 62/75] [Bugfix] Defer offload reads while transfers are pending (#46231) Signed-off-by: test test <2260891073@qq.com> --- .../offloading_connector/test_scheduler.py | 103 +++++++++++++++++- .../kv_connector/v1/offloading/scheduler.py | 7 ++ 2 files changed, 106 insertions(+), 4 deletions(-) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index f6011ebac4e..7b8f6119f57 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Iterable +from types import SimpleNamespace from unittest.mock import MagicMock import pytest @@ -1278,11 +1279,11 @@ def test_reset_cache_finalizes_finished_request_with_pending_store( ) finalized: list[str] = [] - runner.manager.on_request_finished.side_effect = ( - lambda req_context: finalized.append(req_context.req_id) + runner.manager.on_request_finished.side_effect = lambda req_context: ( + finalized.append(req_context.req_id) ) - runner.manager.prepare_store.side_effect = ( - lambda keys, req_context: generate_store_output(keys) + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) ) # Decode a couple of blocks and keep every transfer in flight, so the @@ -1314,6 +1315,100 @@ def test_reset_cache_finalizes_finished_request_with_pending_store( assert req_id not in cs._req_status +def test_pending_transfer_defers_prefix_lookup(): + """A request with an in-flight store must not issue a load on re-admission. + + With async scheduling, a preempted request's store can be flushed by the + worker before the scheduler consumes its completion. If the request is + re-admitted in that window, the connector should defer it instead of + looking up offloaded blocks and later asserting when a load is queued while + the store job is still tracked. + """ + scheduler = object.__new__(OffloadingConnectorScheduler) + scheduler.manager = MagicMock(spec=OffloadingManager) + + request = SimpleNamespace(request_id="req-0") + group_state = SimpleNamespace(block_ids=[1, 2, 3]) + req_status = SimpleNamespace( + group_states=[group_state], + transfer_jobs={123}, + ) + scheduler._req_status = {request.request_id: req_status} + + matched_tokens, is_async = scheduler.get_num_new_matched_tokens( + request, + num_computed_tokens=0, + ) + + assert matched_tokens is None + assert is_async is False + assert group_state.block_ids == [] + scheduler.manager.lookup.assert_not_called() + + +def test_async_preempt_readmit_before_transfer_output_is_deferred(request_runner): + """A preempted request can be scheduled again before flush output is read. + + EngineCore.step_with_batch_queue() may schedule a new batch while a prior + preemption batch is still queued. The store completion from jobs_to_flush is + only cleared when that queued output reaches update_from_output(), so the + re-admission path must defer while the scheduler still tracks the store. + """ + block_size = 4 + block_size_factor = 3 + offloaded_block_size = block_size * block_size_factor + + runner = request_runner( + block_size=block_size, + num_gpu_blocks=100, + async_scheduling=True, + block_size_factor=block_size_factor, + ) + free_block_queue = runner.scheduler.kv_cache_manager.block_pool.free_block_queue + num_free_blocks_empty = free_block_queue.num_free_blocks + + req_id = "0" + runner.new_request(token_ids=[0] * offloaded_block_size * 2) + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + + runner.run(decoded_tokens=[0], complete_transfers=False) + runner.run( + decoded_tokens=[0] * (2 * offloaded_block_size - block_size), + complete_transfers=False, + ) + + req_status = runner.connector_scheduler._req_status[req_id] + pending_store_jobs = set(req_status.transfer_jobs) + assert pending_store_jobs + assert all( + runner.connector_scheduler._jobs[jid].is_store for jid in pending_store_jobs + ) + + free_block_queue.num_free_blocks = 0 + preempt_output = runner.scheduler.schedule() + assert preempt_output.preempted_req_ids == {req_id} + assert preempt_output.kv_connector_metadata is not None + assert pending_store_jobs <= preempt_output.kv_connector_metadata.jobs_to_flush + assert req_status.transfer_jobs == pending_store_jobs + + # Simulate the async batch-queue window: schedule again before the + # preemption batch's ModelRunnerOutput is consumed by update_from_output(). + free_block_queue.num_free_blocks = num_free_blocks_empty + assert runner.scheduler.reset_prefix_cache() + runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: len( + key + ) + + readmit_output = runner.scheduler.schedule() + + assert readmit_output.num_scheduled_tokens == {} + assert readmit_output.kv_connector_metadata is not None + assert readmit_output.kv_connector_metadata.load_jobs == {} + assert req_status.transfer_jobs == pending_store_jobs + + @pytest.mark.parametrize("async_scheduling", [True, False]) def test_swa_alignment_skip(request_runner, async_scheduling: bool): """SWA blocks unreachable by the load path are skipped during store. diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 21be16e486f..55277727889 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -647,6 +647,13 @@ class OffloadingConnectorScheduler: for group_state in req_status.group_states: group_state.block_ids.clear() + if req_status.transfer_jobs: + logger.debug( + "Delaying request %s since it still has in-flight transfers", + request.request_id, + ) + return None, False + req_status.update_offload_keys() req_status.num_locally_computed_tokens = num_computed_tokens From b91b7726e068d5eed2aa3bd084cdeee9456da424 Mon Sep 17 00:00:00 2001 From: junkang1991 <97102394+junkang1991@users.noreply.github.com> Date: Sun, 21 Jun 2026 20:55:19 +0800 Subject: [PATCH 63/75] [ROCm][P/D] Support MiniMax-M3 mixed KV layouts in MoRIIO READ mode (#46039) Signed-off-by: Jun Kang Chow Signed-off-by: tjtanaa Co-authored-by: Hongxia Yang Co-authored-by: Tan Pin Siang Co-authored-by: vllmellm Co-authored-by: Chun Fang Co-authored-by: TianDi101 Co-authored-by: functionstackx <47992694+functionstackx@users.noreply.github.com> Co-authored-by: tjtanaa Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../unit/test_moriio_connector.py | 37 ++- .../unit/test_moriio_kv_layout.py | 228 ++++++++++++++++++ .../v1/moriio/moriio_connector.py | 160 ++++++------ .../kv_connector/v1/moriio/moriio_layout.py | 213 ++++++++++++++++ 4 files changed, 566 insertions(+), 72 deletions(-) create mode 100644 tests/v1/kv_connector/unit/test_moriio_kv_layout.py create mode 100644 vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py diff --git a/tests/v1/kv_connector/unit/test_moriio_connector.py b/tests/v1/kv_connector/unit/test_moriio_connector.py index cfac6fa5a36..a8da6cf36d1 100644 --- a/tests/v1/kv_connector/unit/test_moriio_connector.py +++ b/tests/v1/kv_connector/unit/test_moriio_connector.py @@ -36,13 +36,33 @@ from vllm.utils.network_utils import ( get_ip, make_zmq_path, ) -from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, + KVCacheTensor, +) from .utils import create_request, create_scheduler def _make_test_kv_cache_config() -> KVCacheConfig: - return KVCacheConfig(num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]) + layer_names = ["layer0", "layer1", "layer2"] + return KVCacheConfig( + num_blocks=2, + kv_cache_tensors=[KVCacheTensor(size=0, shared_by=layer_names)], + kv_cache_groups=[ + KVCacheGroupSpec( + layer_names=layer_names, + kv_cache_spec=FullAttentionSpec( + block_size=16, + num_kv_heads=4, + head_size=64, + dtype=torch.float16, + ), + ) + ], + ) aiter_available = importlib.util.find_spec("aiter") is not None @@ -175,9 +195,18 @@ class FakeMoRIIOConnectorWorker(MoRIIOConnectorWorker): REMOTE_ENGINE_ID = "remote_engine" def __init__( - self, *args, hand_shake_latency: float = 1.8, kv_cache_layout="HND", **kwargs + self, + vllm_config, + engine_id, + *args, + hand_shake_latency: float = 1.8, + kv_cache_layout="HND", + kv_cache_config=None, + **kwargs, ): - super().__init__(*args, **kwargs) + super().__init__( + vllm_config, engine_id, kv_cache_config or _make_test_kv_cache_config() + ) def create_vllm_config( diff --git a/tests/v1/kv_connector/unit/test_moriio_kv_layout.py b/tests/v1/kv_connector/unit/test_moriio_kv_layout.py new file mode 100644 index 00000000000..5b3219db867 --- /dev/null +++ b/tests/v1/kv_connector/unit/test_moriio_kv_layout.py @@ -0,0 +1,228 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import importlib.util +from types import SimpleNamespace + +import pytest +import torch + +from vllm.platforms import current_platform +from vllm.v1.kv_cache_interface import FullAttentionSpec, MLAAttentionSpec + +aiter_available = importlib.util.find_spec("aiter") is not None +mori_available = importlib.util.find_spec("mori") is not None + +if not (current_platform.is_rocm() and mori_available): + pytest.skip( + "MoRIIOs are only available on ROCm with mori package installed", + allow_module_level=True, + ) + +moriio_layout = importlib.import_module( + "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_layout" +) + + +def _full_spec(block_size: int = 4) -> FullAttentionSpec: + return FullAttentionSpec( + block_size=block_size, + num_kv_heads=2, + head_size=3, + dtype=torch.bfloat16, + ) + + +def _mla_spec(block_size: int = 4) -> MLAAttentionSpec: + return MLAAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=3, + dtype=torch.bfloat16, + ) + + +def _worker( + kv_caches: dict[str, torch.Tensor], + layer_to_spec: dict[str, object], + num_blocks: int = 8, +) -> SimpleNamespace: + return SimpleNamespace( + kv_caches=kv_caches, + layer_to_spec=layer_to_spec, + num_blocks=num_blocks, + block_size=4, + ) + + +def _remote_meta(num_blocks: int = 16) -> SimpleNamespace: + return SimpleNamespace(num_blocks=num_blocks) + + +def test_separated_kv_layout_uses_kv_axis_zero_and_block_axis_one(): + cache = torch.empty((2, 8, 4, 2, 3), dtype=torch.bfloat16) + worker = _worker({"layer": cache}, {"layer": _full_spec()}) + + geometry = moriio_layout.get_layer_transfer_geometry( + "layer", cache, worker.layer_to_spec, remote_num_blocks=16 + ) + assert geometry.block_stride == 24 + assert geometry.local_kv_stride == 192 + assert geometry.remote_kv_stride == 384 + assert geometry.split_kv_regions + + assert moriio_layout.compute_block_transfer_offsets( + "layer", cache, worker.layer_to_spec, [1, 3], [4, 5], _remote_meta().num_blocks + ) == ([48, 144, 432, 528], [192, 240, 960, 1008], [48, 48, 48, 48]) + + +def test_interleaved_kv_layout_uses_block_axis_zero_and_kv_axis_one(): + cache = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16) + worker = _worker({"layer": cache}, {"layer": _full_spec()}) + + geometry = moriio_layout.get_layer_transfer_geometry( + "layer", cache, worker.layer_to_spec, remote_num_blocks=16 + ) + assert geometry.block_stride == 48 + assert geometry.local_kv_stride == 24 + assert geometry.remote_kv_stride == 24 + assert not geometry.split_kv_regions + + assert moriio_layout.compute_block_transfer_offsets( + "layer", cache, worker.layer_to_spec, [1, 3], [4, 5], _remote_meta().num_blocks + ) == ([96, 288], [384, 480], [96, 96]) + + +def test_mla_key_only_layout_transfers_one_slab_per_block(): + cache = torch.empty((8, 4, 3), dtype=torch.bfloat16) + worker = _worker({"layer": cache}, {"layer": _mla_spec()}) + + geometry = moriio_layout.get_layer_transfer_geometry( + "layer", cache, worker.layer_to_spec, remote_num_blocks=16 + ) + assert geometry.block_stride == 12 + assert geometry.local_kv_stride is None + assert geometry.remote_kv_stride is None + assert geometry.transfers_per_block == 1 + + assert moriio_layout.compute_block_transfer_offsets( + "layer", cache, worker.layer_to_spec, [1, 3], [4, 5], _remote_meta().num_blocks + ) == ([24, 72], [96, 120], [24, 24]) + + +def test_mixed_layers_compute_distinct_offsets_per_layer(): + kv_caches = { + "separated": torch.empty((2, 8, 4, 2, 3), dtype=torch.bfloat16), + "interleaved": torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16), + "indexer": torch.empty((8, 4, 3), dtype=torch.bfloat16), + } + worker = _worker( + kv_caches, + { + "separated": _full_spec(), + "interleaved": _full_spec(), + "indexer": _mla_spec(), + }, + ) + + separated = moriio_layout.compute_block_transfer_offsets( + "separated", + kv_caches["separated"], + worker.layer_to_spec, + [1, 3], + [4, 5], + _remote_meta().num_blocks, + ) + interleaved = moriio_layout.compute_block_transfer_offsets( + "interleaved", + kv_caches["interleaved"], + worker.layer_to_spec, + [1, 3], + [4, 5], + _remote_meta().num_blocks, + ) + indexer = moriio_layout.compute_block_transfer_offsets( + "indexer", + kv_caches["indexer"], + worker.layer_to_spec, + [1, 3], + [4, 5], + _remote_meta().num_blocks, + ) + + assert separated != interleaved + assert separated != indexer + assert interleaved != indexer + + +def test_block_id_length_mismatch_raises_value_error(): + cache = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16) + worker = _worker({"layer": cache}, {"layer": _full_spec()}) + + with pytest.raises(ValueError, match="must have the same length"): + moriio_layout.compute_block_transfer_offsets( + "layer", cache, worker.layer_to_spec, [1, 3], [4], _remote_meta().num_blocks + ) + + +def test_registration_regions_do_not_split_interleaved_or_mla_cache(): + separated = torch.empty((2, 8, 4, 2, 3), dtype=torch.bfloat16) + interleaved = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16) + indexer = torch.empty((8, 4, 3), dtype=torch.bfloat16) + worker = _worker( + { + "separated": separated, + "interleaved": interleaved, + "indexer": indexer, + }, + { + "separated": _full_spec(), + "interleaved": _full_spec(), + "indexer": _mla_spec(), + }, + ) + + separated_regions = moriio_layout.iter_layer_registration_regions( + "separated", separated, worker.layer_to_spec + ) + interleaved_regions = moriio_layout.iter_layer_registration_regions( + "interleaved", interleaved, worker.layer_to_spec + ) + indexer_regions = moriio_layout.iter_layer_registration_regions( + "indexer", indexer, worker.layer_to_spec + ) + + assert [region[0].data_ptr() for region in separated_regions] == [ + separated[0].data_ptr(), + separated[1].data_ptr(), + ] + assert separated_regions[0][1] == 8 * 48 + assert separated_regions[1][1] == 8 * 48 + + assert len(interleaved_regions) == 1 + assert interleaved_regions[0][0].data_ptr() == interleaved.data_ptr() + assert interleaved_regions[0][1] == 8 * 2 * 48 + + assert len(indexer_regions) == 1 + assert indexer_regions[0][0].data_ptr() == indexer.data_ptr() + assert indexer_regions[0][1] == 8 * 24 + + +def test_registration_regions_use_layer_num_blocks(): + cache = torch.empty((4, 2, 4, 2, 3), dtype=torch.bfloat16) + worker = _worker({"layer": cache}, {"layer": _full_spec()}, num_blocks=8) + + regions = moriio_layout.iter_layer_registration_regions( + "layer", cache, worker.layer_to_spec + ) + + assert len(regions) == 1 + assert regions[0][1] == 4 * 2 * 48 + + +def test_unsupported_shape_raises_value_error(): + cache = torch.empty((8, 4, 2, 3), dtype=torch.bfloat16) + worker = _worker({"layer": cache}, {"layer": _full_spec()}) + + with pytest.raises(ValueError, match="Unsupported MoRIIO K/V cache shape"): + moriio_layout.get_layer_transfer_geometry("layer", cache, worker.layer_to_spec) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index b5552f72046..a41bb5789f0 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -47,6 +47,14 @@ from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_engine import ( MoRIIOWrapper, MoRIIOWriter, ) +from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_layout import ( + LayerTransferGeometry, + build_layer_to_spec, + compute_block_transfer_offsets, + get_layer_transfer_geometry, + is_mla_cache_layer, + iter_layer_registration_regions, +) from vllm.distributed.parallel_state import ( get_tensor_model_parallel_world_size, get_tp_group, @@ -71,6 +79,7 @@ if TYPE_CHECKING: logger = init_logger(__name__) + try: from mori.io import ( BackendType, @@ -117,7 +126,9 @@ class MoRIIOConnector(KVConnectorBase_V1): self.connector_worker: MoRIIOConnectorWorker | None = None elif role == KVConnectorRole.WORKER: self.connector_scheduler = None - self.connector_worker = MoRIIOConnectorWorker(vllm_config, self.engine_id) + self.connector_worker = MoRIIOConnectorWorker( + vllm_config, self.engine_id, kv_cache_config + ) logger.info( "Initialized MoRIIO Connector,engine_id:%s,role: %s", self.engine_id, @@ -683,7 +694,12 @@ class MoRIIOConnectorScheduler: class MoRIIOConnectorWorker: """Implementation of Worker side methods""" - def __init__(self, vllm_config: VllmConfig, engine_id: str): + def __init__( + self, + vllm_config: VllmConfig, + engine_id: str, + kv_cache_config: "KVCacheConfig", + ): if not is_moriio_available(): raise RuntimeError( "MoRIIO is not available. Please ensure the 'mori' package " @@ -707,6 +723,7 @@ class MoRIIOConnectorWorker: ) self.kv_transfer_config = vllm_config.kv_transfer_config self.is_producer = self.kv_transfer_config.is_kv_producer + self.layer_to_spec = build_layer_to_spec(kv_cache_config) if self.is_producer: set_role(ROLE.PRODUCER) @@ -809,6 +826,8 @@ class MoRIIOConnectorWorker: self.kv_cache_shape = None self.block_shape = None self.kv_element_size = 0 + self.kv_cache_shapes: dict[str, torch.Size] = {} + self.block_lens: dict[str, int] = {} # Map of engine_id -> {agent_name0, agent_name1..}. self._remote_agents: dict[EngineId, set[str]] = {} @@ -1218,51 +1237,86 @@ class MoRIIOConnectorWorker: all_done_future = self._handshake_initiation_executor.submit(wait_all_dp) all_done_future.add_done_callback(request_ready) + def _is_mla_cache_layer(self, layer_name: str) -> bool: + return is_mla_cache_layer(self.layer_to_spec, layer_name) + + def _get_layer_transfer_geometry( + self, layer_name: str, remote_num_blocks: int | None = None + ) -> LayerTransferGeometry: + return get_layer_transfer_geometry( + layer_name, + self.kv_caches[layer_name], + self.layer_to_spec, + remote_num_blocks, + ) + + def _iter_layer_registration_regions( + self, layer_name: str + ) -> list[tuple[torch.Tensor, int]]: + return iter_layer_registration_regions( + layer_name, + self.kv_caches[layer_name], + self.layer_to_spec, + ) + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): """Register the KV Cache data in moriio.""" - _, first_kv_cache = next(iter(kv_caches.items())) + self.kv_caches = kv_caches # layer name to kv cache + self.kv_cache_shapes = { + layer_name: kv_cache.shape for layer_name, kv_cache in kv_caches.items() + } + + first_layer_name, first_kv_cache = next( + ( + (layer_name, kv_cache) + for layer_name, kv_cache in kv_caches.items() + if ( + not self._is_mla_cache_layer(layer_name) + and len(kv_cache.shape) == 5 + and (kv_cache.shape[0] == 2 or kv_cache.shape[1] == 2) + ) + ), + next(iter(kv_caches.items())), + ) kv_elem_size = first_kv_cache.element_size() - use_mla = len(first_kv_cache.shape) == 3 - assert use_mla == self.use_mla + use_mla = self._is_mla_cache_layer(first_layer_name) + first_geometry = self._get_layer_transfer_geometry(first_layer_name) if use_mla: # MLA case. - self.num_blocks = first_kv_cache.shape[0] block_rank = 2 # [block_size, latent_dim] block_shape = first_kv_cache.shape[-block_rank:] - block_size, kv_latent_dim = block_shape - self.slot_size_bytes = kv_elem_size * kv_latent_dim else: - # [2 (k and v), num_blocks, ...] - self.num_blocks = first_kv_cache.shape[1] + # [2, num_blocks, ...] or [num_blocks, 2, ...] block_rank = 3 # [block_size, kv_heads, head_dim] block_shape = first_kv_cache.shape[-block_rank:] - block_size, n_kv_heads, head_dim = block_shape[-3:] - # head size in bytes. - self.slot_size_bytes = ( - kv_elem_size * n_kv_heads * head_dim - ) # 1 token 1 layer size , slot size - assert block_size == self.block_size + self.num_blocks = first_geometry.num_blocks + self.slot_size_bytes = first_geometry.slot_size_bytes + assert first_geometry.block_size == self.block_size # TODO(tms): self.block_len needs to be per-layer for sliding window, # hybrid attn, etc # block size in bytes - self.block_len = kv_elem_size * math.prod(block_shape) + self.block_len = first_geometry.block_len self.kv_cache_shape = first_kv_cache.shape self.block_shape = block_shape self.kv_element_size = kv_elem_size self.dst_num_blocks[self.engine_id] = self.num_blocks - self.kv_caches = kv_caches # layer name to kv cache kv_caches_base_addr = [] caches_data = [] - for cache_or_caches in kv_caches.values(): - cache_list = [cache_or_caches] if use_mla else cache_or_caches - for cache in cache_list: + for layer_name in kv_caches: + geometry = self._get_layer_transfer_geometry(layer_name) + if geometry.block_size != self.block_size: + raise ValueError( + "MoRIIO KV cache block size mismatch for layer " + f"{layer_name}: {geometry.block_size} != {self.block_size}" + ) + self.block_lens[layer_name] = geometry.block_len + for cache, region_len in self._iter_layer_registration_regions(layer_name): base_addr = cache.data_ptr() - region_len = self.num_blocks * self.block_len caches_data.append((base_addr, region_len, cache.device.index, "")) kv_caches_base_addr.append(base_addr) @@ -1275,7 +1329,9 @@ class MoRIIOConnectorWorker: moriio_mem_metadata ) - self.local_kv_cache_size.append(cache.nelement() * cache.element_size()) + self.local_kv_cache_size.append( + kv_cache.nelement() * kv_cache.element_size() + ) self.kv_caches_base_addr[self.engine_id] = kv_caches_base_addr self.num_regions = len(caches_data) @@ -1666,47 +1722,17 @@ class MoRIIOConnectorWorker: Returns: Tuple of (local_offsets, remote_offsets, transfer_sizes) """ - assert self.kv_cache_shape is not None, "KV caches shape not initialized" - is_mla = len(self.kv_cache_shape) == 3 - stride = self.kv_caches[layer_name].stride() - sz = self.kv_caches[layer_name].element_size() - if is_mla: - blknum, blksize, hs = self.kv_cache_shape - hn = 1 - block_stride = stride[0] - else: - _, blknum, blksize, hn, hs = self.kv_cache_shape - local_ktov_stride = stride[0] - block_stride = stride[1] - remote_ktov_stride = block_stride * remote_moriio_meta.num_blocks - - transfer_size_byte = blksize * hn * hs * sz - per_block = 1 if is_mla else 2 - total = len(local_block_ids) * per_block - offset_local = [0] * total - offset_remote = [0] * total - sizes = [transfer_size_byte] * total - - w = 0 - for i, lb in enumerate(local_block_ids): - rb = remote_block_ids[i] - # K - offset_local[w] = sz * (lb * block_stride) - offset_remote[w] = sz * (rb * block_stride) - w += 1 - if not is_mla: - # V - # Handle num_block variations originating from PD (different kv strides) - # TODO: address block_sz differences in heterogeneous TP scenarios - # In MLA, we don't need to consider these two cases. - offset_local[w] = sz * (1 * local_ktov_stride + lb * block_stride) - offset_remote[w] = sz * (1 * remote_ktov_stride + rb * block_stride) - w += 1 - - merged_l, merged_r, merged_s = self.merge_contiguous_blocks( - offset_local, offset_remote, sizes, assume_sorted=False + return compute_block_transfer_offsets( + layer_name=layer_name, + kv_cache=self.kv_caches[layer_name], + layer_to_spec=self.layer_to_spec, + local_block_ids=local_block_ids, + remote_block_ids=remote_block_ids, + remote_num_blocks=remote_moriio_meta.num_blocks, + merge_fn=lambda local, remote, sizes: self.merge_contiguous_blocks( + local, remote, sizes, assume_sorted=False + ), ) - return merged_l, merged_r, merged_s def _read_blocks( self, @@ -1724,15 +1750,13 @@ class MoRIIOConnectorWorker: dp0_engine_id = self.get_engine_name_with_dp(dst_engine_id, 0) sessions, remote_moriio_meta = self._get_built_session(dp0_engine_id) - first_layer = list(self.layer_name_to_local_kv_cache_metadata.keys())[0] - offs = self._compute_block_transfer_offsets( - first_layer, local_block_ids, remote_block_ids, remote_moriio_meta - ) - for layer_name in self.layer_name_to_local_kv_cache_metadata: sess_idx = list(self.layer_name_to_local_kv_cache_metadata.keys()).index( layer_name ) + offs = self._compute_block_transfer_offsets( + layer_name, local_block_ids, remote_block_ids, remote_moriio_meta + ) # TODO : apply multi-session batch-read when moriio support it transfer_status = self.moriio_wrapper.read_remote_data( offs[2], offs[0], offs[1], sessions[sess_idx] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py new file mode 100644 index 00000000000..8a6aced9daa --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py @@ -0,0 +1,213 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Callable, Mapping +from typing import NamedTuple + +import torch + +from vllm.v1.kv_cache_interface import ( + KVCacheConfig, + KVCacheSpec, + MLAAttentionSpec, + SlidingWindowMLASpec, + UniformTypeKVCacheSpecs, +) + + +class LayerTransferGeometry(NamedTuple): + num_blocks: int + block_size: int + block_len: int + slot_size_bytes: int + block_stride: int + local_kv_stride: int | None + remote_kv_stride: int | None + transfers_per_block: int + regions_per_block: int + split_kv_regions: bool + + +def build_layer_to_spec(kv_cache_config: KVCacheConfig) -> dict[str, KVCacheSpec]: + layer_to_spec: dict[str, KVCacheSpec] = {} + for group in kv_cache_config.kv_cache_groups: + group_spec = group.kv_cache_spec + if isinstance(group_spec, UniformTypeKVCacheSpecs): + layer_to_spec.update( + { + layer_name: group_spec.kv_cache_specs[layer_name] + for layer_name in group.layer_names + } + ) + else: + layer_to_spec.update( + {layer_name: group_spec for layer_name in group.layer_names} + ) + return layer_to_spec + + +def is_mla_cache_layer( + layer_to_spec: Mapping[str, KVCacheSpec], layer_name: str +) -> bool: + try: + spec = layer_to_spec[layer_name] + except KeyError as e: + raise ValueError(f"Missing KV cache spec for layer {layer_name}") from e + return isinstance(spec, (MLAAttentionSpec, SlidingWindowMLASpec)) + + +def get_layer_transfer_geometry( + layer_name: str, + kv_cache: torch.Tensor, + layer_to_spec: Mapping[str, KVCacheSpec], + remote_num_blocks: int | None = None, +) -> LayerTransferGeometry: + shape = kv_cache.shape + stride = kv_cache.stride() + element_size = kv_cache.element_size() + is_mla_cache = is_mla_cache_layer(layer_to_spec, layer_name) + + if is_mla_cache and len(shape) == 3: + num_blocks, block_size, latent_dim = shape + slot_size_bytes = latent_dim * element_size + block_len = block_size * slot_size_bytes + return LayerTransferGeometry( + num_blocks=num_blocks, + block_size=block_size, + block_len=block_len, + slot_size_bytes=slot_size_bytes, + block_stride=stride[0], + local_kv_stride=None, + remote_kv_stride=None, + transfers_per_block=1, + regions_per_block=1, + split_kv_regions=False, + ) + + if not is_mla_cache and len(shape) == 5 and shape[0] == 2: + _, num_blocks, block_size, num_kv_heads, head_dim = shape + slot_size_bytes = num_kv_heads * head_dim * element_size + block_len = block_size * slot_size_bytes + remote_kv_stride = stride[1] * (remote_num_blocks or num_blocks) + return LayerTransferGeometry( + num_blocks=num_blocks, + block_size=block_size, + block_len=block_len, + slot_size_bytes=slot_size_bytes, + block_stride=stride[1], + local_kv_stride=stride[0], + remote_kv_stride=remote_kv_stride, + transfers_per_block=2, + regions_per_block=1, + split_kv_regions=True, + ) + + if not is_mla_cache and len(shape) == 5 and shape[1] == 2: + num_blocks, _, block_size, num_kv_heads, head_dim = shape + slot_size_bytes = num_kv_heads * head_dim * element_size + block_len = block_size * slot_size_bytes + return LayerTransferGeometry( + num_blocks=num_blocks, + block_size=block_size, + block_len=block_len, + slot_size_bytes=slot_size_bytes, + block_stride=stride[0], + local_kv_stride=stride[1], + remote_kv_stride=stride[1], + transfers_per_block=2, + regions_per_block=2, + split_kv_regions=False, + ) + + cache_kind = "MLA" if is_mla_cache else "K/V" + raise ValueError( + f"Unsupported MoRIIO {cache_kind} cache shape for layer " + f"{layer_name}: {tuple(shape)}" + ) + + +def iter_layer_registration_regions( + layer_name: str, + kv_cache: torch.Tensor, + layer_to_spec: Mapping[str, KVCacheSpec], +) -> list[tuple[torch.Tensor, int]]: + geometry = get_layer_transfer_geometry(layer_name, kv_cache, layer_to_spec) + region_len = geometry.num_blocks * geometry.regions_per_block * geometry.block_len + if geometry.split_kv_regions: + return [(cache, region_len) for cache in kv_cache] + return [(kv_cache, region_len)] + + +def merge_contiguous_offsets( + offsets_local: list[int], + offsets_remote: list[int], + sizes: list[int], +) -> tuple[list[int], list[int], list[int]]: + if not offsets_local: + return [], [], [] + if not (len(offsets_local) == len(offsets_remote) == len(sizes)): + raise ValueError("Input list lengths mismatch") + + rows = sorted(zip(offsets_local, offsets_remote, sizes), key=lambda row: row[0]) + merged: list[list[int]] = [] + for local, remote, size in rows: + if ( + merged + and local == merged[-1][0] + merged[-1][2] + and remote == merged[-1][1] + merged[-1][2] + ): + merged[-1][2] += size + else: + merged.append([local, remote, size]) + + return ( + [row[0] for row in merged], + [row[1] for row in merged], + [row[2] for row in merged], + ) + + +def compute_block_transfer_offsets( + layer_name: str, + kv_cache: torch.Tensor, + layer_to_spec: Mapping[str, KVCacheSpec], + local_block_ids: list[int], + remote_block_ids: list[int], + remote_num_blocks: int, + merge_fn: Callable[ + [list[int], list[int], list[int]], tuple[list[int], list[int], list[int]] + ] = merge_contiguous_offsets, +) -> tuple[list[int], list[int], list[int]]: + if len(local_block_ids) != len(remote_block_ids): + raise ValueError( + "local_block_ids and remote_block_ids must have the same length: " + f"{len(local_block_ids)} != {len(remote_block_ids)}" + ) + geometry = get_layer_transfer_geometry( + layer_name, kv_cache, layer_to_spec, remote_num_blocks + ) + element_size = kv_cache.element_size() + transfer_size_byte = geometry.block_len + per_block = geometry.transfers_per_block + total = len(local_block_ids) * per_block + offset_local = [0] * total + offset_remote = [0] * total + sizes = [transfer_size_byte] * total + + w = 0 + for lb, rb in zip(local_block_ids, remote_block_ids): + offset_local[w] = element_size * (lb * geometry.block_stride) + offset_remote[w] = element_size * (rb * geometry.block_stride) + w += 1 + if per_block == 2: + assert geometry.local_kv_stride is not None + assert geometry.remote_kv_stride is not None + offset_local[w] = element_size * ( + geometry.local_kv_stride + lb * geometry.block_stride + ) + offset_remote[w] = element_size * ( + geometry.remote_kv_stride + rb * geometry.block_stride + ) + w += 1 + + return merge_fn(offset_local, offset_remote, sizes) From 3e6e33526da729fbe30ebec86be9049e1899ce67 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Sun, 21 Jun 2026 22:37:10 +0800 Subject: [PATCH 64/75] [Disagg] return routed_experts on streaming generate responses (#44638) Signed-off-by: aoshen02 Co-authored-by: Claude Opus 4.8 (1M context) Co-authored-by: Roger Wang --- vllm/entrypoints/serve/disagg/protocol.py | 1 + vllm/entrypoints/serve/disagg/serving.py | 9 +++++++++ 2 files changed, 10 insertions(+) diff --git a/vllm/entrypoints/serve/disagg/protocol.py b/vllm/entrypoints/serve/disagg/protocol.py index 7e776ae7178..2e98f5e811c 100644 --- a/vllm/entrypoints/serve/disagg/protocol.py +++ b/vllm/entrypoints/serve/disagg/protocol.py @@ -181,6 +181,7 @@ class GenerateResponseStreamChoice(BaseModel): logprobs: ChatCompletionLogProbs | None = None finish_reason: str | None = None token_ids: list[int] | None = None + routed_experts: str | None = None class GenerateStreamResponse(BaseModel): diff --git a/vllm/entrypoints/serve/disagg/serving.py b/vllm/entrypoints/serve/disagg/serving.py index 0bb29c68d01..5031627ea01 100644 --- a/vllm/entrypoints/serve/disagg/serving.py +++ b/vllm/entrypoints/serve/disagg/serving.py @@ -400,6 +400,14 @@ class ServingTokens(OpenAIServing): else: logprobs = None + routed_experts_b64 = None + if output.routed_experts is not None: + buf = io.BytesIO() + np.save(buf, output.routed_experts) + routed_experts_b64 = base64.b64encode(buf.getvalue()).decode( + "ascii" + ) + chunk = GenerateStreamResponse( request_id=request_id, choices=[ @@ -408,6 +416,7 @@ class ServingTokens(OpenAIServing): logprobs=logprobs, finish_reason=finish_reason, token_ids=as_list(delta_token_ids), + routed_experts=routed_experts_b64, ) ], ) From 2cac89f9da865dfaceb6d337d97aaff5c9195e48 Mon Sep 17 00:00:00 2001 From: Alex Steiner Date: Sun, 21 Jun 2026 07:45:14 -0700 Subject: [PATCH 65/75] [Spec Decode] Support mixed KV page sizes for DFlash (#45181) Signed-off-by: Alex Steiner Signed-off-by: Giancarlo Delfin Signed-off-by: Yifan Qiao Co-authored-by: Claude Opus 4.8 Co-authored-by: Giancarlo Delfin Co-authored-by: Yifan Qiao --- tests/v1/core/test_kv_cache_utils.py | 109 +++++++- tests/v1/worker/test_attn_utils.py | 242 ++++++++++++++++++ vllm/v1/attention/backend.py | 32 +++ vllm/v1/core/kv_cache_utils.py | 33 ++- vllm/v1/kv_cache_interface.py | 16 +- vllm/v1/worker/gpu/attn_utils.py | 118 ++++++--- vllm/v1/worker/gpu_model_runner.py | 69 ++--- .../worker/kv_connector_model_runner_mixin.py | 33 +-- 8 files changed, 511 insertions(+), 141 deletions(-) create mode 100644 tests/v1/worker/test_attn_utils.py diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 3be24d7fb34..3f5b7a12433 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -117,6 +117,7 @@ def new_kv_cache_spec( page_size_padded=None, sliding_window=None, attention_chunk_size=None, + indexes_kv_by_block_stride=False, ): return FullAttentionSpec( block_size=block_size, @@ -126,6 +127,7 @@ def new_kv_cache_spec( page_size_padded=page_size_padded, sliding_window=sliding_window, attention_chunk_size=attention_chunk_size, + indexes_kv_by_block_stride=indexes_kv_by_block_stride, ) @@ -136,6 +138,7 @@ def new_sliding_window_spec( dtype=torch.float32, page_size_padded=None, sliding_window=1, + indexes_kv_by_block_stride=False, ): return SlidingWindowSpec( block_size=block_size, @@ -144,6 +147,7 @@ def new_sliding_window_spec( dtype=dtype, page_size_padded=page_size_padded, sliding_window=sliding_window, + indexes_kv_by_block_stride=indexes_kv_by_block_stride, ) @@ -1799,16 +1803,38 @@ def test_get_kv_cache_config_one_worker(): ], ) - # different hidden size that cannot be aligned by using different block size + # different hidden size that cannot be aligned by using different block size, + # but can be aligned by padding the smaller physical page. + swa_spec = new_sliding_window_spec(head_size=96, indexes_kv_by_block_stride=True) kv_cache_specs_hybrid = { - "layer_1": new_kv_cache_spec(head_size=64), - "layer_2": new_sliding_window_spec(head_size=96), + "layer_1": new_kv_cache_spec(head_size=64, indexes_kv_by_block_stride=True), + "layer_2": swa_spec, } - with pytest.raises(NotImplementedError): - get_kv_cache_configs( - vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] - )[0] + kv_cache_config_hybrid = get_kv_cache_configs( + vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] + )[0] + padded_page_size = swa_spec.page_size_bytes + assert kv_cache_config_hybrid == KVCacheConfig( + num_blocks=42, + kv_cache_tensors=[ + KVCacheTensor(size=padded_page_size * 42, shared_by=["layer_1", "layer_2"]), + ], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer_1"], + new_kv_cache_spec( + head_size=64, + page_size_padded=padded_page_size, + indexes_kv_by_block_stride=True, + ), + ), + KVCacheGroupSpec( + ["layer_2"], + new_sliding_window_spec(head_size=96, indexes_kv_by_block_stride=True), + ), + ], + ) # Test num_gpu_blocks_override vllm_config.cache_config.num_gpu_blocks_override = 16 @@ -2322,6 +2348,75 @@ def test_check_enough_kv_cache_memory_respects_num_gpu_blocks_override(): get_kv_cache_configs(vllm_config, [kv_cache_specs], [large_available_memory]) +def test_unify_kv_cache_page_size_uses_padding_for_non_divisible_sizes(): + """DFlash drafters can have a smaller head size than the target model. + + For example, MiMo uses 192-dim target KV heads while its DFlash draft uses + 128-dim KV heads. The resulting page sizes are 3:2 rather than an integer + block-size multiple, so the smaller page must be padded instead. + """ + # Both layers' backends opt into the padded-page strided view (e.g. + # FlashAttention / its DiffKV subclass), so padding is allowed. + target_spec = new_kv_cache_spec( + block_size=16, + num_kv_heads=1, + head_size=192, + dtype=torch.bfloat16, + indexes_kv_by_block_stride=True, + ) + draft_spec = new_sliding_window_spec( + block_size=16, + num_kv_heads=1, + head_size=128, + dtype=torch.bfloat16, + sliding_window=1024, + indexes_kv_by_block_stride=True, + ) + + unified_specs = kv_cache_utils.unify_kv_cache_spec_page_size( + { + "target_attn": target_spec, + "draft_attn": draft_spec, + } + ) + + assert unified_specs["target_attn"] == target_spec + unified_draft_spec = unified_specs["draft_attn"] + assert unified_draft_spec.block_size == draft_spec.block_size + assert unified_draft_spec.real_page_size_bytes == draft_spec.real_page_size_bytes + assert unified_draft_spec.page_size_padded == target_spec.page_size_bytes + assert unified_draft_spec.page_size_bytes == target_spec.page_size_bytes + + +def test_unify_kv_cache_page_size_padding_requires_backend_support(): + """Padding is gated on the backend declaring ``indexes_kv_by_block_stride``. + + A backend that does not support the strided padded-page view must raise + rather than silently padding (and misreading KV at runtime). + """ + target_spec = new_kv_cache_spec( + block_size=16, + num_kv_heads=1, + head_size=192, + dtype=torch.bfloat16, + indexes_kv_by_block_stride=True, + ) + # The non-divisible draft layer needs padding but its backend does not + # support the strided padded-page view -> must raise, not silently pad. + draft_spec = new_sliding_window_spec( + block_size=16, + num_kv_heads=1, + head_size=128, + dtype=torch.bfloat16, + sliding_window=1024, + indexes_kv_by_block_stride=False, + ) + specs = {"target_attn": target_spec, "draft_attn": draft_spec} + + with pytest.raises(NotImplementedError): + kv_cache_utils.unify_kv_cache_spec_page_size(specs) + + def test_unify_hybrid_kv_cache_specs(): # 1. has_full_attention and has_sliding_window before_spec_1 = new_kv_cache_spec() diff --git a/tests/v1/worker/test_attn_utils.py b/tests/v1/worker/test_attn_utils.py new file mode 100644 index 00000000000..7e65d650f7e --- /dev/null +++ b/tests/v1/worker/test_attn_utils.py @@ -0,0 +1,242 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVQuantMode +from vllm.v1.worker.gpu.attn_utils import _reshape_kv_cache +from vllm.v1.worker.utils import AttentionGroup + + +class FakeFlashAttentionBackend: + @staticmethod + def get_kv_cache_shape( + num_blocks: int, + block_size: int, + num_kv_heads: int, + head_size: int, + cache_dtype_str: str = "auto", + ) -> tuple[int, ...]: + return (num_blocks, 2, block_size, num_kv_heads, head_size) + + @staticmethod + def get_kv_cache_stride_order( + include_num_layers_dimension: bool = False, + ) -> tuple[int, ...]: + assert not include_num_layers_dimension + return (0, 1, 2, 3, 4) + + +class FakeHNDFlashAttentionBackend(FakeFlashAttentionBackend): + @staticmethod + def get_kv_cache_stride_order( + include_num_layers_dimension: bool = False, + ) -> tuple[int, ...]: + assert not include_num_layers_dimension + return (0, 1, 3, 2, 4) + + +def test_reshape_padded_flash_attention_kv_cache_strides_by_page(): + num_blocks = 3 + spec = FullAttentionSpec( + block_size=16, + num_kv_heads=1, + head_size=2, + dtype=torch.float32, + page_size_padded=384, + ) + assert spec.real_page_size_bytes == 256 + + raw_tensors = { + "layer": torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8) + } + attn_groups = [ + AttentionGroup( + backend=FakeFlashAttentionBackend, + layer_names=["layer"], + kv_cache_spec=spec, + kv_cache_group_id=0, + ) + ] + + kv_cache = _reshape_kv_cache( + attn_groups, + raw_tensors, + "auto", + [spec.block_size], + {}, + )["layer"] + + assert kv_cache.shape == (num_blocks, 2, 16, 1, 2) + assert kv_cache.stride(0) == spec.page_size_bytes // 4 + assert kv_cache.stride(1) == spec.real_page_size_bytes // 2 // 4 + assert kv_cache[1, 0].storage_offset() == spec.page_size_bytes // 4 + assert ( + kv_cache[1, 1].storage_offset() + == (spec.page_size_bytes + spec.real_page_size_bytes // 2) // 4 + ) + + +def test_reshape_padded_hnd_flash_attention_kv_cache_strides_by_page(): + num_blocks = 3 + spec = FullAttentionSpec( + block_size=16, + num_kv_heads=3, + head_size=2, + dtype=torch.float32, + page_size_padded=1024, + ) + assert spec.real_page_size_bytes == 768 + + raw_tensors = { + "layer": torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8) + } + attn_groups = [ + AttentionGroup( + backend=FakeHNDFlashAttentionBackend, + layer_names=["layer"], + kv_cache_spec=spec, + kv_cache_group_id=0, + ) + ] + + kv_cache = _reshape_kv_cache( + attn_groups, + raw_tensors, + "auto", + [spec.block_size], + {}, + )["layer"] + + assert kv_cache.shape == (num_blocks, 2, 16, 3, 2) + assert kv_cache.stride(0) == spec.page_size_bytes // 4 + assert kv_cache.stride(1) == spec.real_page_size_bytes // 2 // 4 + assert kv_cache.stride(2) == 2 + assert kv_cache.stride(3) == spec.block_size * spec.head_size + assert kv_cache[1, 0].storage_offset() == spec.page_size_bytes // 4 + assert ( + kv_cache[1, 1].storage_offset() + == (spec.page_size_bytes + spec.real_page_size_bytes // 2) // 4 + ) + assert ( + kv_cache[1, 1, 3, 2].storage_offset() + == ( + spec.page_size_bytes + + spec.real_page_size_bytes // 2 + + 3 * spec.head_size * 4 + + 2 * spec.block_size * spec.head_size * 4 + ) + // 4 + ) + + +class FakeDiffKVBackend: + @staticmethod + def get_kv_cache_shape( + num_blocks: int, + block_size: int, + num_kv_heads: int, + head_size: int, + cache_dtype_str: str = "auto", + ) -> tuple[int, ...]: + return (num_blocks, block_size, num_kv_heads, head_size * 2) + + @staticmethod + def get_kv_cache_stride_order( + include_num_layers_dimension: bool = False, + ) -> tuple[int, ...]: + assert not include_num_layers_dimension + return (0, 1, 2, 3) + + +def test_reshape_padded_diff_kv_cache_does_not_infer_kv_dim(): + num_blocks = 3 + spec = FullAttentionSpec( + block_size=16, + num_kv_heads=1, + head_size=2, + dtype=torch.float32, + page_size_padded=384, + ) + + raw_tensors = { + "layer": torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8) + } + attn_groups = [ + AttentionGroup( + backend=FakeDiffKVBackend, + layer_names=["layer"], + kv_cache_spec=spec, + kv_cache_group_id=0, + ) + ] + + kv_cache = _reshape_kv_cache( + attn_groups, + raw_tensors, + "auto", + [spec.block_size], + {}, + )["layer"] + + assert kv_cache.shape == (num_blocks, 16, 1, 4) + assert kv_cache.stride(0) == spec.page_size_bytes // 4 + assert kv_cache.stride(1) == 4 + + +class FakePerTokenScaleBackend: + @staticmethod + def get_kv_cache_shape( + num_blocks: int, + block_size: int, + num_kv_heads: int, + head_size: int, + cache_dtype_str: str = "auto", + ) -> tuple[int, ...]: + return (num_blocks, 2, block_size, num_kv_heads, head_size + 4) + + @staticmethod + def get_kv_cache_stride_order( + include_num_layers_dimension: bool = False, + ) -> tuple[int, ...]: + assert not include_num_layers_dimension + return (0, 1, 2, 3, 4) + + +def test_reshape_padded_quantized_kv_cache_preserves_scale_stride(): + num_blocks = 3 + spec = FullAttentionSpec( + block_size=16, + num_kv_heads=1, + head_size=4, + dtype=torch.int8, + kv_quant_mode=KVQuantMode.INT8_PER_TOKEN_HEAD, + page_size_padded=384, + ) + assert spec.real_page_size_bytes == 128 + assert spec.page_size_bytes == 384 + + raw_tensors = { + "layer": torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8) + } + attn_groups = [ + AttentionGroup( + backend=FakePerTokenScaleBackend, + layer_names=["layer"], + kv_cache_spec=spec, + kv_cache_group_id=0, + ) + ] + + kv_cache = _reshape_kv_cache( + attn_groups, + raw_tensors, + "int8_per_token_head", + [spec.block_size], + {}, + )["layer"] + + assert kv_cache.shape == (num_blocks, 2, 16, 1, 8) + assert kv_cache.stride(0) == spec.page_size_bytes + assert kv_cache.stride(1) == 16 * 1 * 8 + assert kv_cache[1, 1].storage_offset() == spec.page_size_bytes + 16 * 1 * 8 diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 03a203a1bcf..ebf607b65a7 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -201,6 +201,38 @@ class AttentionBackend(ABC): return min(s.base if isinstance(s, MultipleOf) else s for s in supported_sizes) + @classmethod + def indexes_kv_by_block_stride(cls) -> bool: + """Whether the backend reads KV pages by the runtime block stride. + + True when ``num_blocks`` is the outermost physical dimension of the KV + cache, so the backend tolerates a non-contiguous block dim. This gates + page size padding and cross-layer uniform KV layout. + + Returns: + True if the backend's physical KV layout is num-blocks-first. False + otherwise, including when the backend does not define a layered + stride order. + """ + try: + kv_cache_stride_order = cls.get_kv_cache_stride_order( + include_num_layers_dimension=False + ) + layered_kv_cache_stride_order = cls.get_kv_cache_stride_order( + include_num_layers_dimension=True + ) + except (AttributeError, NotImplementedError): + return False + + # Check that attention backend includes a layers dimension. + if len(layered_kv_cache_stride_order) != len(kv_cache_stride_order) + 1: + return False + + # stride_order[0] == 0 means num_layers stays first in physical + # layout (identity permutation), so indexing by block stride is + # not supported. + return layered_kv_cache_stride_order[0] != 0 + @classmethod def is_mla(cls) -> bool: return False diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 4e1d28d7d5d..95b8fba4ccf 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -20,6 +20,7 @@ from vllm.utils.math_utils import cdiv, round_up from vllm.utils.mem_utils import format_gib from vllm.utils.torch_utils import get_dtype_size from vllm.v1.kv_cache_interface import ( + AttentionSpec, ChunkedLocalAttentionSpec, FullAttentionSpec, HiddenStateCacheSpec, @@ -1029,9 +1030,14 @@ def unify_kv_cache_spec_page_size( ) -> dict[str, KVCacheSpec]: """ Unify the page size of the given KVCacheSpec. If the page size of all layers - are the same, return the original KVCacheSpec. If not same, unify the page - size by increasing the block size of layers with smaller page size. Raise - NotImplementedError if failed to unify the page size. + are the same, return the original KVCacheSpec. If not same, first try to + unify page size by increasing the block size of layers with smaller page + size. If a smaller attention page does not evenly divide the maximum page + size, keep its logical block size and pad its physical page instead --- but + only for attention layers whose backend opts in via + ``AttentionSpec.indexes_kv_by_block_stride`` (the padded page is read through + a strided view, which not every backend handles). Raise NotImplementedError + if failed to unify the page size. Args: kv_cache_spec: The KVCacheSpec of each attention layer in the model @@ -1051,14 +1057,23 @@ def unify_kv_cache_spec_page_size( new_kv_cache_spec[layer_name] = layer_spec else: layer_page_size = layer_spec.page_size_bytes - if max_page_size % layer_page_size != 0: + if max_page_size % layer_page_size == 0: + ratio = max_page_size // layer_page_size + new_block_size = layer_spec.block_size * ratio + new_spec = replace(layer_spec, block_size=new_block_size) + elif ( + isinstance(layer_spec, AttentionSpec) + and layer_spec.indexes_kv_by_block_stride + ): + new_spec = replace(layer_spec, page_size_padded=max_page_size) + else: raise NotImplementedError( - "The page size of the layer is not divisible by the " - "maximum page size. Cannot unify by adjusting block_size." + f"Layer {layer_name}: page size is not divisible by the " + "maximum page size and cannot be padded. Padding is only " + "supported for attention layers whose backend indexes KV " + "pages by the block stride (indexes_kv_by_block_stride is " + "True)." ) - ratio = max_page_size // layer_page_size - new_block_size = layer_spec.block_size * ratio - new_spec = replace(layer_spec, block_size=new_block_size) assert new_spec.page_size_bytes == max_page_size new_kv_cache_spec[layer_name] = new_spec return new_kv_cache_spec diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 2e779b2c2a4..5a2a5c5e298 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -163,6 +163,7 @@ class AttentionSpec(KVCacheSpec): dtype: torch.dtype kv_quant_mode: KVQuantMode = KVQuantMode.NONE page_size_padded: int | None = None + indexes_kv_by_block_stride: bool = False @property def page_size_bytes(self) -> int: @@ -283,6 +284,7 @@ class FullAttentionSpec(AttentionSpec): dtype=specs[0].dtype, kv_quant_mode=specs[0].kv_quant_mode, page_size_padded=specs[0].page_size_padded, + indexes_kv_by_block_stride=specs[0].indexes_kv_by_block_stride, sliding_window=cls.merge_window_sizes(sliding_window), attention_chunk_size=cls.merge_window_sizes(attention_chunk_size), # If any layer in the group is non-causal, treat the group as @@ -403,13 +405,16 @@ class MLAAttentionSpec(FullAttentionSpec): cache_dtype_str_set = set(spec.cache_dtype_str for spec in specs) compress_ratio_set = set(spec.compress_ratio for spec in specs) model_version_set = set(spec.model_version for spec in specs) + block_stride_set = set(spec.indexes_kv_by_block_stride for spec in specs) assert ( len(cache_dtype_str_set) == 1 and len(compress_ratio_set) == 1 and len(model_version_set) == 1 + and len(block_stride_set) == 1 ), ( "All attention layers in the same KV cache group must use the same " - "quantization method, compress ratio, and model version." + "quantization method, compress ratio, model version, and KV block " + "stride indexing." ) return cls( block_size=specs[0].block_size, @@ -418,6 +423,7 @@ class MLAAttentionSpec(FullAttentionSpec): dtype=specs[0].dtype, kv_quant_mode=specs[0].kv_quant_mode, page_size_padded=specs[0].page_size_padded, + indexes_kv_by_block_stride=block_stride_set.pop(), cache_dtype_str=cache_dtype_str_set.pop(), compress_ratio=compress_ratio_set.pop(), model_version=model_version_set.pop(), @@ -584,15 +590,17 @@ class SlidingWindowMLASpec(SlidingWindowSpec): compress_ratio_set = set(spec.compress_ratio for spec in specs) model_version_set = set(spec.model_version for spec in specs) sliding_window_set = set(spec.sliding_window for spec in specs) + block_stride_set = set(spec.indexes_kv_by_block_stride for spec in specs) assert ( len(cache_dtype_str_set) == 1 and len(compress_ratio_set) == 1 and len(model_version_set) == 1 and len(sliding_window_set) == 1 + and len(block_stride_set) == 1 ), ( "All attention layers in the same KV cache group must use the same " - "quantization method, compress ratio, model version and sliding " - "window size." + "quantization method, compress ratio, model version, sliding " + "window size, and KV block stride indexing." ) return cls( block_size=specs[0].block_size, @@ -600,6 +608,7 @@ class SlidingWindowMLASpec(SlidingWindowSpec): head_size=specs[0].head_size, dtype=specs[0].dtype, page_size_padded=specs[0].page_size_padded, + indexes_kv_by_block_stride=block_stride_set.pop(), sliding_window=sliding_window_set.pop(), cache_dtype_str=cache_dtype_str_set.pop(), compress_ratio=compress_ratio_set.pop(), @@ -711,6 +720,7 @@ class SinkFullAttentionSpec(FullAttentionSpec): dtype=specs[0].dtype, kv_quant_mode=specs[0].kv_quant_mode, page_size_padded=specs[0].page_size_padded, + indexes_kv_by_block_stride=specs[0].indexes_kv_by_block_stride, sliding_window=cls.merge_window_sizes(sliding_window), attention_chunk_size=cls.merge_window_sizes(attention_chunk_size), non_causal=any(spec.non_causal for spec in specs), diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 7b85e6fa316..737feb7d277 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -1,13 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Iterable, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, replace from math import prod from typing import Any, cast import torch -from vllm.config import VllmConfig, get_layers_from_vllm_config +from vllm.config import ( + VllmConfig, + get_layers_from_vllm_config, + set_current_vllm_config, +) from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.utils.torch_utils import get_dtype_size @@ -47,6 +51,13 @@ def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]: continue # Skip modules that don't need KV cache (eg encoder-only attention) if spec := attn_module.get_kv_cache_spec(vllm_config): + if isinstance(spec, AttentionSpec): + backend = attn_module.get_attn_backend() + # indexes_kv_by_block_stride() -> get_kv_cache_stride_order() -> + # get_kv_cache_layout() needs the current vLLM config. + with set_current_vllm_config(vllm_config): + indexes = backend.indexes_kv_by_block_stride() + spec = replace(spec, indexes_kv_by_block_stride=indexes) kv_cache_spec[layer_name] = spec return kv_cache_spec @@ -180,6 +191,62 @@ def _allocate_kv_cache( return kv_cache_raw_tensors +def _reshape_attention_kv_cache( + kv_raw_tensor: torch.Tensor, + kv_cache_spec: AttentionSpec, + kv_cache_shape: tuple[int, ...], + kv_cache_stride_order: tuple[int, ...], + num_blocks: int, + packing: tuple[int, int] | None, +) -> torch.Tensor: + permuted_kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order) + inv_order = [ + kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order)) + ] + dtype = kv_cache_spec.dtype + + if packing is not None: + offset, block_stride = packing + assert inv_order[0] == 0 + page_bytes = prod(kv_cache_shape[1:]) * get_dtype_size(dtype) + kv_cache = ( + kv_raw_tensor.view(-1, block_stride)[:, offset : offset + page_bytes] + .view(dtype) + .view(kv_cache_shape) + ) + elif kv_cache_spec.page_size_padded is not None: + # Use a strided view to skip the padding between physical pages. + # + # Only num-blocks-first layouts are supported (the block dimension is + # dim 0 of the unpermuted shape). kv-first layouts such as ROCm's + # ``(2, num_blocks, ...)`` are intentionally not supported here. For a + # num-blocks-first layout the only stride that must change is the block + # stride: every other (contiguous) stride already steps within the + # unpadded region of a page, so no further adjustment is needed. + assert kv_cache_shape[0] == num_blocks, ( + "Padded KV pages require a num-blocks-first KV cache layout (got " + f"shape {kv_cache_shape} with num_blocks={num_blocks}); " + "kv-first layouts are not supported." + ) + dtype_size = get_dtype_size(kv_cache_spec.dtype) + page_stride = kv_cache_spec.page_size_bytes // dtype_size + + num_blocks_dim = inv_order[0] + strides = list(torch.empty(permuted_kv_cache_shape).stride()) + strides[num_blocks_dim] = page_stride + + kv_cache = torch.as_strided( + kv_raw_tensor.view(dtype), + size=permuted_kv_cache_shape, + stride=tuple(strides), + ) + else: + # No padding — safe to use a contiguous view. + kv_cache = kv_raw_tensor.view(dtype).view(permuted_kv_cache_shape) + + return kv_cache.permute(*inv_order) + + def _reshape_kv_cache( attn_groups: Sequence[AttentionGroup], kv_cache_raw_tensors: dict[str, torch.Tensor], @@ -248,45 +315,14 @@ def _reshape_kv_cache( except (AttributeError, NotImplementedError): kv_cache_stride_order = tuple(range(len(kv_cache_shape))) - kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order) - inv_order = [ - kv_cache_stride_order.index(i) - for i in range(len(kv_cache_stride_order)) - ] - - dtype = kv_cache_spec.dtype - if packing is not None: - offset, block_stride = packing - assert inv_order[0] == 0 - page_bytes = prod(kv_cache_shape[1:]) * get_dtype_size(dtype) - kv_cache = ( - kv_raw_tensor.view(-1, block_stride)[ - :, offset : offset + page_bytes - ] - .view(dtype) - .view(kv_cache_shape) - ) - elif kv_cache_spec.page_size_padded is not None: - # Use strided view to handle page_size_bytes that - # include padding. This follows the same pattern as - # MambaSpec handling in gpu_model_runner.py. - # NOTE: This assumes kv_cache_shape[0] == num_blocks - # (i.e. the first physical dimension is the block - # index), which holds for all current backends - # (MLA, FlashAttention, TritonAttention, etc.). - dtype_size = get_dtype_size(dtype) - page_stride = kv_cache_spec.page_size_bytes // dtype_size - strides = list(torch.empty(kv_cache_shape).stride()) - strides[inv_order[0]] = page_stride - kv_cache = torch.as_strided( - kv_raw_tensor.view(dtype), - size=kv_cache_shape, - stride=tuple(strides), - ) - else: - # No padding — safe to use a contiguous view. - kv_cache = kv_raw_tensor.view(dtype).view(kv_cache_shape) - kv_caches[layer_name] = kv_cache.permute(*inv_order) + kv_caches[layer_name] = _reshape_attention_kv_cache( + kv_raw_tensor, + kv_cache_spec, + kv_cache_shape, + kv_cache_stride_order, + kernel_num_blocks, + packing, + ) elif isinstance(kv_cache_spec, MambaSpec): has_mamba = True diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index b554542e65d..0b72870fc4d 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -12,7 +12,6 @@ from contextlib import contextmanager from copy import copy, deepcopy from dataclasses import dataclass, replace from functools import reduce -from math import prod from typing import TYPE_CHECKING, Any, NamedTuple, TypeAlias, cast import numpy as np @@ -203,6 +202,7 @@ from vllm.v1.worker.cp_utils import ( ) from vllm.v1.worker.dp_utils import coordinate_batch_across_dp from vllm.v1.worker.ec_connector_model_runner_mixin import ECConnectorModelRunnerMixin +from vllm.v1.worker.gpu.attn_utils import _reshape_attention_kv_cache from vllm.v1.worker.gpu.pool.late_interaction_runner import LateInteractionRunner from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch from vllm.v1.worker.gpu_ubatch_wrapper import UBatchWrapper @@ -7125,62 +7125,20 @@ class GPUModelRunner( kv_cache_spec.head_size, cache_dtype_str=self.cache_config.cache_dtype, ) - dtype = kv_cache_spec.dtype try: kv_cache_stride_order = attn_backend.get_kv_cache_stride_order() assert len(kv_cache_stride_order) == len(kv_cache_shape) except (AttributeError, NotImplementedError): kv_cache_stride_order = tuple(range(len(kv_cache_shape))) - # The allocation respects the backend-defined stride order - # to ensure the semantic remains consistent for each - # backend. We first obtain the generic kv cache shape and - # then permute it according to the stride order which could - # result in a non-contiguous tensor. - kv_cache_shape = tuple( - kv_cache_shape[i] for i in kv_cache_stride_order + raw_tensor = kv_cache_raw_tensors[layer_name] + kv_caches[layer_name] = _reshape_attention_kv_cache( + raw_tensor, + kv_cache_spec, + kv_cache_shape, + kv_cache_stride_order, + kernel_num_blocks, + packing, ) - # Maintain original KV shape view. - inv_order = [ - kv_cache_stride_order.index(i) - for i in range(len(kv_cache_stride_order)) - ] - - if packing is not None: - offset, block_stride = packing - assert inv_order[0] == 0 - page_bytes = prod(kv_cache_shape[1:]) * get_dtype_size(dtype) - kv_cache = ( - kv_cache_raw_tensors[layer_name] - .view(-1, block_stride)[:, offset : offset + page_bytes] - .view(dtype) - .view(kv_cache_shape) - ) - elif kv_cache_spec.page_size_padded is not None: - # Use strided view to handle page_size_bytes that - # include padding. This follows - # the same pattern as MambaSpec handling below. - # NOTE: This assumes kv_cache_shape[0] == num_blocks - # (i.e. the first physical dimension is the block - # index), which holds for MLA backends but NOT for - # standard attention backends whose shape starts with - # a K/V dimension of size 2. - dtype_size = get_dtype_size(dtype) - page_stride = kv_cache_spec.page_size_bytes // dtype_size - strides = list(torch.empty(kv_cache_shape).stride()) - strides[inv_order[0]] = page_stride - kv_cache = torch.as_strided( - kv_cache_raw_tensors[layer_name].view(dtype), - size=kv_cache_shape, - stride=tuple(strides), - ) - else: - # No padding — safe to use a contiguous view. - kv_cache = ( - kv_cache_raw_tensors[layer_name] - .view(dtype) - .view(kv_cache_shape) - ) - kv_caches[layer_name] = kv_cache.permute(*inv_order) elif isinstance(kv_cache_spec, MambaSpec): has_mamba = True @@ -7265,7 +7223,7 @@ class GPUModelRunner( # Try creating KV caches optimized for kv-connector transfers cache_dtype = self.cache_config.cache_dtype - if self.use_uniform_kv_cache(self.attn_groups, cache_dtype): + if self.use_uniform_kv_cache(self.attn_groups): kv_caches, cross_layers_kv_cache, attn_backend = ( self.allocate_uniform_kv_caches( kv_cache_config, @@ -7515,6 +7473,13 @@ class GPUModelRunner( continue # Skip modules that don't need KV cache (eg encoder-only attention) if spec := attn_module.get_kv_cache_spec(self.vllm_config): + if isinstance(spec, AttentionSpec): + backend = attn_module.get_attn_backend() + # indexes_kv_by_block_stride() -> get_kv_cache_stride_order() + # -> get_kv_cache_layout() needs the current vLLM config. + with set_current_vllm_config(self.vllm_config): + indexes = backend.indexes_kv_by_block_stride() + spec = replace(spec, indexes_kv_by_block_stride=indexes) kv_cache_spec[layer_name] = spec return kv_cache_spec diff --git a/vllm/v1/worker/kv_connector_model_runner_mixin.py b/vllm/v1/worker/kv_connector_model_runner_mixin.py index 797e59c0290..c2c54e647df 100644 --- a/vllm/v1/worker/kv_connector_model_runner_mixin.py +++ b/vllm/v1/worker/kv_connector_model_runner_mixin.py @@ -114,7 +114,6 @@ class KVConnectorModelRunnerMixin: @staticmethod def use_uniform_kv_cache( attn_groups: list[list[AttentionGroup]], - cache_dtype: CacheDType, ) -> bool: """ Determines whether a uniform KV layout should be used. @@ -128,9 +127,9 @@ class KVConnectorModelRunnerMixin: have the same page size. 2. A KV connector is configured, and the KV connector instance prefers to use this layout (prefer_cross_layer_blocks() returns True) - 2. The flash attention backend supports this layout - (get_kv_cache_stride_order(True) includes a placement for a - num_layers dimension) + 3. The attention backend indexes KV by the block stride + (kv_cache_spec.indexes_kv_by_block_stride), i.e. num_blocks is the + outermost physical dim so per-block all-layers data is contiguous. Note that the actual placement of the num_layers dimensions in the unified layers tensors will be determined by the attention @@ -140,7 +139,6 @@ class KVConnectorModelRunnerMixin: Args: attn_groups: The list of attention groups for this model - cache_dtype: The KV cache dtype Returns: True if we should use a uniform KV cache layout. """ @@ -157,30 +155,7 @@ class KVConnectorModelRunnerMixin: kv_cache_spec = attn_group.kv_cache_spec if not isinstance(kv_cache_spec, AttentionSpec): return False - - attn_backend = attn_group.backend - kv_cache_shape = attn_backend.get_kv_cache_shape( - 1234, - kv_cache_spec.block_size, - kv_cache_spec.num_kv_heads, - kv_cache_spec.head_size, - cache_dtype_str=cache_dtype, - ) - - try: - kv_cache_stride_order = attn_backend.get_kv_cache_stride_order( - include_num_layers_dimension=True - ) - except (AttributeError, NotImplementedError): - return False - - # check that attention backend includes a layers dimension - if len(kv_cache_stride_order) != len(kv_cache_shape) + 1: - return False - - # stride_order[0] == 0 means num_layers stays first in physical - # layout (identity permutation), so cross-layer is unsupported. - return kv_cache_stride_order[0] != 0 + return kv_cache_spec.indexes_kv_by_block_stride @staticmethod def allocate_uniform_kv_caches( From 745bba5ea8fa17dd6ae3751daf43c3d1bd8522df Mon Sep 17 00:00:00 2001 From: Jee Jee Li Date: Mon, 22 Jun 2026 00:28:52 +0800 Subject: [PATCH 66/75] [Model]Fix MiniMaxM2ForCausalLM perf regression (#45935) Signed-off-by: Jee Jee Li --- tests/kernels/core/test_minimax_reduce_rms.py | 62 +++++- .../layers/minimax_rms_norm/rms_norm_tp.py | 185 ++++++++++++++++-- 2 files changed, 223 insertions(+), 24 deletions(-) diff --git a/tests/kernels/core/test_minimax_reduce_rms.py b/tests/kernels/core/test_minimax_reduce_rms.py index de9fc2bbb4f..b9c591bb93f 100644 --- a/tests/kernels/core/test_minimax_reduce_rms.py +++ b/tests/kernels/core/test_minimax_reduce_rms.py @@ -10,8 +10,12 @@ from torch.multiprocessing import spawn from tests.kernels.utils import opcheck from tests.utils import ensure_current_vllm_config, init_test_distributed_environment from vllm.distributed import cleanup_dist_env_and_memory -from vllm.model_executor.layers.minimax_rms_norm import MiniMaxText01RMSNormTP +from vllm.model_executor.layers.minimax_rms_norm import ( + MiniMaxText01RMSNormTP, + rms_norm_tp, +) from vllm.platforms import current_platform +from vllm.triton_utils import HAS_TRITON from vllm.utils.network_utils import get_open_port from vllm.utils.torch_utils import set_random_seed @@ -54,8 +58,19 @@ def _worker_forward_qk( torch.manual_seed(seed + 1000 + local_rank) qkv = torch.randn(num_tokens, hq + hk + hk, dtype=dtype, device="cuda") - q_ref, k_ref, v_ref = qkv.clone().split([hq, hk, hk], dim=-1) - ref_q, ref_k = MiniMaxText01RMSNormTP.forward_qk(q_norm, k_norm, q_ref, k_ref) + # Reference: eager all-reduce path. ``forward_qk`` no longer all-reduces + # the variance (it is the tp==1 / already-reduced building block), so the + # multi-rank reference must use the eager path that performs the global + # variance all-reduce, matching the fused kernel below. + ref_q, ref_k = rms_norm_tp._minimax_qk_norm_tp_eager( + qkv.clone(), + q_norm.weight, + k_norm.weight, + hq, + hk, + world_size, + eps, + ) # Set up Lamport workspace. from vllm.distributed.parallel_state import get_tp_group @@ -150,3 +165,44 @@ def test_minimax_reduce_rms_qk( nprocs=world_size, join=True, ) + + +@pytest.mark.skipif( + not current_platform.is_cuda() or not HAS_TRITON, + reason="CUDA and Triton required", +) +@pytest.mark.parametrize("num_tokens", [1, 7, 128, 333, 2049]) +@pytest.mark.parametrize("hidden_dims", [(3072, 512), (768, 256), (3000, 500)]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("tp_world", [1, 4, 8]) +@pytest.mark.parametrize("eps", [1e-6]) +@pytest.mark.parametrize("seed", [42]) +def test_minimax_qk_norm_triton_fallback( + monkeypatch, num_tokens, hidden_dims, dtype, tp_world, eps, seed +): + """Single-GPU check: Triton fallback kernels vs the pure-torch reference. + + The all-reduce is a TP communication barrier, so it is monkeypatched to + identity here; both the Triton path and the reference see the same + (patched) reduction. This validates the kernel math and the folded + ``/ tp_world`` scaling without needing multiple ranks -- ``hidden_dims`` + are the per-rank q/k segment widths. + """ + monkeypatch.setattr(rms_norm_tp, "_all_reduce_variance", lambda v: v) + + q_size, kv_size = hidden_dims + device = "cuda" + torch.manual_seed(seed) + qkv = torch.randn(num_tokens, q_size + 2 * kv_size, dtype=dtype, device=device) + q_weight = torch.randn(q_size, dtype=dtype, device=device) + k_weight = torch.randn(kv_size, dtype=dtype, device=device) + + q_triton, k_triton = rms_norm_tp._minimax_qk_norm_tp_fallback( + qkv, q_weight, k_weight, q_size, kv_size, 0, tp_world, eps + ) + q_ref, k_ref = rms_norm_tp._minimax_qk_norm_tp_eager( + qkv, q_weight, k_weight, q_size, kv_size, tp_world, eps + ) + + torch.testing.assert_close(q_triton, q_ref, atol=3e-2, rtol=3e-2) + torch.testing.assert_close(k_triton, k_ref, atol=3e-2, rtol=3e-2) diff --git a/vllm/model_executor/layers/minimax_rms_norm/rms_norm_tp.py b/vllm/model_executor/layers/minimax_rms_norm/rms_norm_tp.py index e2c938ddad1..e48d9c01354 100644 --- a/vllm/model_executor/layers/minimax_rms_norm/rms_norm_tp.py +++ b/vllm/model_executor/layers/minimax_rms_norm/rms_norm_tp.py @@ -14,7 +14,7 @@ from vllm.distributed.parallel_state import ( ) from vllm.logger import init_logger from vllm.model_executor.custom_op import CustomOp -from vllm.platforms import current_platform +from vllm.triton_utils import HAS_TRITON, tl, triton from vllm.utils.torch_utils import direct_register_custom_op logger = init_logger(__name__) @@ -40,8 +40,116 @@ def _all_reduce_variance(var: torch.Tensor) -> torch.Tensor: return tensor_model_parallel_all_reduce(var.flatten()).view_as(var) -@torch.compile(backend=current_platform.simple_compile_backend, dynamic=True) -def _minimax_qk_norm_fallback( +@triton.jit +def _minimax_qk_var_kernel( + qkv_ptr, # [num_tokens, hidden], 16-bit activations + var_ptr, # [num_tokens, 2], fp32 + row_stride, # element stride between tokens in qkv + q_size: tl.constexpr, # constant per deployment -> loops unroll, mask elides + kv_size: tl.constexpr, + BLOCK: tl.constexpr, +): + """TP-pre stage: per-token mean-of-squares for the q and k segments. + + Accumulates in fp32 while reading the 16-bit qkv in place, so no fp32 + copy of q/k is materialized. ``var[:, 0]`` is the q variance and + ``var[:, 1]`` the k variance; both are the local-shard means, ready for + the all-reduce that follows. + """ + token = tl.program_id(0) + base = qkv_ptr + token * row_stride + + q_acc = 0.0 + for off in range(0, q_size, BLOCK): + idx = off + tl.arange(0, BLOCK) + mask = idx < q_size + x = tl.load(base + idx, mask=mask, other=0.0).to(tl.float32) + q_acc += tl.sum(x * x, axis=0) + + k_acc = 0.0 + for off in range(0, kv_size, BLOCK): + idx = off + tl.arange(0, BLOCK) + mask = idx < kv_size + x = tl.load(base + q_size + idx, mask=mask, other=0.0).to(tl.float32) + k_acc += tl.sum(x * x, axis=0) + + tl.store(var_ptr + token * 2 + 0, q_acc / q_size) + tl.store(var_ptr + token * 2 + 1, k_acc / kv_size) + + +@triton.jit +def _minimax_rms_apply_kernel( + qkv_ptr, # [num_tokens, hidden] + var_ptr, # [num_tokens, 2], fp32, all-reduced sum of per-shard means + q_w_ptr, # [q_size], q per-channel weight + k_w_ptr, # [kv_size], k per-channel weight + q_out_ptr, # [num_tokens, q_size], contiguous + k_out_ptr, # [num_tokens, kv_size], contiguous + row_stride, # element stride between tokens in qkv + q_size: tl.constexpr, # constant per deployment -> loops unroll, mask elides + kv_size: tl.constexpr, + tp_world: tl.constexpr, # folds the post-all-reduce /tp_world into rsqrt + eps: tl.constexpr, + BLOCK: tl.constexpr, +): + """TP-post stage: ``x * rsqrt(var / tp_world + eps) * weight``. + + A single program normalizes both the q and k segments of one token, so q + and k share one launch instead of two. The all-reduce yields the sum of + per-shard means, so the ``/ tp_world`` that recovers the global + mean-of-squares is folded into the ``rsqrt`` here rather than run as a + separate elementwise pass over the ``[num_tokens, 2]`` variance tensor. + """ + token = tl.program_id(0) + base = qkv_ptr + token * row_stride + + q_inv = tl.rsqrt(tl.load(var_ptr + token * 2 + 0) / tp_world + eps) + q_out_row = q_out_ptr + token * q_size + for off in range(0, q_size, BLOCK): + idx = off + tl.arange(0, BLOCK) + mask = idx < q_size + x = tl.load(base + idx, mask=mask, other=0.0).to(tl.float32) + w = tl.load(q_w_ptr + idx, mask=mask, other=0.0).to(tl.float32) + y = x * q_inv * w + tl.store(q_out_row + idx, y.to(q_out_ptr.dtype.element_ty), mask=mask) + + k_inv = tl.rsqrt(tl.load(var_ptr + token * 2 + 1) / tp_world + eps) + k_out_row = k_out_ptr + token * kv_size + for off in range(0, kv_size, BLOCK): + idx = off + tl.arange(0, BLOCK) + mask = idx < kv_size + x = tl.load(base + q_size + idx, mask=mask, other=0.0).to(tl.float32) + w = tl.load(k_w_ptr + idx, mask=mask, other=0.0).to(tl.float32) + y = x * k_inv * w + tl.store(k_out_row + idx, y.to(k_out_ptr.dtype.element_ty), mask=mask) + + +def _minimax_qk_norm_tp_eager( + qkv: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + q_size: int, + kv_size: int, + tp_world: int, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """Pure-torch reference path used when Triton is unavailable.""" + q, k, _ = qkv.split([q_size, kv_size, kv_size], dim=-1) + orig_dtype = q.dtype + q = q.to(torch.float32) + k = k.to(torch.float32) + q_var = q.pow(2).mean(dim=-1, keepdim=True) + k_var = k.pow(2).mean(dim=-1, keepdim=True) + + qk_var = torch.cat([q_var, k_var], dim=-1) + qk_var = _all_reduce_variance(qk_var) / tp_world + q_var, k_var = qk_var.chunk(2, dim=-1) + q = q * torch.rsqrt(q_var + eps) * q_weight + k = k * torch.rsqrt(k_var + eps) * k_weight + return q.to(orig_dtype), k.to(orig_dtype) + + +def _minimax_qk_norm_tp_fallback( qkv: torch.Tensor, q_weight: torch.Tensor, k_weight: torch.Tensor, @@ -51,19 +159,50 @@ def _minimax_qk_norm_fallback( tp_world: int, eps: float, ) -> tuple[torch.Tensor, torch.Tensor]: - q, k, _ = qkv.split([q_size, kv_size, kv_size], dim=-1) - orig_dtype = q.dtype - q = q.to(torch.float32) - k = k.to(torch.float32) - q_var = q.pow(2).mean(dim=-1, keepdim=True) - k_var = k.pow(2).mean(dim=-1, keepdim=True) - if tp_world > 1: - qk_var = torch.cat([q_var, k_var], dim=-1) - qk_var = _all_reduce_variance(qk_var) / tp_world - q_var, k_var = qk_var.chunk(2, dim=-1) - q = q * torch.rsqrt(q_var + eps) * q_weight - k = k * torch.rsqrt(k_var + eps) * k_weight - return q.to(orig_dtype), k.to(orig_dtype) + """All-reduce + QK RMSNorm without the Lamport fused kernel. + + The all-reduce is a TP communication barrier and cannot live inside a + single kernel, so the eager-torch path is split into two Triton kernels + around it: a variance reduction before the all-reduce and a normalize + after. Compared to the ``torch.compile`` path this avoids materializing + fp32 copies of q/k and the ``cat``/``chunk`` temporaries. + """ + if not HAS_TRITON: + return _minimax_qk_norm_tp_eager( + qkv, q_weight, k_weight, q_size, kv_size, tp_world, eps + ) + + num_tokens = qkv.shape[0] + row_stride = qkv.stride(0) + BLOCK = 1024 + grid = (num_tokens,) + + qk_var = torch.empty(num_tokens, 2, dtype=torch.float32, device=qkv.device) + _minimax_qk_var_kernel[grid]( + qkv, qk_var, row_stride, q_size=q_size, kv_size=kv_size, BLOCK=BLOCK + ) + + # All-reduce sums the per-shard means; the /tp_world that turns this back + # into the global mean is folded into the apply kernel's rsqrt below. + qk_var = _all_reduce_variance(qk_var) + + q_out = torch.empty(num_tokens, q_size, dtype=qkv.dtype, device=qkv.device) + k_out = torch.empty(num_tokens, kv_size, dtype=qkv.dtype, device=qkv.device) + _minimax_rms_apply_kernel[grid]( + qkv, + qk_var, + q_weight, + k_weight, + q_out, + k_out, + row_stride, + q_size=q_size, + kv_size=kv_size, + tp_world=tp_world, + eps=eps, + BLOCK=BLOCK, + ) + return q_out, k_out def _minimax_qk_norm_fusion( @@ -96,7 +235,7 @@ def _minimax_qk_norm_fusion( tp_world, eps, ) - return _minimax_qk_norm_fallback( + return _minimax_qk_norm_tp_fallback( qkv, q_weight, k_weight, q_size, kv_size, tp_rank, tp_world, eps ) @@ -231,10 +370,7 @@ class MiniMaxText01RMSNormTP(CustomOp): k = k.to(torch.float32) q_var = q.pow(2).mean(dim=-1, keepdim=True) k_var = k.pow(2).mean(dim=-1, keepdim=True) - if q_norm.tp_world > 1: - qk_var = torch.cat([q_var, k_var], dim=-1) - qk_var = _all_reduce_variance(qk_var) / q_norm.tp_world - q_var, k_var = qk_var.chunk(2, dim=-1) + q = q * torch.rsqrt(q_var + q_norm.variance_epsilon) * q_norm.weight k = k * torch.rsqrt(k_var + k_norm.variance_epsilon) * k_norm.weight q = q.to(orig_dtype) @@ -250,7 +386,14 @@ class MiniMaxText01RMSNormTP(CustomOp): kv_size: int, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: assert qkv.ndim == 2 + assert q_norm.variance_epsilon == k_norm.variance_epsilon + # Case 0： tp_size=1 + if get_tensor_model_parallel_world_size() == 1: + q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1) + q, k = MiniMaxText01RMSNormTP.forward_qk(q_norm, k_norm, q, k) + return q, k, v + # Case : tp_size>1 q, k = torch.ops.vllm.minimax_qk_norm_fusion( qkv, q_norm.weight, From c441ad1c07cbfe0240a5699e4386fd0b5cc8aa82 Mon Sep 17 00:00:00 2001 From: Srinivas Krovvidi <194645829+Srinivasoo7@users.noreply.github.com> Date: Sun, 21 Jun 2026 13:04:01 -0500 Subject: [PATCH 67/75] [KV Offloading] Add labeled metrics support (#45957) Signed-off-by: srinivas_oo7 Co-authored-by: srinivas_oo7 --- .../unit/offloading_connector/test_metrics.py | 315 ++++++++++++++---- .../kv_connector/v1/offloading/metrics.py | 191 +++++++---- vllm/v1/kv_offload/base.py | 1 + 3 files changed, 384 insertions(+), 123 deletions(-) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_metrics.py b/tests/v1/kv_connector/unit/offloading_connector/test_metrics.py index f9a4b377959..6f36a6c8149 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_metrics.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_metrics.py @@ -22,7 +22,7 @@ from vllm.v1.kv_offload.base import ( OffloadingGaugeMetadata, OffloadingHistogramMetadata, ) -from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec +from vllm.v1.kv_offload.factory import OffloadingSpecFactory LOAD_BYTES = _TransferMetricName.LOAD_BYTES LOAD_TIME = _TransferMetricName.LOAD_TIME @@ -33,6 +33,8 @@ STORE_SIZE = _TransferMetricName.STORE_SIZE STORES_SKIPPED = "vllm:kv_offload_stores_skipped" PENDING_STORES = "vllm:kv_offload_pending_stores" LOOKUP_LATENCY = "vllm:kv_offload_lookup_latency_seconds" +MY_COUNTER = "my_counter" +MY_LABEL = "my_label" class _FakeMetric: @@ -67,6 +69,20 @@ class _FakeVllmConfig: ) +def _spec_cls_with_metric_definitions( + metric_definitions: dict[str, Any], +) -> type: + """Build a fake offloading spec class reporting the given metric + definitions, so tests don't need to patch the real CPU spec.""" + + class _FakeOffloadingSpec: + @staticmethod + def build_metric_definitions(extra_config): + return metric_definitions + + return _FakeOffloadingSpec + + def _metric_metadata(): return { LOAD_BYTES: OffloadingCounterMetadata( @@ -96,9 +112,17 @@ def _metric_metadata(): LOOKUP_LATENCY: OffloadingHistogramMetadata( documentation="lookup latency", ), + MY_COUNTER: OffloadingCounterMetadata( + documentation="counter with a label", + labelnames=(MY_LABEL,), + ), } +def _unlabeled(values: dict[str, Any], metric_name: str) -> Any: + return values[metric_name][()] + + def test_build_kv_connector_stats_with_none(): """Test that build_kv_connector_stats returns empty stats when given None.""" stats = OffloadingConnector.build_kv_connector_stats(data=None) @@ -131,13 +155,13 @@ def test_build_kv_connector_stats_reconstructs_offload_stats(): STORES_SKIPPED: _MetricType.COUNTER, }, _StatsKey.DATA: { - LOAD_BYTES: 24, - LOAD_TIME: 1.5, - LOAD_SIZE: [16, 8], - STORE_BYTES: 3, - STORE_TIME: 0.3, - STORE_SIZE: [1, 2], - STORES_SKIPPED: 5, + LOAD_BYTES: {(): 24}, + LOAD_TIME: {(): 1.5}, + LOAD_SIZE: {(): [16, 8]}, + STORE_BYTES: {(): 3}, + STORE_TIME: {(): 0.3}, + STORE_SIZE: {(): [1, 2]}, + STORES_SKIPPED: {(): 5}, }, } @@ -145,22 +169,28 @@ def test_build_kv_connector_stats_reconstructs_offload_stats(): assert isinstance(stats, OffloadingConnectorStats) values = stats.data[_StatsKey.DATA] - assert values[LOAD_BYTES] == 24 - assert values[LOAD_TIME] == 1.5 - assert values[LOAD_SIZE] == [16, 8] - assert values[STORE_BYTES] == 3 - assert values[STORE_TIME] == 0.3 - assert values[STORE_SIZE] == [1, 2] - assert values[STORES_SKIPPED] == 5 + assert _unlabeled(values, LOAD_BYTES) == 24 + assert _unlabeled(values, LOAD_TIME) == 1.5 + assert _unlabeled(values, LOAD_SIZE) == [16, 8] + assert _unlabeled(values, STORE_BYTES) == 3 + assert _unlabeled(values, STORE_TIME) == 0.3 + assert _unlabeled(values, STORE_SIZE) == [1, 2] + assert _unlabeled(values, STORES_SKIPPED) == 5 def _make_stats_data( metric_data: dict[str, Any], metric_metadata: dict[str, Any], ) -> dict[str, Any]: - """Build a structured data dict from flat metric data and metadata.""" + """Build a structured data dict from flat metric data and metadata. + + Values for unlabeled metrics may be passed flat (wrapped here under the + empty label tuple); values for labeled metrics must already be passed as + a ``{labelvalues: value}`` map. + """ metric_types = {} - for key in metric_data: + data = {} + for key, value in metric_data.items(): md = metric_metadata[key] if isinstance(md, OffloadingCounterMetadata): metric_types[key] = _MetricType.COUNTER @@ -168,9 +198,10 @@ def _make_stats_data( metric_types[key] = _MetricType.GAUGE elif isinstance(md, OffloadingHistogramMetadata): metric_types[key] = _MetricType.HISTOGRAM + data[key] = value if md.labelnames else {(): value} return { _StatsKey.TYPES: metric_types, - _StatsKey.DATA: metric_data, + _StatsKey.DATA: data, } @@ -215,34 +246,106 @@ def test_aggregate_same_connector(): assert result is stats1 # Should return self values = result.data[_StatsKey.DATA] - assert values[LOAD_BYTES] == 34 - assert values[LOAD_TIME] == 2.6 - assert values[LOAD_SIZE] == [16, 8, 3, 7] - assert values[STORE_BYTES] == 19 - assert values[STORE_TIME] == 2.3 - assert values[STORE_SIZE] == [1, 2, 16] - assert values[STORES_SKIPPED] == 4 - assert values[PENDING_STORES] == 1 - assert values[LOOKUP_LATENCY] == [0.1, 0.2, 0.3] + assert _unlabeled(values, LOAD_BYTES) == 34 + assert _unlabeled(values, LOAD_TIME) == 2.6 + assert _unlabeled(values, LOAD_SIZE) == [16, 8, 3, 7] + assert _unlabeled(values, STORE_BYTES) == 19 + assert _unlabeled(values, STORE_TIME) == 2.3 + assert _unlabeled(values, STORE_SIZE) == [1, 2, 16] + assert _unlabeled(values, STORES_SKIPPED) == 4 + assert _unlabeled(values, PENDING_STORES) == 1 + assert _unlabeled(values, LOOKUP_LATENCY) == [0.1, 0.2, 0.3] + + +def test_aggregate_labeled_metrics(): + metadata = _metric_metadata() + stats1 = OffloadingConnectorStats( + data=_make_stats_data( + { + MY_COUNTER: { + ("a",): 10, + ("b",): 3, + }, + }, + metadata, + ), + ) + stats2 = OffloadingConnectorStats( + data=_make_stats_data( + { + MY_COUNTER: { + ("a",): 7, + ("c",): 5, + }, + }, + metadata, + ), + ) + + stats1.aggregate(stats2) + + values = stats1.data[_StatsKey.DATA][MY_COUNTER] + assert values[("a",)] == 17 + assert values[("b",)] == 3 + assert values[("c",)] == 5 + + +def test_aggregate_labeled_metric_missing_from_self(): + """Aggregating a labeled metric that self doesn't have at all yet.""" + metadata = _metric_metadata() + stats1 = OffloadingConnectorStats() + stats2 = OffloadingConnectorStats( + data=_make_stats_data( + { + MY_COUNTER: { + ("a",): 7, + ("b",): 5, + }, + }, + metadata, + ), + ) + + stats1.aggregate(stats2) + + values = stats1.data[_StatsKey.DATA][MY_COUNTER] + assert values[("a",)] == 7 + assert values[("b",)] == 5 + assert stats1.data[_StatsKey.TYPES][MY_COUNTER] == _MetricType.COUNTER + + +def test_helper_methods_accept_labeled_metrics(): + stats = OffloadingConnectorStats() + + stats.increase_counter(MY_COUNTER, 3, ("a",)) + stats.increase_counter(MY_COUNTER, 4, ("a",)) + stats.set_gauge(PENDING_STORES, 2, ("b",)) + stats.observe_histogram(LOOKUP_LATENCY, 0.1, ("b",)) + stats.observe_histogram(LOOKUP_LATENCY, 0.2, ("b",)) + + values = stats.data[_StatsKey.DATA] + assert values[MY_COUNTER][("a",)] == 7 + assert values[PENDING_STORES][("b",)] == 2 + assert values[LOOKUP_LATENCY][("b",)] == [0.1, 0.2] def test_aggregate_merges_types(): stats1 = OffloadingConnectorStats( data={ _StatsKey.TYPES: {LOAD_BYTES: _MetricType.COUNTER}, - _StatsKey.DATA: {LOAD_BYTES: 1}, + _StatsKey.DATA: {LOAD_BYTES: {(): 1}}, }, ) stats2 = OffloadingConnectorStats( data={ _StatsKey.TYPES: {PENDING_STORES: _MetricType.GAUGE}, - _StatsKey.DATA: {PENDING_STORES: 2}, + _StatsKey.DATA: {PENDING_STORES: {(): 2}}, }, ) result = stats1.aggregate(stats2) - assert result.data[_StatsKey.DATA][PENDING_STORES] == 2 + assert _unlabeled(result.data[_StatsKey.DATA], PENDING_STORES) == 2 assert result.data[_StatsKey.TYPES][PENDING_STORES] == _MetricType.GAUGE @@ -283,6 +386,26 @@ def test_reduce(): assert reduced[f"{LOOKUP_LATENCY}_sum"] == sum([0.1, 0.2, 0.3]) +def test_reduce_labeled_metrics(): + metadata = _metric_metadata() + stats = OffloadingConnectorStats( + data=_make_stats_data( + { + MY_COUNTER: { + ("a",): 17, + ("b",): 3, + }, + }, + metadata, + ), + ) + + reduced = stats.reduce() + + assert reduced[f"{MY_COUNTER}:{('a',)}"] == 17 + assert reduced[f"{MY_COUNTER}:{('b',)}"] == 3 + + def test_reset(): """Test that reset() resets all connector stats.""" metadata = _metric_metadata() @@ -326,11 +449,11 @@ def test_prom_metrics_observes_manager_counter(): prom_metrics.observe( { _StatsKey.TYPES: {STORES_SKIPPED: _MetricType.COUNTER}, - _StatsKey.DATA: {STORES_SKIPPED: 7}, + _StatsKey.DATA: {STORES_SKIPPED: {(): 7}}, } ) - counter = prom_metrics.offloading_metrics[(0, STORES_SKIPPED)] + counter = prom_metrics.offloading_metrics[(0, STORES_SKIPPED, ())] assert counter.increments == [7] counter_def = prom_metrics._offloading_metric_defs[STORES_SKIPPED] assert counter_def.kwargs["name"] == "vllm:kv_offload_stores_skipped" @@ -360,22 +483,22 @@ def test_prom_metrics_observes_flat_transfer_metrics_and_legacy_metrics(): STORE_SIZE: _MetricType.HISTOGRAM, }, _StatsKey.DATA: { - LOAD_BYTES: 24, - LOAD_TIME: 1.5, - LOAD_SIZE: [16, 8], - STORE_BYTES: 3, - STORE_TIME: 0.3, - STORE_SIZE: [1, 2], + LOAD_BYTES: {(): 24}, + LOAD_TIME: {(): 1.5}, + LOAD_SIZE: {(): [16, 8]}, + STORE_BYTES: {(): 3}, + STORE_TIME: {(): 0.3}, + STORE_SIZE: {(): [1, 2]}, }, } ) - assert prom_metrics.offloading_metrics[(0, LOAD_BYTES)].increments == [24] - assert prom_metrics.offloading_metrics[(0, LOAD_TIME)].increments == [1.5] - assert prom_metrics.offloading_metrics[(0, LOAD_SIZE)].observed == [16, 8] - assert prom_metrics.offloading_metrics[(0, STORE_BYTES)].increments == [3] - assert prom_metrics.offloading_metrics[(0, STORE_TIME)].increments == [0.3] - assert prom_metrics.offloading_metrics[(0, STORE_SIZE)].observed == [1, 2] + assert prom_metrics.offloading_metrics[(0, LOAD_BYTES, ())].increments == [24] + assert prom_metrics.offloading_metrics[(0, LOAD_TIME, ())].increments == [1.5] + assert prom_metrics.offloading_metrics[(0, LOAD_SIZE, ())].observed == [16, 8] + assert prom_metrics.offloading_metrics[(0, STORE_BYTES, ())].increments == [3] + assert prom_metrics.offloading_metrics[(0, STORE_TIME, ())].increments == [0.3] + assert prom_metrics.offloading_metrics[(0, STORE_SIZE, ())].observed == [1, 2] assert prom_metrics.counter_kv_bytes[(0, "CPU_to_GPU")].increments == [24] assert prom_metrics.counter_kv_transfer_time[(0, "CPU_to_GPU")].increments == [1.5] @@ -396,7 +519,9 @@ def test_prom_metrics_observes_manager_gauge_and_histogram(): ), } with patch.object( - CPUOffloadingSpec, "build_metric_definitions", return_value=metric_definitions + OffloadingSpecFactory, + "get_spec_cls", + return_value=_spec_cls_with_metric_definitions(metric_definitions), ): prom_metrics = OffloadPromMetrics( vllm_config=_FakeVllmConfig(store_threshold=0), # type: ignore[arg-type] @@ -416,20 +541,91 @@ def test_prom_metrics_observes_manager_gauge_and_histogram(): LOOKUP_LATENCY: _MetricType.HISTOGRAM, }, _StatsKey.DATA: { - PENDING_STORES: 5, - LOOKUP_LATENCY: [0.2, 0.4], + PENDING_STORES: {(): 5}, + LOOKUP_LATENCY: {(): [0.2, 0.4]}, }, } ) - gauge = prom_metrics.offloading_metrics[(0, PENDING_STORES)] - histogram = prom_metrics.offloading_metrics[(0, LOOKUP_LATENCY)] + gauge = prom_metrics.offloading_metrics[(0, PENDING_STORES, ())] + histogram = prom_metrics.offloading_metrics[(0, LOOKUP_LATENCY, ())] assert gauge.set_values == [5] assert histogram.observed == [0.2, 0.4] histogram_def = prom_metrics._offloading_metric_defs[LOOKUP_LATENCY] assert histogram_def.kwargs["buckets"] == (0.1, 1.0) +def test_prom_metrics_lazily_observes_labeled_metric(): + metric_definitions = { + MY_COUNTER: OffloadingCounterMetadata( + documentation="counter with a label", + labelnames=(MY_LABEL,), + ), + } + with patch.object( + OffloadingSpecFactory, + "get_spec_cls", + return_value=_spec_cls_with_metric_definitions(metric_definitions), + ): + prom_metrics = OffloadPromMetrics( + vllm_config=_FakeVllmConfig(store_threshold=0), # type: ignore[arg-type] + metric_types={ + Gauge: _FakeMetric, + Counter: _FakeMetric, + Histogram: _FakeMetric, + }, + labelnames=["model_name", "engine"], + per_engine_labelvalues={0: ["model", "0"]}, + ) + + assert (0, MY_COUNTER, ("a",)) not in prom_metrics.offloading_metrics + + prom_metrics.observe( + { + _StatsKey.TYPES: {MY_COUNTER: _MetricType.COUNTER}, + _StatsKey.DATA: {MY_COUNTER: {("a",): 7}}, + } + ) + + counter = prom_metrics.offloading_metrics[(0, MY_COUNTER, ("a",))] + assert counter.increments == [7] + assert counter.labelvalues == ("model", "0", "a") + counter_def = prom_metrics._offloading_metric_defs[MY_COUNTER] + assert counter_def.kwargs["labelnames"] == ["model_name", "engine", MY_LABEL] + + +def test_prom_metrics_rejects_wrong_label_count(): + metric_definitions = { + MY_COUNTER: OffloadingCounterMetadata( + documentation="counter with a label", + labelnames=(MY_LABEL,), + ), + } + with patch.object( + OffloadingSpecFactory, + "get_spec_cls", + return_value=_spec_cls_with_metric_definitions(metric_definitions), + ): + prom_metrics = OffloadPromMetrics( + vllm_config=_FakeVllmConfig(store_threshold=0), # type: ignore[arg-type] + metric_types={ + Gauge: _FakeMetric, + Counter: _FakeMetric, + Histogram: _FakeMetric, + }, + labelnames=["model_name", "engine"], + per_engine_labelvalues={0: ["model", "0"]}, + ) + + with pytest.raises(AssertionError, match="expects 1 labels"): + prom_metrics.observe( + { + _StatsKey.TYPES: {MY_COUNTER: _MetricType.COUNTER}, + _StatsKey.DATA: {MY_COUNTER: {("a", "extra"): 7}}, + } + ) + + def test_prom_metrics_uses_configured_manager_metrics(): prom_metrics = OffloadPromMetrics( vllm_config=_FakeVllmConfig(store_threshold=0), # type: ignore[arg-type] @@ -458,9 +654,9 @@ def test_aggregate_into_empty_stats(): PENDING_STORES: _MetricType.GAUGE, }, _StatsKey.DATA: { - LOAD_BYTES: 42, - LOAD_SIZE: [10, 20], - PENDING_STORES: 3, + LOAD_BYTES: {(): 42}, + LOAD_SIZE: {(): [10, 20]}, + PENDING_STORES: {(): 3}, }, }, ) @@ -469,9 +665,9 @@ def test_aggregate_into_empty_stats(): assert result is empty values = result.data[_StatsKey.DATA] - assert values[LOAD_BYTES] == 42 - assert values[LOAD_SIZE] == [10, 20] - assert values[PENDING_STORES] == 3 + assert _unlabeled(values, LOAD_BYTES) == 42 + assert _unlabeled(values, LOAD_SIZE) == [10, 20] + assert _unlabeled(values, PENDING_STORES) == 3 def test_prom_metrics_multi_engine_routing(): @@ -490,14 +686,13 @@ def test_prom_metrics_multi_engine_routing(): prom_metrics.observe( { _StatsKey.TYPES: {LOAD_BYTES: _MetricType.COUNTER}, - _StatsKey.DATA: {LOAD_BYTES: 100}, + _StatsKey.DATA: {LOAD_BYTES: {(): 100}}, }, engine_idx=1, ) - engine0 = prom_metrics.offloading_metrics[(0, LOAD_BYTES)] - engine1 = prom_metrics.offloading_metrics[(1, LOAD_BYTES)] - assert engine0.increments == [] + assert (0, LOAD_BYTES, ()) not in prom_metrics.offloading_metrics + engine1 = prom_metrics.offloading_metrics[(1, LOAD_BYTES, ())] assert engine1.increments == [100] @@ -518,6 +713,6 @@ def test_prom_metrics_rejects_undeclared_metric(): prom_metrics.observe( { _StatsKey.TYPES: {"unknown:metric": _MetricType.COUNTER}, - _StatsKey.DATA: {"unknown:metric": 1}, + _StatsKey.DATA: {"unknown:metric": {(): 1}}, } ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py index 3e4463924b7..a90250d285e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py @@ -112,7 +112,7 @@ class _StatsKey: # Maps metric name -> _MetricType value TYPES = "types" - # Maps metric name -> observed value (number or list) + # Maps metric name -> {label values tuple -> observed value (number or list)} DATA = "data" @@ -125,15 +125,17 @@ class OffloadingConnectorStats(KVConnectorStats): { _StatsKey.TYPES: {name: _MetricType.*, ...}, - _StatsKey.DATA: {name: value, ...}, + _StatsKey.DATA: {name: {labelvalues: value, ...}, ...}, } This structure is self-describing: it survives IPC serialization without needing the full ``OffloadingMetricMetadata`` objects on the receiving side. - Counter values are aggregated by summing, gauge values use the latest - snapshot, and histogram values are lists of observed samples. + Counter values are aggregated by summing per-label-tuple, gauge values + use the latest snapshot per-label-tuple, and histogram values are lists of + observed samples per-label-tuple. Unlabeled metrics use ``()`` as their + labelvalues tuple. """ def __post_init__(self): @@ -160,26 +162,32 @@ class OffloadingConnectorStats(KVConnectorStats): assert isinstance(other, OffloadingConnectorStats) other_types = other._types other_values = other._values - for key, value in other_values.items(): + for key, other_label_values in other_values.items(): type_str = other_types.get(key) if type_str is None: raise AssertionError(f"Unknown offloading stats key: {key}") self._types.setdefault(key, type_str) - if type_str == _MetricType.HISTOGRAM: - assert isinstance(value, list) - if key not in self._values: - self._values[key] = value + current_label_values = self._values.setdefault(key, {}) + for labelvalues, value in other_label_values.items(): + if type_str == _MetricType.HISTOGRAM: + assert isinstance(value, list) + if labelvalues not in current_label_values: + current_label_values[labelvalues] = list(value) + else: + assert isinstance(current_label_values[labelvalues], list) + current_label_values[labelvalues].extend(value) + elif type_str == _MetricType.COUNTER: + assert isinstance(value, int | float) + current_label_values[labelvalues] = ( + current_label_values.get(labelvalues, 0) + value + ) + elif type_str == _MetricType.GAUGE: + assert isinstance(value, int | float) + current_label_values[labelvalues] = value else: - assert isinstance(self._values[key], list) - self._values[key].extend(value) - elif type_str == _MetricType.COUNTER: - assert isinstance(value, int | float) - self._values[key] = self._values.get(key, 0) + value - elif type_str == _MetricType.GAUGE: - assert isinstance(value, int | float) - self._values[key] = value - else: - raise AssertionError(f"Unknown metric type '{type_str}' for key: {key}") + raise AssertionError( + f"Unknown metric type '{type_str}' for key: {key}" + ) return self def reduce(self) -> dict[str, int | float]: @@ -190,44 +198,62 @@ class OffloadingConnectorStats(KVConnectorStats): stats for the last time interval. """ return_dict: dict[str, int | float] = {} - for key, value in self._values.items(): + for key, label_value_map in self._values.items(): type_str = self._types.get(key) if type_str is None: raise AssertionError(f"Unknown offloading stats key: {key}") - if type_str == _MetricType.HISTOGRAM: - assert isinstance(value, list) - return_dict[f"{key}_count"] = len(value) - return_dict[f"{key}_sum"] = sum(value) - elif type_str in (_MetricType.COUNTER, _MetricType.GAUGE): - assert isinstance(value, int | float) - return_dict[key] = value - else: - raise AssertionError(f"Unknown metric type '{type_str}' for key: {key}") + for labelvalues, value in label_value_map.items(): + key_with_labels = f"{key}:{labelvalues}" if labelvalues else key + if type_str == _MetricType.HISTOGRAM: + assert isinstance(value, list) + return_dict[f"{key_with_labels}_count"] = len(value) + return_dict[f"{key_with_labels}_sum"] = sum(value) + elif type_str in (_MetricType.COUNTER, _MetricType.GAUGE): + assert isinstance(value, int | float) + return_dict[key_with_labels] = value + else: + raise AssertionError( + f"Unknown metric type '{type_str}' for key: {key}" + ) return return_dict def is_empty(self) -> bool: return not self.data.get(_StatsKey.DATA) def increase_counter( - self, counter_name: str, counter_increase_value: int | float + self, + counter_name: str, + counter_increase_value: int | float, + labelvalues: tuple[str, ...] = (), ) -> None: """Increase a counter on the stats payload.""" self._types.setdefault(counter_name, _MetricType.COUNTER) - self._values[counter_name] = ( - self._values.get(counter_name, 0) + counter_increase_value + counter_values = self._values.setdefault(counter_name, {}) + counter_values[labelvalues] = ( + counter_values.get(labelvalues, 0) + counter_increase_value ) - def set_gauge(self, gauge_name: str, gauge_value: int | float) -> None: + def set_gauge( + self, + gauge_name: str, + gauge_value: int | float, + labelvalues: tuple[str, ...] = (), + ) -> None: """Set a gauge snapshot on the stats payload.""" self._types.setdefault(gauge_name, _MetricType.GAUGE) - self._values[gauge_name] = gauge_value + gauge_values = self._values.setdefault(gauge_name, {}) + gauge_values[labelvalues] = gauge_value def observe_histogram( - self, histogram_name: str, histogram_value: int | float + self, + histogram_name: str, + histogram_value: int | float, + labelvalues: tuple[str, ...] = (), ) -> None: """Record a histogram observation on the stats payload.""" self._types.setdefault(histogram_name, _MetricType.HISTOGRAM) - self._values.setdefault(histogram_name, []).append(histogram_value) + histogram_values = self._values.setdefault(histogram_name, {}) + histogram_values.setdefault(labelvalues, []).append(histogram_value) class OffloadPromMetrics(KVConnectorPromMetrics): @@ -255,7 +281,10 @@ class OffloadPromMetrics(KVConnectorPromMetrics): self._observe_deprecated_metrics = issubclass(spec_cls, CPUOffloadingSpec) self._offloading_metric_defs: dict[str, PromMetricT] = {} - self.offloading_metrics: dict[tuple[int, str], PromMetricT] = {} + # (engine_idx, metric_name, labelvalues) -> metric with bound labels + self.offloading_metrics: dict[ + tuple[int, str, tuple[str, ...]], PromMetricT + ] = {} self._counter_kv_bytes = self._counter_cls( name=_DEPRECATED_TOTAL_BYTES, @@ -301,10 +330,6 @@ class OffloadPromMetrics(KVConnectorPromMetrics): self._offloading_metric_defs[metric_name] = self._create_metric( metric_name, metadata ) - for engine_idx, labelvalues in per_engine_labelvalues.items(): - self.offloading_metrics[(engine_idx, metric_name)] = ( - self._offloading_metric_defs[metric_name].labels(*labelvalues) - ) def _create_metric( self, metric_name: str, metadata: OffloadingMetricMetadata @@ -312,7 +337,7 @@ class OffloadPromMetrics(KVConnectorPromMetrics): kwargs: dict[str, Any] = { "name": metric_name, "documentation": metadata.documentation, - "labelnames": self._labelnames, + "labelnames": self._labelnames + list(metadata.labelnames), } if isinstance(metadata, OffloadingCounterMetadata): metric_cls = self._counter_cls @@ -326,11 +351,37 @@ class OffloadPromMetrics(KVConnectorPromMetrics): raise AssertionError(f"Unknown offloading metric metadata: {metadata}") return metric_cls(**kwargs) + def _get_prometheus_metric( + self, + metric_name: str, + labelvalues: tuple[str, ...], + engine_idx: int, + ) -> PromMetric: + metadata = self._offloading_metric_metadata[metric_name] + if len(labelvalues) != len(metadata.labelnames): + raise AssertionError( + f"Metric {metric_name} expects {len(metadata.labelnames)} labels, " + f"got {len(labelvalues)}" + ) + key = (engine_idx, metric_name, labelvalues) + prom_metric = self.offloading_metrics.get(key) + if prom_metric is None: + engine_labelvalues = self.per_engine_labelvalues[engine_idx] + prom_metric = self._offloading_metric_defs[metric_name].labels( + *(engine_labelvalues + list(labelvalues)) + ) + self.offloading_metrics[key] = prom_metric + return prom_metric + def _increase_counter( - self, metric_name: str, value: int | float, engine_idx: int + self, + metric_name: str, + value: int | float, + labelvalues: tuple[str, ...], + engine_idx: int, ) -> None: - self.offloading_metrics[(engine_idx, metric_name)].inc(value) - if not self._observe_deprecated_metrics: + self._get_prometheus_metric(metric_name, labelvalues, engine_idx).inc(value) + if labelvalues or not self._observe_deprecated_metrics: return # Keep deprecated CPU offload transfer metrics updated during the # transition to flat metric names. @@ -343,15 +394,26 @@ class OffloadPromMetrics(KVConnectorPromMetrics): elif metric_name == _TransferMetricName.STORE_TIME: self.counter_kv_transfer_time[(engine_idx, _TransferType.STORE)].inc(value) - def _set_gauge(self, metric_name: str, value: int | float, engine_idx: int) -> None: - self.offloading_metrics[(engine_idx, metric_name)].set(value) + def _set_gauge( + self, + metric_name: str, + value: int | float, + labelvalues: tuple[str, ...], + engine_idx: int, + ) -> None: + self._get_prometheus_metric(metric_name, labelvalues, engine_idx).set(value) def _observe_histogram( - self, metric_name: str, value: list[int | float], engine_idx: int + self, + metric_name: str, + value: list[int | float], + labelvalues: tuple[str, ...], + engine_idx: int, ) -> None: + prom_metric = self._get_prometheus_metric(metric_name, labelvalues, engine_idx) for observation in value: - self.offloading_metrics[(engine_idx, metric_name)].observe(observation) - if not self._observe_deprecated_metrics: + prom_metric.observe(observation) + if labelvalues or not self._observe_deprecated_metrics: continue # Keep deprecated CPU offload transfer metrics updated during the # transition to flat metric names. @@ -368,20 +430,23 @@ class OffloadPromMetrics(KVConnectorPromMetrics): """Observe transfer statistics.""" metric_types = transfer_stats_data.get(_StatsKey.TYPES, {}) metric_data = transfer_stats_data.get(_StatsKey.DATA, {}) - for key, value in metric_data.items(): + for key, label_value_map in metric_data.items(): type_str = metric_types.get(key) if type_str is None: raise AssertionError(f"Unknown offloading stats key: {key}") assert key in self._offloading_metric_defs - if type_str == _MetricType.COUNTER: - assert isinstance(value, int | float) - self._increase_counter(key, value, engine_idx) - elif type_str == _MetricType.GAUGE: - assert isinstance(value, int | float) - self._set_gauge(key, value, engine_idx) - elif type_str == _MetricType.HISTOGRAM: - assert isinstance(value, list) - assert all(isinstance(v, int | float) for v in value) - self._observe_histogram(key, value, engine_idx) - else: - raise AssertionError(f"Unknown metric type '{type_str}' for key: {key}") + for labelvalues, value in label_value_map.items(): + if type_str == _MetricType.COUNTER: + assert isinstance(value, int | float) + self._increase_counter(key, value, labelvalues, engine_idx) + elif type_str == _MetricType.GAUGE: + assert isinstance(value, int | float) + self._set_gauge(key, value, labelvalues, engine_idx) + elif type_str == _MetricType.HISTOGRAM: + assert isinstance(value, list) + assert all(isinstance(v, int | float) for v in value) + self._observe_histogram(key, value, labelvalues, engine_idx) + else: + raise AssertionError( + f"Unknown metric type '{type_str}' for key: {key}" + ) diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index 2d27c14fe81..904003bdcb0 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -129,6 +129,7 @@ The class provides the following primitives: @dataclass(frozen=True) class OffloadingMetricMetadata: documentation: str + labelnames: tuple[str, ...] = () @dataclass(frozen=True) From 635c38338afe132f9555ebf4a7c8ac7dfb05b2a0 Mon Sep 17 00:00:00 2001 From: Ranran Date: Sun, 21 Jun 2026 13:56:50 -0500 Subject: [PATCH 68/75] [Multimodal] Add Qwen2-VL/Qwen2.5-VL processor-mapped video loader (#45555) Signed-off-by: Ranran Signed-off-by: Ranran Haoran Zhang Signed-off-by: Isotr0py Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Isotr0py --- .../pooling/classify/test_online_vision.py | 4 +- tests/multimodal/test_video.py | 27 ++++++- vllm/multimodal/video.py | 80 +++++++++++++++++++ vllm/transformers_utils/processor.py | 10 +++ 4 files changed, 117 insertions(+), 4 deletions(-) diff --git a/tests/entrypoints/pooling/classify/test_online_vision.py b/tests/entrypoints/pooling/classify/test_online_vision.py index 2776dc8d806..ce60e01ebe3 100644 --- a/tests/entrypoints/pooling/classify/test_online_vision.py +++ b/tests/entrypoints/pooling/classify/test_online_vision.py @@ -25,7 +25,7 @@ def server(): "--runner", "pooling", "--max-model-len", - "5000", + "16384", "--enforce-eager", "--limit-mm-per-prompt", json.dumps({"video": MAXIMUM_VIDEOS}), @@ -143,4 +143,4 @@ def test_chat_video_url_request(server: RemoteOpenAIServer, model_name: str): assert output.model == model_name assert len(output.data) == 1 assert len(output.data[0].probs) == 2 - assert output.usage.prompt_tokens == 4807 + assert output.usage.prompt_tokens == 8993 diff --git a/tests/multimodal/test_video.py b/tests/multimodal/test_video.py index d9f5413b635..694eb392c48 100644 --- a/tests/multimodal/test_video.py +++ b/tests/multimodal/test_video.py @@ -15,6 +15,7 @@ from vllm.multimodal.video import ( DynamicVideoBackend, GLM46VVideoBackend, Molmo2VideoBackend, + Qwen2VLVideoBackend, Qwen3VLVideoBackend, VideoLoader, VideoSourceMetadata, @@ -70,11 +71,12 @@ def test_video_loader_type_doesnt_exist(): @pytest.mark.parametrize( - "model_repo, expected_loader_cls", + "model_repo, expected_loader_cls, hf_sample_kwargs", [ pytest.param( "allenai/Molmo2-4B", Molmo2VideoBackend, + None, marks=pytest.mark.skip( reason="Video processor not aligned, investigate later.", ), @@ -83,23 +85,44 @@ def test_video_loader_type_doesnt_exist(): pytest.param( "zai-org/GLM-4.1V-9B-Thinking", DynamicVideoBackend, + None, id="glm4v", ), pytest.param( "zai-org/GLM-4.6V-Flash", GLM46VVideoBackend, + None, id="glm46v", ), pytest.param( "Qwen/Qwen3-VL-4B-Instruct", Qwen3VLVideoBackend, + None, id="qwen3vl", ), + # Qwen2-VL/Qwen2.5-VL ship no ``video_processor_type`` in their + # preprocessor config, so resolution relies on the model_type -> + # video processor fallback in get_video_processor_cls_name_from_config. + # They also ship no default fps/num_frames, so the HF sampler needs an + # explicit target rate; pass fps=2 to match the loader default. + pytest.param( + "Qwen/Qwen2-VL-7B-Instruct", + Qwen2VLVideoBackend, + {"fps": 2}, + id="qwen2vl", + ), + pytest.param( + "Qwen/Qwen2.5-VL-7B-Instruct", + Qwen2VLVideoBackend, + {"fps": 2}, + id="qwen2_5_vl", + ), ], ) def test_video_processor_from_model_repo( model_repo: str, expected_loader_cls: type, + hf_sample_kwargs: dict[str, int | float] | None, ): """Test that a model repo resolves to the correct video loader backend. @@ -143,7 +166,7 @@ def test_video_processor_from_model_repo( fps=vllm_meta["fps"], duration=vllm_meta["duration"], ) - hf_indices = processor.sample_frames(hf_metadata) + hf_indices = processor.sample_frames(hf_metadata, **(hf_sample_kwargs or {})) vllm_indices = np.array(vllm_meta["frames_indices"]) np.testing.assert_array_equal( hf_indices, diff --git a/vllm/multimodal/video.py b/vllm/multimodal/video.py index bb74f073fbc..4a82dd24e75 100644 --- a/vllm/multimodal/video.py +++ b/vllm/multimodal/video.py @@ -7,6 +7,7 @@ from typing import Any, ClassVar, Literal, NamedTuple, cast import numpy as np import numpy.typing as npt +import torch from vllm.logger import init_logger from vllm.utils.import_utils import PlaceholderModule @@ -653,6 +654,85 @@ class Qwen3VLVideoBackend(VideoBackend): ) +@VIDEO_LOADER_REGISTRY.register( + "qwen2_vl", + video_processor="Qwen2VLVideoProcessor", +) +class Qwen2VLVideoBackend(VideoBackend): + """Qwen2-VL / Qwen2.5-VL fps-based video backend. + + Ports transformers' ``Qwen2VLVideoProcessor.sample_frames`` (fps mode), + shared by Qwen2-VL and Qwen2.5-VL (the latter has no video processor of its + own): sample ``total / original_fps * fps`` frames, clamp to + ``[min_frames, max_frames]`` (4 and 768), floor to a multiple of + ``temporal_patch_size`` (2), and take indices with the exact + ``torch.arange(0, total, total / n)`` call so they match HF byte-for-byte. + + ``num_frames`` is ignored (fps-driven, like the Qwen3-VL loader). The + float32 step can emit an out-of-range tail index (e.g. 451 for a 451-frame + clip); it is clamped to the last valid frame. + """ + + @classmethod + def compute_frames_index_to_sample( + cls, + source: VideoSourceMetadata, + target: VideoTargetMetadata, + **kwargs, + ) -> list[int]: + # Refer to: + # https://github.com/huggingface/transformers/blob/v5.7.0/src/transformers/models/qwen2_vl/video_processing_qwen2_vl.py#L122-L190 + total_frames_num = source.total_frames_num + original_fps = source.original_fps + temporal_patch_size = kwargs.get("temporal_patch_size", 2) + min_frames = kwargs.get("min_frames", 4) + max_frames = kwargs.get("max_frames", 768) + + # vLLM reports original_fps == 0 for clips with unknown/variable fps + # (VFR, malformed, streaming); fail loudly instead of dividing by zero. + if original_fps <= 0: + raise ValueError( + "Qwen2-VL video sampling needs a known source fps, but the " + "container reported 0 (variable or unknown frame rate)." + ) + + max_frames = ( + math.floor(min(max_frames, total_frames_num) / temporal_patch_size) + * temporal_patch_size + ) + n = total_frames_num / original_fps * target.fps + n = min(max(n, min_frames), max_frames, total_frames_num) + n = math.floor(n / temporal_patch_size) * temporal_patch_size + + # ``torch.arange`` matches transformers' float32 index math exactly + # (numpy's float64 diverges by a frame on some inputs); clamp the tail + # because that step can emit an index == total_frames_num. + indices = torch.arange(0, total_frames_num, total_frames_num / n).int() + return torch.clamp(indices, max=total_frames_num - 1).tolist() + + @classmethod + def load_bytes( + cls, + data: bytes, + num_frames: int = -1, + fps: int = 2, + max_duration: int = 300, + frame_recovery: bool = False, + *, + backend: Literal["opencv", "pyav"] = "opencv", + **kwargs, + ) -> tuple[npt.NDArray, dict[str, Any]]: + return super().load_bytes( + data, + num_frames=num_frames, + fps=fps, + max_duration=max_duration, + frame_recovery=frame_recovery, + backend=backend, + **kwargs, + ) + + @VIDEO_LOADER_REGISTRY.register( "opencv_dynamic", video_processor="Glm4vVideoProcessor", diff --git a/vllm/transformers_utils/processor.py b/vllm/transformers_utils/processor.py index 462a6582ed4..fa4c558a739 100644 --- a/vllm/transformers_utils/processor.py +++ b/vllm/transformers_utils/processor.py @@ -18,6 +18,7 @@ from transformers.audio_utils import AudioInput from transformers.feature_extraction_utils import FeatureExtractionMixin from transformers.image_processing_utils import BaseImageProcessor from transformers.image_utils import ImageInput +from transformers.models.auto.video_processing_auto import VIDEO_PROCESSOR_MAPPING_NAMES from transformers.processing_utils import ProcessorMixin from transformers.video_processing_utils import BaseVideoProcessor from transformers.video_utils import VideoInput @@ -169,6 +170,15 @@ def get_video_processor_cls_name_from_config( config = get_hf_file_to_dict(file, processor_name, revision=revision) if config and "video_processor_type" in config: return config["video_processor_type"] + + # Some models ship no explicit ``video_processor_type`` in their + # preprocessor config. Fall back to transformers' ``model_type`` -> video + # processor mapping so these still resolve to their registered loader + # instead of the generic opencv fallback. The mapping is ``None`` for a + # given type when torchvision is unavailable; callers then use opencv. + model_config = get_hf_file_to_dict("config.json", processor_name, revision=revision) + if model_config and "model_type" in model_config: + return VIDEO_PROCESSOR_MAPPING_NAMES.get(model_config["model_type"]) return None From 9c450b102788bb271c1710354db22290ed57f543 Mon Sep 17 00:00:00 2001 From: ZedongLiu <113341356+Zedong-Liu@users.noreply.github.com> Date: Mon, 22 Jun 2026 03:59:40 +0800 Subject: [PATCH 69/75] [Kernel][Bugfix] Fix INT8 per-token-head KV cache rounding in Triton reshape-and-cache (#45361) Signed-off-by: ZedongLiu <113341356+Zedong-Liu@users.noreply.github.com> --- tests/quantization/test_per_token_kv_cache.py | 60 ++++++++++++++++--- .../ops/triton_reshape_and_cache_flash.py | 14 ++++- 2 files changed, 65 insertions(+), 9 deletions(-) diff --git a/tests/quantization/test_per_token_kv_cache.py b/tests/quantization/test_per_token_kv_cache.py index 254e284efb5..b657c77a29a 100644 --- a/tests/quantization/test_per_token_kv_cache.py +++ b/tests/quantization/test_per_token_kv_cache.py @@ -61,8 +61,8 @@ class QuantConfig: quant_max: float quant_min: float kv_quant_mode: KVQuantMode - # INT8 Triton stores truncate; FP8 hardware casts round. - uses_trunc: bool + # INT8 rounds explicitly; FP8 relies on dtype cast rounding. + rounds_before_store: bool INT8_CONFIG = QuantConfig( @@ -71,7 +71,7 @@ INT8_CONFIG = QuantConfig( quant_max=127.0, quant_min=-128.0, kv_quant_mode=KVQuantMode.INT8_PER_TOKEN_HEAD, - uses_trunc=True, + rounds_before_store=True, ) FP8_CONFIG = QuantConfig( cache_dtype=FP8_DTYPE, @@ -79,7 +79,7 @@ FP8_CONFIG = QuantConfig( quant_max=FP8_MAX, quant_min=FP8_MIN, kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD, - uses_trunc=False, + rounds_before_store=False, ) QUANT_CONFIGS = [INT8_CONFIG, FP8_CONFIG] @@ -104,7 +104,7 @@ def _quantize_per_token_head_ref( absmax = data.float().abs().amax(dim=2) # [num_tokens, num_heads] scales = (absmax / cfg.quant_max).clamp(min=1e-6) scaled = data.float() * (1.0 / scales[:, :, None]) - if cfg.uses_trunc: + if cfg.rounds_before_store: q = scaled.round().clamp(cfg.quant_min, cfg.quant_max).to(cfg.cache_dtype) else: q = scaled.clamp(cfg.quant_min, cfg.quant_max).to(cfg.cache_dtype) @@ -255,7 +255,7 @@ def test_per_token_head_round_trip_accuracy( ): """Verify per-token-head round-trip: kernel dequant matches reference. - INT8: Triton truncates on float->int8 store. + INT8: round-to-nearest before int8 store. FP8: hardware cast (clamp then cast). """ from vllm.v1.attention.ops.triton_reshape_and_cache_flash import ( @@ -315,6 +315,52 @@ def test_per_token_head_round_trip_accuracy( ) +@torch.inference_mode() +def test_int8_per_token_head_raw_cache_matches_round_reference(): + """INT8 cache writes should match round-to-nearest quantization exactly.""" + from vllm.v1.attention.ops.triton_reshape_and_cache_flash import ( + triton_reshape_and_cache_flash_per_token_head_quant, + ) + + torch.set_default_device(DEVICE_TYPE) + + head_size = 8 + block_size = 4 + + key = torch.tensor( + [[[-127.0, -2.6, -2.4, -1.6, -1.4, -0.6, -0.4, 127.0]]], + dtype=torch.bfloat16, + ) + value = -key + + key_cache = torch.zeros(1, block_size, 1, head_size, dtype=torch.int8) + value_cache = torch.zeros_like(key_cache) + k_scale_cache = torch.ones(1, block_size, 1, dtype=torch.float32) + v_scale_cache = torch.ones_like(k_scale_cache) + slot_mapping = torch.tensor([2], dtype=torch.long) + + triton_reshape_and_cache_flash_per_token_head_quant( + key, + value, + key_cache, + value_cache, + k_scale_cache, + v_scale_cache, + slot_mapping, + ) + + ref_k_quant, ref_k_scales = _quantize_per_token_head_ref(key, INT8_CONFIG) + ref_v_quant, ref_v_scales = _quantize_per_token_head_ref(value, INT8_CONFIG) + + slot = slot_mapping.item() + blk = slot // block_size + off = slot % block_size + assert torch.equal(key_cache[blk, off], ref_k_quant[0]) + assert torch.equal(value_cache[blk, off], ref_v_quant[0]) + torch.testing.assert_close(k_scale_cache[blk, off], ref_k_scales[0]) + torch.testing.assert_close(v_scale_cache[blk, off], ref_v_scales[0]) + + # =========================================================================== # 4. Negative slot mapping (padding tokens should be skipped) # =========================================================================== @@ -461,7 +507,7 @@ def test_triton_unified_attention_per_token_head_scale( scaled_k = key_cache_bf16.float() / k_scale_cache[:, :, :, None] scaled_v = value_cache_bf16.float() / v_scale_cache[:, :, :, None] - if qcfg.uses_trunc: + if qcfg.rounds_before_store: key_cache_q = ( scaled_k.round().clamp(qcfg.quant_min, qcfg.quant_max).to(qcfg.cache_dtype) ) diff --git a/vllm/v1/attention/ops/triton_reshape_and_cache_flash.py b/vllm/v1/attention/ops/triton_reshape_and_cache_flash.py index 320b7aa597f..fb0c9230551 100644 --- a/vllm/v1/attention/ops/triton_reshape_and_cache_flash.py +++ b/vllm/v1/attention/ops/triton_reshape_and_cache_flash.py @@ -181,6 +181,7 @@ def _reshape_cache_per_token_head( HEAD_SIZE_PADDED: tl.constexpr, # next_power_of_2(max(head_size, head_size_v)) QUANT_MAX: tl.constexpr = 127.0, QUANT_MIN: tl.constexpr = -128.0, + IS_INT_QUANT: tl.constexpr = False, ): tok = tl.program_id(0) head = tl.program_id(1) @@ -211,7 +212,11 @@ def _reshape_cache_per_token_head( k_scale, ) - k_q = tl.clamp(k_h * (1.0 / k_scale), QUANT_MIN, QUANT_MAX) + k_q = k_h * (1.0 / k_scale) + if IS_INT_QUANT: + # Round half away from zero before the int8 store truncates. + k_q = tl.where(k_q >= 0, k_q + 0.5, k_q - 0.5) + k_q = tl.clamp(k_q, QUANT_MIN, QUANT_MAX) tl.store( key_cache_ptr + blk * stride_kc_blk @@ -239,7 +244,11 @@ def _reshape_cache_per_token_head( v_scale, ) - v_q = tl.clamp(v_h * (1.0 / v_scale), QUANT_MIN, QUANT_MAX) + v_q = v_h * (1.0 / v_scale) + if IS_INT_QUANT: + # Round half away from zero before the int8 store truncates. + v_q = tl.where(v_q >= 0, v_q + 0.5, v_q - 0.5) + v_q = tl.clamp(v_q, QUANT_MIN, QUANT_MAX) tl.store( value_cache_ptr + blk * stride_vc_blk @@ -327,6 +336,7 @@ def triton_reshape_and_cache_flash_per_token_head_quant( HEAD_SIZE_PADDED=head_size_padded, QUANT_MAX=quant_max, QUANT_MIN=quant_min, + IS_INT_QUANT=cache_dtype == torch.int8, num_warps=num_warps, ) From 89bd2c14d39075a6109ff188b896b600732dc348 Mon Sep 17 00:00:00 2001 From: Benjamin Chislett Date: Sun, 21 Jun 2026 16:55:26 -0400 Subject: [PATCH 70/75] [Spec Decode] Add Qwen3 architecture support for EAGLE3 (#43132) Signed-off-by: Benjamin Chislett --- tests/models/registry.py | 19 + .../test_speculators_correctness.py | 31 ++ vllm/model_executor/models/qwen3_eagle3.py | 453 ++++++++++++++++++ vllm/model_executor/models/registry.py | 2 + .../configs/speculators/algos.py | 18 +- vllm/v1/spec_decode/llm_base_proposer.py | 2 + 6 files changed, 523 insertions(+), 2 deletions(-) create mode 100644 vllm/model_executor/models/qwen3_eagle3.py diff --git a/tests/models/registry.py b/tests/models/registry.py index ec2c52db567..e865f8efe85 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -1530,6 +1530,16 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = { "Qwen/Qwen3-VL-8B-Instruct", speculative_model="taobao-mnn/Qwen3-VL-8B-Instruct-Eagle3", ), + "Eagle3Qwen3ForCausalLM": _HfExamplesInfo( + "Qwen/Qwen3-8B", + trust_remote_code=True, + speculative_model=( + "inference-optimization/" + "Qwen3-8B-from-Qwen3-8B_regen-speculators.eagle3-qwen3arch-ckpt1" + ), + tokenizer="Qwen/Qwen3-8B", + use_original_num_layers=True, + ), # [PEagle] "PEagleDraftModel": _HfExamplesInfo( "Qwen/Qwen3-8B", @@ -1545,6 +1555,15 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = { tokenizer="Qwen/Qwen3-8B", use_original_num_layers=True, ), + "PeagleQwen3ForCausalLM": _HfExamplesInfo( + "Qwen/Qwen3-8B", + trust_remote_code=True, + speculative_model=( + "inference-optimization/Qwen3-8B-speculators.peagle-qwen3arch-ckpt4" + ), + tokenizer="Qwen/Qwen3-8B", + use_original_num_layers=True, + ), # [MTP] "DeepSeekMTPModel": _HfExamplesInfo( "luccafong/deepseek_mtp_main_random", diff --git a/tests/v1/spec_decode/test_speculators_correctness.py b/tests/v1/spec_decode/test_speculators_correctness.py index e133d9eaf9e..5a92fe00f3f 100644 --- a/tests/v1/spec_decode/test_speculators_correctness.py +++ b/tests/v1/spec_decode/test_speculators_correctness.py @@ -53,9 +53,39 @@ PEAGLE_CONFIG = SpeculatorTestConfig( parallel_drafting=True, ) +QWEN3_EAGLE3_CONFIG = SpeculatorTestConfig( + model_path=( + "inference-optimization/" + "Qwen3-8B-from-Qwen3-8B_regen-speculators.eagle3-qwen3arch-ckpt1" + ), + method="eagle3", + display_name="Qwen3 Eagle3", + expected_gsm8k_accuracy=0.88, + accuracy_rtol=0.05, + expected_acceptance_len=2.67, + acceptance_len_rtol=0.10, + expected_per_pos_acceptance_rates=(0.76, 0.55, 0.36), + per_pos_rtol=0.10, +) + +QWEN3_PEAGLE_CONFIG = SpeculatorTestConfig( + model_path="inference-optimization/Qwen3-8B-speculators.peagle-qwen3arch-ckpt4", + method="eagle3", + display_name="Qwen3 PEagle", + expected_gsm8k_accuracy=0.88, + accuracy_rtol=0.05, + expected_acceptance_len=3.42, + acceptance_len_rtol=0.15, + expected_per_pos_acceptance_rates=(0.78, 0.59, 0.43, 0.29, 0.18, 0.10, 0.05), + per_pos_rtol=0.10, + parallel_drafting=True, +) + SPECULATOR_CONFIGS = [ pytest.param(DFLASH_CONFIG, id="dflash"), pytest.param(PEAGLE_CONFIG, id="peagle"), + pytest.param(QWEN3_EAGLE3_CONFIG, id="qwen3arch_eagle3"), + pytest.param(QWEN3_PEAGLE_CONFIG, id="qwen3arch_peagle"), ] @@ -176,6 +206,7 @@ def test_speculators_correctness(monkeypatch, config): results = evaluate_gsm8k_offline(spec_llm) accuracy = results["accuracy"] + print(f"GSM8K Accuracy: {accuracy:.4f}") accuracy_threshold = config.expected_gsm8k_accuracy * (1 - config.accuracy_rtol) assert accuracy >= accuracy_threshold, ( f"Expected GSM8K accuracy >= {accuracy_threshold:.3f}, got {accuracy:.3f}" diff --git a/vllm/model_executor/models/qwen3_eagle3.py b/vllm/model_executor/models/qwen3_eagle3.py new file mode 100644 index 00000000000..6b03dfcdbdd --- /dev/null +++ b/vllm/model_executor/models/qwen3_eagle3.py @@ -0,0 +1,453 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Iterable + +import torch +import torch.nn as nn +from transformers import Qwen3Config + +from vllm.compilation.decorators import support_torch_compile +from vllm.config import VllmConfig, get_current_vllm_config +from vllm.logger import init_logger +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import QKVParallelLinear, ReplicatedLinear +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from vllm.model_executor.models.qwen3 import Qwen3DecoderLayer, Qwen3ForCausalLM +from vllm.multimodal.inputs import NestedTensors + +from .utils import ( + AutoWeightsLoader, + get_draft_quant_config, + maybe_prefix, + process_eagle_weight, +) + +logger = init_logger(__name__) + + +class Qwen3Eagle3DecoderLayer(Qwen3DecoderLayer): + def __init__( + self, + vllm_config: VllmConfig, + prefix: str = "", + config: Qwen3Config | None = None, + layer_idx: int = 0, + ) -> None: + config = config or vllm_config.model_config.hf_config + cache_config = vllm_config.cache_config + quant_config = get_draft_quant_config(vllm_config) + + super().__init__( + config=config, + cache_config=cache_config, + quant_config=quant_config, + prefix=prefix, + ) + + # First layer uses 2*hidden_size (embeds + hidden_states concatenated) + # 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 + + # 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=qkv_bias, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "qkv_proj"), + ) + + self.hidden_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.layer_idx = layer_idx + + if getattr(config, "norm_before_residual", False): + self._residual_norm = self._norm_before_residual + else: + self._residual_norm = self._norm_after_residual + + def _norm_before_residual( + self, hidden_states: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + hidden_states = self.hidden_norm(hidden_states) + residual = hidden_states + return hidden_states, residual + + def _norm_after_residual( + self, hidden_states: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + residual = hidden_states + hidden_states = self.hidden_norm(hidden_states) + return hidden_states, residual + + def forward( + self, + positions: torch.Tensor, + embeds: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.layer_idx == 0: + # First layer: concatenate embeds with hidden_states + embeds = self.input_layernorm(embeds) + hidden_states, residual = self._residual_norm(hidden_states=hidden_states) + hidden_states = torch.cat([embeds, hidden_states], dim=-1) + else: + # Subsequent layers: process hidden_states and residuals only + hidden_states, residual = self.input_layernorm(hidden_states, residual) + + # Self Attention + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + ) + + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + + # Fully Connected + hidden_states = self.mlp(hidden_states) + + return hidden_states, residual + + +@support_torch_compile( + dynamic_arg_dims={ + "input_ids": 0, + "positions": -1, + "hidden_states": 0, + "input_embeds": 0, + } +) +class Qwen3Eagle3Model(nn.Module): + def __init__( + self, + *, + vllm_config: VllmConfig, + start_layer_id: int = 0, + prefix: str = "", + ) -> None: + super().__init__() + self.config = vllm_config.speculative_config.draft_model_config.hf_config + self.vocab_size = self.config.vocab_size + + # Get drafter's quantization config + self.quant_config = get_draft_quant_config(vllm_config) + + eagle_config = getattr(self.config, "eagle_config", None) or {} + if "use_aux_hidden_state" in eagle_config: + self.use_aux_hidden_state = eagle_config["use_aux_hidden_state"] + else: + self.use_aux_hidden_state = True + self.norm_before_fc = bool( + eagle_config.get( + "norm_before_fc", getattr(self.config, "norm_before_fc", False) + ) + ) + self.fc_input_size = self.config.hidden_size + + current_vllm_config = get_current_vllm_config() + + self.embed_tokens = VocabParallelEmbedding( + self.config.vocab_size, + self.config.hidden_size, + prefix=maybe_prefix(prefix, "embed_tokens"), + ) + + self.layers = nn.ModuleList( + [ + Qwen3Eagle3DecoderLayer( + current_vllm_config, + prefix=maybe_prefix(prefix, f"layers.{layer_idx + start_layer_id}"), + config=self.config, + layer_idx=layer_idx, + ) + for layer_idx in range(self.config.num_hidden_layers) + ] + ) + if self.use_aux_hidden_state: + num_aux_features = getattr(self.config, "num_aux_layers", None) + if num_aux_features is None: + num_aux_features = getattr(self.config, "num_aux_hidden_states", None) + if num_aux_features is None: + aux_ids = getattr( + self.config, "eagle_aux_hidden_state_layer_ids", None + ) or eagle_config.get("eagle_aux_hidden_state_layer_ids") + num_aux_features = len(aux_ids) if aux_ids else 3 + self.num_aux_layers = num_aux_features + target_hidden_size = getattr( + self.config, "target_hidden_size", self.config.hidden_size + ) + self.fc_input_size = target_hidden_size * num_aux_features + if self.norm_before_fc: + self.input_norm = RMSNorm( + self.fc_input_size, + eps=self.config.rms_norm_eps, + ) + else: + self.input_norm = None + + use_fc_norm = getattr(self.config, "fc_norm", False) + if use_fc_norm: + self.fc_norm = nn.ModuleList( + [ + RMSNorm(target_hidden_size, eps=self.config.rms_norm_eps) + for _ in range(num_aux_features) + ] + ) + else: + self.fc_norm = None + + self.fc = ReplicatedLinear( + input_size=self.fc_input_size, + output_size=self.config.hidden_size, + bias=False, + params_dtype=vllm_config.model_config.dtype, + quant_config=self.quant_config, + prefix=maybe_prefix(prefix, "fc"), + return_bias=False, + ) + + self.norm_output = getattr(self.config, "norm_output", False) + self.norm = RMSNorm( + self.config.hidden_size, + eps=self.config.rms_norm_eps, + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + hidden_states: torch.Tensor, + input_embeds: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if input_embeds is None: + input_embeds = self.embed_input_ids(input_ids) + assert hidden_states.shape[-1] == input_embeds.shape[-1] + + residual = None + for layer in self.layers: + hidden_states, residual = layer( + positions=positions, + embeds=input_embeds, + hidden_states=hidden_states, + residual=residual, + ) + hidden_states, hidden_prenorm = self.norm(hidden_states, residual) + + # norm_output variant uses the post-norm hidden states. + aux_output = hidden_states if self.norm_output else hidden_prenorm + + return hidden_states, aux_output + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + for name, loaded_weight in weights: + if "midlayer." in name: + name = name.replace("midlayer.", "layers.0.") + # Remapping the name FP8 kv-scale or zero point. + if "scale" in name or "zero_point" in name: + 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 + name = name.replace(weight_name, param_name) + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + 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 + + +class Eagle3Qwen3ForCausalLM(Qwen3ForCausalLM): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + nn.Module.__init__(self) + self.config = vllm_config.speculative_config.draft_model_config.hf_config + # Ensure draft_vocab_size is set + # default to the base vocab size when absent + if getattr(self.config, "draft_vocab_size", None) is None: + base_vocab_size = getattr(self.config, "vocab_size", None) + self.config.draft_vocab_size = base_vocab_size + target_layer_num = vllm_config.model_config.get_num_layers( + vllm_config.parallel_config + ) + + # Store target layer count in draft config for + # proper layer_types indexing in draft models + self.config.target_layer_count = target_layer_num + self.model = Qwen3Eagle3Model( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "model"), + start_layer_id=target_layer_num, + ) + + logit_scale = getattr(self.config, "logit_scale", 1.0) + self.lm_head = ParallelLMHead( + self.config.draft_vocab_size, + self.config.hidden_size, + quant_config=get_draft_quant_config(vllm_config), + prefix=maybe_prefix(prefix, "lm_head"), + ) + self.logits_processor = LogitsProcessor( + self.config.draft_vocab_size, scale=logit_scale + ) + self.draft_id_to_target_id = nn.Parameter( + torch.zeros(self.config.draft_vocab_size, dtype=torch.long), + 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, self.model.fc_input_size), + persistent=False, + ) + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: NestedTensors | None = None, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + hidden_states: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + return self.model(input_ids, positions, hidden_states, inputs_embeds) + + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor | None: + logits = self.logits_processor(self.lm_head, hidden_states) + if self.draft_id_to_target_id is None: + assert logits.shape[1] == self.config.vocab_size, ( + "Expected logits to have shape " + f"(*, {self.config.vocab_size}), but got {logits.shape}" + ) + return logits + + base = torch.arange(self.config.draft_vocab_size, device=logits.device) + targets = base + self.draft_id_to_target_id + logits_new = logits.new_full( + ( + logits.shape[0], + self.config.vocab_size, + ), + float("-inf"), + ) + logits_new[:, targets] = logits + return logits_new + + def combine_hidden_states( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + if not self.model.use_aux_hidden_state: + return hidden_states + # combine multiple auxiliary hidden states returned by eagle3 + + if self.model.norm_before_fc: + hidden_states = self.model.input_norm(hidden_states) + + # `norm_before_fc` adds a single RMSNorm before the FC layer, whereas `fc_norm` + # applies separate RMSNorms to each chunk of the hidden states. + if self.model.fc_norm is not None: + chunks = hidden_states.chunk(self.model.num_aux_layers, dim=-1) + hidden_states = torch.cat( + [norm(chunk) for norm, chunk in zip(self.model.fc_norm, chunks)], + dim=-1, + ) + + return self.model.fc(hidden_states) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + 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: + includes_embed_tokens = True + model_weights[name] = loaded_weight + process_eagle_weight(self, name) + + 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: + skip_substrs.append("embed_tokens") + if not self.model.use_aux_hidden_state: + skip_substrs.append("fc.") + if not self.model.norm_before_fc: + skip_substrs.append("input_norm.") + loader = AutoWeightsLoader( + self, + skip_prefixes=None, + skip_substrs=skip_substrs, + ) + loader.load_weights(model_weights.items()) diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index f6286439e63..5c023b3f41e 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -617,6 +617,8 @@ _SPECULATIVE_DECODING_MODELS = { "LlamaForCausalLMEagle3": ("llama_eagle3", "Eagle3LlamaForCausalLM"), "Eagle3Qwen2_5vlForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"), "Eagle3Qwen3vlForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"), + "Eagle3Qwen3ForCausalLM": ("qwen3_eagle3", "Eagle3Qwen3ForCausalLM"), + "PeagleQwen3ForCausalLM": ("qwen3_eagle3", "Eagle3Qwen3ForCausalLM"), "EagleMistralForCausalLM": ("mistral_eagle", "EagleMistralForCausalLM"), "EagleMistralLarge3ForCausalLM": ( "mistral_large_3_eagle", diff --git a/vllm/transformers_utils/configs/speculators/algos.py b/vllm/transformers_utils/configs/speculators/algos.py index 650f09c39fb..0dc3ccce089 100644 --- a/vllm/transformers_utils/configs/speculators/algos.py +++ b/vllm/transformers_utils/configs/speculators/algos.py @@ -36,7 +36,14 @@ def update_eagle3(config_dict: dict, pre_trained_config: dict) -> None: "norm_before_residual", True ) pre_trained_config["norm_before_fc"] = config_dict.get("norm_before_fc", False) - pre_trained_config["architectures"] = ["Eagle3LlamaForCausalLM"] + eagle3_arch_map = { + "qwen3": "Eagle3Qwen3ForCausalLM", + "llama": "Eagle3LlamaForCausalLM", + } + model_type = pre_trained_config.get("model_type", "llama") + if model_type not in eagle3_arch_map: + raise ValueError(f"Unsupported model_type {model_type} for Eagle3 speculator") + pre_trained_config["architectures"] = [eagle3_arch_map[model_type]] if config_dict.get("eagle_aux_hidden_state_layer_ids"): pre_trained_config["eagle_aux_hidden_state_layer_ids"] = config_dict[ "eagle_aux_hidden_state_layer_ids" @@ -59,7 +66,6 @@ def update_peagle(config_dict: dict, pre_trained_config: dict) -> None: - eagle_aux_hidden_state_layer_ids: Layer indices from the target model whose intermediate hidden states are used as auxiliary inputs """ - pre_trained_config["architectures"] = ["PeagleLlamaForCausalLM"] pre_trained_config["draft_vocab_size"] = config_dict.get("draft_vocab_size") if config_dict.get("target_hidden_size") is not None: pre_trained_config["target_hidden_size"] = config_dict["target_hidden_size"] @@ -67,6 +73,14 @@ def update_peagle(config_dict: dict, pre_trained_config: dict) -> None: "norm_before_residual", False ) pre_trained_config["norm_before_fc"] = config_dict.get("norm_before_fc", False) + peagle_arch_map = { + "qwen3": "PeagleQwen3ForCausalLM", + "llama": "PeagleLlamaForCausalLM", + } + model_type = pre_trained_config.get("model_type", "llama") + if model_type not in peagle_arch_map: + raise ValueError(f"Unsupported model_type {model_type} for PEagle speculator") + pre_trained_config["architectures"] = [peagle_arch_map[model_type]] pre_trained_config["pard_token"] = config_dict["mask_token_id"] if config_dict.get("eagle_aux_hidden_state_layer_ids"): pre_trained_config["eagle_aux_hidden_state_layer_ids"] = config_dict[ diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index bdc10313f4a..9f46cbd2423 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -24,6 +24,7 @@ from vllm.model_executor.models.deepseek_eagle3 import Eagle3DeepseekV2ForCausal from vllm.model_executor.models.interfaces import SupportsMultiModal from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM from vllm.model_executor.models.qwen3_dflash import DFlashQwen3ForCausalLM +from vllm.model_executor.models.qwen3_eagle3 import Eagle3Qwen3ForCausalLM from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.platforms import current_platform from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d @@ -473,6 +474,7 @@ class SpecDecodeBaseProposer: Eagle3LlamaForCausalLM, Eagle3DeepseekV2ForCausalLM, DFlashQwen3ForCausalLM, + Eagle3Qwen3ForCausalLM, ), ) target_hidden_states = self.model.combine_hidden_states( From 12fe2a9aac8e0284ff1dfbd53857f7e6f7f50da1 Mon Sep 17 00:00:00 2001 From: Ting SUN Date: Mon, 22 Jun 2026 04:31:23 +0700 Subject: [PATCH 71/75] [Bugfix][Qwen3-VL] Fix multi-video crash with list-valued fps/num_frames (#46305) Signed-off-by: Ting Sun --- .../multimodal/processing/test_qwen3_vl.py | 46 +++++++++++++++++++ vllm/model_executor/models/qwen3_vl.py | 15 +++++- 2 files changed, 59 insertions(+), 2 deletions(-) diff --git a/tests/models/multimodal/processing/test_qwen3_vl.py b/tests/models/multimodal/processing/test_qwen3_vl.py index 9155fde5033..a1ab94ca43a 100644 --- a/tests/models/multimodal/processing/test_qwen3_vl.py +++ b/tests/models/multimodal/processing/test_qwen3_vl.py @@ -138,3 +138,49 @@ def test_processor_multi_video( assert video_phs[i].offset >= prev_end, ( f"Placeholder {i} overlaps with placeholder {i - 1}" ) + + +@pytest.mark.parametrize("model_id", [MODEL_ID]) +@pytest.mark.parametrize( + "hf_mm_kwargs", + [{"num_frames": [8, 16]}, {"fps": [2.0, 4.0]}], +) +def test_processor_multi_video_list_kwargs( + model_id: str, + hf_mm_kwargs: dict[str, Any], +) -> None: + """Regression test: a multi-video request with list-valued per-video + ``mm_processor_kwargs`` (one ``fps``/``num_frames`` per video) must not + crash. + + Before the fix, ``_call_hf_processor`` copied the whole kwargs to every + video without slicing, so ``_get_video_second_idx`` received the list + where a scalar was expected and raised ``TypeError``. + """ + ctx = build_model_context( + model_id, + limit_mm_per_prompt={"image": 0, "video": 2}, + ) + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) + + prompt = ( + "<|vision_start|><|video_pad|><|vision_end|>" + "<|vision_start|><|video_pad|><|vision_end|>" + ) + mm_data = { + "video": [ + _build_video_mm_data(num_frames=16)["video"][0], + _build_video_mm_data(num_frames=32)["video"][0], + ] + } + + processed = processor( + prompt, + mm_items=processor.info.parse_mm_data(mm_data), + hf_processor_mm_kwargs=hf_mm_kwargs, + ) + + video_phs = processed["mm_placeholders"].get("video", []) + assert len(video_phs) == 2, ( + f"Expected exactly 2 video placeholders, got {len(video_phs)}" + ) diff --git a/vllm/model_executor/models/qwen3_vl.py b/vllm/model_executor/models/qwen3_vl.py index 3183a23ffde..0e6ddef7f36 100644 --- a/vllm/model_executor/models/qwen3_vl.py +++ b/vllm/model_executor/models/qwen3_vl.py @@ -1271,7 +1271,7 @@ class Qwen3VLMultiModalProcessor(BaseMultiModalProcessor[Qwen3VLProcessingInfo]) vision_end_token_id = hf_config.vision_end_token_id video_token_id = hf_config.video_token_id - for item in videos: + for item_idx, item in enumerate(videos): video_array, metadata = item # NOTE: @JJJYmmm new attr metadata.frames_indices indicates @@ -1282,6 +1282,12 @@ class Qwen3VLMultiModalProcessor(BaseMultiModalProcessor[Qwen3VLProcessingInfo]) # NOTE: a copy of is created to update do_sample_frames, # otherwise mm_hash for the object will be incorrect. video_mm_kwargs = dict(**mm_kwargs) + sampled_fps = video_mm_kwargs.get("fps") + if is_list_of(sampled_fps, float): + video_mm_kwargs["fps"] = sampled_fps[item_idx] + sampled_num_frames = video_mm_kwargs.get("num_frames") + if is_list_of(sampled_num_frames, int): + video_mm_kwargs["num_frames"] = sampled_num_frames[item_idx] if "do_sample_frames" not in video_mm_kwargs: # qwen_vl_utils already has "do_sample_frames" in # mm_kwargs, don't overwrite it. @@ -1363,10 +1369,15 @@ class Qwen3VLMultiModalProcessor(BaseMultiModalProcessor[Qwen3VLProcessingInfo]) else: video_outputs = dict() + # fps/num_frames are video-only kwargs already consumed by the loop; + # exclude them so the text/image processor call below never gets a list. + non_video_mm_kwargs = { + k: v for k, v in mm_kwargs.items() if k not in ("fps", "num_frames") + } processed_outputs = super()._call_hf_processor( prompt=prompt, mm_data=mm_data, - mm_kwargs=mm_kwargs, + mm_kwargs=non_video_mm_kwargs, tok_kwargs=tok_kwargs, ) From 50241602fd7b672751dfd9c034b806640255e599 Mon Sep 17 00:00:00 2001 From: Matt <156021403+mawong-amd@users.noreply.github.com> Date: Sun, 21 Jun 2026 16:45:37 -0500 Subject: [PATCH 72/75] [Hardware][AMD][CI] Fix gfx942 Kernels MoE test group (#46298) Signed-off-by: Matthew Wong --- .buildkite/test-amd.yaml | 9 +++++---- .buildkite/test_areas/kernels.yaml | 16 ++++++++++++++++ tests/kernels/moe/test_ocp_mx_moe.py | 16 ++++++++-------- 3 files changed, 29 insertions(+), 12 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 5550c0a0c18..03b94f6e403 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -1627,10 +1627,11 @@ steps: - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py - label: Kernels MoE Test %N # TBD - timeout_in_minutes: 180 + timeout_in_minutes: 50 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] agent_pool: mi300_1 - parallelism: 4 + optional: true + parallelism: 5 working_dir: "/vllm-workspace/tests" source_file_dependencies: - csrc/quantization/cutlass_w8a8/moe/ @@ -3082,10 +3083,10 @@ steps: - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - label: Kernels MoE Test %N # TBD - timeout_in_minutes: 180 + timeout_in_minutes: 50 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] agent_pool: mi355_1 - parallelism: 4 + parallelism: 5 working_dir: "/vllm-workspace/tests" source_file_dependencies: - csrc/quantization/cutlass_w8a8/moe/ diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index ebcb95a9d82..4c1117440a4 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -128,6 +128,22 @@ steps: - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT parallelism: 5 + mirror: + amd: + device: mi325_1 + timeout_in_minutes: 50 + source_file_dependencies: + - csrc/quantization/cutlass_w8a8/moe/ + - csrc/moe/ + - tests/kernels/moe + - vllm/model_executor/layers/fused_moe/ + - vllm/distributed/device_communicators/ + - vllm/envs.py + - vllm/config + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + depends_on: + - image-build-amd - label: Kernels Mamba Test key: kernels-mamba-test diff --git a/tests/kernels/moe/test_ocp_mx_moe.py b/tests/kernels/moe/test_ocp_mx_moe.py index 5c52c8af6a8..fb5ae527b2c 100644 --- a/tests/kernels/moe/test_ocp_mx_moe.py +++ b/tests/kernels/moe/test_ocp_mx_moe.py @@ -9,6 +9,7 @@ import pytest import torch from packaging import version +from vllm._aiter_ops import is_aiter_found from vllm.platforms import current_platform from vllm.utils.flashinfer import has_flashinfer @@ -31,17 +32,15 @@ HOPPER_MXFP4_BF16_AVAILABLE = ( # ROCm platform and dependencies ROCM_AVAILABLE = current_platform.is_rocm() ROCM_TRITON_KERNELS_AVAILABLE = False -ROCM_AITER_AVAILABLE = False +ROCM_AITER_AVAILABLE = is_aiter_found() ROCM_GFX950 = False if ROCM_AVAILABLE: - from vllm._aiter_ops import rocm_aiter_ops from vllm.platforms.rocm import on_gfx950 from vllm.utils.import_utils import has_triton_kernels ROCM_TRITON_KERNELS_AVAILABLE = has_triton_kernels() ROCM_GFX950 = on_gfx950() - ROCM_AITER_AVAILABLE = rocm_aiter_ops.is_enabled() if ROCM_AITER_AVAILABLE: from aiter.ops.triton.moe.quant_moe import upcast_from_mxfp @@ -83,7 +82,7 @@ def enable_pickle(monkeypatch): [ ModelCase("fxmarty/qwen_1.5-moe-a2.7b-mxfp4", tp=2), ModelCase("fxmarty/deepseek_r1_3_layers_mxfp4", tp=8), - ModelCase("fxmarty/Llama-4-Scout-17B-16E-Instruct-2-layers-mxfp4", tp=1), + ModelCase("mawong-amd/Llama-4-Scout-17B-16E-Instruct-2-layers-mxfp4", tp=1), ModelCase("fxmarty/Llama-3.1-70B-Instruct-2-layers-mxfp6", tp=1), ModelCase("fxmarty/Llama-3.1-70B-Instruct-2-layers-mxfp6", tp=4), ], @@ -102,6 +101,7 @@ def test_mxfp4_loading_and_execution_moe(vllm_runner, model_case: ModelCase): tensor_parallel_size=model_case.tp, load_format="dummy", compilation_config={"cudagraph_capture_sizes": [16]}, + gpu_memory_utilization=0.8, # mxfp6 models use more scratch space ) as llm: # Disabled as check_model is broken: https://github.com/vllm-project/vllm/pull/18465#issuecomment-3329880562 # def check_model(model): @@ -1267,7 +1267,7 @@ def test_rocm_mxfp4_moe_oracle( This test validates that the oracle functions work end-to-end: - select_mxfp4_moe_backend() selects a valid backend - - convert_to_mxfp4_moe_kernel_format() converts weights without error + - convert_gpt_oss_weight_to_mxfp4_moe_kernel_format() converts weights without error - make_mxfp4_moe_quant_config() builds a valid quant config - make_mxfp4_moe_kernel() creates a kernel that runs without error - The kernel output is within accuracy tolerance of reference @@ -1287,7 +1287,7 @@ def test_rocm_mxfp4_moe_oracle( from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( Mxfp4MoeBackend, backend_to_kernel_cls, - convert_to_mxfp4_moe_kernel_format, + convert_gpt_oss_weight_to_mxfp4_moe_kernel_format, make_mxfp4_moe_kernel, make_mxfp4_moe_quant_config, ) @@ -1387,7 +1387,7 @@ def test_rocm_mxfp4_moe_oracle( # Convert weights using oracle w13_conv, w2_conv, w13_scale_conv, w2_scale_conv, w13_bias_conv, w2_bias_conv = ( - convert_to_mxfp4_moe_kernel_format( + convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( mxfp4_backend=backend, layer=layer, # type: ignore[arg-type] w13_weight=w13_quant, @@ -1423,7 +1423,7 @@ def test_rocm_mxfp4_moe_oracle( mxfp4_backend=backend, experts_cls=experts_cls, routing_tables=None, - shared_experts=None, + layer=None, ) # Create inputs From 13b83d77ad21cb9351417ad4c80b2427e2e21f4f Mon Sep 17 00:00:00 2001 From: Charlie Fu Date: Sun, 21 Jun 2026 16:53:11 -0500 Subject: [PATCH 73/75] [ROCm][CI] skip test_double_aiter_rms_quant_fusion (#45967) Signed-off-by: charlifu Co-authored-by: Andreas Karatzas --- .buildkite/test_areas/pytorch.yaml | 6 ++++++ tests/compile/passes/test_double_aiter_rms_quant_fusion.py | 7 +++---- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/.buildkite/test_areas/pytorch.yaml b/.buildkite/test_areas/pytorch.yaml index 6866d5e3695..a33c7f48016 100644 --- a/.buildkite/test_areas/pytorch.yaml +++ b/.buildkite/test_areas/pytorch.yaml @@ -107,6 +107,12 @@ steps: - tests/compile/passes commands: - pytest -s -v compile/passes --ignore compile/passes/distributed + mirror: + amd: + device: mi300_1 + timeout_in_minutes: 180 + depends_on: + - image-build-amd - label: PyTorch Fullgraph Smoke Test key: pytorch-fullgraph-smoke-test diff --git a/tests/compile/passes/test_double_aiter_rms_quant_fusion.py b/tests/compile/passes/test_double_aiter_rms_quant_fusion.py index 161c956548a..6a620d11a49 100644 --- a/tests/compile/passes/test_double_aiter_rms_quant_fusion.py +++ b/tests/compile/passes/test_double_aiter_rms_quant_fusion.py @@ -22,7 +22,7 @@ import torch import vllm.config from tests.compile.backend import TestBackend -from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops +from vllm._aiter_ops import rocm_aiter_ops from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass from vllm.config import ( @@ -83,9 +83,8 @@ class _ViewDoubleQuantModel(torch.nn.Module): [_NoViewDoubleQuantModel, _ViewDoubleQuantModel], ids=["no_view", "with_view"], ) -@pytest.mark.skipif( - not is_aiter_found_and_supported(), - reason="Only test on ROCm with AITER installed and supported", +@pytest.mark.skip( + reason="Skipping for now because pytorch compiler removes one the two quant ops" ) def test_double_aiter_rms_fp8_group_quant_fusion( model_cls: type[torch.nn.Module], From 4f0d0049a0a2a188fcfaa2d07317d629b79b81d9 Mon Sep 17 00:00:00 2001 From: Matt <156021403+mawong-amd@users.noreply.github.com> Date: Sun, 21 Jun 2026 17:10:51 -0500 Subject: [PATCH 74/75] [Hardware][AMD][CI] Fix Kernels Attention test groups (#46080) Signed-off-by: Matthew Wong --- .buildkite/test-amd.yaml | 8 ++--- .buildkite/test_areas/kernels.yaml | 14 +++++++++ .../attention/test_attention_selector.py | 29 +++++++++++-------- .../kernels/attention/test_prefix_prefill.py | 26 +++++++++++------ .../attention/test_rocm_triton_attn_dsv4.py | 22 +++++++++----- .../test_triton_unified_attention.py | 6 +--- 6 files changed, 68 insertions(+), 37 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 03b94f6e403..191537276e4 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -1594,9 +1594,10 @@ steps: #---------------------------------------------------------- mi300 · kernels ----------------------------------------------------------# - label: Kernels Attention Test %N # TBD - timeout_in_minutes: 180 + timeout_in_minutes: 55 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] agent_pool: mi300_1 + optional: true parallelism: 2 working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -3041,7 +3042,7 @@ steps: #---------------------------------------------------------- mi355 · kernels ----------------------------------------------------------# - label: Kernels (B200-MI355) # TBD - timeout_in_minutes: 180 + timeout_in_minutes: 15 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] agent_pool: mi355_1 working_dir: "/vllm-workspace/" @@ -3065,11 +3066,10 @@ steps: - pytest -v -s tests/kernels/attention/test_attention_selector.py - label: Kernels Attention Test %N # TBD - timeout_in_minutes: 180 + timeout_in_minutes: 60 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] agent_pool: mi355_1 parallelism: 2 - optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - csrc/attention/ diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 4c1117440a4..4953b4d441c 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -74,6 +74,20 @@ steps: commands: - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT parallelism: 2 + mirror: + amd: + device: mi325_1 + timeout_in_minutes: 55 + depends_on: + - image-build-amd + source_file_dependencies: + - csrc/attention/ + - vllm/v1/attention + - vllm/model_executor/layers/attention + - tests/kernels/attention + - vllm/_aiter_ops.py + - vllm/envs.py + - vllm/platforms/rocm.py - label: Kernels Attention DiffKV Test (H100) key: kernels-attention-diffkv-test-h100 diff --git a/tests/kernels/attention/test_attention_selector.py b/tests/kernels/attention/test_attention_selector.py index db4dcc8a636..447b502293b 100644 --- a/tests/kernels/attention/test_attention_selector.py +++ b/tests/kernels/attention/test_attention_selector.py @@ -15,16 +15,14 @@ from vllm.config import ( from vllm.platforms import current_platform from vllm.platforms.cpu import CpuPlatform -# CudaPlatform and RocmPlatform import their respective compiled C extensions -# at module level, raising ModuleNotFoundError on incompatible builds. -try: +if current_platform.is_cuda(): from vllm.platforms.cuda import CudaPlatform -except (ImportError, ModuleNotFoundError): +else: CudaPlatform = None -try: +if current_platform.is_rocm(): from vllm.platforms.rocm import RocmPlatform -except (ImportError, ModuleNotFoundError): +else: RocmPlatform = None from vllm.v1.attention.backends.registry import AttentionBackendEnum @@ -434,9 +432,15 @@ def test_per_head_quant_scales_backend_selection( [ ("FLASH_ATTN", True, True), # FlashAttn supports non-causal ("FLASH_ATTN", False, True), # FlashAttn also works with causal - ("FLASHINFER", True, False), # FlashInfer does not support non-causal - ("FLASHINFER", False, True), # FlashInfer works with causal - ], + ] + + ( + [ + ("FLASHINFER", True, False), # FlashInfer does not support non-causal + ("FLASHINFER", False, True), # FlashInfer works with causal + ] + if CudaPlatform is not None + else [] + ), ) def test_non_causal_backend_selection( backend_name: str, use_non_causal: bool, should_succeed: bool @@ -459,11 +463,12 @@ def test_non_causal_backend_selection( attention_config=attention_config, cache_config=cache_config ) - if CudaPlatform is None: - pytest.skip("CudaPlatform not available") + platform = CudaPlatform or RocmPlatform + if platform is None: + pytest.skip("CudaPlatform and RocmPlatform are not available") with ( set_current_vllm_config(vllm_config), - patch("vllm.platforms.current_platform", CudaPlatform()), + patch("vllm.platforms.current_platform", platform()), ): if should_succeed: backend = get_attn_backend( diff --git a/tests/kernels/attention/test_prefix_prefill.py b/tests/kernels/attention/test_prefix_prefill.py index de63b4548f2..f1c591fb671 100644 --- a/tests/kernels/attention/test_prefix_prefill.py +++ b/tests/kernels/attention/test_prefix_prefill.py @@ -5,10 +5,12 @@ import math import random import time from collections.abc import Callable +from contextlib import nullcontext import pytest import torch import torch.nn.functional as F +from torch.nn.attention import SDPBackend, sdpa_kernel from vllm.platforms import current_platform from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed @@ -557,15 +559,21 @@ def test_contexted_kv_attention_alibi( query_len, seq_len, alibi_slopes, device, dtype ) - # Compute attention - out = F.scaled_dot_product_attention( - q_sdpa, - k_sdpa, - v_sdpa, - attn_mask=alibi_mask, - dropout_p=0.0, - scale=scale, - ) + # Compute attention. On ROCm we force use of the Math SDPA backend rather than + # the Flash or Mem-Efficient backends for increased numerical accuracy + if current_platform.is_rocm(): + sdpa_context = sdpa_kernel(SDPBackend.MATH) + else: + sdpa_context = nullcontext() + with sdpa_context: + out = F.scaled_dot_product_attention( + q_sdpa, + k_sdpa, + v_sdpa, + attn_mask=alibi_mask, + dropout_p=0.0, + scale=scale, + ) # Reshape output back to [query_len, num_heads, head_size] out = out.view(num_heads, query_len, head_size).permute(1, 0, 2) diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index daf73b82e61..e00726f64d8 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -90,7 +90,9 @@ def _ref_sparse_prefill_ragged( return out.to(torch.bfloat16) -def _pack_fp8_ds_mla_cache(kv: torch.Tensor, block_size: int) -> torch.Tensor: +def _pack_fp8_ds_mla_cache( + kv: torch.Tensor, block_size: int, is_extra: bool = False +) -> torch.Tensor: assert kv.shape[-1] == HEAD_DIM num_tokens = kv.shape[0] num_blocks = (num_tokens + block_size - 1) // block_size @@ -101,7 +103,9 @@ def _pack_fp8_ds_mla_cache(kv: torch.Tensor, block_size: int) -> torch.Tensor: ) cache_flat = cache.view(torch.uint8).flatten() kv_nope_fp8 = ( - kv[:, :NOPE_HEAD_DIM].to(current_platform.fp8_dtype()).view(torch.uint8) + kv[:, :NOPE_HEAD_DIM] + .to(torch.float8_e4m3fn if is_extra else current_platform.fp8_dtype()) + .view(torch.uint8) ) kv_rope_u8 = kv[:, NOPE_HEAD_DIM:].contiguous().view(torch.uint8) @@ -120,7 +124,7 @@ def _pack_fp8_ds_mla_cache(kv: torch.Tensor, block_size: int) -> torch.Tensor: def _read_fp8_ds_mla_cache( - cache: torch.Tensor, slot: int, block_size: int + cache: torch.Tensor, slot: int, block_size: int, is_extra: bool = False ) -> torch.Tensor: cache_flat = cache.view(torch.uint8).flatten() block_idx = slot // block_size @@ -129,7 +133,9 @@ def _read_fp8_ds_mla_cache( token_base = block_base + pos * 576 nope_u8 = cache_flat[token_base : token_base + NOPE_HEAD_DIM] - nope = nope_u8.view(current_platform.fp8_dtype()).to(torch.float32) + nope = nope_u8.view( + torch.float8_e4m3fn if is_extra else current_platform.fp8_dtype() + ).to(torch.float32) rope_u8 = cache_flat[ token_base + NOPE_HEAD_DIM : token_base + NOPE_HEAD_DIM + ROPE_HEAD_DIM * 2 ] @@ -157,7 +163,9 @@ def _ref_sparse_decode_ragged( ] if extra_cache is not None and extra_rows is not None: row_kv.extend( - _read_fp8_ds_mla_cache(extra_cache, int(slot), block_size) + _read_fp8_ds_mla_cache( + extra_cache, int(slot), block_size, is_extra=True + ) for slot in extra_rows[query_idx] ) @@ -326,7 +334,7 @@ def test_sparse_attn_decode_ragged_kernel() -> None: main_kv = torch.randn(6, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125 extra_kv = torch.randn(5, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125 main_cache = _pack_fp8_ds_mla_cache(main_kv, block_size) - extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size) + extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size, is_extra=True) main_indices = torch.tensor([0, 2, 4, 1], dtype=torch.int32, device=device) main_indptr = torch.tensor([0, 2, 4], dtype=torch.int32, device=device) extra_indices = torch.tensor([1, 3, 0], dtype=torch.int32, device=device) @@ -477,7 +485,7 @@ def test_sparse_attn_decode_split_k_kernel( rows = [[1, 3, 0, 5, 2, 4], [3, 0, 6]] extra_kv = torch.randn(7, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125 extra_rows = rows - extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size) + extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size, is_extra=True) extra_indices, extra_indptr = _ragged_from_rows(rows, device) attn_sink = ( diff --git a/tests/kernels/attention/test_triton_unified_attention.py b/tests/kernels/attention/test_triton_unified_attention.py index 6440ba3156e..d3435ea665d 100644 --- a/tests/kernels/attention/test_triton_unified_attention.py +++ b/tests/kernels/attention/test_triton_unified_attention.py @@ -18,11 +18,7 @@ HEAD_SIZES = [128, 256] BLOCK_SIZES = [16] DTYPES = [torch.bfloat16] -QDTYPES = ( - [None, torch.float8_e4m3fn] - if not current_platform.is_rocm() - else [None, torch.float8_e4m3fnuz] -) +QDTYPES = [None, current_platform.fp8_dtype()] FP8_DTYPE = current_platform.fp8_dtype() # one value large enough to test overflow in index calculation. From a19ff2218a79a99fcd9ebf3a2cf202c9f7eeb9f9 Mon Sep 17 00:00:00 2001 From: Matt <156021403+mawong-amd@users.noreply.github.com> Date: Sun, 21 Jun 2026 17:40:02 -0500 Subject: [PATCH 75/75] [Hardware][AMD][CI] Fix Spec Decode Eagle test group (#46018) Signed-off-by: Matthew Wong --- .buildkite/test-amd.yaml | 3 ++- .buildkite/test_areas/spec_decode.yaml | 14 ++++++++++++++ tests/v1/e2e/spec_decode/test_spec_decode.py | 2 +- 3 files changed, 17 insertions(+), 2 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 191537276e4..bc0687db3a5 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -2122,9 +2122,10 @@ steps: - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference" - label: Spec Decode Eagle # TBD - timeout_in_minutes: 180 + timeout_in_minutes: 45 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] agent_pool: mi300_1 + optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/v1/spec_decode/ diff --git a/.buildkite/test_areas/spec_decode.yaml b/.buildkite/test_areas/spec_decode.yaml index bc73a53a359..6e532eddc71 100644 --- a/.buildkite/test_areas/spec_decode.yaml +++ b/.buildkite/test_areas/spec_decode.yaml @@ -12,6 +12,20 @@ steps: - tests/v1/e2e/spec_decode/ commands: - pytest -v -s v1/e2e/spec_decode -k "eagle_correctness" + mirror: + amd: + device: mi325_1 + timeout_in_minutes: 45 + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/v1/spec_decode/ + - vllm/v1/worker/gpu/spec_decode/ + - vllm/model_executor/model_loader/ + - vllm/v1/sample/ + - vllm/model_executor/layers/ + - tests/v1/e2e/spec_decode/ + - vllm/platforms/rocm.py - label: Spec Decode Eagle Nightly B200 key: spec-decode-eagle-nightly-b200 diff --git a/tests/v1/e2e/spec_decode/test_spec_decode.py b/tests/v1/e2e/spec_decode/test_spec_decode.py index 06e8b3bf0e3..532a3e8d6a7 100644 --- a/tests/v1/e2e/spec_decode/test_spec_decode.py +++ b/tests/v1/e2e/spec_decode/test_spec_decode.py @@ -425,7 +425,7 @@ def _run_eagle_correctness( if "deepseek" in model_setup[1].lower(): m.setenv("VLLM_ROCM_USE_AITER", "1") m.delenv("VLLM_MLA_DISABLE", raising=False) - attention_config = {"backend": "TRITON_MLA"} + attention_config = {"backend": "ROCM_AITER_MLA"} else: m.setenv("VLLM_ROCM_USE_AITER", "1")